modelcontextprotocol/csharp-sdk
Watch
1
0
Fork
You've already forked csharp-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

Add ISessionMigrationHandler (#1270)

Browse source
This commit is contained in:
Mackinnon Buck 2026-02-19 15:31:34 -08:00 • committed by GitHub
commit 27a3196c7f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 589 additions and 26 deletions
Download patch file Download diff file Expand all files Collapse all files

62
src/ModelContextProtocol.AspNetCore/ISessionMigrationHandler.cs Normal file
View file

@ -0,0 +1,62 @@
using Microsoft.AspNetCore.Http;
using ModelContextProtocol.Protocol;
namespace ModelContextProtocol.AspNetCore;
/// <summary>
/// Provides hooks for persisting and restoring MCP session initialization data,
/// enabling session migration across server instances.
/// </summary>
/// <remarks>
/// <para>
/// When an MCP server is horizontally scaled, stateful sessions are bound to a single process.
/// If that process restarts or scales down, the session is lost. By implementing this interface
/// and registering it with DI, you can persist the initialization handshake data and restore it
/// when a client reconnects to a different server instance with its existing <c>Mcp-Session-Id</c>.
/// </para>
/// <para>
/// This does <strong>not</strong> solve the session-affinity problem for in-flight server-to-client
/// requests (such as sampling or elicitation). Responses to those requests must still be routed to
/// the process that created the request. This interface only enables migration of idle sessions
/// by persisting the data established during the initialization handshake.
/// </para>
/// </remarks>
public interface ISessionMigrationHandler
{
/// <summary>
/// Called after a session has been successfully initialized via the MCP initialization handshake.
/// </summary>
/// <remarks>
/// Use this to persist the <paramref name="initializeParams"/> (which includes client capabilities,
/// client info, and protocol version) to an external store so the session can be migrated to
/// another server instance later via <see cref="AllowSessionMigrationAsync"/>.
/// </remarks>
/// <param name="context">The <see cref="HttpContext"/> for the initialization request.</param>
/// <param name="sessionId">The unique identifier for the session.</param>
/// <param name="initializeParams">The initialization parameters sent by the client during the handshake.</param>
/// <param name="cancellationToken">A cancellation token.</param>
/// <returns>A <see cref="ValueTask"/> representing the asynchronous operation.</returns>
ValueTask OnSessionInitializedAsync(HttpContext context, string sessionId, InitializeRequestParams initializeParams, CancellationToken cancellationToken);
/// <summary>
/// Called when a request arrives with an <c>Mcp-Session-Id</c> that the current server doesn't recognize.
/// </summary>
/// <remarks>
/// <para>
/// Return the original <see cref="InitializeRequestParams"/> to allow the session to be migrated
/// to this server instance, or <see langword="null"/> to reject the request (returning a 404 to the client).
/// </para>
/// <para>
/// Implementations should validate that the request is authorized, for example by checking
/// <see cref="HttpContext.User"/>, to ensure the caller is permitted to migrate the session.
/// </para>
/// </remarks>
/// <param name="context">The <see cref="HttpContext"/> for the request with the unrecognized session ID.</param>
/// <param name="sessionId">The session ID from the request that was not found on this server.</param>
/// <param name="cancellationToken">A cancellation token.</param>
/// <returns>
/// The original <see cref="InitializeRequestParams"/> if migration is allowed,
/// or <see langword="null"/> to reject the request.
/// </returns>
ValueTask<InitializeRequestParams?> AllowSessionMigrationAsync(HttpContext context, string sessionId, CancellationToken cancellationToken);
}

129
src/ModelContextProtocol.AspNetCore/StreamableHttpHandler.cs
View file

@ -7,6 +7,7 @@ using Microsoft.Extensions.Options;
using Microsoft.Net.Http.Headers;
using ModelContextProtocol.Protocol;
using ModelContextProtocol.Server;
using System.Collections.Concurrent;
using System.Security.Claims;
using System.Security.Cryptography;
using System.Text.Json.Serialization.Metadata;
@ -20,7 +21,8 @@ internal sealed class StreamableHttpHandler(
StatefulSessionManager sessionManager,
IHostApplicationLifetime hostApplicationLifetime,
IServiceProvider applicationServices,
ILoggerFactory loggerFactory)
ILoggerFactory loggerFactory,
ISessionMigrationHandler? sessionMigrationHandler = null)
{
private const string McpSessionIdHeaderName = "Mcp-Session-Id";
private const string McpProtocolVersionHeaderName = "MCP-Protocol-Version";
@ -41,6 +43,11 @@ internal sealed class StreamableHttpHandler(
private static readonly JsonTypeInfo<JsonRpcMessage> s_messageTypeInfo = GetRequiredJsonTypeInfo<JsonRpcMessage>();
private static readonly JsonTypeInfo<JsonRpcError> s_errorTypeInfo = GetRequiredJsonTypeInfo<JsonRpcError>();
private static bool AllowNewSessionForNonInitializeRequests { get; } =
AppContext.TryGetSwitch("ModelContextProtocol.AspNetCore.AllowNewSessionForNonInitializeRequests", out var enabled) && enabled;
private readonly ConcurrentDictionary<string, SemaphoreSlim> _migrationLocks = new(StringComparer.Ordinal);
public HttpServerTransportOptions HttpServerTransportOptions => httpServerTransportOptions.Value;
public async Task HandlePostRequestAsync(HttpContext context)
@ -64,14 +71,6 @@ internal sealed class StreamableHttpHandler(
return;
}
var session = await GetOrCreateSessionAsync(context);
if (session is null)
{
return;
}
await using var _ = await session.AcquireReferenceAsync(context.RequestAborted);
var message = await ReadJsonRpcMessageAsync(context);
if (message is null)
{
@ -81,6 +80,14 @@ internal sealed class StreamableHttpHandler(
return;
}
var session = await GetOrCreateSessionAsync(context, message);
if (session is null)
{
return;
}
await using var _ = await session.AcquireReferenceAsync(context.RequestAborted);
InitializeSseResponse(context);
var wroteResponse = await session.Transport.HandlePostRequestAsync(message, context.Response.Body, context.RequestAborted);
if (!wroteResponse)
@ -219,12 +226,18 @@ internal sealed class StreamableHttpHandler(
if (!sessionManager.TryGetValue(sessionId, out var session))
{
// -32001 isn't part of the MCP standard, but this is what the typescript-sdk currently does.
// One of the few other usages I found was from some Ethereum JSON-RPC documentation and this
// JSON-RPC library from Microsoft called StreamJsonRpc where it's called JsonRpcErrorCode.NoMarshaledObjectFound
// https://learn.microsoft.com/dotnet/api/streamjsonrpc.protocol.jsonrpcerrorcode?view=streamjsonrpc-2.9#fields
await WriteJsonRpcErrorAsync(context, "Session not found", StatusCodes.Status404NotFound, -32001);
return null;
// Session not found locally. Attempt migration if a handler is registered.
session = await TryMigrateSessionAsync(context, sessionId);
if (session is null)
{
// -32001 isn't part of the MCP standard, but this is what the typescript-sdk currently does.
// One of the few other usages I found was from some Ethereum JSON-RPC documentation and this
// JSON-RPC library from Microsoft called StreamJsonRpc where it's called JsonRpcErrorCode.NoMarshaledObjectFound
// https://learn.microsoft.com/dotnet/api/streamjsonrpc.protocol.jsonrpcerrorcode?view=streamjsonrpc-2.9#fields
await WriteJsonRpcErrorAsync(context, "Session not found", StatusCodes.Status404NotFound, -32001);
return null;
}
}
if (!session.HasSameUserId(context.User))
@ -240,12 +253,61 @@ internal sealed class StreamableHttpHandler(
return session;
}
private async ValueTask<StreamableHttpSession?> GetOrCreateSessionAsync(HttpContext context)
private async ValueTask<StreamableHttpSession?> TryMigrateSessionAsync(HttpContext context, string sessionId)
{
if (sessionMigrationHandler is not { } handler)
{
return null;
}
var migrationLock = _migrationLocks.GetOrAdd(sessionId, static _ => new SemaphoreSlim(1, 1));
await migrationLock.WaitAsync(context.RequestAborted);
try
{
// Re-check after acquiring the lock - another thread may have already completed migration.
if (sessionManager.TryGetValue(sessionId, out var session))
{
return session;
}
var initParams = await handler.AllowSessionMigrationAsync(context, sessionId, context.RequestAborted);
if (initParams is null)
{
return null;
}
var migratedSession = await MigrateSessionAsync(context, sessionId, initParams);
// Register the session with the session manager while still holding the lock
// so concurrent requests for the same session ID find it via sessionManager.TryGetValue.
await migratedSession.EnsureStartedAsync(context.RequestAborted);
return migratedSession;
}
finally
{
migrationLock.Release();
_migrationLocks.TryRemove(sessionId, out _);
}
}
private async ValueTask<StreamableHttpSession?> GetOrCreateSessionAsync(HttpContext context, JsonRpcMessage message)
{
var sessionId = context.Request.Headers[McpSessionIdHeaderName].ToString();
if (string.IsNullOrEmpty(sessionId))
{
// In stateful mode, only allow creating new sessions for initialize requests.
// In stateless mode, every request is independent, so we always create a new session.
if (!HttpServerTransportOptions.Stateless && !AllowNewSessionForNonInitializeRequests
&& message is not JsonRpcRequest { Method: RequestMethods.Initialize })
{
await WriteJsonRpcErrorAsync(context,
"Bad Request: A new session can only be created by an initialize request. Include a valid Mcp-Session-Id header for non-initialize requests.",
StatusCodes.Status400BadRequest);
return null;
}
return await StartNewSessionAsync(context);
}
else if (HttpServerTransportOptions.Stateless)
@ -274,7 +336,11 @@ internal sealed class StreamableHttpHandler(
SessionId = sessionId,
FlowExecutionContextFromRequests = !HttpServerTransportOptions.PerSessionExecutionContext,
EventStreamStore = HttpServerTransportOptions.EventStreamStore,
OnSessionInitialized = sessionMigrationHandler is { } handler
? (initParams, ct) => handler.OnSessionInitializedAsync(context, sessionId, initParams, ct)
: null,
};
context.Response.Headers[McpSessionIdHeaderName] = sessionId;
}
else
@ -295,11 +361,12 @@ internal sealed class StreamableHttpHandler(
private async ValueTask<StreamableHttpSession> CreateSessionAsync(
HttpContext context,
StreamableHttpServerTransport transport,
string sessionId)
string sessionId,
Action<McpServerOptions>? configureOptions = null)
{
var mcpServerServices = applicationServices;
var mcpServerOptions = mcpServerOptionsSnapshot.Value;
if (HttpServerTransportOptions.Stateless || HttpServerTransportOptions.ConfigureSessionOptions is not null)
if (HttpServerTransportOptions.Stateless || HttpServerTransportOptions.ConfigureSessionOptions is not null || configureOptions is not null)
{
mcpServerOptions = mcpServerOptionsFactory.Create(Options.DefaultName);
@ -310,6 +377,8 @@ internal sealed class StreamableHttpHandler(
mcpServerOptions.ScopeRequests = false;
}
configureOptions?.Invoke(mcpServerOptions);
if (HttpServerTransportOptions.ConfigureSessionOptions is { } configureSessionOptions)
{
await configureSessionOptions(context, mcpServerOptions, context.RequestAborted);
@ -328,6 +397,30 @@ internal sealed class StreamableHttpHandler(
return session;
}
private async ValueTask<StreamableHttpSession> MigrateSessionAsync(
HttpContext context,
string sessionId,
InitializeRequestParams initializeParams)
{
var transport = new StreamableHttpServerTransport(loggerFactory)
{
SessionId = sessionId,
FlowExecutionContextFromRequests = !HttpServerTransportOptions.PerSessionExecutionContext,
EventStreamStore = HttpServerTransportOptions.EventStreamStore,
};
// Initialize the transport with the migrated session's init params.
await transport.HandleInitializeRequestAsync(initializeParams);
context.Response.Headers[McpSessionIdHeaderName] = sessionId;
return await CreateSessionAsync(context, transport, sessionId, options =>
{
options.KnownClientInfo = initializeParams.ClientInfo;
options.KnownClientCapabilities = initializeParams.Capabilities;
});
}
private async ValueTask<ISseEventStreamReader?> GetEventStreamReaderAsync(HttpContext context, string lastEventId)
{
if (HttpServerTransportOptions.EventStreamStore is not { } eventStreamStore)

25
src/ModelContextProtocol.AspNetCore/StreamableHttpSession.cs
View file

@ -74,6 +74,31 @@ internal sealed class StreamableHttpSession(
return new UnreferenceDisposable(this);
}
/// <summary>
/// Ensures the session is registered with the session manager without acquiring a reference.
/// No-ops if the session is already started.
/// </summary>
public async ValueTask EnsureStartedAsync(CancellationToken cancellationToken)
{
bool needsStart;
lock (_stateLock)
{
needsStart = _state == SessionState.Uninitialized;
if (needsStart)
{
_state = SessionState.Started;
}
}
if (needsStart)
{
await sessionManager.StartNewSessionAsync(this, cancellationToken);
// Session is registered with 0 references (idle), so reflect that in the idle count.
sessionManager.IncrementIdleSessionCount();
}
}
public bool TryStartGetRequest() => Interlocked.Exchange(ref _getRequestStarted, 1) == 0;
public bool HasSameUserId(ClaimsPrincipal user) => userId == StreamableHttpHandler.GetUserIdClaim(user);