diff --git a/mcp/mrtr.go b/mcp/mrtr.go index 291d4641..231971fa 100644 --- a/mcp/mrtr.go +++ b/mcp/mrtr.go @@ -317,16 +317,30 @@ func fulfillInputRequests(ctx context.Context, cs *ClientSession, requests Input } func fulfillInputRequest(ctx context.Context, cs *ClientSession, ir InputRequest) (InputResponse, error) { + var ( + method string + params Params + ) switch p := ir.(type) { case *ElicitParams: - return cs.client.elicit(ctx, newClientRequest(cs, p)) + method, params = methodElicit, p case *CreateMessageParams: - return cs.client.createMessage(ctx, &CreateMessageWithToolsRequest{Session: cs, Params: createMessageParamsToWithTools(p)}) + method, params = methodCreateMessage, createMessageParamsToWithTools(p) case *CreateMessageWithToolsParams: - return cs.client.createMessage(ctx, &CreateMessageWithToolsRequest{Session: cs, Params: p}) + method, params = methodCreateMessage, p case *ListRootsParams: - return cs.client.listRoots(ctx, newClientRequest(cs, p)) + method, params = methodListRoots, p default: return nil, fmt.Errorf("unknown input request type: %T", ir) } + req := cs.receivingMethodInfos()[method].newRequest(cs, params, nil, jsonrpc.ID{}) + res, err := cs.receivingMethodHandler()(ctx, method, req) + if err != nil { + return nil, err + } + resp, ok := res.(InputResponse) + if !ok { + return nil, fmt.Errorf("%s: result of type %T is not an input response", method, res) + } + return resp, nil } diff --git a/mcp/mrtr_test.go b/mcp/mrtr_test.go index fd96cd10..d442f790 100644 --- a/mcp/mrtr_test.go +++ b/mcp/mrtr_test.go @@ -12,6 +12,7 @@ import ( "fmt" "slices" "strings" + "sync" "sync/atomic" "testing" @@ -200,6 +201,81 @@ func TestMultiRoundTrip_AutoRetry(t *testing.T) { } } +// inputRequestServer returns a server with a tool "ask" that asks the client +// for an elicitation, a sampling message and its roots before it completes. +func inputRequestServer() *Server { + 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{ + "confirm": &ElicitParams{Message: "OK?"}, + "draft": &CreateMessageParams{ + Messages: []*SamplingMessage{{Role: "user", Content: &TextContent{Text: "write"}}}, + MaxTokens: 50, + }, + "roots": &ListRootsParams{}, + }}, nil, nil + } + return &CallToolResult{Content: []Content{&TextContent{Text: "done"}}}, nil, nil + }) + return srv +} + +var inputRequestClientOptions = &ClientOptions{ + ElicitationHandler: func(context.Context, *ElicitRequest) (*ElicitResult, error) { + return &ElicitResult{Action: "accept"}, nil + }, + CreateMessageHandler: func(context.Context, *CreateMessageRequest) (*CreateMessageResult, error) { + return &CreateMessageResult{Model: "test-model", Role: "assistant", Content: &TextContent{Text: "response"}}, nil + }, +} + +// TestMultiRoundTrip_InputRequestsReachReceivingMiddleware checks that client +// receiving middleware sees the requests a server makes for input in the same +// way whether they arrive as server-to-client requests (2025-11-25) or as +// input requests in a multi round-trip result (2026-07-28). +func TestMultiRoundTrip_InputRequestsReachReceivingMiddleware(t *testing.T) { + for _, tc := range []struct { + name string + connect func(*testing.T, *Server, *ClientOptions) *ClientSession + }{ + {"2025-11-25", mustConnectOldProtocol}, + {"2026-07-28", mustConnect}, + } { + t.Run(tc.name, func(t *testing.T) { + cs := tc.connect(t, inputRequestServer(), inputRequestClientOptions) + cs.client.AddRoots(&Root{URI: "file:///workspace", Name: "workspace"}) + var ( + mu sync.Mutex + seen = map[string]string{} // method -> params type + ) + cs.client.AddReceivingMiddleware(func(next MethodHandler) MethodHandler { + return func(ctx context.Context, method string, req Request) (Result, error) { + switch method { + case methodElicit, methodCreateMessage, methodListRoots: + mu.Lock() + seen[method] = fmt.Sprintf("%T", req.GetParams()) + mu.Unlock() + } + return next(ctx, method, req) + } + }) + + if _, err := cs.CallTool(context.Background(), &CallToolParams{Name: "ask"}); err != nil { + t.Fatalf("CallTool() error = %v", err) + } + want := map[string]string{ + methodElicit: "*mcp.ElicitParams", + methodCreateMessage: "*mcp.CreateMessageWithToolsParams", + methodListRoots: "*mcp.ListRootsParams", + } + if diff := cmp.Diff(want, seen); diff != "" { + t.Errorf("receiving middleware saw (-want +got):\n%s", diff) + } + }) + } +} + // TestMultiRoundTrip_AutoRetryDoesNotMutateCallerParams verifies that // the client middleware carries inputResponses and requestState on a // copy of the caller's params: after CallTool returns, the caller's