Pass session id's to MCP endpoints. (#466)

* Pass session id's to MCP endpoints.

* Update src/ModelContextProtocol/Client/McpClient.cs

Co-authored-by: Stephen Halter <halter73@gmail.com>

* Update src/ModelContextProtocol.AspNetCore/StreamableHttpHandler.cs

Co-authored-by: Stephen Halter <halter73@gmail.com>

* Make init callback asynchronous.

* Address feedback.

---------

Co-authored-by: Stephen Halter <halter73@gmail.com>

src/ModelContextProtocol.AspNetCore/SseHandler.cs

@@ -31,7 +31,7 @@ internal sealed class SseHandler(
 
         var requestPath = (context.Request.PathBase + context.Request.Path).ToString();
         var endpointPattern = requestPath[..(requestPath.LastIndexOf('/') + 1)];
-        await using var transport = new SseResponseStreamTransport(context.Response.Body, $"{endpointPattern}message?sessionId={sessionId}");
+        await using var transport = new SseResponseStreamTransport(context.Response.Body, $"{endpointPattern}message?sessionId={sessionId}", sessionId);
 
         var userIdClaim = StreamableHttpHandler.GetUserIdClaim(context.User);
         await using var httpMcpSession = new HttpMcpSession<SseResponseStreamTransport>(sessionId, transport, userIdClaim, httpMcpServerOptions.Value.TimeProvider);

src/ModelContextProtocol.AspNetCore/StreamableHttpHandler.cs

@@ -4,7 +4,6 @@ using Microsoft.AspNetCore.Http.Features;
 using Microsoft.AspNetCore.WebUtilities;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
-using Microsoft.Extensions.Primitives;
 using Microsoft.Net.Http.Headers;
 using ModelContextProtocol.AspNetCore.Stateless;
 using ModelContextProtocol.Protocol;
@@ -136,6 +135,7 @@ internal sealed class StreamableHttpHandler(
             var transport = new StreamableHttpServerTransport
             {
                 Stateless = true,
+                SessionId = sessionId,
             };
             session = await CreateSessionAsync(context, transport, sessionId, statelessSessionId);
         }
@@ -184,7 +184,10 @@ internal sealed class StreamableHttpHandler(
         if (!HttpServerTransportOptions.Stateless)
         {
             sessionId = MakeNewSessionId();
-            transport = new();
+            transport = new()
+            {
+                SessionId = sessionId,
+            };
             context.Response.Headers["mcp-session-id"] = sessionId;
         }
         else
@@ -286,21 +289,19 @@ internal sealed class StreamableHttpHandler(
 
     private void ScheduleStatelessSessionIdWrite(HttpContext context, StreamableHttpServerTransport transport)
     {
-        context.Response.OnStarting(() =>
+        transport.OnInitRequestReceived = initRequestParams =>
         {
             var statelessId = new StatelessSessionId
             {
-                ClientInfo = transport?.InitializeRequest?.ClientInfo,
+                ClientInfo = initRequestParams?.ClientInfo,
                 UserIdClaim = GetUserIdClaim(context.User),
             };
 
             var sessionJson = JsonSerializer.Serialize(statelessId, StatelessSessionIdJsonContext.Default.StatelessSessionId);
-            var sessionId = Protector.Protect(sessionJson);
-
-            context.Response.Headers["mcp-session-id"] = sessionId;
-
-            return Task.CompletedTask;
-        });
+            transport.SessionId = Protector.Protect(sessionJson);
+            context.Response.Headers["mcp-session-id"] = transport.SessionId;
+            return ValueTask.CompletedTask;
+        };
     }
 
     internal static Task RunSessionAsync(HttpContext httpContext, IMcpServer session, CancellationToken requestAborted)

src/ModelContextProtocol.Core/Client/AutoDetectingClientSessionTransport.cs

@@ -45,6 +45,8 @@ internal sealed partial class AutoDetectingClientSessionTransport : ITransport
 
     public ChannelReader<JsonRpcMessage> MessageReader => _messageChannel.Reader;
 
+    string? ITransport.SessionId => ActiveTransport?.SessionId;
+
     /// <inheritdoc/>
     public Task SendMessageAsync(JsonRpcMessage message, CancellationToken cancellationToken = default)
     {

src/ModelContextProtocol.Core/Client/McpClient.cs

@@ -1,5 +1,6 @@
 using Microsoft.Extensions.Logging;
 using ModelContextProtocol.Protocol;
+using System.Diagnostics;
 using System.Text.Json;
 
 namespace ModelContextProtocol.Client;
@@ -93,6 +94,20 @@ internal sealed partial class McpClient : McpEndpoint, IMcpClient
         }
     }
 
+    /// <inheritdoc/>
+    public string? SessionId
+    {
+        get
+        {
+            if (_sessionTransport is null)
+            {
+                throw new InvalidOperationException("Must have already initialized a session when invoking this property.");
+            }
+
+            return _sessionTransport.SessionId;
+        }
+    }
+
     /// <inheritdoc/>
     public ServerCapabilities ServerCapabilities => _serverCapabilities ?? throw new InvalidOperationException("The client is not connected.");
 

src/ModelContextProtocol.Core/Client/StreamableHttpClientSessionTransport.cs

@@ -27,7 +27,6 @@ internal sealed partial class StreamableHttpClientSessionTransport : TransportBa
     private readonly CancellationTokenSource _connectionCts;
     private readonly ILogger _logger;
 
-    private string? _mcpSessionId;
     private Task? _getReceiveTask;
 
     public StreamableHttpClientSessionTransport(
@@ -85,7 +84,7 @@ internal sealed partial class StreamableHttpClientSessionTransport : TransportBa
             },
         };
 
-        CopyAdditionalHeaders(httpRequestMessage.Headers, _options.AdditionalHeaders, _mcpSessionId);
+        CopyAdditionalHeaders(httpRequestMessage.Headers, _options.AdditionalHeaders, SessionId);
 
         var response = await _httpClient.SendAsync(httpRequestMessage, HttpCompletionOption.ResponseHeadersRead, cancellationToken).ConfigureAwait(false);
 
@@ -124,7 +123,7 @@ internal sealed partial class StreamableHttpClientSessionTransport : TransportBa
             // We've successfully initialized! Copy session-id and start GET request if any.
             if (response.Headers.TryGetValues("mcp-session-id", out var sessionIdValues))
             {
-                _mcpSessionId = sessionIdValues.FirstOrDefault();
+                SessionId = sessionIdValues.FirstOrDefault();
             }
 
             _getReceiveTask = ReceiveUnsolicitedMessagesAsync();
@@ -170,7 +169,7 @@ internal sealed partial class StreamableHttpClientSessionTransport : TransportBa
         // Send a GET request to handle any unsolicited messages not sent over a POST response.
         using var request = new HttpRequestMessage(HttpMethod.Get, _options.Endpoint);
         request.Headers.Accept.Add(s_textEventStreamMediaType);
-        CopyAdditionalHeaders(request.Headers, _options.AdditionalHeaders, _mcpSessionId);
+        CopyAdditionalHeaders(request.Headers, _options.AdditionalHeaders, SessionId);
 
         using var response = await _httpClient.SendAsync(request, HttpCompletionOption.ResponseHeadersRead, _connectionCts.Token).ConfigureAwait(false);
 

src/ModelContextProtocol.Core/IMcpEndpoint.cs

@@ -28,6 +28,14 @@ namespace ModelContextProtocol;
 /// </remarks>
 public interface IMcpEndpoint : IAsyncDisposable
 {
+    /// <summary>Gets an identifier associated with the current MCP session.</summary>
+    /// <remarks>
+    /// Typically populated in transports supporting multiple sessions such as Streamable HTTP or SSE.
+    /// Can return <see langword="null"/> if the session hasn't initialized or if the transport doesn't
+    /// support multiple sessions (as is the case with STDIO).
+    /// </remarks>
+    string? SessionId { get; }
+
     /// <summary>
     /// Sends a JSON-RPC request to the connected endpoint and waits for a response.
     /// </summary>

src/ModelContextProtocol.Core/Protocol/ITransport.cs

@@ -25,6 +25,14 @@ namespace ModelContextProtocol.Protocol;
 /// </remarks>
 public interface ITransport : IAsyncDisposable
 {
+    /// <summary>Gets an identifier associated with the current MCP session.</summary>
+    /// <remarks>
+    /// Typically populated in transports supporting multiple sessions such as Streamable HTTP or SSE.
+    /// Can return <see langword="null"/> if the session hasn't initialized or if the transport doesn't
+    /// support multiple sessions (as is the case with STDIO).
+    /// </remarks>
+    string? SessionId { get; }
+
     /// <summary>
     /// Gets a channel reader for receiving messages from the transport.
     /// </summary>

src/ModelContextProtocol.Core/Protocol/TransportBase.cs

@@ -59,6 +59,9 @@ public abstract partial class TransportBase : ITransport
     /// <summary>Gets the logger used by this transport.</summary>
     private protected ILogger Logger => _logger;
 
+    /// <inheritdoc/>
+    public virtual string? SessionId { get; protected set; }
+
     /// <summary>
     /// Gets the name that identifies this transport endpoint in logs.
     /// </summary>

src/ModelContextProtocol.Core/Server/DestinationBoundMcpServer.cs

@@ -6,6 +6,7 @@ namespace ModelContextProtocol.Server;
 internal sealed class DestinationBoundMcpServer(McpServer server, ITransport? transport) : IMcpServer
 {
     public string EndpointName => server.EndpointName;
+    public string? SessionId => transport?.SessionId ?? server.SessionId;
     public ClientCapabilities? ClientCapabilities => server.ClientCapabilities;
     public Implementation? ClientInfo => server.ClientInfo;
     public McpServerOptions ServerOptions => server.ServerOptions;

src/ModelContextProtocol.Core/Server/McpServer.cs

@@ -96,6 +96,10 @@ internal sealed class McpServer : McpEndpoint, IMcpServer
         InitializeSession(transport);
     }
 
+    /// <inheritdoc/>
+    public string? SessionId => _sessionTransport.SessionId;
+
+    /// <inheritdoc/>
     public ServerCapabilities ServerCapabilities { get; } = new();
 
     /// <inheritdoc />

src/ModelContextProtocol.Core/Server/SseResponseStreamTransport.cs

@@ -23,7 +23,8 @@ namespace ModelContextProtocol.Server;
 /// These messages should be passed to <see cref="OnMessageReceivedAsync(JsonRpcMessage, CancellationToken)"/>.
 /// Defaults to "/message".
 /// </param>
-public sealed class SseResponseStreamTransport(Stream sseResponseStream, string? messageEndpoint = "/message") : ITransport
+/// <param name="sessionId">The identifier corresponding to the current MCP session.</param>
+public sealed class SseResponseStreamTransport(Stream sseResponseStream, string? messageEndpoint = "/message", string? sessionId = null) : ITransport
 {
     private readonly SseWriter _sseWriter = new(messageEndpoint);
     private readonly Channel<JsonRpcMessage> _incomingChannel = Channel.CreateBounded<JsonRpcMessage>(new BoundedChannelOptions(1)
@@ -49,6 +50,9 @@ public sealed class SseResponseStreamTransport(Stream sseResponseStream, string?
     /// <inheritdoc/>
     public ChannelReader<JsonRpcMessage> MessageReader => _incomingChannel.Reader;
 
+    /// <inheritdoc/>
+    public string? SessionId { get; } = sessionId;
+
     /// <inheritdoc/>
     public async ValueTask DisposeAsync()
     {

src/ModelContextProtocol.Core/Server/StreamableHttpPostTransport.cs

@@ -18,6 +18,8 @@ internal sealed class StreamableHttpPostTransport(StreamableHttpServerTransport
 
     public ChannelReader<JsonRpcMessage> MessageReader => throw new NotSupportedException("JsonRpcMessage.RelatedTransport should only be used for sending messages.");
 
+    string? ITransport.SessionId => parentTransport.SessionId;
+
     /// <returns>
     /// True, if data was written to the respond body.
     /// False, if nothing was written because the request body did not contain any <see cref="JsonRpcRequest"/> messages to respond to.
@@ -79,10 +81,11 @@ internal sealed class StreamableHttpPostTransport(StreamableHttpServerTransport
         {
             _pendingRequest = request.Id;
 
-            // Store client capabilities so they can be serialized by "stateless" callers for use in later requests.
-            if (parentTransport.Stateless && request.Method == RequestMethods.Initialize)
+            // Invoke the initialize request callback if applicable.
+            if (parentTransport.OnInitRequestReceived is { } onInitRequest && request.Method == RequestMethods.Initialize)
             {
-                parentTransport.InitializeRequest = JsonSerializer.Deserialize(request.Params, McpJsonUtilities.JsonContext.Default.InitializeRequestParams);
+                var initializeRequest = JsonSerializer.Deserialize(request.Params, McpJsonUtilities.JsonContext.Default.InitializeRequestParams);
+                await onInitRequest(initializeRequest).ConfigureAwait(false);
             }
         }
 

src/ModelContextProtocol.Core/Server/StreamableHttpServerTransport.cs

@@ -46,15 +46,18 @@ public sealed class StreamableHttpServerTransport : ITransport
     public bool Stateless { get; init; }
 
     /// <summary>
-    /// Gets the initialize request if it was received by <see cref="HandlePostRequest(IDuplexPipe, CancellationToken)"/> and <see cref="Stateless"/> is set to <see langword="true"/>.
+    /// Gets or sets a callback to be invoked before handling the initialize request.
     /// </summary>
-    public InitializeRequestParams? InitializeRequest { get; internal set; }
+    public Func<InitializeRequestParams?, ValueTask>? OnInitRequestReceived { get; set; }
 
     /// <inheritdoc/>
     public ChannelReader<JsonRpcMessage> MessageReader => _incomingChannel.Reader;
 
     internal ChannelWriter<JsonRpcMessage> MessageWriter => _incomingChannel.Writer;
 
+    /// <inheritdoc/>
+    public string? SessionId { get; set; }
+
     /// <summary>
     /// Handles an optional SSE GET request a client using the Streamable HTTP transport might make by
     /// writing any unsolicited JSON-RPC messages sent via <see cref="SendMessageAsync"/>