Flow ExecutionContext with JsonRpcMessage (#616)

The primary goal of this change is to support IHttpContextAccessor in tool calls when the Streamable HTTP is in its default non-Stateless mode.

Author: halter73

src/ModelContextProtocol.AspNetCore/HttpServerTransportOptions.cs

@@ -35,6 +35,20 @@ public class HttpServerTransportOptions
     /// </remarks>
     public bool Stateless { get; set; }
 
+    /// <summary>
+    /// Gets or sets whether the server should use a single execution context for the entire session.
+    /// If <see langword="false"/>, handlers like tools get called with the <see cref="ExecutionContext"/>
+    /// belonging to the corresponding HTTP request which can change throughout the MCP session.
+    /// If <see langword="true"/>, handlers will get called with the same <see cref="ExecutionContext"/>
+    /// used to call <see cref="ConfigureSessionOptions" /> and <see cref="RunSessionHandler"/>.
+    /// </summary>
+    /// <remarks>
+    /// Enabling a per-session <see cref="ExecutionContext"/> can be useful for setting <see cref="AsyncLocal{T}"/> variables
+    /// that persist for the entire session, but it prevents you from using IHttpContextAccessor in handlers.
+    /// Defaults to <see langword="false"/>.
+    /// </remarks>
+    public bool PerSessionExecutionContext { get; set; }
+
     /// <summary>
     /// Gets or sets the duration of time the server will wait between any active requests before timing out an MCP session.
     /// </summary>

src/ModelContextProtocol.AspNetCore/StreamableHttpHandler.cs

@@ -188,6 +188,7 @@ private async ValueTask<HttpMcpSession<StreamableHttpServerTransport>> StartNewS
             transport = new()
             {
                 SessionId = sessionId,
+                FlowExecutionContextFromRequests = !HttpServerTransportOptions.PerSessionExecutionContext,
             };
             context.Response.Headers[McpSessionIdHeaderName] = sessionId;
         }

src/ModelContextProtocol.Core/McpSession.cs

@@ -115,7 +115,16 @@ public async Task ProcessMessagesAsync(CancellationToken cancellationToken)
                 LogMessageRead(EndpointName, message.GetType().Name);
 
                 // Fire and forget the message handling to avoid blocking the transport.
-                _ = ProcessMessageAsync();
+                if (message.ExecutionContext is null)
+                {
+                    _ = ProcessMessageAsync();
+                }
+                else
+                {
+                    // Flow the execution context from the HTTP request corresponding to this message if provided.
+                    ExecutionContext.Run(message.ExecutionContext, _ => _ = ProcessMessageAsync(), null);
+                }
+
                 async Task ProcessMessageAsync()
                 {
                     JsonRpcMessageWithId? messageWithId = message as JsonRpcMessageWithId;
@@ -609,9 +618,9 @@ private static void AddExceptionTags(ref TagList tags, Activity? activity, Excep
             e = ae.InnerException;
         }
 
-        int? intErrorCode = 
+        int? intErrorCode =
             (int?)((e as McpException)?.ErrorCode) is int errorCode ? errorCode :
-            e is JsonException ? (int)McpErrorCode.ParseError : 
+            e is JsonException ? (int)McpErrorCode.ParseError :
             null;
 
         string? errorType = intErrorCode?.ToString() ?? e.GetType().FullName;

src/ModelContextProtocol.Core/Protocol/JsonRpcMessage.cs

@@ -1,3 +1,4 @@
+using ModelContextProtocol.Server;
 using System.ComponentModel;
 using System.Text.Json;
 using System.Text.Json.Serialization;
@@ -38,6 +39,19 @@ private protected JsonRpcMessage()
     [JsonIgnore]
     public ITransport? RelatedTransport { get; set; }
 
+    /// <summary>
+    /// Gets or sets the <see cref="ExecutionContext"/> that should be used to run any handlers
+    /// </summary>
+    /// <remarks>
+    /// This is used to support the Streamable HTTP transport in its default stateful mode. In this mode,
+    /// the <see cref="IMcpServer"/> outlives the initial HTTP request context it was created on, and new
+    /// JSON-RPC messages can originate from future HTTP requests. This allows the transport to flow the
+    /// context with the JSON-RPC message. This is particularly useful for enabling IHttpContextAccessor
+    /// in tool calls.
+    /// </remarks>
+    [JsonIgnore]
+    public ExecutionContext? ExecutionContext { get; set; }
+
     /// <summary>
     /// Provides a <see cref="JsonConverter"/> for <see cref="JsonRpcMessage"/> messages,
     /// handling polymorphic deserialization of different message types.

src/ModelContextProtocol.Core/Server/StreamableHttpPostTransport.cs

@@ -91,6 +91,11 @@ private async ValueTask OnMessageReceivedAsync(JsonRpcMessage? message, Cancella
 
         message.RelatedTransport = this;
 
+        if (parentTransport.FlowExecutionContextFromRequests)
+        {
+            message.ExecutionContext = ExecutionContext.Capture();
+        }
+
         await parentTransport.MessageWriter.WriteAsync(message, cancellationToken).ConfigureAwait(false);
     }
 }

src/ModelContextProtocol.Core/Server/StreamableHttpServerTransport.cs

@@ -10,7 +10,7 @@ namespace ModelContextProtocol.Server;
 /// <remarks>
 /// <para>
 /// This transport provides one-way communication from server to client using the SSE protocol over HTTP,
-/// while receiving client messages through a separate mechanism. It writes messages as 
+/// while receiving client messages through a separate mechanism. It writes messages as
 /// SSE events to a response stream, typically associated with an HTTP response.
 /// </para>
 /// <para>
@@ -36,6 +36,9 @@ public sealed class StreamableHttpServerTransport : ITransport
 
     private int _getRequestStarted;
 
+    /// <inheritdoc/>
+    public string? SessionId { get; set; }
+
     /// <summary>
     /// Configures whether the transport should be in stateless mode that does not require all requests for a given session
     /// to arrive to the same ASP.NET Core application process. Unsolicited server-to-client messages are not supported in this mode,
@@ -45,6 +48,15 @@ public sealed class StreamableHttpServerTransport : ITransport
     /// </summary>
     public bool Stateless { get; init; }
 
+    /// <summary>
+    /// Gets a value indicating whether the execution context should flow from the calls to <see cref="HandlePostRequest(IDuplexPipe, CancellationToken)"/>
+    /// to the corresponding <see cref="JsonRpcMessage.ExecutionContext"/> emitted by the <see cref="MessageReader"/>.
+    /// </summary>
+    /// <remarks>
+    /// Defaults to <see langword="false"/>.
+    /// </remarks>
+    public bool FlowExecutionContextFromRequests { get; init; }
+
     /// <summary>
     /// Gets or sets a callback to be invoked before handling the initialize request.
     /// </summary>
@@ -55,9 +67,6 @@ public sealed class StreamableHttpServerTransport : ITransport
 
     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"/>

tests/ModelContextProtocol.AspNetCore.Tests/MapMcpTests.cs

@@ -52,12 +52,6 @@ public async Task MapMcp_ThrowsInvalidOperationException_IfWithHttpTransportIsNo
     [Fact]
     public async Task Can_UseIHttpContextAccessor_InTool()
     {
-        Assert.SkipWhen(UseStreamableHttp && !Stateless,
-            """
-            IHttpContextAccessor is not currently supported with non-stateless Streamable HTTP.
-            TODO: Support it in stateless mode by manually capturing and flowing execution context.
-            """);
-
         Builder.Services.AddMcpServer().WithHttpTransport(ConfigureStateless).WithTools<EchoHttpContextUserTools>();
 
         Builder.Services.AddHttpContextAccessor();

tests/ModelContextProtocol.AspNetCore.Tests/StreamableHttpServerConformanceTests.cs

@@ -387,14 +387,15 @@ public async Task Progress_IsReported_InSameSseResponseAsRpcResponse()
     }
 
     [Fact]
-    public async Task AsyncLocalSetInRunSessionHandlerCallback_Flows_ToAllToolCalls()
+    public async Task AsyncLocalSetInRunSessionHandlerCallback_Flows_ToAllToolCalls_IfPerSessionExecutionContextEnabled()
     {
         var asyncLocal = new AsyncLocal<string>();
         var totalSessionCount = 0;
 
         Builder.Services.AddMcpServer()
             .WithHttpTransport(options =>
             {
+                options.PerSessionExecutionContext = true;
                 options.RunSessionHandler = async (httpContext, mcpServer, cancellationToken) =>
                 {
                     asyncLocal.Value = $"RunSessionHandler ({totalSessionCount++})";