From feaae683a40246c83f06140ec34f0bd6693093c3 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Sat, 27 Jun 2026 12:49:24 -0400 Subject: [PATCH] [codex] Fix event store replay isolation (#4402) --- fastmcp_slim/fastmcp/server/event_store.py | 52 ++++++++++++++ fastmcp_slim/fastmcp/server/http.py | 42 ++++++++++- tests/server/test_event_store.py | 82 +++++++++++++++++++++- 3 files changed, 173 insertions(+), 3 deletions(-) diff --git a/fastmcp_slim/fastmcp/server/event_store.py b/fastmcp_slim/fastmcp/server/event_store.py index 9304a931e..4e6ec81f1 100644 --- a/fastmcp_slim/fastmcp/server/event_store.py +++ b/fastmcp_slim/fastmcp/server/event_store.py @@ -37,6 +37,58 @@ class StreamEventList(FastMCPBaseModel): event_ids: list[str] +class SessionScopedEventStore(SDKEventStore): + """EventStore adapter that isolates stream IDs to one transport session.""" + + def __init__(self, event_store: SDKEventStore, session_id: str): + self._event_store = event_store + self._stream_prefix = f"{len(session_id)}:{session_id}:" + + def _scope_stream_id(self, stream_id: StreamId) -> StreamId: + return f"{self._stream_prefix}{stream_id}" + + def _unscope_stream_id(self, stream_id: StreamId) -> StreamId | None: + if not stream_id.startswith(self._stream_prefix): + return None + return stream_id[len(self._stream_prefix) :] + + async def store_event( + self, stream_id: StreamId, message: JSONRPCMessage | None + ) -> EventId: + return await self._event_store.store_event( + self._scope_stream_id(stream_id), message + ) + + async def replay_events_after( + self, + last_event_id: EventId, + send_callback: EventCallback, + ) -> StreamId | None: + replayed_events: list[EventMessage] = [] + + async def buffer_event(event: EventMessage) -> None: + replayed_events.append(event) + + scoped_stream_id = await self._event_store.replay_events_after( + last_event_id, buffer_event + ) + if scoped_stream_id is None: + return None + + stream_id = self._unscope_stream_id(scoped_stream_id) + if stream_id is None: + logger.warning( + "Event ID %s does not belong to this session-scoped event store", + last_event_id, + ) + return None + + for event in replayed_events: + await send_callback(event) + + return stream_id + + class EventStore(SDKEventStore): """EventStore implementation backed by AsyncKeyValue. diff --git a/fastmcp_slim/fastmcp/server/http.py b/fastmcp_slim/fastmcp/server/http.py index 6459ee845..17ff78ad9 100644 --- a/fastmcp_slim/fastmcp/server/http.py +++ b/fastmcp_slim/fastmcp/server/http.py @@ -3,7 +3,8 @@ from __future__ import annotations from collections.abc import AsyncGenerator, Callable, Generator from contextlib import asynccontextmanager, contextmanager from contextvars import ContextVar -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any +from uuid import uuid4 from mcp.server.auth.routes import build_resource_metadata_url from mcp.server.lowlevel.server import LifespanResultT @@ -12,6 +13,7 @@ from mcp.server.streamable_http import ( EventStore, ) from mcp.server.streamable_http_manager import StreamableHTTPSessionManager +from mcp.server.transport_security import TransportSecuritySettings from starlette.applications import Starlette from starlette.middleware import Middleware from starlette.requests import Request @@ -21,6 +23,7 @@ from starlette.types import Lifespan, Receive, Scope, Send from fastmcp.server.auth import AuthProvider from fastmcp.server.auth.middleware import RequireAuthMiddleware +from fastmcp.server.event_store import SessionScopedEventStore from fastmcp.utilities.logging import get_logger if TYPE_CHECKING: @@ -29,6 +32,41 @@ if TYPE_CHECKING: logger = get_logger(__name__) +class FastMCPStreamableHTTPSessionManager(StreamableHTTPSessionManager): + """Session manager that scopes resumability storage per transport session.""" + + def __init__( + self, + app: Any, + event_store: EventStore | None = None, + json_response: bool = False, + stateless: bool = False, + security_settings: TransportSecuritySettings | None = None, + retry_interval: int | None = None, + ) -> None: + self._shared_event_store: EventStore | None = None + super().__init__( + app=app, + event_store=event_store, + json_response=json_response, + stateless=stateless, + security_settings=security_settings, + retry_interval=retry_interval, + ) + + @property + def event_store(self) -> EventStore | None: + if self._shared_event_store is None: + return None + # The SDK reads `self.event_store` once when constructing each transport. + # A fresh adapter gives that transport a private stream namespace. + return SessionScopedEventStore(self._shared_event_store, session_id=uuid4().hex) + + @event_store.setter + def event_store(self, event_store: EventStore | None) -> None: + self._shared_event_store = event_store + + class StreamableHTTPASGIApp: """ASGI application wrapper for Streamable HTTP server transport.""" @@ -359,7 +397,7 @@ def create_streamable_http_app( # Create a lifespan manager to start and stop the session manager @asynccontextmanager async def lifespan(app: Starlette) -> AsyncGenerator[None, None]: - streamable_http_app.session_manager = StreamableHTTPSessionManager( + streamable_http_app.session_manager = FastMCPStreamableHTTPSessionManager( app=server._mcp_server, event_store=event_store, retry_interval=retry_interval, diff --git a/tests/server/test_event_store.py b/tests/server/test_event_store.py index 533844365..41eb18206 100644 --- a/tests/server/test_event_store.py +++ b/tests/server/test_event_store.py @@ -4,7 +4,13 @@ import pytest from mcp.server.streamable_http import EventMessage from mcp.types import JSONRPCMessage, JSONRPCRequest -from fastmcp.server.event_store import EventEntry, EventStore, StreamEventList +from fastmcp.server.event_store import ( + EventEntry, + EventStore, + SessionScopedEventStore, + StreamEventList, +) +from fastmcp.server.http import FastMCPStreamableHTTPSessionManager class TestEventEntry: @@ -179,6 +185,80 @@ class TestEventStore: assert stream_id == "stream-2" assert len(stream2_replayed) == 1 + @pytest.mark.parametrize("stream_id", ["_GET_stream", "1"]) + async def test_session_scoped_stores_isolate_overlapping_stream_ids( + self, event_store, stream_id + ): + session_a_store = SessionScopedEventStore(event_store, "session-a") + session_b_store = SessionScopedEventStore(event_store, "session-b") + msg_a1 = JSONRPCMessage( + root=JSONRPCRequest(jsonrpc="2.0", method="session-a-1", id=1) + ) + msg_a2 = JSONRPCMessage( + root=JSONRPCRequest(jsonrpc="2.0", method="session-a-2", id=2) + ) + msg_b = JSONRPCMessage( + root=JSONRPCRequest(jsonrpc="2.0", method="session-b", id=3) + ) + + session_a_event = await session_a_store.store_event(stream_id, msg_a1) + await session_b_store.store_event(stream_id, msg_b) + session_a_second_event = await session_a_store.store_event(stream_id, msg_a2) + + replayed_events: list[EventMessage] = [] + + async def callback(event: EventMessage): + replayed_events.append(event) + + replayed_stream_id = await session_a_store.replay_events_after( + session_a_event, callback + ) + + assert replayed_stream_id == stream_id + assert [event.event_id for event in replayed_events] == [session_a_second_event] + replayed_message = replayed_events[0].message.root + assert isinstance(replayed_message, JSONRPCRequest) + assert replayed_message.method == "session-a-2" + + async def test_session_scoped_replay_rejects_foreign_last_event_id( + self, event_store + ): + session_a_store = SessionScopedEventStore(event_store, "session-a") + session_b_store = SessionScopedEventStore(event_store, "session-b") + msg_b1 = JSONRPCMessage( + root=JSONRPCRequest(jsonrpc="2.0", method="session-b-1", id=1) + ) + msg_b2 = JSONRPCMessage( + root=JSONRPCRequest(jsonrpc="2.0", method="session-b-2", id=2) + ) + + foreign_event_id = await session_b_store.store_event("_GET_stream", msg_b1) + await session_b_store.store_event("_GET_stream", msg_b2) + + replayed_events: list[EventMessage] = [] + + async def callback(event: EventMessage): + replayed_events.append(event) + + replayed_stream_id = await session_a_store.replay_events_after( + foreign_event_id, callback + ) + + assert replayed_stream_id is None + assert replayed_events == [] + + def test_session_manager_returns_scoped_event_stores(self, event_store): + session_manager = FastMCPStreamableHTTPSessionManager( + app=object(), event_store=event_store + ) + + first_transport_store = session_manager.event_store + second_transport_store = session_manager.event_store + + assert isinstance(first_transport_store, SessionScopedEventStore) + assert isinstance(second_transport_store, SessionScopedEventStore) + assert first_transport_store is not second_transport_store + async def test_default_storage_is_memory(self): """Test that EventStore defaults to in-memory storage.""" event_store = EventStore()