Refactor transports to help enable graceful shutdown

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,17 @@ 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);
+                var transportTask = transport.RunAsync(cancellationToken: requestAborted);
+                await server.RunAsync(cancellationToken: requestAborted);
+                await transportTask;
             }
             catch (OperationCanceledException) when (requestAborted.IsCancellationRequested)
             {

src/ModelContextProtocol/Client/McpClient.cs

@@ -1,35 +1,36 @@
-using ModelContextProtocol.Configuration;
+using Microsoft.Extensions.Logging;
+using ModelContextProtocol.Configuration;
 using ModelContextProtocol.Logging;
 using ModelContextProtocol.Protocol.Messages;
 using ModelContextProtocol.Protocol.Transport;
 using ModelContextProtocol.Protocol.Types;
 using ModelContextProtocol.Shared;
 using ModelContextProtocol.Utils.Json;
-using Microsoft.Extensions.Logging;
 using System.Text.Json;
 
 namespace ModelContextProtocol.Client;
 
 /// <inheritdoc/>
 internal sealed class McpClient : McpJsonRpcEndpoint, IMcpClient
 {
-    private readonly McpClientOptions _options;
     private readonly IClientTransport _clientTransport;
+    private readonly McpClientOptions _options;
 
-    private volatile bool _isInitializing;
+    private ITransport? _sessionTransport;
+    private CancellationTokenSource? _connectCts;
 
     /// <summary>
     /// Initializes a new instance of the <see cref="McpClient"/> class.
     /// </summary>
-    /// <param name="transport">The transport to use for communication with the server.</param>
+    /// <param name="clientTransport">The transport to use for communication with the server.</param>
     /// <param name="options">Options for the client, defining protocol version and capabilities.</param>
     /// <param name="serverConfig">The server configuration.</param>
     /// <param name="loggerFactory">The logger factory.</param>
-    public McpClient(IClientTransport transport, McpClientOptions options, McpServerConfig serverConfig, ILoggerFactory? loggerFactory)
-        : base(transport, loggerFactory)
+    public McpClient(IClientTransport clientTransport, McpClientOptions options, McpServerConfig serverConfig, ILoggerFactory? loggerFactory)
+        : base(loggerFactory)
     {
+        _clientTransport = clientTransport;
         _options = options;
-        _clientTransport = transport;
 
         EndpointName = $"Client ({serverConfig.Id}: {serverConfig.Name})";
 
@@ -70,95 +71,90 @@ public McpClient(IClientTransport transport, McpClientOptions options, McpServer
     /// <inheritdoc/>
     public override string EndpointName { get; }
 
-    /// <inheritdoc/>
     public async Task ConnectAsync(CancellationToken cancellationToken = default)
     {
-        if (IsInitialized)
-        {
-            _logger.ClientAlreadyInitialized(EndpointName);
-            return;
-        }
-
-        if (_isInitializing)
-        {
-            _logger.ClientAlreadyInitializing(EndpointName);
-            throw new InvalidOperationException("Client is already initializing");
-        }
+        _connectCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+        cancellationToken = _connectCts.Token;
 
-        _isInitializing = true;
         try
         {
-            CancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
-
             // Connect transport
-            await _clientTransport.ConnectAsync(CancellationTokenSource.Token).ConfigureAwait(false);
-
-            // Start processing messages
-            MessageProcessingTask = ProcessMessagesAsync(CancellationTokenSource.Token);
+            _sessionTransport = await _clientTransport.ConnectAsync(cancellationToken).ConfigureAwait(false);
+            InitializeSession(_sessionTransport);
+            // We don't want the ConnectAsync token to cancel the session after we've successfully connected.
+            // The base class handles cleaning up the session in DisposeAsync without our help.
+            StartSession(fullSessionCancellationToken: CancellationToken.None);
 
             // 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 InitializeRequestParams()
+                        {
+                            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);
+            await DisposeAsync().ConfigureAwait(false);
             throw;
         }
-        finally
-        {
-            _isInitializing = false;
-        }
     }
 
-    private async Task InitializeAsync(CancellationToken cancellationToken)
+    /// <inheritdoc/>
+    public override async ValueTask DisposeUnsynchronizedAsync()
     {
-        using var initializationCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
-        initializationCts.CancelAfter(_options.InitializationTimeout);
-
-        try
+        if (_connectCts is not null)
         {
-            // Send initialize request
-            var initializeResponse = await SendRequestAsync<InitializeResult>(
-                new JsonRpcRequest
-                {
-                    Method = "initialize",
-                    Params = new InitializeRequestParams()
-                    {
-                        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);
+            await _connectCts.CancelAsync().ConfigureAwait(false);
         }
-        catch (OperationCanceledException) when (initializationCts.IsCancellationRequested)
+
+        await base.DisposeUnsynchronizedAsync().ConfigureAwait(false);
+
+        if (_sessionTransport is not null)
         {
-            _logger.ClientInitializationTimeout(EndpointName);
-            throw new McpClientException("Initialization timed out");
+            await _sessionTransport.DisposeAsync().ConfigureAwait(false);
         }
+
+        _connectCts?.Dispose();
     }
 }

src/ModelContextProtocol/Client/McpClientFactory.cs

@@ -64,7 +64,6 @@ public static async Task<IMcpClient> CreateAsync(
         var transport = 
             createTransportFunc(serverConfig, loggerFactory) ??
             throw new InvalidOperationException($"{nameof(createTransportFunc)} returned a null transport.");
-
         try
         {
             McpClient client = new(transport, clientOptions, serverConfig, loggerFactory);
@@ -82,7 +81,14 @@ public static async Task<IMcpClient> CreateAsync(
         }
         catch
         {
-            await transport.DisposeAsync().ConfigureAwait(false);
+            if (transport is IAsyncDisposable asyncDisposableTransport)
+            {
+                await asyncDisposableTransport.DisposeAsync().ConfigureAwait(false);
+            }
+            else if (transport is IDisposable disposableTransport)
+            {
+                disposableTransport.Dispose();
+            }
             throw;
         }
     }

src/ModelContextProtocol/Configuration/McpServerBuilderExtensions.Transports.cs

@@ -3,6 +3,9 @@
 using ModelContextProtocol.Protocol.Transport;
 using ModelContextProtocol.Utils;
 using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Logging;
+using Microsoft.Extensions.Options;
+using ModelContextProtocol.Server;
 
 namespace ModelContextProtocol;
 
@@ -19,8 +22,18 @@ public static IMcpServerBuilder WithStdioServerTransport(this IMcpServerBuilder
     {
         Throw.IfNull(builder);
 
-        builder.Services.AddSingleton<IServerTransport, StdioServerTransport>();
-        builder.Services.AddHostedService<McpServerHostedService>();
+        builder.Services.AddSingleton<ITransport, StdioServerTransport>();
+        builder.Services.AddHostedService<McpServerSingleSessionHostedService>();
+
+        builder.Services.AddSingleton(services =>
+        {
+            ITransport serverTransport = services.GetRequiredService<ITransport>();
+            IOptions<McpServerOptions> options = services.GetRequiredService<IOptions<McpServerOptions>>();
+            ILoggerFactory? loggerFactory = services.GetService<ILoggerFactory>();
+
+            return McpServerFactory.Create(serverTransport, options.Value, loggerFactory, services);
+        });
+
         return builder;
     }
 
@@ -33,7 +46,7 @@ public static IMcpServerBuilder WithHttpListenerSseServerTransport(this IMcpServ
         Throw.IfNull(builder);
 
         builder.Services.AddSingleton<IServerTransport, HttpListenerSseServerTransport>();
-        builder.Services.AddHostedService<McpServerHostedService>();
+        builder.Services.AddHostedService<McpServerMultiSessionHostedService>();
         return builder;
     }
 }

src/ModelContextProtocol/Configuration/McpServerServiceCollectionExtension.cs

@@ -8,7 +8,7 @@
 namespace ModelContextProtocol;
 
 /// <summary>
-/// Extension to host the MCP server
+/// Extension to host an MCP server
 /// </summary>
 public static class McpServerServiceCollectionExtension
 {
@@ -20,15 +20,6 @@ public static class McpServerServiceCollectionExtension
     /// <returns></returns>
     public static IMcpServerBuilder AddMcpServer(this IServiceCollection services, Action<McpServerOptions>? configureOptions = null)
     {
-        services.AddSingleton(services =>
-        {
-            IServerTransport serverTransport = services.GetRequiredService<IServerTransport>();
-            IOptions<McpServerOptions> options = services.GetRequiredService<IOptions<McpServerOptions>>();
-            ILoggerFactory? loggerFactory = services.GetService<ILoggerFactory>();
-
-            return McpServerFactory.Create(serverTransport, options.Value, loggerFactory, services);
-        });
-
         services.AddOptions();
         services.AddTransient<IConfigureOptions<McpServerOptions>, McpServerOptionsSetup>();
         if (configureOptions is not null)