69c17de411Fixes #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>
913 lines
30 KiB
Go
913 lines
30 KiB
Go
// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
|
|
// Use of this source code is governed by the license
|
|
// that can be found in the LICENSE file.
|
|
|
|
//lint:file-ignore SA1019 tests exercise deprecated SEP-2577 APIs (roots, sampling, logging).
|
|
|
|
package mcp
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"slices"
|
|
"strings"
|
|
"sync/atomic"
|
|
"testing"
|
|
|
|
"github.com/google/go-cmp/cmp"
|
|
"github.com/google/jsonschema-go/jsonschema"
|
|
)
|
|
|
|
func TestMultiRoundTrip_ManualRetry(t *testing.T) {
|
|
type deployResult struct {
|
|
Deployed bool `json:"deployed"`
|
|
Reason string `json:"reason,omitempty"`
|
|
}
|
|
|
|
ctx := context.Background()
|
|
|
|
srv := NewServer(testImpl, nil)
|
|
AddTool(srv, &Tool{Name: "deploy"}, func(ctx context.Context, req *CallToolRequest, input struct{}) (*CallToolResult, *deployResult, error) {
|
|
if len(req.Params.InputResponses) == 0 {
|
|
return &CallToolResult{
|
|
InputRequests: InputRequestMap{"confirm": &ElicitParams{Message: "Deploy to production?"}},
|
|
RequestState: "deployment-123",
|
|
}, nil, nil
|
|
}
|
|
|
|
resp, ok := req.Params.InputResponses["confirm"]
|
|
if !ok {
|
|
return &CallToolResult{
|
|
InputRequests: InputRequestMap{"confirm": &ElicitParams{Message: "Please confirm (retry)"}},
|
|
}, nil, nil
|
|
}
|
|
|
|
if req.Params.RequestState == "" {
|
|
return &CallToolResult{}, &deployResult{Deployed: false, Reason: "no_state"}, nil
|
|
}
|
|
if elicitResult := resp.(*ElicitResult); elicitResult != nil && elicitResult.Action != "accept" {
|
|
return &CallToolResult{}, &deployResult{Deployed: false, Reason: "cancelled"}, nil
|
|
}
|
|
|
|
return &CallToolResult{}, &deployResult{Deployed: true}, nil
|
|
})
|
|
|
|
conn := mustConnect(t, srv, &ClientOptions{
|
|
MultiRoundTrip: &MultiRoundTripOptions{Disabled: true},
|
|
})
|
|
|
|
// Round 1: initiate deployment
|
|
res, err := conn.CallTool(ctx, &CallToolParams{Name: "deploy"})
|
|
if err != nil {
|
|
t.Fatalf("CallTool() error = %v", err)
|
|
}
|
|
if !res.NeedsInput() {
|
|
t.Fatal("NeedsInput() = false, want true")
|
|
}
|
|
if got := len(res.InputRequests); got != 1 {
|
|
t.Fatalf("len(res.InputRequests) = %d, want 1", got)
|
|
}
|
|
if _, ok := res.InputRequests["confirm"].(*ElicitParams); !ok {
|
|
t.Fatalf("res.InputRequests[confirm] type = %T, want *ElicitParams", res.InputRequests["confirm"])
|
|
}
|
|
|
|
// Round 2: retry with confirmation
|
|
res, err = conn.CallTool(ctx, &CallToolParams{
|
|
Name: "deploy",
|
|
InputResponses: InputResponseMap{
|
|
"confirm": &ElicitResult{Action: "accept", Content: map[string]any{"ok": true}},
|
|
},
|
|
RequestState: res.RequestState,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("CallTool() follow-up error = %v", err)
|
|
}
|
|
if res.NeedsInput() {
|
|
t.Fatal("NeedsInput() = true after follow-up, want false")
|
|
}
|
|
|
|
if diff := cmp.Diff(map[string]any{"deployed": true}, res.StructuredContent, ctrCmpOpts...); diff != "" {
|
|
t.Errorf("result mismatch (-want +got):\n%s", diff)
|
|
}
|
|
}
|
|
|
|
func TestMultiRoundTrip_AutoRetry(t *testing.T) {
|
|
|
|
tests := []struct {
|
|
name string
|
|
inputRequests InputRequestMap
|
|
wantResult map[string]any
|
|
}{
|
|
{
|
|
name: "elicit",
|
|
inputRequests: InputRequestMap{
|
|
"confirm": &ElicitParams{Message: "Deploy?"},
|
|
},
|
|
wantResult: map[string]any{"ids": []any{"confirm"}},
|
|
},
|
|
{
|
|
name: "createMessage",
|
|
inputRequests: InputRequestMap{
|
|
"summarize": &CreateMessageParams{
|
|
Messages: []*SamplingMessage{{Role: "user", Content: &TextContent{Text: "summarize"}}},
|
|
MaxTokens: 100,
|
|
},
|
|
},
|
|
wantResult: map[string]any{"ids": []any{"summarize"}},
|
|
},
|
|
{
|
|
name: "listRoots",
|
|
inputRequests: InputRequestMap{
|
|
"roots": &ListRootsParams{},
|
|
},
|
|
wantResult: map[string]any{"ids": []any{"roots"}},
|
|
},
|
|
{
|
|
name: "all three",
|
|
inputRequests: InputRequestMap{
|
|
"confirm": &ElicitParams{Message: "OK?"},
|
|
"draft": &CreateMessageParams{
|
|
Messages: []*SamplingMessage{{Role: "user", Content: &TextContent{Text: "write"}}},
|
|
MaxTokens: 50,
|
|
},
|
|
"roots": &ListRootsParams{},
|
|
},
|
|
wantResult: map[string]any{"ids": []any{"confirm", "draft", "roots"}},
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
srv := NewServer(testImpl, nil)
|
|
inputRequests := tt.inputRequests
|
|
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: inputRequests,
|
|
RequestState: "state-1",
|
|
}, nil, nil
|
|
}
|
|
// Collect the IDs of fulfilled responses.
|
|
var ids []string
|
|
for id := range req.Params.InputResponses {
|
|
ids = append(ids, id)
|
|
}
|
|
slices.Sort(ids)
|
|
return &CallToolResult{}, map[string]any{"ids": ids}, nil
|
|
})
|
|
|
|
conn := mustConnect(t, srv, &ClientOptions{
|
|
ElicitationHandler: func(_ context.Context, req *ElicitRequest) (*ElicitResult, error) {
|
|
return &ElicitResult{Action: "accept"}, nil
|
|
},
|
|
CreateMessageHandler: func(_ context.Context, req *CreateMessageRequest) (*CreateMessageResult, error) {
|
|
return &CreateMessageResult{
|
|
Model: "test-model",
|
|
Role: "assistant",
|
|
Content: &TextContent{Text: "response"},
|
|
}, nil
|
|
},
|
|
})
|
|
conn.client.AddRoots(&Root{URI: "file:///workspace", Name: "workspace"})
|
|
|
|
res, err := conn.CallTool(ctx, &CallToolParams{Name: "act"})
|
|
if err != nil {
|
|
t.Fatalf("CallTool() error = %v", err)
|
|
}
|
|
if res.NeedsInput() {
|
|
t.Fatal("NeedsInput() = true after auto-retry, want false")
|
|
}
|
|
|
|
// Sort the expected IDs for stable comparison.
|
|
if wantIDs, ok := tt.wantResult["ids"].([]any); ok {
|
|
slices.SortFunc(wantIDs, func(a, b any) int {
|
|
if a.(string) < b.(string) {
|
|
return -1
|
|
}
|
|
return 1
|
|
})
|
|
}
|
|
|
|
if diff := cmp.Diff(tt.wantResult, res.StructuredContent, ctrCmpOpts...); diff != "" {
|
|
t.Errorf("result mismatch (-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
|
|
// 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
|
|
requests InputRequestMap
|
|
wantRetries int
|
|
}{
|
|
{
|
|
name: "load shedding",
|
|
requests: InputRequestMap{},
|
|
wantRetries: maxLoadSheddingMultiRoundTripRetries,
|
|
},
|
|
{
|
|
name: "input request",
|
|
requests: InputRequestMap{"confirm": &ElicitParams{Message: "Again?"}},
|
|
wantRetries: maxMultiRoundTripRetries,
|
|
},
|
|
}
|
|
|
|
for _, tc := range testCases {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
var serverCalls atomic.Int32
|
|
srv := NewServer(testImpl, nil)
|
|
AddTool(srv, &Tool{Name: "loop"}, func(ctx context.Context, req *CallToolRequest, input struct{}) (*CallToolResult, any, error) {
|
|
serverCalls.Add(1)
|
|
return &CallToolResult{InputRequests: tc.requests, RequestState: "loop-state"}, nil, nil
|
|
})
|
|
|
|
conn := mustConnect(t, srv, &ClientOptions{
|
|
ElicitationHandler: func(_ context.Context, req *ElicitRequest) (*ElicitResult, error) {
|
|
return &ElicitResult{Action: "accept"}, nil
|
|
},
|
|
})
|
|
|
|
_, err := conn.CallTool(ctx, &CallToolParams{Name: "loop"})
|
|
if err == nil {
|
|
t.Fatal("CallTool() err = nil, want error for exceeded max retries")
|
|
}
|
|
if serverCalls.Load() != int32(tc.wantRetries) {
|
|
t.Errorf("serverCalls = %d, want %d", serverCalls.Load(), tc.wantRetries)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestMultiRoundTrip_ServerMiddleware(t *testing.T) {
|
|
// multiRoundTripToolHandler returns a ToolHandler (plain, non-generic) that requests
|
|
// the given inputRequests on the first call and returns the fulfilled
|
|
// response IDs on the second.
|
|
multiRoundTripToolHandler := func(inputRequests InputRequestMap) ToolHandler {
|
|
return func(ctx context.Context, req *CallToolRequest) (*CallToolResult, error) {
|
|
if len(req.Params.InputResponses) == 0 {
|
|
return &CallToolResult{
|
|
InputRequests: inputRequests,
|
|
RequestState: "state-1",
|
|
}, nil
|
|
}
|
|
var ids []string
|
|
for id := range req.Params.InputResponses {
|
|
ids = append(ids, id)
|
|
}
|
|
slices.Sort(ids)
|
|
content := &TextContent{Text: fmt.Sprintf("%v", ids)}
|
|
return &CallToolResult{Content: []Content{content}}, nil
|
|
}
|
|
}
|
|
|
|
tests := []struct {
|
|
name string
|
|
inputRequests InputRequestMap
|
|
wantText string
|
|
}{
|
|
{
|
|
name: "elicit via ToolHandler",
|
|
inputRequests: InputRequestMap{
|
|
"confirm": &ElicitParams{Message: "Sure?"},
|
|
},
|
|
wantText: "[confirm]",
|
|
},
|
|
{
|
|
name: "elicit and listRoots via ToolHandler",
|
|
inputRequests: InputRequestMap{
|
|
"confirm": &ElicitParams{Message: "OK?"},
|
|
"roots": &ListRootsParams{},
|
|
},
|
|
wantText: "[confirm roots]",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
srv := NewServer(testImpl, nil)
|
|
srv.AddTool(
|
|
&Tool{Name: "act", InputSchema: &jsonschema.Schema{Type: "object"}},
|
|
multiRoundTripToolHandler(tt.inputRequests),
|
|
)
|
|
|
|
// Connect with an OLD protocol version where multi-round-trip is not supported.
|
|
// The server middleware should handle it transparently.
|
|
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{
|
|
MultiRoundTrip: &MultiRoundTripOptions{Disabled: true},
|
|
ElicitationHandler: func(_ context.Context, req *ElicitRequest) (*ElicitResult, error) {
|
|
return &ElicitResult{Action: "accept"}, nil
|
|
},
|
|
})
|
|
c.AddRoots(&Root{URI: "file:///workspace", Name: "workspace"})
|
|
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() })
|
|
|
|
res, err := cs.CallTool(ctx, &CallToolParams{Name: "act"})
|
|
if err != nil {
|
|
t.Fatalf("CallTool() error = %v", err)
|
|
}
|
|
if got := res.Content[0].(*TextContent).Text; got != tt.wantText {
|
|
t.Errorf("result text = %q, want %q", got, tt.wantText)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// 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()
|
|
|
|
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{
|
|
Description: "Code review prompt",
|
|
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
|
|
},
|
|
})
|
|
|
|
res, err := conn.GetPrompt(ctx, &GetPromptParams{Name: "review"})
|
|
if err != nil {
|
|
t.Fatalf("GetPrompt() error = %v", err)
|
|
}
|
|
if res.NeedsInput() {
|
|
t.Fatal("NeedsInput() = true after auto-retry, want false")
|
|
}
|
|
if len(res.Messages) != 1 {
|
|
t.Fatalf("len(res.Messages) = %d, want 1", len(res.Messages))
|
|
}
|
|
if got := res.Messages[0].Content.(*TextContent).Text; got != "review this code" {
|
|
t.Errorf("message text = %q, want %q", got, "review this code")
|
|
}
|
|
}
|
|
|
|
func TestMultiRoundTrip_GetPrompt_ManualRetry(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{
|
|
Description: "Code review prompt",
|
|
Messages: []*PromptMessage{{Role: "user", Content: &TextContent{Text: "review this code"}}},
|
|
}, nil
|
|
})
|
|
|
|
conn := mustConnect(t, srv, &ClientOptions{
|
|
MultiRoundTrip: &MultiRoundTripOptions{Disabled: true},
|
|
})
|
|
|
|
res, err := conn.GetPrompt(ctx, &GetPromptParams{Name: "review"})
|
|
if err != nil {
|
|
t.Fatalf("GetPrompt() error = %v", err)
|
|
}
|
|
if !res.NeedsInput() {
|
|
t.Fatal("NeedsInput() = false, want true")
|
|
}
|
|
if _, ok := res.InputRequests["confirm"].(*ElicitParams); !ok {
|
|
t.Fatalf("InputRequests[confirm] type = %T, want *ElicitParams", res.InputRequests["confirm"])
|
|
}
|
|
|
|
res, err = conn.GetPrompt(ctx, &GetPromptParams{
|
|
Name: "review",
|
|
InputResponses: InputResponseMap{"confirm": &ElicitResult{Action: "accept"}},
|
|
RequestState: res.RequestState,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("GetPrompt() follow-up error = %v", err)
|
|
}
|
|
if res.NeedsInput() {
|
|
t.Fatal("NeedsInput() = true after follow-up, want false")
|
|
}
|
|
if len(res.Messages) != 1 {
|
|
t.Fatalf("len(res.Messages) = %d, want 1", len(res.Messages))
|
|
}
|
|
}
|
|
|
|
func TestMultiRoundTrip_ReadResource_AutoRetry(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
|
|
},
|
|
})
|
|
|
|
res, err := conn.ReadResource(ctx, &ReadResourceParams{URI: "test://data"})
|
|
if err != nil {
|
|
t.Fatalf("ReadResource() error = %v", err)
|
|
}
|
|
if res.NeedsInput() {
|
|
t.Fatal("NeedsInput() = true after auto-retry, want false")
|
|
}
|
|
if len(res.Contents) != 1 {
|
|
t.Fatalf("len(res.Contents) = %d, want 1", len(res.Contents))
|
|
}
|
|
if got := res.Contents[0].Text; got != "resource data" {
|
|
t.Errorf("resource text = %q, want %q", got, "resource data")
|
|
}
|
|
}
|
|
|
|
func TestMultiRoundTrip_ReadResource_ManualRetry(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{
|
|
MultiRoundTrip: &MultiRoundTripOptions{Disabled: true},
|
|
})
|
|
|
|
res, err := conn.ReadResource(ctx, &ReadResourceParams{URI: "test://data"})
|
|
if err != nil {
|
|
t.Fatalf("ReadResource() error = %v", err)
|
|
}
|
|
if !res.NeedsInput() {
|
|
t.Fatal("NeedsInput() = false, want true")
|
|
}
|
|
if _, ok := res.InputRequests["auth"].(*ElicitParams); !ok {
|
|
t.Fatalf("InputRequests[auth] type = %T, want *ElicitParams", res.InputRequests["auth"])
|
|
}
|
|
|
|
res, err = conn.ReadResource(ctx, &ReadResourceParams{
|
|
URI: "test://data",
|
|
InputResponses: InputResponseMap{"auth": &ElicitResult{Action: "accept"}},
|
|
RequestState: res.RequestState,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("ReadResource() follow-up error = %v", err)
|
|
}
|
|
if res.NeedsInput() {
|
|
t.Fatal("NeedsInput() = true after follow-up, want false")
|
|
}
|
|
if len(res.Contents) != 1 {
|
|
t.Fatalf("len(res.Contents) = %d, want 1", len(res.Contents))
|
|
}
|
|
}
|
|
|
|
// 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()
|
|
|
|
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: protocolVersion20260728})
|
|
if err != nil {
|
|
t.Fatalf("client.Connect() error = %v", err)
|
|
}
|
|
t.Cleanup(func() {
|
|
_ = cs.Close()
|
|
})
|
|
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
|
|
}
|