From e082a340be847cc05e0f05cacdf720c570ded8f5 Mon Sep 17 00:00:00 2001 From: Rob Findley Date: Tue, 29 Apr 2025 18:26:00 +0000 Subject: [PATCH] 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 LUCI-TryBot-Result: Go LUCI --- examples/hello/main.go | 6 +- internal/jsonschema/infer.go | 44 +++++++++++-- internal/jsonschema/infer_test.go | 1 + mcp/mcp.go | 2 +- mcp/mcp_test.go | 3 +- mcp/tool.go | 105 +++++++++++++++++++++++++++++- mcp/tool_test.go | 89 +++++++++++++++++++++++++ 7 files changed, 236 insertions(+), 14 deletions(-) create mode 100644 mcp/tool_test.go diff --git a/examples/hello/main.go b/examples/hello/main.go index 7ea8353..6a42bbe 100644 --- a/examples/hello/main.go +++ b/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 { diff --git a/internal/jsonschema/infer.go b/internal/jsonschema/infer.go index 34309be..c0a11e5 100644 --- a/internal/jsonschema/infer.go +++ b/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 } diff --git a/internal/jsonschema/infer_test.go b/internal/jsonschema/infer_test.go index a77c15e..0921cbc 100644 --- a/internal/jsonschema/infer_test.go +++ b/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{}}, }}, } diff --git a/mcp/mcp.go b/mcp/mcp.go index 315fba4..c278234 100644 --- a/mcp/mcp.go +++ b/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 diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index e178e3c..c5d8f88 100644 --- a/mcp/mcp_test.go +++ b/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"}, }, diff --git a/mcp/tool.go b/mcp/tool.go index bd19539..ed37cae 100644 --- a/mcp/tool.go +++ b/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 + }) +} diff --git a/mcp/tool_test.go b/mcp/tool_test.go new file mode 100644 index 0000000..36dd8f4 --- /dev/null +++ b/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) + } + } +}