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 <jba@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>

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

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"},
 			},

mcp/tool.go

@@ -7,6 +7,7 @@ package mcp
 import (
 	"context"
 	"encoding/json"
+	"slices"
 
 	"github.com/modelcontextprotocol/go-sdk/internal/jsonschema"
 	"github.com/modelcontextprotocol/go-sdk/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,7 +55,7 @@ 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,
@@ -59,6 +63,10 @@ func MakeTool[TReq any](name, description string, handler func(context.Context,
 		},
 		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
+	})
+}

mcp/tool_test.go

@@ -0,0 +1,89 @@
+// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
+// Use of this source code is governed by an MIT-style
+// license that can be found in the LICENSE file.
+
+package mcp_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/modelcontextprotocol/go-sdk/mcp"
+	"github.com/modelcontextprotocol/go-sdk/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)
+		}
+	}
+}