diff --git a/mcp/client.go b/mcp/client.go index ab1c3353..2055f5fd 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -752,7 +752,11 @@ 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 { - return c.opts.CreateMessageWithToolsHandler(ctx, req) + res, err := c.opts.CreateMessageWithToolsHandler(ctx, req) + if err == nil && res == nil { + return nil, fmt.Errorf("%w: CreateMessageWithToolsHandler returned a nil result", jsonrpc2.ErrInternal) + } + return res, err } if c.opts.CreateMessageHandler != nil { // Downconvert the request for the basic handler. @@ -768,6 +772,9 @@ func (c *Client) createMessage(ctx context.Context, req *CreateMessageWithToolsR if err != nil { return nil, err } + if res == nil { + return nil, fmt.Errorf("%w: CreateMessageHandler returned a nil result", jsonrpc2.ErrInternal) + } return res.toWithTools(), nil } return nil, &jsonrpc.Error{Code: codeUnsupportedMethod, Message: "client does not support CreateMessage"} @@ -901,6 +908,9 @@ func (c *Client) elicit(ctx context.Context, req *ElicitRequest) (*ElicitResult, if err != nil { return nil, err } + if res == nil { + return nil, fmt.Errorf("%w: ElicitationHandler returned a nil result", jsonrpc2.ErrInternal) + } // Validate elicitation result content against requested schema. if res.Action == "accept" && schema != nil && res.Content != nil { resolved, err := schema.Resolve(nil) @@ -923,8 +933,12 @@ func (c *Client) elicit(ctx context.Context, req *ElicitRequest) (*ElicitResult, if req.Params.URL == "" { return nil, &jsonrpc.Error{Code: jsonrpc.CodeInvalidParams, Message: "URL must be set for URL elicitation"} } - // No schema validation for URL mode, just pass through to handler. - return c.opts.ElicitationHandler(ctx, req) + // No schema validation for URL mode. + res, err := c.opts.ElicitationHandler(ctx, req) + if err == nil && res == nil { + return nil, fmt.Errorf("%w: ElicitationHandler returned a nil result", jsonrpc2.ErrInternal) + } + return res, err default: return nil, &jsonrpc.Error{Code: jsonrpc.CodeInvalidParams, Message: fmt.Sprintf("unsupported elicitation mode: %q", mode)} } diff --git a/mcp/mrtr_test.go b/mcp/mrtr_test.go index 432838fd..e15339a7 100644 --- a/mcp/mrtr_test.go +++ b/mcp/mrtr_test.go @@ -8,6 +8,7 @@ package mcp import ( "context" + "errors" "fmt" "slices" "strings" @@ -16,6 +17,8 @@ import ( "github.com/google/go-cmp/cmp" "github.com/google/jsonschema-go/jsonschema" + "github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2" + "github.com/modelcontextprotocol/go-sdk/jsonrpc" ) func TestMultiRoundTrip_ManualRetry(t *testing.T) { @@ -1311,6 +1314,99 @@ func mustConnectOldProtocol(t *testing.T, s *Server, clientOpts *ClientOptions) return cs } +// TestMultiRoundTrip_NilHandlerResult verifies that a sampling or elicitation handler +// returning (nil, nil) fails the request that needed it, rather than +// panicking in the client or sending a null result to the server. +func TestMultiRoundTrip_NilHandlerResult(t *testing.T) { + sampling := &CreateMessageParams{ + Messages: []*SamplingMessage{{Role: "user", Content: &TextContent{Text: "hi"}}}, + MaxTokens: 10, + } + nilElicit := func(context.Context, *ElicitRequest) (*ElicitResult, error) { return nil, nil } + tests := []struct { + name string + input InputRequest + opts *ClientOptions + }{ + { + name: "CreateMessageHandler", + input: sampling, + opts: &ClientOptions{CreateMessageHandler: func(context.Context, *CreateMessageRequest) (*CreateMessageResult, error) { + return nil, nil + }}, + }, + { + name: "CreateMessageWithToolsHandler", + input: sampling, + opts: &ClientOptions{CreateMessageWithToolsHandler: func(context.Context, *CreateMessageWithToolsRequest) (*CreateMessageWithToolsResult, error) { + return nil, nil + }}, + }, + { + name: "ElicitationHandler form", + input: &ElicitParams{Message: "OK?"}, + opts: &ClientOptions{ElicitationHandler: nilElicit}, + }, + { + name: "ElicitationHandler url", + input: &ElicitParams{Mode: "url", Message: "Sign in", URL: "https://example.com/login", ElicitationID: "e1"}, + opts: &ClientOptions{ + ElicitationHandler: nilElicit, + Capabilities: &ClientCapabilities{Elicitation: &ElicitationCapabilities{URL: &URLElicitationCapabilities{}}}, + }, + }, + } + connections := []struct { + version string + connect func(*testing.T, *Server, *ClientOptions) *ClientSession + }{ + {protocolVersion20260728, mustConnect}, + {protocolVersion20251125, mustConnectOldProtocol}, + } + for _, tt := range tests { + t.Run(tt.name+" error code", func(t *testing.T) { + c := NewClient(testImpl, tt.opts) + var err error + switch input := tt.input.(type) { + case *CreateMessageParams: + _, err = c.createMessage(t.Context(), &CreateMessageWithToolsRequest{Params: createMessageParamsToWithTools(input)}) + case *ElicitParams: + _, err = c.elicit(t.Context(), &ElicitRequest{Params: input}) + default: + t.Fatalf("unexpected input type %T", input) + } + if code := errorCode(err); code != jsonrpc.CodeInternalError { + t.Errorf("got error %v (code %d), want code %d", err, code, jsonrpc.CodeInternalError) + } + if !errors.Is(err, jsonrpc2.ErrInternal) { + t.Errorf("got error %v, want it to wrap jsonrpc2.ErrInternal", err) + } + }) + for _, conn := range connections { + t.Run(tt.name+" "+conn.version, func(t *testing.T) { + srv := NewServer(testImpl, nil) + AddTool(srv, &Tool{Name: "ask"}, func(ctx context.Context, req *CallToolRequest, _ struct{}) (*CallToolResult, any, error) { + if len(req.Params.InputResponses) == 0 { + return &CallToolResult{InputRequests: InputRequestMap{"q": tt.input}}, nil, nil + } + return &CallToolResult{Content: []Content{&TextContent{Text: fmt.Sprintf("got %v", req.Params.InputResponses)}}}, nil, nil + }) + cs := conn.connect(t, srv, tt.opts) + res, err := cs.CallTool(t.Context(), &CallToolParams{Name: "ask"}) + if err == nil { + t.Fatalf("CallTool succeeded with %+v, want an error", res) + } + if want := "returned a nil result"; !strings.Contains(err.Error(), want) { + t.Errorf("CallTool error = %v, want it to mention %q", err, want) + } + if code := errorCode(err); code != jsonrpc.CodeInternalError { + t.Errorf("CallTool error = %v (code %d), want code %d", err, code, jsonrpc.CodeInternalError) + } + }) + } + } +} + // gatedResourceServer returns a server whose "test://gated" resource asks for // an elicitation before it reveals its contents, and whose "test://plain" // resource answers at once. Every result is marked cacheable for a minute.