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.

src/ModelContextProtocol.Core/McpSession.cs

@@ -115,7 +115,16 @@ internal sealed partial class McpSession : IDisposable
                 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 @@ internal sealed partial class McpSession : IDisposable
             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 @@ public abstract class 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 @@ internal sealed class StreamableHttpPostTransport(StreamableHttpServerTransport
 
         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"/>