From d7d9eab3f85a093d561ec9922e8394cef7b07945 Mon Sep 17 00:00:00 2001 From: Eirik Tsarpalis Date: Thu, 12 Jun 2025 16:56:57 +0200 Subject: [PATCH] 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 * Update src/ModelContextProtocol.AspNetCore/StreamableHttpHandler.cs Co-authored-by: Stephen Halter * Make init callback asynchronous. * Address feedback. --------- Co-authored-by: Stephen Halter --- .../SseHandler.cs | 2 +- .../StreamableHttpHandler.cs | 21 +++++----- .../AutoDetectingClientSessionTransport.cs | 2 + .../Client/McpClient.cs | 15 ++++++++ .../StreamableHttpClientSessionTransport.cs | 7 ++-- src/ModelContextProtocol.Core/IMcpEndpoint.cs | 8 ++++ .../Protocol/ITransport.cs | 8 ++++ .../Protocol/TransportBase.cs | 3 ++ .../Server/DestinationBoundMcpServer.cs | 1 + .../Server/McpServer.cs | 4 ++ .../Server/SseResponseStreamTransport.cs | 6 ++- .../Server/StreamableHttpPostTransport.cs | 9 +++-- .../Server/StreamableHttpServerTransport.cs | 7 +++- tests/Common/Utils/TestServerTransport.cs | 2 + .../HttpServerIntegrationTests.cs | 38 +++++++++++++++++++ .../Program.cs | 17 +++++++++ .../Program.cs | 17 +++++++++ .../Client/McpClientFactoryTests.cs | 1 + .../ClientIntegrationTests.cs | 18 +++++++++ .../Server/McpServerTests.cs | 1 + 20 files changed, 166 insertions(+), 21 deletions(-) diff --git a/src/ModelContextProtocol.AspNetCore/SseHandler.cs b/src/ModelContextProtocol.AspNetCore/SseHandler.cs index 251e3bf..c5ac5a9 100644 --- a/src/ModelContextProtocol.AspNetCore/SseHandler.cs +++ b/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(sessionId, transport, userIdClaim, httpMcpServerOptions.Value.TimeProvider); diff --git a/src/ModelContextProtocol.AspNetCore/StreamableHttpHandler.cs b/src/ModelContextProtocol.AspNetCore/StreamableHttpHandler.cs index 6077efa..a38fa7c 100644 --- a/src/ModelContextProtocol.AspNetCore/StreamableHttpHandler.cs +++ b/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) diff --git a/src/ModelContextProtocol.Core/Client/AutoDetectingClientSessionTransport.cs b/src/ModelContextProtocol.Core/Client/AutoDetectingClientSessionTransport.cs index 50601f6..39ae7e8 100644 --- a/src/ModelContextProtocol.Core/Client/AutoDetectingClientSessionTransport.cs +++ b/src/ModelContextProtocol.Core/Client/AutoDetectingClientSessionTransport.cs @@ -45,6 +45,8 @@ internal sealed partial class AutoDetectingClientSessionTransport : ITransport public ChannelReader MessageReader => _messageChannel.Reader; + string? ITransport.SessionId => ActiveTransport?.SessionId; + /// public Task SendMessageAsync(JsonRpcMessage message, CancellationToken cancellationToken = default) { diff --git a/src/ModelContextProtocol.Core/Client/McpClient.cs b/src/ModelContextProtocol.Core/Client/McpClient.cs index 5c5ad2f..43639db 100644 --- 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 @@ internal sealed partial class McpClient : McpEndpoint, IMcpClient } } + /// + public string? SessionId + { + get + { + if (_sessionTransport is null) + { + throw new InvalidOperationException("Must have already initialized a session when invoking this property."); + } + + return _sessionTransport.SessionId; + } + } + /// 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 index e99aa2a..771f0cf 100644 --- 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 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); diff --git a/src/ModelContextProtocol.Core/IMcpEndpoint.cs b/src/ModelContextProtocol.Core/IMcpEndpoint.cs index a38b8c9..ea825e6 100644 --- a/src/ModelContextProtocol.Core/IMcpEndpoint.cs +++ b/src/ModelContextProtocol.Core/IMcpEndpoint.cs @@ -28,6 +28,14 @@ namespace ModelContextProtocol; /// public interface IMcpEndpoint : IAsyncDisposable { + /// Gets an identifier associated with the current MCP session. + /// + /// Typically populated in transports supporting multiple sessions such as Streamable HTTP or SSE. + /// Can return if the session hasn't initialized or if the transport doesn't + /// support multiple sessions (as is the case with STDIO). + /// + string? SessionId { get; } + /// /// Sends a JSON-RPC request to the connected endpoint and waits for a response. /// diff --git a/src/ModelContextProtocol.Core/Protocol/ITransport.cs b/src/ModelContextProtocol.Core/Protocol/ITransport.cs index 4fc36e9..e35b3a6 100644 --- a/src/ModelContextProtocol.Core/Protocol/ITransport.cs +++ b/src/ModelContextProtocol.Core/Protocol/ITransport.cs @@ -25,6 +25,14 @@ namespace ModelContextProtocol.Protocol; /// public interface ITransport : IAsyncDisposable { + /// Gets an identifier associated with the current MCP session. + /// + /// Typically populated in transports supporting multiple sessions such as Streamable HTTP or SSE. + /// Can return if the session hasn't initialized or if the transport doesn't + /// support multiple sessions (as is the case with STDIO). + /// + string? SessionId { get; } + /// /// Gets a channel reader for receiving messages from the transport. /// diff --git a/src/ModelContextProtocol.Core/Protocol/TransportBase.cs b/src/ModelContextProtocol.Core/Protocol/TransportBase.cs index 9be9c6f..a5ed1d3 100644 --- a/src/ModelContextProtocol.Core/Protocol/TransportBase.cs +++ b/src/ModelContextProtocol.Core/Protocol/TransportBase.cs @@ -59,6 +59,9 @@ public abstract partial class TransportBase : ITransport /// Gets the logger used by this transport. private protected ILogger Logger => _logger; + /// + public virtual string? SessionId { get; protected set; } + /// /// Gets the name that identifies this transport endpoint in logs. /// diff --git a/src/ModelContextProtocol.Core/Server/DestinationBoundMcpServer.cs b/src/ModelContextProtocol.Core/Server/DestinationBoundMcpServer.cs index db594da..d286d1e 100644 --- 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 index 808300f..b715edd 100644 --- a/src/ModelContextProtocol.Core/Server/McpServer.cs +++ b/src/ModelContextProtocol.Core/Server/McpServer.cs @@ -96,6 +96,10 @@ internal sealed class McpServer : McpEndpoint, IMcpServer InitializeSession(transport); } + /// + public string? SessionId => _sessionTransport.SessionId; + + /// public ServerCapabilities ServerCapabilities { get; } = new(); /// diff --git a/src/ModelContextProtocol.Core/Server/SseResponseStreamTransport.cs b/src/ModelContextProtocol.Core/Server/SseResponseStreamTransport.cs index 82562d9..438421f 100644 --- 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 . /// Defaults to "/message". /// -public sealed class SseResponseStreamTransport(Stream sseResponseStream, string? messageEndpoint = "/message") : ITransport +/// The identifier corresponding to the current MCP session. +public sealed class SseResponseStreamTransport(Stream sseResponseStream, string? messageEndpoint = "/message", string? sessionId = null) : ITransport { private readonly SseWriter _sseWriter = new(messageEndpoint); private readonly Channel _incomingChannel = Channel.CreateBounded(new BoundedChannelOptions(1) @@ -49,6 +50,9 @@ public sealed class SseResponseStreamTransport(Stream sseResponseStream, string? /// public ChannelReader MessageReader => _incomingChannel.Reader; + /// + public string? SessionId { get; } = sessionId; + /// public async ValueTask DisposeAsync() { diff --git a/src/ModelContextProtocol.Core/Server/StreamableHttpPostTransport.cs b/src/ModelContextProtocol.Core/Server/StreamableHttpPostTransport.cs index 174cf9b..343b574 100644 --- 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 MessageReader => throw new NotSupportedException("JsonRpcMessage.RelatedTransport should only be used for sending messages."); + string? ITransport.SessionId => parentTransport.SessionId; + /// /// True, if data was written to the respond body. /// False, if nothing was written because the request body did not contain any 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); } } diff --git a/src/ModelContextProtocol.Core/Server/StreamableHttpServerTransport.cs b/src/ModelContextProtocol.Core/Server/StreamableHttpServerTransport.cs index 6b024ff..1f5775e 100644 --- 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; } /// - /// Gets the initialize request if it was received by and is set to . + /// Gets or sets a callback to be invoked before handling the initialize request. /// - public InitializeRequestParams? InitializeRequest { get; internal set; } + public Func? OnInitRequestReceived { get; set; } /// public ChannelReader MessageReader => _incomingChannel.Reader; internal ChannelWriter MessageWriter => _incomingChannel.Writer; + /// + public string? SessionId { get; set; } + /// /// Handles an optional SSE GET request a client using the Streamable HTTP transport might make by /// writing any unsolicited JSON-RPC messages sent via diff --git a/tests/Common/Utils/TestServerTransport.cs b/tests/Common/Utils/TestServerTransport.cs index cd12504..1679fee 100644 --- a/tests/Common/Utils/TestServerTransport.cs +++ b/tests/Common/Utils/TestServerTransport.cs @@ -16,6 +16,8 @@ public class TestServerTransport : ITransport public Action? OnMessageSent { get; set; } + public string? SessionId => null; + public TestServerTransport() { _messageChannel = Channel.CreateUnbounded(new UnboundedChannelOptions diff --git a/tests/ModelContextProtocol.AspNetCore.Tests/HttpServerIntegrationTests.cs b/tests/ModelContextProtocol.AspNetCore.Tests/HttpServerIntegrationTests.cs index 30187fa..8fc7fb3 100644 --- a/tests/ModelContextProtocol.AspNetCore.Tests/HttpServerIntegrationTests.cs +++ b/tests/ModelContextProtocol.AspNetCore.Tests/HttpServerIntegrationTests.cs @@ -52,6 +52,15 @@ public abstract class HttpServerIntegrationTests : LoggedTest, IClassFixture(""" + { + "type": "object" + } + """, McpJsonUtilities.DefaultOptions), + }, + new Tool() { Name = "sampleLLM", Description = "Samples from an LLM using MCP's sampling feature.", @@ -170,6 +180,13 @@ internal static class Program 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 index 2e37c1d..cf078c2 100644 --- a/tests/ModelContextProtocol.TestSseServer/Program.cs +++ b/tests/ModelContextProtocol.TestSseServer/Program.cs @@ -128,6 +128,16 @@ public class Program """, McpJsonUtilities.DefaultOptions), }, new Tool() + { + Name = "echoSessionId", + Description = "Echoes the session id back to the client.", + InputSchema = JsonSerializer.Deserialize(""" + { + "type": "object" + } + """, McpJsonUtilities.DefaultOptions), + }, + new Tool() { Name = "sampleLLM", Description = "Samples from an LLM using MCP's sampling feature.", @@ -168,6 +178,13 @@ public class Program 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 index 2b8b756..9b7f445 100644 --- a/tests/ModelContextProtocol.Tests/Client/McpClientFactoryTests.cs +++ b/tests/ModelContextProtocol.Tests/Client/McpClientFactoryTests.cs @@ -109,6 +109,7 @@ public class McpClientFactoryTests private readonly Channel _channel = Channel.CreateUnbounded(); public bool IsConnected => true; + public string? SessionId => null; public ChannelReader MessageReader => _channel.Reader; diff --git a/tests/ModelContextProtocol.Tests/ClientIntegrationTests.cs b/tests/ModelContextProtocol.Tests/ClientIntegrationTests.cs index 14fd02e..dfd8d76 100644 --- a/tests/ModelContextProtocol.Tests/ClientIntegrationTests.cs +++ b/tests/ModelContextProtocol.Tests/ClientIntegrationTests.cs @@ -54,6 +54,8 @@ public partial class ClientIntegrationTests : LoggedTest, IClassFixture c.Type == "text"); + Assert.Null(textContent.Text); + } + [Theory] [MemberData(nameof(GetClients))] public async Task CallTool_Stdio_ViaAIFunction_EchoServer(string clientId) diff --git a/tests/ModelContextProtocol.Tests/Server/McpServerTests.cs b/tests/ModelContextProtocol.Tests/Server/McpServerTests.cs index 75966b6..0e1bb42 100644 --- a/tests/ModelContextProtocol.Tests/Server/McpServerTests.cs +++ b/tests/ModelContextProtocol.Tests/Server/McpServerTests.cs @@ -650,6 +650,7 @@ public class McpServerTests : LoggedTest public ValueTask DisposeAsync() => default; + public string? SessionId => throw new NotImplementedException(); public Implementation? ClientInfo => throw new NotImplementedException(); public IServiceProvider? Services => throw new NotImplementedException(); public LoggingLevel? LoggingLevel => throw new NotImplementedException();