small fixes

Author: ihrpr

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

@@ -8,9 +8,10 @@
 import logging
 import time
 from collections.abc import Awaitable, Callable
+from operator import itemgetter
 from uuid import uuid4
 
-from mcp.server.streamableHttp import EventId, EventStore, StreamId
+from mcp.server.streamable_http import EventId, EventStore, StreamId
 from mcp.types import JSONRPCMessage
 
 logger = logging.getLogger(__name__)
@@ -54,15 +55,19 @@ async def replay_events_after(
         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, (sid, message, timestamp) in self.events.items()
-            if sid == stream_id and timestamp > last_timestamp
+            (event_id, message) for event_id, message, _ in events_sorted
         ]
 
-        # Sort by timestamp to ensure chronological order
-        events_to_replay.sort(key=lambda x: self.events[x[0]][2])
-
         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]}"

src/mcp/server/streamable_http.py

@@ -373,8 +373,7 @@ async def _handle_post_request(
                     response_message = None
 
                     # Use similar approach to SSE writer for consistency
-                    async for item in request_stream_reader:
-                        received_message, _ = item  # Extract message, ignore event_id
+                    async for received_message, _ in request_stream_reader:
                         # If it's a response, this is what we're waiting for
                         if isinstance(
                             received_message.root, JSONRPCResponse | JSONRPCError
@@ -425,9 +424,10 @@ async def sse_writer():
                     try:
                         async with sse_stream_writer, request_stream_reader:
                             # Process messages from the request-specific stream
-                            async for item in request_stream_reader:
-                                received_message, event_id = item
-
+                            async for (
+                                received_message,
+                                event_id,
+                            ) in request_stream_reader:
                                 # Build the event data
                                 event_data = {
                                     "event": "message",