diff --git a/docs/server.md b/docs/server.md index 764c96bf..19674b0c 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 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 1b24c1d6..e161184a 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 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 90c888c8..5262cfd4 100644 --- a/mcp/mrtr.go +++ b/mcp/mrtr.go @@ -34,20 +34,28 @@ type multiRoundTripResponse interface { inputRequests() map[string]InputRequest requestState() string hasContent() bool + 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 @@ -79,10 +87,15 @@ 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 { + // 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 { @@ -112,7 +125,9 @@ 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 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) { @@ -137,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 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 fa968c51..d33f56c2 100644 --- a/mcp/mrtr_test.go +++ b/mcp/mrtr_test.go @@ -338,6 +338,226 @@ func TestMultiRoundTrip_ReadResource_AutoRetryDoesNotMutateCallerParams(t *testi // retry does not recognize. type wrappedCallToolParams struct{ *CallToolParams } +// 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) + } +} + +// 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) + } +} + // 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 @@ -393,6 +613,7 @@ func TestMultiRoundTrip_MaxRetries(t *testing.T) { testCases := []struct { name string requests InputRequestMap + state string wantRetries int }{ { @@ -400,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, }, } @@ -415,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{ @@ -650,6 +878,84 @@ 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 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 + srv := NewServer(testImpl, nil) + srv.AddTool( + &Tool{Name: "act", InputSchema: &jsonschema.Schema{Type: "object"}}, + func(_ context.Context, req *CallToolRequest) (*CallToolResult, error) { + calls.Add(1) + return &CallToolResult{InputRequests: InputRequestMap{}}, nil + }, + ) + + cs := mustConnectOldProtocol(t, srv, &ClientOptions{ + MultiRoundTrip: &MultiRoundTripOptions{Disabled: true}, + }) + + _, 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 != 1 { + t.Errorf("tool calls = %d, want 1", 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 6f3566ec..3efa3b83 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 {