Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 10 additions & 1 deletion src/ModelContextProtocol.Core/McpSessionHandler.cs
Original file line number Diff line number Diff line change
Expand Up @@ -213,6 +213,11 @@ private async Task ProcessMessagesCoreAsync(CancellationToken cancellationToken)
int inFlightCount = 1;
var allHandlersCompleted = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);

// subscriptions/listen handlers only complete when cancelled. Once the transport stops delivering messages,
// no notifications/cancelled can arrive, so their requests are linked to this source and cancelled in the
// finally below. Other requests are left to run to completion.
using var subscriptionsCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);

try
{
await foreach (var message in _transport.MessageReader.ReadAllAsync(cancellationToken).ConfigureAwait(false))
Expand Down Expand Up @@ -245,7 +250,10 @@ async Task ProcessMessageAsync()
// canceling it.
if (messageWithId is not null)
{
combinedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
combinedCts = CancellationTokenSource.CreateLinkedTokenSource(
message is JsonRpcRequest { Method: RequestMethods.SubscriptionsListen }
? subscriptionsCts.Token
: cancellationToken);
if (message is not JsonRpcRequest { Method: RequestMethods.Initialize })
{
_handlingRequests[messageWithId.Id] = combinedCts;
Expand Down Expand Up @@ -359,6 +367,7 @@ ex is OperationCanceledException &&
// Decrement our own count. If all handlers have already completed, this will signal completion.
if (Interlocked.Decrement(ref inFlightCount) != 0)
{
await subscriptionsCts.CancelAsync().ConfigureAwait(false);
await allHandlersCompleted.Task.ConfigureAwait(false);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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