From 570838cb9405bd7b567c61e2e4da24a29f852120 Mon Sep 17 00:00:00 2001 From: Dmitry Sergeev Date: Fri, 9 Oct 2026 10:20:33 +0200 Subject: [PATCH 1/3] mcp: retry input-required results that carry only requestState SEP-2322 requires an InputRequiredResult to carry at least one of inputRequests and requestState, so a result with only requestState is valid, and the TypeScript SDK server sends one from inputRequired({ requestState }). clientMultiRoundTripMiddleware stopped as soon as InputRequests was nil and returned such a result to the caller as the final one. Also check the result type, so these results are retried with the requestState echoed, like an empty inputRequests map. A result with neither field is now retried until the load-shedding limit instead of being returned. --- mcp/mrtr.go | 6 +++- mcp/mrtr_test.go | 87 ++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 92 insertions(+), 1 deletion(-) diff --git a/mcp/mrtr.go b/mcp/mrtr.go index fdf98ece..161a1f84 100644 --- a/mcp/mrtr.go +++ b/mcp/mrtr.go @@ -34,6 +34,7 @@ type multiRoundTripResponse interface { inputRequests() map[string]InputRequest requestState() string hasContent() bool + NeedsInput() bool } // validateMultiRoundTripResult rejects a result that carries both content and @@ -79,7 +80,10 @@ func clientMultiRoundTripMiddleware() Middleware { return res, nil } reqMap := mrtrResult.inputRequests() - if reqMap == nil { + // An input-required result may carry only a requestState + // (SEP-2322 requires at least one of the two fields), so the + // result type, not the map alone, decides whether to retry. + if reqMap == nil && !mrtrResult.NeedsInput() { return res, nil } if len(reqMap) == 0 { diff --git a/mcp/mrtr_test.go b/mcp/mrtr_test.go index d9662198..0f5353ce 100644 --- a/mcp/mrtr_test.go +++ b/mcp/mrtr_test.go @@ -342,6 +342,93 @@ type wrappedCallToolParams struct{ *CallToolParams } // 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. +// rawJSONResult is a result sent exactly as given, standing in for a server +// built with another SDK. +type rawJSONResult struct { + ResultBase + raw string +} + +func (r *rawJSONResult) MarshalJSON() ([]byte, error) { return []byte(r.raw), nil } + +// TestMultiRoundTrip_AutoRetry_RequestStateOnly verifies that the client +// retries an input-required result that carries a requestState but no +// inputRequests, echoing the state, instead of returning it as the final +// result. The spec requires at least one of the two fields, so a result with +// only requestState is valid. +func TestMultiRoundTrip_AutoRetry_RequestStateOnly(t *testing.T) { + ctx := context.Background() + + srv := NewServer(testImpl, nil) + srv.AddTool( + &Tool{Name: "act", InputSchema: &jsonschema.Schema{Type: "object"}}, + func(_ context.Context, req *CallToolRequest) (*CallToolResult, error) { + return &CallToolResult{Content: []Content{&TextContent{Text: "state=" + req.Params.RequestState}}}, nil + }, + ) + // The SDK's own server never sends a requestState-only result, so send + // one directly on the first call. + srv.AddReceivingMiddleware(func(next MethodHandler) MethodHandler { + return func(ctx context.Context, method string, req Request) (Result, error) { + if method == methodCallTool && req.GetParams().(*CallToolParamsRaw).RequestState == "" { + return &rawJSONResult{raw: `{"resultType":"input_required","requestState":"continue"}`}, nil + } + return next(ctx, method, req) + } + }) + + cs := mustConnect(t, srv, nil) + + res, err := cs.CallTool(ctx, &CallToolParams{Name: "act"}) + if err != nil { + t.Fatalf("CallTool() error = %v", err) + } + if res.NeedsInput() { + t.Fatalf("CallTool() returned an input-required result: requestState = %q", res.RequestState) + } + if len(res.Content) != 1 { + t.Fatalf("CallTool() content = %v, want one content block", res.Content) + } + if got := res.Content[0].(*TextContent).Text; got != "state=continue" { + t.Errorf("result text = %q, want %q", got, "state=continue") + } +} + +// TestMultiRoundTrip_AutoRetry_InputRequiredWithoutFields verifies that an +// input-required result with neither inputRequests nor requestState is not +// returned as the final result. It is retried like load shedding until the +// load-shedding limit is reached. +func TestMultiRoundTrip_AutoRetry_InputRequiredWithoutFields(t *testing.T) { + ctx := context.Background() + + var calls atomic.Int32 + srv := NewServer(testImpl, nil) + srv.AddTool( + &Tool{Name: "act", InputSchema: &jsonschema.Schema{Type: "object"}}, + func(context.Context, *CallToolRequest) (*CallToolResult, error) { + return &CallToolResult{Content: []Content{&TextContent{Text: "done"}}}, nil + }, + ) + srv.AddReceivingMiddleware(func(next MethodHandler) MethodHandler { + return func(ctx context.Context, method string, req Request) (Result, error) { + if method == methodCallTool { + calls.Add(1) + return &rawJSONResult{raw: `{"resultType":"input_required"}`}, nil + } + return next(ctx, method, req) + } + }) + + cs := mustConnect(t, srv, nil) + + if _, err := cs.CallTool(ctx, &CallToolParams{Name: "act"}); err == nil { + t.Fatal("CallTool() err = nil, want error for an input-required result without inputRequests or requestState") + } + if got := calls.Load(); got != maxLoadSheddingMultiRoundTripRetries { + t.Errorf("tools/call requests = %d, want %d", got, maxLoadSheddingMultiRoundTripRetries) + } +} + func TestMultiRoundTrip_UnsupportedRetryParamsType(t *testing.T) { ctx := context.Background() From 67e17de5357f23b4491800f7ec1eec5b2dd21cb2 Mon Sep 17 00:00:00 2001 From: Dmitry Sergeev Date: Fri, 9 Oct 2026 17:11:46 +0200 Subject: [PATCH 2/3] mcp: let servers return input-required results with only requestState The server treated a result as input-required only when InputRequests was set. A handler that returned only a RequestState sent a complete result, and a resources/read handler doing so failed with "read handler returned nil information". Use one check, inputRequests or requestState, for the result type, the content validation, the empty content and messages fill, the resources/read early return and the legacy shim. For clients before 2026-07-28 the shim now calls the handler again with the state echoed. On the client, a requestState-only result counts as a round rather than load shedding, the same as in the shim. --- docs/server.md | 9 +- internal/docs/server.src.md | 9 +- mcp/mrtr.go | 31 ++++--- mcp/mrtr_test.go | 171 ++++++++++++++++++++++++++++++++++++ mcp/protocol.go | 2 +- mcp/server.go | 12 +-- mcp/server_test.go | 25 +++++- 7 files changed, 234 insertions(+), 25 deletions(-) diff --git a/docs/server.md b/docs/server.md index 764c96bf..97d1ac3f 100644 --- a/docs/server.md +++ b/docs/server.md @@ -854,9 +854,11 @@ The SDK supports this pattern from both sides without requiring callers to choose: - **Returning input requests from a handler.** Set `InputRequests` on the - result you return. + result you return. To have the client call again without asking it for + input, return only a `RequestState`; the client echoes it on the retry. Leave `Content` / `StructuredContent` empty — returning - both at once is a server bug and the SDK returns `-32603 InternalError`. + content together with `InputRequests` or `RequestState` is a server bug and + the SDK returns `-32603 InternalError`. - **Calling the legacy APIs.** `ServerSession.Elicit`, `ServerSession.CreateMessage(WithTools)`, and `ServerSession.ListRoots` remain available and work for both new and legacy clients. @@ -868,7 +870,8 @@ clients on a protocol version earlier than `2026-07-28`, the middleware intercepts any `InputRequiredResult` your handler returns, fulfils each input request itself by calling the legacy server-initiated APIs (`Elicit`, `CreateMessage`, `ListRoots`), and re-invokes your handler -with the responses already populated. If the handler asks for more input, +with the responses already populated. A result with only a `RequestState` +re-invokes it with that state and no responses. If the handler asks for more input, the middleware repeats this until the handler completes, up to the same limit of 10 rounds that the client-side middleware applies. This means a handler written in the MRTR style works against both old and new clients diff --git a/internal/docs/server.src.md b/internal/docs/server.src.md index 1b24c1d6..4c7825f9 100644 --- a/internal/docs/server.src.md +++ b/internal/docs/server.src.md @@ -335,9 +335,11 @@ The SDK supports this pattern from both sides without requiring callers to choose: - **Returning input requests from a handler.** Set `InputRequests` on the - result you return. + result you return. To have the client call again without asking it for + input, return only a `RequestState`; the client echoes it on the retry. Leave `Content` / `StructuredContent` empty — returning - both at once is a server bug and the SDK returns `-32603 InternalError`. + content together with `InputRequests` or `RequestState` is a server bug and + the SDK returns `-32603 InternalError`. - **Calling the legacy APIs.** `ServerSession.Elicit`, `ServerSession.CreateMessage(WithTools)`, and `ServerSession.ListRoots` remain available and work for both new and legacy clients. @@ -349,7 +351,8 @@ clients on a protocol version earlier than `2026-07-28`, the middleware intercepts any `InputRequiredResult` your handler returns, fulfils each input request itself by calling the legacy server-initiated APIs (`Elicit`, `CreateMessage`, `ListRoots`), and re-invokes your handler -with the responses already populated. If the handler asks for more input, +with the responses already populated. A result with only a `RequestState` +re-invokes it with that state and no responses. If the handler asks for more input, the middleware repeats this until the handler completes, up to the same limit of 10 rounds that the client-side middleware applies. This means a handler written in the MRTR style works against both old and new clients diff --git a/mcp/mrtr.go b/mcp/mrtr.go index 449acf9a..c73f3cca 100644 --- a/mcp/mrtr.go +++ b/mcp/mrtr.go @@ -37,18 +37,25 @@ type multiRoundTripResponse interface { NeedsInput() bool } +// asksForInput reports whether a handler's result is an input-required result. +// SEP-2322 requires at least one of inputRequests and requestState, so a +// result with only a requestState asks the client to continue too. +func asksForInput(res multiRoundTripResponse) bool { + return res.inputRequests() != nil || res.requestState() != "" +} + // validateMultiRoundTripResult rejects a result that carries both content and -// input requests. [annotateResultType] labels the result later, in -// [ServerSession.handle], so this function does not set resultType. +// input requests or a request state. [annotateResultType] labels the result +// later, in [ServerSession.handle], so this function does not set resultType. func validateMultiRoundTripResult(logger *slog.Logger, res multiRoundTripResponse) error { if res == nil { return nil } - if res.inputRequests() != nil && res.hasContent() { - logger.Warn("handler returned both content and inputRequests") + if asksForInput(res) && res.hasContent() { + logger.Warn("handler returned both content and inputRequests or requestState") return &jsonrpc.Error{ Code: jsonrpc.CodeInternalError, - Message: "server bug: result has both content and inputRequests", + Message: "server bug: result has both content and inputRequests or requestState", } } return nil @@ -86,7 +93,10 @@ func clientMultiRoundTripMiddleware() Middleware { if reqMap == nil && !mrtrResult.NeedsInput() { return res, nil } - if len(reqMap) == 0 { + // A requestState-only result asks to continue, so it counts as a + // round, not as load shedding. + continuing := reqMap == nil && mrtrResult.requestState() != "" + if len(reqMap) == 0 && !continuing { loadSheddingFailures++ } if loadSheddingFailures >= maxLoadSheddingMultiRoundTripRetries { @@ -116,7 +126,8 @@ func clientMultiRoundTripMiddleware() Middleware { // 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 // the requests by calling the client directly and reinvokes the handler with the responses, -// repeating until the handler completes or maxMultiRoundTripRetries is reached. +// repeating until the handler completes or maxMultiRoundTripRetries is reached. A result +// with only a RequestState reinvokes the handler with that state and no responses. func serverMultiRoundTripMiddleware() Middleware { return func(next MethodHandler) MethodHandler { return func(ctx context.Context, method string, req Request) (Result, error) { @@ -141,11 +152,11 @@ func serverMultiRoundTripMiddleware() Middleware { if !ok { return res, nil } - reqMap := mrtrResult.inputRequests() - if reqMap == nil { + if !asksForInput(mrtrResult) { return res, nil } - if len(reqMap) == 0 { + reqMap := mrtrResult.inputRequests() + if reqMap != nil && len(reqMap) == 0 { return nil, fmt.Errorf("the server is busy, retry later") } if retries >= maxMultiRoundTripRetries { diff --git a/mcp/mrtr_test.go b/mcp/mrtr_test.go index e073da91..6d5f5aef 100644 --- a/mcp/mrtr_test.go +++ b/mcp/mrtr_test.go @@ -429,6 +429,139 @@ func TestMultiRoundTrip_AutoRetry_InputRequiredWithoutFields(t *testing.T) { } } +// TestMultiRoundTrip_ServerRequestStateOnly verifies that a handler can return +// a result with only a RequestState for each multi round-trip method: the +// server labels it input_required, and the client retries it with the state +// echoed. +func TestMultiRoundTrip_ServerRequestStateOnly(t *testing.T) { + ctx := context.Background() + + srv := NewServer(testImpl, nil) + srv.AddTool( + &Tool{Name: "act", InputSchema: &jsonschema.Schema{Type: "object"}}, + func(_ context.Context, req *CallToolRequest) (*CallToolResult, error) { + if req.Params.RequestState == "" { + return &CallToolResult{RequestState: "continue"}, nil + } + return &CallToolResult{Content: []Content{&TextContent{Text: "state=" + req.Params.RequestState}}}, nil + }, + ) + srv.AddPrompt(&Prompt{Name: "review"}, func(_ context.Context, req *GetPromptRequest) (*GetPromptResult, error) { + if req.Params.RequestState == "" { + return &GetPromptResult{RequestState: "continue"}, nil + } + return &GetPromptResult{ + Messages: []*PromptMessage{{Role: "user", Content: &TextContent{Text: "state=" + req.Params.RequestState}}}, + }, nil + }) + srv.AddResource(&Resource{URI: "test://data", Name: "data"}, func(_ context.Context, req *ReadResourceRequest) (*ReadResourceResult, error) { + if req.Params.RequestState == "" { + return &ReadResourceResult{RequestState: "continue"}, nil + } + return &ReadResourceResult{ + Contents: []*ResourceContents{{URI: "test://data", Text: "state=" + req.Params.RequestState}}, + }, nil + }) + + cs := mustConnect(t, srv, nil) + const want = "state=continue" + + t.Run("tools/call", func(t *testing.T) { + res, err := cs.CallTool(ctx, &CallToolParams{Name: "act"}) + if err != nil { + t.Fatalf("CallTool() error = %v", err) + } + if len(res.Content) != 1 { + t.Fatalf("CallTool() content = %v, want one content block", res.Content) + } + if got := res.Content[0].(*TextContent).Text; got != want { + t.Errorf("result text = %q, want %q", got, want) + } + }) + t.Run("prompts/get", func(t *testing.T) { + res, err := cs.GetPrompt(ctx, &GetPromptParams{Name: "review"}) + if err != nil { + t.Fatalf("GetPrompt() error = %v", err) + } + if len(res.Messages) != 1 { + t.Fatalf("GetPrompt() messages = %v, want one message", res.Messages) + } + if got := res.Messages[0].Content.(*TextContent).Text; got != want { + t.Errorf("message text = %q, want %q", got, want) + } + }) + t.Run("resources/read", func(t *testing.T) { + res, err := cs.ReadResource(ctx, &ReadResourceParams{URI: "test://data"}) + if err != nil { + t.Fatalf("ReadResource() error = %v", err) + } + if len(res.Contents) != 1 { + t.Fatalf("ReadResource() contents = %v, want one entry", res.Contents) + } + if got := res.Contents[0].Text; got != want { + t.Errorf("contents text = %q, want %q", got, want) + } + }) +} + +// TestMultiRoundTrip_RequestStateOnly_MaxRetries verifies that the client +// counts requestState-only results as rounds, not as load shedding: a server +// that keeps asking to continue is called maxMultiRoundTripRetries times. +func TestMultiRoundTrip_RequestStateOnly_MaxRetries(t *testing.T) { + ctx := context.Background() + + var calls atomic.Int32 + srv := NewServer(testImpl, nil) + srv.AddTool( + &Tool{Name: "loop", InputSchema: &jsonschema.Schema{Type: "object"}}, + func(context.Context, *CallToolRequest) (*CallToolResult, error) { + calls.Add(1) + return &CallToolResult{RequestState: "again"}, nil + }, + ) + + cs := mustConnect(t, srv, nil) + + _, err := cs.CallTool(ctx, &CallToolParams{Name: "loop"}) + if err == nil { + t.Fatal("CallTool() err = nil, want error for exceeded max retries") + } + if !strings.Contains(err.Error(), "exceeded maximum retries") { + t.Errorf("CallTool() err = %v, want it to mention the retry limit", err) + } + if got := calls.Load(); got != maxMultiRoundTripRetries { + t.Errorf("tool calls = %d, want %d", got, maxMultiRoundTripRetries) + } +} + +// TestMultiRoundTrip_ContentWithRequestState verifies that a result with both +// content and a RequestState is rejected as a server bug, like content with +// InputRequests. +func TestMultiRoundTrip_ContentWithRequestState(t *testing.T) { + ctx := context.Background() + + srv := NewServer(testImpl, nil) + srv.AddTool( + &Tool{Name: "mixed", InputSchema: &jsonschema.Schema{Type: "object"}}, + func(context.Context, *CallToolRequest) (*CallToolResult, error) { + return &CallToolResult{ + Content: []Content{&TextContent{Text: "done"}}, + RequestState: "stray", + }, nil + }, + ) + + cs := mustConnect(t, srv, nil) + + _, err := cs.CallTool(ctx, &CallToolParams{Name: "mixed"}) + if err == nil { + t.Fatal("CallTool() err = nil, want a server error for content with a requestState") + } + if !strings.Contains(err.Error(), "both content and inputRequests or requestState") { + t.Errorf("CallTool() err = %v, want the content/requestState server error", err) + } +} + func TestMultiRoundTrip_UnsupportedRetryParamsType(t *testing.T) { ctx := context.Background() @@ -737,6 +870,44 @@ func TestMultiRoundTrip_ServerMiddleware_MaxRetries(t *testing.T) { } } +// TestMultiRoundTrip_ServerMiddleware_RequestStateOnly verifies that, for a +// client on an older protocol version, the server-side shim calls the handler +// again with the state echoed when it returns only a RequestState. +func TestMultiRoundTrip_ServerMiddleware_RequestStateOnly(t *testing.T) { + ctx := context.Background() + + var calls atomic.Int32 + srv := NewServer(testImpl, nil) + srv.AddTool( + &Tool{Name: "act", InputSchema: &jsonschema.Schema{Type: "object"}}, + func(_ context.Context, req *CallToolRequest) (*CallToolResult, error) { + calls.Add(1) + if req.Params.RequestState == "" { + return &CallToolResult{RequestState: "continue"}, nil + } + return &CallToolResult{Content: []Content{&TextContent{Text: "state=" + req.Params.RequestState}}}, nil + }, + ) + + cs := mustConnectOldProtocol(t, srv, &ClientOptions{ + MultiRoundTrip: &MultiRoundTripOptions{Disabled: true}, + }) + + res, err := cs.CallTool(ctx, &CallToolParams{Name: "act"}) + if err != nil { + t.Fatalf("CallTool() error = %v", err) + } + if len(res.Content) != 1 { + t.Fatalf("CallTool() content = %v, want one content block", res.Content) + } + if got := res.Content[0].(*TextContent).Text; got != "state=continue" { + t.Errorf("result text = %q, want %q", got, "state=continue") + } + if got := calls.Load(); got != 2 { + t.Errorf("tool calls = %d, want 2", got) + } +} + func TestMultiRoundTrip_GetPrompt_AutoRetry(t *testing.T) { ctx := context.Background() diff --git a/mcp/protocol.go b/mcp/protocol.go index ff6fb516..22844de2 100644 --- a/mcp/protocol.go +++ b/mcp/protocol.go @@ -49,7 +49,7 @@ func annotateResultType(res Result) { case multiRoundTripResponse: // These results are complete or input_required, so label them by // whether the handler asked for more client input. - if r.inputRequests() != nil { + if asksForInput(r) { r.setResultType(resultTypeInputRequired) } else { r.setResultType(resultTypeComplete) diff --git a/mcp/server.go b/mcp/server.go index 16aad6a9..d9b9ad50 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -441,10 +441,10 @@ func toolForErr[In, Out any](t *Tool, h ToolHandlerFor[In, Out], cache *SchemaCa } // Marshal the output and put the RawMessage in the StructuredContent field. - // Skip when the handler returned input requests (multi round-trip): content and - // inputRequests are mutually exclusive on the wire. + // Skip when the handler asked for input (multi round-trip): content and + // an input-required result are mutually exclusive on the wire. var outval any = out - if res.InputRequests != nil { + if asksForInput(res) { outval = nil } else if elemZero != nil { // Avoid typed nil, which will serialize as JSON null. @@ -957,7 +957,7 @@ func (s *Server) getPrompt(ctx context.Context, req *GetPromptRequest) (*GetProm if err := validateMultiRoundTripResult(s.opts.Logger, res); err != nil { return nil, err } - if res.Messages == nil && res.InputRequests == nil { + if res.Messages == nil && !asksForInput(res) { res2 := *res res2.Messages = []*PromptMessage{} // avoid "null" res = &res2 @@ -1065,7 +1065,7 @@ func (s *Server) callTool(ctx context.Context, req *CallToolRequest) (*CallToolR if err := validateMultiRoundTripResult(s.opts.Logger, res); err != nil { return nil, err } - if res.Content == nil && res.InputRequests == nil { + if res.Content == nil && !asksForInput(res) { res2 := *res res2.Content = []Content{} // avoid "null" res = &res2 @@ -1134,7 +1134,7 @@ func (s *Server) readResource(ctx context.Context, req *ReadResourceRequest) (*R return nil, err } s.resolveCacheable(ctx, req, &res.Cacheable) - if res.InputRequests != nil { + if asksForInput(res) { return res, nil } if res.Contents == nil { diff --git a/mcp/server_test.go b/mcp/server_test.go index 64e7de0d..314afefa 100644 --- a/mcp/server_test.go +++ b/mcp/server_test.go @@ -1637,8 +1637,8 @@ func TestServerSessionHandle_SetsResultTypeWhenMiddlewareShortCircuits(t *testin want: resultTypeComplete, }, { - // setResultType is unexported, so populating InputRequests is - // the only way a middleware can ask for input. + // setResultType is unexported, so populating InputRequests or + // RequestState is the only way a middleware can ask for input. name: "tools/call input required", method: methodCallTool, params: newProtocolParams(map[string]any{"name": "tool", "arguments": map[string]any{}}), @@ -1647,6 +1647,27 @@ func TestServerSessionHandle_SetsResultTypeWhenMiddlewareShortCircuits(t *testin }, want: resultTypeInputRequired, }, + { + name: "tools/call request state only", + method: methodCallTool, + params: newProtocolParams(map[string]any{"name": "tool", "arguments": map[string]any{}}), + result: &CallToolResult{RequestState: "continue"}, + want: resultTypeInputRequired, + }, + { + name: "prompts/get request state only", + method: methodGetPrompt, + params: newProtocolParams(map[string]any{"name": "prompt"}), + result: &GetPromptResult{RequestState: "continue"}, + want: resultTypeInputRequired, + }, + { + name: "resources/read request state only", + method: methodReadResource, + params: newProtocolParams(map[string]any{"uri": "test://resource"}), + result: &ReadResourceResult{RequestState: "continue"}, + want: resultTypeInputRequired, + }, } for _, tc := range tests { From ca9222f595e7dfe3032285a3342fd6e48aa42efa Mon Sep 17 00:00:00 2001 From: Dmitry Sergeev Date: Sat, 10 Oct 2026 15:44:05 +0200 Subject: [PATCH 3/3] mcp: treat any input-required result with a requestState as a round Apply the review suggestions: the client counts a round as load shedding, and the legacy server path returns "busy", only when the result has neither input requests nor a requestState. An empty inputRequests map with a requestState now continues like a requestState-only result. Also move the TestMultiRoundTrip_UnsupportedRetryParamsType comment back above its function. --- docs/server.md | 4 +- internal/docs/server.src.md | 4 +- mcp/mrtr.go | 12 ++--- mcp/mrtr_test.go | 90 ++++++++++++++++++++++++++++--------- 4 files changed, 79 insertions(+), 31 deletions(-) diff --git a/docs/server.md b/docs/server.md index 97d1ac3f..19674b0c 100644 --- a/docs/server.md +++ b/docs/server.md @@ -870,8 +870,8 @@ clients on a protocol version earlier than `2026-07-28`, the middleware intercepts any `InputRequiredResult` your handler returns, fulfils each input request itself by calling the legacy server-initiated APIs (`Elicit`, `CreateMessage`, `ListRoots`), and re-invokes your handler -with the responses already populated. A result with only a `RequestState` -re-invokes it with that state and no responses. If the handler asks for more input, +with the responses already populated. A result with a `RequestState` and no +input requests re-invokes it with that state and no responses. If the handler asks for more input, the middleware repeats this until the handler completes, up to the same limit of 10 rounds that the client-side middleware applies. This means a handler written in the MRTR style works against both old and new clients diff --git a/internal/docs/server.src.md b/internal/docs/server.src.md index 4c7825f9..e161184a 100644 --- a/internal/docs/server.src.md +++ b/internal/docs/server.src.md @@ -351,8 +351,8 @@ clients on a protocol version earlier than `2026-07-28`, the middleware intercepts any `InputRequiredResult` your handler returns, fulfils each input request itself by calling the legacy server-initiated APIs (`Elicit`, `CreateMessage`, `ListRoots`), and re-invokes your handler -with the responses already populated. A result with only a `RequestState` -re-invokes it with that state and no responses. If the handler asks for more input, +with the responses already populated. A result with a `RequestState` and no +input requests re-invokes it with that state and no responses. If the handler asks for more input, the middleware repeats this until the handler completes, up to the same limit of 10 rounds that the client-side middleware applies. This means a handler written in the MRTR style works against both old and new clients diff --git a/mcp/mrtr.go b/mcp/mrtr.go index c73f3cca..5262cfd4 100644 --- a/mcp/mrtr.go +++ b/mcp/mrtr.go @@ -93,10 +93,9 @@ func clientMultiRoundTripMiddleware() Middleware { if reqMap == nil && !mrtrResult.NeedsInput() { return res, nil } - // A requestState-only result asks to continue, so it counts as a - // round, not as load shedding. - continuing := reqMap == nil && mrtrResult.requestState() != "" - if len(reqMap) == 0 && !continuing { + // A result with a requestState asks to continue, so it counts as + // a round, not as load shedding. + if len(reqMap) == 0 && mrtrResult.requestState() == "" { loadSheddingFailures++ } if loadSheddingFailures >= maxLoadSheddingMultiRoundTripRetries { @@ -127,7 +126,8 @@ func clientMultiRoundTripMiddleware() Middleware { // InputRequests and the client does not support multi-round-trip, the middleware fulfills // the requests by calling the client directly and reinvokes the handler with the responses, // repeating until the handler completes or maxMultiRoundTripRetries is reached. A result -// with only a RequestState reinvokes the handler with that state and no responses. +// with a RequestState and no input requests reinvokes the handler with that state and no +// responses. func serverMultiRoundTripMiddleware() Middleware { return func(next MethodHandler) MethodHandler { return func(ctx context.Context, method string, req Request) (Result, error) { @@ -156,7 +156,7 @@ func serverMultiRoundTripMiddleware() Middleware { return res, nil } reqMap := mrtrResult.inputRequests() - if reqMap != nil && len(reqMap) == 0 { + if len(reqMap) == 0 && mrtrResult.requestState() == "" { return nil, fmt.Errorf("the server is busy, retry later") } if retries >= maxMultiRoundTripRetries { diff --git a/mcp/mrtr_test.go b/mcp/mrtr_test.go index 6d5f5aef..d33f56c2 100644 --- a/mcp/mrtr_test.go +++ b/mcp/mrtr_test.go @@ -338,10 +338,6 @@ func TestMultiRoundTrip_ReadResource_AutoRetryDoesNotMutateCallerParams(t *testi // 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. // rawJSONResult is a result sent exactly as given, standing in for a server // built with another SDK. type rawJSONResult struct { @@ -562,6 +558,10 @@ func TestMultiRoundTrip_ContentWithRequestState(t *testing.T) { } } +// 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() @@ -613,6 +613,7 @@ func TestMultiRoundTrip_MaxRetries(t *testing.T) { testCases := []struct { name string requests InputRequestMap + state string wantRetries int }{ { @@ -620,9 +621,16 @@ func TestMultiRoundTrip_MaxRetries(t *testing.T) { requests: InputRequestMap{}, wantRetries: maxLoadSheddingMultiRoundTripRetries, }, + { + name: "empty requests with state", + requests: InputRequestMap{}, + state: "loop-state", + wantRetries: maxMultiRoundTripRetries, + }, { name: "input request", requests: InputRequestMap{"confirm": &ElicitParams{Message: "Again?"}}, + state: "loop-state", wantRetries: maxMultiRoundTripRetries, }, } @@ -635,7 +643,7 @@ func TestMultiRoundTrip_MaxRetries(t *testing.T) { 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 + return &CallToolResult{InputRequests: tc.requests, RequestState: tc.state}, nil, nil }) conn := mustConnect(t, srv, &ClientOptions{ @@ -872,8 +880,57 @@ func TestMultiRoundTrip_ServerMiddleware_MaxRetries(t *testing.T) { // TestMultiRoundTrip_ServerMiddleware_RequestStateOnly verifies that, for a // client on an older protocol version, the server-side shim calls the handler -// again with the state echoed when it returns only a RequestState. +// again with the state echoed when it returns a RequestState and no input +// requests. func TestMultiRoundTrip_ServerMiddleware_RequestStateOnly(t *testing.T) { + for _, tc := range []struct { + name string + requests InputRequestMap + }{ + {name: "nil requests", requests: nil}, + {name: "empty requests", requests: InputRequestMap{}}, + } { + t.Run(tc.name, func(t *testing.T) { + ctx := context.Background() + + var calls atomic.Int32 + srv := NewServer(testImpl, nil) + srv.AddTool( + &Tool{Name: "act", InputSchema: &jsonschema.Schema{Type: "object"}}, + func(_ context.Context, req *CallToolRequest) (*CallToolResult, error) { + calls.Add(1) + if req.Params.RequestState == "" { + return &CallToolResult{InputRequests: tc.requests, RequestState: "continue"}, nil + } + return &CallToolResult{Content: []Content{&TextContent{Text: "state=" + req.Params.RequestState}}}, nil + }, + ) + + cs := mustConnectOldProtocol(t, srv, &ClientOptions{ + MultiRoundTrip: &MultiRoundTripOptions{Disabled: true}, + }) + + res, err := cs.CallTool(ctx, &CallToolParams{Name: "act"}) + if err != nil { + t.Fatalf("CallTool() error = %v", err) + } + if len(res.Content) != 1 { + t.Fatalf("CallTool() content = %v, want one content block", res.Content) + } + if got := res.Content[0].(*TextContent).Text; got != "state=continue" { + t.Errorf("result text = %q, want %q", got, "state=continue") + } + if got := calls.Load(); got != 2 { + t.Errorf("tool calls = %d, want 2", got) + } + }) + } +} + +// TestMultiRoundTrip_ServerMiddleware_Busy verifies that, for a client on an +// older protocol version, the server-side shim reports an empty input request +// map without a RequestState as busy instead of calling the handler again. +func TestMultiRoundTrip_ServerMiddleware_Busy(t *testing.T) { ctx := context.Background() var calls atomic.Int32 @@ -882,10 +939,7 @@ func TestMultiRoundTrip_ServerMiddleware_RequestStateOnly(t *testing.T) { &Tool{Name: "act", InputSchema: &jsonschema.Schema{Type: "object"}}, func(_ context.Context, req *CallToolRequest) (*CallToolResult, error) { calls.Add(1) - if req.Params.RequestState == "" { - return &CallToolResult{RequestState: "continue"}, nil - } - return &CallToolResult{Content: []Content{&TextContent{Text: "state=" + req.Params.RequestState}}}, nil + return &CallToolResult{InputRequests: InputRequestMap{}}, nil }, ) @@ -893,18 +947,12 @@ func TestMultiRoundTrip_ServerMiddleware_RequestStateOnly(t *testing.T) { MultiRoundTrip: &MultiRoundTripOptions{Disabled: true}, }) - res, err := cs.CallTool(ctx, &CallToolParams{Name: "act"}) - if err != nil { - t.Fatalf("CallTool() error = %v", err) - } - if len(res.Content) != 1 { - t.Fatalf("CallTool() content = %v, want one content block", res.Content) - } - if got := res.Content[0].(*TextContent).Text; got != "state=continue" { - t.Errorf("result text = %q, want %q", got, "state=continue") + _, err := cs.CallTool(ctx, &CallToolParams{Name: "act"}) + if err == nil || !strings.Contains(err.Error(), "busy") { + t.Fatalf("CallTool() error = %v, want a busy error", err) } - if got := calls.Load(); got != 2 { - t.Errorf("tool calls = %d, want 2", got) + if got := calls.Load(); got != 1 { + t.Errorf("tool calls = %d, want 1", got) } }