Enable graceful shutdown of servers

Author: stephentoub

README.md

@@ -198,11 +198,7 @@ McpServerOptions options = new()
 };
 
 await using IMcpServer server = McpServerFactory.Create(new StdioServerTransport("MyServer"), options);
-
-await server.StartAsync();
-
-// Run until process is stopped by the client (parent process)
-await Task.Delay(Timeout.Infinite);
+await server.RunAsync();
 ```
 
 ## Acknowledgements

samples/AspNetCoreSseServer/McpEndpointRouteBuilderExtensions.cs

@@ -10,7 +10,6 @@ public static class McpEndpointRouteBuilderExtensions
 {
     public static IEndpointConventionBuilder MapMcpSse(this IEndpointRouteBuilder endpoints)
     {
-        IMcpServer? server = null;
         SseResponseStreamTransport? transport = null;
         var loggerFactory = endpoints.ServiceProvider.GetRequiredService<ILoggerFactory>();
         var mcpServerOptions = endpoints.ServiceProvider.GetRequiredService<IOptions<McpServerOptions>>();
@@ -19,17 +18,15 @@ public static IEndpointConventionBuilder MapMcpSse(this IEndpointRouteBuilder en
 
         routeGroup.MapGet("/sse", async (HttpResponse response, CancellationToken requestAborted) =>
         {
-            await using var localTransport = transport = new SseResponseStreamTransport(response.Body);
-            await using var localServer = server = McpServerFactory.Create(transport, mcpServerOptions.Value, loggerFactory, endpoints.ServiceProvider);
-
-            await localServer.StartAsync(requestAborted);
-
             response.Headers.ContentType = "text/event-stream";
             response.Headers.CacheControl = "no-cache";
 
+            await using var localTransport = transport = new SseResponseStreamTransport(response.Body);
+            await using var server = McpServerFactory.Create(transport, mcpServerOptions.Value, loggerFactory, endpoints.ServiceProvider);
+
             try
             {
-                await transport.RunAsync(requestAborted);
+                await transport.RunAsync(cancellationToken: requestAborted);
             }
             catch (OperationCanceledException) when (requestAborted.IsCancellationRequested)
             {

src/ModelContextProtocol/Client/McpClient.cs

@@ -6,7 +6,6 @@
 using ModelContextProtocol.Shared;
 using ModelContextProtocol.Utils.Json;
 using Microsoft.Extensions.Logging;
-using Microsoft.Extensions.Logging.Abstractions;
 using System.Text.Json;
 
 namespace ModelContextProtocol.Client;
@@ -17,7 +16,7 @@ internal sealed class McpClient : McpJsonRpcEndpoint, IMcpClient
     private readonly McpClientOptions _options;
     private readonly IClientTransport _clientTransport;
 
-    private volatile bool _isInitializing;
+    private int _connecting;
 
     /// <summary>
     /// Initializes a new instance of the <see cref="McpClient"/> class.
@@ -74,92 +73,75 @@ public McpClient(IClientTransport transport, McpClientOptions options, McpServer
     /// <inheritdoc/>
     public async Task ConnectAsync(CancellationToken cancellationToken = default)
     {
-        if (IsInitialized)
-        {
-            _logger.ClientAlreadyInitialized(EndpointName);
-            return;
-        }
-
-        if (_isInitializing)
+        if (Interlocked.Exchange(ref _connecting, 1) != 0)
         {
             _logger.ClientAlreadyInitializing(EndpointName);
-            throw new InvalidOperationException("Client is already initializing");
+            throw new InvalidOperationException("Client is already in use.");
         }
 
-        _isInitializing = true;
+        CancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+        cancellationToken = CancellationTokenSource.Token;
+
         try
         {
-            CancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
-
             // Connect transport
-            await _clientTransport.ConnectAsync(CancellationTokenSource.Token).ConfigureAwait(false);
+            await _clientTransport.ConnectAsync(cancellationToken).ConfigureAwait(false);
 
             // Start processing messages
-            MessageProcessingTask = ProcessMessagesAsync(CancellationTokenSource.Token);
+            MessageProcessingTask = ProcessMessagesAsync(cancellationToken);
 
             // Perform initialization sequence
-            await InitializeAsync(CancellationTokenSource.Token).ConfigureAwait(false);
+            using var initializationCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+            initializationCts.CancelAfter(_options.InitializationTimeout);
 
-            IsInitialized = true;
+            try
+            {
+                // Send initialize request
+                var initializeResponse = await SendRequestAsync<InitializeResult>(
+                    new JsonRpcRequest
+                    {
+                        Method = "initialize",
+                        Params = new
+                        {
+                            protocolVersion = _options.ProtocolVersion,
+                            capabilities = _options.Capabilities ?? new ClientCapabilities(),
+                            clientInfo = _options.ClientInfo
+                        }
+                    },
+                    initializationCts.Token).ConfigureAwait(false);
+
+                // Store server information
+                _logger.ServerCapabilitiesReceived(EndpointName,
+                    capabilities: JsonSerializer.Serialize(initializeResponse.Capabilities, McpJsonUtilities.JsonContext.Default.ServerCapabilities),
+                    serverInfo: JsonSerializer.Serialize(initializeResponse.ServerInfo, McpJsonUtilities.JsonContext.Default.Implementation));
+
+                ServerCapabilities = initializeResponse.Capabilities;
+                ServerInfo = initializeResponse.ServerInfo;
+                ServerInstructions = initializeResponse.Instructions;
+
+                // Validate protocol version
+                if (initializeResponse.ProtocolVersion != _options.ProtocolVersion)
+                {
+                    _logger.ServerProtocolVersionMismatch(EndpointName, _options.ProtocolVersion, initializeResponse.ProtocolVersion);
+                    throw new McpClientException($"Server protocol version mismatch. Expected {_options.ProtocolVersion}, got {initializeResponse.ProtocolVersion}");
+                }
+
+                // Send initialized notification
+                await SendMessageAsync(
+                    new JsonRpcNotification { Method = "notifications/initialized" },
+                    initializationCts.Token).ConfigureAwait(false);
+            }
+            catch (OperationCanceledException) when (initializationCts.IsCancellationRequested)
+            {
+                _logger.ClientInitializationTimeout(EndpointName);
+                throw new McpClientException("Initialization timed out");
+            }
         }
         catch (Exception e)
         {
             _logger.ClientInitializationError(EndpointName, e);
             await CleanupAsync().ConfigureAwait(false);
             throw;
         }
-        finally
-        {
-            _isInitializing = false;
-        }
-    }
-
-    private async Task InitializeAsync(CancellationToken cancellationToken)
-    {
-        using var initializationCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
-        initializationCts.CancelAfter(_options.InitializationTimeout);
-
-        try
-        {
-            // Send initialize request
-            var initializeResponse = await SendRequestAsync<InitializeResult>(
-                new JsonRpcRequest
-                {
-                    Method = "initialize",
-                    Params = new
-                    {
-                        protocolVersion = _options.ProtocolVersion,
-                        capabilities = _options.Capabilities ?? new ClientCapabilities(),
-                        clientInfo = _options.ClientInfo
-                    }
-                },
-                initializationCts.Token).ConfigureAwait(false);
-
-            // Store server information
-            _logger.ServerCapabilitiesReceived(EndpointName, 
-                capabilities: JsonSerializer.Serialize(initializeResponse.Capabilities, McpJsonUtilities.JsonContext.Default.ServerCapabilities),
-                serverInfo: JsonSerializer.Serialize(initializeResponse.ServerInfo, McpJsonUtilities.JsonContext.Default.Implementation));
-
-            ServerCapabilities = initializeResponse.Capabilities;
-            ServerInfo = initializeResponse.ServerInfo;
-            ServerInstructions = initializeResponse.Instructions;
-
-            // Validate protocol version
-            if (initializeResponse.ProtocolVersion != _options.ProtocolVersion)
-            {
-                _logger.ServerProtocolVersionMismatch(EndpointName, _options.ProtocolVersion, initializeResponse.ProtocolVersion);
-                throw new McpClientException($"Server protocol version mismatch. Expected {_options.ProtocolVersion}, got {initializeResponse.ProtocolVersion}");
-            }
-
-            // Send initialized notification
-            await SendMessageAsync(
-                new JsonRpcNotification { Method = "notifications/initialized" },
-                initializationCts.Token).ConfigureAwait(false);
-        }
-        catch (OperationCanceledException) when (initializationCts.IsCancellationRequested)
-        {
-            _logger.ClientInitializationTimeout(EndpointName);
-            throw new McpClientException("Initialization timed out");
-        }
     }
 }

src/ModelContextProtocol/Hosting/McpServerHostedService.cs

@@ -26,6 +26,6 @@ public McpServerHostedService(IMcpServer server)
     /// <inheritdoc />
     protected override async Task ExecuteAsync(CancellationToken stoppingToken)
     {
-        await _server.StartAsync(stoppingToken).ConfigureAwait(false);
+        await _server.RunAsync(cancellationToken: stoppingToken).ConfigureAwait(false);
     }
 }