From 2df1d5187095455ad76c79414060ebc903832cb0 Mon Sep 17 00:00:00 2001 From: King Star Date: Sun, 6 Sep 2026 22:51:01 +0800 Subject: [PATCH] fix(server): close the register-after-sweep race on server disposal MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit DisposeAsync cancelled the MRTR continuations present when it iterated _mrtrContinuations. A continuation published after that iteration escaped cancellation entirely and could outlive the server. Add a server-lifetime CancellationTokenSource, cancel it at the very start of DisposeAsync — before the session handler is disposed and before the sweep — and re-check it immediately after publishing a continuation. Either the sweep sees the entry or the publisher sees the cancellation, so every interleaving cancels the handler. The issue proposed linking each handler's CTS to the server-lifetime token instead. That would leak: the per-handler source is deliberately never disposed, so a registration on a server-lifetime token is never released and accumulates for every MRTR handler invocation the process ever runs. The existing link to the per-request token is released when the request's source is disposed, which is why it does not accumulate today. Publish-then-recheck closes the same race without holding any per-request state on a process-lifetime token. --- .../Server/McpServerImpl.cs | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/src/ModelContextProtocol.Core/Server/McpServerImpl.cs b/src/ModelContextProtocol.Core/Server/McpServerImpl.cs index 2ce838713..66ee87fc3 100644 --- a/src/ModelContextProtocol.Core/Server/McpServerImpl.cs +++ b/src/ModelContextProtocol.Core/Server/McpServerImpl.cs @@ -32,6 +32,16 @@ internal sealed partial class McpServerImpl : McpServer private readonly string[] _perRequestMetadataProtocolVersions; private readonly SemaphoreSlim _disposeLock = new(1, 1); private readonly ConcurrentDictionary _mrtrContinuations = new(); + + /// + /// Cancelled at the start of , before the sweep that cancels + /// the continuations present at that moment. A continuation published after the sweep + /// observes this and cancels itself, so no MRTR handler can outlive the server. + /// Like the per-handler sources, it is intentionally never disposed: nothing registers + /// on it, and not disposing keeps usable + /// from the publish path without an ObjectDisposedException race against disposal. + /// + private readonly CancellationTokenSource _serverLifetimeCts = new(); private readonly ConcurrentDictionary _mrtrContextsByRequestId = new(); private static readonly string[] s_perRequestMetadataKeys = [ @@ -609,6 +619,10 @@ public override async ValueTask DisposeAsync() _disposed = true; + // Close the register-after-sweep race before anything else: a continuation + // published from here on cancels itself instead of escaping the sweep below. + _serverLifetimeCts.Cancel(); + // Dispose the session handler - cancels message processing and waits for all // in-flight request handlers (including retries in AwaitMrtrHandlerAsync) to complete. // After this returns, no new requests can be processed and no new MRTR continuations @@ -2451,6 +2465,14 @@ private void WrapHandlerWithMrtr(string method) continuation.PendingExchange = exchange; _mrtrContinuations[correlationId] = continuation; + // Publish-then-recheck: DisposeAsync cancels _serverLifetimeCts before iterating + // _mrtrContinuations, so either the sweep sees this entry or this check sees the + // cancellation. Both interleavings cancel the handler. + if (_serverLifetimeCts.IsCancellationRequested) + { + continuation.CancelHandler(); + } + return SerializeInputRequiredResult(inputRequiredResult); }