From 7bfde44981f838739645a1a4df900c70c831be98 Mon Sep 17 00:00:00 2001 From: Jonathan Amsterdam Date: Sat, 30 Aug 2025 08:16:06 -0400 Subject: [PATCH] mcp: statically type the server-side tool params (#378) Introduce CallToolParamsRaw on the server side, so that tool handlers see in the type system that a tool's arguments are a json.RawMessage. Fixes #377. --- mcp/mcp_test.go | 5 ++--- mcp/protocol.go | 26 ++++++++++++-------------- mcp/requests.go | 2 +- mcp/server.go | 6 ++---- 4 files changed, 17 insertions(+), 22 deletions(-) diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index 8529be5..a16ac83 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -494,7 +494,6 @@ func TestEndToEnd(t *testing.T) { if result.Action != "accept" { t.Errorf("got action %q, want %q", result.Action, "accept") } - }) // Disconnect. @@ -1638,7 +1637,7 @@ func TestPointerArgEquivalence(t *testing.T) { // // We handle a few different types of results, to assert they behave the // same in all cases. - AddTool(s, &Tool{Name: "pointer"}, func(_ context.Context, req *ServerRequest[*CallToolParams], in *input) (*CallToolResult, *output, error) { + AddTool(s, &Tool{Name: "pointer"}, func(_ context.Context, req *CallToolRequest, in *input) (*CallToolResult, *output, error) { switch in.In { case "": return nil, nil, fmt.Errorf("must provide input") @@ -1652,7 +1651,7 @@ func TestPointerArgEquivalence(t *testing.T) { panic("unreachable") } }) - AddTool(s, &Tool{Name: "nonpointer"}, func(_ context.Context, req *ServerRequest[*CallToolParams], in input) (*CallToolResult, output, error) { + AddTool(s, &Tool{Name: "nonpointer"}, func(_ context.Context, req *CallToolRequest, in input) (*CallToolResult, output, error) { switch in.In { case "": return nil, output{}, fmt.Errorf("must provide input") diff --git a/mcp/protocol.go b/mcp/protocol.go index 382f745..2786065 100644 --- a/mcp/protocol.go +++ b/mcp/protocol.go @@ -40,6 +40,7 @@ type Annotations struct { Priority float64 `json:"priority,omitempty"` } +// CallToolParams is used by clients to call a tool. type CallToolParams struct { // This property is reserved by the protocol to allow clients and servers to // attach additional metadata to their responses. @@ -48,20 +49,13 @@ type CallToolParams struct { Arguments any `json:"arguments,omitempty"` } -// When unmarshalling CallToolParams on the server side, we need to delay unmarshaling of the arguments. -func (c *CallToolParams) UnmarshalJSON(data []byte) error { - var raw struct { - Meta `json:"_meta,omitempty"` - Name string `json:"name"` - RawArguments json.RawMessage `json:"arguments,omitempty"` - } - if err := json.Unmarshal(data, &raw); err != nil { - return err - } - c.Meta = raw.Meta - c.Name = raw.Name - c.Arguments = raw.RawArguments - return nil +// CallToolParamsRaw is passed to tool handlers on the server. +type CallToolParamsRaw struct { + // This property is reserved by the protocol to allow clients and servers to + // attach additional metadata to their responses. + Meta `json:"_meta,omitempty"` + Name string `json:"name"` + Arguments json.RawMessage `json:"arguments,omitempty"` } // The server's response to a tool call. @@ -115,6 +109,10 @@ func (x *CallToolParams) isParams() {} func (x *CallToolParams) GetProgressToken() any { return getProgressToken(x) } func (x *CallToolParams) SetProgressToken(t any) { setProgressToken(x, t) } +func (x *CallToolParamsRaw) isParams() {} +func (x *CallToolParamsRaw) GetProgressToken() any { return getProgressToken(x) } +func (x *CallToolParamsRaw) SetProgressToken(t any) { setProgressToken(x, t) } + type CancelledParams struct { // This property is reserved by the protocol to allow clients and servers to // attach additional metadata to their responses. diff --git a/mcp/requests.go b/mcp/requests.go index 3afaac5..82b700f 100644 --- a/mcp/requests.go +++ b/mcp/requests.go @@ -7,7 +7,7 @@ package mcp type ( - CallToolRequest = ServerRequest[*CallToolParams] + CallToolRequest = ServerRequest[*CallToolParamsRaw] CompleteRequest = ServerRequest[*CompleteParams] GetPromptRequest = ServerRequest[*GetPromptParams] InitializedRequest = ServerRequest[*InitializedParams] diff --git a/mcp/server.go b/mcp/server.go index dd5d807..c496b33 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -9,7 +9,6 @@ import ( "context" "encoding/base64" "encoding/gob" - "encoding/json" "fmt" "iter" "maps" @@ -234,10 +233,9 @@ func toolForErr[In, Out any](t *Tool, h ToolHandlerFor[In, Out]) (*Tool, ToolHan th := func(ctx context.Context, req *CallToolRequest) (*CallToolResult, error) { // Unmarshal and validate args. - rawArgs := req.Params.Arguments.(json.RawMessage) var in In - if rawArgs != nil { - if err := unmarshalSchema(rawArgs, inputResolved, &in); err != nil { + if req.Params.Arguments != nil { + if err := unmarshalSchema(req.Params.Arguments, inputResolved, &in); err != nil { return nil, err } }