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 <[email protected]>
Reviewed-by: Sam Thanawalla <[email protected]>
Reviewed-by: Robert Findley <[email protected]>

Author: jba

internal/mcp/client.go

@@ -9,6 +9,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"slices"
 	"sync"
 
 	jsonrpc2 "golang.org/x/tools/internal/jsonrpc2_v2"
@@ -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

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

internal/mcp/protocol/generate.go

@@ -75,15 +75,22 @@ 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"}},
 	},
 	"ListToolsResult": {Name: "ListToolsResult"},
 	"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)
 	}

internal/mcp/protocol/protocol.go

@@ -0,0 +1,5 @@
+// Copyright 2025 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package mcp

internal/mcp/root.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