refactor(server): generalize lifetime registrations

Signed-off-by: King Star <mcxin.y@gmail.com>

Author: jstar0

src/ModelContextProtocol.Core/Server/DestinationBoundMcpServer.cs

@@ -73,11 +73,11 @@ public override Implementation? ClientInfo
 
     public override bool IsMrtrSupported => server.IsMrtrSupported;
 
-    CancellationToken IMcpServerLifetimeFeature.BackgroundTaskCancellationToken =>
-        ((IMcpServerLifetimeFeature)server).BackgroundTaskCancellationToken;
+    CancellationToken IMcpServerLifetimeFeature.ServerCancellationToken =>
+        ((IMcpServerLifetimeFeature)server).ServerCancellationToken;
 
-    void IMcpServerLifetimeFeature.RegisterBackgroundTask(Task backgroundTask) =>
-        ((IMcpServerLifetimeFeature)server).RegisterBackgroundTask(backgroundTask);
+    IDisposable IMcpServerLifetimeFeature.RegisterForDisposeAsync(IAsyncDisposable disposable) =>
+        ((IMcpServerLifetimeFeature)server).RegisterForDisposeAsync(disposable);
 
     public override ValueTask DisposeAsync() => server.DisposeAsync();
 

src/ModelContextProtocol.Core/Server/IMcpServerLifetimeFeature.cs

@@ -8,15 +8,19 @@ namespace ModelContextProtocol.Server;
 [EditorBrowsable(EditorBrowsableState.Never)]
 public interface IMcpServerLifetimeFeature
 {
-    /// <summary>Gets the token that should cancel background work owned by this server.</summary>
+    /// <summary>Gets the token that is cancelled when this server starts disposing.</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.
+    /// The token is <see cref="CancellationToken.None"/> when work intentionally outlives the server,
+    /// as it does for per-request servers in stateless HTTP mode.
     /// </remarks>
-    CancellationToken BackgroundTaskCancellationToken { get; }
+    CancellationToken ServerCancellationToken { 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);
+    /// <summary>Registers an asynchronously disposable resource that server disposal must await.</summary>
+    /// <param name="disposable">The resource to dispose when this server is disposed.</param>
+    /// <returns>A handle that unregisters the resource without disposing it.</returns>
+    /// <remarks>
+    /// Dispose the returned handle when the resource completes independently so the server does not
+    /// retain it until shutdown. Registration is a no-op when the server does not own the resource.
+    /// </remarks>
+    IDisposable RegisterForDisposeAsync(IAsyncDisposable disposable);
 }

src/ModelContextProtocol.Core/Server/McpServerImpl.cs

@@ -32,8 +32,8 @@ internal sealed partial class McpServerImpl : McpServer, IMcpServerLifetimeFeatu
     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 object _serverLifetimeRegistrationsLock = new();
+    private readonly HashSet<ServerLifetimeRegistration> _serverLifetimeRegistrations = [];
     private readonly ConcurrentDictionary<string, MrtrContinuation> _mrtrContinuations = new();
     private readonly ConcurrentDictionary<RequestId, MrtrContext> _mrtrContextsByRequestId = new();
     private static readonly string[] s_perRequestMetadataKeys =
@@ -58,7 +58,7 @@ internal sealed partial class McpServerImpl : McpServer, IMcpServerLifetimeFeatu
     private int _started;
 
     private bool _disposed;
-    private bool _backgroundTaskRegistrationClosed;
+    private bool _serverLifetimeRegistrationClosed;
 
     /// <summary>Holds a boxed <see cref="LoggingLevel"/> value for the server.</summary>
     /// <remarks>
@@ -508,36 +508,40 @@ 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 =>
+    CancellationToken IMcpServerLifetimeFeature.ServerCancellationToken =>
         HasStatefulTransport() ? _serverLifetimeCts.Token : CancellationToken.None;
 
-    void IMcpServerLifetimeFeature.RegisterBackgroundTask(Task backgroundTask)
+    IDisposable IMcpServerLifetimeFeature.RegisterForDisposeAsync(IAsyncDisposable disposable)
     {
-        Throw.IfNull(backgroundTask);
+        Throw.IfNull(disposable);
 
         // 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;
+            return NoopRegistration.Instance;
         }
 
-        lock (_backgroundTasksLock)
+        var registration = new ServerLifetimeRegistration(this, disposable);
+        lock (_serverLifetimeRegistrationsLock)
         {
-            if (_backgroundTaskRegistrationClosed)
+            if (_serverLifetimeRegistrationClosed)
             {
                 throw new ObjectDisposedException(nameof(McpServer));
             }
 
-            _backgroundTasks.TryAdd(backgroundTask, 0);
+            _serverLifetimeRegistrations.Add(registration);
         }
 
-        _ = backgroundTask.ContinueWith(
-            static (task, state) => ((ConcurrentDictionary<Task, byte>)state!).TryRemove(task, out _),
-            _backgroundTasks,
-            CancellationToken.None,
-            TaskContinuationOptions.ExecuteSynchronously,
-            TaskScheduler.Default);
+        return registration;
+    }
+
+    private void UnregisterServerLifetime(ServerLifetimeRegistration registration)
+    {
+        lock (_serverLifetimeRegistrationsLock)
+        {
+            _serverLifetimeRegistrations.Remove(registration);
+        }
     }
 
     /// <inheritdoc/>
@@ -560,11 +564,11 @@ public override async ValueTask DisposeAsync()
         _disposables.ForEach(d => d());
         await _sessionHandler.DisposeAsync().ConfigureAwait(false);
 
-        Task[] backgroundTasks;
-        lock (_backgroundTasksLock)
+        ServerLifetimeRegistration[] serverLifetimeRegistrations;
+        lock (_serverLifetimeRegistrationsLock)
         {
-            _backgroundTaskRegistrationClosed = true;
-            backgroundTasks = [.. _backgroundTasks.Keys];
+            _serverLifetimeRegistrationClosed = true;
+            serverLifetimeRegistrations = [.. _serverLifetimeRegistrations];
         }
 
         // Cancel all orphaned MRTR handlers still suspended in continuations (waiting for
@@ -589,9 +593,34 @@ public override async ValueTask DisposeAsync()
             await _allMrtrHandlersCompleted.Task.ConfigureAwait(false);
         }
 
-        if (backgroundTasks.Length > 0)
+        if (serverLifetimeRegistrations.Length > 0)
+        {
+            await Task.WhenAll(
+                serverLifetimeRegistrations.Select(static registration => registration.DisposeResourceAsync().AsTask())
+            ).ConfigureAwait(false);
+        }
+    }
+
+    private sealed class ServerLifetimeRegistration(
+        McpServerImpl server,
+        IAsyncDisposable resource) : IDisposable
+    {
+        private McpServerImpl? _server = server;
+
+        public ValueTask DisposeResourceAsync() => resource.DisposeAsync();
+
+        public void Dispose()
+        {
+            Interlocked.Exchange(ref _server, null)?.UnregisterServerLifetime(this);
+        }
+    }
+
+    private sealed class NoopRegistration : IDisposable
+    {
+        public static NoopRegistration Instance { get; } = new();
+
+        public void Dispose()
         {
-            await Task.WhenAll(backgroundTasks).ConfigureAwait(false);
         }
     }
 

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

@@ -148,14 +148,25 @@ private async ValueTask<ResultOrAlternate<CallToolResult>> RunAsTaskAsync(
             executionRequest.Server = request.Server.WithMcpTaskOutgoingRequestInterceptor(taskId, _store);
             var serverLifetime = request.Server as IMcpServerLifetimeFeature;
             var cancellationState = new TaskCancellationState(
-                serverLifetime?.BackgroundTaskCancellationToken ?? CancellationToken.None);
+                serverLifetime?.ServerCancellationToken ?? CancellationToken.None);
             _cancellationStates[taskId] = cancellationState;
 
             var taskCancellationToken = cancellationState.Token;
             var backgroundTask = Task.Run(
                 () => ExecuteTaskAsync(next, executionRequest, taskId, taskCancellationToken, executionScope),
                 CancellationToken.None);
-            serverLifetime?.RegisterBackgroundTask(backgroundTask);
+            cancellationState.SetBackgroundTask(backgroundTask);
+            try
+            {
+                cancellationState.SetServerLifetimeRegistration(
+                    serverLifetime?.RegisterForDisposeAsync(cancellationState));
+            }
+            catch
+            {
+                cancellationState.Cancel();
+                await backgroundTask.ConfigureAwait(false);
+                throw;
+            }
 
             return ResultOrAlternate<CallToolResult>.FromAlternate(
                 ToCreateTaskResult(taskInfo),
@@ -326,27 +337,63 @@ private async Task ExecuteToolPipelineAsync(
             return JsonSerializer.SerializeToNode(new CancelTaskResult(), McpTasksJsonContext.Default.CancelTaskResult);
         }
 
-        private sealed class TaskCancellationState
+        private sealed class TaskCancellationState : IAsyncDisposable
         {
             private readonly CancellationTokenSource _source = new();
             private readonly CancellationTokenRegistration _serverLifetimeRegistration;
+            private Task? _backgroundTask;
+            private IDisposable? _serverLifetimeUnregistration;
+            private int _completed;
 
             public TaskCancellationState(CancellationToken serverLifetimeToken)
             {
                 _serverLifetimeRegistration = serverLifetimeToken.Register(
-                    static state => ((CancellationTokenSource)state!).Cancel(),
-                    _source);
+                    static state => ((TaskCancellationState)state!).Cancel(),
+                    this);
             }
 
             public CancellationToken Token => _source.Token;
 
             public void Cancel() => _source.Cancel();
 
+            public void SetBackgroundTask(Task backgroundTask) => _backgroundTask = backgroundTask;
+
+            public void SetServerLifetimeRegistration(IDisposable? registration)
+            {
+                if (registration is null)
+                {
+                    return;
+                }
+
+                if (Volatile.Read(ref _completed) != 0)
+                {
+                    registration.Dispose();
+                    return;
+                }
+
+                Interlocked.CompareExchange(ref _serverLifetimeUnregistration, registration, null);
+                if (Volatile.Read(ref _completed) != 0)
+                {
+                    Interlocked.Exchange(ref _serverLifetimeUnregistration, null)?.Dispose();
+                }
+            }
+
             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.
+                Interlocked.Exchange(ref _completed, 1);
                 _serverLifetimeRegistration.Dispose();
+                Interlocked.Exchange(ref _serverLifetimeUnregistration, null)?.Dispose();
+            }
+
+            public async ValueTask DisposeAsync()
+            {
+                Cancel();
+                if (_backgroundTask is { } backgroundTask)
+                {
+                    await backgroundTask.ConfigureAwait(false);
+                }
             }
         }
 

tests/ModelContextProtocol.Tests/Server/TaskCancellationIntegrationTests.cs

@@ -179,6 +179,45 @@ protected override void ConfigureServices(ServiceCollection services, IMcpServer
             })]);
     }
 
+    [Fact]
+    public async Task DisposeAsync_DisposesAndWaitsForRegisteredLifetimeResource()
+    {
+        await using var client = await CreateMcpClientForServer();
+        var ct = TestContext.Current.CancellationToken;
+        var serverLifetime = Assert.IsAssignableFrom<IMcpServerLifetimeFeature>(Server);
+        var resource = new BlockingAsyncDisposable();
+        using var registration = serverLifetime.RegisterForDisposeAsync(resource);
+
+        Task disposeTask = Server.DisposeAsync().AsTask();
+
+        try
+        {
+            await resource.DisposeStarted.WaitAsync(TestConstants.DefaultTimeout, ct);
+            Assert.False(disposeTask.IsCompleted, "DisposeAsync should await the registered resource.");
+
+            resource.Release();
+            await disposeTask.WaitAsync(TestConstants.DefaultTimeout, ct);
+        }
+        finally
+        {
+            resource.Release();
+        }
+    }
+
+    [Fact]
+    public async Task LifetimeRegistration_DisposeUnregistersResource()
+    {
+        await using var client = await CreateMcpClientForServer();
+        var serverLifetime = Assert.IsAssignableFrom<IMcpServerLifetimeFeature>(Server);
+        var resource = new RecordingAsyncDisposable();
+        using var registration = serverLifetime.RegisterForDisposeAsync(resource);
+
+        registration.Dispose();
+        await Server.DisposeAsync();
+
+        Assert.False(resource.IsDisposed);
+    }
+
     [Fact]
     public async Task DisposeAsync_CancelsAndWaitsForTaskStoreRunner()
     {
@@ -230,7 +269,7 @@ public async Task DisposeAsync_CancelsAndWaitsForRunnerRegisteredDuringDisposal(
 
         var serverLifetime = Assert.IsAssignableFrom<IMcpServerLifetimeFeature>(Server);
         var serverCancellationFired = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
-        using var registration = serverLifetime.BackgroundTaskCancellationToken.Register(
+        using var registration = serverLifetime.ServerCancellationToken.Register(
             static state => ((TaskCompletionSource<bool>)state!).TrySetResult(true), serverCancellationFired);
 
         Task disposeTask = Server.DisposeAsync().AsTask();
@@ -283,6 +322,33 @@ async Task<bool> IMcpTaskStore.SetCancelledAsync(string taskId, CancellationToke
             return await base.SetCancelledAsync(taskId, cancellationToken);
         }
     }
+
+    private sealed class BlockingAsyncDisposable : IAsyncDisposable
+    {
+        private readonly TaskCompletionSource<bool> _disposeStarted = new(TaskCreationOptions.RunContinuationsAsynchronously);
+        private readonly TaskCompletionSource<bool> _release = new(TaskCreationOptions.RunContinuationsAsynchronously);
+
+        public Task DisposeStarted => _disposeStarted.Task;
+
+        public void Release() => _release.TrySetResult(true);
+
+        public async ValueTask DisposeAsync()
+        {
+            _disposeStarted.TrySetResult(true);
+            await _release.Task;
+        }
+    }
+
+    private sealed class RecordingAsyncDisposable : IAsyncDisposable
+    {
+        public bool IsDisposed { get; private set; }
+
+        public ValueTask DisposeAsync()
+        {
+            IsDisposed = true;
+            return default;
+        }
+    }
 }
 
 public class McpServerLifetimeFeatureTests(ITestOutputHelper testOutputHelper) : LoggedTest(testOutputHelper)
@@ -299,24 +365,26 @@ public async Task DisposeAsync_DoesNotCancelOrWaitForStatelessBackgroundTask()
             },
             LoggerFactory);
         var serverLifetime = Assert.IsAssignableFrom<IMcpServerLifetimeFeature>(statelessServer);
-        var releaseBackgroundTask = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
-        Task backgroundTask = releaseBackgroundTask.Task;
+        var backgroundResource = new RecordingAsyncDisposable();
 
-        serverLifetime.RegisterBackgroundTask(backgroundTask);
+        using var registration = serverLifetime.RegisterForDisposeAsync(backgroundResource);
 
-        try
-        {
-            await statelessServer.DisposeAsync().AsTask()
-                .WaitAsync(TestConstants.DefaultTimeout, TestContext.Current.CancellationToken);
+        await statelessServer.DisposeAsync().AsTask()
+            .WaitAsync(TestConstants.DefaultTimeout, TestContext.Current.CancellationToken);
 
-            Assert.False(serverLifetime.BackgroundTaskCancellationToken.CanBeCanceled);
-            Assert.False(backgroundTask.IsCompleted,
-                "A stateless per-request server should not own background work that outlives the request.");
-        }
-        finally
+        Assert.False(serverLifetime.ServerCancellationToken.CanBeCanceled);
+        Assert.False(backgroundResource.IsDisposed,
+            "A stateless per-request server should not own background work that outlives the request.");
+    }
+
+    private sealed class RecordingAsyncDisposable : IAsyncDisposable
+    {
+        public bool IsDisposed { get; private set; }
+
+        public ValueTask DisposeAsync()
         {
-            releaseBackgroundTask.TrySetResult(true);
-            await backgroundTask;
+            IsDisposed = true;
+            return default;
         }
     }
 }