Wait for in-flight message handlers before ProcessMessagesCoreAsync returns

[Copilot]

ProcessMessagesCoreAsync fire-and-forgets ProcessMessageAsync tasks but returns as soon as the transport channel completes — without waiting for in-flight handlers. The caller (RunAsync) then disposes service scopes, causing ObjectDisposedException in still-running tool handlers. This is easily triggered by OpenAI's stateless HTTP transport closing the connection while a long-running tool is executing.

Changes

• McpSessionHandler.ProcessMessagesCoreAsync: Track in-flight handlers with an Interlocked counter (starts at 1 for the loop itself) and a TaskCompletionSource<bool>. Increment before each ProcessMessageAsync dispatch, decrement in its finally. When the count hits zero, signal the TCS. The finally block of ProcessMessagesCoreAsync decrements its own count and awaits the TCS if handlers are still running.

• McpServerTests.RunAsync_WaitsForInFlightHandlersBeforeReturning: Regression test that blocks a tool handler, disposes the transport, then verifies RunAsync doesn't return until the handler completes.

Uses TaskCompletionSource<bool> instead of non-generic TaskCompletionSource for .NET Standard 2.0 compatibility.

Original prompt

---

This section details on the original issue you should resolve

<issue_title>ObjectDisposedException - IServiceProvider disposed while tool handler execution.</issue_title>
<issue_description>Describe the bug

Sometimes I get ObjectDisposedException from DI services like the IServiceProvider itself or DbContext
inside the tool call handler.

It looks like OpenAI closes the connection for too-long-running tools,
and instead of canceling the CancellationToken, the service scope gets disposed.
Or the scope gets disposed of without waiting for the handler to cancel cleanly.

To Reproduce
Steps to reproduce the behavior:

1. Setup: Stateless, Http Transport
2. Call a long-running tool
3. Close the connection to simulate a read-timeout.

[!NOTE]
I use it with OpenAI
OpenAI requires statelessness and closes the connection after some timeout.

services.AddMcpServer(options => options.ScopeRequests = true)
    .WithHttpTransport(options => {
        options.Stateless = true;
    })
    .AddAuthorizationFilters()
    .WithToolsFromAssemblyFixed(toolAssembly)
    .AddCallToolFilter(next => (context, token) => ToolCallErrorHandler(next, context, token));

// ...

app.MapMcp("/mcp").RequireAuthorization("mcp");

static async ValueTask<CallToolResult> ToolCallErrorHandler(McpRequestHandler<CallToolRequestParams, CallToolResult> next, RequestContext<CallToolRequestParams> context, CancellationToken cancellationToken)
{
    try {
        return await next(context, cancellationToken);
    }
    catch (Exception ex) {
        if (x is not OperationCanceledException) {
            // Error here inside `GetRequiredService` - the `ex` error is also an ObjectDisposedException with origin in DbContext.SaveChanges()
            context.Services!.GetRequiredService<ILogger<McpFeatureAttribute>>()
                  .LogCritical(ex, "MCP Error {tool}: {message}", context.Params?.Name, ex.Message);
        }

        return new CallToolResult {
            Content = [new TextContentBlock { Text = $"Error: {ex.Message}" }],
            IsError = true
        };
    }
}

Expected behavior

I expect the handler to either complete or cancel through CancellationToken before the service scope is disposed.

Logs

ERROR  ModelContextProtocol.Server.McpServer  "{ToolName}" threw an unhandled exception.
System.ObjectDisposedException: Cannot access a disposed object.
Object name: 'IServiceProvider'.
   at Microsoft.Extensions.DependencyInjection.ServiceLookup.ThrowHelper.ThrowObjectDisposedException()
   at ModelContextProtocol.Server.RequestServiceProvider`1.GetService(Type serviceType)
   at Microsoft.Extensions.DependencyInjection.ServiceProviderServiceExtensions.GetRequiredService[T](IServiceProvider provider)
   at Nomos.Web.Features.Mcp.Abstructions.McpFeatureInstaller.ToolCallErrorHandler(McpRequestHandler`2 next, RequestContext`1 context, CancellationToken cancellationToken) in /src/Nomos.Web/Features/Mcp/Abstructions/McpFeatureInstaller.cs:line 42
   at ModelContextProtocol.AspNetCore.AuthorizationFilterSetup.<>c__DisplayClass7_0.<<ConfigureCallToolFilter>b__1>d.MoveNext()
--- End of stack trace from previous location ---
   at ModelContextProtocol.Server.McpServerImpl.<>c__DisplayClass51_2.<<ConfigureTools>b__5>d.MoveNext()

Additional context
Add any other context about the problem here.
</issue_description>

<agent_instructions>Both #1401 and #1400 are trying to fix this issue. I'd like you to create a PR to fix it as well, in a similar manner, but you shouldn't need to actually track the task instances / maintain a list. It should be enough to just have a counter that's incremented each time and just before ProcessMessageAsync is called, and then ProcessMessageAsync decrements it at the end of its finally block. When the count drops to zero, ProcessMessagesCoreAsync can be allowed to continue / complete. The count can start at 1 and then be Interlocked.Decremented just before it waits on a TaskCompletionSource Task that whoever Interlocked.Decrements the count to 0 will set. Start by writing one or more tests that will fail before the fix and will pass after.</agent_instructions>

Comments on the Issue (you are @copilot in this section)

• Fixes ObjectDisposedException - IServiceProvider disposed while tool handler execution. #1269

---

🔒 GitHub Advanced Security automatically protects Copilot coding agent pull requests. You can protect all pull requests by enabling Advanced Security for your repositories. Learn more about Advanced Security.