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

eiriktsarpalis · halter73 · web-flow · commit d7d9eab3f85a · 2025-06-12T17:56:57.000+03:00
* Pass session id's to MCP endpoints.

* Update src/ModelContextProtocol/Client/McpClient.cs

Co-authored-by: Stephen Halter &lt;halter73@gmail.com&gt;

* Update src/ModelContextProtocol.AspNetCore/StreamableHttpHandler.cs

Co-authored-by: Stephen Halter &lt;halter73@gmail.com&gt;

* Make init callback asynchronous.

* Address feedback.

---------

Co-authored-by: Stephen Halter &lt;halter73@gmail.com&gt;
diff --git a/src/ModelContextProtocol.AspNetCore/SseHandler.cs b/src/ModelContextProtocol.AspNetCore/SseHandler.cs
@@ -31,7 +31,7 @@ public async Task HandleSseRequestAsync(HttpContext context)
 
         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);
diff --git a/src/ModelContextProtocol.AspNetCore/StreamableHttpHandler.cs b/src/ModelContextProtocol.AspNetCore/StreamableHttpHandler.cs
@@ -4,7 +4,6 @@
 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 @@ public async Task HandleDeleteRequestAsync(HttpContext context)
             var transport = new StreamableHttpServerTransport
             {
                 Stateless = true,
+                SessionId = sessionId,
             };
             session = await CreateSessionAsync(context, transport, sessionId, statelessSessionId);
         }
@@ -184,7 +184,10 @@ private async ValueTask<HttpMcpSession<StreamableHttpServerTransport>> StartNewS
         if (!HttpServerTransportOptions.Stateless)
         {
             sessionId = MakeNewSessionId();
-            transport = new();
+            transport = new()
+            {
+                SessionId = sessionId,
+            };
             context.Response.Headers["mcp-session-id"] = sessionId;
         }
         else
@@ -286,21 +289,19 @@ internal static string MakeNewSessionId()
 
     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)
diff --git a/src/ModelContextProtocol.Core/Client/AutoDetectingClientSessionTransport.cs b/src/ModelContextProtocol.Core/Client/AutoDetectingClientSessionTransport.cs
@@ -45,6 +45,8 @@ public AutoDetectingClientSessionTransport(SseClientTransportOptions transportOp
 
     public ChannelReader<JsonRpcMessage> MessageReader => _messageChannel.Reader;
 
+    string? ITransport.SessionId => ActiveTransport?.SessionId;
+
     /// <inheritdoc/>
     public Task SendMessageAsync(JsonRpcMessage message, CancellationToken cancellationToken = default)
     {
diff --git a/src/ModelContextProtocol.Core/Client/McpClient.cs b/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 @@ public McpClient(IClientTransport clientTransport, McpClientOptions? options, IL
         }
     }
 
+    /// <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.");
 
diff --git a/src/ModelContextProtocol.Core/Client/StreamableHttpClientSessionTransport.cs b/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 async Task<HttpResponseMessage> SendHttpRequestAsync(JsonRpcMessage mes
             },
         };
 
-        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 async Task<HttpResponseMessage> SendHttpRequestAsync(JsonRpcMessage mes
             // 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 @@ private async Task ReceiveUnsolicitedMessagesAsync()
         // 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);
 
diff --git a/src/ModelContextProtocol.Core/IMcpEndpoint.cs b/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>
diff --git a/src/ModelContextProtocol.Core/Protocol/ITransport.cs b/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>
diff --git a/src/ModelContextProtocol.Core/Protocol/TransportBase.cs b/src/ModelContextProtocol.Core/Protocol/TransportBase.cs
@@ -59,6 +59,9 @@ internal TransportBase(string name, Channel<JsonRpcMessage>? messageChannel, ILo
     /// <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>
diff --git a/src/ModelContextProtocol.Core/Server/DestinationBoundMcpServer.cs b/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;
diff --git a/src/ModelContextProtocol.Core/Server/McpServer.cs b/src/ModelContextProtocol.Core/Server/McpServer.cs
@@ -96,6 +96,10 @@ void Register<TPrimitive>(McpServerPrimitiveCollection<TPrimitive>? collection,
         InitializeSession(transport);
     }
 
+    /// <inheritdoc/>
+    public string? SessionId => _sessionTransport.SessionId;
+
+    /// <inheritdoc/>
     public ServerCapabilities ServerCapabilities { get; } = new();
 
     /// <inheritdoc />
diff --git a/src/ModelContextProtocol.Core/Server/SseResponseStreamTransport.cs b/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 async Task RunAsync(CancellationToken cancellationToken)
     /// <inheritdoc/>
     public ChannelReader<JsonRpcMessage> MessageReader => _incomingChannel.Reader;
 
+    /// <inheritdoc/>
+    public string? SessionId { get; } = sessionId;
+
     /// <inheritdoc/>
     public async ValueTask DisposeAsync()
     {
diff --git a/src/ModelContextProtocol.Core/Server/StreamableHttpPostTransport.cs b/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 @@ private async ValueTask OnMessageReceivedAsync(JsonRpcMessage? message, Cancella
         {
             _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);
             }
         }
 
diff --git a/src/ModelContextProtocol.Core/Server/StreamableHttpServerTransport.cs b/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"/>
diff --git a/tests/Common/Utils/TestServerTransport.cs b/tests/Common/Utils/TestServerTransport.cs
@@ -16,6 +16,8 @@ public class TestServerTransport : ITransport
 
     public Action<JsonRpcMessage>? OnMessageSent { get; set; }
 
+    public string? SessionId => null;
+
     public TestServerTransport()
     {
         _messageChannel = Channel.CreateUnbounded<JsonRpcMessage>(new UnboundedChannelOptions
diff --git a/tests/ModelContextProtocol.AspNetCore.Tests/HttpServerIntegrationTests.cs b/tests/ModelContextProtocol.AspNetCore.Tests/HttpServerIntegrationTests.cs
@@ -52,6 +52,15 @@ public async Task Connect_TestServer_ShouldProvideServerFields()
         // Assert
         Assert.NotNull(client.ServerCapabilities);
         Assert.NotNull(client.ServerInfo);
+
+        if (ClientTransportOptions.Endpoint.AbsolutePath.EndsWith("/sse"))
+        {
+            Assert.Null(client.SessionId);
+        }
+        else
+        {
+            Assert.NotNull(client.SessionId);
+        }
     }
 
     [Fact]
@@ -90,6 +99,35 @@ public async Task CallTool_Sse_EchoServer()
         Assert.Equal("Echo: Hello MCP!", textContent.Text);
     }
 
+    [Fact]
+    public async Task CallTool_EchoSessionId_ReturnsTheSameSessionId()
+    {
+        // arrange
+
+        // act
+        await using var client = await GetClientAsync();
+        var result1 = await client.CallToolAsync("echoSessionId", cancellationToken: TestContext.Current.CancellationToken);
+        var result2 = await client.CallToolAsync("echoSessionId", cancellationToken: TestContext.Current.CancellationToken);
+        var result3 = await client.CallToolAsync("echoSessionId", cancellationToken: TestContext.Current.CancellationToken);
+
+        // assert
+        Assert.NotNull(result1);
+        Assert.NotNull(result2);
+        Assert.NotNull(result3);
+
+        Assert.False(result1.IsError);
+        Assert.False(result2.IsError);
+        Assert.False(result3.IsError);
+        
+        var textContent1 = Assert.Single(result1.Content);
+        var textContent2 = Assert.Single(result2.Content);
+        var textContent3 = Assert.Single(result3.Content);
+
+        Assert.NotNull(textContent1.Text);
+        Assert.Equal(textContent1.Text, textContent2.Text);
+        Assert.Equal(textContent1.Text, textContent3.Text);
+    }
+
     [Fact]
     public async Task ListResources_Sse_TestServer()
     {
diff --git a/tests/ModelContextProtocol.TestServer/Program.cs b/tests/ModelContextProtocol.TestServer/Program.cs
@@ -133,6 +133,16 @@ private static ToolsCapability ConfigureTools()
                                 """),
                         },
                         new Tool()
+                        {
+                            Name = "echoSessionId",
+                            Description = "Echoes the session id back to the client.",
+                            InputSchema = JsonSerializer.Deserialize<JsonElement>("""
+                                {
+                                    "type": "object"
+                                }
+                                """, McpJsonUtilities.DefaultOptions),
+                        },
+                        new Tool()
                         {
                             Name = "sampleLLM",
                             Description = "Samples from an LLM using MCP's sampling feature.",
@@ -170,6 +180,13 @@ private static ToolsCapability ConfigureTools()
                         Content = [new Content() { Text = "Echo: " + message.ToString(), Type = "text" }]
                     };
                 }
+                else if (request.Params?.Name == "echoSessionId")
+                {
+                    return new CallToolResponse()
+                    {
+                        Content = [new Content() { Text = request.Server.SessionId, Type = "text" }]
+                    };
+                }
                 else if (request.Params?.Name == "sampleLLM")
                 {
                     if (request.Params?.Arguments is null ||
diff --git a/tests/ModelContextProtocol.TestSseServer/Program.cs b/tests/ModelContextProtocol.TestSseServer/Program.cs
@@ -128,6 +128,16 @@ static CreateMessageRequestParams CreateRequestSamplingParams(string context, st
                                     """, McpJsonUtilities.DefaultOptions),
                             },
                             new Tool()
+                            {
+                                Name = "echoSessionId",
+                                Description = "Echoes the session id back to the client.",
+                                InputSchema = JsonSerializer.Deserialize<JsonElement>("""
+                                    {
+                                        "type": "object"
+                                    }
+                                    """, McpJsonUtilities.DefaultOptions),
+                            },
+                            new Tool()
                             {
                                 Name = "sampleLLM",
                                 Description = "Samples from an LLM using MCP's sampling feature.",
@@ -168,6 +178,13 @@ static CreateMessageRequestParams CreateRequestSamplingParams(string context, st
                             Content = [new Content() { Text = "Echo: " + message.ToString(), Type = "text" }]
                         };
                     }
+                    else if (request.Params.Name == "echoSessionId")
+                    {
+                        return new CallToolResponse()
+                        {
+                            Content = [new Content() { Text = request.Server.SessionId, Type = "text" }]
+                        };
+                    }
                     else if (request.Params.Name == "sampleLLM")
                     {
                         if (request.Params.Arguments is null ||
diff --git a/tests/ModelContextProtocol.Tests/Client/McpClientFactoryTests.cs b/tests/ModelContextProtocol.Tests/Client/McpClientFactoryTests.cs
@@ -109,6 +109,7 @@ private class NopTransport : ITransport, IClientTransport
         private readonly Channel<JsonRpcMessage> _channel = Channel.CreateUnbounded<JsonRpcMessage>();
 
         public bool IsConnected => true;
+        public string? SessionId => null;
 
         public ChannelReader<JsonRpcMessage> MessageReader => _channel.Reader;
 
diff --git a/tests/ModelContextProtocol.Tests/ClientIntegrationTests.cs b/tests/ModelContextProtocol.Tests/ClientIntegrationTests.cs
diff --git a/tests/ModelContextProtocol.Tests/Server/McpServerTests.cs b/tests/ModelContextProtocol.Tests/Server/McpServerTests.cs

Original file line number	Diff line number	Diff line change
`@@ -45,6 +45,8 @@ public AutoDetectingClientSessionTransport(SseClientTransportOptions transportOp`
`45`	`45`
`46`	`46`	`public ChannelReader<JsonRpcMessage> MessageReader => _messageChannel.Reader;`
`47`	`47`
	`48`	`+ string? ITransport.SessionId => ActiveTransport?.SessionId;`
	`49`	`+`
`48`	`50`	`/// <inheritdoc/>`
`49`	`51`	`public Task SendMessageAsync(JsonRpcMessage message, CancellationToken cancellationToken = default)`
`50`	`52`	`{`
Original file line number	Diff line number	Diff line change
`@@ -6,6 +6,7 @@ namespace ModelContextProtocol.Server;`
`6`	`6`	`internal sealed class DestinationBoundMcpServer(McpServer server, ITransport? transport) : IMcpServer`
`7`	`7`	`{`
`8`	`8`	`public string EndpointName => server.EndpointName;`
	`9`	`+ public string? SessionId => transport?.SessionId ?? server.SessionId;`
`9`	`10`	`public ClientCapabilities? ClientCapabilities => server.ClientCapabilities;`
`10`	`11`	`public Implementation? ClientInfo => server.ClientInfo;`
`11`	`12`	`public McpServerOptions ServerOptions => server.ServerOptions;`
Original file line number	Diff line number	Diff line change
`@@ -96,6 +96,10 @@ void Register<TPrimitive>(McpServerPrimitiveCollection<TPrimitive>? collection,`
`96`	`96`	`InitializeSession(transport);`
`97`	`97`	`}`
`98`	`98`
	`99`	`+ /// <inheritdoc/>`
	`100`	`+ public string? SessionId => _sessionTransport.SessionId;`
	`101`	`+`
	`102`	`+ /// <inheritdoc/>`
`99`	`103`	`public ServerCapabilities ServerCapabilities { get; } = new();`
`100`	`104`
`101`	`105`	`/// <inheritdoc />`