Add DistributedCacheEventStreamStore (#1136)

Co-authored-by: Stephen Halter <halter73@gmail.com>

Directory.Packages.props

@@ -14,6 +14,8 @@
   <ItemGroup>
     <PackageVersion Include="Microsoft.Extensions.AI" Version="$(MicrosoftExtensionsVersion)" />
     <PackageVersion Include="Microsoft.Extensions.AI.Abstractions" Version="$(MicrosoftExtensionsVersion)" />
+    <PackageVersion Include="Microsoft.Extensions.Caching.Abstractions" Version="$(System10Version)" />
+    <PackageVersion Include="Microsoft.Extensions.Caching.Memory" Version="$(System10Version)" />
     <PackageVersion Include="Microsoft.Extensions.Hosting.Abstractions" Version="$(System10Version)" />
     <PackageVersion Include="Microsoft.Extensions.Logging.Abstractions" Version="$(System10Version)" />
   </ItemGroup>

src/ModelContextProtocol/ModelContextProtocol.csproj

@@ -24,6 +24,7 @@
   </ItemGroup>
 
   <ItemGroup>
+    <PackageReference Include="Microsoft.Extensions.Caching.Abstractions" />
     <PackageReference Include="Microsoft.Extensions.Hosting.Abstractions" />
   </ItemGroup>
 

src/ModelContextProtocol/Server/DistributedCacheEventIdFormatter.cs

@@ -0,0 +1,117 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+// This is a shared source file included in both ModelContextProtocol and the test project.
+// Do not reference symbols internal to the core project, as they won't be available in tests.
+#if NET
+using System.Buffers;
+using System.Buffers.Text;
+using System.Diagnostics.CodeAnalysis;
+
+#endif
+using System.Text;
+
+namespace ModelContextProtocol.Server;
+
+/// <summary>
+/// Provides methods for formatting and parsing event IDs used by <see cref="DistributedCacheEventStreamStore"/>.
+/// </summary>
+/// <remarks>
+/// Event IDs are formatted as "{base64(sessionId)}:{base64(streamId)}:{sequence}".
+/// </remarks>
+internal static class DistributedCacheEventIdFormatter
+{
+    private const char Separator = ':';
+
+    /// <summary>
+    /// Formats session ID, stream ID, and sequence number into an event ID string.
+    /// </summary>
+    public static string Format(string sessionId, string streamId, long sequence)
+    {
+        // Base64-encode session and stream IDs so the event ID can be parsed
+        // even if the original IDs contain the ':' separator character
+        var sessionBase64 = Convert.ToBase64String(Encoding.UTF8.GetBytes(sessionId));
+        var streamBase64 = Convert.ToBase64String(Encoding.UTF8.GetBytes(streamId));
+        return $"{sessionBase64}{Separator}{streamBase64}{Separator}{sequence}";
+    }
+
+    /// <summary>
+    /// Attempts to parse an event ID into its component parts.
+    /// </summary>
+    public static bool TryParse(string eventId, out string sessionId, out string streamId, out long sequence)
+    {
+        sessionId = string.Empty;
+        streamId = string.Empty;
+        sequence = 0;
+
+#if NET
+        ReadOnlySpan<char> eventIdSpan = eventId.AsSpan();
+        Span<Range> partRanges = stackalloc Range[4];
+        int rangeCount = eventIdSpan.Split(partRanges, Separator);
+        if (rangeCount != 3)
+        {
+            return false;
+        }
+
+        try
+        {
+            ReadOnlySpan<char> sessionBase64 = eventIdSpan[partRanges[0]];
+            ReadOnlySpan<char> streamBase64 = eventIdSpan[partRanges[1]];
+            ReadOnlySpan<char> sequenceSpan = eventIdSpan[partRanges[2]];
+
+            if (!TryDecodeBase64ToString(sessionBase64, out sessionId!) ||
+                !TryDecodeBase64ToString(streamBase64, out streamId!))
+            {
+                return false;
+            }
+
+            return long.TryParse(sequenceSpan, out sequence);
+        }
+        catch
+        {
+            return false;
+        }
+#else
+        var parts = eventId.Split(Separator);
+        if (parts.Length != 3)
+        {
+            return false;
+        }
+
+        try
+        {
+            sessionId = Encoding.UTF8.GetString(Convert.FromBase64String(parts[0]));
+            streamId = Encoding.UTF8.GetString(Convert.FromBase64String(parts[1]));
+            return long.TryParse(parts[2], out sequence);
+        }
+        catch
+        {
+            return false;
+        }
+#endif
+    }
+
+#if NET
+    private static bool TryDecodeBase64ToString(ReadOnlySpan<char> base64Chars, [NotNullWhen(true)] out string? result)
+    {
+        // Use a single buffer: base64 chars are ASCII (1:1 with UTF8 bytes),
+        // and decoded data is always smaller than encoded, so we can decode in-place.
+        int bufferLength = base64Chars.Length;
+        Span<byte> buffer = bufferLength <= 256
+            ? stackalloc byte[bufferLength]
+            : new byte[bufferLength];
+
+        Encoding.UTF8.GetBytes(base64Chars, buffer);
+
+        OperationStatus status = Base64.DecodeFromUtf8InPlace(buffer, out int bytesWritten);
+        if (status != OperationStatus.Done)
+        {
+            result = null;
+            return false;
+        }
+
+        result = Encoding.UTF8.GetString(buffer[..bytesWritten]);
+        return true;
+    }
+#endif
+}

src/ModelContextProtocol/Server/DistributedCacheEventStreamStore.cs

@@ -0,0 +1,401 @@
+using Microsoft.Extensions.Caching.Distributed;
+using Microsoft.Extensions.Logging;
+using Microsoft.Extensions.Logging.Abstractions;
+using ModelContextProtocol.Protocol;
+using System.Net.ServerSentEvents;
+using System.Runtime.CompilerServices;
+using System.Text.Json;
+
+namespace ModelContextProtocol.Server;
+
+/// <summary>
+/// An <see cref="ISseEventStreamStore"/> implementation backed by <see cref="IDistributedCache"/>.
+/// </summary>
+/// <remarks>
+/// <para>
+/// This implementation stores SSE events in a distributed cache, enabling resumability across
+/// multiple server instances. Event IDs are encoded with session, stream, and sequence information
+/// to allow efficient retrieval of events after a given point.
+/// </para>
+/// <para>
+/// The writer maintains in-memory state for sequence number generation, as there is guaranteed
+/// to be only one writer per stream. Readers may be created from separate processes.
+/// </para>
+/// </remarks>
+public sealed partial class DistributedCacheEventStreamStore : ISseEventStreamStore
+{
+    private readonly IDistributedCache _cache;
+    private readonly DistributedCacheEventStreamStoreOptions _options;
+    private readonly ILogger _logger;
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="DistributedCacheEventStreamStore"/> class.
+    /// </summary>
+    /// <param name="cache">The distributed cache to use for storage.</param>
+    /// <param name="options">Optional configuration options for the store.</param>
+    /// <param name="logger">Optional logger for diagnostic output.</param>
+    public DistributedCacheEventStreamStore(IDistributedCache cache, DistributedCacheEventStreamStoreOptions? options = null, ILogger<DistributedCacheEventStreamStore>? logger = null)
+    {
+        Throw.IfNull(cache);
+        _cache = cache;
+        _options = options ?? new();
+        _logger = logger ?? NullLogger<DistributedCacheEventStreamStore>.Instance;
+    }
+
+    /// <inheritdoc />
+    public ValueTask<ISseEventStreamWriter> CreateStreamAsync(SseEventStreamOptions options, CancellationToken cancellationToken = default)
+    {
+        Throw.IfNull(options);
+        LogStreamCreated(options.SessionId, options.StreamId, options.Mode);
+        var writer = new DistributedCacheEventStreamWriter(_cache, options.SessionId, options.StreamId, options.Mode, _options, _logger);
+        return new ValueTask<ISseEventStreamWriter>(writer);
+    }
+
+    /// <inheritdoc />
+    public async ValueTask<ISseEventStreamReader?> GetStreamReaderAsync(string lastEventId, CancellationToken cancellationToken = default)
+    {
+        Throw.IfNull(lastEventId);
+
+        // Parse the event ID to get session, stream, and sequence information
+        if (!DistributedCacheEventIdFormatter.TryParse(lastEventId, out var sessionId, out var streamId, out var sequence))
+        {
+            LogEventIdParsingFailed(lastEventId);
+            return null;
+        }
+
+        // Check if the stream exists by looking for its metadata
+        var metadataKey = CacheKeys.StreamMetadata(sessionId, streamId);
+        var metadataBytes = await _cache.GetAsync(metadataKey, cancellationToken).ConfigureAwait(false);
+        if (metadataBytes is null)
+        {
+            LogStreamMetadataNotFound(sessionId, streamId);
+            return null;
+        }
+
+        var metadata = JsonSerializer.Deserialize(metadataBytes, DistributedCacheEventStreamStoreJsonUtilities.StreamMetadataJsonTypeInfo);
+        if (metadata is null)
+        {
+            LogStreamMetadataDeserializationFailed(sessionId, streamId);
+            return null;
+        }
+
+        var startSequence = sequence + 1;
+        LogStreamReaderCreated(sessionId, streamId, startSequence, metadata.LastSequence);
+        return new DistributedCacheEventStreamReader(_cache, sessionId, streamId, startSequence, metadata, _options, _logger);
+    }
+
+    /// <summary>
+    /// Provides methods for generating cache keys.
+    /// </summary>
+    /// <remarks>
+    /// Cache keys are versioned to allow format changes without conflicts with existing entries.
+    /// When the cache format changes, increment <see cref="Version"/> to invalidate old entries.
+    /// </remarks>
+    internal static class CacheKeys
+    {
+        /// <summary>
+        /// The current cache key version. Increment this when changing the cache format
+        /// to ensure old entries are ignored.
+        /// </summary>
+        private const string Version = "v1";
+        private const string Prefix = $"mcp:sse:{Version}:";
+
+        public static string StreamMetadata(string sessionId, string streamId)
+        {
+            var sessionIdBase64 = Convert.ToBase64String(System.Text.Encoding.UTF8.GetBytes(sessionId));
+            var streamIdBase64 = Convert.ToBase64String(System.Text.Encoding.UTF8.GetBytes(streamId));
+            return $"{Prefix}meta:{sessionIdBase64}:{streamIdBase64}";
+        }
+
+        public static string Event(string eventId)
+            => $"{Prefix}event:{eventId}";
+    }
+
+    /// <summary>
+    /// Metadata about a stream stored in the cache.
+    /// </summary>
+    internal sealed class StreamMetadata
+    {
+        public SseEventStreamMode Mode { get; set; }
+        public bool IsCompleted { get; set; }
+        public long LastSequence { get; set; }
+    }
+
+    /// <summary>
+    /// Serialized representation of an SSE event stored in the cache.
+    /// </summary>
+    internal sealed class StoredEvent
+    {
+        public string? EventType { get; set; }
+        public string? EventId { get; set; }
+        public int? ReconnectionIntervalMs { get; set; }
+        public JsonRpcMessage? Data { get; set; }
+    }
+
+    private sealed partial class DistributedCacheEventStreamWriter : ISseEventStreamWriter
+    {
+        private readonly IDistributedCache _cache;
+        private readonly string _sessionId;
+        private readonly string _streamId;
+        private SseEventStreamMode _mode;
+        private readonly DistributedCacheEventStreamStoreOptions _options;
+        private readonly ILogger _logger;
+        private long _sequence;
+        private bool _disposed;
+
+        public DistributedCacheEventStreamWriter(
+            IDistributedCache cache,
+            string sessionId,
+            string streamId,
+            SseEventStreamMode mode,
+            DistributedCacheEventStreamStoreOptions options,
+            ILogger logger)
+        {
+            _cache = cache;
+            _sessionId = sessionId;
+            _streamId = streamId;
+            _mode = mode;
+            _options = options;
+            _logger = logger;
+        }
+
+        public async ValueTask SetModeAsync(SseEventStreamMode mode, CancellationToken cancellationToken = default)
+        {
+            ThrowIfDisposed();
+
+            LogStreamModeChanged(_sessionId, _streamId, mode);
+            _mode = mode;
+            await UpdateMetadataAsync(isCompleted: false, cancellationToken).ConfigureAwait(false);
+        }
+
+        public async ValueTask<SseItem<JsonRpcMessage?>> WriteEventAsync(SseItem<JsonRpcMessage?> sseItem, CancellationToken cancellationToken = default)
+        {
+            ThrowIfDisposed();
+
+            // Skip if already has an event ID
+            if (sseItem.EventId is not null)
+            {
+                LogEventAlreadyHasId(_sessionId, _streamId, sseItem.EventId);
+                return sseItem;
+            }
+
+            // Generate a new sequence number and event ID
+            var sequence = Interlocked.Increment(ref _sequence);
+            var eventId = DistributedCacheEventIdFormatter.Format(_sessionId, _streamId, sequence);
+            var newItem = sseItem with { EventId = eventId };
+
+            // Store the event in the cache
+            var storedEvent = new StoredEvent
+            {
+                EventType = newItem.EventType,
+                EventId = eventId,
+                ReconnectionIntervalMs = newItem.ReconnectionInterval.HasValue
+                    ? (int)newItem.ReconnectionInterval.Value.TotalMilliseconds
+                    : null,
+                Data = newItem.Data,
+            };
+
+            var eventBytes = JsonSerializer.SerializeToUtf8Bytes(storedEvent, DistributedCacheEventStreamStoreJsonUtilities.StoredEventJsonTypeInfo);
+            var eventKey = CacheKeys.Event(eventId);
+
+            await _cache.SetAsync(eventKey, eventBytes, new DistributedCacheEntryOptions
+            {
+                SlidingExpiration = _options.EventSlidingExpiration,
+                AbsoluteExpirationRelativeToNow = _options.EventAbsoluteExpiration,
+            }, cancellationToken).ConfigureAwait(false);
+
+            // Update metadata with the latest sequence
+            await UpdateMetadataAsync(isCompleted: false, cancellationToken).ConfigureAwait(false);
+
+            LogEventWritten(_sessionId, _streamId, eventId, sequence);
+            return newItem;
+        }
+
+        private async ValueTask UpdateMetadataAsync(bool isCompleted, CancellationToken cancellationToken)
+        {
+            var metadata = new StreamMetadata
+            {
+                Mode = _mode,
+                IsCompleted = isCompleted,
+                LastSequence = Interlocked.Read(ref _sequence),
+            };
+
+            var metadataBytes = JsonSerializer.SerializeToUtf8Bytes(metadata, DistributedCacheEventStreamStoreJsonUtilities.StreamMetadataJsonTypeInfo);
+            var metadataKey = CacheKeys.StreamMetadata(_sessionId, _streamId);
+
+            await _cache.SetAsync(metadataKey, metadataBytes, new DistributedCacheEntryOptions
+            {
+                SlidingExpiration = _options.MetadataSlidingExpiration,
+                AbsoluteExpirationRelativeToNow = _options.MetadataAbsoluteExpiration,
+            }, cancellationToken).ConfigureAwait(false);
+        }
+
+        private void ThrowIfDisposed()
+        {
+#if NET
+            ObjectDisposedException.ThrowIf(_disposed, this);
+#else
+            if (_disposed)
+            {
+                throw new ObjectDisposedException(nameof(DistributedCacheEventStreamWriter));
+            }
+#endif
+        }
+
+        public async ValueTask DisposeAsync()
+        {
+            if (_disposed)
+            {
+                return;
+            }
+
+            _disposed = true;
+
+            // Mark the stream as completed in the metadata
+            await UpdateMetadataAsync(isCompleted: true, CancellationToken.None).ConfigureAwait(false);
+            LogStreamWriterDisposed(_sessionId, _streamId, Interlocked.Read(ref _sequence));
+        }
+
+        [LoggerMessage(Level = LogLevel.Debug, Message = "Stream mode changed for session '{SessionId}', stream '{StreamId}' to {Mode}.")]
+        private partial void LogStreamModeChanged(string sessionId, string streamId, SseEventStreamMode mode);
+
+        [LoggerMessage(Level = LogLevel.Trace, Message = "Event already has ID '{EventId}' for session '{SessionId}', stream '{StreamId}'. Skipping ID generation.")]
+        private partial void LogEventAlreadyHasId(string sessionId, string streamId, string eventId);
+
+        [LoggerMessage(Level = LogLevel.Debug, Message = "Event written to session '{SessionId}', stream '{StreamId}' with ID '{EventId}' (sequence {Sequence}).")]
+        private partial void LogEventWritten(string sessionId, string streamId, string eventId, long sequence);
+
+        [LoggerMessage(Level = LogLevel.Information, Message = "Stream writer disposed for session '{SessionId}', stream '{StreamId}'. Total events written: {TotalEvents}.")]
+        private partial void LogStreamWriterDisposed(string sessionId, string streamId, long totalEvents);
+    }
+
+    private sealed partial class DistributedCacheEventStreamReader : ISseEventStreamReader
+    {
+        private readonly IDistributedCache _cache;
+        private readonly long _startSequence;
+        private readonly StreamMetadata _initialMetadata;
+        private readonly DistributedCacheEventStreamStoreOptions _options;
+        private readonly ILogger _logger;
+
+        public DistributedCacheEventStreamReader(
+            IDistributedCache cache,
+            string sessionId,
+            string streamId,
+            long startSequence,
+            StreamMetadata initialMetadata,
+            DistributedCacheEventStreamStoreOptions options,
+            ILogger logger)
+        {
+            _cache = cache;
+            SessionId = sessionId;
+            StreamId = streamId;
+            _startSequence = startSequence;
+            _initialMetadata = initialMetadata;
+            _options = options;
+            _logger = logger;
+        }
+
+        public string SessionId { get; }
+        public string StreamId { get; }
+
+        public async IAsyncEnumerable<SseItem<JsonRpcMessage?>> ReadEventsAsync([EnumeratorCancellation] CancellationToken cancellationToken = default)
+        {
+            // Start from the sequence after the last received event
+            var currentSequence = _startSequence;
+
+            // Use the initial metadata passed to the constructor for the first read.
+            var lastSequence = _initialMetadata.LastSequence;
+            var isCompleted = _initialMetadata.IsCompleted;
+            var mode = _initialMetadata.Mode;
+
+            LogReadingEventsStarted(SessionId, StreamId, _startSequence, lastSequence);
+
+            while (!cancellationToken.IsCancellationRequested)
+            {
+                // Read all available events from currentSequence + 1 to lastSequence
+                for (; currentSequence <= lastSequence; currentSequence++)
+                {
+                    cancellationToken.ThrowIfCancellationRequested();
+
+                    var eventId = DistributedCacheEventIdFormatter.Format(SessionId, StreamId, currentSequence);
+                    var eventKey = CacheKeys.Event(eventId);
+                    var eventBytes = await _cache.GetAsync(eventKey, cancellationToken).ConfigureAwait(false)
+                        ?? throw new McpException($"SSE event with ID '{eventId}' was not found in the cache. The event may have expired.");
+
+                    var storedEvent = JsonSerializer.Deserialize(eventBytes, DistributedCacheEventStreamStoreJsonUtilities.StoredEventJsonTypeInfo);
+                    if (storedEvent is not null)
+                    {
+                        LogEventRead(SessionId, StreamId, eventId, currentSequence);
+                        yield return new SseItem<JsonRpcMessage?>(storedEvent.Data, storedEvent.EventType)
+                        {
+                            EventId = storedEvent.EventId,
+                            ReconnectionInterval = storedEvent.ReconnectionIntervalMs.HasValue
+                                ? TimeSpan.FromMilliseconds(storedEvent.ReconnectionIntervalMs.Value)
+                                : null,
+                        };
+                    }
+                }
+
+                // If in polling mode, stop after returning currently available events
+                if (mode == SseEventStreamMode.Polling)
+                {
+                    LogReadingEventsCompletedPolling(SessionId, StreamId, currentSequence - 1);
+                    yield break;
+                }
+
+                // If the stream is completed and we've read all events, stop
+                if (isCompleted)
+                {
+                    LogReadingEventsCompletedStreamEnded(SessionId, StreamId, currentSequence - 1);
+                    yield break;
+                }
+
+                // Wait before polling again for new events
+                LogWaitingForNewEvents(SessionId, StreamId, _options.StreamReaderPollingInterval);
+                await Task.Delay(_options.StreamReaderPollingInterval, cancellationToken).ConfigureAwait(false);
+
+                // Refresh metadata to get the latest sequence and completion status
+                var metadataKey = CacheKeys.StreamMetadata(SessionId, StreamId);
+                var metadataBytes = await _cache.GetAsync(metadataKey, cancellationToken).ConfigureAwait(false)
+                    ?? throw new McpException($"Stream metadata for session '{SessionId}' and stream '{StreamId}' was not found in the cache. The metadata may have expired.");
+
+                var currentMetadata = JsonSerializer.Deserialize(metadataBytes, DistributedCacheEventStreamStoreJsonUtilities.StreamMetadataJsonTypeInfo)
+                    ?? throw new McpException($"Stream metadata for session '{SessionId}' and stream '{StreamId}' could not be deserialized.");
+
+                lastSequence = currentMetadata.LastSequence;
+                isCompleted = currentMetadata.IsCompleted;
+                mode = currentMetadata.Mode;
+            }
+        }
+
+        [LoggerMessage(Level = LogLevel.Debug, Message = "Starting to read events for session '{SessionId}', stream '{StreamId}' starting at sequence {StartSequence}. Last available sequence: {LastSequence}.")]
+        private partial void LogReadingEventsStarted(string sessionId, string streamId, long startSequence, long lastSequence);
+
+        [LoggerMessage(Level = LogLevel.Trace, Message = "Event read from session '{SessionId}', stream '{StreamId}' with ID '{EventId}' (sequence {Sequence}).")]
+        private partial void LogEventRead(string sessionId, string streamId, string eventId, long sequence);
+
+        [LoggerMessage(Level = LogLevel.Debug, Message = "Reading events completed for session '{SessionId}', stream '{StreamId}' in polling mode. Last sequence read: {LastSequence}.")]
+        private partial void LogReadingEventsCompletedPolling(string sessionId, string streamId, long lastSequence);
+
+        [LoggerMessage(Level = LogLevel.Debug, Message = "Reading events completed for session '{SessionId}', stream '{StreamId}' as stream has ended. Last sequence read: {LastSequence}.")]
+        private partial void LogReadingEventsCompletedStreamEnded(string sessionId, string streamId, long lastSequence);
+
+        [LoggerMessage(Level = LogLevel.Trace, Message = "Waiting for new events on session '{SessionId}', stream '{StreamId}'. Polling interval: {PollingInterval}.")]
+        private partial void LogWaitingForNewEvents(string sessionId, string streamId, TimeSpan pollingInterval);
+    }
+
+    [LoggerMessage(Level = LogLevel.Information, Message = "Stream created for session '{SessionId}', stream '{StreamId}' with mode {Mode}.")]
+    private partial void LogStreamCreated(string sessionId, string streamId, SseEventStreamMode mode);
+
+    [LoggerMessage(Level = LogLevel.Debug, Message = "Stream reader created for session '{SessionId}', stream '{StreamId}' starting at sequence {StartSequence}. Last available sequence: {LastSequence}.")]
+    private partial void LogStreamReaderCreated(string sessionId, string streamId, long startSequence, long lastSequence);
+
+    [LoggerMessage(Level = LogLevel.Warning, Message = "Failed to parse event ID '{EventId}'. Unable to create stream reader.")]
+    private partial void LogEventIdParsingFailed(string eventId);
+
+    [LoggerMessage(Level = LogLevel.Debug, Message = "Stream metadata not found for session '{SessionId}', stream '{StreamId}'.")]
+    private partial void LogStreamMetadataNotFound(string sessionId, string streamId);
+
+    [LoggerMessage(Level = LogLevel.Warning, Message = "Failed to deserialize stream metadata for session '{SessionId}', stream '{StreamId}'.")]
+    private partial void LogStreamMetadataDeserializationFailed(string sessionId, string streamId);
+}

src/ModelContextProtocol/Server/DistributedCacheEventStreamStoreJsonContext.cs

@@ -0,0 +1,60 @@
+using System.Text.Json;
+using System.Text.Json.Serialization;
+using System.Text.Json.Serialization.Metadata;
+
+namespace ModelContextProtocol.Server;
+
+/// <summary>
+/// Provides JSON serialization utilities for <see cref="DistributedCacheEventStreamStore"/> types.
+/// </summary>
+/// <remarks>
+/// This class provides source-generated serialization for the internal types used by
+/// <see cref="DistributedCacheEventStreamStore"/> to persist SSE events and stream metadata.
+/// It combines with <see cref="McpJsonUtilities.DefaultOptions"/> to support the full object
+/// graph including <see cref="Protocol.JsonRpcMessage"/> and its derived types.
+/// </remarks>
+internal static partial class DistributedCacheEventStreamStoreJsonUtilities
+{
+    /// <summary>
+    /// Gets the <see cref="JsonSerializerOptions"/> instance for serializing distributed cache event stream store types.
+    /// </summary>
+    /// <remarks>
+    /// This options instance combines the source-generated context for <see cref="DistributedCacheEventStreamStore.StreamMetadata"/>
+    /// and <see cref="DistributedCacheEventStreamStore.StoredEvent"/> with <see cref="McpJsonUtilities.DefaultOptions"/>
+    /// to support the full object graph including <see cref="Protocol.JsonRpcMessage"/> types.
+    /// </remarks>
+    public static JsonSerializerOptions DefaultOptions { get; } = CreateDefaultOptions();
+
+    /// <summary>
+    /// Gets the <see cref="JsonTypeInfo{T}"/> for <see cref="DistributedCacheEventStreamStore.StreamMetadata"/>.
+    /// </summary>
+    public static JsonTypeInfo<DistributedCacheEventStreamStore.StreamMetadata> StreamMetadataJsonTypeInfo { get; } =
+        (JsonTypeInfo<DistributedCacheEventStreamStore.StreamMetadata>)DefaultOptions.GetTypeInfo(typeof(DistributedCacheEventStreamStore.StreamMetadata));
+
+    /// <summary>
+    /// Gets the <see cref="JsonTypeInfo{T}"/> for <see cref="DistributedCacheEventStreamStore.StoredEvent"/>.
+    /// </summary>
+    public static JsonTypeInfo<DistributedCacheEventStreamStore.StoredEvent> StoredEventJsonTypeInfo { get; } =
+        (JsonTypeInfo<DistributedCacheEventStreamStore.StoredEvent>)DefaultOptions.GetTypeInfo(typeof(DistributedCacheEventStreamStore.StoredEvent));
+
+    private static JsonSerializerOptions CreateDefaultOptions()
+    {
+        // Copy the configuration from McpJsonUtilities.DefaultOptions.
+        var options = new JsonSerializerOptions(McpJsonUtilities.DefaultOptions);
+
+        // Add our source-generated context for StreamMetadata and StoredEvent.
+        options.TypeInfoResolverChain.Insert(0, JsonContext.Default);
+
+        options.MakeReadOnly();
+        return options;
+    }
+
+    [JsonSourceGenerationOptions(
+        JsonSerializerDefaults.Web,
+        NumberHandling = JsonNumberHandling.AllowReadingFromString,
+        DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull,
+        GenerationMode = JsonSourceGenerationMode.Metadata)]
+    [JsonSerializable(typeof(DistributedCacheEventStreamStore.StreamMetadata))]
+    [JsonSerializable(typeof(DistributedCacheEventStreamStore.StoredEvent))]
+    private sealed partial class JsonContext : JsonSerializerContext;
+}

src/ModelContextProtocol/Server/DistributedCacheEventStreamStoreOptions.cs

@@ -0,0 +1,51 @@
+namespace ModelContextProtocol.Server;
+
+/// <summary>
+/// Configuration options for <see cref="DistributedCacheEventStreamStore"/>.
+/// </summary>
+public sealed class DistributedCacheEventStreamStoreOptions
+{
+    /// <summary>
+    /// Gets or sets the sliding expiration for individual events in the cache.
+    /// </summary>
+    /// <remarks>
+    /// Events are refreshed on each access. If an event is not accessed within this
+    /// time period, it may be evicted from the cache.
+    /// </remarks>
+    public TimeSpan? EventSlidingExpiration { get; set; } = TimeSpan.FromMinutes(30);
+
+    /// <summary>
+    /// Gets or sets the absolute expiration for individual events in the cache.
+    /// </summary>
+    /// <remarks>
+    /// Events will be evicted from the cache after this time period, regardless of access.
+    /// </remarks>
+    public TimeSpan? EventAbsoluteExpiration { get; set; } = TimeSpan.FromHours(2);
+
+    /// <summary>
+    /// Gets or sets the sliding expiration for stream metadata in the cache.
+    /// </summary>
+    /// <remarks>
+    /// Stream metadata includes mode and completion status. This should typically be
+    /// set to a longer duration than event expiration to allow for resumability.
+    /// </remarks>
+    public TimeSpan? MetadataSlidingExpiration { get; set; } = TimeSpan.FromHours(1);
+
+    /// <summary>
+    /// Gets or sets the absolute expiration for stream metadata in the cache.
+    /// </summary>
+    /// <remarks>
+    /// Stream metadata will be evicted from the cache after this time period, regardless of access.
+    /// </remarks>
+    public TimeSpan? MetadataAbsoluteExpiration { get; set; } = TimeSpan.FromHours(4);
+
+    /// <summary>
+    /// Gets or sets the interval between polling attempts when a stream reader is waiting for new events
+    /// in the default <see cref="SseEventStreamMode.Streaming"/> mode.
+    /// </summary>
+    /// <remarks>
+    /// This only affects stream readers. A shorter interval provides lower latency for new events
+    /// but increases cache access frequency.
+    /// </remarks>
+    public TimeSpan StreamReaderPollingInterval { get; set; } = TimeSpan.FromSeconds(1);
+}

tests/ModelContextProtocol.AspNetCore.Tests/DistributedCacheResumabilityIntegrationTests.cs

@@ -0,0 +1,49 @@
+using Microsoft.Extensions.Caching.Distributed;
+using Microsoft.Extensions.Caching.Memory;
+using Microsoft.Extensions.Logging;
+using Microsoft.Extensions.Options;
+using ModelContextProtocol.Server;
+
+namespace ModelContextProtocol.AspNetCore.Tests;
+
+/// <summary>
+/// Integration tests for SSE resumability using <see cref="DistributedCacheEventStreamStore"/>.
+/// </summary>
+/// <remarks>
+/// <para>
+/// This class runs the same resumability tests as <see cref="ResumabilityIntegrationTests"/> but
+/// using the production <see cref="DistributedCacheEventStreamStore"/> implementation backed by
+/// an in-memory distributed cache.
+/// </para>
+/// <para>
+/// These tests verify that the distributed cache implementation correctly stores and retrieves
+/// events for resumability, including across simulated disconnections.
+/// </para>
+/// </remarks>
+public class DistributedCacheResumabilityIntegrationTests(ITestOutputHelper testOutputHelper) : ResumabilityIntegrationTestsBase(testOutputHelper)
+{
+    private MemoryDistributedCache? _cache;
+
+    /// <inheritdoc />
+    protected override ValueTask<ISseEventStreamStore> CreateEventStreamStoreAsync()
+    {
+        // Create a new in-memory distributed cache for each test
+        _cache = new MemoryDistributedCache(Options.Create(new MemoryDistributedCacheOptions()));
+
+        // Configure the store with shorter expiration times suitable for testing
+        var options = new DistributedCacheEventStreamStoreOptions
+        {
+            // Use shorter polling interval for faster test execution
+            StreamReaderPollingInterval = TimeSpan.FromMilliseconds(50),
+
+            // Use shorter expiration times for tests
+            EventSlidingExpiration = TimeSpan.FromMinutes(5),
+            EventAbsoluteExpiration = TimeSpan.FromMinutes(10),
+            MetadataSlidingExpiration = TimeSpan.FromMinutes(5),
+            MetadataAbsoluteExpiration = TimeSpan.FromMinutes(10),
+        };
+
+        var store = new DistributedCacheEventStreamStore(_cache, options, LoggerFactory.CreateLogger<DistributedCacheEventStreamStore>());
+        return new ValueTask<ISseEventStreamStore>(store);
+    }
+}

tests/ModelContextProtocol.AspNetCore.Tests/ModelContextProtocol.AspNetCore.Tests.csproj

@@ -31,6 +31,7 @@
       <PrivateAssets>all</PrivateAssets>
     </PackageReference>
     <PackageReference Include="Microsoft.AspNetCore.Authentication.JwtBearer" />
+    <PackageReference Include="Microsoft.Extensions.Caching.Memory" />
     <PackageReference Include="Microsoft.Extensions.AI" />
     <PackageReference Include="Microsoft.Extensions.AI.OpenAI" />
     <PackageReference Include="Microsoft.Extensions.Logging" />

tests/ModelContextProtocol.AspNetCore.Tests/ResumabilityIntegrationTests.cs

@@ -1,70 +1,73 @@
-using System.ComponentModel;
-using System.Diagnostics;
-using System.Net;
 using System.Net.ServerSentEvents;
-using System.Text;
-using System.Text.Json.Nodes;
-using Microsoft.AspNetCore.Builder;
-using Microsoft.Extensions.DependencyInjection;
 using ModelContextProtocol.AspNetCore.Tests.Utils;
-using ModelContextProtocol.Client;
 using ModelContextProtocol.Protocol;
 using ModelContextProtocol.Server;
-using ModelContextProtocol.Tests.Utils;
 
 namespace ModelContextProtocol.AspNetCore.Tests;
 
 /// <summary>
-/// Integration tests for SSE resumability with full client-server flow.
+/// Integration tests for SSE resumability using the in-memory <see cref="TestSseEventStreamStore"/>.
-/// These tests use McpClient for end-to-end testing and only use raw HTTP
-/// for SSE format verification where McpClient abstracts away the details.
 /// </summary>
-public class ResumabilityIntegrationTests(ITestOutputHelper testOutputHelper) : KestrelInMemoryTest(testOutputHelper)
+/// <remarks>
+/// <para>
+/// This class extends <see cref="ResumabilityIntegrationTestsBase"/> to add assertions specific to
+/// <see cref="TestSseEventStreamStore"/>, such as verifying event counts and stored data.
+/// </para>
+/// <para>
+/// Tests that don't require <see cref="TestSseEventStreamStore"/>-specific assertions are inherited
+/// from the base class.
+/// </para>
+/// </remarks>
+public class ResumabilityIntegrationTests(ITestOutputHelper testOutputHelper) : ResumabilityIntegrationTestsBase(testOutputHelper)
 {
-    private const string InitializeRequest = """
+    /// <summary>
-        {"jsonrpc":"2.0","id":"1","method":"initialize","params":{"protocolVersion":"2025-11-25","capabilities":{},"clientInfo":{"name":"TestClient","version":"1.0.0"}}}
+    /// Gets the test event stream store, cast to the concrete type for test-specific assertions.
-        """;
+    /// </summary>
+    private TestSseEventStreamStore TestEventStreamStore => (TestSseEventStreamStore)EventStreamStore!;
+
+    /// <inheritdoc />
+    protected override ValueTask<ISseEventStreamStore> CreateEventStreamStoreAsync()
+        => new(new TestSseEventStreamStore());
 
     [Fact]
-    public async Task Server_StoresEvents_WhenEventStoreConfigured()
+    public override async Task Server_StoresEvents_WhenEventStoreConfigured()
     {
-        // Arrange
+        await base.Server_StoresEvents_WhenEventStoreConfigured();
-        var eventStreamStore = new TestSseEventStreamStore();
-        await using var app = await CreateServerAsync(eventStreamStore);
-        await using var client = await ConnectClientAsync();
 
-        // Act - Make a tool call which generates events
+        // Additional assertion: verify events were actually stored
-        var result = await client.CallToolAsync("echo",
+        Assert.True(TestEventStreamStore.StoreEventCallCount > 0, "Expected events to be stored when EventStore is configured");
-            new Dictionary<string, object?> { ["message"] = "test" },
+    }
-            cancellationToken: TestContext.Current.CancellationToken);
 
-        // Assert - Events were stored
+    [Fact]
-        Assert.NotNull(result);
+    public override async Task Client_CanMakeMultipleRequests_WithResumabilityEnabled()
-        Assert.True(eventStreamStore.StoreEventCallCount > 0, "Expected events to be stored when EventStore is configured");
+    {
+        await base.Client_CanMakeMultipleRequests_WithResumabilityEnabled();
+
+        // Additional assertion: verify events were stored for each request
+        Assert.True(TestEventStreamStore.StoreEventCallCount >= 5, "Expected events to be stored for each request");
     }
 
     [Fact]
     public async Task Server_StoresMultipleEvents_ForMultipleToolCalls()
     {
         // Arrange
-        var eventStreamStore = new TestSseEventStreamStore();
+        await using var app = await CreateServerAsync();
-        await using var app = await CreateServerAsync(eventStreamStore);
         await using var client = await ConnectClientAsync();
 
         // Act - Make multiple tool calls
-        var initialCount = eventStreamStore.StoreEventCallCount;
+        var initialCount = TestEventStreamStore.StoreEventCallCount;
 
         await client.CallToolAsync("echo",
             new Dictionary<string, object?> { ["message"] = "test1" },
             cancellationToken: TestContext.Current.CancellationToken);
 
-        var countAfterFirst = eventStreamStore.StoreEventCallCount;
+        var countAfterFirst = TestEventStreamStore.StoreEventCallCount;
 
         await client.CallToolAsync("echo",
             new Dictionary<string, object?> { ["message"] = "test2" },
             cancellationToken: TestContext.Current.CancellationToken);
 
-        var countAfterSecond = eventStreamStore.StoreEventCallCount;
+        var countAfterSecond = TestEventStreamStore.StoreEventCallCount;
 
         // Assert - More events were stored for each call
         Assert.True(countAfterFirst > initialCount, "Expected more events after first call");
@@ -72,61 +75,20 @@ public class ResumabilityIntegrationTests(ITestOutputHelper testOutputHelper) :
     }
 
     [Fact]
-    public async Task Client_CanMakeMultipleRequests_WithResumabilityEnabled()
+    public override async Task EnablePollingAsync_SendsSseItemWithRetryField()
     {
-        // Arrange
+        await base.EnablePollingAsync_SendsSseItemWithRetryField();
-        var eventStreamStore = new TestSseEventStreamStore();
-        await using var app = await CreateServerAsync(eventStreamStore);
-        await using var client = await ConnectClientAsync();
 
-        // Act - Make many requests to verify stability
+        // Additional assertion: verify that the event store received the retry interval
-        for (int i = 0; i < 5; i++)
+        var expectedRetryInterval = TimeSpan.FromSeconds(5);
-        {
+        Assert.Contains(expectedRetryInterval, TestEventStreamStore.StoredReconnectionIntervals);
-            var result = await client.CallToolAsync("echo",
-                new Dictionary<string, object?> { ["message"] = $"test{i}" },
-                cancellationToken: TestContext.Current.CancellationToken);
-
-            var textContent = Assert.Single(result.Content.OfType<TextContentBlock>());
-            Assert.Equal($"Echo: test{i}", textContent.Text);
-        }
-
-        // Assert - All requests succeeded and events were stored
-        Assert.True(eventStreamStore.StoreEventCallCount >= 5, "Expected events to be stored for each request");
-    }
-
-    [Fact]
-    public async Task Ping_WorksWithResumabilityEnabled()
-    {
-        // Arrange
-        var eventStreamStore = new TestSseEventStreamStore();
-        await using var app = await CreateServerAsync(eventStreamStore);
-        await using var client = await ConnectClientAsync();
-
-        // Act & Assert - Ping should work
-        await client.PingAsync(cancellationToken: TestContext.Current.CancellationToken);
-    }
-
-    [Fact]
-    public async Task ListTools_WorksWithResumabilityEnabled()
-    {
-        // Arrange
-        var eventStreamStore = new TestSseEventStreamStore();
-        await using var app = await CreateServerAsync(eventStreamStore);
-        await using var client = await ConnectClientAsync();
-
-        // Act
-        var tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
-
-        // Assert
-        Assert.NotNull(tools);
-        Assert.Single(tools);
     }
 
     [Fact]
     public async Task Server_WithoutEventStore_DoesNotIncludeEventId()
     {
-        // Arrange - Server without event store
+        // Arrange - Server without event store (pass null explicitly)
-        await using var app = await CreateServerAsync();
+        await using var app = await CreateServerAsync(eventStreamStore: null);
 
         // Act
         var sseResponse = await SendInitializeAndReadSseResponseAsync(InitializeRequest);
@@ -139,8 +101,7 @@ public class ResumabilityIntegrationTests(ITestOutputHelper testOutputHelper) :
     public async Task Server_DoesNotSendPrimingEvents_ToOlderProtocolVersionClients()
     {
         // Arrange - Server with resumability enabled
-        var eventStreamStore = new TestSseEventStreamStore();
+        await using var app = await CreateServerAsync();
-        await using var app = await CreateServerAsync(eventStreamStore);
 
         // Use an older protocol version that doesn't support resumability
         const string OldProtocolInitRequest = """
@@ -153,358 +114,7 @@ public class ResumabilityIntegrationTests(ITestOutputHelper testOutputHelper) :
         Assert.True(sseResponse.LastEventId is null, "Old protocol clients should not receive event IDs");
 
         // Event store should not have been called for old clients
-        Assert.Equal(0, eventStreamStore.StoreEventCallCount);
+        Assert.Equal(0, TestEventStreamStore.StoreEventCallCount);
-    }
-
-    [Fact]
-    public async Task Client_CanPollResponse_FromServer()
-    {
-        const string ProgressToolName = "progress_tool";
-        var clientReceivedInitialValueTcs = new TaskCompletionSource();
-        var clientReceivedPolledValueTcs = new TaskCompletionSource();
-        var progressTool = McpServerTool.Create(async (RequestContext<CallToolRequestParams> context, IProgress<ProgressNotificationValue> progress) =>
-        {
-            progress.Report(new() { Progress = 0, Message = "Initial value" });
-
-            await clientReceivedInitialValueTcs.Task.WaitAsync(TestContext.Current.CancellationToken);
-
-            await context.EnablePollingAsync(retryInterval: TimeSpan.FromSeconds(1));
-
-            progress.Report(new() { Progress = 50, Message = "Polled value" });
-
-            await clientReceivedPolledValueTcs.Task.WaitAsync(TestContext.Current.CancellationToken); ;
-
-            return "Complete";
-        }, options: new() { Name = ProgressToolName });
-        var eventStreamStore = new TestSseEventStreamStore();
-        await using var app = await CreateServerAsync(eventStreamStore, configureServer: builder =>
-        {
-            builder.WithTools([progressTool]);
-        });
-        await using var client = await ConnectClientAsync();
-
-        var progressHandler = new Progress<ProgressNotificationValue>(value =>
-        {
-            switch (value.Message)
-            {
-                case "Initial value":
-                    Assert.True(clientReceivedInitialValueTcs.TrySetResult(), "Received the initial value more than once.");
-                    break;
-                case "Polled value":
-                    Assert.True(clientReceivedPolledValueTcs.TrySetResult(), "Received the polled value more than once.");
-                    break;
-                default:
-                    throw new UnreachableException($"Unknown progress message '{value.Message}'");
-            }
-        });
-
-        var result = await client.CallToolAsync(ProgressToolName, progress: progressHandler, cancellationToken: TestContext.Current.CancellationToken);
-
-        Assert.False(result.IsError is true);
-        Assert.Equal("Complete", result.Content.OfType<TextContentBlock>().Single().Text);
-    }
-
-    [Fact]
-    public async Task Client_CanResumePostResponseStream_AfterDisconnection()
-    {
-        using var faultingStreamHandler = new FaultingStreamHandler()
-        {
-            InnerHandler = SocketsHttpHandler,
-        };
-
-        HttpClient = new(faultingStreamHandler);
-        ConfigureHttpClient(HttpClient);
-
-        const string ProgressToolName = "progress_tool";
-        const string InitialMessage = "Initial notification";
-        const string ReplayedMessage = "Replayed notification";
-        const string ResultMessage = "Complete";
-
-        var clientReceivedInitialValueTcs = new TaskCompletionSource();
-        var clientReceivedReconnectValueTcs = new TaskCompletionSource();
-        var progressTool = McpServerTool.Create(async (RequestContext<CallToolRequestParams> context, IProgress<ProgressNotificationValue> progress, CancellationToken cancellationToken) =>
-        {
-            progress.Report(new() { Progress = 0, Message = InitialMessage });
-
-            // Make sure the client receives one message before we disconnect.
-            await clientReceivedInitialValueTcs.Task.WaitAsync(TestContext.Current.CancellationToken);
-
-            // Simulate a network disconnection by faulting the response stream.
-            var reconnectAttempt = await faultingStreamHandler.TriggerFaultAsync(TestContext.Current.CancellationToken);
-
-            // Send another message that the client should receive after reconnecting.
-            progress.Report(new() { Progress = 50, Message = ReplayedMessage });
-
-            reconnectAttempt.Continue();
-
-            // Wait for the client to receive the message via replay.
-            await clientReceivedReconnectValueTcs.Task.WaitAsync(TestContext.Current.CancellationToken);
-
-            // Return the final result with the client still connected.
-            return ResultMessage;
-        }, options: new() { Name = ProgressToolName });
-        var eventStreamStore = new TestSseEventStreamStore();
-        await using var app = await CreateServerAsync(eventStreamStore, configureServer: builder =>
-        {
-            builder.WithTools([progressTool]);
-        });
-        await using var client = await ConnectClientAsync();
-
-        var initialNotificationReceivedCount = 0;
-        var replayedNotificationReceivedCount = 0;
-        var progressHandler = new Progress<ProgressNotificationValue>(value =>
-        {
-            switch (value.Message)
-            {
-                case InitialMessage:
-                    initialNotificationReceivedCount++;
-                    clientReceivedInitialValueTcs.TrySetResult();
-                    break;
-                case ReplayedMessage:
-                    replayedNotificationReceivedCount++;
-                    clientReceivedReconnectValueTcs.TrySetResult();
-                    break;
-                default:
-                    throw new UnreachableException($"Unknown progress message '{value.Message}'");
-            }
-        });
-
-        var result = await client.CallToolAsync(ProgressToolName, progress: progressHandler, cancellationToken: TestContext.Current.CancellationToken);
-
-        Assert.False(result.IsError is true);
-        Assert.Equal(1, initialNotificationReceivedCount);
-        Assert.Equal(1, replayedNotificationReceivedCount);
-        Assert.Equal(ResultMessage, result.Content.OfType<TextContentBlock>().Single().Text);
-    }
-
-    [Fact]
-    public async Task Client_CanResumeUnsolicitedMessageStream_AfterDisconnection()
-    {
-        var timeout = TestConstants.DefaultTimeout;
-        using var faultingStreamHandler = new FaultingStreamHandler()
-        {
-            InnerHandler = SocketsHttpHandler,
-        };
-
-        HttpClient = new(faultingStreamHandler);
-        ConfigureHttpClient(HttpClient);
-
-        var eventStreamStore = new TestSseEventStreamStore();
-
-        // Capture the server instance via RunSessionHandler
-        var serverTcs = new TaskCompletionSource<McpServer>();
-
-        await using var app = await CreateServerAsync(eventStreamStore, configureTransport: options =>
-        {
-            options.RunSessionHandler = (httpContext, mcpServer, cancellationToken) =>
-            {
-                serverTcs.TrySetResult(mcpServer);
-                return mcpServer.RunAsync(cancellationToken);
-            };
-        });
-
-        await using var client = await ConnectClientAsync();
-
-        // Get the server instance
-        var server = await serverTcs.Task.WaitAsync(timeout, TestContext.Current.CancellationToken);
-
-        // Set up notification tracking with unique messages
-        var clientReceivedInitialNotificationTcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
-        var clientReceivedReplayedNotificationTcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
-        var clientReceivedReconnectNotificationTcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
-
-        const string CustomNotificationMethod = "test/custom_notification";
-        const string InitialMessage = "Initial notification";
-        const string ReplayedMessage = "Replayed notification";
-        const string ReconnectMessage = "Reconnect notification";
-
-        var initialNotificationReceivedCount = 0;
-        var replayedNotificationReceivedCount = 0;
-        var reconnectNotificationReceivedCount = 0;
-
-        await using var _ = client.RegisterNotificationHandler(CustomNotificationMethod, (notification, cancellationToken) =>
-        {
-            var message = notification.Params?["message"]?.GetValue<string>();
-            switch (message)
-            {
-                case InitialMessage:
-                    initialNotificationReceivedCount++;
-                    clientReceivedInitialNotificationTcs.TrySetResult();
-                    break;
-                case ReplayedMessage:
-                    replayedNotificationReceivedCount++;
-                    clientReceivedReplayedNotificationTcs.TrySetResult();
-                    break;
-                case ReconnectMessage:
-                    reconnectNotificationReceivedCount++;
-                    clientReceivedReconnectNotificationTcs.TrySetResult();
-                    break;
-                default:
-                    throw new UnreachableException($"Unknown notification message '{message}'");
-            }
-            return default;
-        });
-
-        // Wait for the client's unsolicited message stream to be established before sending notifications
-        await faultingStreamHandler.WaitForUnsolicitedMessageStreamAsync(TestContext.Current.CancellationToken);
-
-        // Send a custom notification to the client on the unsolicited message stream
-        await server.SendNotificationAsync(CustomNotificationMethod, new JsonObject { ["message"] = InitialMessage }, cancellationToken: TestContext.Current.CancellationToken);
-
-        // Wait for client to receive the first notification
-        await clientReceivedInitialNotificationTcs.Task.WaitAsync(timeout, TestContext.Current.CancellationToken);
-
-        // Fault the unsolicited message stream (GET SSE)
-        var reconnectAttempt = await faultingStreamHandler.TriggerFaultAsync(TestContext.Current.CancellationToken);
-
-        // Send another notification while the client is disconnected - this should be stored
-        await server.SendNotificationAsync(CustomNotificationMethod, new JsonObject { ["message"] = ReplayedMessage }, cancellationToken: TestContext.Current.CancellationToken);
-
-        // Allow the client to reconnect
-        reconnectAttempt.Continue();
-
-        // Wait for client to receive the notification via replay
-        await clientReceivedReplayedNotificationTcs.Task.WaitAsync(timeout, TestContext.Current.CancellationToken);
-
-        // Send a final notification while the client has reconnected - this should be handled by the transport
-        await server.SendNotificationAsync(CustomNotificationMethod, new JsonObject { ["message"] = ReconnectMessage }, cancellationToken: TestContext.Current.CancellationToken);
-
-        // Wait for the client to receive the final notification
-        await clientReceivedReconnectNotificationTcs.Task.WaitAsync(timeout, TestContext.Current.CancellationToken);
-
-        // Assert each notification was received exactly once
-        Assert.Equal(1, initialNotificationReceivedCount);
-        Assert.Equal(1, replayedNotificationReceivedCount);
-        Assert.Equal(1, reconnectNotificationReceivedCount);
-    }
-
-    [Fact]
-    public async Task Server_Returns400_WhenLastEventIdRefersToWrongSession()
-    {
-        // Arrange - Create server with event store
-        var eventStreamStore = new TestSseEventStreamStore();
-        await using var app = await CreateServerAsync(eventStreamStore);
-
-        // First, initialize a session and make a call to generate some events
-        using var initRequest = new HttpRequestMessage(HttpMethod.Post, "/")
-        {
-            Headers =
-            {
-                Accept = { new("application/json"), new("text/event-stream") }
-            },
-            Content = new StringContent(InitializeRequest, Encoding.UTF8, "application/json"),
-        };
-        var initResponse = await HttpClient.SendAsync(initRequest, HttpCompletionOption.ResponseHeadersRead, TestContext.Current.CancellationToken);
-        initResponse.EnsureSuccessStatusCode();
-
-        // Get the session ID from the response
-        var sessionId = initResponse.Headers.GetValues("Mcp-Session-Id").First();
-
-        // Read the SSE response to get an event ID
-        await using var initStream = await initResponse.Content.ReadAsStreamAsync(TestContext.Current.CancellationToken);
-        string? eventId = null;
-        await foreach (var sseItem in SseParser.Create(initStream).EnumerateAsync(TestContext.Current.CancellationToken))
-        {
-            if (!string.IsNullOrEmpty(sseItem.EventId))
-            {
-                eventId = sseItem.EventId;
-            }
-        }
-
-        Assert.NotNull(eventId);
-
-        // Act - Try to resume with a different session ID but the same event ID
-        var wrongSessionId = "wrong-session-id";
-        using var resumeRequest = new HttpRequestMessage(HttpMethod.Get, "/")
-        {
-            Headers =
-            {
-                Accept = { new("text/event-stream") },
-            }
-        };
-        resumeRequest.Headers.Add("Mcp-Session-Id", wrongSessionId);
-        resumeRequest.Headers.Add("Last-Event-ID", eventId);
-
-        var resumeResponse = await HttpClient.SendAsync(resumeRequest, HttpCompletionOption.ResponseHeadersRead, TestContext.Current.CancellationToken);
-
-        // Assert - First we get 404 because the wrong session doesn't exist
-        Assert.Equal(HttpStatusCode.NotFound, resumeResponse.StatusCode);
-
-        // Now test with an existing session but event ID from a different session
-        // Create a second session
-        using var initRequest2 = new HttpRequestMessage(HttpMethod.Post, "/")
-        {
-            Headers =
-            {
-                Accept = { new("application/json"), new("text/event-stream") }
-            },
-            Content = new StringContent(InitializeRequest, Encoding.UTF8, "application/json"),
-        };
-        var initResponse2 = await HttpClient.SendAsync(initRequest2, HttpCompletionOption.ResponseHeadersRead, TestContext.Current.CancellationToken);
-        initResponse2.EnsureSuccessStatusCode();
-
-        var sessionId2 = initResponse2.Headers.GetValues("Mcp-Session-Id").First();
-        Assert.NotEqual(sessionId, sessionId2);
-
-        // Read the second session's response
-        await using var initStream2 = await initResponse2.Content.ReadAsStreamAsync(TestContext.Current.CancellationToken);
-        await foreach (var _ in SseParser.Create(initStream2).EnumerateAsync(TestContext.Current.CancellationToken))
-        {
-            // Consume the stream
-        }
-
-        // Try to use session 2's ID but with an event ID from session 1
-        using var mismatchRequest = new HttpRequestMessage(HttpMethod.Get, "/")
-        {
-            Headers =
-            {
-                Accept = { new("text/event-stream") },
-            }
-        };
-        mismatchRequest.Headers.Add("Mcp-Session-Id", sessionId2);
-        mismatchRequest.Headers.Add("Last-Event-ID", eventId);  // This event ID belongs to session 1
-
-        var mismatchResponse = await HttpClient.SendAsync(mismatchRequest, HttpCompletionOption.ResponseHeadersRead, TestContext.Current.CancellationToken);
-
-        // Assert - Should get 400 Bad Request because the event ID doesn't match the session
-        Assert.Equal(HttpStatusCode.BadRequest, mismatchResponse.StatusCode);
-
-        // Verify the error message
-        var responseBody = await mismatchResponse.Content.ReadAsStringAsync(TestContext.Current.CancellationToken);
-        var errorResponse = JsonNode.Parse(responseBody);
-        Assert.NotNull(errorResponse);
-        var errorMessage = errorResponse["error"]?["message"]?.GetValue<string>();
-        Assert.Equal("Bad Request: The Last-Event-ID header refers to a session with a different session ID.", errorMessage);
-    }
-
-    [Fact]
-    public async Task EnablePollingAsync_SendsSseItemWithRetryField()
-    {
-        // Arrange
-        const string PollingToolName = "polling_tool";
-        var expectedRetryInterval = TimeSpan.FromSeconds(5);
-        var pollingTool = McpServerTool.Create(async (RequestContext<CallToolRequestParams> context) =>
-        {
-            await context.EnablePollingAsync(retryInterval: expectedRetryInterval);
-            return "Polling enabled";
-        }, options: new() { Name = PollingToolName });
-
-        var eventStreamStore = new TestSseEventStreamStore();
-        await using var app = await CreateServerAsync(eventStreamStore, configureServer: builder =>
-        {
-            builder.WithTools([pollingTool]);
-        });
-        await using var client = await ConnectClientAsync();
-
-        // Act - Call the tool that enables polling
-        var result = await client.CallToolAsync(PollingToolName, cancellationToken: TestContext.Current.CancellationToken);
-
-        // Assert - The result should be successful
-        Assert.False(result.IsError is true);
-        Assert.Equal("Polling enabled", result.Content.OfType<TextContentBlock>().Single().Text);
-
-        // Verify that the event store received the retry interval
-        Assert.Contains(expectedRetryInterval, eventStreamStore.StoredReconnectionIntervals);
     }
 
     [Fact]
@@ -542,82 +152,6 @@ public class ResumabilityIntegrationTests(ITestOutputHelper testOutputHelper) :
         await blockingStore.DisposedTask.WaitAsync(timeoutCts.Token);
     }
 
-    [McpServerToolType]
-    private class ResumabilityTestTools
-    {
-        [McpServerTool(Name = "echo"), Description("Echoes the message back")]
-        public static string Echo(string message) => $"Echo: {message}";
-    }
-
-    private async Task<WebApplication> CreateServerAsync(
-        ISseEventStreamStore? eventStreamStore = null,
-        Action<IMcpServerBuilder>? configureServer = null,
-        Action<HttpServerTransportOptions>? configureTransport = null)
-    {
-        var serverBuilder = Builder.Services.AddMcpServer()
-            .WithHttpTransport(options =>
-            {
-                options.EventStreamStore = eventStreamStore;
-                configureTransport?.Invoke(options);
-            })
-            .WithTools<ResumabilityTestTools>();
-
-        configureServer?.Invoke(serverBuilder);
-
-        var app = Builder.Build();
-        app.MapMcp();
-
-        await app.StartAsync(TestContext.Current.CancellationToken);
-        return app;
-    }
-
-    private async Task<McpClient> ConnectClientAsync()
-    {
-        var transport = new HttpClientTransport(new HttpClientTransportOptions
-        {
-            Endpoint = new Uri("http://localhost:5000/"),
-            TransportMode = HttpTransportMode.StreamableHttp,
-        }, HttpClient, LoggerFactory);
-
-        return await McpClient.CreateAsync(transport, loggerFactory: LoggerFactory,
-            cancellationToken: TestContext.Current.CancellationToken);
-    }
-
-    private async Task<SseResponse> SendInitializeAndReadSseResponseAsync(string initializeRequest)
-    {
-        using var requestContent = new StringContent(initializeRequest, Encoding.UTF8, "application/json");
-        using var request = new HttpRequestMessage(HttpMethod.Post, "/")
-        {
-            Headers =
-            {
-                Accept = { new("application/json"), new("text/event-stream") }
-            },
-            Content = requestContent,
-        };
-
-        var response = await HttpClient.SendAsync(request, HttpCompletionOption.ResponseHeadersRead,
-            TestContext.Current.CancellationToken);
-
-        response.EnsureSuccessStatusCode();
-
-        var sseResponse = new SseResponse();
-        await using var stream = await response.Content.ReadAsStreamAsync(TestContext.Current.CancellationToken);
-        await foreach (var sseItem in SseParser.Create(stream).EnumerateAsync(TestContext.Current.CancellationToken))
-        {
-            if (!string.IsNullOrEmpty(sseItem.EventId))
-            {
-                sseResponse.LastEventId = sseItem.EventId;
-            }
-        }
-
-        return sseResponse;
-    }
-
-    private struct SseResponse
-    {
-        public string? LastEventId { get; set; }
-    }
-
     /// <summary>
     /// A test event stream store that blocks on WriteEventAsync for response messages,
     /// allowing the test to cancel the operation and verify proper cleanup.

tests/ModelContextProtocol.AspNetCore.Tests/ResumabilityIntegrationTestsBase.cs

@@ -0,0 +1,560 @@
+using System.ComponentModel;
+using System.Diagnostics;
+using System.Net;
+using System.Net.ServerSentEvents;
+using System.Text;
+using System.Text.Json.Nodes;
+using Microsoft.AspNetCore.Builder;
+using Microsoft.Extensions.DependencyInjection;
+using ModelContextProtocol.AspNetCore.Tests.Utils;
+using ModelContextProtocol.Client;
+using ModelContextProtocol.Protocol;
+using ModelContextProtocol.Server;
+using ModelContextProtocol.Tests.Utils;
+
+namespace ModelContextProtocol.AspNetCore.Tests;
+
+/// <summary>
+/// Base class for SSE resumability integration tests that can be run against different
+/// <see cref="ISseEventStreamStore"/> implementations.
+/// </summary>
+/// <remarks>
+/// <para>
+/// Tests in this class verify resumability behavior without relying on implementation-specific
+/// internals of the event store. Derived classes can override virtual tests to add additional
+/// assertions specific to their event store implementation.
+/// </para>
+/// <para>
+/// The <see cref="CreateEventStreamStoreAsync"/> method must be implemented by derived classes
+/// to provide the specific event store implementation to test.
+/// </para>
+/// </remarks>
+public abstract class ResumabilityIntegrationTestsBase(ITestOutputHelper testOutputHelper) : KestrelInMemoryTest(testOutputHelper)
+{
+    /// <summary>
+    /// The initialize request JSON for the current protocol version.
+    /// </summary>
+    protected const string InitializeRequest = """
+        {"jsonrpc":"2.0","id":"1","method":"initialize","params":{"protocolVersion":"2025-11-25","capabilities":{},"clientInfo":{"name":"TestClient","version":"1.0.0"}}}
+        """;
+
+    /// <summary>
+    /// Gets the event stream store created for the current test.
+    /// </summary>
+    /// <remarks>
+    /// This is set after <see cref="CreateServerAsync"/> is called.
+    /// </remarks>
+    protected ISseEventStreamStore? EventStreamStore { get; private set; }
+
+    /// <summary>
+    /// Creates the event stream store implementation to use for this test.
+    /// </summary>
+    /// <returns>The event stream store instance.</returns>
+    protected abstract ValueTask<ISseEventStreamStore> CreateEventStreamStoreAsync();
+
+    [Fact]
+    public virtual async Task Server_StoresEvents_WhenEventStoreConfigured()
+    {
+        // Arrange
+        await using var app = await CreateServerAsync();
+        await using var client = await ConnectClientAsync();
+
+        // Act - Make a tool call which generates events
+        var result = await client.CallToolAsync("echo",
+            new Dictionary<string, object?> { ["message"] = "test" },
+            cancellationToken: TestContext.Current.CancellationToken);
+
+        // Assert - The call succeeded
+        Assert.NotNull(result);
+        var textContent = Assert.Single(result.Content.OfType<TextContentBlock>());
+        Assert.Equal("Echo: test", textContent.Text);
+    }
+
+    [Fact]
+    public virtual async Task Client_CanMakeMultipleRequests_WithResumabilityEnabled()
+    {
+        // Arrange
+        await using var app = await CreateServerAsync();
+        await using var client = await ConnectClientAsync();
+
+        // Act - Make many requests to verify stability
+        for (int i = 0; i < 5; i++)
+        {
+            var result = await client.CallToolAsync("echo",
+                new Dictionary<string, object?> { ["message"] = $"test{i}" },
+                cancellationToken: TestContext.Current.CancellationToken);
+
+            var textContent = Assert.Single(result.Content.OfType<TextContentBlock>());
+            Assert.Equal($"Echo: test{i}", textContent.Text);
+        }
+    }
+
+    [Fact]
+    public virtual async Task Ping_WorksWithResumabilityEnabled()
+    {
+        // Arrange
+        await using var app = await CreateServerAsync();
+        await using var client = await ConnectClientAsync();
+
+        // Act & Assert - Ping should work
+        await client.PingAsync(cancellationToken: TestContext.Current.CancellationToken);
+    }
+
+    [Fact]
+    public virtual async Task ListTools_WorksWithResumabilityEnabled()
+    {
+        // Arrange
+        await using var app = await CreateServerAsync();
+        await using var client = await ConnectClientAsync();
+
+        // Act
+        var tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
+
+        // Assert
+        Assert.NotNull(tools);
+        Assert.Single(tools);
+    }
+
+    [Fact]
+    public virtual async Task Client_CanPollResponse_FromServer()
+    {
+        const string ProgressToolName = "progress_tool";
+        var clientReceivedInitialValueTcs = new TaskCompletionSource();
+        var clientReceivedPolledValueTcs = new TaskCompletionSource();
+        var progressTool = McpServerTool.Create(async (RequestContext<CallToolRequestParams> context, IProgress<ProgressNotificationValue> progress) =>
+        {
+            progress.Report(new() { Progress = 0, Message = "Initial value" });
+
+            await clientReceivedInitialValueTcs.Task.WaitAsync(TestContext.Current.CancellationToken);
+
+            await context.EnablePollingAsync(retryInterval: TimeSpan.FromSeconds(1));
+
+            progress.Report(new() { Progress = 50, Message = "Polled value" });
+
+            await clientReceivedPolledValueTcs.Task.WaitAsync(TestContext.Current.CancellationToken);
+
+            return "Complete";
+        }, options: new() { Name = ProgressToolName });
+
+        await using var app = await CreateServerAsync(configureServer: builder =>
+        {
+            builder.WithTools([progressTool]);
+        });
+        await using var client = await ConnectClientAsync();
+
+        var progressHandler = new Progress<ProgressNotificationValue>(value =>
+        {
+            switch (value.Message)
+            {
+                case "Initial value":
+                    Assert.True(clientReceivedInitialValueTcs.TrySetResult(), "Received the initial value more than once.");
+                    break;
+                case "Polled value":
+                    Assert.True(clientReceivedPolledValueTcs.TrySetResult(), "Received the polled value more than once.");
+                    break;
+                default:
+                    throw new UnreachableException($"Unknown progress message '{value.Message}'");
+            }
+        });
+
+        var result = await client.CallToolAsync(ProgressToolName, progress: progressHandler, cancellationToken: TestContext.Current.CancellationToken);
+
+        Assert.False(result.IsError is true);
+        Assert.Equal("Complete", result.Content.OfType<TextContentBlock>().Single().Text);
+    }
+
+    [Fact]
+    public virtual async Task Client_CanResumePostResponseStream_AfterDisconnection()
+    {
+        using var faultingStreamHandler = new FaultingStreamHandler()
+        {
+            InnerHandler = SocketsHttpHandler,
+        };
+
+        HttpClient = new(faultingStreamHandler);
+        ConfigureHttpClient(HttpClient);
+
+        const string ProgressToolName = "progress_tool";
+        const string InitialMessage = "Initial notification";
+        const string ReplayedMessage = "Replayed notification";
+        const string ResultMessage = "Complete";
+
+        var clientReceivedInitialValueTcs = new TaskCompletionSource();
+        var clientReceivedReconnectValueTcs = new TaskCompletionSource();
+        var progressTool = McpServerTool.Create(async (RequestContext<CallToolRequestParams> context, IProgress<ProgressNotificationValue> progress, CancellationToken cancellationToken) =>
+        {
+            progress.Report(new() { Progress = 0, Message = InitialMessage });
+
+            // Make sure the client receives one message before we disconnect.
+            await clientReceivedInitialValueTcs.Task.WaitAsync(TestContext.Current.CancellationToken);
+
+            // Simulate a network disconnection by faulting the response stream.
+            var reconnectAttempt = await faultingStreamHandler.TriggerFaultAsync(TestContext.Current.CancellationToken);
+
+            // Send another message that the client should receive after reconnecting.
+            progress.Report(new() { Progress = 50, Message = ReplayedMessage });
+
+            reconnectAttempt.Continue();
+
+            // Wait for the client to receive the message via replay.
+            await clientReceivedReconnectValueTcs.Task.WaitAsync(TestContext.Current.CancellationToken);
+
+            // Return the final result with the client still connected.
+            return ResultMessage;
+        }, options: new() { Name = ProgressToolName });
+
+        await using var app = await CreateServerAsync(configureServer: builder =>
+        {
+            builder.WithTools([progressTool]);
+        });
+        await using var client = await ConnectClientAsync();
+
+        var initialNotificationReceivedCount = 0;
+        var replayedNotificationReceivedCount = 0;
+        var progressHandler = new Progress<ProgressNotificationValue>(value =>
+        {
+            switch (value.Message)
+            {
+                case InitialMessage:
+                    initialNotificationReceivedCount++;
+                    clientReceivedInitialValueTcs.TrySetResult();
+                    break;
+                case ReplayedMessage:
+                    replayedNotificationReceivedCount++;
+                    clientReceivedReconnectValueTcs.TrySetResult();
+                    break;
+                default:
+                    throw new UnreachableException($"Unknown progress message '{value.Message}'");
+            }
+        });
+
+        var result = await client.CallToolAsync(ProgressToolName, progress: progressHandler, cancellationToken: TestContext.Current.CancellationToken);
+
+        Assert.False(result.IsError is true);
+        Assert.Equal(1, initialNotificationReceivedCount);
+        Assert.Equal(1, replayedNotificationReceivedCount);
+        Assert.Equal(ResultMessage, result.Content.OfType<TextContentBlock>().Single().Text);
+    }
+
+    [Fact]
+    public virtual async Task Client_CanResumeUnsolicitedMessageStream_AfterDisconnection()
+    {
+        var timeout = TestConstants.DefaultTimeout;
+        using var faultingStreamHandler = new FaultingStreamHandler()
+        {
+            InnerHandler = SocketsHttpHandler,
+        };
+
+        HttpClient = new(faultingStreamHandler);
+        ConfigureHttpClient(HttpClient);
+
+        // Capture the server instance via RunSessionHandler
+        var serverTcs = new TaskCompletionSource<McpServer>();
+
+        await using var app = await CreateServerAsync(configureTransport: options =>
+        {
+            options.RunSessionHandler = (httpContext, mcpServer, cancellationToken) =>
+            {
+                serverTcs.TrySetResult(mcpServer);
+                return mcpServer.RunAsync(cancellationToken);
+            };
+        });
+
+        await using var client = await ConnectClientAsync();
+
+        // Get the server instance
+        var server = await serverTcs.Task.WaitAsync(timeout, TestContext.Current.CancellationToken);
+
+        // Set up notification tracking with unique messages
+        var clientReceivedInitialNotificationTcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+        var clientReceivedReplayedNotificationTcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+        var clientReceivedReconnectNotificationTcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+
+        const string CustomNotificationMethod = "test/custom_notification";
+        const string InitialMessage = "Initial notification";
+        const string ReplayedMessage = "Replayed notification";
+        const string ReconnectMessage = "Reconnect notification";
+
+        var initialNotificationReceivedCount = 0;
+        var replayedNotificationReceivedCount = 0;
+        var reconnectNotificationReceivedCount = 0;
+
+        await using var _ = client.RegisterNotificationHandler(CustomNotificationMethod, (notification, cancellationToken) =>
+        {
+            var message = notification.Params?["message"]?.GetValue<string>();
+            switch (message)
+            {
+                case InitialMessage:
+                    initialNotificationReceivedCount++;
+                    clientReceivedInitialNotificationTcs.TrySetResult();
+                    break;
+                case ReplayedMessage:
+                    replayedNotificationReceivedCount++;
+                    clientReceivedReplayedNotificationTcs.TrySetResult();
+                    break;
+                case ReconnectMessage:
+                    reconnectNotificationReceivedCount++;
+                    clientReceivedReconnectNotificationTcs.TrySetResult();
+                    break;
+                default:
+                    throw new UnreachableException($"Unknown notification message '{message}'");
+            }
+            return default;
+        });
+
+        // Wait for the client's unsolicited message stream to be established before sending notifications
+        await faultingStreamHandler.WaitForUnsolicitedMessageStreamAsync(TestContext.Current.CancellationToken);
+
+        // Send a custom notification to the client on the unsolicited message stream
+        await server.SendNotificationAsync(CustomNotificationMethod, new JsonObject { ["message"] = InitialMessage }, cancellationToken: TestContext.Current.CancellationToken);
+
+        // Wait for client to receive the first notification
+        await clientReceivedInitialNotificationTcs.Task.WaitAsync(timeout, TestContext.Current.CancellationToken);
+
+        // Fault the unsolicited message stream (GET SSE)
+        var reconnectAttempt = await faultingStreamHandler.TriggerFaultAsync(TestContext.Current.CancellationToken);
+
+        // Send another notification while the client is disconnected - this should be stored
+        await server.SendNotificationAsync(CustomNotificationMethod, new JsonObject { ["message"] = ReplayedMessage }, cancellationToken: TestContext.Current.CancellationToken);
+
+        // Allow the client to reconnect
+        reconnectAttempt.Continue();
+
+        // Wait for client to receive the notification via replay
+        await clientReceivedReplayedNotificationTcs.Task.WaitAsync(timeout, TestContext.Current.CancellationToken);
+
+        // Send a final notification while the client has reconnected - this should be handled by the transport
+        await server.SendNotificationAsync(CustomNotificationMethod, new JsonObject { ["message"] = ReconnectMessage }, cancellationToken: TestContext.Current.CancellationToken);
+
+        // Wait for the client to receive the final notification
+        await clientReceivedReconnectNotificationTcs.Task.WaitAsync(timeout, TestContext.Current.CancellationToken);
+
+        // Assert each notification was received exactly once
+        Assert.Equal(1, initialNotificationReceivedCount);
+        Assert.Equal(1, replayedNotificationReceivedCount);
+        Assert.Equal(1, reconnectNotificationReceivedCount);
+    }
+
+    [Fact]
+    public virtual async Task Server_Returns400_WhenLastEventIdRefersToWrongSession()
+    {
+        // Arrange - Create server with event store
+        await using var app = await CreateServerAsync();
+
+        // First, initialize a session and make a call to generate some events
+        using var initRequest = new HttpRequestMessage(HttpMethod.Post, "/")
+        {
+            Headers =
+            {
+                Accept = { new("application/json"), new("text/event-stream") }
+            },
+            Content = new StringContent(InitializeRequest, Encoding.UTF8, "application/json"),
+        };
+        var initResponse = await HttpClient.SendAsync(initRequest, HttpCompletionOption.ResponseHeadersRead, TestContext.Current.CancellationToken);
+        initResponse.EnsureSuccessStatusCode();
+
+        // Get the session ID from the response
+        var sessionId = initResponse.Headers.GetValues("Mcp-Session-Id").First();
+
+        // Read the SSE response to get an event ID
+        await using var initStream = await initResponse.Content.ReadAsStreamAsync(TestContext.Current.CancellationToken);
+        string? eventId = null;
+        await foreach (var sseItem in SseParser.Create(initStream).EnumerateAsync(TestContext.Current.CancellationToken))
+        {
+            if (!string.IsNullOrEmpty(sseItem.EventId))
+            {
+                eventId = sseItem.EventId;
+            }
+        }
+
+        Assert.NotNull(eventId);
+
+        // Act - Try to resume with a different session ID but the same event ID
+        var wrongSessionId = "wrong-session-id";
+        using var resumeRequest = new HttpRequestMessage(HttpMethod.Get, "/")
+        {
+            Headers =
+            {
+                Accept = { new("text/event-stream") },
+            }
+        };
+        resumeRequest.Headers.Add("Mcp-Session-Id", wrongSessionId);
+        resumeRequest.Headers.Add("Last-Event-ID", eventId);
+
+        var resumeResponse = await HttpClient.SendAsync(resumeRequest, HttpCompletionOption.ResponseHeadersRead, TestContext.Current.CancellationToken);
+
+        // Assert - First we get 404 because the wrong session doesn't exist
+        Assert.Equal(HttpStatusCode.NotFound, resumeResponse.StatusCode);
+
+        // Now test with an existing session but event ID from a different session
+        // Create a second session
+        using var initRequest2 = new HttpRequestMessage(HttpMethod.Post, "/")
+        {
+            Headers =
+            {
+                Accept = { new("application/json"), new("text/event-stream") }
+            },
+            Content = new StringContent(InitializeRequest, Encoding.UTF8, "application/json"),
+        };
+        var initResponse2 = await HttpClient.SendAsync(initRequest2, HttpCompletionOption.ResponseHeadersRead, TestContext.Current.CancellationToken);
+        initResponse2.EnsureSuccessStatusCode();
+
+        var sessionId2 = initResponse2.Headers.GetValues("Mcp-Session-Id").First();
+        Assert.NotEqual(sessionId, sessionId2);
+
+        // Read the second session's response
+        await using var initStream2 = await initResponse2.Content.ReadAsStreamAsync(TestContext.Current.CancellationToken);
+        await foreach (var _ in SseParser.Create(initStream2).EnumerateAsync(TestContext.Current.CancellationToken))
+        {
+            // Consume the stream
+        }
+
+        // Try to use session 2's ID but with an event ID from session 1
+        using var mismatchRequest = new HttpRequestMessage(HttpMethod.Get, "/")
+        {
+            Headers =
+            {
+                Accept = { new("text/event-stream") },
+            }
+        };
+        mismatchRequest.Headers.Add("Mcp-Session-Id", sessionId2);
+        mismatchRequest.Headers.Add("Last-Event-ID", eventId);  // This event ID belongs to session 1
+
+        var mismatchResponse = await HttpClient.SendAsync(mismatchRequest, HttpCompletionOption.ResponseHeadersRead, TestContext.Current.CancellationToken);
+
+        // Assert - Should get 400 Bad Request because the event ID doesn't match the session
+        Assert.Equal(HttpStatusCode.BadRequest, mismatchResponse.StatusCode);
+
+        // Verify the error message
+        var responseBody = await mismatchResponse.Content.ReadAsStringAsync(TestContext.Current.CancellationToken);
+        var errorResponse = JsonNode.Parse(responseBody);
+        Assert.NotNull(errorResponse);
+        var errorMessage = errorResponse["error"]?["message"]?.GetValue<string>();
+        Assert.Equal("Bad Request: The Last-Event-ID header refers to a session with a different session ID.", errorMessage);
+    }
+
+    [Fact]
+    public virtual async Task EnablePollingAsync_SendsSseItemWithRetryField()
+    {
+        // Arrange
+        const string PollingToolName = "polling_tool";
+        var expectedRetryInterval = TimeSpan.FromSeconds(5);
+        var pollingTool = McpServerTool.Create(async (RequestContext<CallToolRequestParams> context) =>
+        {
+            await context.EnablePollingAsync(retryInterval: expectedRetryInterval);
+            return "Polling enabled";
+        }, options: new() { Name = PollingToolName });
+
+        await using var app = await CreateServerAsync(configureServer: builder =>
+        {
+            builder.WithTools([pollingTool]);
+        });
+        await using var client = await ConnectClientAsync();
+
+        // Act - Call the tool that enables polling
+        var result = await client.CallToolAsync(PollingToolName, cancellationToken: TestContext.Current.CancellationToken);
+
+        // Assert - The result should be successful
+        Assert.False(result.IsError is true);
+        Assert.Equal("Polling enabled", result.Content.OfType<TextContentBlock>().Single().Text);
+    }
+
+    [McpServerToolType]
+    protected class ResumabilityTestTools
+    {
+        [McpServerTool(Name = "echo"), Description("Echoes the message back")]
+        public static string Echo(string message) => $"Echo: {message}";
+    }
+
+    /// <summary>
+    /// Creates a server with the event stream store from <see cref="CreateEventStreamStoreAsync"/>.
+    /// </summary>
+    protected async Task<WebApplication> CreateServerAsync(
+        Action<IMcpServerBuilder>? configureServer = null,
+        Action<HttpServerTransportOptions>? configureTransport = null)
+    {
+        EventStreamStore = await CreateEventStreamStoreAsync();
+        return await CreateServerAsync(EventStreamStore, configureServer, configureTransport);
+    }
+
+    /// <summary>
+    /// Creates a server with the specified event stream store.
+    /// </summary>
+    protected async Task<WebApplication> CreateServerAsync(
+        ISseEventStreamStore? eventStreamStore,
+        Action<IMcpServerBuilder>? configureServer = null,
+        Action<HttpServerTransportOptions>? configureTransport = null)
+    {
+        var serverBuilder = Builder.Services.AddMcpServer()
+            .WithHttpTransport(options =>
+            {
+                options.EventStreamStore = eventStreamStore;
+                configureTransport?.Invoke(options);
+            })
+            .WithTools<ResumabilityTestTools>();
+
+        configureServer?.Invoke(serverBuilder);
+
+        var app = Builder.Build();
+        app.MapMcp();
+
+        await app.StartAsync(TestContext.Current.CancellationToken);
+        return app;
+    }
+
+    /// <summary>
+    /// Connects a client to the server.
+    /// </summary>
+    protected async Task<McpClient> ConnectClientAsync()
+    {
+        var transport = new HttpClientTransport(new HttpClientTransportOptions
+        {
+            Endpoint = new Uri("http://localhost:5000/"),
+            TransportMode = HttpTransportMode.StreamableHttp,
+        }, HttpClient, LoggerFactory);
+
+        return await McpClient.CreateAsync(transport, loggerFactory: LoggerFactory,
+            cancellationToken: TestContext.Current.CancellationToken);
+    }
+
+    /// <summary>
+    /// Sends an initialize request and reads the SSE response.
+    /// </summary>
+    protected async Task<SseResponse> SendInitializeAndReadSseResponseAsync(string initializeRequest)
+    {
+        using var requestContent = new StringContent(initializeRequest, Encoding.UTF8, "application/json");
+        using var request = new HttpRequestMessage(HttpMethod.Post, "/")
+        {
+            Headers =
+            {
+                Accept = { new("application/json"), new("text/event-stream") }
+            },
+            Content = requestContent,
+        };
+
+        var response = await HttpClient.SendAsync(request, HttpCompletionOption.ResponseHeadersRead,
+            TestContext.Current.CancellationToken);
+
+        response.EnsureSuccessStatusCode();
+
+        var sseResponse = new SseResponse();
+        await using var stream = await response.Content.ReadAsStreamAsync(TestContext.Current.CancellationToken);
+        await foreach (var sseItem in SseParser.Create(stream).EnumerateAsync(TestContext.Current.CancellationToken))
+        {
+            if (!string.IsNullOrEmpty(sseItem.EventId))
+            {
+                sseResponse.LastEventId = sseItem.EventId;
+            }
+        }
+
+        return sseResponse;
+    }
+
+    /// <summary>
+    /// Response data from an SSE stream.
+    /// </summary>
+    protected struct SseResponse
+    {
+        public string? LastEventId { get; set; }
+    }
+}

tests/ModelContextProtocol.Tests/ModelContextProtocol.Tests.csproj

@@ -41,6 +41,10 @@
     <Compile Include="..\..\src\Common\Experimentals.cs" Link="Experimentals.cs" />
   </ItemGroup>
 
+  <ItemGroup>
+    <Compile Include="..\..\src\ModelContextProtocol\Server\DistributedCacheEventIdFormatter.cs" Link="DistributedCacheEventIdFormatter.cs" />
+  </ItemGroup>
+
   <ItemGroup Condition="'$(TargetFramework)' == 'net472'">
     <Compile Include="..\..\src\Common\Throw.cs" Link="Throw.cs" />
     <Compile Include="..\..\src\Common\Polyfills\**\*.cs" />
@@ -57,6 +61,7 @@
     </PackageReference>
     <PackageReference Include="Microsoft.Extensions.AI" />
     <PackageReference Include="Microsoft.Extensions.AI.OpenAI" />
+    <PackageReference Include="Microsoft.Extensions.Caching.Memory" />
     <PackageReference Include="Microsoft.Extensions.Logging" />
     <PackageReference Include="Microsoft.Extensions.Logging.Console" />
     <PackageReference Include="Microsoft.Extensions.TimeProvider.Testing" />
@@ -78,6 +83,7 @@
   </ItemGroup>
 
   <ItemGroup>
+    <ProjectReference Include="..\..\src\ModelContextProtocol\ModelContextProtocol.csproj" />
     <ProjectReference Include="..\..\src\ModelContextProtocol.Core\ModelContextProtocol.Core.csproj" />
     <ProjectReference Include="..\ModelContextProtocol.TestServer\ModelContextProtocol.TestServer.csproj" />
     <ProjectReference Include="..\..\samples\TestServerWithHosting\TestServerWithHosting.csproj" />