Cancel background task runners on server disposal

src/ModelContextProtocol.Core/Server/DestinationBoundMcpServer.cs

@@ -5,7 +5,7 @@
 namespace ModelContextProtocol.Server;
 
 #pragma warning disable MCPEXP002
-internal sealed class DestinationBoundMcpServer(McpServerImpl server, ITransport? transport, JsonRpcMessageContext? requestContext = null) : McpServer
+internal sealed class DestinationBoundMcpServer(McpServerImpl server, ITransport? transport, JsonRpcMessageContext? requestContext = null) : McpServer, IMcpServerLifetimeFeature
 #pragma warning restore MCPEXP002
 {
     private readonly bool _isJuly2026OrLaterRequest = server.IsJuly2026OrLaterProtocolRequest(requestContext);
@@ -73,6 +73,12 @@ public override Implementation? ClientInfo
 
     public override bool IsMrtrSupported => server.IsMrtrSupported;
 
+    CancellationToken IMcpServerLifetimeFeature.BackgroundTaskCancellationToken =>
+        ((IMcpServerLifetimeFeature)server).BackgroundTaskCancellationToken;
+
+    void IMcpServerLifetimeFeature.RegisterBackgroundTask(Task backgroundTask) =>
+        ((IMcpServerLifetimeFeature)server).RegisterBackgroundTask(backgroundTask);
+
     public override ValueTask DisposeAsync() => server.DisposeAsync();
 
     public override IAsyncDisposable RegisterNotificationHandler(string method, Func<JsonRpcNotification, CancellationToken, ValueTask> handler) => server.RegisterNotificationHandler(method, handler);

src/ModelContextProtocol.Core/Server/IMcpServerLifetimeFeature.cs

@@ -0,0 +1,22 @@
+using System.ComponentModel;
+
+namespace ModelContextProtocol.Server;
+
+/// <summary>
+/// Provides server-lifetime services used by MCP extension infrastructure.
+/// </summary>
+[EditorBrowsable(EditorBrowsableState.Never)]
+public interface IMcpServerLifetimeFeature
+{
+    /// <summary>Gets the token that should cancel background work owned by this server.</summary>
+    /// <remarks>
+    /// The token is <see cref="CancellationToken.None"/> when background work intentionally outlives
+    /// the server instance, as it does for per-request servers in stateless HTTP mode.
+    /// </remarks>
+    CancellationToken BackgroundTaskCancellationToken { get; }
+
+    /// <summary>Registers background work that server disposal must await.</summary>
+    /// <param name="backgroundTask">The background work to track.</param>
+    /// <remarks>This is a no-op when background work intentionally outlives the server instance.</remarks>
+    void RegisterBackgroundTask(Task backgroundTask);
+}

src/ModelContextProtocol.Core/Server/McpServerImpl.cs

@@ -12,7 +12,7 @@ namespace ModelContextProtocol.Server;
 
 /// <inheritdoc />
 #pragma warning disable MCPEXP001, MCPEXP002
-internal sealed partial class McpServerImpl : McpServer
+internal sealed partial class McpServerImpl : McpServer, IMcpServerLifetimeFeature
 {
     internal static Implementation DefaultImplementation { get; } = new()
     {
@@ -31,6 +31,9 @@ internal sealed partial class McpServerImpl : McpServer
     private readonly string[] _initializeHandshakeProtocolVersions;
     private readonly string[] _perRequestMetadataProtocolVersions;
     private readonly SemaphoreSlim _disposeLock = new(1, 1);
+    private readonly CancellationTokenSource _serverLifetimeCts = new();
+    private readonly object _backgroundTasksLock = new();
+    private readonly ConcurrentDictionary<Task, byte> _backgroundTasks = new();
     private readonly ConcurrentDictionary<string, MrtrContinuation> _mrtrContinuations = new();
     private readonly ConcurrentDictionary<RequestId, MrtrContext> _mrtrContextsByRequestId = new();
 
@@ -56,6 +59,7 @@ internal sealed partial class McpServerImpl : McpServer
     private int _started;
 
     private bool _disposed;
+    private bool _backgroundTaskRegistrationClosed;
 
     /// <summary>Holds a boxed <see cref="LoggingLevel"/> value for the server.</summary>
     /// <remarks>
@@ -505,6 +509,38 @@ public override Task SendMessageAsync(JsonRpcMessage message, CancellationToken
     public override IAsyncDisposable RegisterNotificationHandler(string method, Func<JsonRpcNotification, CancellationToken, ValueTask> handler)
         => _sessionHandler.RegisterNotificationHandler(method, handler);
 
+    CancellationToken IMcpServerLifetimeFeature.BackgroundTaskCancellationToken =>
+        HasStatefulTransport() ? _serverLifetimeCts.Token : CancellationToken.None;
+
+    void IMcpServerLifetimeFeature.RegisterBackgroundTask(Task backgroundTask)
+    {
+        Throw.IfNull(backgroundTask);
+
+        // Stateless HTTP servers are request-scoped, while Tasks runners intentionally outlive
+        // the originating request and are governed by tasks/cancel and task-store retention.
+        if (!HasStatefulTransport())
+        {
+            return;
+        }
+
+        lock (_backgroundTasksLock)
+        {
+            if (_backgroundTaskRegistrationClosed)
+            {
+                throw new ObjectDisposedException(nameof(McpServer));
+            }
+
+            _backgroundTasks.TryAdd(backgroundTask, 0);
+        }
+
+        _ = backgroundTask.ContinueWith(
+            static (task, state) => ((ConcurrentDictionary<Task, byte>)state!).TryRemove(task, out _),
+            _backgroundTasks,
+            CancellationToken.None,
+            TaskContinuationOptions.ExecuteSynchronously,
+            TaskScheduler.Default);
+    }
+
     /// <inheritdoc/>
     public override async ValueTask DisposeAsync()
     {
@@ -516,6 +552,7 @@ public override async ValueTask DisposeAsync()
         }
 
         _disposed = true;
+        _serverLifetimeCts.Cancel();
 
         // Dispose the session handler - cancels message processing and waits for all
         // in-flight request handlers (including retries in AwaitMrtrHandlerAsync) to complete.
@@ -524,6 +561,13 @@ public override async ValueTask DisposeAsync()
         _disposables.ForEach(d => d());
         await _sessionHandler.DisposeAsync().ConfigureAwait(false);
 
+        Task[] backgroundTasks;
+        lock (_backgroundTasksLock)
+        {
+            _backgroundTaskRegistrationClosed = true;
+            backgroundTasks = [.. _backgroundTasks.Keys];
+        }
+
         // Cancel all orphaned MRTR handlers still suspended in continuations (waiting for
         // retries that will never arrive now that the session handler is disposed).
         int cancelledCount = _mrtrContinuations.Count;
@@ -545,6 +589,11 @@ public override async ValueTask DisposeAsync()
         {
             await _allMrtrHandlersCompleted.Task.ConfigureAwait(false);
         }
+
+        if (backgroundTasks.Length > 0)
+        {
+            await Task.WhenAll(backgroundTasks).ConfigureAwait(false);
+        }
     }
 
     private void ConfigureInitialize(McpServerOptions options)
@@ -2124,6 +2173,9 @@ private void WrapHandlerWithMrtr(string method)
             // is thread-safe with itself, and not disposing avoids deadlock risks from
             // calling Cancel/Dispose inside locks or Interlocked guards.
             var handlerCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+            var serverLifetimeRegistration = _serverLifetimeCts.Token.Register(
+                static state => ((CancellationTokenSource)state!).Cancel(),
+                handlerCts);
 
             // Store the MrtrContext so CreateDestinationBoundServer can pick it up and set it
             // on the per-request DestinationBoundMcpServer. This is picked up synchronously
@@ -2134,6 +2186,11 @@ private void WrapHandlerWithMrtr(string method)
             {
                 handlerTask = originalHandler(request, handlerCts.Token);
             }
+            catch
+            {
+                serverLifetimeRegistration.Dispose();
+                throw;
+            }
             finally
             {
                 _mrtrContextsByRequestId.TryRemove(request.Id, out _);
@@ -2146,7 +2203,7 @@ private void WrapHandlerWithMrtr(string method)
             // exceptions and decrements _mrtrInFlightCount when the handler completes,
             // mirroring how McpSessionHandler tracks in-flight handlers.
             Interlocked.Increment(ref _mrtrInFlightCount);
-            _ = ObserveHandlerCompletionAsync(handlerTask);
+            _ = ObserveHandlerCompletionAsync(handlerTask, serverLifetimeRegistration);
 
             return await AwaitMrtrHandlerAsync(
                 handlerTask, continuation, mrtrContext.InitialExchangeTask, cancellationToken).ConfigureAwait(false);
@@ -2205,7 +2262,9 @@ private void WrapHandlerWithMrtr(string method)
     /// double-reporting at Error) and decrements <see cref="_mrtrInFlightCount"/> when the
     /// handler completes, following the same in-flight tracking pattern as <see cref="McpSessionHandler"/>.
     /// </summary>
-    private async Task ObserveHandlerCompletionAsync(Task<JsonNode?> handlerTask)
+    private async Task ObserveHandlerCompletionAsync(
+        Task<JsonNode?> handlerTask,
+        CancellationTokenRegistration serverLifetimeRegistration)
     {
         try
         {
@@ -2225,6 +2284,8 @@ private async Task ObserveHandlerCompletionAsync(Task<JsonNode?> handlerTask)
         }
         finally
         {
+            serverLifetimeRegistration.Dispose();
+
             if (Interlocked.Decrement(ref _mrtrInFlightCount) == 0)
             {
                 _allMrtrHandlersCompleted.TrySetResult(true);

src/ModelContextProtocol.Extensions.Tasks/Server/McpTasksBuilderExtensions.cs

@@ -44,7 +44,7 @@ private sealed class McpTasksPostConfigureOptions(IMcpTaskStore store, ILoggerFa
     {
         private readonly IMcpTaskStore _store = store;
         private readonly ILogger _logger = (loggerFactory ?? NullLoggerFactory.Instance).CreateLogger<McpTasksPostConfigureOptions>();
-        private readonly ConcurrentDictionary<string, CancellationTokenSource> _cancellationSources = new(StringComparer.Ordinal);
+        private readonly ConcurrentDictionary<string, TaskCancellationState> _cancellationStates = new(StringComparer.Ordinal);
 
         public void PostConfigure(string? name, McpServerOptions options)
         {
@@ -75,11 +75,13 @@ public void PostConfigure(string? name, McpServerOptions options)
                 {
                     var taskInfo = await _store.CreateTaskAsync(cancellationToken).ConfigureAwait(false);
                     var taskId = taskInfo.TaskId;
-                    var cts = new CancellationTokenSource();
-                    _cancellationSources[taskId] = cts;
-                    var taskCancellationToken = cts.Token;
+                    var serverLifetime = request.Server as IMcpServerLifetimeFeature;
+                    var cancellationState = new TaskCancellationState(
+                        serverLifetime?.BackgroundTaskCancellationToken ?? CancellationToken.None);
+                    _cancellationStates[taskId] = cancellationState;
+                    var taskCancellationToken = cancellationState.Token;
 
-                    _ = Task.Run(async () =>
+                    var backgroundTask = Task.Run(async () =>
                     {
                         try
                         {
@@ -142,13 +144,6 @@ public void PostConfigure(string? name, McpServerOptions options)
                                     var resultJson = JsonSerializer.SerializeToElement(errorResult, McpJsonUtilities.DefaultOptions.GetTypeInfo<CallToolResult>());
                                     await _store.SetCompletedAsync(taskId, resultJson).ConfigureAwait(false);
                                 }
-                                finally
-                                {
-                                    if (_cancellationSources.TryRemove(taskId, out var registeredCts))
-                                    {
-                                        registeredCts.Dispose();
-                                    }
-                                }
                             }
                         }
                         catch (Exception outer)
@@ -169,14 +164,18 @@ public void PostConfigure(string? name, McpServerOptions options)
                             {
                                 _logger.LogError(storeEx, "Failed to record the failure of background task '{TaskId}'.", taskId);
                             }
-
-                            if (_cancellationSources.TryRemove(taskId, out var leftoverCts))
+                        }
+                        finally
+                        {
+                            if (_cancellationStates.TryRemove(taskId, out var registeredState))
                             {
-                                leftoverCts.Dispose();
+                                registeredState.UnregisterServerLifetime();
                             }
                         }
                     }, CancellationToken.None);
 
+                    serverLifetime?.RegisterBackgroundTask(backgroundTask);
+
                     return ResultOrAlternate<CallToolResult>.FromAlternate(
                         ToCreateTaskResult(taskInfo),
                         McpTasksJsonContext.Default.CreateTaskResult);
@@ -230,15 +229,38 @@ public void PostConfigure(string? name, McpServerOptions options)
 
             await _store.SetCancelledAsync(requestParams.TaskId, cancellationToken).ConfigureAwait(false);
 
-            if (_cancellationSources.TryRemove(requestParams.TaskId, out var cts))
+            if (_cancellationStates.TryGetValue(requestParams.TaskId, out var cancellationState))
             {
-                cts.Cancel();
-                cts.Dispose();
+                cancellationState.Cancel();
             }
 
             return JsonSerializer.SerializeToNode(new CancelTaskResult(), McpTasksJsonContext.Default.CancelTaskResult);
         }
 
+        private sealed class TaskCancellationState
+        {
+            private readonly CancellationTokenSource _source = new();
+            private readonly CancellationTokenRegistration _serverLifetimeRegistration;
+
+            public TaskCancellationState(CancellationToken serverLifetimeToken)
+            {
+                _serverLifetimeRegistration = serverLifetimeToken.Register(
+                    static state => ((CancellationTokenSource)state!).Cancel(),
+                    _source);
+            }
+
+            public CancellationToken Token => _source.Token;
+
+            public void Cancel() => _source.Cancel();
+
+            public void UnregisterServerLifetime()
+            {
+                // Cancellation can arrive concurrently from tasks/cancel and server disposal.
+                // Once the dictionary entry and server registration are gone, the CTS is collectible.
+                _serverLifetimeRegistration.Dispose();
+            }
+        }
+
         private static void GateToJuly2026OrLaterProtocol(JsonRpcRequest request, string method)
         {
             if (!IsJuly2026OrLaterProtocolRequest(request))