internal/mcp: add tool and schema options

Implement variadic options for building tools and schemas, to allow
easier customization of tool input schemas.

Also, make fields required, unless they are marked as "omitempty".

Change-Id: I7a29258fae3ec5e7ea4906ed30ed6750979de962
Reviewed-on: https://go-review.googlesource.com/c/tools/+/668955
Reviewed-by: Jonathan Amsterdam <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>

Author: findleyr

internal/mcp/examples/hello/main.go

@@ -17,7 +17,7 @@ import (
 var httpAddr = flag.String("http", "", "if set, use SSE HTTP at this address, instead of stdin/stdout")
 
 type SayHiParams struct {
-	Name string `json:"name" mcp:"the name to say hi to"`
+	Name string `json:"name"`
 }
 
 func SayHi(ctx context.Context, cc *mcp.ClientConnection, params *SayHiParams) ([]mcp.Content, error) {
@@ -30,7 +30,9 @@ func main() {
 	flag.Parse()
 
 	server := mcp.NewServer("greeter", "v0.0.1", nil)
-	server.AddTools(mcp.MakeTool("greet", "say hi", SayHi))
+	server.AddTools(mcp.MakeTool("greet", "say hi", SayHi, mcp.Input(
+		mcp.Property("name", mcp.Description("the name to say hi to")),
+	)))
 
 	if *httpAddr != "" {
 		handler := mcp.NewSSEHandler(func(*http.Request) *mcp.Server {

internal/mcp/internal/jsonschema/infer.go

@@ -9,6 +9,7 @@ package jsonschema
 import (
 	"fmt"
 	"reflect"
+	"slices"
 	"strings"
 )
 
@@ -19,6 +20,21 @@ func For[T any]() (*Schema, error) {
 	return ForType(reflect.TypeFor[T]())
 }
 
+// ForType constructs a JSON schema object for the given type.
+// It translates Go types into compatible JSON schema types, as follows:
+//   - strings have schema type "string"
+//   - bools have schema type "boolean"
+//   - signed and unsigned integer types have schema type "integer"
+//   - floating point types have schema type "number"
+//   - slices and arrays have schema type "array", and a corresponding schema
+//     for items
+//   - maps with string key have schema type "object", and corresponding
+//     schema for additionalProperties
+//   - structs have schema type "object", and disallow additionalProperties.
+//     Their properties are derived from exported struct fields, using the
+//     struct field json name. Fields that are marked "omitempty" are
+//     considered optional; all other fields become required properties.
+//
 // It returns an error if t contains (possibly recursively) any of the following Go
 // types, as they are incompatible with the JSON schema spec.
 //   - maps with key other than 'string'
@@ -75,6 +91,10 @@ func typeSchema(t reflect.Type, seen map[reflect.Type]*Schema) (*Schema, error)
 		if err != nil {
 			return nil, fmt.Errorf("computing element schema: %v", err)
 		}
+		if t.Kind() == reflect.Array {
+			s.MinItems = Ptr(float64(t.Len()))
+			s.MaxItems = Ptr(float64(t.Len()))
+		}
 
 	case reflect.String:
 		s.Type = "string"
@@ -86,8 +106,8 @@ func typeSchema(t reflect.Type, seen map[reflect.Type]*Schema) (*Schema, error)
 
 		for i := range t.NumField() {
 			field := t.Field(i)
-			name, ok := jsonName(field)
-			if !ok {
+			name, required, include := parseField(field)
+			if !include {
 				continue
 			}
 			if s.Properties == nil {
@@ -97,6 +117,9 @@ func typeSchema(t reflect.Type, seen map[reflect.Type]*Schema) (*Schema, error)
 			if err != nil {
 				return nil, err
 			}
+			if required {
+				s.Required = append(s.Required, name)
+			}
 		}
 
 	default:
@@ -105,14 +128,21 @@ func typeSchema(t reflect.Type, seen map[reflect.Type]*Schema) (*Schema, error)
 	return s, nil
 }
 
-func jsonName(f reflect.StructField) (string, bool) {
+func parseField(f reflect.StructField) (name string, required, include bool) {
 	if !f.IsExported() {
-		return "", false
+		return "", false, false
 	}
+	name = f.Name
+	required = true
 	if tag, ok := f.Tag.Lookup("json"); ok {
-		if name, _, _ := strings.Cut(tag, ","); name != "" {
-			return name, name != "-"
+		props := strings.Split(tag, ",")
+		if props[0] != "" {
+			if props[0] == "-" {
+				return "", false, false
+			}
+			name = props[0]
 		}
+		required = !slices.Contains(props[1:], "omitempty")
 	}
-	return f.Name, true
+	return name, required, true
 }

internal/mcp/internal/jsonschema/infer_test.go

@@ -56,6 +56,7 @@ func TestForType(t *testing.T) {
 				"P":      {Type: "boolean"},
 				"NoSkip": {Type: "string"},
 			},
+			Required:             []string{"f", "G", "P"},
 			AdditionalProperties: &jsonschema.Schema{Not: &jsonschema.Schema{}},
 		}},
 	}

internal/mcp/mcp.go

@@ -67,6 +67,6 @@
 //   - Support all client/server operations.
 //   - Support streamable HTTP transport.
 //   - Support multiple versions of the spec.
-//   - Implement proper JSON schema support, with both client-side and
+//   - Implement full JSON schema support, with both client-side and
 //     server-side validation.
 package mcp

internal/mcp/mcp_test.go

@@ -88,7 +88,8 @@ func TestEndToEnd(t *testing.T) {
 		Name:        "greet",
 		Description: "say hi",
 		InputSchema: &jsonschema.Schema{
-			Type: "object",
+			Type:     "object",
+			Required: []string{"Name"},
 			Properties: map[string]*jsonschema.Schema{
 				"Name": {Type: "string"},
 			},

internal/mcp/tool.go

@@ -7,6 +7,7 @@ package mcp
 import (
 	"context"
 	"encoding/json"
+	"slices"
 
 	"golang.org/x/tools/internal/mcp/internal/jsonschema"
 	"golang.org/x/tools/internal/mcp/internal/protocol"
@@ -23,13 +24,16 @@ type Tool struct {
 
 // MakeTool is a helper to make a tool using reflection on the given handler.
 //
+// If provided, variadic [ToolOption] values may be used to customize the tool.
+//
 // The input schema for the tool is extracted from the request type for the
-// handler, and used to unmmarshal and validate requests to the handler.
+// handler, and used to unmmarshal and validate requests to the handler. This
+// schema may be customized using the [Input] option.
 //
 // It is the caller's responsibility that the handler request type can produce
 // a valid schema, as documented by [jsonschema.ForType]; otherwise, MakeTool
 // panics.
-func MakeTool[TReq any](name, description string, handler func(context.Context, *ClientConnection, TReq) ([]Content, error)) *Tool {
+func MakeTool[TReq any](name, description string, handler func(context.Context, *ClientConnection, TReq) ([]Content, error), opts ...ToolOption) *Tool {
 	schema, err := jsonschema.For[TReq]()
 	if err != nil {
 		panic(err)
@@ -51,14 +55,18 @@ func MakeTool[TReq any](name, description string, handler func(context.Context,
 		}
 		return res, nil
 	}
-	return &Tool{
+	t := &Tool{
 		Definition: protocol.Tool{
 			Name:        name,
 			Description: description,
 			InputSchema: schema,
 		},
 		Handler: wrapped,
 	}
+	for _, opt := range opts {
+		opt.set(t)
+	}
+	return t
 }
 
 // unmarshalSchema unmarshals data into v and validates the result according to
@@ -68,3 +76,94 @@ func unmarshalSchema(data json.RawMessage, _ *jsonschema.Schema, v any) error {
 	// Separate validation from assignment.
 	return json.Unmarshal(data, v)
 }
+
+// A ToolOption configures the behavior of a Tool.
+type ToolOption interface {
+	set(*Tool)
+}
+
+type toolSetter func(*Tool)
+
+func (s toolSetter) set(t *Tool) { s(t) }
+
+// Input applies the provided [SchemaOption] configuration to the tool's input
+// schema.
+func Input(opts ...SchemaOption) ToolOption {
+	return toolSetter(func(t *Tool) {
+		for _, opt := range opts {
+			opt.set(t.Definition.InputSchema)
+		}
+	})
+}
+
+// A SchemaOption configures a jsonschema.Schema.
+type SchemaOption interface {
+	set(s *jsonschema.Schema)
+}
+
+type schemaSetter func(*jsonschema.Schema)
+
+func (s schemaSetter) set(schema *jsonschema.Schema) { s(schema) }
+
+// Property configures the schema for the property of the given name.
+// If there is no such property in the schema, it is created.
+func Property(name string, opts ...SchemaOption) SchemaOption {
+	return schemaSetter(func(schema *jsonschema.Schema) {
+		propSchema, ok := schema.Properties[name]
+		if !ok {
+			propSchema = new(jsonschema.Schema)
+			schema.Properties[name] = propSchema
+		}
+		// Apply the options, with special handling for Required, as it needs to be
+		// set on the parent schema.
+		for _, opt := range opts {
+			if req, ok := opt.(required); ok {
+				if req {
+					if !slices.Contains(schema.Required, name) {
+						schema.Required = append(schema.Required, name)
+					}
+				} else {
+					schema.Required = slices.DeleteFunc(schema.Required, func(s string) bool {
+						return s == name
+					})
+				}
+			} else {
+				opt.set(propSchema)
+			}
+		}
+	})
+}
+
+// Required sets whether the associated property is required. It is only valid
+// when used in a [Property] option: using Required outside of Property panics.
+func Required(v bool) SchemaOption {
+	return required(v)
+}
+
+type required bool
+
+func (required) set(s *jsonschema.Schema) {
+	panic("use of required outside of Property")
+}
+
+// Enum sets the provided values as the "enum" value of the schema.
+func Enum(values ...any) SchemaOption {
+	return schemaSetter(func(s *jsonschema.Schema) {
+		s.Enum = values
+	})
+}
+
+// Description sets the provided schema description.
+func Description(description string) SchemaOption {
+	return schemaSetter(func(schema *jsonschema.Schema) {
+		schema.Description = description
+	})
+}
+
+// Schema overrides the inferred schema with a shallow copy of the given
+// schema.
+func Schema(schema *jsonschema.Schema) SchemaOption {
+	return schemaSetter(func(s *jsonschema.Schema) {
+		*s = *schema
+	})
+}

internal/mcp/tool_test.go

@@ -0,0 +1,89 @@
+// Copyright 2025 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package mcp_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"golang.org/x/tools/internal/mcp"
+	"golang.org/x/tools/internal/mcp/internal/jsonschema"
+)
+
+// testHandler is used for type inference in TestMakeTool.
+func testHandler[T any](context.Context, *mcp.ClientConnection, T) ([]mcp.Content, error) {
+	panic("not implemented")
+}
+
+func TestMakeTool(t *testing.T) {
+	tests := []struct {
+		tool *mcp.Tool
+		want *jsonschema.Schema
+	}{
+		{
+			mcp.MakeTool("basic", "", testHandler[struct {
+				Name string `json:"name"`
+			}]),
+			&jsonschema.Schema{
+				Type:     "object",
+				Required: []string{"name"},
+				Properties: map[string]*jsonschema.Schema{
+					"name": {Type: "string"},
+				},
+				AdditionalProperties: &jsonschema.Schema{Not: new(jsonschema.Schema)},
+			},
+		},
+		{
+			mcp.MakeTool("enum", "", testHandler[struct{ Name string }], mcp.Input(
+				mcp.Property("Name", mcp.Enum("x", "y", "z")),
+			)),
+			&jsonschema.Schema{
+				Type:     "object",
+				Required: []string{"Name"},
+				Properties: map[string]*jsonschema.Schema{
+					"Name": {Type: "string", Enum: []any{"x", "y", "z"}},
+				},
+				AdditionalProperties: &jsonschema.Schema{Not: new(jsonschema.Schema)},
+			},
+		},
+		{
+			mcp.MakeTool("required", "", testHandler[struct {
+				Name     string `json:"name"`
+				Language string `json:"language"`
+				X        int    `json:"x,omitempty"`
+				Y        int    `json:"y,omitempty"`
+			}], mcp.Input(
+				mcp.Property("x", mcp.Required(true)))),
+			&jsonschema.Schema{
+				Type:     "object",
+				Required: []string{"name", "language", "x"},
+				Properties: map[string]*jsonschema.Schema{
+					"language": {Type: "string"},
+					"name":     {Type: "string"},
+					"x":        {Type: "integer"},
+					"y":        {Type: "integer"},
+				},
+				AdditionalProperties: &jsonschema.Schema{Not: new(jsonschema.Schema)},
+			},
+		},
+		{
+			mcp.MakeTool("set_schema", "", testHandler[struct {
+				X int `json:"x,omitempty"`
+				Y int `json:"y,omitempty"`
+			}], mcp.Input(
+				mcp.Schema(&jsonschema.Schema{Type: "object"})),
+			),
+			&jsonschema.Schema{
+				Type: "object",
+			},
+		},
+	}
+	for _, test := range tests {
+		if diff := cmp.Diff(test.want, test.tool.Definition.InputSchema); diff != "" {
+			t.Errorf("MakeTool(%v) mismatch (-want +got):\n%s", test.tool.Definition.Name, diff)
+		}
+	}
+}