From 5004d3e1bbd734300893349a596f9c48f1760618 Mon Sep 17 00:00:00 2001 From: Jonathan Amsterdam Date: Fri, 9 May 2025 22:04:50 -0400 Subject: [PATCH] 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 Reviewed-by: Sam Thanawalla Reviewed-by: Robert Findley --- mcp/client.go | 62 +++++++++++++++++++++++++++++++--------- mcp/mcp_test.go | 12 ++++++-- mcp/protocol/generate.go | 13 ++++++++- mcp/protocol/protocol.go | 53 ++++++++++++++++++++++++++++++---- mcp/root.go | 5 ++++ mcp/server.go | 16 +++++++---- 6 files changed, 133 insertions(+), 28 deletions(-) create mode 100644 mcp/root.go diff --git a/mcp/client.go b/mcp/client.go index 76794c1..30a1ac3 100644 --- a/mcp/client.go +++ b/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 diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index 60ac84c..177e9c1 100644 --- a/mcp/mcp_test.go +++ b/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() diff --git a/mcp/protocol/generate.go b/mcp/protocol/generate.go index e082502..09b5bd3 100644 --- a/mcp/protocol/generate.go +++ b/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) } diff --git a/mcp/protocol/protocol.go b/mcp/protocol/protocol.go index a993a62..af88ea4 100644 --- a/mcp/protocol/protocol.go +++ b/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. diff --git a/mcp/root.go b/mcp/root.go new file mode 100644 index 0000000..b56ad99 --- /dev/null +++ b/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 diff --git a/mcp/server.go b/mcp/server.go index 8060941..6e4dddb 100644 --- a/mcp/server.go +++ b/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