Merge branch 'main' into localden/experimental

dend · web-flow · commit 84a7253847bb · 2025-06-12T22:49:38.000-07:00
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;
     private int _disposed;
 
@@ -86,7 +85,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);
 
@@ -125,7 +124,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();
@@ -176,7 +175,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/CustomizableJsonStringEnumConverter.cs b/src/ModelContextProtocol.Core/CustomizableJsonStringEnumConverter.cs
@@ -5,8 +5,8 @@
 using System.Diagnostics.CodeAnalysis;
 #if !NET9_0_OR_GREATER
 using System.Reflection;
-using System.Text.Json;
 #endif
+using System.Text.Json;
 using System.Text.Json.Serialization;
 #if !NET9_0_OR_GREATER
 using ModelContextProtocol;
@@ -66,6 +66,31 @@ public override string ConvertName(string name) =>
         }
 #endif
     }
+
+    /// <summary>
+    /// A JSON converter for enums that allows customizing the serialized string value of enum members
+    /// using the <see cref="JsonStringEnumMemberNameAttribute"/>.
+    /// </summary>
+    /// <remarks>
+    /// This is a temporary workaround for lack of System.Text.Json's JsonStringEnumConverter&lt;T&gt;
+    /// 9.x support for custom enum member naming. It will be replaced by the built-in functionality
+    /// once .NET 9 is fully adopted.
+    /// </remarks>
+    [EditorBrowsable(EditorBrowsableState.Never)]
+    [RequiresUnreferencedCode("Requires unreferenced code to instantiate the generic enum converter.")]
+    [RequiresDynamicCode("Requires dynamic code to instantiate the generic enum converter.")]
+    public sealed class CustomizableJsonStringEnumConverter : JsonConverterFactory
+    {
+        /// <inheritdoc/>
+        public override bool CanConvert(Type typeToConvert) => typeToConvert.IsEnum;
+        /// <inheritdoc/>
+        public override JsonConverter? CreateConverter(Type typeToConvert, JsonSerializerOptions options)
+        {
+            Type converterType = typeof(CustomizableJsonStringEnumConverter<>).MakeGenericType(typeToConvert)!;
+            var factory = (JsonConverterFactory)Activator.CreateInstance(converterType)!;
+            return factory.CreateConverter(typeToConvert, options);
+        }
+    }
 }
 
 #if !NET9_0_OR_GREATER
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/McpJsonUtilities.cs b/src/ModelContextProtocol.Core/McpJsonUtilities.cs
@@ -3,6 +3,7 @@
 using ModelContextProtocol.Protocol;
 using ModelContextProtocol.Server;
 using System.Diagnostics.CodeAnalysis;
+using System.Reflection;
 using System.Text.Json;
 using System.Text.Json.Serialization;
 using System.Text.Json.Serialization.Metadata;
@@ -35,14 +36,22 @@ public static partial class McpJsonUtilities
     /// Creates default options to use for MCP-related serialization.
     /// </summary>
     /// <returns>The configured options.</returns>
+    [UnconditionalSuppressMessage("ReflectionAnalysis", "IL3050:RequiresDynamicCode", Justification = "Converter is guarded by IsReflectionEnabledByDefault check.")]
+    [UnconditionalSuppressMessage("Trimming", "IL2026:Members annotated with 'RequiresUnreferencedCodeAttribute' require dynamic access", Justification = "Converter is guarded by IsReflectionEnabledByDefault check.")]
     private static JsonSerializerOptions CreateDefaultOptions()
     {
         // Copy the configuration from the source generated context.
         JsonSerializerOptions options = new(JsonContext.Default.Options);
 
-        // Chain with all supported types from MEAI
+        // Chain with all supported types from MEAI.
         options.TypeInfoResolverChain.Add(AIJsonUtilities.DefaultOptions.TypeInfoResolver!);
 
+        // Add a converter for user-defined enums, if reflection is enabled by default.
+        if (JsonSerializer.IsReflectionEnabledByDefault)
+        {
+            options.Converters.Add(new CustomizableJsonStringEnumConverter());
+        }
+
         options.MakeReadOnly();
         return options;
     }
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
diff --git a/tests/ModelContextProtocol.TestSseServer/Program.cs b/tests/ModelContextProtocol.TestSseServer/Program.cs
diff --git a/tests/ModelContextProtocol.Tests/Client/McpClientFactoryTests.cs b/tests/ModelContextProtocol.Tests/Client/McpClientFactoryTests.cs
diff --git a/tests/ModelContextProtocol.Tests/ClientIntegrationTests.cs b/tests/ModelContextProtocol.Tests/ClientIntegrationTests.cs
diff --git a/tests/ModelContextProtocol.Tests/McpJsonUtilitiesTests.cs b/tests/ModelContextProtocol.Tests/McpJsonUtilitiesTests.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`	`{`