Cancel open subscriptions/listen requests when the transport stops delivering messages

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.
     /// </summary>
     private readonly ConcurrentDictionary<RequestId, CancellationTokenSource> _handlingRequests = new();
+    /// <summary>
+    /// The <c>subscriptions/listen</c> requests in <see cref="_handlingRequests"/>. Their handlers only complete when
+    /// cancelled, so they're cancelled once the transport stops delivering messages.
+    /// </summary>
+    private readonly ConcurrentDictionary<RequestId, CancellationTokenSource> _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);
             }
 

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<string>());
     }
+
+    [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<string>());
+
+        _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<bool>(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<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+                    using var registration = cancellationToken.Register(
+                        static state => ((TaskCompletionSource<bool>)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