mcp: new tool API

Author: jba

mcp/client_list_test.go

@@ -22,12 +22,14 @@ func TestList(t *testing.T) {
 	defer serverSession.Close()
 
 	t.Run("tools", func(t *testing.T) {
-		toolA := mcp.NewServerTool("apple", "apple tool", SayHi)
-		toolB := mcp.NewServerTool("banana", "banana tool", SayHi)
-		toolC := mcp.NewServerTool("cherry", "cherry tool", SayHi)
-		tools := []*mcp.ServerTool{toolA, toolB, toolC}
-		wantTools := []*mcp.Tool{toolA.Tool, toolB.Tool, toolC.Tool}
-		server.AddTools(tools...)
+		var wantTools []*mcp.Tool
+		for _, name := range []string{"apple", "banana", "cherry"} {
+			wantTools = append(wantTools, &mcp.Tool{Name: name, Description: name + " tool"})
+		}
+
+		for _, t := range wantTools {
+			mcp.AddTool(server, t, SayHi)
+		}
 		t.Run("list", func(t *testing.T) {
 			res, err := clientSession.ListTools(ctx, nil)
 			if err != nil {

mcp/cmd_test.go

@@ -31,8 +31,7 @@ func runServer() {
 	ctx := context.Background()
 
 	server := mcp.NewServer("greeter", "v0.0.1", nil)
-	server.AddTools(mcp.NewServerTool("greet", "say hi", SayHi))
-
+	mcp.AddTool(server, &mcp.Tool{Name: "greet", Description: "say hi"}, SayHi)
 	if err := server.Run(ctx, mcp.NewStdioTransport()); err != nil {
 		log.Fatal(err)
 	}

mcp/features_test.go

@@ -27,9 +27,9 @@ func SayHi(ctx context.Context, cc *ServerSession, params *CallToolParamsFor[Say
 }
 
 func TestFeatureSetOrder(t *testing.T) {
-	toolA := NewServerTool("apple", "apple tool", SayHi).Tool
-	toolB := NewServerTool("banana", "banana tool", SayHi).Tool
-	toolC := NewServerTool("cherry", "cherry tool", SayHi).Tool
+	toolA := &Tool{Name: "apple", Description: "apple tool"}
+	toolB := &Tool{Name: "banana", Description: "banana tool"}
+	toolC := &Tool{Name: "cherry", Description: "cherry tool"}
 
 	testCases := []struct {
 		tools []*Tool
@@ -52,9 +52,9 @@ func TestFeatureSetOrder(t *testing.T) {
 }
 
 func TestFeatureSetAbove(t *testing.T) {
-	toolA := NewServerTool("apple", "apple tool", SayHi).Tool
-	toolB := NewServerTool("banana", "banana tool", SayHi).Tool
-	toolC := NewServerTool("cherry", "cherry tool", SayHi).Tool
+	toolA := &Tool{Name: "apple", Description: "apple tool"}
+	toolB := &Tool{Name: "banana", Description: "banana tool"}
+	toolC := &Tool{Name: "cherry", Description: "cherry tool"}
 
 	testCases := []struct {
 		tools []*Tool

mcp/mcp_test.go

@@ -21,7 +21,6 @@ import (
 	"time"
 
 	"github.com/google/go-cmp/cmp"
-	"github.com/google/go-cmp/cmp/cmpopts"
 	"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
 	"github.com/modelcontextprotocol/go-sdk/jsonschema"
 )
@@ -30,6 +29,9 @@ type hiParams struct {
 	Name string
 }
 
+// TODO(jba): after schemas are stateless (WIP), this can be a variable.
+func greetTool() *Tool { return &Tool{Name: "greet", Description: "say hi"} }
+
 func sayHi(ctx context.Context, ss *ServerSession, params *CallToolParamsFor[hiParams]) (*CallToolResultFor[any], error) {
 	if err := ss.Ping(ctx, nil); err != nil {
 		return nil, fmt.Errorf("ping failed: %v", err)
@@ -63,7 +65,14 @@ func TestEndToEnd(t *testing.T) {
 		},
 	}
 	s := NewServer("testServer", "v1.0.0", sopts)
-	add(tools, s.AddTools, "greet", "fail")
+	AddTool(s, &Tool{
+		Name:        "greet",
+		Description: "say hi",
+	}, sayHi)
+	s.AddTool(&Tool{Name: "fail", InputSchema: &jsonschema.Schema{}},
+		func(context.Context, *ServerSession, *CallToolParamsFor[map[string]any]) (*CallToolResult, error) {
+			return nil, errTestFailure
+		})
 	add(prompts, s.AddPrompts, "code_review", "fail")
 	add(resources, s.AddResources, "info.txt", "fail.txt")
 
@@ -161,32 +170,7 @@ func TestEndToEnd(t *testing.T) {
 	})
 
 	t.Run("tools", func(t *testing.T) {
-		res, err := cs.ListTools(ctx, nil)
-		if err != nil {
-			t.Errorf("tools/list failed: %v", err)
-		}
-		wantTools := []*Tool{
-			{
-				Name:        "fail",
-				InputSchema: nil,
-			},
-			{
-				Name:        "greet",
-				Description: "say hi",
-				InputSchema: &jsonschema.Schema{
-					Type:     "object",
-					Required: []string{"Name"},
-					Properties: map[string]*jsonschema.Schema{
-						"Name": {Type: "string"},
-					},
-					AdditionalProperties: falseSchema(),
-				},
-			},
-		}
-		if diff := cmp.Diff(wantTools, res.Tools, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
-			t.Fatalf("tools/list mismatch (-want +got):\n%s", diff)
-		}
-
+		// ListTools is tested in client_list_test.go.
 		gotHi, err := cs.CallTool(ctx, &CallToolParams{
 			Name:      "greet",
 			Arguments: map[string]any{"name": "user"},
@@ -222,7 +206,7 @@ func TestEndToEnd(t *testing.T) {
 			t.Errorf("tools/call 'fail' mismatch (-want +got):\n%s", diff)
 		}
 
-		s.AddTools(&ServerTool{Tool: &Tool{Name: "T"}, Handler: nopHandler})
+		s.AddTool(&Tool{Name: "T", InputSchema: &jsonschema.Schema{}}, nopHandler)
 		waitForNotification(t, "tools")
 		s.RemoveTools("T")
 		waitForNotification(t, "tools")
@@ -434,16 +418,6 @@ func TestEndToEnd(t *testing.T) {
 var (
 	errTestFailure = errors.New("mcp failure")
 
-	tools = map[string]*ServerTool{
-		"greet": NewServerTool("greet", "say hi", sayHi),
-		"fail": {
-			Tool: &Tool{Name: "fail"},
-			Handler: func(context.Context, *ServerSession, *CallToolParamsFor[map[string]any]) (*CallToolResult, error) {
-				return nil, errTestFailure
-			},
-		},
-	}
-
 	prompts = map[string]*ServerPrompt{
 		"code_review": {
 			Prompt: &Prompt{
@@ -540,21 +514,21 @@ func errorCode(err error) int64 {
 	return -1
 }
 
-// basicConnection returns a new basic client-server connection configured with
-// the provided tools.
+// basicConnection returns a new basic client-server connection, with the server
+// configured via the provided function.
 //
 // The caller should cancel either the client connection or server connection
 // when the connections are no longer needed.
-func basicConnection(t *testing.T, tools ...*ServerTool) (*ServerSession, *ClientSession) {
+func basicConnection(t *testing.T, config func(*Server)) (*ServerSession, *ClientSession) {
 	t.Helper()
 
 	ctx := context.Background()
 	ct, st := NewInMemoryTransports()
 
 	s := NewServer("testServer", "v1.0.0", nil)
-
-	// The 'greet' tool says hi.
-	s.AddTools(tools...)
+	if config != nil {
+		config(s)
+	}
 	ss, err := s.Connect(ctx, st)
 	if err != nil {
 		t.Fatal(err)
@@ -569,7 +543,9 @@ func basicConnection(t *testing.T, tools ...*ServerTool) (*ServerSession, *Clien
 }
 
 func TestServerClosing(t *testing.T) {
-	cc, cs := basicConnection(t, NewServerTool("greet", "say hi", sayHi))
+	cc, cs := basicConnection(t, func(s *Server) {
+		AddTool(s, greetTool(), sayHi)
+	})
 	defer cs.Close()
 
 	ctx := context.Background()
@@ -651,11 +627,9 @@ func TestCancellation(t *testing.T) {
 		}
 		return nil, nil
 	}
-	st := &ServerTool{
-		Tool:    &Tool{Name: "slow"},
-		Handler: slowRequest,
-	}
-	_, cs := basicConnection(t, st)
+	_, cs := basicConnection(t, func(s *Server) {
+		s.AddTool(&Tool{Name: "slow"}, slowRequest)
+	})
 	defer cs.Close()
 
 	ctx, cancel := context.WithCancel(context.Background())
@@ -852,7 +826,7 @@ func TestKeepAlive(t *testing.T) {
 		KeepAlive: 100 * time.Millisecond,
 	}
 	s := NewServer("testServer", "v1.0.0", serverOpts)
-	s.AddTools(NewServerTool("greet", "say hi", sayHi))
+	AddTool(s, greetTool(), sayHi)
 
 	ss, err := s.Connect(ctx, st)
 	if err != nil {
@@ -897,7 +871,7 @@ func TestKeepAliveFailure(t *testing.T) {
 
 	// Server without keepalive (to test one-sided keepalive)
 	s := NewServer("testServer", "v1.0.0", nil)
-	s.AddTools(NewServerTool("greet", "say hi", sayHi))
+	AddTool(s, greetTool(), sayHi)
 	ss, err := s.Connect(ctx, st)
 	if err != nil {
 		t.Fatal(err)

mcp/server.go

@@ -12,6 +12,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"iter"
+	"log"
 	"net/url"
 	"path/filepath"
 	"slices"
@@ -20,7 +21,6 @@ import (
 
 	"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
 	"github.com/modelcontextprotocol/go-sdk/internal/util"
-	"github.com/modelcontextprotocol/go-sdk/jsonschema"
 )
 
 const DefaultPageSize = 1000
@@ -37,7 +37,7 @@ type Server struct {
 
 	mu                      sync.Mutex
 	prompts                 *featureSet[*ServerPrompt]
-	tools                   *featureSet[*ServerTool]
+	tools                   *featureSet[*serverTool]
 	resources               *featureSet[*ServerResource]
 	resourceTemplates       *featureSet[*ServerResourceTemplate]
 	sessions                []*ServerSession
@@ -88,7 +88,7 @@ func NewServer(name, version string, opts *ServerOptions) *Server {
 		version:                 version,
 		opts:                    *opts,
 		prompts:                 newFeatureSet(func(p *ServerPrompt) string { return p.Prompt.Name }),
-		tools:                   newFeatureSet(func(t *ServerTool) string { return t.Tool.Name }),
+		tools:                   newFeatureSet(func(t *serverTool) string { return t.tool.Name }),
 		resources:               newFeatureSet(func(r *ServerResource) string { return r.Resource.URI }),
 		resourceTemplates:       newFeatureSet(func(t *ServerResourceTemplate) string { return t.ResourceTemplate.URITemplate }),
 		sendingMethodHandler_:   defaultSendingMethodHandler[*ServerSession],
@@ -118,55 +118,44 @@ func (s *Server) RemovePrompts(names ...string) {
 		func() bool { return s.prompts.remove(names...) })
 }
 
-// AddTools adds the given tools to the server,
-// replacing any with the same names.
-// The arguments must not be modified after this call.
-//
-// AddTools panics if errors are detected.
-func (s *Server) AddTools(tools ...*ServerTool) {
-	if err := s.addToolsErr(tools...); err != nil {
+// AddTool adds a [Tool] to the server, or replaces one with the same name.
+// The tool's input schema must be non-nil.
+// The Tool argument must not be modified after this call.
+func (s *Server) AddTool(t *Tool, h ToolHandler) {
+	// TODO(jba): This is a breaking behavior change. Add before v0.2.0?
+	if t.InputSchema == nil {
+		log.Printf("mcp: tool %q has a nil input schema. This will panic in a future release.", t.Name)
+		// panic(fmt.Sprintf("adding tool %q: nil input schema", t.Name))
+	}
+	if err := addToolErr(s, t, h); err != nil {
 		panic(err)
 	}
 }
 
-// addToolsErr is like [AddTools], but returns an error instead of panicking.
-func (s *Server) addToolsErr(tools ...*ServerTool) error {
-	// Only notify if something could change.
-	if len(tools) == 0 {
-		return nil
+// AddTool adds a [Tool] to the server, or replaces one with the same name.
+// If the tool's input schema is nil, it is set to the schema inferred from the In
+// type parameter, using [jsonschema.For].
+// If the tool's output schema is nil and the Out type parameter is not the empty
+// interface, then the output schema is set to the schema inferred from Out.
+// The Tool argument must not be modified after this call.
+func AddTool[In, Out any](s *Server, t *Tool, h ToolHandlerFor[In, Out]) {
+	if err := addToolErr(s, t, h); err != nil {
+		panic(err)
 	}
-	// Wrap the user's Handlers with rawHandlers that take a json.RawMessage.
-	for _, st := range tools {
-		if st.rawHandler == nil {
-			// This ServerTool was not created with NewServerTool.
-			if st.Handler == nil {
-				return fmt.Errorf("AddTools: tool %q has no handler", st.Tool.Name)
-			}
-			st.rawHandler = newRawHandler(st)
-			// Resolve the schemas, with no base URI. We don't expect tool schemas to
-			// refer outside of themselves.
-			if st.Tool.InputSchema != nil {
-				r, err := st.Tool.InputSchema.Resolve(&jsonschema.ResolveOptions{ValidateDefaults: true})
-				if err != nil {
-					return err
-				}
-				st.inputResolved = r
-			}
+}
 
-			// if st.Tool.OutputSchema != nil {
-			// 	st.outputResolved, err := st.Tool.OutputSchema.Resolve(&jsonschema.ResolveOptions{ValidateDefaults: true})
-			// 	if err != nil {
-			// 		return err
-			// 	}
-			// }
-		}
+func addToolErr[In, Out any](s *Server, t *Tool, h ToolHandlerFor[In, Out]) (err error) {
+	defer util.Wrapf(&err, "adding tool %q", t.Name)
+	st, err := newServerTool(t, h)
+	if err != nil {
+		return err
 	}
-
 	// Assume there was a change, since add replaces existing tools.
 	// (It's possible a tool was replaced with an identical one, but not worth checking.)
-	// TODO: surface notify error here?
+	// TODO: Batch these changes by size and time? The typescript SDK doesn't.
+	// TODO: Surface notify error here? best not, in case we need to batch.
 	s.changeAndNotify(notificationToolListChanged, &ToolListChangedParams{},
-		func() bool { s.tools.add(tools...); return true })
+		func() bool { s.tools.add(st); return true })
 	return nil
 }
 
@@ -293,22 +282,22 @@ func (s *Server) listTools(_ context.Context, _ *ServerSession, params *ListTool
 	if params == nil {
 		params = &ListToolsParams{}
 	}
-	return paginateList(s.tools, s.opts.PageSize, params, &ListToolsResult{}, func(res *ListToolsResult, tools []*ServerTool) {
+	return paginateList(s.tools, s.opts.PageSize, params, &ListToolsResult{}, func(res *ListToolsResult, tools []*serverTool) {
 		res.Tools = []*Tool{} // avoid JSON null
 		for _, t := range tools {
-			res.Tools = append(res.Tools, t.Tool)
+			res.Tools = append(res.Tools, t.tool)
 		}
 	})
 }
 
 func (s *Server) callTool(ctx context.Context, cc *ServerSession, params *CallToolParamsFor[json.RawMessage]) (*CallToolResult, error) {
 	s.mu.Lock()
-	tool, ok := s.tools.get(params.Name)
+	st, ok := s.tools.get(params.Name)
 	s.mu.Unlock()
 	if !ok {
 		return nil, fmt.Errorf("%s: unknown tool %q", jsonrpc2.ErrInvalidParams, params.Name)
 	}
-	return tool.rawHandler(ctx, cc, params)
+	return st.handler(ctx, cc, params)
 }
 
 func (s *Server) listResources(_ context.Context, _ *ServerSession, params *ListResourcesParams) (*ListResourcesResult, error) {

mcp/server_example_test.go

@@ -29,7 +29,7 @@ func ExampleServer() {
 	clientTransport, serverTransport := mcp.NewInMemoryTransports()
 
 	server := mcp.NewServer("greeter", "v0.0.1", nil)
-	server.AddTools(mcp.NewServerTool("greet", "say hi", SayHi))
+	mcp.AddTool(server, &mcp.Tool{Name: "greet", Description: "say hi"}, SayHi)
 
 	serverSession, err := server.Connect(ctx, serverTransport)
 	if err != nil {