diff --git a/docs/concepts/tasks/tasks.md b/docs/concepts/tasks/tasks.md index c1ab0c23a..a388197b5 100644 --- a/docs/concepts/tasks/tasks.md +++ b/docs/concepts/tasks/tasks.md @@ -210,6 +210,14 @@ Supported input request methods: | --- | --- | | `elicitation/create` | | | `sampling/createMessage` | | +| `roots/list` | | + +Tools that use MRTR directly by throwing +also compose with task execution. When a task-enabled call throws with input requests, the SDK +publishes them through the task store and reruns the tool after `tasks/update`, with + and + populated for the retry. Calls that +do not opt in to Tasks keep the normal MRTR behavior. Per SEP-2663: @@ -358,10 +366,6 @@ compatibility bridge for the previous experimental API. synchronously and then transition its remaining work to a background task. Use a custom if you need that pattern. -- **`roots/list` as an input request**: the server SDK routes `RequestRootsAsync` through the - task channel when called from inside a task scope, but the client SDK does not currently - dispatch a handler for that method. Avoid calling `server.RequestRootsAsync` from within a - task scope until client-side support is added. - **`ServerCapabilities.Extensions` round-trip**: the dictionary is typed as `IDictionary` so its values cannot be deserialized by the source generator. The negotiated extension surfaces correctly at the wire level, but round-tripping arbitrary diff --git a/src/Common/InputRequiredRequestRunner.cs b/src/Common/InputRequiredRequestRunner.cs new file mode 100644 index 000000000..454d71fe3 --- /dev/null +++ b/src/Common/InputRequiredRequestRunner.cs @@ -0,0 +1,179 @@ +using System.Text.Json; +using System.Text.Json.Nodes; +using System.Text.Json.Serialization.Metadata; + +namespace ModelContextProtocol.Protocol; + +/// +/// Runs request handlers that can require one or more rounds of additional input. +/// +internal static class InputRequiredRequestRunner +{ + private const int MaxRetries = 10; + + private static readonly JsonTypeInfo> s_inputResponsesTypeInfo = + (JsonTypeInfo>)McpJsonUtilities.DefaultOptions.GetTypeInfo( + typeof(IDictionary)); + + private static readonly JsonTypeInfo s_inputRequiredResultTypeInfo = + (JsonTypeInfo)McpJsonUtilities.DefaultOptions.GetTypeInfo(typeof(InputRequiredResult)); + + /// + /// Invokes a handler until it returns a final result, normalizing both thrown and returned + /// values into the same retry flow. + /// + internal static async Task RunAsync( + TRequest request, + Func> invoke, + Func getReturnedInputRequiredResult, + Func? createDirectResult, + Func> prepareRetry, + Func createFailure, + CancellationToken cancellationToken) + { + for (int retry = 0; ; retry++) + { + InputRequiredResult inputRequiredResult; + Exception? inputRequiredException = null; + + try + { + TResult result = await invoke(request, cancellationToken).ConfigureAwait(false); + if (getReturnedInputRequiredResult(result) is not { } returnedInputRequiredResult) + { + return result; + } + + inputRequiredResult = returnedInputRequiredResult; + } + catch (InputRequiredException ex) + { + inputRequiredResult = ex.Result; + inputRequiredException = ex; + } + + if (createDirectResult is not null) + { + return createDirectResult(inputRequiredResult); + } + + if (inputRequiredResult.InputRequests is not { Count: > 0 } && + inputRequiredResult.RequestState is null) + { + throw createFailure( + "A tool returned an input-required result without input requests or request state.", + inputRequiredException); + } + + if (retry >= MaxRetries) + { + throw createFailure( + $"MRTR-native tool exceeded {MaxRetries} retry rounds without completing.", + inputRequiredException); + } + + request = await prepareRetry( + request, + inputRequiredResult, + inputRequiredException, + cancellationToken).ConfigureAwait(false); + } + } + + /// + /// Resolves a batch concurrently, cancelling sibling requests when any resolver fails. + /// + internal static async Task> ResolveInputRequestsAsync( + IDictionary inputRequests, + Func> resolveInputRequest, + CancellationToken cancellationToken) + { + using var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken); + var keyedTasks = new (string Key, Task ResponseTask)[inputRequests.Count]; + + int index = 0; + foreach (var pair in inputRequests) + { + keyedTasks[index++] = (pair.Key, ResolveAndCancelSiblingsAsync(pair.Value)); + } + + await Task.WhenAll(Array.ConvertAll(keyedTasks, static item => item.ResponseTask)).ConfigureAwait(false); + + var responses = new Dictionary(keyedTasks.Length); + foreach (var (key, responseTask) in keyedTasks) + { + responses[key] = responseTask.Result; + } + + return responses; + + async Task ResolveAndCancelSiblingsAsync(InputRequest inputRequest) + { + try + { + return await resolveInputRequest(inputRequest, linkedCts.Token).ConfigureAwait(false); + } + catch + { + try + { + linkedCts.Cancel(); + } + catch + { + // Preserve the resolver failure. Awaiting Task.WhenAll observes every sibling outcome. + } + + throw; + } + } + } + + /// + /// Clones request parameters and applies the response and state for the next round, removing + /// values left over from the previous round when the current result omits them. + /// + internal static JsonObject CreateRetryParams( + JsonNode? requestParams, + IDictionary? inputResponses, + string? requestState) + { + var paramsObject = requestParams?.DeepClone() as JsonObject ?? new JsonObject(); + + if (inputResponses is not null) + { + paramsObject["inputResponses"] = JsonSerializer.SerializeToNode(inputResponses, s_inputResponsesTypeInfo); + } + else + { + paramsObject.Remove("inputResponses"); + } + + if (requestState is not null) + { + paramsObject["requestState"] = requestState; + } + else + { + paramsObject.Remove("requestState"); + } + + return paramsObject; + } + + /// + /// Detects a serialized returned through an alternate-result path. + /// + internal static InputRequiredResult? GetReturnedInputRequiredResult(JsonNode? result) + { + if (result is JsonObject resultObject && + resultObject.TryGetPropertyValue("resultType", out var resultTypeNode) && + resultTypeNode?.GetValueKind() == JsonValueKind.String && + resultTypeNode.GetValue() == "input_required") + { + return JsonSerializer.Deserialize(result, s_inputRequiredResultTypeInfo); + } + + return null; + } +} diff --git a/src/ModelContextProtocol.Core/Client/McpClientImpl.cs b/src/ModelContextProtocol.Core/Client/McpClientImpl.cs index a62ad0eea..3cbced5bb 100644 --- a/src/ModelContextProtocol.Core/Client/McpClientImpl.cs +++ b/src/ModelContextProtocol.Core/Client/McpClientImpl.cs @@ -187,42 +187,10 @@ public override async ValueTask> ResolveInput IDictionary inputRequests, CancellationToken cancellationToken) { - // Resolve all input requests concurrently. If any fails, cancel the rest so user-facing - // handlers (sampling/elicitation prompts) don't keep running for a request whose caller - // has already given up, and ensure exceptions from late-completing tasks are observed. - using var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken); - - var keyed = new (string Key, Task Task)[inputRequests.Count]; - int i = 0; - foreach (var kvp in inputRequests) - { - keyed[i++] = (kvp.Key, ResolveInputRequestAsync(kvp.Value, linkedCts.Token)); - } - - try - { - await Task.WhenAll(Array.ConvertAll(keyed, k => k.Task)).ConfigureAwait(false); - } - catch - { - linkedCts.Cancel(); - try - { - await Task.WhenAll(Array.ConvertAll(keyed, k => k.Task)).ConfigureAwait(false); - } - catch - { - // Observed; the original exception is the one we want to surface. - } - throw; - } - - var responses = new Dictionary(keyed.Length); - foreach (var (key, task) in keyed) - { - responses[key] = task.Result; - } - return responses; + return await InputRequiredRequestRunner.ResolveInputRequestsAsync( + inputRequests, + ResolveInputRequestAsync, + cancellationToken).ConfigureAwait(false); } private async Task ResolveInputRequestAsync(InputRequest inputRequest, CancellationToken cancellationToken) @@ -693,74 +661,44 @@ request.Params is System.Text.Json.Nodes.JsonObject paramsObjForHeaders && } } - const int maxRetries = 10; - InjectRequestMetaIfNeeded(request); - - for (int attempt = 0; attempt <= maxRetries; attempt++) - { - JsonRpcResponse response = await _sessionHandler.SendRequestAsync(request, cancellationToken).ConfigureAwait(false); - - // Check if the result is an InputRequiredResult by looking at result_type. - if (response.Result is JsonObject resultObj && - resultObj.TryGetPropertyValue("resultType", out var resultTypeNode) && - resultTypeNode?.GetValue() is "input_required") + return await InputRequiredRequestRunner.RunAsync( + request, + (currentRequest, token) => _sessionHandler.SendRequestAsync(currentRequest, token), + static response => InputRequiredRequestRunner.GetReturnedInputRequiredResult(response.Result), + (Func?)null, + PrepareRetryAsync, + static (message, innerException) => new McpException(message, innerException), + cancellationToken).ConfigureAwait(false); + + async Task PrepareRetryAsync( + JsonRpcRequest currentRequest, + InputRequiredResult inputRequiredResult, + Exception? _, + CancellationToken retryCancellationToken) + { + WarnIfInputRequiredResultOnNonMrtrSession(currentRequest.Method); + + IDictionary? inputResponses = null; + if (inputRequiredResult.InputRequests is { Count: > 0 } inputRequests) { - WarnIfInputRequiredResultOnNonMrtrSession(request.Method); - - var inputRequiredResult = JsonSerializer.Deserialize(response.Result, McpJsonUtilities.JsonContext.Default.InputRequiredResult) - ?? throw new JsonException("Failed to deserialize InputRequiredResult."); - - if (inputRequiredResult.InputRequests is { Count: > 0 } inputRequests) - { - IDictionary inputResponses = - await ResolveInputRequestsAsync(inputRequests, cancellationToken).ConfigureAwait(false); - - // Clone the original request params and add inputResponses + requestState for the retry. - var paramsObj = request.Params?.DeepClone() as JsonObject ?? new JsonObject(); - - paramsObj["inputResponses"] = JsonSerializer.SerializeToNode( - inputResponses, McpJsonUtilities.JsonContext.Default.IDictionaryStringInputResponse); - - if (inputRequiredResult.RequestState is { } requestState) - { - paramsObj["requestState"] = requestState; - } - else - { - // Strip any stale requestState carried over from the previous round's clone so - // the server doesn't see a continuation token the current round is not using. - paramsObj.Remove("requestState"); - } - - request = new JsonRpcRequest { Method = request.Method, Params = paramsObj, Context = request.Context }; - InjectRequestMetaIfNeeded(request); - } - else if (inputRequiredResult.RequestState is not null) - { - // No input requests but has requestState (e.g., load shedding) - just retry with state. - var paramsObj = request.Params?.DeepClone() as JsonObject ?? new JsonObject(); - paramsObj["requestState"] = inputRequiredResult.RequestState; - paramsObj.Remove("inputResponses"); - - request = new JsonRpcRequest { Method = request.Method, Params = paramsObj, Context = request.Context }; - InjectRequestMetaIfNeeded(request); - } - else - { - // An input_required result carrying neither inputRequests nor requestState is - // malformed: there is nothing to resolve and nothing to continue, so retrying the - // unchanged request would just loop until maxRetries. Fail fast instead. - throw new McpException("Server returned an InputRequiredResult without inputRequests or requestState."); - } - - continue; // retry with the updated request + inputResponses = await ResolveInputRequestsAsync( + inputRequests, + retryCancellationToken).ConfigureAwait(false); } - return response; + var retryRequest = new JsonRpcRequest + { + Method = currentRequest.Method, + Params = InputRequiredRequestRunner.CreateRetryParams( + currentRequest.Params, + inputResponses, + inputRequiredResult.RequestState), + Context = currentRequest.Context, + }; + InjectRequestMetaIfNeeded(retryRequest); + return retryRequest; } - - throw new McpException($"Server returned InputRequiredResult more than {maxRetries} times."); } /// diff --git a/src/ModelContextProtocol.Core/ModelContextProtocol.Core.csproj b/src/ModelContextProtocol.Core/ModelContextProtocol.Core.csproj index 3fbef0377..9fae024ff 100644 --- a/src/ModelContextProtocol.Core/ModelContextProtocol.Core.csproj +++ b/src/ModelContextProtocol.Core/ModelContextProtocol.Core.csproj @@ -34,6 +34,7 @@ + diff --git a/src/ModelContextProtocol.Core/Server/McpServerImpl.cs b/src/ModelContextProtocol.Core/Server/McpServerImpl.cs index 2ce838713..5617fdb6b 100644 --- a/src/ModelContextProtocol.Core/Server/McpServerImpl.cs +++ b/src/ModelContextProtocol.Core/Server/McpServerImpl.cs @@ -2087,148 +2087,61 @@ internal bool IsJuly2026OrLaterProtocolRequest(JsonRpcMessageContext? requestCon JsonRpcRequest request, CancellationToken cancellationToken) { - const int MaxRetries = 10; - - for (int retry = 0; ; retry++) - { - InputRequiredResult inputRequiredResult; - Exception? inputRequiredException = null; - - try - { - var result = await handler(request, cancellationToken).ConfigureAwait(false); - - // A handler can surface an input-required result two ways: by throwing InputRequiredException, - // or by RETURNING an InputRequiredResult through the alternate result path (ResultOrAlternate). - // Normalize both forms so a client that doesn't natively support MRTR gets the same server-side - // resolution either way. - if (GetReturnedInputRequiredResult(result) is not { } returnedInputRequired) - { - return result; - } - - inputRequiredResult = returnedInputRequired; - } - catch (InputRequiredException ex) - { - inputRequiredResult = ex.Result; - inputRequiredException = ex; - } - - // If the client natively supports MRTR, serialize and return directly - - // the client will drive the retry loop. - if (ClientSupportsMrtr()) - { - return SerializeInputRequiredResult(inputRequiredResult); - } - - // In stateless mode without MRTR, the server can't resolve input requests via - // JSON-RPC (no persistent session for server-to-client requests), and the client - // won't recognize the InputRequiredResult. This is the one unsupported configuration. + Func? createDirectResult = ClientSupportsMrtr() + ? static inputRequiredResult => SerializeInputRequiredResult(inputRequiredResult) + : null; + + return await InputRequiredRequestRunner.RunAsync( + request, + handler, + InputRequiredRequestRunner.GetReturnedInputRequiredResult, + createDirectResult, + PrepareRetryAsync, + static (message, innerException) => new McpException(message, innerException), + cancellationToken).ConfigureAwait(false); + + async Task PrepareRetryAsync( + JsonRpcRequest currentRequest, + InputRequiredResult inputRequiredResult, + Exception? inputRequiredException, + CancellationToken retryCancellationToken) + { + // In stateless mode without MRTR, the server cannot resolve input requests via JSON-RPC + // because there is no persistent session for server-to-client requests. if (!HasStatefulTransport()) { throw new McpException( "A tool handler returned an incomplete result, but the server is stateless and the client does not support MRTR. " + - "MRTR-native tools require either an MRTR-capable client or a stateful server for backward-compatible resolution.", inputRequiredException); - } - - // Backcompat: resolve input requests via standard JSON-RPC calls and retry the handler. - if (inputRequiredResult.InputRequests is not { Count: > 0 } inputRequests) - { - throw new McpException( - "A tool handler returned an incomplete result without input requests, and the client does not support MRTR.", inputRequiredException); + "MRTR-native tools require either an MRTR-capable client or a stateful server for backward-compatible resolution.", + inputRequiredException); } - if (retry >= MaxRetries) + IDictionary? inputResponses = null; + if (inputRequiredResult.InputRequests is { Count: > 0 } inputRequests) { - throw new McpException( - $"MRTR-native tool exceeded {MaxRetries} retry rounds without completing.", inputRequiredException); + // Route outgoing requests through the originating POST response stream, matching + // normal tool-initiated requests and avoiding a race with a separate GET stream. + var destinationServer = CreateDestinationBoundServer(currentRequest); + inputResponses = await InputRequiredRequestRunner.ResolveInputRequestsAsync( + inputRequests, + (inputRequest, requestCancellationToken) => + ResolveInputRequestAsync(destinationServer, inputRequest, requestCancellationToken), + retryCancellationToken).ConfigureAwait(false); } - // Resolve each input request by sending the corresponding JSON-RPC call to the client. - // Route the outgoing requests via the same DestinationBoundMcpServer used for normal tool - // handlers, so they go through the POST's response stream (RelatedTransport) rather than - // the session-level transport. Without this, the messages can race with the client's GET - // stream startup and be silently dropped by StreamableHttpServerTransport.SendMessageAsync - // when no GET request has arrived yet. - var destinationServer = CreateDestinationBoundServer(request); - var inputResponses = await ResolveInputRequestsAsync(destinationServer, inputRequests, cancellationToken).ConfigureAwait(false); - - // Reconstruct request params with inputResponses and requestState for the retry. - var paramsObj = request.Params?.DeepClone() as JsonObject ?? new JsonObject(); - paramsObj["inputResponses"] = JsonSerializer.SerializeToNode( - (IDictionary)inputResponses, McpJsonUtilities.JsonContext.Default.IDictionaryStringInputResponse); - - if (inputRequiredResult.RequestState is { } requestState) + return new JsonRpcRequest { - paramsObj["requestState"] = requestState; - } - else - { - // Strip any stale requestState carried over from the previous round's clone so - // the next tool invocation doesn't see a continuation token the current round is not using. - paramsObj.Remove("requestState"); - } - - request = new JsonRpcRequest - { - Id = request.Id, - Method = request.Method, - Params = paramsObj, - Context = request.Context, + Id = currentRequest.Id, + Method = currentRequest.Method, + Params = InputRequiredRequestRunner.CreateRetryParams( + currentRequest.Params, + inputResponses, + inputRequiredResult.RequestState), + Context = currentRequest.Context, }; } } - /// - /// Resolves a batch of MRTR input requests concurrently by dispatching each as a standard - /// JSON-RPC request to the client. The requests are routed via - /// so they go out through the POST's response stream (matching the behavior of tool-initiated - /// server-to-client requests like server.SampleAsync) and avoid racing with the client's - /// GET stream startup. On the first failure all remaining handlers are cancelled so user-facing - /// flows (sampling/elicitation prompts) don't keep running once the caller has given up, and - /// exceptions from late-completing tasks are observed before the original exception is rethrown. - /// - private static async Task> ResolveInputRequestsAsync( - McpServer destinationServer, - IDictionary inputRequests, - CancellationToken cancellationToken) - { - using var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken); - - var keyed = new (string Key, Task Task)[inputRequests.Count]; - int i = 0; - foreach (var kvp in inputRequests) - { - keyed[i++] = (kvp.Key, ResolveInputRequestAsync(destinationServer, kvp.Value, linkedCts.Token)); - } - - try - { - await Task.WhenAll(Array.ConvertAll(keyed, k => k.Task)).ConfigureAwait(false); - } - catch - { - linkedCts.Cancel(); - try - { - await Task.WhenAll(Array.ConvertAll(keyed, k => k.Task)).ConfigureAwait(false); - } - catch - { - // Observed; the original exception is the one we want to surface. - } - throw; - } - - var responses = new Dictionary(keyed.Length); - foreach (var (key, task) in keyed) - { - responses[key] = task.Result; - } - return responses; - } - /// /// Resolves a single MRTR by dispatching it as a standard JSON-RPC /// request to the client via . This is the server-side mirror @@ -2264,24 +2177,6 @@ private static async Task ResolveInputRequestAsync(McpServer dest private static JsonNode? SerializeInputRequiredResult(InputRequiredResult inputRequiredResult) => JsonSerializer.SerializeToNode(inputRequiredResult, McpJsonUtilities.JsonContext.Default.InputRequiredResult); - /// - /// Detects an that a handler surfaced by RETURNING it through the alternate - /// result path (rather than throwing ), so both forms can be resolved - /// identically for clients that don't natively support MRTR. Returns for any other result. - /// - private static InputRequiredResult? GetReturnedInputRequiredResult(JsonNode? result) - { - if (result is JsonObject resultObject && - resultObject.TryGetPropertyValue("resultType", out var resultTypeNode) && - resultTypeNode?.GetValueKind() == JsonValueKind.String && - resultTypeNode.GetValue() == "input_required") - { - return JsonSerializer.Deserialize(result, McpJsonUtilities.JsonContext.Default.InputRequiredResult); - } - - return null; - } - /// /// Wraps MRTR-eligible request handlers so that when a handler calls ElicitAsync/SampleAsync/RequestRootsAsync, /// an is returned early and the handler is suspended until the retry arrives. diff --git a/src/ModelContextProtocol.Extensions.Tasks/ModelContextProtocol.Extensions.Tasks.csproj b/src/ModelContextProtocol.Extensions.Tasks/ModelContextProtocol.Extensions.Tasks.csproj index cfeef00f4..4c482e7b9 100644 --- a/src/ModelContextProtocol.Extensions.Tasks/ModelContextProtocol.Extensions.Tasks.csproj +++ b/src/ModelContextProtocol.Extensions.Tasks/ModelContextProtocol.Extensions.Tasks.csproj @@ -24,6 +24,7 @@ + diff --git a/src/ModelContextProtocol.Extensions.Tasks/Server/McpTasksBuilderExtensions.cs b/src/ModelContextProtocol.Extensions.Tasks/Server/McpTasksBuilderExtensions.cs index e61466a69..5b595af16 100644 --- a/src/ModelContextProtocol.Extensions.Tasks/Server/McpTasksBuilderExtensions.cs +++ b/src/ModelContextProtocol.Extensions.Tasks/Server/McpTasksBuilderExtensions.cs @@ -213,6 +213,10 @@ private async Task ExecuteTaskAsync( } finally { + // A tool may start an outgoing input request and then fail before awaiting it. + // Cancel the task token before disposing the scope so that the request waiter + // observes cancellation and removes its store subscription. + CancelTaskExecution(taskId); await executionScope.DisposeAsync().ConfigureAwait(false); } } @@ -235,10 +239,29 @@ private async Task ExecuteTaskAsync( } finally { - if (_cancellationSources.TryRemove(taskId, out var registeredCts)) - { - registeredCts.Dispose(); - } + // Also cover races with tasks/cancel and failures during scope disposal. + CancelTaskExecution(taskId); + } + } + + private void CancelTaskExecution(string taskId) + { + if (!_cancellationSources.TryRemove(taskId, out var cts)) + { + return; + } + + try + { + cts.Cancel(); + } + catch (Exception ex) + { + _logger.LogError(ex, "Failed to cancel background task '{TaskId}' during cleanup.", taskId); + } + finally + { + cts.Dispose(); } } @@ -250,7 +273,16 @@ private async Task ExecuteToolPipelineAsync( { try { - var augmented = await next(request, taskCancellationToken).ConfigureAwait(false); + var augmented = await InputRequiredRequestRunner.RunAsync( + request, + async (currentRequest, cancellationToken) => + await next(currentRequest, cancellationToken).ConfigureAwait(false), + static result => result.IsAlternate ? result.Alternate as InputRequiredResult : null, + (Func>?)null, + PrepareRetryAsync, + static (message, innerException) => + new McpProtocolException(message, innerException, McpErrorCode.InvalidRequest), + taskCancellationToken).ConfigureAwait(false); if (augmented.IsAlternate) { @@ -266,21 +298,43 @@ private async Task ExecuteToolPipelineAsync( var resultJson = JsonSerializer.SerializeToElement(augmented.Result!, McpJsonUtilities.DefaultOptions.GetTypeInfo()); await _store.SetCompletedAsync(taskId, resultJson).ConfigureAwait(false); + + async Task> PrepareRetryAsync( + RequestContext currentRequest, + InputRequiredResult inputRequiredResult, + Exception? _, + CancellationToken retryCancellationToken) + { + IDictionary? inputResponses = null; + if (inputRequiredResult.InputRequests is { Count: > 0 } inputRequests) + { + inputResponses = await InputRequiredRequestRunner.ResolveInputRequestsAsync( + inputRequests, + (inputRequest, requestCancellationToken) => + ResolveInputRequestAsync(currentRequest.Server, inputRequest, requestCancellationToken), + retryCancellationToken).ConfigureAwait(false); + } + + currentRequest.Params.InputResponses = inputResponses; + currentRequest.Params.RequestState = inputRequiredResult.RequestState; + currentRequest.JsonRpcRequest = new JsonRpcRequest + { + Id = currentRequest.JsonRpcRequest.Id, + Method = currentRequest.JsonRpcRequest.Method, + Params = InputRequiredRequestRunner.CreateRetryParams( + currentRequest.JsonRpcRequest.Params, + inputResponses, + inputRequiredResult.RequestState), + Context = currentRequest.JsonRpcRequest.Context, + }; + + return currentRequest; + } } catch (OperationCanceledException) when (taskCancellationToken.IsCancellationRequested) { await _store.SetCancelledAsync(taskId, CancellationToken.None).ConfigureAwait(false); } - catch (InputRequiredException) - { - var error = new JsonRpcErrorDetail - { - Code = (int)McpErrorCode.InvalidRequest, - Message = "MRTR and tasks cannot be composed via [McpServerTool] yet.", - }; - var errorJson = JsonSerializer.SerializeToElement(error, McpJsonUtilities.DefaultOptions.GetTypeInfo()); - await _store.SetFailedAsync(taskId, errorJson).ConfigureAwait(false); - } catch (McpProtocolException mcpEx) { // SEP-2663 §186: protocol exceptions store as failed with JSON-RPC error shape. @@ -307,6 +361,54 @@ private async Task ExecuteToolPipelineAsync( } } +#pragma warning disable MCP9005 // Tasks still supports the deprecated Sampling and Roots MRTR input methods. + private static async Task ResolveInputRequestAsync( + McpServer server, + InputRequest inputRequest, + CancellationToken cancellationToken) + { + switch (inputRequest.Method) + { + case RequestMethods.ElicitationCreate: + _ = inputRequest.ElicitationParams + ?? throw new McpException("Failed to deserialize elicitation parameters from MRTR input request."); + break; + + case RequestMethods.SamplingCreateMessage: + _ = inputRequest.SamplingParams + ?? throw new McpException("Failed to deserialize sampling parameters from MRTR input request."); + break; + + case RequestMethods.RootsList: + _ = inputRequest.RootsParams ?? new ListRootsRequestParams(); + break; + + default: + throw new McpException($"Unsupported input request method: '{inputRequest.Method}'."); + } + + var response = await server.SendRequestAsync(new JsonRpcRequest + { + Method = inputRequest.Method, + Params = inputRequest.Params is { } paramsElement + ? JsonSerializer.SerializeToNode( + paramsElement, + McpJsonUtilities.DefaultOptions.GetTypeInfo()) + : null, + }, cancellationToken).ConfigureAwait(false); + + var result = response.Result + ?? throw new McpException($"The '{inputRequest.Method}' input request returned no result."); + + return new InputResponse + { + RawValue = JsonSerializer.SerializeToElement( + result, + McpJsonUtilities.DefaultOptions.GetTypeInfo()), + }; + } +#pragma warning restore MCP9005 + private async ValueTask HandleGetTask(JsonRpcRequest request, CancellationToken cancellationToken) { GateToJuly2026OrLaterProtocol(request, TasksProtocol.MethodTasksGet); diff --git a/tests/ModelContextProtocol.AspNetCore.Tests/MapMcpTests.Mrtr.cs b/tests/ModelContextProtocol.AspNetCore.Tests/MapMcpTests.Mrtr.cs index 03af131b4..2845310b5 100644 --- a/tests/ModelContextProtocol.AspNetCore.Tests/MapMcpTests.Mrtr.cs +++ b/tests/ModelContextProtocol.AspNetCore.Tests/MapMcpTests.Mrtr.cs @@ -766,12 +766,20 @@ public async Task Mrtr_Backcompat_AlwaysIncomplete_FailsAfterMaxRetries() } [Fact] - public async Task Mrtr_Backcompat_EmptyInputRequests_FailsWithError() + public async Task Mrtr_Backcompat_EmptyInputRequestsWithState_RetriesWithoutResponses() { Assert.SkipWhen(Stateless, "Backcompat requires stateful server for legacy JSON-RPC."); + int attempt = 0; ConfigureServer( [McpServerTool(Name = "mrtr-empty-inputs")] (RequestContext context) => { + Interlocked.Increment(ref attempt); + if (context.Params.RequestState is "empty") + { + Assert.Null(context.Params.InputResponses); + return "state-only-resolved"; + } + throw new InputRequiredException( inputRequests: new Dictionary(), requestState: "empty"); @@ -782,11 +790,35 @@ public async Task Mrtr_Backcompat_EmptyInputRequests_FailsWithError() await using var client = await ConnectLegacyAsync(); Assert.Equal("2025-11-25", client.NegotiatedProtocolVersion); + var result = await client.CallToolAsync( + "mrtr-empty-inputs", + cancellationToken: TestContext.Current.CancellationToken); + + Assert.Equal(2, attempt); + Assert.Equal( + "state-only-resolved", + Assert.IsType(Assert.Single(result.Content)).Text); + } + + [Fact] + public async Task Mrtr_Backcompat_InputRequiredWithoutRequestsOrState_FailsWithError() + { + Assert.SkipWhen(Stateless, "Backcompat requires stateful server for legacy JSON-RPC."); + ConfigureServer( + [McpServerTool(Name = "mrtr-malformed-input-required")] static string () => + throw new InputRequiredException(new InputRequiredResult())); + await using var app = Builder.Build(); + app.MapMcp(); + await app.StartAsync(TestContext.Current.CancellationToken); + await using var client = await ConnectLegacyAsync(); + var ex = await Assert.ThrowsAsync(() => - client.CallToolAsync("mrtr-empty-inputs", + client.CallToolAsync( + "mrtr-malformed-input-required", cancellationToken: TestContext.Current.CancellationToken).AsTask()); Assert.Contains("without input requests", ex.Message, StringComparison.OrdinalIgnoreCase); + Assert.Contains("request state", ex.Message, StringComparison.OrdinalIgnoreCase); Assert.Equal(McpErrorCode.InternalError, ex.ErrorCode); } diff --git a/tests/ModelContextProtocol.Tests/Server/McpTaskStoreTests.cs b/tests/ModelContextProtocol.Tests/Server/McpTaskStoreTests.cs index ebf7ca4ad..4f8d7cf9c 100644 --- a/tests/ModelContextProtocol.Tests/Server/McpTaskStoreTests.cs +++ b/tests/ModelContextProtocol.Tests/Server/McpTaskStoreTests.cs @@ -3,9 +3,11 @@ using ModelContextProtocol.Client; using ModelContextProtocol.Protocol; using ModelContextProtocol.Server; +using ModelContextProtocol.Tests.Utils; using Microsoft.Extensions.DependencyInjection; using System.Runtime.InteropServices; using System.Text.Json; +using System.Text.Json.Serialization.Metadata; using System.Threading.Channels; #pragma warning disable MCPEXP001 @@ -19,6 +21,14 @@ namespace ModelContextProtocol.Tests.Server; /// public class McpTaskStoreTests : ClientServerTestBase { + private static readonly JsonTypeInfo s_inputRequiredResultTypeInfo = + (JsonTypeInfo)McpJsonUtilities.DefaultOptions.GetTypeInfo(typeof(InputRequiredResult)); + + private readonly TrackingTaskStore _taskStore = new() + { + DefaultPollIntervalMs = 50, + }; + public McpTaskStoreTests(ITestOutputHelper testOutputHelper) : base(testOutputHelper) { #if !NET @@ -30,16 +40,140 @@ protected override void ConfigureServices(ServiceCollection services, IMcpServer { mcpServerBuilder .WithTools() - .WithTasks(new InMemoryMcpTaskStore - { - DefaultPollIntervalMs = 50, - }, options => options.ExecutionModeSelector = request => + .WithTasks(_taskStore, options => options.ExecutionModeSelector = request => request.Params?.Name switch { "sync-tool" => McpTaskExecutionMode.Synchronous, "required-tool" => McpTaskExecutionMode.Required, _ => McpTaskExecutionMode.Optional, }); + +#pragma warning disable MCPEXP002 // Exercises returned InputRequiredResult from an alternate-result filter. + services.Configure(options => + options.Filters.Request.CallToolWithAlternateFilters.Add((request, next, cancellationToken) => + { + if (request.Params?.Name is not "returned-mrtr-tool") + { + if (request.Params?.Name is "stateful-mrtr-tool") + { + if (request.Params.RequestState is null && request.Params.InputResponses is null) + { + return new ValueTask>( + ResultOrAlternate.FromAlternate( + new InputRequiredResult + { + InputRequests = new Dictionary + { + ["first"] = InputRequest.ForElicitation(new ElicitRequestParams + { + Message = "first", + RequestedSchema = new(), + }), + }, + RequestState = "round-one", + }, + s_inputRequiredResultTypeInfo)); + } + + if (request.Params.RequestState is "round-one") + { + Assert.True(request.Params.InputResponses?.ContainsKey("first")); + Assert.Equal("round-one", request.JsonRpcRequest.Params?["requestState"]?.GetValue()); + Assert.NotNull(request.JsonRpcRequest.Params?["inputResponses"]?["first"]); + throw new InputRequiredException(requestState: "round-two"); + } + + if (request.Params.RequestState is "round-two") + { + Assert.Null(request.Params.InputResponses); + Assert.Equal("round-two", request.JsonRpcRequest.Params?["requestState"]?.GetValue()); + Assert.Null(request.JsonRpcRequest.Params?["inputResponses"]); + return new ValueTask>( + ResultOrAlternate.FromAlternate( + new InputRequiredResult + { + InputRequests = new Dictionary + { + ["second"] = InputRequest.ForElicitation(new ElicitRequestParams + { + Message = "second", + RequestedSchema = new(), + }), + }, + }, + s_inputRequiredResultTypeInfo)); + } + + Assert.Null(request.Params.RequestState); + Assert.True(request.Params.InputResponses?.ContainsKey("second")); + Assert.Null(request.JsonRpcRequest.Params?["requestState"]); + Assert.NotNull(request.JsonRpcRequest.Params?["inputResponses"]?["second"]); + return new ValueTask>(new CallToolResult + { + Content = [new TextContentBlock { Text = "stateful-resolved" }], + }); + } + + if (request.Params?.Name is "malformed-mrtr-tool") + { + return new ValueTask>( + ResultOrAlternate.FromAlternate( + new InputRequiredResult(), + s_inputRequiredResultTypeInfo)); + } + + if (request.Params?.Name is "max-retry-mrtr-tool") + { + if (request.Params.Arguments?.ContainsKey("completeAtLimit") is true && + request.Params.RequestState is "10") + { + return new ValueTask>(new CallToolResult + { + Content = [new TextContentBlock { Text = "completed-at-limit" }], + }); + } + + int nextRound = int.TryParse(request.Params.RequestState, out var round) + ? round + 1 + : 1; + return new ValueTask>( + ResultOrAlternate.FromAlternate( + new InputRequiredResult { RequestState = nextRound.ToString() }, + s_inputRequiredResultTypeInfo)); + } + + return next(request, cancellationToken); + } + + if (request.Params.RequestState is "returned-state") + { + Assert.True(request.Params.InputResponses?.ContainsKey("input")); + Assert.Equal("returned-state", request.JsonRpcRequest.Params?["requestState"]?.GetValue()); + Assert.NotNull(request.JsonRpcRequest.Params?["inputResponses"]?["input"]); + + return new ValueTask>(new CallToolResult + { + Content = [new TextContentBlock { Text = "returned-resolved" }], + }); + } + + var inputRequired = new InputRequiredResult + { + InputRequests = new Dictionary + { + ["input"] = InputRequest.ForElicitation(new ElicitRequestParams + { + Message = "Confirm?", + RequestedSchema = new(), + }), + }, + RequestState = "returned-state", + }; + + return new ValueTask>( + ResultOrAlternate.FromAlternate(inputRequired, s_inputRequiredResultTypeInfo)); + })); +#pragma warning restore MCPEXP002 } [Fact] @@ -277,40 +411,542 @@ public async Task McpProtocolException_FromTool_StoresAsFailedWithJsonRpcErrorSh Assert.Equal("custom-protocol-message", failed.Error.GetProperty("message").GetString()); } + [Theory] + [InlineData(RequestMethods.ElicitationCreate, "accept")] + [InlineData(RequestMethods.SamplingCreateMessage, "sampled response")] + [InlineData(RequestMethods.RootsList, "file:///workspace")] + public async Task InputRequiredException_FromTool_UsesTaskInputRequests( + string method, + string expectedResult) + { + await using var client = await CreateMcpClientForServer(); + var ct = TestContext.Current.CancellationToken; + + var augmented = await client.CallToolAsTaskAsync( + new CallToolRequestParams + { + Name = "mrtr-tool", + Arguments = new Dictionary + { + ["method"] = JsonSerializer.SerializeToElement(method, McpJsonUtilities.DefaultOptions.GetTypeInfo()), + }, + }, ct); + + var taskId = augmented.TaskCreated!.TaskId; + + var pending = await _taskStore.WaitForTaskAsync( + taskId, + static task => task.Status is McpTaskStatus.InputRequired, + ct); + Assert.NotNull(pending.InputRequests); + var inputRequest = Assert.Single(pending.InputRequests); + Assert.Equal(method, inputRequest.Value.Method); + + var inputResponse = method switch + { + RequestMethods.ElicitationCreate => InputResponse.FromElicitResult(new ElicitResult { Action = "accept" }), + RequestMethods.SamplingCreateMessage => InputResponse.FromSamplingResult(new CreateMessageResult + { + Content = [new TextContentBlock { Text = "sampled response" }], + Model = "test-model", + }), + RequestMethods.RootsList => InputResponse.FromRootsResult(new ListRootsResult + { + Roots = [new Root { Uri = "file:///workspace" }], + }), + _ => throw new InvalidOperationException($"Unexpected method: {method}"), + }; + + await client.UpdateTaskAsync(new UpdateTaskRequestParams + { + TaskId = taskId, + InputResponses = new Dictionary { [inputRequest.Key] = inputResponse }, + }, ct); + + await _taskStore.WaitForTaskAsync( + taskId, + static task => task.Status is McpTaskStatus.Completed, + ct); + var completed = Assert.IsType(await client.GetTaskAsync(taskId, ct)); + var result = JsonSerializer.Deserialize(completed.Result, McpJsonUtilities.DefaultOptions.GetTypeInfo()); + Assert.Equal(expectedResult, Assert.IsType(result!.Content[0]).Text); + } + [Fact] - public async Task InputRequiredException_FromTool_FailsTaskWithActionableMessage() + public async Task ReturnedInputRequiredResult_FromTaskPipeline_UsesTaskInputRequests() { - // [McpServerTool] methods that throw InputRequiredException can't compose with the task-store - // wrapper today: the taskId was already returned synchronously and there's no way to surface - // InputRequiredResult retroactively. The wrapper must fail the task with a clear, actionable - // message instead of leaking the raw exception through the generic catch. await using var client = await CreateMcpClientForServer(); var ct = TestContext.Current.CancellationToken; var augmented = await client.CallToolAsTaskAsync( - new CallToolRequestParams { Name = "mrtr-tool" }, ct); + new CallToolRequestParams { Name = "returned-mrtr-tool" }, ct); + var taskId = augmented.TaskCreated!.TaskId; + + var pending = await _taskStore.WaitForTaskAsync( + taskId, + static task => task.Status is not McpTaskStatus.Working, + ct); + Assert.Equal(McpTaskStatus.InputRequired, pending.Status); + var inputRequest = Assert.Single(pending.InputRequests!); + + await client.UpdateTaskAsync(new UpdateTaskRequestParams + { + TaskId = taskId, + InputResponses = new Dictionary + { + [inputRequest.Key] = InputResponse.FromElicitResult(new ElicitResult { Action = "accept" }), + }, + }, ct); + + await _taskStore.WaitForTaskAsync( + taskId, + static task => task.Status is McpTaskStatus.Completed, + ct); + var completed = Assert.IsType(await client.GetTaskAsync(taskId, ct)); + var result = JsonSerializer.Deserialize( + completed.Result, + McpJsonUtilities.DefaultOptions.GetTypeInfo()); + Assert.Equal("returned-resolved", Assert.IsType(result!.Content[0]).Text); + } + + [Fact] + public async Task TaskPipeline_ThrownAndReturnedRoundsReplaceStateAndResponses() + { + await using var client = await CreateMcpClientForServer(new McpClientOptions + { + Handlers = new McpClientHandlers + { + ElicitationHandler = (_, _) => + new ValueTask(new ElicitResult { Action = "accept" }), + }, + }); + var ct = TestContext.Current.CancellationToken; + var augmented = await client.CallToolAsTaskAsync( + new CallToolRequestParams { Name = "stateful-mrtr-tool" }, ct); var taskId = augmented.TaskCreated!.TaskId; - GetTaskResult? taskResult = null; - for (int i = 0; i < 20; i++) + var first = await _taskStore.WaitForTaskAsync( + taskId, + static task => task.Status is not McpTaskStatus.Working, + ct); + Assert.Equal(McpTaskStatus.InputRequired, first.Status); + Assert.NotNull(first.InputRequests); + Assert.Single(first.InputRequests!); + var firstRequest = first.InputRequests!.Single(); + Assert.Equal("first", firstRequest.Value.ElicitationParams?.Message); + + await client.UpdateTaskAsync(new UpdateTaskRequestParams { - await Task.Delay(50, ct); - taskResult = await client.GetTaskAsync(taskId, ct); - if (taskResult is FailedTaskResult) + TaskId = taskId, + InputResponses = new Dictionary { - break; - } - } + [firstRequest.Key] = InputResponse.FromElicitResult(new ElicitResult { Action = "accept" }), + }, + }, ct); + + var second = await _taskStore.WaitForTaskAsync( + taskId, + static task => task.Status is McpTaskStatus.InputRequired && + task.InputRequests is { Count: 1 } requests && + requests.Values.Any(static request => request.ElicitationParams?.Message == "second"), + ct); + Assert.Single(second.InputRequests!); + var secondRequest = second.InputRequests!.Single(); + Assert.Equal("second", secondRequest.Value.ElicitationParams?.Message); + + await client.UpdateTaskAsync(new UpdateTaskRequestParams + { + TaskId = taskId, + InputResponses = new Dictionary + { + [secondRequest.Key] = InputResponse.FromElicitResult(new ElicitResult { Action = "accept" }), + }, + }, ct); + + await _taskStore.WaitForTaskAsync( + taskId, + static task => task.Status is McpTaskStatus.Completed, + ct); + var completed = Assert.IsType(await client.GetTaskAsync(taskId, ct)); + var result = JsonSerializer.Deserialize( + completed.Result, + McpJsonUtilities.DefaultOptions.GetTypeInfo()); + Assert.Equal("stateful-resolved", Assert.IsType(result!.Content[0]).Text); + } - var failed = Assert.IsType(taskResult); - Assert.Equal(JsonValueKind.Object, failed.Error.ValueKind); + [Fact] + public async Task TaskPipeline_MultipleInputRequestsSupportsPartialUpdates() + { + await using var client = await CreateMcpClientForServer(new McpClientOptions + { + Handlers = new McpClientHandlers + { + ElicitationHandler = (request, _) => + new ValueTask(new ElicitResult + { + Action = request?.Message == "first" ? "first" : "second", + }), + }, + }); + var ct = TestContext.Current.CancellationToken; + + var augmented = await client.CallToolAsTaskAsync( + new CallToolRequestParams { Name = "multi-mrtr-tool" }, ct); + var taskId = augmented.TaskCreated!.TaskId; + + var pending = await _taskStore.WaitForTaskAsync( + taskId, + static task => task.Status is McpTaskStatus.InputRequired && task.InputRequests?.Count == 2, + ct); + var firstKey = pending.InputRequests!.Single( + pair => pair.Value.ElicitationParams?.Message == "first").Key; + var secondKey = pending.InputRequests!.Single( + pair => pair.Value.ElicitationParams?.Message == "second").Key; + Assert.NotEqual(firstKey, secondKey); + + await client.UpdateTaskAsync(new UpdateTaskRequestParams + { + TaskId = taskId, + InputResponses = new Dictionary + { + [firstKey] = InputResponse.FromElicitResult(new ElicitResult { Action = "first" }), + }, + }, ct); + + var partial = await _taskStore.WaitForTaskAsync( + taskId, + static task => task.Status is McpTaskStatus.InputRequired && + task.InputRequests is { Count: 1 } requests && + requests.Values.Any(static request => request.ElicitationParams?.Message == "second"), + ct); + Assert.Single(partial.InputRequests!); + var remainingKey = partial.InputRequests!.Single().Key; + Assert.Equal(secondKey, remainingKey); + + await client.UpdateTaskAsync(new UpdateTaskRequestParams + { + TaskId = taskId, + InputResponses = new Dictionary + { + [remainingKey] = InputResponse.FromElicitResult(new ElicitResult { Action = "second" }), + }, + }, ct); + + await _taskStore.WaitForTaskAsync( + taskId, + static task => task.Status is McpTaskStatus.Completed, + ct); + var completed = Assert.IsType(await client.GetTaskAsync(taskId, ct)); + var result = JsonSerializer.Deserialize( + completed.Result, + McpJsonUtilities.DefaultOptions.GetTypeInfo()); + Assert.Equal("first|second", Assert.IsType(result!.Content[0]).Text); + await _taskStore.WaitForNoInputResponseSubscribersAsync(ct); + Assert.Equal(0, _taskStore.InputResponseSubscriberCount); + } + + [Fact] + public async Task TaskPipeline_OutOfOrderDuplicateResponsesKeepOriginalValue() + { + await using var client = await CreateMcpClientForServer(); + var ct = TestContext.Current.CancellationToken; + + var augmented = await client.CallToolAsTaskAsync( + new CallToolRequestParams { Name = "multi-mrtr-tool" }, ct); + var taskId = augmented.TaskCreated!.TaskId; + var pending = await _taskStore.WaitForTaskAsync( + taskId, + static task => task.Status is McpTaskStatus.InputRequired && task.InputRequests?.Count == 2, + ct); + var firstKey = pending.InputRequests!.Single( + pair => pair.Value.ElicitationParams?.Message == "first").Key; + var secondKey = pending.InputRequests!.Single( + pair => pair.Value.ElicitationParams?.Message == "second").Key; + + await client.UpdateTaskAsync(new UpdateTaskRequestParams + { + TaskId = taskId, + InputResponses = new Dictionary + { + [secondKey] = InputResponse.FromElicitResult(new ElicitResult { Action = "second-original" }), + }, + }, ct); + + await _taskStore.WaitForTaskAsync( + taskId, + task => task.Status is McpTaskStatus.InputRequired && + task.InputRequests is { Count: 1 } requests && + requests.ContainsKey(firstKey), + ct); + + await client.UpdateTaskAsync(new UpdateTaskRequestParams + { + TaskId = taskId, + InputResponses = new Dictionary + { + [secondKey] = InputResponse.FromElicitResult(new ElicitResult { Action = "second-duplicate" }), + [firstKey] = InputResponse.FromElicitResult(new ElicitResult { Action = "first" }), + }, + }, ct); + + await _taskStore.WaitForTaskAsync(taskId, static task => task.Status is McpTaskStatus.Completed, ct); + var completed = Assert.IsType(await client.GetTaskAsync(taskId, ct)); + var result = JsonSerializer.Deserialize( + completed.Result, + McpJsonUtilities.DefaultOptions.GetTypeInfo()); + Assert.Equal("first|second-original", Assert.IsType(result!.Content[0]).Text); + await _taskStore.WaitForInputResponseSubscriberCountAsync(0, ct); + } + + [Fact] + public async Task TaskPipeline_CompletesOnTenthRetryRound() + { + await using var client = await CreateMcpClientForServer(); + var ct = TestContext.Current.CancellationToken; + + var augmented = await client.CallToolAsTaskAsync(new CallToolRequestParams + { + Name = "max-retry-mrtr-tool", + Arguments = new Dictionary + { + ["completeAtLimit"] = JsonSerializer.SerializeToElement( + true, + McpJsonUtilities.DefaultOptions.GetTypeInfo()), + }, + }, ct); + var taskId = augmented.TaskCreated!.TaskId; + + await _taskStore.WaitForTaskAsync(taskId, static task => task.Status is McpTaskStatus.Completed, ct); + var completed = Assert.IsType(await client.GetTaskAsync(taskId, ct)); + var result = JsonSerializer.Deserialize( + completed.Result, + McpJsonUtilities.DefaultOptions.GetTypeInfo()); + Assert.Equal("completed-at-limit", Assert.IsType(result!.Content[0]).Text); + } + + [Fact] + public async Task TaskPipeline_PartialStoreFailureCancelsRegisteredSiblingWaiter() + { + _taskStore.FailSetInputRequestsOnCall = 2; + await using var client = await CreateMcpClientForServer(); + var ct = TestContext.Current.CancellationToken; + + var augmented = await client.CallToolAsTaskAsync( + new CallToolRequestParams { Name = "multi-mrtr-tool" }, ct); + var taskId = augmented.TaskCreated!.TaskId; + + await _taskStore.WaitForTaskAsync(taskId, static task => task.Status is McpTaskStatus.Completed, ct); + await _taskStore.WaitForInputResponseSubscriberCountAsync(0, ct); + var completed = Assert.IsType(await client.GetTaskAsync(taskId, ct)); + var result = JsonSerializer.Deserialize( + completed.Result, + McpJsonUtilities.DefaultOptions.GetTypeInfo()); + Assert.True(result!.IsError); + } + + [Fact] + public async Task TaskPipeline_UnexpectedFailureCancelsOutstandingInputWaiter() + { + await using var client = await CreateMcpClientForServer(); + var ct = TestContext.Current.CancellationToken; + + var augmented = await client.CallToolAsTaskAsync( + new CallToolRequestParams { Name = "unawaited-input-then-fail" }, ct); + var taskId = augmented.TaskCreated!.TaskId; + + await _taskStore.WaitForTaskAsync(taskId, static task => task.Status is McpTaskStatus.Completed, ct); + await _taskStore.WaitForInputResponseSubscriberCountAsync(0, ct); + var completed = Assert.IsType(await client.GetTaskAsync(taskId, ct)); + Assert.True(completed.Result.GetProperty("isError").GetBoolean()); + } + + [Fact] + public async Task TaskPipeline_MalformedInputRequiredResultStoresProtocolFailure() + { + await using var client = await CreateMcpClientForServer(); + var ct = TestContext.Current.CancellationToken; + + var augmented = await client.CallToolAsTaskAsync( + new CallToolRequestParams { Name = "malformed-mrtr-tool" }, ct); + var taskId = augmented.TaskCreated!.TaskId; + + await _taskStore.WaitForTaskAsync( + taskId, + static task => task.Status is McpTaskStatus.Failed, + ct); + var failed = Assert.IsType(await client.GetTaskAsync(taskId, ct)); + Assert.Equal((int)McpErrorCode.InvalidRequest, failed.Error.GetProperty("code").GetInt32()); + Assert.Contains("without input requests", failed.Error.GetProperty("message").GetString()); + } + + [Fact] + public async Task TaskPipeline_StateOnlyInputRequiredResultHonorsRetryLimit() + { + await using var client = await CreateMcpClientForServer(); + var ct = TestContext.Current.CancellationToken; + + var augmented = await client.CallToolAsTaskAsync( + new CallToolRequestParams { Name = "max-retry-mrtr-tool" }, ct); + var taskId = augmented.TaskCreated!.TaskId; + + await _taskStore.WaitForTaskAsync( + taskId, + static task => task.Status is McpTaskStatus.Failed, + ct); + var failed = Assert.IsType(await client.GetTaskAsync(taskId, ct)); Assert.Equal((int)McpErrorCode.InvalidRequest, failed.Error.GetProperty("code").GetInt32()); + Assert.Contains("exceeded 10 retry rounds", failed.Error.GetProperty("message").GetString()); + } + + [Fact] + public async Task TaskPipeline_InputResponseFailureCancelsSiblingWaiter() + { + await using var client = await CreateMcpClientForServer(); + var ct = TestContext.Current.CancellationToken; + + var augmented = await client.CallToolAsTaskAsync( + new CallToolRequestParams { Name = "multi-mrtr-tool" }, ct); + var taskId = augmented.TaskCreated!.TaskId; + var pending = await _taskStore.WaitForTaskAsync( + taskId, + static task => task.Status is McpTaskStatus.InputRequired && task.InputRequests?.Count == 2, + ct); + var firstKey = pending.InputRequests!.Single( + pair => pair.Value.ElicitationParams?.Message == "first").Key; + var secondKey = pending.InputRequests!.Single( + pair => pair.Value.ElicitationParams?.Message == "second").Key; + + await client.UpdateTaskAsync(new UpdateTaskRequestParams + { + TaskId = taskId, + InputResponses = new Dictionary + { + [firstKey] = new InputResponse { RawValue = JsonElement.Parse("null") }, + }, + }, ct); + + await _taskStore.WaitForNoInputResponseSubscribersAsync(ct); + await _taskStore.WaitForTaskAsync( + taskId, + static task => task.Status is McpTaskStatus.Completed, + ct); + var completed = Assert.IsType(await client.GetTaskAsync(taskId, ct)); + var result = JsonSerializer.Deserialize( + completed.Result, + McpJsonUtilities.DefaultOptions.GetTypeInfo()); + Assert.True(result!.IsError); + + // A response that arrives after the failed execution must not resurrect the terminal task. + await client.UpdateTaskAsync(new UpdateTaskRequestParams + { + TaskId = taskId, + InputResponses = new Dictionary + { + [secondKey] = InputResponse.FromElicitResult(new ElicitResult { Action = "late" }), + }, + }, ct); + Assert.IsType(await client.GetTaskAsync(taskId, ct)); + } + + [Fact] + public async Task InputRequiredException_FromTool_CancelledTaskDoesNotResume() + { + await using var client = await CreateMcpClientForServer(); + var ct = TestContext.Current.CancellationToken; + + var augmented = await client.CallToolAsTaskAsync( + new CallToolRequestParams + { + Name = "mrtr-tool", + Arguments = new Dictionary + { + ["method"] = JsonSerializer.SerializeToElement( + RequestMethods.ElicitationCreate, + McpJsonUtilities.DefaultOptions.GetTypeInfo()), + }, + }, ct); + + var taskId = augmented.TaskCreated!.TaskId; + var pending = await _taskStore.WaitForTaskAsync( + taskId, + static task => task.Status is McpTaskStatus.InputRequired, + ct); + var inputRequest = Assert.Single(pending.InputRequests!); + + await client.CancelTaskAsync(taskId, ct); + + // A late response must not wake the cancelled tool invocation or resurrect its task. + await client.UpdateTaskAsync(new UpdateTaskRequestParams + { + TaskId = taskId, + InputResponses = new Dictionary + { + [inputRequest.Key] = InputResponse.FromElicitResult(new ElicitResult { Action = "accept" }), + }, + }, ct); + + await _taskStore.WaitForNoInputResponseSubscribersAsync(ct); + Assert.Equal(0, _taskStore.InputResponseSubscriberCount); + Assert.IsType(await client.GetTaskAsync(taskId, ct)); + } + + [Fact] + public async Task InputRequiredException_FromTool_InputHandlerFailurePropagatesPromptly() + { + await using var client = await CreateMcpClientForServer(new McpClientOptions + { + Handlers = new McpClientHandlers + { + ElicitationHandler = (_, _) => throw new InvalidOperationException("handler-failed"), + }, + }); + var ct = TestContext.Current.CancellationToken; - var message = failed.Error.GetProperty("message").GetString(); - Assert.NotNull(message); - Assert.Contains("MRTR", message); - Assert.Contains("tasks", message); + var callTask = client.CallToolWithPollingAsync( + new CallToolRequestParams + { + Name = "mrtr-tool", + Arguments = new Dictionary + { + ["method"] = JsonSerializer.SerializeToElement( + RequestMethods.ElicitationCreate, + McpJsonUtilities.DefaultOptions.GetTypeInfo()), + }, + }, + cancellationToken: ct); + var exception = await Assert.ThrowsAsync( + () => callTask.AsTask().WaitAsync(TestConstants.DefaultTimeout, ct)); + + Assert.Equal("handler-failed", exception.Message); + await _taskStore.WaitForNoInputResponseSubscribersAsync(ct); + Assert.Equal(0, _taskStore.InputResponseSubscriberCount); + } + + [Fact] + public async Task InputRequiredException_FromTool_LegacyClientUsesSynchronousBackcompat() + { + var clientOptions = new McpClientOptions + { + ProtocolVersion = McpProtocolVersions.June2025ProtocolVersion, + Capabilities = new ClientCapabilities { Elicitation = new() }, + Handlers = new McpClientHandlers + { + ElicitationHandler = (_, _) => + new ValueTask(new ElicitResult { Action = "accept" }), + }, + }; + await using var client = await CreateMcpClientForServer(clientOptions); + + var result = await client.CallToolAsync( + "mrtr-tool", + new Dictionary + { + ["method"] = RequestMethods.ElicitationCreate, + }, + cancellationToken: TestContext.Current.CancellationToken); + + Assert.Equal("accept", Assert.IsType(result.Content[0]).Text); } [Fact] @@ -712,8 +1348,81 @@ public static string ThrowsMcpProtocol() => throw new McpProtocolException("custom-protocol-message", McpErrorCode.InvalidParams); [McpServerTool(Name = "mrtr-tool"), System.ComponentModel.Description("A tool that throws InputRequiredException (MRTR)")] - public static string MrtrTool() => - throw new InputRequiredException(requestState: "test-state"); + public static string MrtrTool(string method, RequestContext context) + { + if (context.Params.InputResponses?.TryGetValue("input", out var response) is true) + { + Assert.Equal("test-state", context.Params.RequestState); + return method switch + { + RequestMethods.ElicitationCreate => response.Deserialize(InputResponse.ElicitResultJsonTypeInfo)!.Action, + RequestMethods.SamplingCreateMessage => response.Deserialize(InputResponse.CreateMessageResultJsonTypeInfo)! + .Content.OfType().Single().Text, + RequestMethods.RootsList => response.Deserialize(InputResponse.ListRootsResultJsonTypeInfo)!.Roots.Single().Uri, + _ => throw new InvalidOperationException($"Unexpected method: {method}"), + }; + } + + var inputRequest = method switch + { + RequestMethods.ElicitationCreate => InputRequest.ForElicitation(new ElicitRequestParams + { + Message = "Confirm?", + RequestedSchema = new(), + }), + RequestMethods.SamplingCreateMessage => InputRequest.ForSampling(new CreateMessageRequestParams + { + Messages = [new SamplingMessage { Role = Role.User, Content = [new TextContentBlock { Text = "hello" }] }], + MaxTokens = 100, + }), + RequestMethods.RootsList => InputRequest.ForRootsList(new ListRootsRequestParams()), + _ => throw new InvalidOperationException($"Unexpected method: {method}"), + }; + + throw new InputRequiredException( + new Dictionary { ["input"] = inputRequest }, + requestState: "test-state"); + } + + [McpServerTool(Name = "returned-mrtr-tool")] + public static string ReturnedMrtrTool() => "unexpected"; + + [McpServerTool(Name = "stateful-mrtr-tool")] + public static string StatefulMrtrTool() => "unexpected"; + + [McpServerTool(Name = "malformed-mrtr-tool")] + public static string MalformedMrtrTool() => "unexpected"; + + [McpServerTool(Name = "max-retry-mrtr-tool")] + public static string MaxRetryMrtrTool() => "unexpected"; + + [McpServerTool(Name = "multi-mrtr-tool")] + public static string MultiMrtrTool(RequestContext context) + { + if (context.Params.InputResponses is { } responses && + responses.ContainsKey("first") && responses.ContainsKey("second")) + { + var first = responses["first"].Deserialize(InputResponse.ElicitResultJsonTypeInfo)!.Action; + var second = responses["second"].Deserialize(InputResponse.ElicitResultJsonTypeInfo)!.Action; + return $"{first}|{second}"; + } + + throw new InputRequiredException( + new Dictionary + { + ["first"] = InputRequest.ForElicitation(new ElicitRequestParams + { + Message = "first", + RequestedSchema = new(), + }), + ["second"] = InputRequest.ForElicitation(new ElicitRequestParams + { + Message = "second", + RequestedSchema = new(), + }), + }, + requestState: "multi-state"); + } [McpServerTool(Name = "elicit-tool"), System.ComponentModel.Description("A tool that elicits")] public static async Task ElicitTool(McpServer server, CancellationToken cancellationToken) @@ -817,5 +1526,144 @@ public static async Task MultiElicitTool(McpServer server, CancellationT return $"{first.Result.Action}|{second.Result.Action}"; } + + [McpServerTool(Name = "unawaited-input-then-fail"), System.ComponentModel.Description("A tool that leaves an input waiter while failing")] + public static string UnawaitedInputThenFail(McpServer server, CancellationToken cancellationToken) + { + _ = server.ElicitAsync(new ElicitRequestParams + { + Message = "orphaned input", + RequestedSchema = new(), + }, cancellationToken); + + throw new InvalidOperationException("unexpected-tool-failure"); + } + } + + private sealed class TrackingTaskStore : InMemoryMcpTaskStore, IMcpTaskStore + { + private readonly Channel _taskUpdates = Channel.CreateUnbounded(); + private readonly Channel _inputResponseSubscriberCounts = Channel.CreateUnbounded(); + private int _inputResponseSubscriberCount; + private int _setInputRequestsCallCount; + + public int InputResponseSubscriberCount => Volatile.Read(ref _inputResponseSubscriberCount); + public int FailSetInputRequestsOnCall { get; set; } + + event Action? IMcpTaskStore.InputResponseReceived + { + add + { + InputResponseReceived += value; + _inputResponseSubscriberCounts.Writer.TryWrite( + Interlocked.Increment(ref _inputResponseSubscriberCount)); + } + remove + { + InputResponseReceived -= value; + _inputResponseSubscriberCounts.Writer.TryWrite( + Interlocked.Decrement(ref _inputResponseSubscriberCount)); + } + } + + public Task WaitForNoInputResponseSubscribersAsync(CancellationToken cancellationToken) => + WaitForInputResponseSubscriberCountAsync(0, cancellationToken); + + public async Task WaitForInputResponseSubscriberCountAsync( + int expectedCount, + CancellationToken cancellationToken) + { + if (InputResponseSubscriberCount == expectedCount) + { + return; + } + + while (true) + { + var count = await _inputResponseSubscriberCounts.Reader.ReadAsync(cancellationToken).AsTask() + .WaitAsync(TestConstants.DefaultTimeout, cancellationToken); + if (count == expectedCount) + { + return; + } + } + } + + public async Task WaitForTaskAsync( + string taskId, + Func predicate, + CancellationToken cancellationToken) + { + if (await GetTaskAsync(taskId, cancellationToken) is { } current && predicate(current)) + { + return current; + } + + while (true) + { + var update = await _taskUpdates.Reader.ReadAsync(cancellationToken).AsTask() + .WaitAsync(TestConstants.DefaultTimeout, cancellationToken); + if (update.TaskId == taskId && predicate(update)) + { + return update; + } + } + } + + async Task IMcpTaskStore.SetCompletedAsync( + string taskId, + JsonElement result, + CancellationToken cancellationToken) + { + await SetCompletedAsync(taskId, result, cancellationToken); + await PublishTaskUpdateAsync(taskId, cancellationToken); + } + + async Task IMcpTaskStore.SetFailedAsync( + string taskId, + JsonElement error, + CancellationToken cancellationToken) + { + await SetFailedAsync(taskId, error, cancellationToken); + await PublishTaskUpdateAsync(taskId, cancellationToken); + } + + async Task IMcpTaskStore.SetCancelledAsync(string taskId, CancellationToken cancellationToken) + { + var cancelled = await SetCancelledAsync(taskId, cancellationToken); + await PublishTaskUpdateAsync(taskId, cancellationToken); + return cancelled; + } + + async Task IMcpTaskStore.ResolveInputRequestsAsync( + string taskId, + IDictionary inputResponses, + CancellationToken cancellationToken) + { + await ResolveInputRequestsAsync(taskId, inputResponses, cancellationToken); + await PublishTaskUpdateAsync(taskId, cancellationToken); + } + + async Task IMcpTaskStore.SetInputRequestsAsync( + string taskId, + IDictionary inputRequests, + CancellationToken cancellationToken) + { + if (Interlocked.Increment(ref _setInputRequestsCallCount) == FailSetInputRequestsOnCall) + { + throw new InvalidOperationException("Injected SetInputRequestsAsync failure."); + } + + await SetInputRequestsAsync(taskId, inputRequests, cancellationToken); + await PublishTaskUpdateAsync(taskId, cancellationToken); + } + + private async Task PublishTaskUpdateAsync(string taskId, CancellationToken cancellationToken) + { + if (await GetTaskAsync(taskId, cancellationToken) is { } task) + { + _taskUpdates.Writer.TryWrite(task); + } + } } } diff --git a/tests/ModelContextProtocol.Tests/Server/MrtrInputRequiredExceptionTests.cs b/tests/ModelContextProtocol.Tests/Server/MrtrInputRequiredExceptionTests.cs index 3d927b163..e4d3a8858 100644 --- a/tests/ModelContextProtocol.Tests/Server/MrtrInputRequiredExceptionTests.cs +++ b/tests/ModelContextProtocol.Tests/Server/MrtrInputRequiredExceptionTests.cs @@ -15,6 +15,7 @@ namespace ModelContextProtocol.Tests.Server; public class MrtrInputRequiredExceptionTests : ClientServerTestBase { private readonly ServerMessageTracker _messageTracker = new(); + private int _inputElicitCallCount; public MrtrInputRequiredExceptionTests(ITestOutputHelper testOutputHelper) : base(testOutputHelper, startServer: false) @@ -40,7 +41,74 @@ static string (McpServer server) => Name = "always-incomplete", Description = "Tool that always throws InputRequiredException" }), - ]); + McpServerTool.Create( + static string (McpServer server) => + { + throw new InputRequiredException( + inputRequests: new Dictionary + { + ["confirm"] = InputRequest.ForElicitation(new ElicitRequestParams + { + Message = "always-incomplete", + RequestedSchema = new(), + }), + }, + requestState: "should-not-work"); + }, + new McpServerToolCreateOptions + { + Name = "always-incomplete-with-input", + Description = "Tool that always requests input" + }), + McpServerTool.Create( + static string (McpServer server) => + { + throw new InputRequiredException( + inputRequests: new Dictionary + { + ["first"] = InputRequest.ForElicitation(new ElicitRequestParams + { + Message = "first", + RequestedSchema = new(), + }), + ["second"] = InputRequest.ForSampling(new CreateMessageRequestParams + { + Messages = [new SamplingMessage { Role = Role.User, Content = [new TextContentBlock { Text = "second" }] }], + MaxTokens = 1, + }), + }, + requestState: "two-inputs"); + }, + new McpServerToolCreateOptions + { + Name = "two-inputs", + Description = "Tool that requests two inputs" + }), + McpServerTool.Create( + static string (McpServer server) => + { + throw new InputRequiredException( + inputRequests: new Dictionary + { + ["blocking-first"] = InputRequest.ForSampling(new CreateMessageRequestParams + { + Messages = [new SamplingMessage { Role = Role.User, Content = [new TextContentBlock { Text = "blocking-first" }] }], + MaxTokens = 1, + }), + ["failing-second"] = InputRequest.ForElicitation(new ElicitRequestParams + { + Message = "failing-second", + RequestedSchema = new(), + }), + }, + requestState: "reversed-two-inputs"); + }, + new McpServerToolCreateOptions + { + Name = "reversed-two-inputs", + Description = "Tool that requests a blocking input before a failing input" + }), + ]); } [Fact] @@ -58,7 +126,121 @@ public async Task InputRequiredException_WithoutInputRequests_ExhaustsRetries() client.CallToolAsync("always-incomplete", cancellationToken: TestContext.Current.CancellationToken).AsTask()); - Assert.Contains("more than", exception.Message); + Assert.Contains("retry", exception.Message, StringComparison.OrdinalIgnoreCase); + Assert.Contains("10", exception.Message); + } + + [Fact] + public async Task InputRequiredException_WithInputRequests_ExhaustsRetriesWithoutExtraResolution() + { + StartServer(); + var clientOptions = new McpClientOptions + { + Capabilities = new ClientCapabilities { Elicitation = new() }, + }; + clientOptions.Handlers.ElicitationHandler = (_, _) => + { + Interlocked.Increment(ref _inputElicitCallCount); + return new ValueTask(new ElicitResult { Action = "accept" }); + }; + + await using var client = await CreateMcpClientForServer(clientOptions); + await Assert.ThrowsAsync(() => + client.CallToolAsync( + "always-incomplete-with-input", + cancellationToken: TestContext.Current.CancellationToken).AsTask()); + + Assert.Equal(10, _inputElicitCallCount); + } + + [Fact] + public async Task InputRequiredException_InputHandlerFailureCancelsSiblingHandler() + { + StartServer(); + var siblingCancelled = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var gate = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var clientOptions = new McpClientOptions + { + Capabilities = new ClientCapabilities + { + Elicitation = new(), + Sampling = new(), + }, + }; + clientOptions.Handlers.ElicitationHandler = (_, _) => + throw new InvalidOperationException("first-input-failed"); + clientOptions.Handlers.SamplingHandler = async (_, _, cancellationToken) => + { + try + { + await gate.Task.WaitAsync(cancellationToken); + return new CreateMessageResult + { + Content = [new TextContentBlock { Text = "unexpected" }], + Model = "unexpected", + }; + } + catch (OperationCanceledException) + { + siblingCancelled.TrySetResult(true); + throw; + } + }; + + await using var client = await CreateMcpClientForServer(clientOptions); + var exception = await Assert.ThrowsAsync(() => + client.CallToolAsync( + "two-inputs", + cancellationToken: TestContext.Current.CancellationToken).AsTask() + .WaitAsync(TestConstants.DefaultTimeout, TestContext.Current.CancellationToken)); + + Assert.Equal("first-input-failed", exception.Message); + await siblingCancelled.Task.WaitAsync(TestConstants.DefaultTimeout, TestContext.Current.CancellationToken); + } + + [Fact] + public async Task InputRequiredException_FailureAfterBlockingSiblingPreservesFailure() + { + StartServer(); + var siblingCancelled = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var gate = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var clientOptions = new McpClientOptions + { + Capabilities = new ClientCapabilities + { + Elicitation = new(), + Sampling = new(), + }, + }; + clientOptions.Handlers.ElicitationHandler = (_, _) => + throw new InvalidOperationException("second-input-failed"); + clientOptions.Handlers.SamplingHandler = async (_, _, cancellationToken) => + { + try + { + await gate.Task.WaitAsync(cancellationToken); + return new CreateMessageResult + { + Content = [new TextContentBlock { Text = "unexpected" }], + Model = "unexpected", + }; + } + catch (OperationCanceledException) + { + siblingCancelled.TrySetResult(true); + throw; + } + }; + + await using var client = await CreateMcpClientForServer(clientOptions); + var exception = await Assert.ThrowsAsync(() => + client.CallToolAsync( + "reversed-two-inputs", + cancellationToken: TestContext.Current.CancellationToken).AsTask() + .WaitAsync(TestConstants.DefaultTimeout, TestContext.Current.CancellationToken)); + + Assert.Equal("second-input-failed", exception.Message); + await siblingCancelled.Task.WaitAsync(TestConstants.DefaultTimeout, TestContext.Current.CancellationToken); } } diff --git a/tests/ModelContextProtocol.Tests/Server/MrtrServerBackcompatTests.cs b/tests/ModelContextProtocol.Tests/Server/MrtrServerBackcompatTests.cs index 3c5a52d92..d9df1f12a 100644 --- a/tests/ModelContextProtocol.Tests/Server/MrtrServerBackcompatTests.cs +++ b/tests/ModelContextProtocol.Tests/Server/MrtrServerBackcompatTests.cs @@ -19,6 +19,7 @@ namespace ModelContextProtocol.Tests.Server; public class MrtrServerBackcompatTests : ClientServerTestBase { private readonly List _observedRequestStates = []; + private readonly List _observedInputResponseKeys = []; private int _attempt; public MrtrServerBackcompatTests(ITestOutputHelper testOutputHelper) @@ -34,6 +35,8 @@ protected override void ConfigureServices(ServiceCollection services, IMcpServer { var attempt = Interlocked.Increment(ref _attempt); _observedRequestStates.Add(context.Params?.RequestState); + _observedInputResponseKeys.Add( + context.Params?.InputResponses?.Keys.OrderBy(static key => key).ToArray() ?? []); return attempt switch { @@ -48,20 +51,22 @@ protected override void ConfigureServices(ServiceCollection services, IMcpServer }) }, requestState: "round1"), - // Round 2: deliberately clear the state by passing requestState: null while still - // asking for another elicitation. This exercises the params clone path that - // previously preserved the stale "round1" carry-over from round 1's deep clone. - 2 => throw new InputRequiredException( + // Round 2: requestState-only. The retry must clear the prior inputResponses + // while preserving this new continuation state. + 2 => throw new InputRequiredException(requestState: "round2"), + // Round 3: request another input while clearing requestState. The retry must + // carry only the new response and must not leak round2. + 3 => throw new InputRequiredException( inputRequests: new Dictionary { - ["confirm"] = InputRequest.ForElicitation(new ElicitRequestParams + ["confirm2"] = InputRequest.ForElicitation(new ElicitRequestParams { - Message = "round2", + Message = "round3", RequestedSchema = new() }) }, requestState: null), - // Round 3 (final): report what the handler observed so the test can assert it. + // Round 4 (final): report what the handler observed so the test can assert it. _ => $"final-state:{context.Params?.RequestState ?? ""}", }; }, @@ -74,7 +79,7 @@ protected override void ConfigureServices(ServiceCollection services, IMcpServer } [Fact] - public async Task InputRequiredException_TransitioningRequestStateToNull_DoesNotLeakStaleState() + public async Task InputRequiredException_StateAndResponsesAreReplacedAcrossRounds() { StartServer(); @@ -100,13 +105,16 @@ public async Task InputRequiredException_TransitioningRequestStateToNull_DoesNot "requeststate-transition", cancellationToken: TestContext.Current.CancellationToken); - // Three attempts: round 1 (no state) → round 2 (state="round1") → round 3 (state=null after fix). - // Without the fix, the third observed state would erroneously remain "round1" because the deep-clone - // of the prior request params carried it forward when InputRequiredException.RequestState was null. - Assert.Equal(3, _observedRequestStates.Count); + // Four attempts: initial input → state-only → new input with cleared state → final result. + Assert.Equal(4, _observedRequestStates.Count); Assert.Null(_observedRequestStates[0]); Assert.Equal("round1", _observedRequestStates[1]); - Assert.Null(_observedRequestStates[2]); + Assert.Equal("round2", _observedRequestStates[2]); + Assert.Null(_observedRequestStates[3]); + Assert.Empty(_observedInputResponseKeys[0]); + Assert.Equal(["confirm"], _observedInputResponseKeys[1]); + Assert.Empty(_observedInputResponseKeys[2]); + Assert.Equal(["confirm2"], _observedInputResponseKeys[3]); var content = Assert.Single(result.Content); var text = Assert.IsType(content).Text;