Add ISessionMigrationHandler (#1270)

src/ModelContextProtocol.AspNetCore/ISessionMigrationHandler.cs

@@ -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);
+}

src/ModelContextProtocol.AspNetCore/StreamableHttpHandler.cs

@@ -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)

src/ModelContextProtocol.AspNetCore/StreamableHttpSession.cs

@@ -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);
 

src/ModelContextProtocol.Core/Server/McpServerImpl.cs

@@ -75,6 +75,7 @@ internal sealed partial class McpServerImpl : McpServer
         }
 
         _clientInfo = options.KnownClientInfo;
+        _clientCapabilities = options.KnownClientCapabilities;
         UpdateEndpointNameWithClientInfo();
 
         _notificationHandlers = new();

src/ModelContextProtocol.Core/Server/McpServerOptions.cs

@@ -82,6 +82,18 @@ public sealed class McpServerOptions
     /// </remarks>
     public Implementation? KnownClientInfo { get; set; }
 
+    /// <summary>
+    /// Gets or sets preexisting knowledge about the client's capabilities to support session migration
+    /// scenarios where the client will not re-send the initialize request.
+    /// </summary>
+    /// <remarks>
+    /// <para>
+    /// When not specified, this information is sourced from the client's initialize request.
+    /// This is typically set during session migration in conjunction with <see cref="KnownClientInfo"/>.
+    /// </para>
+    /// </remarks>
+    public ClientCapabilities? KnownClientCapabilities { get; set; }
+
     /// <summary>
     /// Gets the filter collections for MCP server handlers.
     /// </summary>

src/ModelContextProtocol.Core/Server/StreamableHttpPostTransport.cs

@@ -52,7 +52,7 @@ internal sealed partial class StreamableHttpPostTransport(
             if (request.Method == RequestMethods.Initialize)
             {
                 var initializeRequest = JsonSerializer.Deserialize(request.Params, McpJsonUtilities.JsonContext.Default.InitializeRequestParams);
-                await parentTransport.HandleInitRequestAsync(initializeRequest).ConfigureAwait(false);
+                await parentTransport.HandleInitializeRequestAsync(initializeRequest).ConfigureAwait(false);
             }
         }
 

src/ModelContextProtocol.Core/Server/StreamableHttpServerTransport.cs

@@ -42,6 +42,7 @@ public sealed partial class StreamableHttpServerTransport : ITransport
     private SseEventWriter? _httpSseWriter;
     private ISseEventStreamWriter? _storeSseWriter;
     private TaskCompletionSource<bool>? _httpResponseTcs;
+    private string? _negotiatedProtocolVersion;
     private bool _getHttpRequestStarted;
     private bool _getHttpResponseCompleted;
 
@@ -82,9 +83,13 @@ public sealed partial class StreamableHttpServerTransport : ITransport
     public ISseEventStreamStore? EventStreamStore { get; init; }
 
     /// <summary>
-    /// Gets or sets the negotiated protocol version for this session.
+    /// Gets or sets an optional callback invoked after the initialization handshake completes.
     /// </summary>
-    internal string? NegotiatedProtocolVersion { get; private set; }
+    /// <remarks>
+    /// When set, this callback is invoked with the <see cref="InitializeRequestParams"/> after a successful
+    /// initialization handshake. This can be used to persist session data for cross-instance migration.
+    /// </remarks>
+    public Func<InitializeRequestParams, CancellationToken, ValueTask>? OnSessionInitialized { get; init; }
 
     /// <inheritdoc/>
     public ChannelReader<JsonRpcMessage> MessageReader => _incomingChannel.Reader;
@@ -92,12 +97,23 @@ public sealed partial class StreamableHttpServerTransport : ITransport
     internal ChannelWriter<JsonRpcMessage> MessageWriter => _incomingChannel.Writer;
 
     /// <summary>
-    /// Handles the initialize request by capturing the protocol version and invoking the user callback.
+    /// Handles initialization by capturing the negotiated protocol version and optionally invoking
+    /// <see cref="OnSessionInitialized"/> so session data can be persisted.
     /// </summary>
-    internal async ValueTask HandleInitRequestAsync(InitializeRequestParams? initParams)
+    /// <remarks>
+    /// This is called automatically when an <c>initialize</c> request is processed via
+    /// <see cref="HandlePostRequestAsync(JsonRpcMessage, Stream, CancellationToken)"/>. It can also be called
+    /// directly when restoring a migrated session with known <see cref="InitializeRequestParams"/>.
+    /// </remarks>
+    /// <param name="initParams">The initialization parameters from the client, or <see langword="null"/> if unavailable.</param>
+    public async ValueTask HandleInitializeRequestAsync(InitializeRequestParams? initParams)
     {
-        // Capture the negotiated protocol version for resumability checks
-        NegotiatedProtocolVersion = initParams?.ProtocolVersion;
+        _negotiatedProtocolVersion = initParams?.ProtocolVersion;
+
+        if (initParams is not null && OnSessionInitialized is { } callback)
+        {
+            await callback(initParams, _transportDisposedCts.Token).ConfigureAwait(false);
+        }
     }
 
     /// <summary>
@@ -266,7 +282,7 @@ public sealed partial class StreamableHttpServerTransport : ITransport
 
     internal async ValueTask<ISseEventStreamWriter?> TryCreateEventStreamAsync(string streamId, CancellationToken cancellationToken)
     {
-        if (EventStreamStore is null || !McpSessionHandler.SupportsPrimingEvent(NegotiatedProtocolVersion))
+        if (EventStreamStore is null || !McpSessionHandler.SupportsPrimingEvent(_negotiatedProtocolVersion))
         {
             return null;
         }

tests/ModelContextProtocol.AspNetCore.Tests/SessionMigrationTests.cs

@@ -0,0 +1,341 @@
+using Microsoft.AspNetCore.Builder;
+using Microsoft.AspNetCore.Http;
+using Microsoft.Extensions.DependencyInjection;
+using ModelContextProtocol.AspNetCore.Tests.Utils;
+using ModelContextProtocol.Protocol;
+using ModelContextProtocol.Server;
+using System.Net;
+using System.Net.ServerSentEvents;
+using System.Text;
+using System.Text.Json;
+using System.Text.Json.Serialization.Metadata;
+
+namespace ModelContextProtocol.AspNetCore.Tests;
+
+public class SessionMigrationTests(ITestOutputHelper outputHelper) : KestrelInMemoryTest(outputHelper), IAsyncDisposable
+{
+    private static McpServerTool[] Tools { get; } = [McpServerTool.Create(EchoAsync), McpServerTool.Create(GetClientInfoAsync)];
+
+    private WebApplication? _app;
+
+    private static string InitializeRequest => """
+        {"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{},"clientInfo":{"name":"IntegrationTestClient","version":"1.0.0"}}}
+        """;
+
+    private long _lastRequestId = 1;
+    private string MakeEchoRequest()
+    {
+        var id = Interlocked.Increment(ref _lastRequestId);
+        return $$$$"""
+            {"jsonrpc":"2.0","id":{{{{id}}}},"method":"tools/call","params":{"name":"echo","arguments":{"message":"Hello world! ({{{{id}}}})"}}}
+            """;
+    }
+
+    [Fact]
+    public async Task OnSessionInitializedAsync_IsCalled_AfterInitializeHandshake()
+    {
+        InitializeRequestParams? capturedParams = null;
+        string? capturedSessionId = null;
+
+        var handler = new TestMigrationHandler
+        {
+            OnInitialized = (context, sessionId, initParams, ct) =>
+            {
+                capturedSessionId = sessionId;
+                capturedParams = initParams;
+                return default;
+            },
+        };
+
+        await StartAsync(handler);
+
+        var sessionId = await CallInitializeAndValidateAsync();
+
+        Assert.NotNull(capturedParams);
+        Assert.Equal(sessionId, capturedSessionId);
+        Assert.Equal("IntegrationTestClient", capturedParams.ClientInfo.Name);
+        Assert.Equal("1.0.0", capturedParams.ClientInfo.Version);
+        Assert.NotNull(capturedParams.Capabilities);
+    }
+
+    [Fact]
+    public async Task AllowSessionMigrationAsync_IsCalled_WhenSessionNotFound()
+    {
+        string? requestedSessionId = null;
+
+        var handler = new TestMigrationHandler
+        {
+            OnInitialized = (_, _, _, _) => default,
+            OnMigration = (context, sessionId, ct) =>
+            {
+                requestedSessionId = sessionId;
+                return new ValueTask<InitializeRequestParams?>(new InitializeRequestParams
+                {
+                    ProtocolVersion = "2025-03-26",
+                    Capabilities = new ClientCapabilities(),
+                    ClientInfo = new Implementation { Name = "MigratedClient", Version = "2.0.0" },
+                });
+            },
+        };
+
+        await StartAsync(handler);
+
+        // Send a request with a fake session ID that the server doesn't know about.
+        SetSessionId("migratable-session-id");
+        await CallEchoAndValidateAsync();
+
+        Assert.Equal("migratable-session-id", requestedSessionId);
+
+        // Verify the migrated client info was applied to the session.
+        var clientInfo = await CallGetClientInfoAsync();
+        Assert.NotNull(clientInfo);
+        Assert.Equal("MigratedClient", clientInfo.Name);
+        Assert.Equal("2.0.0", clientInfo.Version);
+    }
+
+    [Fact]
+    public async Task MigratedSession_PreservesSessionId()
+    {
+        var handler = new TestMigrationHandler
+        {
+            OnInitialized = (_, _, _, _) => default,
+            OnMigration = (context, sessionId, ct) =>
+            {
+                return new ValueTask<InitializeRequestParams?>(new InitializeRequestParams
+                {
+                    ProtocolVersion = "2025-03-26",
+                    Capabilities = new ClientCapabilities(),
+                    ClientInfo = new Implementation { Name = "MigratedClient", Version = "2.0.0" },
+                });
+            },
+        };
+
+        await StartAsync(handler);
+
+        const string OriginalSessionId = "preserved-session-id";
+        SetSessionId(OriginalSessionId);
+
+        using var response = await HttpClient.PostAsync("", JsonContent(MakeEchoRequest()), TestContext.Current.CancellationToken);
+        Assert.Equal(HttpStatusCode.OK, response.StatusCode);
+
+        // The response should echo back the same session ID.
+        var returnedSessionId = Assert.Single(response.Headers.GetValues("mcp-session-id"));
+        Assert.Equal(OriginalSessionId, returnedSessionId);
+    }
+
+    [Fact]
+    public async Task MigratedSession_CanHandleSubsequentRequests()
+    {
+        var migrationCount = 0;
+        var handler = new TestMigrationHandler
+        {
+            OnInitialized = (_, _, _, _) => default,
+            OnMigration = (context, sessionId, ct) =>
+            {
+                Interlocked.Increment(ref migrationCount);
+                return new ValueTask<InitializeRequestParams?>(new InitializeRequestParams
+                {
+                    ProtocolVersion = "2025-03-26",
+                    Capabilities = new ClientCapabilities(),
+                    ClientInfo = new Implementation { Name = "MigratedClient", Version = "2.0.0" },
+                });
+            },
+        };
+
+        await StartAsync(handler);
+
+        SetSessionId("multi-request-session");
+
+        // First request triggers migration
+        await CallEchoAndValidateAsync();
+
+        // Second request should use the now-local session without triggering another migration.
+        await CallEchoAndValidateAsync();
+
+        Assert.Equal(1, migrationCount);
+    }
+
+    [Fact]
+    public async Task AllowSessionMigrationAsync_ReturnsNull_ResultsIn404()
+    {
+        var handler = new TestMigrationHandler
+        {
+            OnInitialized = (_, _, _, _) => default,
+            OnMigration = (context, sessionId, ct) =>
+                new ValueTask<InitializeRequestParams?>((InitializeRequestParams?)null),
+        };
+
+        await StartAsync(handler);
+
+        SetSessionId("non-migratable-session");
+
+        using var response = await HttpClient.PostAsync("", JsonContent(MakeEchoRequest()), TestContext.Current.CancellationToken);
+        Assert.Equal(HttpStatusCode.NotFound, response.StatusCode);
+    }
+
+    [Fact]
+    public async Task NoMigrationHandler_UnknownSession_Returns404()
+    {
+        // Start without any migration handler — backward compatibility.
+        await StartAsync(migrationHandler: null);
+
+        SetSessionId("unknown-session");
+
+        using var response = await HttpClient.PostAsync("", JsonContent(MakeEchoRequest()), TestContext.Current.CancellationToken);
+        Assert.Equal(HttpStatusCode.NotFound, response.StatusCode);
+    }
+
+    [Fact]
+    public async Task GetRequest_WithMigratedSession_Works()
+    {
+        var handler = new TestMigrationHandler
+        {
+            OnInitialized = (_, _, _, _) => default,
+            OnMigration = (context, sessionId, ct) =>
+            {
+                return new ValueTask<InitializeRequestParams?>(new InitializeRequestParams
+                {
+                    ProtocolVersion = "2025-03-26",
+                    Capabilities = new ClientCapabilities(),
+                    ClientInfo = new Implementation { Name = "MigratedClient", Version = "2.0.0" },
+                });
+            },
+        };
+
+        await StartAsync(handler);
+
+        // Migrate session via POST first
+        SetSessionId("get-test-session");
+        await CallEchoAndValidateAsync();
+
+        // Now the GET request should work with the migrated session
+        using var getResponse = await HttpClient.GetAsync("", HttpCompletionOption.ResponseHeadersRead, TestContext.Current.CancellationToken);
+        Assert.Equal(HttpStatusCode.OK, getResponse.StatusCode);
+    }
+
+    private async Task StartAsync(ISessionMigrationHandler? migrationHandler = null)
+    {
+        Builder.Services.AddMcpServer(options =>
+        {
+            options.ServerInfo = new Implementation
+            {
+                Name = "SessionMigrationTestServer",
+                Version = "1.0.0",
+            };
+        }).WithTools(Tools).WithHttpTransport();
+
+        if (migrationHandler is not null)
+        {
+            Builder.Services.AddSingleton(migrationHandler);
+        }
+
+        _app = Builder.Build();
+        _app.MapMcp();
+        await _app.StartAsync(TestContext.Current.CancellationToken);
+
+        HttpClient.DefaultRequestHeaders.Accept.Add(new("application/json"));
+        HttpClient.DefaultRequestHeaders.Accept.Add(new("text/event-stream"));
+    }
+
+    public async ValueTask DisposeAsync()
+    {
+        if (_app is not null)
+        {
+            await _app.DisposeAsync();
+        }
+        base.Dispose();
+    }
+
+    private static StringContent JsonContent(string json) => new(json, Encoding.UTF8, "application/json");
+    private static JsonTypeInfo<T> GetJsonTypeInfo<T>() => (JsonTypeInfo<T>)McpJsonUtilities.DefaultOptions.GetTypeInfo(typeof(T));
+
+    private async Task<string> CallInitializeAndValidateAsync()
+    {
+        HttpClient.DefaultRequestHeaders.Remove("mcp-session-id");
+        using var response = await HttpClient.PostAsync("", JsonContent(InitializeRequest), TestContext.Current.CancellationToken);
+        var rpcResponse = await AssertSingleSseResponseAsync(response);
+
+        var sessionId = Assert.Single(response.Headers.GetValues("mcp-session-id"));
+        SetSessionId(sessionId);
+        return sessionId;
+    }
+
+    private void SetSessionId(string sessionId)
+    {
+        HttpClient.DefaultRequestHeaders.Remove("mcp-session-id");
+        HttpClient.DefaultRequestHeaders.Add("mcp-session-id", sessionId);
+    }
+
+    private async Task CallEchoAndValidateAsync()
+    {
+        using var response = await HttpClient.PostAsync("", JsonContent(MakeEchoRequest()), TestContext.Current.CancellationToken);
+        var rpcResponse = await AssertSingleSseResponseAsync(response);
+
+        var callToolResult = JsonSerializer.Deserialize(rpcResponse.Result, GetJsonTypeInfo<CallToolResult>());
+        Assert.NotNull(callToolResult);
+        var content = Assert.Single(callToolResult.Content);
+        Assert.IsType<TextContentBlock>(content);
+    }
+
+    private async Task<Implementation?> CallGetClientInfoAsync()
+    {
+        var id = Interlocked.Increment(ref _lastRequestId);
+        var request = $$$$"""
+            {"jsonrpc":"2.0","id":{{{{id}}}},"method":"tools/call","params":{"name":"getClientInfo","arguments":{}}}
+            """;
+
+        using var response = await HttpClient.PostAsync("", JsonContent(request), TestContext.Current.CancellationToken);
+        var rpcResponse = await AssertSingleSseResponseAsync(response);
+
+        var callToolResult = JsonSerializer.Deserialize(rpcResponse.Result, GetJsonTypeInfo<CallToolResult>());
+        Assert.NotNull(callToolResult);
+        var textContent = Assert.IsType<TextContentBlock>(Assert.Single(callToolResult.Content));
+        return JsonSerializer.Deserialize(textContent.Text, GetJsonTypeInfo<Implementation>());
+    }
+
+    private static async Task<JsonRpcResponse> AssertSingleSseResponseAsync(HttpResponseMessage response)
+    {
+        Assert.Equal(HttpStatusCode.OK, response.StatusCode);
+        Assert.Equal("text/event-stream", response.Content.Headers.ContentType?.MediaType);
+
+        var sseItems = new List<string>();
+        var responseStream = await response.Content.ReadAsStreamAsync(TestContext.Current.CancellationToken);
+        await foreach (var sseItem in SseParser.Create(responseStream).EnumerateAsync(TestContext.Current.CancellationToken))
+        {
+            if (sseItem.EventType == "message")
+            {
+                sseItems.Add(sseItem.Data);
+            }
+        }
+
+        var data = Assert.Single(sseItems);
+        var jsonRpcResponse = JsonSerializer.Deserialize(data, GetJsonTypeInfo<JsonRpcResponse>());
+        Assert.NotNull(jsonRpcResponse);
+        return jsonRpcResponse;
+    }
+
+    [McpServerTool(Name = "echo")]
+    private static async Task<string> EchoAsync(string message)
+    {
+        await Task.Yield();
+        return message;
+    }
+
+    [McpServerTool(Name = "getClientInfo")]
+    private static string GetClientInfoAsync(McpServer server)
+    {
+        return JsonSerializer.Serialize(server.ClientInfo!, GetJsonTypeInfo<Implementation>());
+    }
+
+    private sealed class TestMigrationHandler : ISessionMigrationHandler
+    {
+        public Func<HttpContext, string, InitializeRequestParams, CancellationToken, ValueTask>? OnInitialized { get; set; }
+        public Func<HttpContext, string, CancellationToken, ValueTask<InitializeRequestParams?>>? OnMigration { get; set; }
+
+        public ValueTask OnSessionInitializedAsync(HttpContext context, string sessionId, InitializeRequestParams initializeParams, CancellationToken cancellationToken)
+            => OnInitialized?.Invoke(context, sessionId, initializeParams, cancellationToken) ?? default;
+
+        public ValueTask<InitializeRequestParams?> AllowSessionMigrationAsync(HttpContext context, string sessionId, CancellationToken cancellationToken)
+            => OnMigration?.Invoke(context, sessionId, cancellationToken) ?? new ValueTask<InitializeRequestParams?>((InitializeRequestParams?)null);
+    }
+}

tests/ModelContextProtocol.AspNetCore.Tests/StreamableHttpServerConformanceTests.cs

@@ -234,6 +234,15 @@ public class StreamableHttpServerConformanceTests(ITestOutputHelper outputHelper
         Assert.Equal(HttpStatusCode.NotFound, response.StatusCode);
     }
 
+    [Fact]
+    public async Task PostWithoutSessionId_NonInitializeRequest_Returns400()
+    {
+        await StartAsync();
+
+        using var response = await HttpClient.PostAsync("", JsonContent(ListToolsRequest), TestContext.Current.CancellationToken);
+        Assert.Equal(HttpStatusCode.BadRequest, response.StatusCode);
+    }
+
     [Fact]
     public async Task InitializeRequest_Matches_CustomRoute()
     {
@@ -728,6 +737,10 @@ public class StreamableHttpServerConformanceTests(ITestOutputHelper outputHelper
         {"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{},"clientInfo":{"name":"IntegrationTestClient","version":"1.0.0"}}}
         """;
 
+    private static string ListToolsRequest => """
+        {"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}
+        """;
+
     private long _lastRequestId = 1;
     private string EchoRequest
     {