From 6131e9e7a4cf54d22561a7a37751636762fcb6f3 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sat, 28 Feb 2026 15:47:51 +0000 Subject: [PATCH 1/4] Initial plan From e2ee6ad935e0468ce8df6f9c0ed83d3f25a03e4f Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sat, 28 Feb 2026 16:03:54 +0000 Subject: [PATCH 2/4] Wait for in-flight message handlers before completing ProcessMessagesCoreAsync Use an in-flight counter (starting at 1) and a TaskCompletionSource to ensure ProcessMessagesCoreAsync waits for all fire-and-forget ProcessMessageAsync tasks to complete before returning. This prevents ObjectDisposedException when the transport closes while tool handlers are still executing. Fixes: ObjectDisposedException when IServiceProvider is disposed while tool handler execution is in progress. Co-authored-by: stephentoub <2642209+stephentoub@users.noreply.github.com> --- .../McpSessionHandler.cs | 18 +++++++ .../Server/McpServerTests.cs | 54 +++++++++++++++++++ 2 files changed, 72 insertions(+) diff --git a/src/ModelContextProtocol.Core/McpSessionHandler.cs b/src/ModelContextProtocol.Core/McpSessionHandler.cs index b2f94fb28..389a7cf22 100644 --- a/src/ModelContextProtocol.Core/McpSessionHandler.cs +++ b/src/ModelContextProtocol.Core/McpSessionHandler.cs @@ -173,12 +173,19 @@ public Task ProcessMessagesAsync(CancellationToken cancellationToken) private async Task ProcessMessagesCoreAsync(CancellationToken cancellationToken) { + // Track in-flight message handlers so we can wait for them to complete before returning. + // Start at 1 to represent ProcessMessagesCoreAsync itself; it's decremented after the loop exits. + int inFlightCount = 1; + var allHandlersCompleted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + try { await foreach (var message in _transport.MessageReader.ReadAllAsync(cancellationToken).ConfigureAwait(false)) { LogMessageRead(EndpointName, message.GetType().Name); + Interlocked.Increment(ref inFlightCount); + // Fire and forget the message handling to avoid blocking the transport. if (message.Context?.ExecutionContext is null) { @@ -286,6 +293,11 @@ ex is OperationCanceledException && _handlingRequests.TryRemove(messageWithId.Id, out _); combinedCts!.Dispose(); } + + if (Interlocked.Decrement(ref inFlightCount) == 0) + { + allHandlersCompleted.TrySetResult(true); + } } } } @@ -297,6 +309,12 @@ ex is OperationCanceledException && } finally { + // Decrement our own count. If all handlers have already completed, this will signal completion. + if (Interlocked.Decrement(ref inFlightCount) != 0) + { + await allHandlersCompleted.Task.ConfigureAwait(false); + } + // Fail any pending requests, as they'll never be satisfied. foreach (var entry in _pendingRequests) { diff --git a/tests/ModelContextProtocol.Tests/Server/McpServerTests.cs b/tests/ModelContextProtocol.Tests/Server/McpServerTests.cs index ea680ecf0..7a8a88621 100644 --- a/tests/ModelContextProtocol.Tests/Server/McpServerTests.cs +++ b/tests/ModelContextProtocol.Tests/Server/McpServerTests.cs @@ -1003,6 +1003,60 @@ await transport.SendClientMessageAsync(new JsonRpcNotification await runTask; } + [Fact] + public async Task RunAsync_WaitsForInFlightHandlersBeforeReturning() + { + // Arrange: Create a tool handler that blocks until we release it. + var handlerStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var releaseHandler = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + bool handlerCompleted = false; + + await using var transport = new TestServerTransport(); + var options = CreateOptions(new ServerCapabilities { Tools = new() }); + options.Handlers.CallToolHandler = async (request, ct) => + { + handlerStarted.SetResult(); + await releaseHandler.Task; + handlerCompleted = true; + return new CallToolResult { Content = [new TextContentBlock { Text = "done" }] }; + }; + options.Handlers.ListToolsHandler = (request, ct) => throw new NotImplementedException(); + + await using var server = McpServer.Create(transport, options, LoggerFactory); + var runTask = server.RunAsync(TestContext.Current.CancellationToken); + + // Send a tool call request. + await transport.SendClientMessageAsync( + new JsonRpcRequest + { + Method = RequestMethods.ToolsCall, + Id = new RequestId(1) + }, + TestContext.Current.CancellationToken); + + // Wait for the handler to start executing. + await handlerStarted.Task.WaitAsync(TestConstants.DefaultTimeout, TestContext.Current.CancellationToken); + + // Dispose the transport to simulate client disconnect while the handler is still running. + await transport.DisposeAsync(); + + // Release the handler after a delay, giving ProcessMessagesCoreAsync time to notice the + // channel closed. Without the fix, RunAsync would return before the handler completes. + var ct = TestContext.Current.CancellationToken; + _ = Task.Run(async () => + { + await Task.Delay(500, ct); + releaseHandler.SetResult(); + }, ct); + + // Wait for RunAsync to complete. + await runTask.WaitAsync(TestConstants.DefaultTimeout, TestContext.Current.CancellationToken); + + // With the fix, RunAsync waits for in-flight handlers. Without it, it returns immediately + // after the transport closes (before the 500ms delay releases the handler). + Assert.True(handlerCompleted, "RunAsync should wait for in-flight handlers to complete before returning."); + } + private static async Task InitializeServerAsync(TestServerTransport transport, ClientCapabilities capabilities, CancellationToken cancellationToken = default) { var initializeRequest = new JsonRpcRequest From 9d1fb9f2da5e91ed7cf2970a06ddf1865e65c433 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sat, 28 Feb 2026 16:08:34 +0000 Subject: [PATCH 3/4] Reduce test delay from 500ms to 200ms per code review feedback Co-authored-by: stephentoub <2642209+stephentoub@users.noreply.github.com> --- tests/ModelContextProtocol.Tests/Server/McpServerTests.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/ModelContextProtocol.Tests/Server/McpServerTests.cs b/tests/ModelContextProtocol.Tests/Server/McpServerTests.cs index 7a8a88621..eec967e5f 100644 --- a/tests/ModelContextProtocol.Tests/Server/McpServerTests.cs +++ b/tests/ModelContextProtocol.Tests/Server/McpServerTests.cs @@ -1045,7 +1045,7 @@ await transport.SendClientMessageAsync( var ct = TestContext.Current.CancellationToken; _ = Task.Run(async () => { - await Task.Delay(500, ct); + await Task.Delay(200, ct); releaseHandler.SetResult(); }, ct); From 0a55c5cb3d6b4380f295c35d6620e1faae37c5f8 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 1 Mar 2026 02:15:32 +0000 Subject: [PATCH 4/4] Fix net472 build: use TaskCompletionSource instead of non-generic TaskCompletionSource in test The non-generic TaskCompletionSource doesn't exist in .NET Framework 4.7.2 / .NET Standard 2.0. Use TaskCompletionSource for compatibility. Co-authored-by: stephentoub <2642209+stephentoub@users.noreply.github.com> --- tests/ModelContextProtocol.Tests/Server/McpServerTests.cs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/ModelContextProtocol.Tests/Server/McpServerTests.cs b/tests/ModelContextProtocol.Tests/Server/McpServerTests.cs index eec967e5f..045cbe435 100644 --- a/tests/ModelContextProtocol.Tests/Server/McpServerTests.cs +++ b/tests/ModelContextProtocol.Tests/Server/McpServerTests.cs @@ -1007,15 +1007,15 @@ await transport.SendClientMessageAsync(new JsonRpcNotification public async Task RunAsync_WaitsForInFlightHandlersBeforeReturning() { // Arrange: Create a tool handler that blocks until we release it. - var handlerStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); - var releaseHandler = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var handlerStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var releaseHandler = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); bool handlerCompleted = false; await using var transport = new TestServerTransport(); var options = CreateOptions(new ServerCapabilities { Tools = new() }); options.Handlers.CallToolHandler = async (request, ct) => { - handlerStarted.SetResult(); + handlerStarted.SetResult(true); await releaseHandler.Task; handlerCompleted = true; return new CallToolResult { Content = [new TextContentBlock { Text = "done" }] }; @@ -1046,7 +1046,7 @@ await transport.SendClientMessageAsync( _ = Task.Run(async () => { await Task.Delay(200, ct); - releaseHandler.SetResult(); + releaseHandler.SetResult(true); }, ct); // Wait for RunAsync to complete.