Split StdioClient/ServerTransports into Stdio : Stream

Author: stephentoub

src/ModelContextProtocol.AspNetCore/McpEndpointRouteBuilderExtensions.cs

@@ -27,7 +27,7 @@ public static class McpEndpointRouteBuilderExtensions
     /// <returns>Returns a builder for configuring additional endpoint conventions like authorization policies.</returns>
     public static IEndpointConventionBuilder MapMcp(this IEndpointRouteBuilder endpoints, Func<HttpContext, IMcpServer, CancellationToken, Task>? runSession = null)
     {
-        ConcurrentDictionary<string, SseResponseStreamTransport> _sessions = new(StringComparer.Ordinal);
+        ConcurrentDictionary<string, SseResponseSessionTransport> _sessions = new(StringComparer.Ordinal);
 
         var loggerFactory = endpoints.ServiceProvider.GetRequiredService<ILoggerFactory>();
         var mcpServerOptions = endpoints.ServiceProvider.GetRequiredService<IOptions<McpServerOptions>>();
@@ -43,7 +43,7 @@ public static IEndpointConventionBuilder MapMcp(this IEndpointRouteBuilder endpo
             response.Headers.CacheControl = "no-store";
 
             var sessionId = MakeNewSessionId();
-            await using var transport = new SseResponseStreamTransport(response.Body, $"/message?sessionId={sessionId}");
+            await using var transport = new SseResponseSessionTransport(response.Body, $"/message?sessionId={sessionId}");
             if (!_sessions.TryAdd(sessionId, transport))
             {
                 throw new Exception($"Unreachable given good entropy! Session with ID '{sessionId}' has already been created.");

src/ModelContextProtocol/Protocol/Transport/SseResponseStreamTransport.cs renamed to src/ModelContextProtocol/Protocol/Transport/SseResponseSessionTransport.cs

@@ -13,7 +13,7 @@ namespace ModelContextProtocol.Protocol.Transport;
 /// </summary>
 /// <param name="sseResponseStream">The stream to write the SSE response body to.</param>
 /// <param name="messageEndpoint">The endpoint to send JSON-RPC messages to. Defaults to "/message".</param> 
-public sealed class SseResponseStreamTransport(Stream sseResponseStream, string messageEndpoint = "/message") : ITransport
+public sealed class SseResponseSessionTransport(Stream sseResponseStream, string messageEndpoint = "/message") : ITransport
 {
     private readonly Channel<IJsonRpcMessage> _incomingChannel = CreateBoundedChannel<IJsonRpcMessage>();
     private readonly Channel<SseItem<IJsonRpcMessage?>> _outgoingSseChannel = CreateBoundedChannel<SseItem<IJsonRpcMessage?>>();

src/ModelContextProtocol/Protocol/Transport/StdioClientSessionTransport.cs

@@ -0,0 +1,40 @@
+using Microsoft.Extensions.Logging;
+using ModelContextProtocol.Logging;
+using ModelContextProtocol.Protocol.Messages;
+using System.Diagnostics;
+
+namespace ModelContextProtocol.Protocol.Transport;
+
+/// <summary>Provides the client side of a stdio-based session transport.</summary>
+internal sealed class StdioClientSessionTransport : StreamClientSessionTransport
+{
+    private readonly StdioClientTransportOptions _options;
+    private readonly Process _process;
+
+    public StdioClientSessionTransport(StdioClientTransportOptions options, Process process, string endpointName, ILoggerFactory? loggerFactory)
+        : base(process.StandardInput, process.StandardOutput, endpointName, loggerFactory)
+    {
+        _process = process;
+        _options = options;
+    }
+
+    /// <inheritdoc/>
+    public override async Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default)
+    {
+        if (_process.HasExited)
+        {
+            Logger.TransportNotConnected(EndpointName);
+            throw new McpTransportException("Transport is not connected");
+        }
+
+        await base.SendMessageAsync(message, cancellationToken).ConfigureAwait(false);
+    }
+
+    /// <inheritdoc/>
+    protected override ValueTask CleanupAsync(CancellationToken cancellationToken)
+    {
+        StdioClientTransport.DisposeProcess(_process, processStarted: true, Logger, _options.ShutdownTimeout, EndpointName);
+
+        return base.CleanupAsync(cancellationToken);
+    }
+}