From 1ae23e70315c4efc45f14c16710b32ec6db676df Mon Sep 17 00:00:00 2001 From: Jonathan Amsterdam Date: Sat, 10 May 2025 09:03:39 -0400 Subject: [PATCH] internal/mcp: implement resources Implementation of resources. This follows the design sketched in design/design.md, except for Server.FileResourceHandler, which will be in a forthcoming CL. As the spec recommends, the server verifies an incoming URI from the client against its list of resources, to make sure the URI came from one of them. Some things that aren't in the design doc: - I provide an exported function that creates a "resource not found" jsonrpc2.WireEror error for server authors to return. A wrinkle: the MCP spec defines the error code for "resource not found" as -32002 [1]. That code is reserved for use by implementations of JSON-RPC, not applications. That matters to us because our JSON-RPC implementation already uses that code for the "server closing" error. I filed a bug against MCP about it [2] and chose a different code for the "resource not found" error. - As a slight convenience, I fill in some fields in the ReadResourceResult if the handler doesn't set them. [1] https://modelcontextprotocol.io/specification/2025-03-26/server/resources#error-handling [2] https://github.com/modelcontextprotocol/modelcontextprotocol/issues/509 Change-Id: I06153da982ea00748a3691ded923448a880cea2a Reviewed-on: https://go-review.googlesource.com/c/tools/+/671363 LUCI-TryBot-Result: Go LUCI Reviewed-by: Robert Findley --- mcp/client.go | 10 ++ mcp/mcp_test.go | 259 +++++++++++++++++++++++++-------------- mcp/protocol/generate.go | 17 ++- mcp/protocol/protocol.go | 56 +++++++++ mcp/server.go | 127 +++++++++++++++++-- mcp/transport.go | 4 +- 6 files changed, 369 insertions(+), 104 deletions(-) diff --git a/mcp/client.go b/mcp/client.go index 30a1ac3..b6b1d84 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -221,6 +221,16 @@ func (c *Client) CallTool(ctx context.Context, name string, args map[string]any) return standardCall[protocol.CallToolResult](ctx, c.conn, "tools/call", params) } +// ListResources lists the resources that are currently available on the server. +func (c *Client) ListResources(ctx context.Context, params *protocol.ListResourcesParams) (*protocol.ListResourcesResult, error) { + return standardCall[protocol.ListResourcesResult](ctx, c.conn, "resources/list", params) +} + +// ReadResource ask the server to read a resource and return its contents. +func (c *Client) ReadResource(ctx context.Context, params *protocol.ReadResourceParams) (*protocol.ReadResourceResult, error) { + return standardCall[protocol.ReadResourceResult](ctx, c.conn, "resources/read", 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 { diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index 177e9c1..baa419f 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -16,6 +16,7 @@ import ( "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" "github.com/modelcontextprotocol/go-sdk/mcp/protocol" ) @@ -92,103 +93,179 @@ func TestEndToEnd(t *testing.T) { if err := c.Ping(ctx); err != nil { t.Fatalf("ping failed: %v", err) } - - gotPrompts, err := c.ListPrompts(ctx) - if err != nil { - t.Errorf("prompts/list failed: %v", err) - } - wantPrompts := []protocol.Prompt{ - { - Name: "code_review", - Description: "do a code review", - Arguments: []protocol.PromptArgument{{Name: "Code", Required: true}}, - }, - {Name: "fail"}, - } - if diff := cmp.Diff(wantPrompts, gotPrompts); diff != "" { - t.Fatalf("prompts/list mismatch (-want +got):\n%s", diff) - } - - gotReview, err := c.GetPrompt(ctx, "code_review", map[string]string{"Code": "1+1"}) - if err != nil { - t.Fatal(err) - } - wantReview := &protocol.GetPromptResult{ - Description: "Code review prompt", - Messages: []protocol.PromptMessage{{ - Content: TextContent{Text: "Please review the following code: 1+1"}.ToWire(), - Role: "user", - }}, - } - if diff := cmp.Diff(wantReview, gotReview); diff != "" { - t.Errorf("prompts/get 'code_review' mismatch (-want +got):\n%s", diff) - } - - if _, err := c.GetPrompt(ctx, "fail", map[string]string{}); err == nil || !strings.Contains(err.Error(), failure.Error()) { - t.Errorf("fail returned unexpected error: got %v, want containing %v", err, failure) - } - - gotTools, err := c.ListTools(ctx) - if err != nil { - t.Errorf("tools/list failed: %v", err) - } - wantTools := []protocol.Tool{ - { - Name: "fail", - Description: "just fail", - InputSchema: &jsonschema.Schema{ - Type: "object", - AdditionalProperties: falseSchema, + t.Run("prompts", func(t *testing.T) { + gotPrompts, err := c.ListPrompts(ctx) + if err != nil { + t.Errorf("prompts/list failed: %v", err) + } + wantPrompts := []protocol.Prompt{ + { + Name: "code_review", + Description: "do a code review", + Arguments: []protocol.PromptArgument{{Name: "Code", Required: true}}, }, - }, - { - Name: "greet", - Description: "say hi", - InputSchema: &jsonschema.Schema{ - Type: "object", - Required: []string{"Name"}, - Properties: map[string]*jsonschema.Schema{ - "Name": {Type: "string"}, + {Name: "fail"}, + } + if diff := cmp.Diff(wantPrompts, gotPrompts); diff != "" { + t.Fatalf("prompts/list mismatch (-want +got):\n%s", diff) + } + + gotReview, err := c.GetPrompt(ctx, "code_review", map[string]string{"Code": "1+1"}) + if err != nil { + t.Fatal(err) + } + wantReview := &protocol.GetPromptResult{ + Description: "Code review prompt", + Messages: []protocol.PromptMessage{{ + Content: TextContent{Text: "Please review the following code: 1+1"}.ToWire(), + Role: "user", + }}, + } + if diff := cmp.Diff(wantReview, gotReview); diff != "" { + t.Errorf("prompts/get 'code_review' mismatch (-want +got):\n%s", diff) + } + + if _, err := c.GetPrompt(ctx, "fail", map[string]string{}); err == nil || !strings.Contains(err.Error(), failure.Error()) { + t.Errorf("fail returned unexpected error: got %v, want containing %v", err, failure) + } + }) + + t.Run("tools", func(t *testing.T) { + gotTools, err := c.ListTools(ctx) + if err != nil { + t.Errorf("tools/list failed: %v", err) + } + wantTools := []protocol.Tool{ + { + Name: "fail", + Description: "just fail", + InputSchema: &jsonschema.Schema{ + Type: "object", + AdditionalProperties: falseSchema, }, - AdditionalProperties: falseSchema, }, - }, - } - if diff := cmp.Diff(wantTools, gotTools, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" { - t.Fatalf("tools/list mismatch (-want +got):\n%s", diff) - } + { + 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, gotTools, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" { + t.Fatalf("tools/list mismatch (-want +got):\n%s", diff) + } - gotHi, err := c.CallTool(ctx, "greet", map[string]any{"name": "user"}) - if err != nil { - t.Fatal(err) - } - wantHi := &protocol.CallToolResult{ - Content: []protocol.Content{{Type: "text", Text: "hi user"}}, - } - if diff := cmp.Diff(wantHi, gotHi); diff != "" { - t.Errorf("tools/call 'greet' mismatch (-want +got):\n%s", diff) - } + gotHi, err := c.CallTool(ctx, "greet", map[string]any{"name": "user"}) + if err != nil { + t.Fatal(err) + } + wantHi := &protocol.CallToolResult{ + Content: []protocol.Content{{Type: "text", Text: "hi user"}}, + } + if diff := cmp.Diff(wantHi, gotHi); diff != "" { + t.Errorf("tools/call 'greet' mismatch (-want +got):\n%s", diff) + } - gotFail, err := c.CallTool(ctx, "fail", 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 { - t.Fatal(err) - } - wantFail := &protocol.CallToolResult{ - IsError: true, - Content: []protocol.Content{{Type: "text", Text: failure.Error()}}, - } - if diff := cmp.Diff(wantFail, gotFail); diff != "" { - t.Errorf("tools/call 'fail' mismatch (-want +got):\n%s", diff) - } + gotFail, err := c.CallTool(ctx, "fail", 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 { + t.Fatal(err) + } + wantFail := &protocol.CallToolResult{ + IsError: true, + Content: []protocol.Content{{Type: "text", Text: failure.Error()}}, + } + if diff := cmp.Diff(wantFail, gotFail); diff != "" { + 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) - } + t.Run("resources", func(t *testing.T) { + resource1 := protocol.Resource{ + Name: "public", + MIMEType: "text/plain", + URI: "file:///file1.txt", + } + resource2 := protocol.Resource{ + Name: "public", // names are not unique IDs + MIMEType: "text/plain", + URI: "file:///nonexistent.txt", + } + + readHandler := func(_ context.Context, r protocol.Resource, _ *protocol.ReadResourceParams) (*protocol.ReadResourceResult, error) { + if r.URI == "file:///file1.txt" { + return &protocol.ReadResourceResult{ + Contents: &protocol.ResourceContents{ + Text: "file contents", + }, + }, nil + } + return nil, ResourceNotFoundError(r.URI) + } + s.AddResources( + &ServerResource{resource1, readHandler}, + &ServerResource{resource2, readHandler}) + + lrres, err := c.ListResources(ctx, nil) + if err != nil { + t.Fatal(err) + } + if diff := cmp.Diff([]protocol.Resource{resource1, resource2}, lrres.Resources); diff != "" { + t.Errorf("resources/list mismatch (-want, +got):\n%s", diff) + } + + for _, tt := range []struct { + uri string + mimeType string // "": not found; "text/plain": resource; "text/template": template + }{ + {"file:///file1.txt", "text/plain"}, + {"file:///nonexistent.txt", ""}, + // TODO(jba): add resource template cases when we implement them + } { + rres, err := c.ReadResource(ctx, &protocol.ReadResourceParams{URI: tt.uri}) + if err != nil { + var werr *jsonrpc2.WireError + if errors.As(err, &werr) && werr.Code == codeResourceNotFound { + if tt.mimeType != "" { + t.Errorf("%s: not found but expected it to be", tt.uri) + } + } else { + t.Fatalf("reading %s: %v", tt.uri, err) + } + } else { + if got := rres.Contents.URI; got != tt.uri { + t.Errorf("got uri %q, want %q", got, tt.uri) + } + if got := rres.Contents.MIMEType; got != tt.mimeType { + t.Errorf("%s: got MIME type %q, want %q", tt.uri, got, tt.mimeType) + } + } + } + }) + t.Run("roots", func(t *testing.T) { + // Take the server's first ServerConnection. + var sc *ServerConnection + for sc = range s.Clients() { + break + } + + rootRes, err := sc.ListRoots(ctx, &protocol.ListRootsParams{}) + if err != nil { + t.Fatal(err) + } + 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() diff --git a/mcp/protocol/generate.go b/mcp/protocol/generate.go index b955948..887245f 100644 --- a/mcp/protocol/generate.go +++ b/mcp/protocol/generate.go @@ -82,6 +82,11 @@ var declarations = config{ Fields: config{"Params": {Name: "ListPromptsParams"}}, }, "ListPromptsResult": {}, + "ListResourcesRequest": { + Name: "-", + Fields: config{"Params": {Name: "ListResourcesParams"}}, + }, + "ListResourcesResult": {}, "ListRootsRequest": { Name: "-", Fields: config{"Params": {Name: "ListRootsParams"}}, @@ -97,8 +102,16 @@ var declarations = config{ "PromptArgument": {}, "ProgressToken": {Name: "-", Substitute: "any"}, // null|number|string "RequestId": {Name: "-", Substitute: "any"}, // null|number|string - "Role": {}, - "Root": {}, + "ReadResourceRequest": { + Name: "-", + Fields: config{"Params": {Name: "ReadResourceParams"}}, + }, + "ReadResourceResult": { + Fields: config{"Contents": {Substitute: "*ResourceContents"}}, + }, + "Resource": {}, + "Role": {}, + "Root": {}, "ServerCapabilities": { Fields: config{ diff --git a/mcp/protocol/protocol.go b/mcp/protocol/protocol.go index af88ea4..f43702f 100644 --- a/mcp/protocol/protocol.go +++ b/mcp/protocol/protocol.go @@ -156,6 +156,23 @@ type ListPromptsResult struct { Prompts []Prompt `json:"prompts"` } +type ListResourcesParams struct { + // An opaque token representing the current pagination position. If provided, + // the server should return results starting after this cursor. + Cursor string `json:"cursor,omitempty"` +} + +// The server's response to a resources/list request from the client. +type ListResourcesResult 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"` + // An opaque token representing the pagination position after the last returned + // result. If present, there may be more results available. + NextCursor string `json:"nextCursor,omitempty"` + Resources []Resource `json:"resources"` +} + type ListRootsParams struct { Meta *struct { // If specified, the caller is requesting out-of-band progress notifications for @@ -228,6 +245,45 @@ type PromptMessage struct { Role Role `json:"role"` } +type ReadResourceParams struct { + // The URI of the resource to read. The URI can use any protocol; it is up to + // the server how to interpret it. + URI string `json:"uri"` +} + +// The server's response to a resources/read request from the client. +type ReadResourceResult 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"` + Contents *ResourceContents `json:"contents"` +} + +// A known resource that the server is capable of reading. +type Resource struct { + // Optional annotations for the client. + Annotations *Annotations `json:"annotations,omitempty"` + // A description of what this resource represents. + // + // This can be used by clients to improve the LLM's understanding of available + // resources. It can be thought of like a "hint" to the model. + Description string `json:"description,omitempty"` + // The MIME type of this resource, if known. + MIMEType string `json:"mimeType,omitempty"` + // A human-readable name for this resource. + // + // This can be used by clients to populate UI elements. + Name string `json:"name"` + // The size of the raw resource content, in bytes (i.e., before base64 encoding + // or any tokenization), if known. + // + // This can be used by Hosts to display file sizes and estimate context window + // usage. + Size int64 `json:"size,omitempty"` + // The URI of this resource. + URI string `json:"uri"` +} + // Present if the server offers any resources to read. type ResourceCapabilities struct { // Whether this server supports notifications for changes to the resource list. diff --git a/mcp/server.go b/mcp/server.go index 6e4dddb..a0d0444 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -9,6 +9,8 @@ import ( "encoding/json" "fmt" "iter" + "log" + "net/url" "slices" "sync" @@ -26,10 +28,11 @@ type Server struct { version string opts ServerOptions - mu sync.Mutex - prompts *featureSet[*Prompt] - tools *featureSet[*Tool] - conns []*ServerConnection + mu sync.Mutex + prompts *featureSet[*Prompt] + tools *featureSet[*Tool] + resources *featureSet[*ServerResource] + conns []*ServerConnection } // ServerOptions is used to configure behavior of the server. @@ -49,11 +52,12 @@ func NewServer(name, version string, opts *ServerOptions) *Server { opts = new(ServerOptions) } return &Server{ - name: name, - version: version, - opts: *opts, - prompts: newFeatureSet(func(p *Prompt) string { return p.Definition.Name }), - tools: newFeatureSet(func(t *Tool) string { return t.Definition.Name }), + name: name, + version: version, + opts: *opts, + prompts: newFeatureSet(func(p *Prompt) string { return p.Definition.Name }), + tools: newFeatureSet(func(t *Tool) string { return t.Definition.Name }), + resources: newFeatureSet(func(r *ServerResource) string { return r.Resource.URI }), } } @@ -99,6 +103,62 @@ func (s *Server) RemoveTools(names ...string) { } } +// ResourceNotFoundError returns an error indicating that a resource being read could +// not be found. +func ResourceNotFoundError(uri string) error { + return &jsonrpc2.WireError{ + Code: codeResourceNotFound, + Message: "Resource not found", + Data: json.RawMessage(fmt.Sprintf(`{"uri":%q}`, uri)), + } +} + +// The error code to return when a resource isn't found. +// See https://modelcontextprotocol.io/specification/2025-03-26/server/resources#error-handling +// However, the code they chose in in the wrong space +// (see https://github.com/modelcontextprotocol/modelcontextprotocol/issues/509). +// so we pick a different one, arbirarily for now (until they fix it). +// The immediate problem is that jsonprc2 defines -32002 as "server closing". +const codeResourceNotFound = -31002 + +// A ReadResourceHandler is a function that reads a resource. +// If it cannot find the resource, it should return the result of calling [ResourceNotFoundError]. +type ReadResourceHandler func(context.Context, protocol.Resource, *protocol.ReadResourceParams) (*protocol.ReadResourceResult, error) + +// A ServerResource associates a Resource with its handler. +type ServerResource struct { + Resource protocol.Resource + Handler ReadResourceHandler +} + +// AddResource adds the given resource to the server and associates it with +// a [ReadResourceHandler], which will be called when the client calls [ClientSession.ReadResource]. +// If a resource with the same URI already exists, this one replaces it. +// AddResource panics if a resource URI is invalid or not absolute (has an empty scheme). +func (s *Server) AddResources(resources ...*ServerResource) { + s.mu.Lock() + defer s.mu.Unlock() + for _, r := range resources { + u, err := url.Parse(r.Resource.URI) + if err != nil { + panic(err) // url.Parse includes the URI in the error + } + if !u.IsAbs() { + panic(fmt.Errorf("URI %s needs a scheme", r.Resource.URI)) + } + s.resources.add(r) + } + // TODO: notify +} + +// RemoveResources removes the resources with the given URIs. +// It is not an error to remove a nonexistent resource. +func (s *Server) RemoveResources(uris ...string) { + s.mu.Lock() + defer s.mu.Unlock() + s.resources.remove(uris...) +} + // Clients returns an iterator that yields the current set of client // connections. func (s *Server) Clients() iter.Seq[*ServerConnection] { @@ -149,6 +209,47 @@ func (s *Server) callTool(ctx context.Context, cc *ServerConnection, params *pro return tool.Handler(ctx, cc, params.Arguments) } +func (s *Server) listResources(_ context.Context, _ *ServerConnection, params *protocol.ListResourcesParams) (*protocol.ListResourcesResult, error) { + s.mu.Lock() + defer s.mu.Unlock() + res := new(protocol.ListResourcesResult) + for r := range s.resources.all() { + res.Resources = append(res.Resources, r.Resource) + } + return res, nil +} + +func (s *Server) readResource(ctx context.Context, _ *ServerConnection, params *protocol.ReadResourceParams) (*protocol.ReadResourceResult, error) { + log.Printf("readResource") + defer log.Printf("done") + uri := params.URI + // Look up the resource URI in the list we have. + // This is a security check as well as an information lookup. + s.mu.Lock() + resource, ok := s.resources.get(uri) + s.mu.Unlock() + if !ok { + // Don't expose the server configuration to the client. + // Treat an unregistered resource the same as a registered one that couldn't be found. + return nil, ResourceNotFoundError(uri) + } + res, err := resource.Handler(ctx, resource.Resource, params) + if err != nil { + return nil, err + } + if res == nil || res.Contents == nil { + return nil, fmt.Errorf("reading resource %s: read handler returned nil information", uri) + } + // As a convenience, populate some fields. + if res.Contents.URI == "" { + res.Contents.URI = uri + } + if res.Contents.MIMEType == "" { + res.Contents.MIMEType = resource.Resource.MIMEType + } + return res, nil +} + // Run runs the server over the given transport, which must be persistent. // // Run blocks until the client terminates the connection. @@ -226,7 +327,7 @@ func (cc *ServerConnection) handle(ctx context.Context, req *jsonrpc2.Request) ( case "initialize", "ping": default: if !initialized { - return nil, fmt.Errorf("method %q is invalid during session ininitialization", req.Method) + return nil, fmt.Errorf("method %q is invalid during session initialization", req.Method) } } @@ -254,6 +355,12 @@ func (cc *ServerConnection) handle(ctx context.Context, req *jsonrpc2.Request) ( case "tools/call": return dispatch(ctx, cc, req, cc.server.callTool) + case "resources/list": + return dispatch(ctx, cc, req, cc.server.listResources) + + case "resources/read": + return dispatch(ctx, cc, req, cc.server.readResource) + case "notifications/initialized": } return nil, jsonrpc2.ErrNotHandled diff --git a/mcp/transport.go b/mcp/transport.go index 0b9aa0b..bf474b0 100644 --- a/mcp/transport.go +++ b/mcp/transport.go @@ -147,6 +147,8 @@ func (c *canceller) Preempt(ctx context.Context, req *jsonrpc2.Request) (result // call executes and awaits a jsonrpc2 call on the given connection, // translating errors into the mcp domain. func call(ctx context.Context, conn *jsonrpc2.Connection, method string, params, result any) error { + // TODO: the "%w"s in this function effectively make jsonrpc2.WireError part of the API. + // Consider alternatives. call := conn.Call(ctx, method, params) err := call.Await(ctx, result) switch { @@ -160,7 +162,7 @@ func call(ctx context.Context, conn *jsonrpc2.Connection, method string, params, }) return errors.Join(ctx.Err(), err) case err != nil: - return fmt.Errorf("calling %q: %v", method, err) + return fmt.Errorf("calling %q: %w", method, err) } return nil }