diff --git a/src/Repl.Mcp/McpServerHandler.cs b/src/Repl.Mcp/McpServerHandler.cs index 4f7bf72..738afca 100644 --- a/src/Repl.Mcp/McpServerHandler.cs +++ b/src/Repl.Mcp/McpServerHandler.cs @@ -454,9 +454,8 @@ private async ValueTask GetSnapshotAsync( { var sessionless = IsSessionlessRequest(); var snapshotVersion = Volatile.Read(ref _snapshotState).Version; - if (context.SnapshotCache is { IsStale: false } cached - && cached.Version == snapshotVersion - && cached.Sessionless == sessionless) + if (context.GetSnapshotCache(sessionless) is { IsStale: false } cached + && cached.Version == snapshotVersion) { return cached.Snapshot; } @@ -465,9 +464,8 @@ private async ValueTask GetSnapshotAsync( try { snapshotVersion = Volatile.Read(ref _snapshotState).Version; - if (context.SnapshotCache is { IsStale: false } refreshed - && refreshed.Version == snapshotVersion - && refreshed.Sessionless == sessionless) + if (context.GetSnapshotCache(sessionless) is { IsStale: false } refreshed + && refreshed.Version == snapshotVersion) { return refreshed.Snapshot; } @@ -495,7 +493,7 @@ private async ValueTask BuildOrServePreviousAsync( bool sessionless, CancellationToken cancellationToken) { - var previousSnapshot = context.SnapshotCache?.Snapshot; + var clientHasSchema = context.HasServedSnapshot; try { var built = await BuildCurrentSnapshotAsync(context, snapshotVersion, sessionless, cancellationToken) @@ -512,7 +510,7 @@ private async ValueTask BuildOrServePreviousAsync( } catch (HiddenRequiredOptionException) { - ThrowSanitizedIfAClientAlreadyHasASchema(previousSnapshot); + ThrowSanitizedIfAClientAlreadyHasASchema(clientHasSchema); throw; } catch (Exception) when (IsFallbackEligible(context, sessionless)) @@ -525,7 +523,7 @@ private async ValueTask BuildOrServePreviousAsync( // Re-read rather than reuse the filter's value: nothing holds the retraction watermark // still between the two, and a retraction that lands in between must fail closed with the // original failure rather than serve a catalog it has just withdrawn. - if (context.SnapshotCache is not { } fallback || !IsFallbackEligible(context, sessionless)) + if (context.GetSnapshotCache(sessionless) is not { } fallback || !IsFallbackEligible(context, sessionless)) { throw; } @@ -545,9 +543,9 @@ private async ValueTask BuildOrServePreviousAsync( // was meant to withhold — so it must not reach the client verbatim. A generic McpException (the // pattern this handler already uses for other client-facing failures) reports the failure without // disclosing what triggered it. - private static void ThrowSanitizedIfAClientAlreadyHasASchema(McpGeneratedSnapshot? previousSnapshot) + private static void ThrowSanitizedIfAClientAlreadyHasASchema(bool clientHasSchema) { - if (previousSnapshot is not null) + if (clientHasSchema) { throw new McpException("Tool discovery is temporarily unavailable due to a server configuration error."); } @@ -731,8 +729,7 @@ private void AttachSession(McpSessionContext context, McpServer server) /// private bool IsFallbackEligible(McpSessionContext context, bool sessionless) => !sessionless - && context.SnapshotCache is { } candidate - && candidate.Sessionless == sessionless + && context.GetSnapshotCache(sessionless) is { } candidate && Volatile.Read(ref _snapshotState).LastVisibilityRetractionVersion <= candidate.Version; internal sealed record SnapshotVersionState( diff --git a/src/Repl.Mcp/McpSessionContext.cs b/src/Repl.Mcp/McpSessionContext.cs index e90aacb..f34fac7 100644 --- a/src/Repl.Mcp/McpSessionContext.cs +++ b/src/Repl.Mcp/McpSessionContext.cs @@ -20,7 +20,12 @@ namespace Repl.Mcp; /// internal sealed class McpSessionContext : IAsyncDisposable { - private SnapshotCacheEntry? _snapshotCache; + // One entry per protocol era. A connection is served modern requests, then an initialize and legacy + // ones, and a modern request accepted before initialize can still publish after the legacy one — + // the SDK dispatches a connection's requests concurrently — so one shared slot would let it evict the + // legacy entry the availability fallback reads. + private SnapshotCacheEntry? _legacySnapshot; + private SnapshotCacheEntry? _sessionlessSnapshot; private int _compatibilityIntroServed; private readonly AsyncServiceScope? _scope; @@ -45,17 +50,21 @@ public McpSessionContext( public SemaphoreSlim SnapshotGate { get; } = new(initialCount: 1, maxCount: 1); /// - /// Cached snapshot paired with the routing version it was built at, or - /// before this session's first build. + /// This session's cached snapshot for the era, paired with the routing + /// version it was built at, or before that era's first build. /// - public SnapshotCacheEntry? SnapshotCache => Volatile.Read(ref _snapshotCache); + public SnapshotCacheEntry? GetSnapshotCache(bool sessionless) => Volatile.Read(ref Slot(sessionless)); + + /// Whether any catalog, of either era, has been published on this session. + public bool HasServedSnapshot => + Volatile.Read(ref _legacySnapshot) is not null || Volatile.Read(ref _sessionlessSnapshot) is not null; /// Publishes as current for . public void PublishSnapshot( McpServerHandler.McpGeneratedSnapshot snapshot, long version, bool sessionless) => - Volatile.Write(ref _snapshotCache, new SnapshotCacheEntry(snapshot, version, IsStale: false, sessionless)); + Volatile.Write(ref Slot(sessionless), new SnapshotCacheEntry(snapshot, version, IsStale: false, sessionless)); /// /// Publishes , built at , as serve-able but @@ -65,7 +74,10 @@ public void PublishStaleSnapshot( McpServerHandler.McpGeneratedSnapshot snapshot, long version, bool sessionless) => - Volatile.Write(ref _snapshotCache, new SnapshotCacheEntry(snapshot, version, IsStale: true, sessionless)); + Volatile.Write(ref Slot(sessionless), new SnapshotCacheEntry(snapshot, version, IsStale: true, sessionless)); + + private ref SnapshotCacheEntry? Slot(bool sessionless) => + ref sessionless ? ref _sessionlessSnapshot : ref _legacySnapshot; /// /// Claims this session's one-time compatibility-shim intro; for the first diff --git a/src/Repl.McpTests/Given_McpConcurrentSessions.cs b/src/Repl.McpTests/Given_McpConcurrentSessions.cs index f97a746..d34f7cc 100644 --- a/src/Repl.McpTests/Given_McpConcurrentSessions.cs +++ b/src/Repl.McpTests/Given_McpConcurrentSessions.cs @@ -881,6 +881,26 @@ public async Task When_OneConnectionIsServedBothEras_Then_EachGetsItsOwnCatalog( } } + [TestMethod] + [Description("Publishing one era's catalog must not evict the other's. A connection is served modern requests, then an initialize and legacy ones — and a modern request accepted before initialize can still be in flight, since the SDK dispatches one connection's requests concurrently. With one slot, that late modern publication replaced the legacy entry the availability fallback reads, so a legacy build that then failed transiently found only an opposite-era catalog and failed closed.")] + public async Task When_BothErasPublishOnOneSession_Then_EachKeepsItsOwnEntry() + { + var app = ReplApp.Create(); + var roots = new McpClientRootsService(app.Core, new McpRequestServerAccessor(), McpRootsScope.Connection); + var context = new McpSessionContext(roots, McpTestFixture.EmptyServices, scope: null); + await using var owner = context.ConfigureAwait(false); + var legacy = new McpServerHandler.McpGeneratedSnapshot(null!, [], [], []); + var modern = new McpServerHandler.McpGeneratedSnapshot(null!, [], [], []); + + context.PublishSnapshot(legacy, version: 1, sessionless: false); + context.PublishSnapshot(modern, version: 1, sessionless: true); + + var legacyEntry = context.GetSnapshotCache(sessionless: false); + legacyEntry.Should().NotBeNull("the modern publication must leave the legacy entry in place"); + legacyEntry!.Snapshot.Should().BeSameAs(legacy); + context.GetSnapshotCache(sessionless: true)!.Snapshot.Should().BeSameAs(modern); + } + private static async Task StopRawServerAsync(CancellationTokenSource cts, McpRawIo io, Task serverTask) { await cts.CancelAsync().ConfigureAwait(false);