Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 34 additions & 1 deletion mcp/mrtr.go
Original file line number Diff line number Diff line change
Expand Up @@ -108,12 +108,45 @@ func clientMultiRoundTripMiddleware() Middleware {
if err != nil {
return nil, err
}
setMultiRoundTripRetryParams(req, responses, mrtrResult.requestState())
req = multiRoundTripRetryRequest(cs, req, responses, mrtrResult.requestState())
}
}
}
}

// multiRoundTripRetryRequest builds the request for the next round of a
// client-side multi-round-trip retry: a request whose params carry the
// fulfilled responses and the request state to echo. The retry params
// are a shallow copy of the original, so the caller-owned params passed
// to CallTool, GetPrompt, or ReadResource are never mutated: 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.
func multiRoundTripRetryRequest(cs *ClientSession, req Request, responses InputResponseMap, state string) Request {
var clone Params
switch p := req.GetParams().(type) {
case *CallToolParams:
cp := *p
clone = &cp
case *CallToolParamsRaw:
cp := *p
clone = &cp
case *GetPromptParams:
cp := *p
clone = &cp
case *ReadResourceParams:
cp := *p
clone = &cp
default:
// No retry params to carry; resend the request unchanged.
return req
Comment on lines +141 to +143

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This comment is misleading. This default case should never be reached in general

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we should replace it with an error, or more comprehensive comment

}
retry := newClientRequest(cs, clone)
setMultiRoundTripRetryParams(retry, responses, state)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

what do you think about merging the clone within the setMultiRoundTripRetryParams?
Even if applied to the server middleware it should not be a problem

return retry
}

// serverMultiRoundTripMiddleware is a receiving middleware for servers that transparently
// handles multi-round-trip for clients on older protocol versions. When a handler returns
// InputRequests and the client does not support multi-round-trip, the middleware fulfills
Expand Down
136 changes: 136 additions & 0 deletions mcp/mrtr_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -196,6 +196,142 @@ 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)
}
}

func TestMultiRoundTrip_MaxRetries(t *testing.T) {
testCases := []struct {
name string
Expand Down