From 04ca8ac03b5c2aba427ec9142ba7180dee3db197 Mon Sep 17 00:00:00 2001 From: Rob Findley Date: Tue, 27 May 2025 14:15:09 +0000 Subject: [PATCH] internal/mcp: generic handling of CallToolParams Make Tool APIs future-proof by pushing generics into CallToolParams, and using CallToolParams throughout. This has several advantages: - We no longer need the awkward CallToolOptions type, which is almost but not quite like CallToolParams. - NewTool no longer needs to define its own handler type: it accepts a ToolHandler[TArgs]. This also means its signature is future-proof. Additionally, provide a `CallTool` helper (the analog of NewTool) to serve as a more strongly typed variant of `ClientSession.CallTool`. + design, protocol generation Change-Id: I80172a25bc50734dae84a541aa4776ed9161694e Reviewed-on: https://go-review.googlesource.com/c/tools/+/676495 Reviewed-by: Jonathan Amsterdam Auto-Submit: Robert Findley LUCI-TryBot-Result: Go LUCI --- README.md | 12 ++++++--- design/design.md | 33 ++++++++++++++---------- examples/hello/main.go | 10 +++++--- examples/sse/main.go | 8 +++--- internal/readme/client/client.go | 6 ++++- internal/readme/server/server.go | 6 ++--- mcp/client.go | 39 ++++++++++++---------------- mcp/cmd_test.go | 5 +++- mcp/features_test.go | 8 +++--- mcp/generate.go | 40 ++++++++++++++++++++++------- mcp/mcp_test.go | 44 +++++++++++++++++++++----------- mcp/protocol.go | 12 ++++----- mcp/server.go | 3 ++- mcp/server_example_test.go | 15 +++++++---- mcp/sse_example_test.go | 11 +++++--- mcp/sse_test.go | 5 +++- mcp/tool.go | 21 ++++++++------- mcp/tool_test.go | 2 +- mcp/transport.go | 4 +-- 19 files changed, 173 insertions(+), 111 deletions(-) diff --git a/README.md b/README.md index 9338072..4f27200 100644 --- a/README.md +++ b/README.md @@ -50,7 +50,11 @@ func main() { } defer session.Close() // Call a tool on the server. - if res, err := session.CallTool(ctx, "greet", map[string]any{"name": "you"}, nil); err != nil { + params := &mcp.CallToolParams[map[string]any]{ + Name: "greet", + Arguments: map[string]any{"name": "you"}, + } + if res, err := mcp.CallTool(ctx, session, params); err != nil { log.Printf("CallTool failed: %v", err) } else { if res.IsError { @@ -78,9 +82,9 @@ type HiParams struct { Name string `json:"name"` } -func SayHi(ctx context.Context, cc *mcp.ServerSession, params *HiParams) ([]*mcp.Content, error) { - return []*mcp.Content{ - mcp.NewTextContent("Hi " + params.Name), +func SayHi(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParams[HiParams]) (*mcp.CallToolResult, error) { + return &mcp.CallToolResult{ + Content: []*mcp.Content{mcp.NewTextContent("Hi " + params.Name)}, }, nil } diff --git a/design/design.md b/design/design.md index 1f9dd63..f43e098 100644 --- a/design/design.md +++ b/design/design.md @@ -276,6 +276,7 @@ type Content struct { func NewTextContent(text string) *Content // etc. ``` + The `Meta` type includes a `map[string]any` for arbitrary data, and a `ProgressToken` field. **Differences from mcp-go**: these types are largely similar, but our type generator flattens types rather than using struct embedding. @@ -379,9 +380,9 @@ In our SDK, RPC methods that are defined in the specification take a context and func (*ClientSession) ListTools(context.Context, *ListToolsParams) (*ListToolsResult, error) ``` -Our SDK has a method for every RPC in the spec, and except for `CallTool`, their signatures all share this form. We do this, rather than providing more convenient shortcut signatures, to maintain backward compatibility if the spec makes backward-compatible changes such as adding a new property to the request parameters (as in [this commit](https://github.com/modelcontextprotocol/modelcontextprotocol/commit/2fce8a077688bf8011e80af06348b8fe1dae08ac), for example). To avoid boilerplate, we don't repeat this signature for RPCs defined in the spec; readers may assume it when we mention a "spec method." +Our SDK has a method for every RPC in the spec, their signatures all share this form. We do this, rather than providing more convenient shortcut signatures, to maintain backward compatibility if the spec makes backward-compatible changes such as adding a new property to the request parameters (as in [this commit](https://github.com/modelcontextprotocol/modelcontextprotocol/commit/2fce8a077688bf8011e80af06348b8fe1dae08ac), for example). To avoid boilerplate, we don't repeat this signature for RPCs defined in the spec; readers may assume it when we mention a "spec method." -`CallTool` is the only exception: for convenience, it takes the tool name and arguments, with an options struct for additional request fields. See the section on Tools below for details. +`CallTool` is the only exception: for convenience when binding to Go argument types, `*CallToolParams[TArgs]` is generic, with a type parameter providing the Go type of the tool arguments. The spec method accepts a `*CallToolParams[json.RawMessage]`, but we provide a generic helper function. See the section on Tools below for details. Why do we use params instead of the full JSON-RPC request? As much as possible, we endeavor to hide JSON-RPC details when they are not relevant to the business logic of your client or server. In this case, the additional information in the JSON-RPC request is just the request ID and method name; the request ID is irrelevant, and the method name is implied by the name of the Go method providing the API. @@ -495,7 +496,6 @@ type Meta struct { } ``` - Handlers can notify their peer about progress by calling the `NotifyProgress` method. The notification is only sent if the peer requested it by providing a progress token. ```go @@ -579,7 +579,15 @@ type ClientOptions struct { A `Tool` is a logical MCP tool, generated from the MCP spec, and a `ServerTool` is a tool bound to a tool handler. +A tool handler accepts `CallToolParams` and returns a `CallToolResult`. However, since we want to bind tools to Go input types, it is convenient in associated APIs to make `CallToolParams` generic, with a type parameter `TArgs` for the tool argument type. This allows tool APIs to manage the marshalling and unmarshalling of tool inputs for their caller. The bound `ServerTool` type expects a `json.RawMessage` for its tool arguments, but the `NewTool` constructor described below provides a mechanism to bind a typed handler. + ```go +type CallToolParams[TArgs any] struct { + Meta Meta `json:"_meta,omitempty"` + Arguments TArgs `json:"arguments,omitempty"` + Name string `json:"name"` +} + type Tool struct { Annotations *ToolAnnotations `json:"annotations,omitempty"` Description string `json:"description,omitempty"` @@ -587,11 +595,11 @@ type Tool struct { Name string `json:"name"` } -type ToolHandler func(context.Context, *ServerSession, *CallToolParams) (*CallToolResult, error) +type ToolHandler[TArgs] func(context.Context, *ServerSession, *CallToolParams[TArgs]) (*CallToolResult, error) type ServerTool struct { Tool Tool - Handler ToolHandler + Handler ToolHandler[json.RawMessage] } ``` @@ -620,12 +628,12 @@ We have found that a hybrid model works well, where the _initial_ schema is deri ```go // NewTool creates a Tool using reflection on the given handler. -func NewTool[TInput any](name, description string, handler func(context.Context, *ServerSession, TInput) ([]Content, error), opts …ToolOption) *ServerTool +func NewTool[TArgs any](name, description string, handler ToolHandler[TArgs], opts …ToolOption) *ServerTool type ToolOption interface { /* ... */ } ``` -`NewTool` determines the input schema for a Tool from the struct used in the handler. Each struct field that would be marshaled by `encoding/json.Marshal` becomes a property of the schema. The property is required unless the field's `json` tag specifies "omitempty" or "omitzero" (new in Go 1.24). For example, given this struct: +`NewTool` determines the input schema for a Tool from the `TArgs` type. Each struct field that would be marshaled by `encoding/json.Marshal` becomes a property of the schema. The property is required unless the field's `json` tag specifies "omitempty" or "omitzero" (new in Go 1.24). For example, given this struct: ```go struct { @@ -666,16 +674,12 @@ Schemas are validated on the server before the tool handler is called. Since all the fields of the Tool struct are exported, a Tool can also be created directly with assignment or a struct literal. -Client sessions can call the spec method `ListTools` or an iterator method `Tools` to list the available tools. - -As mentioned above, the client session method `CallTool` has a non-standard signature, so that `CallTool` can handle the marshalling of tool arguments: the type of `CallToolParams.Arguments` is `json.RawMessage`, to delegate unmarshalling to the tool handler. +Client sessions can call the spec method `ListTools` or an iterator method `Tools` to list the available tools, and use spec method `CallTool` to call tools. Similar to `ServerTool.Handler`, `CallTool` expects `*CallToolParams[json.RawMessage]`, but we provide a generic `CallTool` helper to operate on typed arguments. ```go -func (c *ClientSession) CallTool(ctx context.Context, name string, args map[string]any, opts *CallToolOptions) (_ *CallToolResult, err error) +func (cs *ClientSession) CallTool(context.Context, *CallToolParams[json.RawMessage]) (*CallToolResult, error) -type CallToolOptions struct { - ProgressToken any // string or int -} +func CallTool[TArgs any](context.Context, *ClientSession, *CallToolParams[TArgs]) (*CallToolResult, error) ``` **Differences from mcp-go**: using variadic options to configure tools was significantly inspired by mcp-go. However, the distinction between `ToolOption` and `SchemaOption` allows for recursive application of schema options. For example, that limitation is visible in [this code](https://github.com/DCjanus/dida365-mcp-server/blob/master/cmd/mcp/tools.go#L315), which must resort to untyped maps to express a nested schema. @@ -833,6 +837,7 @@ Clients call the spec method `Complete` to request completions. Servers automati MCP specifies a notification for servers to log to clients. Server sessions implement this with the `LoggingMessage` method. It honors the minimum log level established by the client session's `SetLevel` call. As a convenience, we also provide a `slog.Handler` that allows server authors to write logs with the `log/slog` package:: + ```go // A LoggingHandler is a [slog.Handler] for MCP. type LoggingHandler struct {...} diff --git a/examples/hello/main.go b/examples/hello/main.go index 7d1c97d..afe087c 100644 --- a/examples/hello/main.go +++ b/examples/hello/main.go @@ -21,13 +21,15 @@ type HiArgs struct { Name string `json:"name"` } -func SayHi(ctx context.Context, cc *mcp.ServerSession, params *HiArgs) ([]*mcp.Content, error) { - return []*mcp.Content{ - mcp.NewTextContent("Hi " + params.Name), +func SayHi(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParams[HiArgs]) (*mcp.CallToolResult, error) { + return &mcp.CallToolResult{ + Content: []*mcp.Content{ + mcp.NewTextContent("Hi " + params.Name), + }, }, nil } -func PromptHi(ctx context.Context, cc *mcp.ServerSession, args *HiArgs, params *mcp.GetPromptParams) (*mcp.GetPromptResult, error) { +func PromptHi(ctx context.Context, ss *mcp.ServerSession, args *HiArgs, params *mcp.GetPromptParams) (*mcp.GetPromptResult, error) { return &mcp.GetPromptResult{ Description: "Code review prompt", Messages: []*mcp.PromptMessage{ diff --git a/examples/sse/main.go b/examples/sse/main.go index 4c199f8..129f956 100644 --- a/examples/sse/main.go +++ b/examples/sse/main.go @@ -19,9 +19,11 @@ type SayHiParams struct { Name string `json:"name"` } -func SayHi(ctx context.Context, cc *mcp.ServerSession, params *SayHiParams) ([]*mcp.Content, error) { - return []*mcp.Content{ - mcp.NewTextContent("Hi " + params.Name), +func SayHi(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParams[SayHiParams]) (*mcp.CallToolResult, error) { + return &mcp.CallToolResult{ + Content: []*mcp.Content{ + mcp.NewTextContent("Hi " + params.Name), + }, }, nil } diff --git a/internal/readme/client/client.go b/internal/readme/client/client.go index ffe7ec6..b9e00b9 100644 --- a/internal/readme/client/client.go +++ b/internal/readme/client/client.go @@ -25,7 +25,11 @@ func main() { } defer session.Close() // Call a tool on the server. - if res, err := session.CallTool(ctx, "greet", map[string]any{"name": "you"}, nil); err != nil { + params := &mcp.CallToolParams[map[string]any]{ + Name: "greet", + Arguments: map[string]any{"name": "you"}, + } + if res, err := mcp.CallTool(ctx, session, params); err != nil { log.Printf("CallTool failed: %v", err) } else { if res.IsError { diff --git a/internal/readme/server/server.go b/internal/readme/server/server.go index 9644828..269297a 100644 --- a/internal/readme/server/server.go +++ b/internal/readme/server/server.go @@ -15,9 +15,9 @@ type HiParams struct { Name string `json:"name"` } -func SayHi(ctx context.Context, cc *mcp.ServerSession, params *HiParams) ([]*mcp.Content, error) { - return []*mcp.Content{ - mcp.NewTextContent("Hi " + params.Name), +func SayHi(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParams[HiParams]) (*mcp.CallToolResult, error) { + return &mcp.CallToolResult{ + Content: []*mcp.Content{mcp.NewTextContent("Hi " + params.Name)}, }, nil } diff --git a/mcp/client.go b/mcp/client.go index 341142e..c866b8c 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -258,35 +258,30 @@ func (cs *ClientSession) ListTools(ctx context.Context, params *ListToolsParams) // CallTool calls the tool with the given name and arguments. // Pass a [CallToolOptions] to provide additional request fields. -func (cs *ClientSession) CallTool(ctx context.Context, name string, args map[string]any, opts *CallToolOptions) (_ *CallToolResult, err error) { - defer func() { - if err != nil { - err = fmt.Errorf("calling tool %q: %w", name, err) - } - }() - - data, err := json.Marshal(args) - if err != nil { - return nil, fmt.Errorf("marshaling arguments: %w", err) - } - params := &CallToolParams{ - Name: name, - Arguments: json.RawMessage(data), - } +func (cs *ClientSession) CallTool(ctx context.Context, params *CallToolParams[json.RawMessage]) (*CallToolResult, error) { return standardCall[CallToolResult](ctx, cs.conn, methodCallTool, params) } +// CallTool is a helper to call a tool with any argument type. It returns an +// error if params.Arguments fails to marshal to JSON. +func CallTool[TArgs any](ctx context.Context, cs *ClientSession, params *CallToolParams[TArgs]) (*CallToolResult, error) { + data, err := json.Marshal(params.Arguments) + if err != nil { + return nil, fmt.Errorf("failed to marshal arguments: %v", err) + } + // TODO(rfindley): write a test that guarantees this copying is total. + wireParams := &CallToolParams[json.RawMessage]{ + Meta: params.Meta, + Name: params.Name, + Arguments: data, + } + return cs.CallTool(ctx, wireParams) +} + func (cs *ClientSession) SetLevel(ctx context.Context, params *SetLevelParams) error { return call(ctx, cs.conn, methodSetLevel, params, nil) } -// NOTE: the following struct should consist of all fields of callToolParams except name and arguments. - -// CallToolOptions contains options to [ClientSession.CallTool]. -type CallToolOptions struct { - ProgressToken any // string or int -} - // ListResources lists the resources that are currently available on the server. func (cs *ClientSession) ListResources(ctx context.Context, params *ListResourcesParams) (*ListResourcesResult, error) { return standardCall[ListResourcesResult](ctx, cs.conn, methodListResources, params) diff --git a/mcp/cmd_test.go b/mcp/cmd_test.go index 922dbd2..e4f33dd 100644 --- a/mcp/cmd_test.go +++ b/mcp/cmd_test.go @@ -53,7 +53,10 @@ func TestCmdTransport(t *testing.T) { if err != nil { log.Fatal(err) } - got, err := session.CallTool(ctx, "greet", map[string]any{"name": "user"}, nil) + got, err := mcp.CallTool(ctx, session, &mcp.CallToolParams[map[string]any]{ + Name: "greet", + Arguments: map[string]any{"name": "user"}, + }) if err != nil { log.Fatal(err) } diff --git a/mcp/features_test.go b/mcp/features_test.go index 478457d..7ff3536 100644 --- a/mcp/features_test.go +++ b/mcp/features_test.go @@ -18,9 +18,11 @@ type SayHiParams struct { Name string `json:"name"` } -func SayHi(ctx context.Context, cc *ServerSession, params *SayHiParams) ([]*Content, error) { - return []*Content{ - NewTextContent("Hi " + params.Name), +func SayHi(ctx context.Context, cc *ServerSession, params *CallToolParams[SayHiParams]) (*CallToolResult, error) { + return &CallToolResult{ + Content: []*Content{ + NewTextContent("Hi " + params.Name), + }, }, nil } diff --git a/mcp/generate.go b/mcp/generate.go index 809b2ec..dbe89eb 100644 --- a/mcp/generate.go +++ b/mcp/generate.go @@ -38,9 +38,10 @@ var schemaFile = flag.String("schema_file", "", "if set, use this file as the pe // struct field. In others, we may want to extract the type definition to a // name. type typeConfig struct { - Name string // declaration name for the type - Substitute string // type definition to substitute - Fields config // individual field configuration, or nil + Name string // declaration name for the type + TypeParams [][2]string // formatted type parameter list ({name, constraint}), if any + Substitute string // type definition to substitute + Fields config // individual field configuration, or nil } type config map[string]*typeConfig @@ -57,9 +58,10 @@ var declarations = config{ Name: "-", Fields: config{ "Params": { - Name: "CallToolParams", + Name: "CallToolParams", + TypeParams: [][2]string{{"TArgs", "any"}}, Fields: config{ - "Arguments": {Substitute: "json.RawMessage"}, + "Arguments": {Substitute: "TArgs"}, }, }, }, @@ -224,8 +226,6 @@ func main() { package mcp import ( - "encoding/json" - "github.com/modelcontextprotocol/go-sdk/jsonschema" ) `) @@ -320,7 +320,18 @@ func writeDecl(configName string, config typeConfig, def *jsonschema.Schema, nam if def.Description != "" { fmt.Fprintf(buf, "%s\n", toComment(def.Description)) } - fmt.Fprintf(buf, "type %s ", typeName) + typeParams := new(strings.Builder) + if len(config.TypeParams) > 0 { + typeParams.WriteByte('[') + for i, p := range config.TypeParams { + if i > 0 { + typeParams.WriteString(", ") + } + fmt.Fprintf(typeParams, "%s %s", p[0], p[1]) + } + typeParams.WriteByte(']') + } + fmt.Fprintf(buf, "type %s%s ", typeName, typeParams) } if err := writeType(w, &config, def, named); err != nil { return err // Better error here? @@ -329,7 +340,18 @@ func writeDecl(configName string, config typeConfig, def *jsonschema.Schema, nam // Any decl with a _meta field gets a GetMeta method. if _, ok := def.Properties["_meta"]; ok { - fmt.Fprintf(w, "\nfunc (x *%s) GetMeta() *Meta { return &x.Meta }", typeName) + targs := new(strings.Builder) + if len(config.TypeParams) > 0 { + targs.WriteByte('[') + for i, p := range config.TypeParams { + if i > 0 { + targs.WriteString(", ") + } + fmt.Fprintf(targs, "%s", p[0]) + } + targs.WriteByte(']') + } + fmt.Fprintf(w, "\nfunc (x *%s%s) GetMeta() *Meta { return &x.Meta }", typeName, targs) } if _, ok := def.Properties["cursor"]; ok { diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index 34d4b32..1dea894 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -30,16 +30,18 @@ type hiParams struct { Name string } -func sayHi(ctx context.Context, ss *ServerSession, v hiParams) ([]*Content, error) { +func sayHi(ctx context.Context, ss *ServerSession, params *CallToolParams[hiParams]) (*CallToolResult, error) { if err := ss.Ping(ctx, nil); err != nil { return nil, fmt.Errorf("ping failed: %v", err) } - return []*Content{NewTextContent("hi " + v.Name)}, nil + return &CallToolResult{Content: []*Content{NewTextContent("hi " + params.Arguments.Name)}}, nil } func TestEndToEnd(t *testing.T) { ctx := context.Background() - ct, st := NewInMemoryTransports() + var ct, st Transport = NewInMemoryTransports() + // ct = NewLoggingTransport(ct, os.Stderr) + // st = NewLoggingTransport(st, os.Stderr) // Channels to check if notification callbacks happened. notificationChans := map[string]chan int{} @@ -185,7 +187,10 @@ func TestEndToEnd(t *testing.T) { t.Fatalf("tools/list mismatch (-want +got):\n%s", diff) } - gotHi, err := cs.CallTool(ctx, "greet", map[string]any{"name": "user"}, nil) + gotHi, err := CallTool(ctx, cs, &CallToolParams[map[string]any]{ + Name: "greet", + Arguments: map[string]any{"name": "user"}, + }) if err != nil { t.Fatal(err) } @@ -196,7 +201,10 @@ func TestEndToEnd(t *testing.T) { t.Errorf("tools/call 'greet' mismatch (-want +got):\n%s", diff) } - gotFail, err := cs.CallTool(ctx, "fail", map[string]any{}, nil) + gotFail, err := CallTool(ctx, cs, &CallToolParams[map[string]any]{ + Name: "fail", + Arguments: map[string]any{}, + }) // Counter-intuitively, when a tool fails, we don't expect an RPC error for // call tool: instead, the failure is embedded in the result. if err != nil { @@ -387,7 +395,7 @@ var ( tools = map[string]*ServerTool{ "greet": NewTool("greet", "say hi", sayHi), - "fail": NewTool("fail", "just fail", func(context.Context, *ServerSession, struct{}) ([]*Content, error) { + "fail": NewTool("fail", "just fail", func(context.Context, *ServerSession, *CallToolParams[struct{}]) (*CallToolResult, error) { return nil, errTestFailure }), } @@ -506,24 +514,30 @@ func basicConnection(t *testing.T, tools ...*ServerTool) (*ServerSession, *Clien } func TestServerClosing(t *testing.T) { - cc, c := basicConnection(t, NewTool("greet", "say hi", sayHi)) - defer c.Close() + cc, cs := basicConnection(t, NewTool("greet", "say hi", sayHi)) + defer cs.Close() ctx := context.Background() var wg sync.WaitGroup wg.Add(1) go func() { - if err := c.Wait(); err != nil { + if err := cs.Wait(); err != nil { t.Errorf("server connection failed: %v", err) } wg.Done() }() - if _, err := c.CallTool(ctx, "greet", map[string]any{"name": "user"}, nil); err != nil { + if _, err := CallTool(ctx, cs, &CallToolParams[map[string]any]{ + Name: "greet", + Arguments: map[string]any{"name": "user"}, + }); err != nil { t.Fatalf("after connecting: %v", err) } cc.Close() wg.Wait() - if _, err := c.CallTool(ctx, "greet", map[string]any{"name": "user"}, nil); !errors.Is(err, ErrConnectionClosed) { + if _, err := CallTool(ctx, cs, &CallToolParams[map[string]any]{ + Name: "greet", + Arguments: map[string]any{"name": "user"}, + }); !errors.Is(err, ErrConnectionClosed) { t.Errorf("after disconnection, got error %v, want EOF", err) } } @@ -572,7 +586,7 @@ func TestCancellation(t *testing.T) { cancelled = make(chan struct{}, 1) // don't block the request ) - slowRequest := func(ctx context.Context, cc *ServerSession, v struct{}) ([]*Content, error) { + slowRequest := func(ctx context.Context, cc *ServerSession, params *CallToolParams[struct{}]) (*CallToolResult, error) { start <- struct{}{} select { case <-ctx.Done(): @@ -582,11 +596,11 @@ func TestCancellation(t *testing.T) { } return nil, nil } - _, sc := basicConnection(t, NewTool("slow", "a slow request", slowRequest)) - defer sc.Close() + _, cs := basicConnection(t, NewTool("slow", "a slow request", slowRequest)) + defer cs.Close() ctx, cancel := context.WithCancel(context.Background()) - go sc.CallTool(ctx, "slow", map[string]any{}, nil) + go CallTool(ctx, cs, &CallToolParams[struct{}]{Name: "slow"}) <-start cancel() select { diff --git a/mcp/protocol.go b/mcp/protocol.go index 79965ab..1ad8ed5 100644 --- a/mcp/protocol.go +++ b/mcp/protocol.go @@ -7,8 +7,6 @@ package mcp import ( - "encoding/json" - "github.com/modelcontextprotocol/go-sdk/jsonschema" ) @@ -28,15 +26,15 @@ type Annotations struct { Priority float64 `json:"priority,omitempty"` } -type CallToolParams struct { +type CallToolParams[TArgs any] struct { // This property is reserved by the protocol to allow clients and servers to // attach additional metadata to their responses. - Meta Meta `json:"_meta,omitempty"` - Arguments json.RawMessage `json:"arguments,omitempty"` - Name string `json:"name"` + Meta Meta `json:"_meta,omitempty"` + Arguments TArgs `json:"arguments,omitempty"` + Name string `json:"name"` } -func (x *CallToolParams) GetMeta() *Meta { return &x.Meta } +func (x *CallToolParams[TArgs]) GetMeta() *Meta { return &x.Meta } // The server's response to a tool call. // diff --git a/mcp/server.go b/mcp/server.go index 635df94..94c146f 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -9,6 +9,7 @@ import ( "context" "encoding/base64" "encoding/gob" + "encoding/json" "fmt" "iter" "net/url" @@ -227,7 +228,7 @@ func (s *Server) listTools(_ context.Context, _ *ServerSession, params *ListTool return res, nil } -func (s *Server) callTool(ctx context.Context, cc *ServerSession, params *CallToolParams) (*CallToolResult, error) { +func (s *Server) callTool(ctx context.Context, cc *ServerSession, params *CallToolParams[json.RawMessage]) (*CallToolResult, error) { s.mu.Lock() tool, ok := s.tools.get(params.Name) s.mu.Unlock() diff --git a/mcp/server_example_test.go b/mcp/server_example_test.go index 91f7f1b..9855210 100644 --- a/mcp/server_example_test.go +++ b/mcp/server_example_test.go @@ -17,12 +17,14 @@ import ( ) 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.ServerSession, params *SayHiParams) ([]*mcp.Content, error) { - return []*mcp.Content{ - mcp.NewTextContent("Hi " + params.Name), +func SayHi(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParams[SayHiParams]) (*mcp.CallToolResult, error) { + return &mcp.CallToolResult{ + Content: []*mcp.Content{ + mcp.NewTextContent("Hi " + params.Arguments.Name), + }, }, nil } @@ -44,7 +46,10 @@ func ExampleServer() { log.Fatal(err) } - res, err := clientSession.CallTool(ctx, "greet", map[string]any{"name": "user"}, nil) + res, err := mcp.CallTool(ctx, clientSession, &mcp.CallToolParams[map[string]any]{ + Name: "greet", + Arguments: map[string]any{"name": "user"}, + }) if err != nil { log.Fatal(err) } diff --git a/mcp/sse_example_test.go b/mcp/sse_example_test.go index 102fce5..4f0e4c3 100644 --- a/mcp/sse_example_test.go +++ b/mcp/sse_example_test.go @@ -18,9 +18,9 @@ type AddParams struct { X, Y int } -func Add(ctx context.Context, cc *mcp.ServerSession, params *AddParams) ([]*mcp.Content, error) { - return []*mcp.Content{ - mcp.NewTextContent(fmt.Sprintf("%d", params.X+params.Y)), +func Add(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParams[AddParams]) (*mcp.CallToolResult, error) { + return &mcp.CallToolResult{ + Content: []*mcp.Content{mcp.NewTextContent(fmt.Sprintf("%d", params.Arguments.X+params.Arguments.Y))}, }, nil } @@ -41,7 +41,10 @@ func ExampleSSEHandler() { } defer cs.Close() - res, err := cs.CallTool(ctx, "add", map[string]any{"x": 1, "y": 2}, nil) + res, err := mcp.CallTool(ctx, cs, &mcp.CallToolParams[map[string]any]{ + Name: "add", + Arguments: map[string]any{"x": 1, "y": 2}, + }) if err != nil { log.Fatal(err) } diff --git a/mcp/sse_test.go b/mcp/sse_test.go index 8b90958..f122cd3 100644 --- a/mcp/sse_test.go +++ b/mcp/sse_test.go @@ -44,7 +44,10 @@ func TestSSEServer(t *testing.T) { t.Fatal(err) } ss := <-conns - gotHi, err := cs.CallTool(ctx, "greet", map[string]any{"name": "user"}, nil) + gotHi, err := CallTool(ctx, cs, &CallToolParams[map[string]any]{ + Name: "greet", + Arguments: map[string]any{"name": "user"}, + }) if err != nil { t.Fatal(err) } diff --git a/mcp/tool.go b/mcp/tool.go index 47d9040..8f7673e 100644 --- a/mcp/tool.go +++ b/mcp/tool.go @@ -13,12 +13,12 @@ import ( ) // A ToolHandler handles a call to tools/call. -type ToolHandler func(context.Context, *ServerSession, *CallToolParams) (*CallToolResult, error) +type ToolHandler[TArgs any] func(context.Context, *ServerSession, *CallToolParams[TArgs]) (*CallToolResult, error) // A Tool is a tool definition that is bound to a tool handler. type ServerTool struct { Tool *Tool - Handler ToolHandler + Handler ToolHandler[json.RawMessage] } // NewTool is a helper to make a tool using reflection on the given handler. @@ -34,17 +34,19 @@ type ServerTool struct { // // TODO: just have the handler return a CallToolResult: returning []Content is // going to be inconsistent with other server features. -func NewTool[TReq any](name, description string, handler func(context.Context, *ServerSession, TReq) ([]*Content, error), opts ...ToolOption) *ServerTool { +func NewTool[TReq any](name, description string, handler ToolHandler[TReq], opts ...ToolOption) *ServerTool { schema, err := jsonschema.For[TReq]() if err != nil { panic(err) } - wrapped := func(ctx context.Context, cc *ServerSession, params *CallToolParams) (*CallToolResult, error) { - var v TReq - if err := unmarshalSchema(params.Arguments, schema, &v); err != nil { - return nil, err + wrapped := func(ctx context.Context, cc *ServerSession, params *CallToolParams[json.RawMessage]) (*CallToolResult, error) { + var params2 CallToolParams[TReq] + if params.Arguments != nil { + if err := unmarshalSchema(params.Arguments, schema, ¶ms2.Arguments); err != nil { + return nil, err + } } - content, err := handler(ctx, cc, v) + res, err := handler(ctx, cc, ¶ms2) // TODO: investigate why server errors are embedded in this strange way, // rather than returned as jsonrpc2 server errors. if err != nil { @@ -53,9 +55,6 @@ func NewTool[TReq any](name, description string, handler func(context.Context, * IsError: true, }, nil } - res := &CallToolResult{ - Content: content, - } return res, nil } t := &ServerTool{ diff --git a/mcp/tool_test.go b/mcp/tool_test.go index a0cfc76..f34dfd2 100644 --- a/mcp/tool_test.go +++ b/mcp/tool_test.go @@ -15,7 +15,7 @@ import ( ) // testToolHandler is used for type inference in TestNewTool. -func testToolHandler[T any](context.Context, *mcp.ServerSession, T) ([]*mcp.Content, error) { +func testToolHandler[T any](context.Context, *mcp.ServerSession, *mcp.CallToolParams[T]) (*mcp.CallToolResult, error) { panic("not implemented") } diff --git a/mcp/transport.go b/mcp/transport.go index f69299d..ee7c16d 100644 --- a/mcp/transport.go +++ b/mcp/transport.go @@ -200,7 +200,7 @@ func (s *loggingStream) Read(ctx context.Context) (jsonrpc2.Message, int64, erro if err != nil { fmt.Fprintf(s.w, "LoggingTransport: failed to marshal: %v", err) } - fmt.Fprintf(s.w, "read: %s", string(data)) + fmt.Fprintf(s.w, "read: %s\n", string(data)) } return msg, n, err } @@ -215,7 +215,7 @@ func (s *loggingStream) Write(ctx context.Context, msg jsonrpc2.Message) (int64, if err != nil { fmt.Fprintf(s.w, "LoggingTransport: failed to marshal: %v", err) } - fmt.Fprintf(s.w, "write: %s", string(data)) + fmt.Fprintf(s.w, "write: %s\n", string(data)) } return n, err }