modelcontextprotocol/go-sdk
Watch
1
0
Fork
You've already forked go-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 2

mcp: don't mutate caller params on multi-round-trip retries (#1145)

Browse source 
Fixes #1144.

`clientMultiRoundTripMiddleware` calls
`setMultiRoundTripRetryParams(req, ...)`, which writes `InputResponses`
and `RequestState` through the caller's params pointer and never clears
them. A `*CallToolParams` (or `GetPromptParams` / `ReadResourceParams`)
struct reused across calls then carries the first call's answers into
the next call: a server gating on `len(req.Params.InputResponses) == 0`
— the idiom in this repo's own tests and examples — treats the fresh
call as already answered, so the client's `ElicitationHandler` never
fires and an elicitation-gated (e.g. destructive) tool runs without
confirmation.

This change makes the retry loop carry the fulfilled responses and
echoed state on a shallow copy of the params
(`multiRoundTripRetryRequest`) instead of mutating the original request.
Caller-owned params are now left untouched by `CallTool`, `GetPrompt`,
and `ReadResource`. The server-side shim
(`serverMultiRoundTripMiddleware`) is unchanged: the params it mutates
are SDK-owned, decoded from the wire.

Tests: `TestMultiRoundTrip_AutoRetryDoesNotMutateCallerParams`
(CallTool: params unchanged after the call, and reusing the same struct
fulfills the input requests again instead of silently replaying the
first call's answers — it failed with `elicitations after reuse = 1,
want 2` before the fix), plus GetPrompt/ReadResource analogues asserting
params are left clean. All three fail on main and pass with this change;
`go test ./...` and `go vet ./...` are clean.

Note on observable behavior: callers can no longer read the fulfilled
responses off their own params struct after a call returns. That side
effect was undocumented, and relying on it is exactly what makes the
replay hazard above possible.

---------

Co-authored-by: Guglielmo Colombo <guglielmoc@google.com>
This commit is contained in:
Haoxin Yang 2026-08-31 19:53:34 +08:00 • committed by GitHub
commit 69c17de411
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 437 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

62
mcp/mrtr.go
View file

@ -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)
}
}

386
mcp/mrtr_test.go
View file

@ -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
}