improve event store example implementation

Author: ihrpr

examples/servers/simple-streamablehttp/mcp_simple_streamablehttp/event_store.py

@@ -6,75 +6,100 @@
 """
 
 import logging
-import time
-from collections.abc import Awaitable, Callable
-from operator import itemgetter
+from collections import deque
+from dataclasses import dataclass
 from uuid import uuid4
 
-from mcp.server.streamable_http import EventId, EventStore, StreamId
+from mcp.server.streamable_http import (
+    EventCallback,
+    EventId,
+    EventMessage,
+    EventStore,
+    StreamId,
+)
 from mcp.types import JSONRPCMessage
 
 logger = logging.getLogger(__name__)
 
 
+@dataclass
+class EventEntry:
+    """
+    Represents an event entry in the event store.
+    """
+
+    event_id: EventId
+    stream_id: StreamId
+    message: JSONRPCMessage
+
+
 class InMemoryEventStore(EventStore):
     """
     Simple in-memory implementation of the EventStore interface for resumability.
     This is primarily intended for examples and testing, not for production use
     where a persistent storage solution would be more appropriate.
+
+    This implementation keeps only the last N events per stream for memory efficiency.
     """
 
-    def __init__(self):
-        self.events: dict[
-            str, tuple[str, JSONRPCMessage, float]
-        ] = {}  # event_id -> (stream_id, message, timestamp)
+    def __init__(self, max_events_per_stream: int = 100):
+        """Initialize the event store.
+
+        Args:
+            max_events_per_stream: Maximum number of events to keep per stream
+        """
+        self.max_events_per_stream = max_events_per_stream
+        # for maintaining last N events per stream
+        self.streams: dict[StreamId, deque[EventEntry]] = {}
+        # event_id -> EventEntry for quick lookup
+        self.event_index: dict[EventId, EventEntry] = {}
 
     async def store_event(
         self, stream_id: StreamId, message: JSONRPCMessage
     ) -> EventId:
         """Stores an event with a generated event ID."""
         event_id = str(uuid4())
-        self.events[event_id] = (stream_id, message, time.time())
+        event_entry = EventEntry(
+            event_id=event_id, stream_id=stream_id, message=message
+        )
+
+        # Get or create deque for this stream
+        if stream_id not in self.streams:
+            self.streams[stream_id] = deque(maxlen=self.max_events_per_stream)
+
+        # If deque is full, the oldest event will be automatically removed
+        # We need to remove it from the event_index as well
+        if len(self.streams[stream_id]) == self.max_events_per_stream:
+            oldest_event = self.streams[stream_id][0]
+            self.event_index.pop(oldest_event.event_id, None)
+
+        # Add new event
+        self.streams[stream_id].append(event_entry)
+        self.event_index[event_id] = event_entry
+
         return event_id
 
     async def replay_events_after(
         self,
         last_event_id: EventId,
-        send_callback: Callable[[EventId, JSONRPCMessage], Awaitable[None]],
-    ) -> StreamId:
+        send_callback: EventCallback,
+    ) -> StreamId | None:
         """Replays events that occurred after the specified event ID."""
-        logger.debug(f"Attempting to replay events after {last_event_id}")
-        logger.debug(f"Total events in store: {len(self.events)}")
-        logger.debug(f"Event IDs in store: {list(self.events.keys())}")
-
-        if not last_event_id or last_event_id not in self.events:
+        if last_event_id not in self.event_index:
             logger.warning(f"Event ID {last_event_id} not found in store")
-            return ""
-
-        # Get the stream ID and timestamp from the last event
-        stream_id, _, last_timestamp = self.events[last_event_id]
-
-        # Find all events for this stream after the last event
-        events_sorted = sorted(
-            [
-                (event_id, message, timestamp)
-                for event_id, (sid, message, timestamp) in self.events.items()
-                if sid == stream_id and timestamp > last_timestamp
-            ],
-            key=itemgetter(2),
-        )
-
-        events_to_replay = [
-            (event_id, message) for event_id, message, _ in events_sorted
-        ]
-
-        logger.debug(f"Found {len(events_to_replay)} events to replay")
-        logger.debug(
-            f"Events to replay: {[event_id for event_id, _ in events_to_replay]}"
-        )
-
-        # Send all events in order
-        for event_id, message in events_to_replay:
-            await send_callback(event_id, message)
+            return None
+
+        # Get the stream and find events after the last one
+        last_event = self.event_index[last_event_id]
+        stream_id = last_event.stream_id
+        stream_events = self.streams.get(last_event.stream_id, deque())
+
+        # Events in deque are already in chronological order
+        found_last = False
+        for event in stream_events:
+            if found_last:
+                await send_callback(EventMessage(event.message, event.event_id))
+            elif event.event_id == last_event_id:
+                found_last = True
 
         return stream_id

src/mcp/server/streamable_http.py

@@ -63,6 +63,22 @@
 EventId = str
 
 
+class EventMessage:
+    """
+    A JSONRPCMessage with an optional event ID for stream resumability.
+    """
+
+    message: JSONRPCMessage
+    event_id: str | None
+
+    def __init__(self, message: JSONRPCMessage, event_id: str | None = None):
+        self.message = message
+        self.event_id = event_id
+
+
+EventCallback = Callable[[EventMessage], Awaitable[None]]
+
+
 class EventStore(ABC):
     """
     Interface for resumability support via event storage.
@@ -88,8 +104,8 @@ async def store_event(
     async def replay_events_after(
         self,
         last_event_id: EventId,
-        send_callback: Callable[[EventId, JSONRPCMessage], Awaitable[None]],
-    ) -> StreamId:
+        send_callback: EventCallback,
+    ) -> StreamId | None:
         """
         Replays events that occurred after the specified event ID.
 
@@ -149,7 +165,7 @@ def __init__(
         self.is_json_response_enabled = is_json_response_enabled
         self._event_store = event_store
         self._request_streams: dict[
-            RequestId, MemoryObjectSendStream[tuple[JSONRPCMessage, str | None]]
+            RequestId, MemoryObjectSendStream[EventMessage]
         ] = {}
         self._terminated = False
 
@@ -358,7 +374,7 @@ async def _handle_post_request(
             request_id = str(message.root.id)
             # Create promise stream for getting response
             request_stream_writer, request_stream_reader = (
-                anyio.create_memory_object_stream[tuple[JSONRPCMessage, str | None]](0)
+                anyio.create_memory_object_stream[EventMessage](0)
             )
 
             # Register this stream for the request ID
@@ -373,16 +389,18 @@ async def _handle_post_request(
                     response_message = None
 
                     # Use similar approach to SSE writer for consistency
-                    async for received_message, _ in request_stream_reader:
+                    async for event_message in request_stream_reader:
                         # If it's a response, this is what we're waiting for
                         if isinstance(
-                            received_message.root, JSONRPCResponse | JSONRPCError
+                            event_message.message.root, JSONRPCResponse | JSONRPCError
                         ):
-                            response_message = received_message
+                            response_message = event_message.message
                             break
                         # For notifications and request, keep waiting
                         else:
-                            logger.debug(f"received: {received_message.root.method}")
+                            logger.debug(
+                                f"received: {event_message.message.root.method}"
+                            )
 
                     # At this point we should have a response
                     if response_message:
@@ -424,27 +442,24 @@ async def sse_writer():
                     try:
                         async with sse_stream_writer, request_stream_reader:
                             # Process messages from the request-specific stream
-                            async for (
-                                received_message,
-                                event_id,
-                            ) in request_stream_reader:
+                            async for event_message in request_stream_reader:
                                 # Build the event data
                                 event_data = {
                                     "event": "message",
-                                    "data": received_message.model_dump_json(
+                                    "data": event_message.message.model_dump_json(
                                         by_alias=True, exclude_none=True
                                     ),
                                 }
 
                                 # If an event ID was provided, include it
-                                if event_id:
-                                    event_data["id"] = event_id
+                                if event_message.event_id:
+                                    event_data["id"] = event_message.event_id
 
                                 await sse_stream_writer.send(event_data)
 
                                 # If response, remove from pending streams and close
                                 if isinstance(
-                                    received_message.root,
+                                    event_message.message.root,
                                     JSONRPCResponse | JSONRPCError,
                                 ):
                                     if request_id:
@@ -563,20 +578,15 @@ async def standalone_sse_writer():
             try:
                 # Create a standalone message stream for server-initiated messages
                 standalone_stream_writer, standalone_stream_reader = (
-                    anyio.create_memory_object_stream[
-                        tuple[JSONRPCMessage, str | None]
-                    ](0)
+                    anyio.create_memory_object_stream[EventMessage](0)
                 )
 
                 # Register this stream using the special key
                 self._request_streams[GET_STREAM_KEY] = standalone_stream_writer
 
                 async with sse_stream_writer, standalone_stream_reader:
                     # Process messages from the standalone stream
-                    async for item in standalone_stream_reader:
-                        # The message router always sends a tuple of (message, event_id)
-                        received_message, event_id = item
-
+                    async for event_message in standalone_stream_reader:
                         # For the standalone stream, we handle:
                         # - JSONRPCNotification (server sends notifications to client)
                         # - JSONRPCRequest (server sends requests to client)
@@ -585,14 +595,14 @@ async def standalone_sse_writer():
                         # Send the message via SSE
                         event_data = {
                             "event": "message",
-                            "data": received_message.model_dump_json(
+                            "data": event_message.message.model_dump_json(
                                 by_alias=True, exclude_none=True
                             ),
                         }
 
                         # If an event ID was provided, include it in the SSE stream
-                        if event_id:
-                            event_data["id"] = event_id
+                        if event_message.event_id:
+                            event_data["id"] = event_message.event_id
 
                         await sse_stream_writer.send(event_data)
             except Exception as e:
@@ -741,14 +751,12 @@ async def replay_sender():
                 try:
                     async with sse_stream_writer:
                         # Define an async callback for sending events
-                        async def send_event(
-                            event_id: EventId, message: JSONRPCMessage
-                        ) -> None:
+                        async def send_event(event_message: EventMessage) -> None:
                             await sse_stream_writer.send(
                                 {
                                     "event": "message",
-                                    "id": event_id,
-                                    "data": message.model_dump_json(
+                                    "id": event_message.event_id,
+                                    "data": event_message.message.model_dump_json(
                                         by_alias=True, exclude_none=True
                                     ),
                                 }
@@ -762,22 +770,21 @@ async def send_event(
                         # If stream ID not in mapping, create it
                         if stream_id and stream_id not in self._request_streams:
                             msg_writer, msg_reader = anyio.create_memory_object_stream[
-                                tuple[JSONRPCMessage, str | None]
+                                EventMessage
                             ](0)
                             self._request_streams[stream_id] = msg_writer
 
                             # Forward messages to SSE
                             async with msg_reader:
-                                async for item in msg_reader:
-                                    message, event_id = item
-
+                                async for event_message in msg_reader:
+                                    event_data = event_message.message.model_dump_json(
+                                        by_alias=True, exclude_none=True
+                                    )
                                     await sse_stream_writer.send(
                                         {
                                             "event": "message",
-                                            "id": event_id,
-                                            "data": message.model_dump_json(
-                                                by_alias=True, exclude_none=True
-                                            ),
+                                            "id": event_message.event_id,
+                                            "data": event_data,
                                         }
                                     )
                 except Exception as e:
@@ -871,7 +878,7 @@ async def message_router():
                             try:
                                 # Send both the message and the event ID
                                 await self._request_streams[request_stream_id].send(
-                                    (message, event_id)
+                                    EventMessage(message, event_id)
                                 )
                             except (
                                 anyio.BrokenResourceError,