Add back dispose lock

Author: MackinnonBuck

src/ModelContextProtocol.Core/Client/McpClientImpl.cs

@@ -1,6 +1,7 @@
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Logging.Abstractions;
 using ModelContextProtocol.Protocol;
+using System.Runtime.InteropServices;
 using System.Text.Json;
 
 namespace ModelContextProtocol.Client;
@@ -19,14 +20,15 @@ internal sealed partial class McpClientImpl : McpClient
     private readonly string _endpointName;
     private readonly McpClientOptions _options;
     private readonly McpSessionHandler _sessionHandler;
+    private readonly SemaphoreSlim _disposeLock = new(1, 1);
 
     private CancellationTokenSource? _connectCts;
 
     private ServerCapabilities? _serverCapabilities;
     private Implementation? _serverInfo;
     private string? _serverInstructions;
 
-    private int _isDisposed;
+    private bool _disposed;
 
     /// <summary>
     /// Initializes a new instance of the <see cref="McpClientImpl"/> class.
@@ -215,11 +217,15 @@ public override IAsyncDisposable RegisterNotificationHandler(string method, Func
     /// <inheritdoc/>
     public override async ValueTask DisposeAsync()
     {
-        if (Interlocked.CompareExchange(ref _isDisposed, 1, 0) != 0)
+        using var _ = await _disposeLock.LockAsync().ConfigureAwait(false);
+
+        if (_disposed)
         {
             return;
         }
 
+        _disposed = true;
+
         await _sessionHandler.DisposeAsync().ConfigureAwait(false);
         await _transport.DisposeAsync().ConfigureAwait(false);
     }

src/ModelContextProtocol.Core/Server/McpServerImpl.cs

@@ -23,6 +23,7 @@ internal sealed class McpServerImpl : McpServer
     private readonly NotificationHandlers _notificationHandlers;
     private readonly RequestHandlers _requestHandlers;
     private readonly McpSessionHandler _sessionHandler;
+    private readonly SemaphoreSlim _disposeLock = new(1, 1);
 
     private ClientCapabilities? _clientCapabilities;
     private Implementation? _clientInfo;
@@ -31,7 +32,7 @@ internal sealed class McpServerImpl : McpServer
     private string _endpointName;
     private int _started;
 
-    private int _isDisposed;
+    private bool _disposed;
 
     /// <summary>Holds a boxed <see cref="LoggingLevel"/> value for the server.</summary>
     /// <remarks>
@@ -165,11 +166,15 @@ public override IAsyncDisposable RegisterNotificationHandler(string method, Func
     /// <inheritdoc/>
     public override async ValueTask DisposeAsync()
     {
-        if (Interlocked.CompareExchange(ref _isDisposed, 1, 0) != 0)
+        using var _ = await _disposeLock.LockAsync().ConfigureAwait(false);
+
+        if (_disposed)
         {
             return;
         }
 
+        _disposed = true;
+
         _disposables.ForEach(d => d());
         await _sessionHandler.DisposeAsync().ConfigureAwait(false);
     }