internal/mcp: implement roots

Add support for roots.

Change-Id: Ia50abc88f0047238272d698f30ce615b1a8fd486
Reviewed-on: https://go-review.googlesource.com/c/tools/+/671360
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Sam Thanawalla <samthanawalla@google.com>
Reviewed-by: Robert Findley <rfindley@google.com>

mcp/client.go

@@ -9,6 +9,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"slices"
 	"sync"
 
 	"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
@@ -24,6 +25,7 @@ type Client struct {
 	opts             ClientOptions
 	mu               sync.Mutex
 	conn             *jsonrpc2.Connection
+	roots            *featureSet[protocol.Root]
 	initializeResult *protocol.InitializeResult
 }
 
@@ -37,6 +39,7 @@ func NewClient(name, version string, t Transport, opts *ClientOptions) *Client {
 		name:      name,
 		version:   version,
 		transport: t,
+		roots:     newFeatureSet(func(r protocol.Root) string { return r.URI }),
 	}
 	if opts != nil {
 		c.opts = *opts
@@ -106,13 +109,47 @@ func (c *Client) Wait() error {
 	return c.conn.Wait()
 }
 
-func (*Client) handle(ctx context.Context, req *jsonrpc2.Request) (any, error) {
+// AddRoots adds the given roots to the client,
+// replacing any with the same URIs,
+// and notifies any connected servers.
+// TODO: notification
+func (c *Client) AddRoots(roots ...protocol.Root) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.roots.add(roots...)
+}
+
+// RemoveRoots removes the roots with the given URIs,
+// and notifies any connected servers if the list has changed.
+// It is not an error to remove a nonexistent root.
+// TODO: notification
+func (c *Client) RemoveRoots(uris ...string) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.roots.remove(uris...)
+}
+
+func (c *Client) listRoots(_ context.Context, _ *protocol.ListRootsParams) (*protocol.ListRootsResult, error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return &protocol.ListRootsResult{
+		Roots: slices.Collect(c.roots.all()),
+	}, nil
+}
+
+func (c *Client) handle(ctx context.Context, req *jsonrpc2.Request) (any, error) {
+	// TODO: when we switch to ClientSessions, use a copy of the server's dispatch function, or
+	// maybe just add another type parameter.
+	//
 	// No need to check that the connection is initialized, since we initialize
 	// it in Connect.
 	switch req.Method {
 	case "ping":
 		// The spec says that 'ping' expects an empty object result.
 		return struct{}{}, nil
+	case "roots/list":
+		// ListRootsParams happens to be unused.
+		return c.listRoots(ctx, nil)
 	}
 	return nil, jsonrpc2.ErrNotHandled
 }
@@ -162,10 +199,6 @@ func (c *Client) ListTools(ctx context.Context) ([]protocol.Tool, error) {
 }
 
 // CallTool calls the tool with the given name and arguments.
-//
-// TODO(jba): make the following true:
-// If the provided arguments do not conform to the schema for the given tool,
-// the call fails.
 func (c *Client) CallTool(ctx context.Context, name string, args map[string]any) (_ *protocol.CallToolResult, err error) {
 	defer func() {
 		if err != nil {
@@ -180,14 +213,17 @@ func (c *Client) CallTool(ctx context.Context, name string, args map[string]any)
 		}
 		argsJSON[name] = argJSON
 	}
-	var (
-		params = &protocol.CallToolParams{
-			Name:      name,
-			Arguments: argsJSON,
-		}
-		result protocol.CallToolResult
-	)
-	if err := call(ctx, c.conn, "tools/call", params, &result); err != nil {
+
+	params := &protocol.CallToolParams{
+		Name:      name,
+		Arguments: argsJSON,
+	}
+	return standardCall[protocol.CallToolResult](ctx, c.conn, "tools/call", params)
+}
+
+func standardCall[TRes, TParams any](ctx context.Context, conn *jsonrpc2.Connection, method string, params TParams) (*TRes, error) {
+	var result TRes
+	if err := call(ctx, conn, method, params, &result); err != nil {
 		return nil, err
 	}
 	return &result, nil

mcp/mcp_test.go

@@ -63,7 +63,7 @@ func TestEndToEnd(t *testing.T) {
 	)
 
 	// Connect the server.
-	cc, err := s.Connect(ctx, st, nil)
+	sc, err := s.Connect(ctx, st, nil)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -75,13 +75,14 @@ func TestEndToEnd(t *testing.T) {
 	var clientWG sync.WaitGroup
 	clientWG.Add(1)
 	go func() {
-		if err := cc.Wait(); err != nil {
+		if err := sc.Wait(); err != nil {
 			t.Errorf("server failed: %v", err)
 		}
 		clientWG.Done()
 	}()
 
 	c := NewClient("testClient", "v1.0.0", ct, nil)
+	c.AddRoots(protocol.Root{URI: "file:///root"})
 
 	// Connect the client.
 	if err := c.Start(ctx); err != nil {
@@ -182,6 +183,13 @@ func TestEndToEnd(t *testing.T) {
 		t.Errorf("tools/call 'fail' mismatch (-want +got):\n%s", diff)
 	}
 
+	rootRes, err := sc.ListRoots(ctx, &protocol.ListRootsParams{})
+	gotRoots := rootRes.Roots
+	wantRoots := slices.Collect(c.roots.all())
+	if diff := cmp.Diff(wantRoots, gotRoots); diff != "" {
+		t.Errorf("roots/list mismatch (-want +got):\n%s", diff)
+	}
+
 	// Disconnect.
 	c.Close()
 	clientWG.Wait()

mcp/protocol/generate.go

@@ -75,6 +75,10 @@ var declarations = config{
 		Fields: config{"Params": {Name: "ListPromptsParams"}},
 	},
 	"ListPromptsResult": {Name: "ListPromptsResult"},
+	"ListRootsRequest": {
+		Fields: config{"Params": {Name: "ListRootsParams"}},
+	},
+	"ListRootsResult": {Name: "ListRootsResult"},
 	"ListToolsRequest": {
 		Fields: config{"Params": {Name: "ListToolsParams"}},
 	},
@@ -82,8 +86,11 @@ var declarations = config{
 	"Prompt":          {Name: "Prompt"},
 	"PromptMessage":   {Name: "PromptMessage"},
 	"PromptArgument":  {Name: "PromptArgument"},
+	"ProgressToken":   {Substitute: "any"}, // null|number|string
 	"RequestId":       {Substitute: "any"}, // null|number|string
 	"Role":            {Name: "Role"},
+	"Root":            {Name: "Root"},
+
 	"ServerCapabilities": {
 		Name: "ServerCapabilities",
 		Fields: config{
@@ -243,7 +250,11 @@ func writeType(w io.Writer, config *typeConfig, def *jsonschema.Schema, named ma
 	// For types that explicitly allow additional properties, we can either
 	// unmarshal them into a map[string]any, or delay unmarshalling with
 	// json.RawMessage. For now, use json.RawMessage as it defers the choice.
-	if def.Type == "object" && canHaveAdditionalProperties(def) {
+	//
+	// TODO(jba): further refine this classification of object schemas.
+	// For example, the typescript "object" type, which should map to a Go "any",
+	// is represented in schema.json by `{type: object, properties: {}, additionalProperties: true}`.
+	if def.Type == "object" && canHaveAdditionalProperties(def) && def.Properties == nil {
 		w.Write([]byte("map[string]"))
 		return writeType(w, nil, def.AdditionalProperties, named)
 	}

mcp/protocol/protocol.go

@@ -70,14 +70,16 @@ type CancelledParams struct {
 // additional capabilities.
 type ClientCapabilities struct {
 	// Experimental, non-standard capabilities that the client supports.
-	Experimental map[string]map[string]json.RawMessage `json:"experimental,omitempty"`
+	Experimental map[string]struct {
+	} `json:"experimental,omitempty"`
 	// Present if the client supports listing roots.
 	Roots *struct {
 		// Whether the client supports notifications for changes to the roots list.
 		ListChanged bool `json:"listChanged,omitempty"`
 	} `json:"roots,omitempty"`
 	// Present if the client supports sampling from an LLM.
-	Sampling map[string]json.RawMessage `json:"sampling,omitempty"`
+	Sampling struct {
+	} `json:"sampling,omitempty"`
 }
 
 type GetPromptParams struct {
@@ -131,7 +133,11 @@ type InitializeResult struct {
 	ServerInfo      Implementation `json:"serverInfo"`
 }
 
-type InitializedParams map[string]json.RawMessage
+type InitializedParams struct {
+	// This parameter name is reserved by MCP to allow clients and servers to attach
+	// additional metadata to their notifications.
+	Meta map[string]json.RawMessage `json:"_meta,omitempty"`
+}
 
 type ListPromptsParams struct {
 	// An opaque token representing the current pagination position. If provided,
@@ -150,6 +156,26 @@ type ListPromptsResult struct {
 	Prompts    []Prompt `json:"prompts"`
 }
 
+type ListRootsParams struct {
+	Meta *struct {
+		// If specified, the caller is requesting out-of-band progress notifications for
+		// this request (as represented by notifications/progress). The value of this
+		// parameter is an opaque token that will be attached to any subsequent
+		// notifications. The receiver is not obligated to provide these notifications.
+		ProgressToken *any `json:"progressToken,omitempty"`
+	} `json:"_meta,omitempty"`
+}
+
+// The client's response to a roots/list request from the server. This result
+// contains an array of Root objects, each representing a root directory or file
+// that the server can operate on.
+type ListRootsResult struct {
+	// This result property is reserved by the protocol to allow clients and servers
+	// to attach additional metadata to their responses.
+	Meta  map[string]json.RawMessage `json:"_meta,omitempty"`
+	Roots []Root                     `json:"roots"`
+}
+
 type ListToolsParams struct {
 	// An opaque token representing the current pagination position. If provided,
 	// the server should return results starting after this cursor.
@@ -213,16 +239,31 @@ type ResourceCapabilities struct {
 // The sender or recipient of messages and data in a conversation.
 type Role string
 
+// Represents a root directory or file that the server can operate on.
+type Root struct {
+	// An optional name for the root. This can be used to provide a human-readable
+	// identifier for the root, which may be useful for display purposes or for
+	// referencing the root in other parts of the application.
+	Name string `json:"name,omitempty"`
+	// The URI identifying the root. This *must* start with file:// for now. This
+	// restriction may be relaxed in future versions of the protocol to allow other
+	// URI schemes.
+	URI string `json:"uri"`
+}
+
 // Capabilities that a server may support. Known capabilities are defined here,
 // in this schema, but this is not a closed set: any server can define its own,
 // additional capabilities.
 type ServerCapabilities struct {
 	// Present if the server supports argument autocompletion suggestions.
-	Completions map[string]json.RawMessage `json:"completions,omitempty"`
+	Completions struct {
+	} `json:"completions,omitempty"`
 	// Experimental, non-standard capabilities that the server supports.
-	Experimental map[string]map[string]json.RawMessage `json:"experimental,omitempty"`
+	Experimental map[string]struct {
+	} `json:"experimental,omitempty"`
 	// Present if the server supports sending log messages to the client.
-	Logging map[string]json.RawMessage `json:"logging,omitempty"`
+	Logging struct {
+	} `json:"logging,omitempty"`
 	// Present if the server offers any prompt templates.
 	Prompts *PromptCapabilities `json:"prompts,omitempty"`
 	// Present if the server offers any resources to read.

mcp/root.go

@@ -0,0 +1,5 @@
+// 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

mcp/server.go

@@ -57,7 +57,8 @@ func NewServer(name, version string, opts *ServerOptions) *Server {
 	}
 }
 
-// AddPrompts adds the given prompts to the server.
+// AddPrompts adds the given prompts to the server,
+// replacing any with the same names.
 func (s *Server) AddPrompts(prompts ...*Prompt) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
@@ -67,7 +68,7 @@ func (s *Server) AddPrompts(prompts ...*Prompt) {
 	// TODO(rfindley): notify connected clients
 }
 
-// RemovePrompts removes if the prompts with the given names.
+// RemovePrompts removes the prompts with the given names.
 // It is not an error to remove a nonexistent prompt.
 func (s *Server) RemovePrompts(names ...string) {
 	s.mu.Lock()
@@ -77,9 +78,8 @@ func (s *Server) RemovePrompts(names ...string) {
 	}
 }
 
-// AddTools adds the given tools to the server.
-//
-// TODO(rfindley): notify connected clients of any changes.
+// AddTools adds the given tools to the server,
+// replacing any with the same names.
 func (s *Server) AddTools(tools ...*Tool) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
@@ -89,7 +89,7 @@ func (s *Server) AddTools(tools ...*Tool) {
 	// TODO(rfindley): notify connected clients
 }
 
-// RemoveTools removes if the tools with the given names.
+// RemoveTools removes the tools with the given names.
 // It is not an error to remove a nonexistent tool.
 func (s *Server) RemoveTools(names ...string) {
 	s.mu.Lock()
@@ -210,6 +210,10 @@ func (cc *ServerConnection) Ping(ctx context.Context) error {
 	return call(ctx, cc.conn, "ping", nil, nil)
 }
 
+func (cc *ServerConnection) ListRoots(ctx context.Context, params *protocol.ListRootsParams) (*protocol.ListRootsResult, error) {
+	return standardCall[protocol.ListRootsResult](ctx, cc.conn, "roots/list", params)
+}
+
 func (cc *ServerConnection) handle(ctx context.Context, req *jsonrpc2.Request) (any, error) {
 	cc.mu.Lock()
 	initialized := cc.initialized