Serialize the event store's stream list read-modify-write (#4758)

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Co-authored-by: Jeremiah Lowin <153965+jlowin@users.noreply.github.com>

fastmcp_slim/fastmcp/server/event_store.py

@@ -8,6 +8,7 @@ AsyncKeyValue protocol, allowing users to configure any compatible backend
 
 from __future__ import annotations
 
+import asyncio
 from uuid import uuid4
 
 from key_value.aio.adapters.pydantic import PydanticAdapter
@@ -30,6 +31,9 @@ logger = get_logger(__name__)
 # TypeAdapter to validate a stored dict back into the correct member.
 _jsonrpc_message_adapter: TypeAdapter[JSONRPCMessage] = TypeAdapter(JSONRPCMessage)
 
+# Number of striped locks guarding stream event lists. See EventStore.__init__.
+_LOCK_STRIPES = 64
+
 
 class EventEntry(FastMCPBaseModel):
     """Stored event entry."""
@@ -84,6 +88,20 @@ class EventStore(SDKEventStore):
         self._storage: AsyncKeyValue = storage or MemoryStore()
         self._max_events_per_stream = max_events_per_stream
         self._ttl = ttl
+        # Serializes the read-modify-write of each stream's event list. A fixed
+        # set of striped locks rather than one lock per stream: a single store is
+        # shared by every session, so a store-wide lock would serialize unrelated
+        # streams across a Redis round-trip, while a per-stream map would grow
+        # with every session and need its own eviction. Two streams only contend
+        # when their IDs collide on the same stripe.
+        #
+        # In-process locks are enough because a stream list only ever has
+        # in-process writers: every transport gets its own SessionScopedEventStore
+        # with a random per-session prefix, so no two servers sharing one backend
+        # address the same stream key. Coordinating across processes would need a
+        # compare-and-swap or transactional update, which AsyncKeyValue does not
+        # expose -- it offers only get/put/delete/ttl.
+        self._stream_locks = tuple(asyncio.Lock() for _ in range(_LOCK_STRIPES))
 
         # PydanticAdapter for type-safe storage (following OAuth proxy pattern)
         self._event_store: PydanticAdapter[EventEntry] = PydanticAdapter[EventEntry](
@@ -121,7 +139,12 @@ class EventStore(SDKEventStore):
         )
         await self._event_store.put(key=event_id, value=entry, ttl=self._ttl)
 
-        # Update stream's event list
+        # Update stream's event list. A session stores events from more than one
+        # task -- the SSE writer and the message router both do -- so this
+        # read-modify-write has to be serialized. Interleaved, each task reads the
+        # same list, appends only its own ID, and the later write drops the other
+        # event entirely while both tasks evict the same expired IDs.
+        async with self._stream_locks[hash(stream_id) % _LOCK_STRIPES]:
             stream_data = await self._stream_store.get(key=stream_id)
             event_ids = stream_data.event_ids if stream_data else []
             event_ids.append(event_id)

tests/server/test_event_store.py

@@ -1,10 +1,13 @@
 """Tests for the EventStore implementation."""
 
+import asyncio
+
 import pytest
 from mcp.server.streamable_http import EventMessage
 from mcp_types import JSONRPCRequest
 
 from fastmcp.server.event_store import (
+    _LOCK_STRIPES,
     EventEntry,
     EventStore,
     SessionScopedEventStore,
@@ -260,6 +263,87 @@ class TestEventStore:
         assert len(replayed) == 1
 
 
+class TestConcurrentStoreEvent:
+    async def test_concurrent_stores_on_one_stream(self, monkeypatch):
+        """Concurrent stores must not lose events or evict the same ID twice.
+
+        A live session stores events from more than one task (the SSE writer and
+        the message router), so the stream's event list is read and written
+        concurrently. Interleaved, each task appends only its own ID to the list
+        it read, and both evict the same expired IDs -- the second delete is the
+        one that raised `FileNotFoundError` on a file-backed store.
+        """
+        event_store = EventStore(max_events_per_stream=2)
+
+        stream_get = event_store._stream_store.get
+        event_delete = event_store._event_store.delete
+        deleted: list[str] = []
+
+        async def yielding_get(**kwargs):
+            # Suspend between the read and the write so the tasks interleave.
+            stream_data = await stream_get(**kwargs)
+            await asyncio.sleep(0)
+            return stream_data
+
+        async def recording_delete(**kwargs):
+            deleted.append(kwargs["key"])
+            return await event_delete(**kwargs)
+
+        monkeypatch.setattr(event_store._stream_store, "get", yielding_get)
+        monkeypatch.setattr(event_store._event_store, "delete", recording_delete)
+
+        message = JSONRPCRequest(jsonrpc="2.0", method="test", id=1)
+        event_ids = await asyncio.gather(
+            *(event_store.store_event("stream-1", message) for _ in range(5))
+        )
+
+        stream_data = await stream_get(key="stream-1")
+        assert stream_data is not None
+        # The two most recent events are retained; every other ID was evicted
+        # exactly once, and no ID vanished without being evicted.
+        assert len(stream_data.event_ids) == 2
+        assert sorted(stream_data.event_ids + deleted) == sorted(event_ids)
+        assert len(deleted) == len(set(deleted))
+
+    async def test_distinct_streams_are_not_serialized(self, monkeypatch):
+        """Unrelated streams must not wait on each other's backend calls.
+
+        One EventStore is shared by every session, so a store-wide lock would
+        put a Redis round-trip for one session in front of every other one.
+        """
+        event_store = EventStore()
+
+        # hash() is salted per process, so pick the second stream at runtime.
+        first = "stream-a"
+        second = next(
+            candidate
+            for candidate in (f"stream-{i}" for i in range(1000))
+            if hash(candidate) % _LOCK_STRIPES != hash(first) % _LOCK_STRIPES
+        )
+
+        stream_get = event_store._stream_store.get
+        both_inside = asyncio.Event()
+        inside = 0
+
+        async def gate(**kwargs):
+            nonlocal inside
+            inside += 1
+            if inside == 2:
+                both_inside.set()
+            # Both critical sections have to be open at once; a store-wide lock
+            # would keep the second task out until the first finished.
+            await asyncio.wait_for(both_inside.wait(), timeout=2)
+            return await stream_get(**kwargs)
+
+        monkeypatch.setattr(event_store._stream_store, "get", gate)
+
+        message = JSONRPCRequest(jsonrpc="2.0", method="test", id=1)
+        await asyncio.gather(
+            event_store.store_event(first, message),
+            event_store.store_event(second, message),
+        )
+
+
 class TestEventStoreIntegration:
     """Integration tests for EventStore with actual message types."""