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

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)

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)
 	}

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
 }
 

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 {

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 {

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.
 //

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()

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)
 	}

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)
 	}

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)
 			}

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, &params2.Arguments); err != nil {
+				return nil, err
+			}
 		}
-		content, err := handler(ctx, cc, v)
+		res, err := handler(ctx, cc, &params2)
 		// 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{

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

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
 }