From 589e71f5d9e13c78db5977f6a757d68ae694cb16 Mon Sep 17 00:00:00 2001 From: po-et <42566883+po-et@users.noreply.github.com> Date: Fri, 9 Oct 2026 11:09:17 +0800 Subject: [PATCH 1/2] mcp: check requests against declared client capabilities Reject undeclared tool sampling and elicitation modes before invoking client handlers, and check roots and sampling before sending server requests. Fixes #1303 --- mcp/capabilities_test.go | 313 ++++++++++++++++++ mcp/client.go | 61 +++- mcp/server.go | 32 ++ mcp/streamable_test.go | 6 +- .../conformance/server/resources.txtar | 2 +- 5 files changed, 394 insertions(+), 20 deletions(-) diff --git a/mcp/capabilities_test.go b/mcp/capabilities_test.go index 86c997010..eab17d509 100644 --- a/mcp/capabilities_test.go +++ b/mcp/capabilities_test.go @@ -8,12 +8,15 @@ package mcp import ( "context" + "encoding/json" + "strings" "sync/atomic" "testing" "testing/synctest" "time" "github.com/google/jsonschema-go/jsonschema" + "github.com/modelcontextprotocol/go-sdk/jsonrpc" ) // TestServerListChangedNotifications verifies that listChanged notifications @@ -182,3 +185,313 @@ func TestClientListChangedNotifications(t *testing.T) { }) } } + +var ( + calculatorTools = []*Tool{{Name: "calculator", InputSchema: map[string]any{"type": "object"}}} + samplingMessages = []*SamplingMessageV2{{Role: "user", Content: []Content{&TextContent{Text: "hi"}}}} +) + +// A ServerSession doesn't send a request that needs a capability the client +// didn't declare. In particular, servers must not send tool-enabled sampling +// requests to a client that hasn't declared sampling.tools (2025-11-25 +// client/sampling, "Tools in Sampling"). +func TestServerSessionChecksClientCapabilities(t *testing.T) { + basicSampling := func(context.Context, *CreateMessageRequest) (*CreateMessageResult, error) { + return &CreateMessageResult{Model: "m", Content: &TextContent{}}, nil + } + toolSampling := func(context.Context, *CreateMessageWithToolsRequest) (*CreateMessageWithToolsResult, error) { + return &CreateMessageWithToolsResult{Model: "m", Content: []Content{&TextContent{}}}, nil + } + listRoots := func(ctx context.Context, ss *ServerSession) error { + _, err := ss.ListRoots(ctx, nil) + return err + } + createMessage := func(ctx context.Context, ss *ServerSession) error { + _, err := ss.CreateMessage(ctx, &CreateMessageParams{MaxTokens: 1}) + return err + } + createMessageWithTools := func(params *CreateMessageWithToolsParams) func(context.Context, *ServerSession) error { + return func(ctx context.Context, ss *ServerSession) error { + _, err := ss.CreateMessageWithTools(ctx, params) + return err + } + } + withTools := &CreateMessageWithToolsParams{MaxTokens: 1, Messages: samplingMessages, Tools: calculatorTools} + withToolChoice := &CreateMessageWithToolsParams{MaxTokens: 1, Messages: samplingMessages, ToolChoice: &ToolChoice{Mode: "none"}} + + for _, tc := range []struct { + name string + opts *ClientOptions + call func(context.Context, *ServerSession) error + wantErr string // empty if the request goes out + }{ + {name: "roots", call: listRoots}, + {name: "roots not declared", opts: &ClientOptions{Capabilities: &ClientCapabilities{}}, call: listRoots, wantErr: "does not support roots"}, + {name: "sampling", opts: &ClientOptions{CreateMessageHandler: basicSampling}, call: createMessage}, + {name: "sampling not declared", call: createMessage, wantErr: "does not support sampling"}, + {name: "sampling not declared, with tools API", call: createMessageWithTools(&CreateMessageWithToolsParams{MaxTokens: 1}), wantErr: "does not support sampling"}, + {name: "tools", opts: &ClientOptions{CreateMessageWithToolsHandler: toolSampling}, call: createMessageWithTools(withTools)}, + {name: "tools not declared", opts: &ClientOptions{CreateMessageHandler: basicSampling}, call: createMessageWithTools(withTools), wantErr: "does not support sampling with tools"}, + {name: "toolChoice not declared", opts: &ClientOptions{CreateMessageHandler: basicSampling}, call: createMessageWithTools(withToolChoice), wantErr: "does not support sampling with tools"}, + { + name: "tools not in explicit capabilities", + opts: &ClientOptions{ + CreateMessageWithToolsHandler: toolSampling, + Capabilities: &ClientCapabilities{Sampling: &SamplingCapabilities{}}, + }, + call: createMessageWithTools(withTools), + wantErr: "does not support sampling with tools", + }, + } { + t.Run(tc.name, func(t *testing.T) { + ctx := context.Background() + var received atomic.Int32 + client := NewClient(testImpl, tc.opts) + client.AddReceivingMiddleware(func(next MethodHandler) MethodHandler { + return func(ctx context.Context, method string, req Request) (Result, error) { + if method == methodListRoots || method == methodCreateMessage { + received.Add(1) + } + return next(ctx, method, req) + } + }) + ct, st := NewInMemoryTransports() + ss, err := NewServer(testImpl, nil).Connect(ctx, st, nil) + if err != nil { + t.Fatal(err) + } + defer ss.Close() + cs, err := client.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) + if err != nil { + t.Fatal(err) + } + defer cs.Close() + + err = tc.call(ctx, ss) + if tc.wantErr == "" { + if err != nil { + t.Fatalf("got error %v, want none", err) + } + if got := received.Load(); got != 1 { + t.Errorf("client received %d requests, want 1", got) + } + return + } + if err == nil || !strings.Contains(err.Error(), tc.wantErr) { + t.Fatalf("got error %v, want one containing %q", err, tc.wantErr) + } + if got := received.Load(); got != 0 { + t.Errorf("client received %d requests, want none", got) + } + }) + } +} + +// A client answers a request that uses a capability it didn't declare with +// invalid params, before its handler sees it: a sampling request with tools +// (the tools field of CreateMessageRequestParams in the 2025-11-25 schema) or +// an elicitation request for an undeclared mode (2025-11-25 +// client/elicitation, "Error Handling"). +func TestClientChecksDeclaredCapabilities(t *testing.T) { + var emptyTools CreateMessageWithToolsParams + if err := json.Unmarshal([]byte(`{"maxTokens":1,"messages":[],"tools":[]}`), &emptyTools); err != nil { + t.Fatal(err) + } + formOnly := &ClientCapabilities{Elicitation: &ElicitationCapabilities{Form: &FormElicitationCapabilities{}}} + urlOnly := &ClientCapabilities{Elicitation: &ElicitationCapabilities{URL: &URLElicitationCapabilities{}}} + modeless := &ClientCapabilities{Elicitation: &ElicitationCapabilities{}} + form := &ElicitParams{Message: "m", RequestedSchema: &jsonschema.Schema{Type: "object"}} + url := &ElicitParams{Mode: "url", Message: "m", URL: "https://example.com/form", ElicitationID: "e1"} + + for _, tc := range []struct { + name string + samplingAPI string // "basic", "tools" or "" for no sampling handler + caps *ClientCapabilities + sampling *CreateMessageWithToolsParams + elicit *ElicitParams + wantCode int64 // 0 if the handler is called + }{ + {name: "tools to basic handler", samplingAPI: "basic", sampling: &CreateMessageWithToolsParams{MaxTokens: 1, Messages: samplingMessages, Tools: calculatorTools}, wantCode: jsonrpc.CodeInvalidParams}, + {name: "toolChoice to basic handler", samplingAPI: "basic", sampling: &CreateMessageWithToolsParams{MaxTokens: 1, Messages: samplingMessages, ToolChoice: &ToolChoice{Mode: "auto"}}, wantCode: jsonrpc.CodeInvalidParams}, + {name: "empty tools array to basic handler", samplingAPI: "basic", sampling: &emptyTools, wantCode: jsonrpc.CodeInvalidParams}, + {name: "no tools to basic handler", samplingAPI: "basic", sampling: &CreateMessageWithToolsParams{MaxTokens: 1, Messages: samplingMessages}}, + {name: "tools to tools handler", samplingAPI: "tools", sampling: &CreateMessageWithToolsParams{MaxTokens: 1, Messages: samplingMessages, Tools: calculatorTools}}, + {name: "tools to tools handler without declared tools", samplingAPI: "tools", caps: &ClientCapabilities{Sampling: &SamplingCapabilities{}}, sampling: &CreateMessageWithToolsParams{MaxTokens: 1, Messages: samplingMessages, Tools: calculatorTools}, wantCode: jsonrpc.CodeInvalidParams}, + {name: "tools without a sampling handler", sampling: &CreateMessageWithToolsParams{MaxTokens: 1, Messages: samplingMessages, Tools: calculatorTools}, wantCode: codeUnsupportedMethod}, + {name: "url to inferred", elicit: url, wantCode: jsonrpc.CodeInvalidParams}, + {name: "form to inferred", elicit: form}, + {name: "url to form only", caps: formOnly, elicit: url, wantCode: jsonrpc.CodeInvalidParams}, + {name: "url to modeless", caps: modeless, elicit: url, wantCode: jsonrpc.CodeInvalidParams}, + {name: "form to modeless", caps: modeless, elicit: form}, + {name: "form to url only", caps: urlOnly, elicit: form, wantCode: jsonrpc.CodeInvalidParams}, + {name: "url to url only", caps: urlOnly, elicit: url}, + } { + t.Run(tc.name, func(t *testing.T) { + var called bool + opts := &ClientOptions{Capabilities: tc.caps} + switch tc.samplingAPI { + case "basic": + opts.CreateMessageHandler = func(context.Context, *CreateMessageRequest) (*CreateMessageResult, error) { + called = true + return &CreateMessageResult{Model: "m", Content: &TextContent{}}, nil + } + case "tools": + opts.CreateMessageWithToolsHandler = func(context.Context, *CreateMessageWithToolsRequest) (*CreateMessageWithToolsResult, error) { + called = true + return &CreateMessageWithToolsResult{Model: "m", Content: []Content{&TextContent{}}}, nil + } + } + if tc.elicit != nil { + opts.ElicitationHandler = func(context.Context, *ElicitRequest) (*ElicitResult, error) { + called = true + return &ElicitResult{Action: "decline"}, nil + } + } + c := NewClient(testImpl, opts) + + var err error + if tc.sampling != nil { + _, err = c.createMessage(context.Background(), &CreateMessageWithToolsRequest{Params: tc.sampling}) + } else { + _, err = c.elicit(context.Background(), &ElicitRequest{Params: tc.elicit}) + } + if tc.wantCode == 0 { + if err != nil || !called { + t.Errorf("got error %v, handler called = %t; want the handler's result", err, called) + } + return + } + if code := errorCode(err); code != tc.wantCode { + t.Errorf("got error %v (code %d), want code %d", err, code, tc.wantCode) + } + if called { + t.Error("the handler was called") + } + }) + } +} + +// Input requests of a multi-round-trip result go through the same checks: +// the client's own on protocol version 2026-07-28, and the server's where it +// fulfills them for an older client. +func TestInputRequestsChecked(t *testing.T) { + newServer := func() *Server { + s := NewServer(testImpl, nil) + AddTool(s, &Tool{Name: "sample"}, func(context.Context, *CallToolRequest, struct{}) (*CallToolResult, any, error) { + return &CallToolResult{InputRequests: InputRequestMap{ + "s": &CreateMessageWithToolsParams{MaxTokens: 1, Messages: samplingMessages, Tools: calculatorTools}, + }}, nil, nil + }) + return s + } + for _, tc := range []struct { + name string + connect func(*testing.T, *Server, *ClientOptions) *ClientSession + }{ + {"2026-07-28", mustConnect}, + {"2025-11-25", mustConnectOldProtocol}, + } { + t.Run(tc.name, func(t *testing.T) { + var called atomic.Bool + cs := tc.connect(t, newServer(), &ClientOptions{ + CreateMessageHandler: func(context.Context, *CreateMessageRequest) (*CreateMessageResult, error) { + called.Store(true) + return &CreateMessageResult{Model: "m", Content: &TextContent{}}, nil + }, + }) + res, err := cs.CallTool(t.Context(), &CallToolParams{Name: "sample"}) + if err == nil && !res.IsError { + t.Fatalf("CallTool succeeded, want a failure about sampling with tools") + } + if err != nil && !strings.Contains(err.Error(), "does not support sampling with tools") { + t.Errorf("CallTool error = %v, want one about sampling with tools", err) + } + if called.Load() { + t.Error("the client's handler was called") + } + }) + } +} + +func TestMultiRoundTripDeclaredCapabilities(t *testing.T) { + for _, tc := range []struct { + name string + input InputRequest + caps *ClientCapabilities + wantErr bool + }{ + { + name: "sampling with tools", + input: &CreateMessageWithToolsParams{MaxTokens: 1, Messages: samplingMessages, Tools: calculatorTools}, + }, + { + name: "form elicitation", + input: &ElicitParams{Message: "Continue?"}, + }, + { + name: "url elicitation", + input: &ElicitParams{Mode: "url", Message: "Sign in", URL: "https://example.com/login", ElicitationID: "e1"}, + caps: &ClientCapabilities{Elicitation: &ElicitationCapabilities{URL: &URLElicitationCapabilities{}}}, + }, + { + name: "url to form-only client", + input: &ElicitParams{Mode: "url", Message: "Sign in", URL: "https://example.com/login", ElicitationID: "e1"}, + wantErr: true, + }, + { + name: "form to url-only client", + input: &ElicitParams{Message: "Continue?"}, + caps: &ClientCapabilities{Elicitation: &ElicitationCapabilities{URL: &URLElicitationCapabilities{}}}, + wantErr: true, + }, + } { + for _, conn := range []struct { + version string + connect func(*testing.T, *Server, *ClientOptions) *ClientSession + }{ + {protocolVersion20260728, mustConnect}, + {protocolVersion20251125, mustConnectOldProtocol}, + } { + t.Run(tc.name+"/"+conn.version, func(t *testing.T) { + var calls atomic.Int32 + s := NewServer(testImpl, nil) + AddTool(s, &Tool{Name: "ask"}, func(_ context.Context, req *CallToolRequest, _ struct{}) (*CallToolResult, any, error) { + if len(req.Params.InputResponses) == 0 { + return &CallToolResult{InputRequests: InputRequestMap{"q": tc.input}}, nil, nil + } + return &CallToolResult{Content: []Content{&TextContent{Text: "done"}}}, nil, nil + }) + opts := &ClientOptions{Capabilities: tc.caps} + switch tc.input.(type) { + case *CreateMessageWithToolsParams: + opts.CreateMessageWithToolsHandler = func(context.Context, *CreateMessageWithToolsRequest) (*CreateMessageWithToolsResult, error) { + calls.Add(1) + return &CreateMessageWithToolsResult{Model: "m", Content: []Content{&TextContent{Text: "answer"}}}, nil + } + case *ElicitParams: + opts.ElicitationHandler = func(context.Context, *ElicitRequest) (*ElicitResult, error) { + calls.Add(1) + return &ElicitResult{Action: "decline"}, nil + } + } + cs := conn.connect(t, s, opts) + res, err := cs.CallTool(t.Context(), &CallToolParams{Name: "ask"}) + if tc.wantErr { + if err == nil || !strings.Contains(err.Error(), "does not support") { + t.Fatalf("CallTool = %+v, %v; want an unsupported-capability error", res, err) + } + if got := calls.Load(); got != 0 { + t.Errorf("handler called %d times, want 0", got) + } + return + } + if err != nil || res.IsError { + t.Fatalf("CallTool = %+v, %v; want success", res, err) + } + if got := calls.Load(); got != 1 { + t.Errorf("handler called %d times, want 1", got) + } + }) + } + } +} diff --git a/mcp/client.go b/mcp/client.go index 44be0ef6a..ea17d1495 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -93,7 +93,8 @@ type ClientOptions struct { // client to advertise the sampling capability, with default value // &SamplingCapabilities{}. If [ClientOptions.Capabilities] is set and has a // non nil value for [ClientCapabilities.Sampling], that value overrides the - // inferred capability. + // inferred capability. Unless that capability includes tools, a request + // that uses tools is answered with an error. // // Deprecated: the sampling feature is deprecated as of protocol version // 2026-07-28 (SEP-2577). It remains functional during the deprecation @@ -109,7 +110,8 @@ type ClientOptions struct { // Setting this handler causes the client to advertise the sampling // capability with tools support (sampling.tools). As with // [CreateMessageHandler], [ClientOptions.Capabilities].Sampling overrides - // the inferred capability. + // the inferred capability, and a request that uses tools is answered with + // an error unless it includes Tools. // // It is a panic to set both CreateMessageHandler and // CreateMessageWithToolsHandler. @@ -751,26 +753,41 @@ func (c *Client) listRoots(_ context.Context, req *ListRootsRequest) (*ListRoots } func (c *Client) createMessage(ctx context.Context, req *CreateMessageWithToolsRequest) (*CreateMessageWithToolsResult, error) { + if c.opts.CreateMessageWithToolsHandler == nil && c.opts.CreateMessageHandler == nil { + return nil, &jsonrpc.Error{Code: codeUnsupportedMethod, Message: "client does not support CreateMessage"} + } + // A client that didn't declare sampling.tools must return an error for a + // request that uses tools, rather than answer it without them. + if p := req.Params; p != nil && (p.Tools != nil || p.ToolChoice != nil) { + if caps := c.declaredCapabilities(); caps.Sampling == nil || caps.Sampling.Tools == nil { + return nil, &jsonrpc.Error{Code: jsonrpc.CodeInvalidParams, Message: "client does not support sampling with tools"} + } + } if c.opts.CreateMessageWithToolsHandler != nil { return c.opts.CreateMessageWithToolsHandler(ctx, req) } - if c.opts.CreateMessageHandler != nil { - // Downconvert the request for the basic handler. - baseParams, err := req.Params.toBase() - if err != nil { - return nil, err - } - baseReq := &CreateMessageRequest{ - Session: req.Session, - Params: baseParams, - } - res, err := c.opts.CreateMessageHandler(ctx, baseReq) - if err != nil { - return nil, err - } - return res.toWithTools(), nil + // Downconvert the request for the basic handler. + baseParams, err := req.Params.toBase() + if err != nil { + return nil, err + } + baseReq := &CreateMessageRequest{ + Session: req.Session, + Params: baseParams, + } + res, err := c.opts.CreateMessageHandler(ctx, baseReq) + if err != nil { + return nil, err } - return nil, &jsonrpc.Error{Code: codeUnsupportedMethod, Message: "client does not support CreateMessage"} + return res.toWithTools(), nil +} + +// declaredCapabilities returns the capabilities the client declares to +// servers. The protocol version only decides whether an inferred elicitation +// capability spells out form mode, which it implies either way, so it doesn't +// matter to the checks that use this. +func (c *Client) declaredCapabilities() *ClientCapabilities { + return c.capabilities("") } // urlElicitationMiddleware returns middleware that automatically handles URL elicitation @@ -888,6 +905,14 @@ func (c *Client) elicit(ctx context.Context, req *ElicitRequest) (*ElicitResult, mode = "form" } + // A request for a mode the client didn't declare is invalid params. As on + // the server, an elicitation capability without modes means form mode. + if caps := c.declaredCapabilities().Elicitation; caps != nil { + if (mode == "url" && caps.URL == nil) || (mode == "form" && caps.Form == nil && caps.URL != nil) { + return nil, &jsonrpc.Error{Code: jsonrpc.CodeInvalidParams, Message: fmt.Sprintf("client does not support %q elicitation", mode)} + } + } + switch mode { case "form": if req.Params.URL != "" { diff --git a/mcp/server.go b/mcp/server.go index 16aad6a94..f94aa6afa 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -1694,6 +1694,9 @@ func (ss *ServerSession) Ping(ctx context.Context, params *PingParams) error { // ListRoots lists the client roots. // +// It returns an error, without sending the request, if the client did not +// declare the roots capability. +// // Deprecated: the roots feature is deprecated as of protocol version // 2026-07-28 (SEP-2577). It remains functional during the deprecation window // (at least twelve months). Migrate to passing paths via tool parameters, @@ -1706,15 +1709,30 @@ func (ss *ServerSession) ListRoots(ctx context.Context, params *ListRootsParams) if err := ss.assertServerInitiatedRequestAllowed(methodListRoots); err != nil { return nil, err } + if caps := ss.clientCapabilities(); caps == nil || caps.RootsV2 == nil { + return nil, errors.New("client does not support roots") + } return handleSend[*ListRootsResult](ctx, methodListRoots, newServerRequest(ss, orZero[Params](params))) } +// clientCapabilities returns the capabilities the client declared when it +// initialized the session, or nil if it hasn't. +func (ss *ServerSession) clientCapabilities() *ClientCapabilities { + if iparams := ss.InitializeParams(); iparams != nil { + return iparams.Capabilities + } + return nil +} + // CreateMessage sends a sampling request to the client. // // If the client returns multiple content blocks (e.g. parallel tool calls), // CreateMessage returns an error. Use [ServerSession.CreateMessageWithTools] // for tool-enabled sampling. // +// It returns an error, without sending the request, if the client did not +// declare the sampling capability. +// // Deprecated: the sampling feature is deprecated as of protocol version // 2026-07-28 (SEP-2577). It remains functional during the deprecation window // (at least twelve months). Migrate to calling LLM provider APIs directly @@ -1727,6 +1745,9 @@ func (ss *ServerSession) CreateMessage(ctx context.Context, params *CreateMessag if err := ss.assertServerInitiatedRequestAllowed(methodCreateMessage); err != nil { return nil, err } + if caps := ss.clientCapabilities(); caps == nil || caps.Sampling == nil { + return nil, errors.New("client does not support sampling") + } if params == nil { params = &CreateMessageParams{Messages: []*SamplingMessage{}} } @@ -1761,6 +1782,10 @@ func (ss *ServerSession) CreateMessage(ctx context.Context, params *CreateMessag // (for parallel tool calls). Use this instead of [ServerSession.CreateMessage] // when the request includes tools. // +// It returns an error, without sending the request, if the client did not +// declare the sampling capability, or if the request has tools or a tool +// choice and the client did not declare sampling.tools. +// // Deprecated: the sampling feature is deprecated as of protocol version // 2026-07-28 (SEP-2577). It remains functional during the deprecation window // (at least twelve months). Migrate to calling LLM provider APIs directly @@ -1776,6 +1801,13 @@ func (ss *ServerSession) CreateMessageWithTools(ctx context.Context, params *Cre if params == nil { params = &CreateMessageWithToolsParams{Messages: []*SamplingMessageV2{}} } + caps := ss.clientCapabilities() + if caps == nil || caps.Sampling == nil { + return nil, errors.New("client does not support sampling") + } + if (len(params.Tools) > 0 || params.ToolChoice != nil) && caps.Sampling.Tools == nil { + return nil, errors.New("client does not support sampling with tools") + } if params.Messages == nil { p2 := *params p2.Messages = []*SamplingMessageV2{} // avoid JSON "null" diff --git a/mcp/streamable_test.go b/mcp/streamable_test.go index a481e7b7c..df23ec330 100644 --- a/mcp/streamable_test.go +++ b/mcp/streamable_test.go @@ -824,7 +824,11 @@ func TestStreamableServerTransport(t *testing.T) { // requests. // Predefined steps, to avoid repetition below. - initReq := req(1, methodInitialize, &InitializeParams{ProtocolVersion: protocolVersion20250618}) + // The client declares roots, so that tools can call ListRoots. + initReq := req(1, methodInitialize, &InitializeParams{ + ProtocolVersion: protocolVersion20250618, + Capabilities: &ClientCapabilities{RootsV2: &RootCapabilities{}}, + }) initResp := resp(1, &InitializeResult{ Capabilities: &ServerCapabilities{ Logging: &LoggingCapabilities{}, diff --git a/mcp/testdata/conformance/server/resources.txtar b/mcp/testdata/conformance/server/resources.txtar index 1c3ed8097..3c7c9e4f2 100644 --- a/mcp/testdata/conformance/server/resources.txtar +++ b/mcp/testdata/conformance/server/resources.txtar @@ -17,7 +17,7 @@ info.txt "method": "initialize", "params": { "protocolVersion": "2024-11-05", - "capabilities": {}, + "capabilities": { "roots": {} }, "clientInfo": { "name": "ExampleClient", "version": "1.0.0" } } } From 55193839afaa025fb7fbe761c9df1c2151024ce3 Mon Sep 17 00:00:00 2001 From: po-et <42566883+po-et@users.noreply.github.com> Date: Mon, 19 Oct 2026 22:18:34 +0800 Subject: [PATCH 2/2] mcp: reject undeclared roots requests Check the roots capability in the shared client handler so modern input requests respect explicitly disabled roots, as the legacy server API does. Return the standard method-not-found error on the client rejection path. Cover both protocol versions and raw legacy requests in regression tests. --- mcp/capabilities_test.go | 90 ++++++++++++++++++++++++++++++++++++++++ mcp/client.go | 3 ++ 2 files changed, 93 insertions(+) diff --git a/mcp/capabilities_test.go b/mcp/capabilities_test.go index eab17d509..1dc656ef8 100644 --- a/mcp/capabilities_test.go +++ b/mcp/capabilities_test.go @@ -495,3 +495,93 @@ func TestMultiRoundTripDeclaredCapabilities(t *testing.T) { } } } + +func TestMultiRoundTripRootsCapability(t *testing.T) { + for _, tc := range []struct { + name string + caps *ClientCapabilities + wantErr bool + }{ + {name: "inferred"}, + {name: "explicit", caps: &ClientCapabilities{RootsV2: &RootCapabilities{}}}, + {name: "disabled", caps: &ClientCapabilities{}, wantErr: true}, + } { + for _, conn := range []struct { + version string + connect func(*testing.T, *Server, *ClientOptions) *ClientSession + }{ + {protocolVersion20260728, mustConnect}, + {protocolVersion20251125, mustConnectOldProtocol}, + } { + t.Run(tc.name+"/"+conn.version, func(t *testing.T) { + var replies atomic.Int32 + s := NewServer(testImpl, nil) + AddTool(s, &Tool{Name: "ask"}, func(_ context.Context, req *CallToolRequest, _ struct{}) (*CallToolResult, any, error) { + if len(req.Params.InputResponses) == 0 { + return &CallToolResult{InputRequests: InputRequestMap{"roots": &ListRootsParams{}}}, nil, nil + } + roots, ok := req.Params.InputResponses["roots"].(*ListRootsResult) + if !ok || len(roots.Roots) != 1 || roots.Roots[0].URI != "file:///workspace" || roots.Roots[0].Name != "workspace" { + t.Errorf("input response = %+v; want the configured workspace root", req.Params.InputResponses["roots"]) + } + replies.Add(1) + return &CallToolResult{Content: []Content{&TextContent{Text: "done"}}}, nil, nil + }) + cs := conn.connect(t, s, &ClientOptions{Capabilities: tc.caps}) + cs.client.AddRoots(&Root{URI: "file:///workspace", Name: "workspace"}) + res, err := cs.CallTool(t.Context(), &CallToolParams{Name: "ask"}) + if tc.wantErr { + if err == nil || !strings.Contains(err.Error(), "does not support roots") { + t.Fatalf("CallTool = %+v, %v; want an unsupported-roots error", res, err) + } + if got := replies.Load(); got != 0 { + t.Errorf("server received %d input responses, want 0", got) + } + return + } + if err != nil || res.IsError { + t.Fatalf("CallTool = %+v, %v; want success", res, err) + } + if got := replies.Load(); got != 1 { + t.Errorf("server received %d input responses, want 1", got) + } + }) + } + } +} + +func TestClientRejectsUndeclaredRoots(t *testing.T) { + ctx := t.Context() + client := NewClient(testImpl, &ClientOptions{Capabilities: &ClientCapabilities{}}) + client.AddRoots(&Root{URI: "file:///workspace", Name: "workspace"}) + ct, st := NewInMemoryTransports() + ss, err := NewServer(testImpl, nil).Connect(ctx, st, nil) + if err != nil { + t.Fatal(err) + } + defer ss.Close() + cs, err := client.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) + if err != nil { + t.Fatal(err) + } + defer cs.Close() + for _, tc := range []struct { + name string + params any + }{ + {"absent", nil}, + {"empty object", json.RawMessage(`{}`)}, + {"null", json.RawMessage(`null`)}, + } { + t.Run(tc.name, func(t *testing.T) { + var result json.RawMessage + err := ss.conn.Call(ctx, methodListRoots, tc.params).Await(ctx, &result) + if got := errorCode(err); got != jsonrpc.CodeMethodNotFound { + t.Errorf("roots/list error = %v (code %d), want code %d", err, got, jsonrpc.CodeMethodNotFound) + } + if len(result) != 0 { + t.Errorf("roots/list returned %s, want no result", result) + } + }) + } +} diff --git a/mcp/client.go b/mcp/client.go index cf2cb7f79..1b993cb0e 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -741,6 +741,9 @@ func (c *Client) shouldSendListChangedNotification(notification string) bool { } func (c *Client) listRoots(_ context.Context, req *ListRootsRequest) (*ListRootsResult, error) { + if c.declaredCapabilities().RootsV2 == nil { + return nil, &jsonrpc.Error{Code: jsonrpc.CodeMethodNotFound, Message: "client does not support roots"} + } c.mu.Lock() defer c.mu.Unlock() roots := slices.Collect(c.roots.all())