diff --git a/src/ModelContextProtocol.Core/McpSessionHandler.cs b/src/ModelContextProtocol.Core/McpSessionHandler.cs index d2b074547..1c58195d8 100644 --- a/src/ModelContextProtocol.Core/McpSessionHandler.cs +++ b/src/ModelContextProtocol.Core/McpSessionHandler.cs @@ -94,6 +94,11 @@ internal static bool SupportsNaturalOutputSchemas(string? protocolVersion) /// that can be used to request cancellation of the in-flight handler. /// private readonly ConcurrentDictionary _handlingRequests = new(); + /// + /// The subscriptions/listen requests in . Their handlers only complete when + /// cancelled, so they're cancelled once the transport stops delivering messages. + /// + private readonly ConcurrentDictionary _handlingSubscriptionRequests = new(); private readonly ILogger _logger; // This _sessionId is solely used to identify the session in telemetry and logs. @@ -250,6 +255,11 @@ async Task ProcessMessageAsync() { _handlingRequests[messageWithId.Id] = combinedCts; } + + if (message is JsonRpcRequest { Method: RequestMethods.SubscriptionsListen }) + { + _handlingSubscriptionRequests[messageWithId.Id] = combinedCts; + } } // If we await the handler without yielding first, the transport may not be able to read more messages, @@ -338,6 +348,7 @@ ex is OperationCanceledException && if (messageWithId is not null) { _handlingRequests.TryRemove(messageWithId.Id, out _); + _handlingSubscriptionRequests.TryRemove(messageWithId.Id, out _); combinedCts!.Dispose(); } @@ -359,6 +370,21 @@ ex is OperationCanceledException && // Decrement our own count. If all handlers have already completed, this will signal completion. if (Interlocked.Decrement(ref inFlightCount) != 0) { + // No more messages will arrive, so no notifications/cancelled can end an open subscriptions/listen + // request. Its handler only completes when cancelled, so cancel it here or this wait never ends. + // Other handlers are left to run to completion. + foreach (var entry in _handlingSubscriptionRequests) + { + try + { + await entry.Value.CancelAsync().ConfigureAwait(false); + } + catch (ObjectDisposedException) + { + // The handler completed and disposed its token source concurrently. + } + } + await allHandlersCompleted.Task.ConfigureAwait(false); } diff --git a/tests/ModelContextProtocol.Tests/Server/RawStreamConformanceTests.cs b/tests/ModelContextProtocol.Tests/Server/RawStreamConformanceTests.cs index a41adfc22..a17984033 100644 --- a/tests/ModelContextProtocol.Tests/Server/RawStreamConformanceTests.cs +++ b/tests/ModelContextProtocol.Tests/Server/RawStreamConformanceTests.cs @@ -210,5 +210,63 @@ await SendAsync( var call = await ReadAsync(); Assert.Equal("echo:still-modern", call["result"]!["content"]![0]!["text"]!.GetValue()); } + + [Fact] + public async Task InputEndOfStream_WithOpenSubscriptionsListen_CompletesRunAsync() + { + // Closing the server's input is the stdio shutdown signal. An open subscriptions/listen request, + // whose handler only completes when cancelled, must not keep RunAsync alive after that (#1914). + await SendAsync(@"{""jsonrpc"":""2.0"",""id"":1,""method"":""server/discover"",""params"":{" + July2026ProtocolMetaFragment() + "}}"); + Assert.NotNull((await ReadAsync())["result"]); + + await SendAsync( + @"{""jsonrpc"":""2.0"",""id"":""listen:0"",""method"":""subscriptions/listen"",""params"":{""notifications"":{""toolsListChanged"":true}," + + July2026ProtocolMetaFragment() + "}}"); + var acknowledgement = await ReadAsync(); + Assert.Equal(NotificationMethods.SubscriptionsAcknowledgedNotification, acknowledgement["method"]!.GetValue()); + + _clientToServer.Writer.Complete(); + + await _serverTask.WaitAsync(TestConstants.DefaultTimeout, TestContext.Current.CancellationToken); + Assert.False(_cts.IsCancellationRequested); + } + + [Fact] + public async Task InputEndOfStream_WithOpenCustomSubscriptionsListen_CompletesRunAsync() + { + // Same as above, but the stream is owned by a custom SubscriptionsListenHandler that waits for its token. + var clientToServer = new Pipe(); + var serverToClient = new Pipe(); + var handlerWaiting = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + + await using var transport = new StreamServerTransport(clientToServer.Reader.AsStream(), serverToClient.Writer.AsStream(), loggerFactory: LoggerFactory); + await using var server = McpServer.Create(transport, new McpServerOptions + { + Handlers = new() + { + SubscriptionsListenHandler = async (request, cancellationToken) => + { + handlerWaiting.TrySetResult(true); + var cancelled = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + using var registration = cancellationToken.Register( + static state => ((TaskCompletionSource)state!).TrySetResult(true), cancelled); + await cancelled.Task; + return new EmptyResult(); + }, + }, + }, LoggerFactory); + + var runTask = server.RunAsync(TestContext.Current.CancellationToken); + + using var writer = new StreamWriter(clientToServer.Writer.AsStream(), new UTF8Encoding(false)) { AutoFlush = true, NewLine = "\n" }; + await writer.WriteLineAsync( + @"{""jsonrpc"":""2.0"",""id"":""listen:0"",""method"":""subscriptions/listen"",""params"":{""notifications"":{""toolsListChanged"":true}," + + July2026ProtocolMetaFragment() + "}}"); + await handlerWaiting.Task.WaitAsync(TestConstants.DefaultTimeout, TestContext.Current.CancellationToken); + + clientToServer.Writer.Complete(); + + await runTask.WaitAsync(TestConstants.DefaultTimeout, TestContext.Current.CancellationToken); + } } #endif