diff --git a/mcp/mrtr.go b/mcp/mrtr.go index 5f65f569..008c0f84 100644 --- a/mcp/mrtr.go +++ b/mcp/mrtr.go @@ -108,7 +108,10 @@ func clientMultiRoundTripMiddleware() Middleware { if err != nil { return nil, err } - setMultiRoundTripRetryParams(req, responses, mrtrResult.requestState()) + req, err = setMultiRoundTripRetryParams(method, req, responses, mrtrResult.requestState()) + if err != nil { + return nil, err + } } } } @@ -152,7 +155,10 @@ func serverMultiRoundTripMiddleware() Middleware { if err != nil { return nil, err } - setMultiRoundTripRetryParams(req, responses, mrtrResult.requestState()) + req, err = setMultiRoundTripRetryParams(method, req, responses, mrtrResult.requestState()) + if err != nil { + return nil, err + } return next(ctx, method, req) } } @@ -213,20 +219,54 @@ func createMessageParamsToWithTools(p *CreateMessageParams) *CreateMessageWithTo } } -func setMultiRoundTripRetryParams(req Request, responses InputResponseMap, state string) { +// setMultiRoundTripRetryParams returns the request for the next round of a +// multi-round-trip retry: a request whose params are a shallow copy of the +// original, carrying the fulfilled responses and the request state to echo. +// The original params are never mutated: on the client side they are +// caller-owned, and a params struct reused across calls must not carry one +// call's inputResponses and requestState into the next, where a server +// gating on them (for example an elicitation confirmation) would treat the +// new call as already answered. +// +// An unrecognized params type is an error. It cannot occur on the server +// side, but client sending middleware runs before the retry loop and may +// substitute its own params. +func setMultiRoundTripRetryParams(method string, req Request, responses InputResponseMap, state string) (Request, error) { + var retryParams Params switch p := req.GetParams().(type) { case *CallToolParams: - p.InputResponses = responses - p.RequestState = state + cp := *p + cp.InputResponses = responses + cp.RequestState = state + retryParams = &cp case *CallToolParamsRaw: - p.InputResponses = responses - p.RequestState = state + cp := *p + cp.InputResponses = responses + cp.RequestState = state + retryParams = &cp case *GetPromptParams: - p.InputResponses = responses - p.RequestState = state + cp := *p + cp.InputResponses = responses + cp.RequestState = state + retryParams = &cp case *ReadResourceParams: - p.InputResponses = responses - p.RequestState = state + cp := *p + cp.InputResponses = responses + cp.RequestState = state + retryParams = &cp + default: + return nil, fmt.Errorf("multi-round-trip: unsupported params type %T", p) + } + switch s := req.GetSession().(type) { + case *ClientSession: + return newClientRequest(s, retryParams), nil + case *ServerSession: + // Rebuild with the same constructor handleReceive uses: the + // receiving dispatch type-asserts on the request's concrete + // *ServerRequest[P] type, and Extra must carry over. + return s.receivingMethodInfos()[method].newRequest(s, retryParams, req.GetExtra()), nil + default: + return nil, fmt.Errorf("multi-round-trip: unsupported session type %T", s) } } diff --git a/mcp/mrtr_test.go b/mcp/mrtr_test.go index 2ae4fd97..d9662198 100644 --- a/mcp/mrtr_test.go +++ b/mcp/mrtr_test.go @@ -10,6 +10,7 @@ import ( "context" "fmt" "slices" + "strings" "sync/atomic" "testing" @@ -196,6 +197,198 @@ func TestMultiRoundTrip_AutoRetry(t *testing.T) { } } +// 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 +// struct is unchanged, and reusing it for a second call fulfills the +// input requests again instead of silently replaying the first call's +// answers against a server that gates on them. +func TestMultiRoundTrip_AutoRetryDoesNotMutateCallerParams(t *testing.T) { + ctx := context.Background() + + srv := NewServer(testImpl, nil) + AddTool(srv, &Tool{Name: "act"}, func(ctx context.Context, req *CallToolRequest, input struct{}) (*CallToolResult, any, error) { + if len(req.Params.InputResponses) == 0 { + return &CallToolResult{ + InputRequests: InputRequestMap{"confirm": &ElicitParams{Message: "Sure?"}}, + RequestState: "state-1", + }, nil, nil + } + return &CallToolResult{}, map[string]any{"ok": true}, nil + }) + + var elicitations atomic.Int32 + conn := mustConnect(t, srv, &ClientOptions{ + ElicitationHandler: func(_ context.Context, _ *ElicitRequest) (*ElicitResult, error) { + elicitations.Add(1) + return &ElicitResult{Action: "accept"}, nil + }, + }) + + params := &CallToolParams{Name: "act"} + res, err := conn.CallTool(ctx, params) + if err != nil { + t.Fatalf("CallTool() error = %v", err) + } + if res.NeedsInput() { + t.Fatal("NeedsInput() = true after auto-retry, want false") + } + if params.InputResponses != nil { + t.Errorf("params.InputResponses = %v after CallTool, want nil (caller params must not be mutated)", params.InputResponses) + } + if params.RequestState != "" { + t.Errorf("params.RequestState = %q after CallTool, want empty (caller params must not be mutated)", params.RequestState) + } + if got := elicitations.Load(); got != 1 { + t.Fatalf("elicitations = %d, want 1", got) + } + + // Reusing the same params struct must fulfill the input requests + // again, not replay the first call's responses. + if _, err := conn.CallTool(ctx, params); err != nil { + t.Fatalf("CallTool() reuse error = %v", err) + } + if got := elicitations.Load(); got != 2 { + t.Errorf("elicitations after reuse = %d, want 2 (stale responses must not be replayed)", got) + } +} + +// TestMultiRoundTrip_GetPrompt_AutoRetryDoesNotMutateCallerParams is the +// GetPrompt analogue of the CallTool caller-params test. +func TestMultiRoundTrip_GetPrompt_AutoRetryDoesNotMutateCallerParams(t *testing.T) { + ctx := context.Background() + + srv := NewServer(testImpl, nil) + srv.AddPrompt(&Prompt{Name: "review"}, func(_ context.Context, req *GetPromptRequest) (*GetPromptResult, error) { + if len(req.Params.InputResponses) == 0 { + return &GetPromptResult{ + InputRequests: InputRequestMap{"confirm": &ElicitParams{Message: "Include sensitive data?"}}, + RequestState: "prompt-state", + }, nil + } + return &GetPromptResult{ + Messages: []*PromptMessage{{Role: "user", Content: &TextContent{Text: "review this code"}}}, + }, nil + }) + + conn := mustConnect(t, srv, &ClientOptions{ + ElicitationHandler: func(_ context.Context, _ *ElicitRequest) (*ElicitResult, error) { + return &ElicitResult{Action: "accept"}, nil + }, + }) + + params := &GetPromptParams{Name: "review"} + res, err := conn.GetPrompt(ctx, params) + if err != nil { + t.Fatalf("GetPrompt() error = %v", err) + } + if res.NeedsInput() { + t.Fatal("NeedsInput() = true after auto-retry, want false") + } + if params.InputResponses != nil { + t.Errorf("params.InputResponses = %v after GetPrompt, want nil (caller params must not be mutated)", params.InputResponses) + } + if params.RequestState != "" { + t.Errorf("params.RequestState = %q after GetPrompt, want empty (caller params must not be mutated)", params.RequestState) + } +} + +// TestMultiRoundTrip_ReadResource_AutoRetryDoesNotMutateCallerParams is +// the ReadResource analogue of the CallTool caller-params test. +func TestMultiRoundTrip_ReadResource_AutoRetryDoesNotMutateCallerParams(t *testing.T) { + ctx := context.Background() + + srv := NewServer(testImpl, nil) + srv.AddResource(&Resource{URI: "test://data", Name: "data"}, func(_ context.Context, req *ReadResourceRequest) (*ReadResourceResult, error) { + if len(req.Params.InputResponses) == 0 { + return &ReadResourceResult{ + InputRequests: InputRequestMap{"auth": &ElicitParams{Message: "Authenticate?"}}, + RequestState: "resource-state", + }, nil + } + return &ReadResourceResult{ + Contents: []*ResourceContents{{URI: "test://data", Text: "resource data"}}, + }, nil + }) + + conn := mustConnect(t, srv, &ClientOptions{ + ElicitationHandler: func(_ context.Context, _ *ElicitRequest) (*ElicitResult, error) { + return &ElicitResult{Action: "accept"}, nil + }, + }) + + params := &ReadResourceParams{URI: "test://data"} + res, err := conn.ReadResource(ctx, params) + if err != nil { + t.Fatalf("ReadResource() error = %v", err) + } + if res.NeedsInput() { + t.Fatal("NeedsInput() = true after auto-retry, want false") + } + if params.InputResponses != nil { + t.Errorf("params.InputResponses = %v after ReadResource, want nil (caller params must not be mutated)", params.InputResponses) + } + if params.RequestState != "" { + t.Errorf("params.RequestState = %q after ReadResource, want empty (caller params must not be mutated)", params.RequestState) + } +} + +// wrappedCallToolParams embeds *CallToolParams so it satisfies Params and +// marshals identically, while being a concrete type the multi-round-trip +// retry does not recognize. +type wrappedCallToolParams struct{ *CallToolParams } + +// TestMultiRoundTrip_UnsupportedRetryParamsType verifies that the client +// retry loop reports an explicit error when the params carry a type it +// cannot build retry params for, instead of silently resending the request +// unchanged until the retry cap. +func TestMultiRoundTrip_UnsupportedRetryParamsType(t *testing.T) { + ctx := context.Background() + + srv := NewServer(testImpl, nil) + AddTool(srv, &Tool{Name: "act"}, func(ctx context.Context, req *CallToolRequest, input struct{}) (*CallToolResult, any, error) { + return &CallToolResult{ + InputRequests: InputRequestMap{"confirm": &ElicitParams{Message: "Sure?"}}, + RequestState: "state-1", + }, nil, nil + }) + + st, ct := NewInMemoryTransports() + ss, err := srv.Connect(t.Context(), st, nil) + if err != nil { + t.Fatalf("server.Connect() error = %v", err) + } + t.Cleanup(func() { _ = ss.Close() }) + + c := NewClient(testImpl, &ClientOptions{ + ElicitationHandler: func(_ context.Context, _ *ElicitRequest) (*ElicitResult, error) { + return &ElicitResult{Action: "accept"}, nil + }, + }) + // Sending middleware added here wraps outside the multi-round-trip + // middleware installed at client construction, so the retry loop sees + // the substituted params type. + c.AddSendingMiddleware(func(next MethodHandler) MethodHandler { + return func(ctx context.Context, method string, req Request) (Result, error) { + if method == methodCallTool { + cs := req.GetSession().(*ClientSession) + req = newClientRequest[Params](cs, wrappedCallToolParams{req.GetParams().(*CallToolParams)}) + } + return next(ctx, method, req) + } + }) + cs, err := c.Connect(t.Context(), ct, nil) + if err != nil { + t.Fatalf("client.Connect() error = %v", err) + } + t.Cleanup(func() { _ = cs.Close() }) + + _, err = cs.CallTool(ctx, &CallToolParams{Name: "act"}) + if err == nil || !strings.Contains(err.Error(), "unsupported params type") { + t.Fatalf("CallTool() error = %v, want unsupported params type error", err) + } +} + func TestMultiRoundTrip_MaxRetries(t *testing.T) { testCases := []struct { name string @@ -329,6 +522,41 @@ func TestMultiRoundTrip_ServerMiddleware(t *testing.T) { } } +// TestMultiRoundTrip_ServerMiddleware_GetPrompt verifies the server-side +// shim retry for prompts: the rebuilt retry request must keep its concrete +// *ServerRequest[*GetPromptParams] type for the receiving dispatch. +func TestMultiRoundTrip_ServerMiddleware_GetPrompt(t *testing.T) { + ctx := context.Background() + + srv := NewServer(testImpl, nil) + srv.AddPrompt(&Prompt{Name: "review"}, func(_ context.Context, req *GetPromptRequest) (*GetPromptResult, error) { + if len(req.Params.InputResponses) == 0 { + return &GetPromptResult{ + InputRequests: InputRequestMap{"confirm": &ElicitParams{Message: "Sure?"}}, + RequestState: "prompt-state", + }, nil + } + return &GetPromptResult{ + Messages: []*PromptMessage{{Role: "user", Content: &TextContent{Text: "approved"}}}, + }, nil + }) + + conn := mustConnectOldProtocol(t, srv, &ClientOptions{ + MultiRoundTrip: &MultiRoundTripOptions{Disabled: true}, + ElicitationHandler: func(_ context.Context, _ *ElicitRequest) (*ElicitResult, error) { + return &ElicitResult{Action: "accept"}, nil + }, + }) + + res, err := conn.GetPrompt(ctx, &GetPromptParams{Name: "review"}) + if err != nil { + t.Fatalf("GetPrompt() error = %v", err) + } + if got := res.Messages[0].Content.(*TextContent).Text; got != "approved" { + t.Errorf("message text = %q, want %q", got, "approved") + } +} + func TestMultiRoundTrip_GetPrompt_AutoRetry(t *testing.T) { ctx := context.Background() @@ -503,6 +731,138 @@ func TestMultiRoundTrip_ReadResource_ManualRetry(t *testing.T) { } } +// TestSetMultiRoundTripRetryParams pins the properties the retry-request +// builder must uphold: the original params are never mutated, the copy +// carries the responses and state, the server-side rebuild preserves the +// request's concrete *ServerRequest[P] type (the receiving dispatch +// type-asserts on it) and its Extra, and an unrecognized params type is an +// explicit error. The server readResource case matters here because the +// end-to-end shim path cannot reach it (the read handler's nil-contents +// validation rejects input-required results for old-protocol clients). +func TestSetMultiRoundTripRetryParams(t *testing.T) { + ss := &ServerSession{server: NewServer(testImpl, nil)} + extra := &RequestExtra{} + responses := InputResponseMap{"confirm": &ElicitResult{Action: "accept"}} + + // checkRetry asserts the returned params carry the retry fields while + // the original params were left untouched. + checkRetry := func(t *testing.T, orig, retry Params) { + t.Helper() + switch p := retry.(type) { + case *CallToolParams: + if len(p.InputResponses) != 1 || p.RequestState != "s1" { + t.Errorf("retry params = %+v, want responses and state carried", p) + } + case *CallToolParamsRaw: + if len(p.InputResponses) != 1 || p.RequestState != "s1" { + t.Errorf("retry params = %+v, want responses and state carried", p) + } + case *GetPromptParams: + if len(p.InputResponses) != 1 || p.RequestState != "s1" { + t.Errorf("retry params = %+v, want responses and state carried", p) + } + case *ReadResourceParams: + if len(p.InputResponses) != 1 || p.RequestState != "s1" { + t.Errorf("retry params = %+v, want responses and state carried", p) + } + default: + t.Errorf("retry params type = %T, want one of the multi-round-trip params types", p) + } + switch p := orig.(type) { + case *CallToolParams: + if p.InputResponses != nil || p.RequestState != "" { + t.Errorf("original params mutated: %+v", p) + } + case *CallToolParamsRaw: + if p.InputResponses != nil || p.RequestState != "" { + t.Errorf("original params mutated: %+v", p) + } + case *GetPromptParams: + if p.InputResponses != nil || p.RequestState != "" { + t.Errorf("original params mutated: %+v", p) + } + case *ReadResourceParams: + if p.InputResponses != nil || p.RequestState != "" { + t.Errorf("original params mutated: %+v", p) + } + } + } + + t.Run("server callTool", func(t *testing.T) { + orig := &CallToolParamsRaw{Name: "act"} + req, err := setMultiRoundTripRetryParams(methodCallTool, &ServerRequest[*CallToolParamsRaw]{Session: ss, Params: orig, Extra: extra}, responses, "s1") + if err != nil { + t.Fatalf("setMultiRoundTripRetryParams() error = %v", err) + } + got, ok := req.(*ServerRequest[*CallToolParamsRaw]) + if !ok { + t.Fatalf("retry request type = %T, want *ServerRequest[*CallToolParamsRaw]", req) + } + if got.Extra != extra { + t.Error("retry request lost Extra") + } + checkRetry(t, orig, got.Params) + }) + + t.Run("server getPrompt", func(t *testing.T) { + orig := &GetPromptParams{Name: "review"} + req, err := setMultiRoundTripRetryParams(methodGetPrompt, &ServerRequest[*GetPromptParams]{Session: ss, Params: orig, Extra: extra}, responses, "s1") + if err != nil { + t.Fatalf("setMultiRoundTripRetryParams() error = %v", err) + } + got, ok := req.(*ServerRequest[*GetPromptParams]) + if !ok { + t.Fatalf("retry request type = %T, want *ServerRequest[*GetPromptParams]", req) + } + if got.Extra != extra { + t.Error("retry request lost Extra") + } + checkRetry(t, orig, got.Params) + }) + + t.Run("server readResource", func(t *testing.T) { + orig := &ReadResourceParams{URI: "test://data"} + req, err := setMultiRoundTripRetryParams(methodReadResource, &ServerRequest[*ReadResourceParams]{Session: ss, Params: orig, Extra: extra}, responses, "s1") + if err != nil { + t.Fatalf("setMultiRoundTripRetryParams() error = %v", err) + } + got, ok := req.(*ServerRequest[*ReadResourceParams]) + if !ok { + t.Fatalf("retry request type = %T, want *ServerRequest[*ReadResourceParams]", req) + } + if got.Extra != extra { + t.Error("retry request lost Extra") + } + checkRetry(t, orig, got.Params) + }) + + t.Run("client params", func(t *testing.T) { + for _, orig := range []Params{ + &CallToolParams{Name: "act"}, + &GetPromptParams{Name: "review"}, + &ReadResourceParams{URI: "test://data"}, + } { + req, err := setMultiRoundTripRetryParams(methodCallTool, newClientRequest[Params](nil, orig), responses, "s1") + if err != nil { + t.Fatalf("setMultiRoundTripRetryParams(%T) error = %v", orig, err) + } + got, ok := req.(*ClientRequest[Params]) + if !ok { + t.Fatalf("retry request type = %T, want *ClientRequest[Params]", req) + } + checkRetry(t, orig, got.Params) + } + }) + + t.Run("unsupported params type", func(t *testing.T) { + wrapped := wrappedCallToolParams{&CallToolParams{Name: "act"}} + _, err := setMultiRoundTripRetryParams(methodCallTool, newClientRequest[Params](nil, wrapped), responses, "s1") + if err == nil || !strings.Contains(err.Error(), "unsupported params type") { + t.Fatalf("setMultiRoundTripRetryParams() error = %v, want unsupported params type error", err) + } + }) +} + func mustConnect(t *testing.T, s *Server, clientOpts *ClientOptions) *ClientSession { t.Helper() @@ -525,3 +885,29 @@ func mustConnect(t *testing.T, s *Server, clientOpts *ClientOptions) *ClientSess }) return cs } + +// mustConnectOldProtocol is mustConnect for a client on a protocol version +// that predates multi-round-trip: the server-side shim, not the client +// middleware, fulfills input requests. +func mustConnectOldProtocol(t *testing.T, s *Server, clientOpts *ClientOptions) *ClientSession { + t.Helper() + + st, ct := NewInMemoryTransports() + ss, err := s.Connect(t.Context(), st, nil) + if err != nil { + t.Fatalf("server.Connect() error = %v", err) + } + t.Cleanup(func() { + _ = ss.Close() + }) + + c := NewClient(testImpl, clientOpts) + cs, err := c.Connect(t.Context(), ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) + if err != nil { + t.Fatalf("client.Connect() error = %v", err) + } + t.Cleanup(func() { + _ = cs.Close() + }) + return cs +}