1212from collections .abc import AsyncGenerator , Awaitable , Callable
1313from contextlib import asynccontextmanager
1414from dataclasses import dataclass
15+ from functools import partial
1516from http import HTTPStatus
1617from typing import Any , Final
1718
7172# Type aliases
7273StreamId = str
7374EventId = str
75+ # An SSE event-dict as accepted by sse-starlette (`event`, `data`, `id`, `retry`).
76+ SSEEvent = dict [str , Any ]
7477
7578
7679@dataclass
@@ -174,7 +177,7 @@ def __init__(
174177 MemoryObjectReceiveStream [EventMessage ],
175178 ],
176179 ] = {}
177- self ._sse_stream_writers : dict [RequestId , MemoryObjectSendStream [dict [ str , str ] ]] = {}
180+ self ._sse_stream_writers : dict [RequestId , MemoryObjectSendStream [SSEEvent ]] = {}
178181 self ._terminated = False
179182 # Idle timeout cancel scope; managed by the session manager.
180183 self .idle_scope : anyio .CancelScope | None = None
@@ -262,31 +265,48 @@ async def close_standalone_stream_callback() -> None:
262265
263266 return SessionMessage (message , metadata = metadata )
264267
265- async def _maybe_send_priming_event (
266- self ,
267- request_id : RequestId ,
268- sse_stream_writer : MemoryObjectSendStream [dict [str , Any ]],
269- protocol_version : str ,
270- ) -> None :
271- """Send priming event for SSE resumability if event_store is configured.
268+ async def _mint_priming_event (self , stream_id : StreamId , protocol_version : str ) -> SSEEvent | None :
269+ """Store the priming cursor for `stream_id` and return its SSE wire form.
272270
273- Only sends priming events to clients with protocol version >= 2025-11-25,
274- which includes the fix for handling empty SSE data. Older clients would
275- crash trying to parse empty data as JSON.
271+ Called before the request is dispatched so the priming row precedes
272+ anything `message_router` can store for this stream. Returns `None`
273+ when no event store is configured or the client predates 2025-11-25
274+ (older clients cannot parse the empty-data event).
276275 """
277276 if not self ._event_store :
278- return
279- # Priming events have empty data which older clients cannot handle.
277+ return None
280278 if not is_version_at_least (protocol_version , "2025-11-25" ):
281- return
282- priming_event_id = await self ._event_store .store_event (
283- str (request_id ), # Convert RequestId to StreamId (str)
284- None , # Priming event has no payload
285- )
286- priming_event : dict [str , str | int ] = {"id" : priming_event_id , "data" : "" }
279+ return None
280+ priming_event_id = await self ._event_store .store_event (stream_id , None )
281+ priming_event : SSEEvent = {"id" : priming_event_id , "data" : "" }
287282 if self ._retry_interval is not None :
288283 priming_event ["retry" ] = self ._retry_interval
289- await sse_stream_writer .send (priming_event )
284+ return priming_event
285+
286+ async def _run_sse_writer (
287+ self ,
288+ request_id : RequestId ,
289+ sse_stream_writer : MemoryObjectSendStream [SSEEvent ],
290+ request_stream_reader : MemoryObjectReceiveStream [EventMessage ],
291+ priming_event : SSEEvent | None ,
292+ ) -> None :
293+ """Forward `_request_streams[request_id]` onto the SSE wire for one POST."""
294+ try :
295+ async with sse_stream_writer , request_stream_reader :
296+ if priming_event is not None :
297+ await sse_stream_writer .send (priming_event )
298+ async for event_message in request_stream_reader :
299+ await sse_stream_writer .send (self ._create_event_data (event_message ))
300+ if isinstance (event_message .message , JSONRPCResponse | JSONRPCError ):
301+ break
302+ except anyio .ClosedResourceError : # pragma: lax no cover
303+ logger .debug ("SSE stream closed by close_sse_stream()" )
304+ except Exception : # pragma: lax no cover
305+ logger .exception ("Error in SSE writer" )
306+ finally :
307+ logger .debug ("Closing SSE writer" )
308+ self ._sse_stream_writers .pop (request_id , None )
309+ await self ._clean_up_memory_streams (request_id )
290310
291311 def _create_error_response (
292312 self ,
@@ -340,7 +360,7 @@ def _get_session_id(self, request: Request) -> str | None:
340360 """Extract the session ID from request headers."""
341361 return request .headers .get (MCP_SESSION_ID_HEADER )
342362
343- def _create_event_data (self , event_message : EventMessage ) -> dict [ str , str ] :
363+ def _create_event_data (self , event_message : EventMessage ) -> SSEEvent :
344364 """Create event data dictionary from an EventMessage."""
345365 event_data = {
346366 "event" : "message" ,
@@ -583,40 +603,16 @@ async def _handle_post_request(self, scope: Scope, request: Request, receive: Re
583603 await self ._clean_up_memory_streams (request_id )
584604 else :
585605 # Create SSE stream
586- sse_stream_writer , sse_stream_reader = anyio .create_memory_object_stream [dict [ str , str ] ](0 )
606+ sse_stream_writer , sse_stream_reader = anyio .create_memory_object_stream [SSEEvent ](0 )
587607
588608 # Store writer reference so close_sse_stream() can close it
589609 self ._sse_stream_writers [request_id ] = sse_stream_writer
590610
591- async def sse_writer ():
592- # Get the request ID from the incoming request message
593- try :
594- async with sse_stream_writer , request_stream_reader :
595- # Send priming event for SSE resumability
596- await self ._maybe_send_priming_event (request_id , sse_stream_writer , protocol_version )
597-
598- # Process messages from the request-specific stream
599- async for event_message in request_stream_reader :
600- # Build the event data
601- event_data = self ._create_event_data (event_message )
602- await sse_stream_writer .send (event_data )
603-
604- # If response, remove from pending streams and close
605- if isinstance (event_message .message , JSONRPCResponse | JSONRPCError ):
606- break
607- except anyio .ClosedResourceError : # pragma: lax no cover
608- # Expected when close_sse_stream() is called
609- logger .debug ("SSE stream closed by close_sse_stream()" )
610- except Exception : # pragma: lax no cover
611- logger .exception ("Error in SSE writer" )
612- finally :
613- logger .debug ("Closing SSE writer" )
614- self ._sse_stream_writers .pop (request_id , None )
615- await self ._clean_up_memory_streams (request_id )
616-
617- # Create and start EventSourceResponse
618- # SSE stream mode (original behavior)
619- # Set up headers
611+ # Store the priming event before the request is dispatched so its
612+ # event-store position precedes anything message_router can store
613+ # for this id (storage order == wire order by construction).
614+ priming_event = await self ._mint_priming_event (request_id , protocol_version )
615+
620616 headers = {
621617 "Cache-Control" : "no-cache, no-transform" ,
622618 "Connection" : "keep-alive" ,
@@ -625,7 +621,9 @@ async def sse_writer():
625621 }
626622 response = EventSourceResponse (
627623 content = sse_stream_reader ,
628- data_sender_callable = sse_writer ,
624+ data_sender_callable = partial (
625+ self ._run_sse_writer , request_id , sse_stream_writer , request_stream_reader , priming_event
626+ ),
629627 headers = headers ,
630628 )
631629
@@ -708,7 +706,7 @@ async def _handle_get_request(self, request: Request, send: Send) -> None:
708706 return
709707
710708 # Create SSE stream
711- sse_stream_writer , sse_stream_reader = anyio .create_memory_object_stream [dict [ str , str ] ](0 )
709+ sse_stream_writer , sse_stream_reader = anyio .create_memory_object_stream [SSEEvent ](0 )
712710
713711 async def standalone_sse_writer ():
714712 try :
@@ -903,11 +901,10 @@ async def _replay_events(self, last_event_id: str, request: Request, send: Send)
903901 if self .mcp_session_id : # pragma: no branch
904902 headers [MCP_SESSION_ID_HEADER ] = self .mcp_session_id
905903
906- # Get protocol version from header (already validated in _validate_protocol_version)
907904 replay_protocol_version = request .headers .get (MCP_PROTOCOL_VERSION_HEADER , DEFAULT_NEGOTIATED_VERSION )
908905
909906 # Create SSE stream for replay
910- sse_stream_writer , sse_stream_reader = anyio .create_memory_object_stream [dict [ str , str ] ](0 )
907+ sse_stream_writer , sse_stream_reader = anyio .create_memory_object_stream [SSEEvent ](0 )
911908
912909 async def replay_sender ():
913910 try :
@@ -925,8 +922,12 @@ async def send_event(event_message: EventMessage) -> None:
925922 # Register SSE writer so close_sse_stream() can close it
926923 self ._sse_stream_writers [stream_id ] = sse_stream_writer
927924
928- # Send priming event for this new connection
929- await self ._maybe_send_priming_event (stream_id , sse_stream_writer , replay_protocol_version )
925+ # Prime the resumed connection so the client sees the stream
926+ # is re-registered. The replay→live-tail ordering window here
927+ # is pre-existing and tracked separately.
928+ priming_event = await self ._mint_priming_event (stream_id , replay_protocol_version )
929+ if priming_event is not None : # pragma: no branch
930+ await sse_stream_writer .send (priming_event )
930931
931932 # Create new request streams for this connection
932933 self ._request_streams [stream_id ] = anyio .create_memory_object_stream [EventMessage ](
0 commit comments