From 242918cd73bb2d617a84547157f12085d550a159 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 10 Jun 2026 16:50:28 +0000 Subject: [PATCH] wire an EventStore into the stateful MCP session manager for resumability The stateful StreamableHTTPSessionManager ran with event_store=None, so SSE events carried no ids and a client reconnecting with Last-Event-ID (per the Streamable HTTP transport's resumability section) silently lost every message sent while it was disconnected. Add a bounded in-memory EventStore (LRU across streams, capped per stream, sizes overridable via env) and wire it in so the SDK transport replays missed events on reconnect. Storage is per-worker, matching the constraint that the live session object itself only exists on the worker that created it; the stale-session handling for multi-worker deployments is unchanged. --- litellm/constants.py | 9 ++ .../_experimental/mcp_server/event_store.py | 98 ++++++++++++++++ .../proxy/_experimental/mcp_server/server.py | 11 +- .../mcp_server/test_event_store.py | 109 ++++++++++++++++++ 4 files changed, 226 insertions(+), 1 deletion(-) create mode 100644 litellm/proxy/_experimental/mcp_server/event_store.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_event_store.py diff --git a/litellm/constants.py b/litellm/constants.py index 57f55e6c177..fe376ff0e97 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -144,6 +144,15 @@ MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS = int( os.getenv("MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS", "60") ) +# Bounds for the in-memory Streamable HTTP event store backing MCP session +# resumability (Last-Event-ID replay). +MCP_EVENT_STORE_MAX_STREAMS = int( + os.getenv("LITELLM_MCP_EVENT_STORE_MAX_STREAMS", "1000") +) +MCP_EVENT_STORE_MAX_EVENTS_PER_STREAM = int( + os.getenv("LITELLM_MCP_EVENT_STORE_MAX_EVENTS_PER_STREAM", "100") +) + # MCP timeout defaults (seconds). Override via env vars for slow/custom MCP servers. MCP_CLIENT_TIMEOUT = float(os.getenv("LITELLM_MCP_CLIENT_TIMEOUT", "60.0")) MCP_TOOL_LISTING_TIMEOUT = float(os.getenv("LITELLM_MCP_TOOL_LISTING_TIMEOUT", "30.0")) diff --git a/litellm/proxy/_experimental/mcp_server/event_store.py b/litellm/proxy/_experimental/mcp_server/event_store.py new file mode 100644 index 00000000000..fcefbc55965 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/event_store.py @@ -0,0 +1,98 @@ +"""In-memory EventStore enabling Streamable HTTP resumability for the MCP gateway. + +Wired into the stateful ``StreamableHTTPSessionManager`` so the SDK transport +tags SSE events with ids and replays missed events when a client reconnects +with ``Last-Event-ID`` (MCP Streamable HTTP transport, resumability section). + +Storage is per-worker, which matches the constraint that the live session +object itself only exists on the worker that created it. +""" + +from collections import OrderedDict, deque +from dataclasses import dataclass +from typing import Deque, Dict, Optional +from uuid import uuid4 + +from mcp.server.streamable_http import ( + EventCallback, + EventId, + EventMessage, + EventStore, + StreamId, +) +from mcp.types import JSONRPCMessage + +from litellm.constants import ( + MCP_EVENT_STORE_MAX_EVENTS_PER_STREAM, + MCP_EVENT_STORE_MAX_STREAMS, +) + + +@dataclass +class _StoredEvent: + event_id: EventId + message: Optional[JSONRPCMessage] + + +class InMemoryMCPEventStore(EventStore): + """Bounded in-memory event store. + + Streams are evicted least-recently-used once ``max_streams`` is reached, + and each stream keeps at most ``max_events_per_stream`` events, so a + long-lived session cannot grow worker memory without bound. + """ + + def __init__( + self, + max_streams: int = MCP_EVENT_STORE_MAX_STREAMS, + max_events_per_stream: int = MCP_EVENT_STORE_MAX_EVENTS_PER_STREAM, + ) -> None: + self._max_streams = max_streams + self._max_events_per_stream = max_events_per_stream + self._streams: "OrderedDict[StreamId, Deque[_StoredEvent]]" = OrderedDict() + self._event_to_stream: Dict[EventId, StreamId] = {} + + async def store_event( + self, stream_id: StreamId, message: Optional[JSONRPCMessage] + ) -> EventId: + event_id: EventId = uuid4().hex + + stream = self._streams.get(stream_id) + if stream is None: + while len(self._streams) >= self._max_streams: + _, evicted = self._streams.popitem(last=False) + for event in evicted: + self._event_to_stream.pop(event.event_id, None) + stream = deque() + self._streams[stream_id] = stream + else: + self._streams.move_to_end(stream_id) + + while len(stream) >= self._max_events_per_stream: + oldest = stream.popleft() + self._event_to_stream.pop(oldest.event_id, None) + + stream.append(_StoredEvent(event_id=event_id, message=message)) + self._event_to_stream[event_id] = stream_id + return event_id + + async def replay_events_after( + self, + last_event_id: EventId, + send_callback: EventCallback, + ) -> Optional[StreamId]: + stream_id = self._event_to_stream.get(last_event_id) + if stream_id is None: + return None + + found_last = False + for event in list(self._streams.get(stream_id, ())): + if found_last: + if event.message is not None: + await send_callback( + EventMessage(message=event.message, event_id=event.event_id) + ) + elif event.event_id == last_event_id: + found_last = True + + return stream_id diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 3dd4daf75a4..20cc114e5cb 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -346,9 +346,18 @@ if MCP_AVAILABLE: stateless=True, ) + from litellm.proxy._experimental.mcp_server.event_store import ( + InMemoryMCPEventStore, + ) + + # Event store enables spec-compliant resumability: the transport tags SSE + # events with ids and replays missed events when a stateful client + # reconnects with Last-Event-ID. + _stateful_event_store = InMemoryMCPEventStore() + session_manager_stateful = StreamableHTTPSessionManager( app=server, - event_store=None, # TODO: Add EventStore for reconnection/event replay if needed + event_store=_stateful_event_store, json_response=False, # enables SSE streaming stateless=False, ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_event_store.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_event_store.py new file mode 100644 index 00000000000..e2d4a3d5fbb --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_event_store.py @@ -0,0 +1,109 @@ +import pytest + +pytest.importorskip("mcp") + +from mcp.server.streamable_http import EventMessage +from mcp.types import JSONRPCMessage, JSONRPCNotification + +from litellm.proxy._experimental.mcp_server.event_store import InMemoryMCPEventStore + + +def _message(seq: int) -> JSONRPCMessage: + return JSONRPCMessage( + JSONRPCNotification( + jsonrpc="2.0", + method="notifications/message", + params={"seq": seq}, + ) + ) + + +def _collector(): + received: list = [] + + async def send(event: EventMessage) -> None: + received.append(event) + + return received, send + + +@pytest.mark.asyncio +async def test_replay_returns_events_after_last_event_id_in_order(): + store = InMemoryMCPEventStore() + ids = [await store.store_event("stream-1", _message(i)) for i in range(5)] + + received, send = _collector() + stream_id = await store.replay_events_after(ids[1], send) + + assert stream_id == "stream-1" + assert [e.event_id for e in received] == ids[2:] + assert [e.message.root.params["seq"] for e in received] == [2, 3, 4] + + +@pytest.mark.asyncio +async def test_replay_unknown_event_id_returns_none_and_sends_nothing(): + store = InMemoryMCPEventStore() + await store.store_event("stream-1", _message(0)) + + received, send = _collector() + assert await store.replay_events_after("nonexistent", send) is None + assert received == [] + + +@pytest.mark.asyncio +async def test_replay_is_scoped_to_the_stream_of_the_last_event(): + store = InMemoryMCPEventStore() + first = await store.store_event("stream-1", _message(1)) + await store.store_event("stream-2", _message(100)) + await store.store_event("stream-1", _message(2)) + + received, send = _collector() + stream_id = await store.replay_events_after(first, send) + + assert stream_id == "stream-1" + assert [e.message.root.params["seq"] for e in received] == [2] + + +@pytest.mark.asyncio +async def test_priming_events_are_not_replayed(): + store = InMemoryMCPEventStore() + first = await store.store_event("stream-1", _message(1)) + await store.store_event("stream-1", None) + await store.store_event("stream-1", _message(2)) + + received, send = _collector() + await store.replay_events_after(first, send) + + assert [e.message.root.params["seq"] for e in received] == [2] + + +@pytest.mark.asyncio +async def test_per_stream_bound_evicts_oldest_events(): + store = InMemoryMCPEventStore(max_events_per_stream=3) + ids = [await store.store_event("stream-1", _message(i)) for i in range(5)] + + received, send = _collector() + assert await store.replay_events_after(ids[0], send) is None + + received, send = _collector() + stream_id = await store.replay_events_after(ids[2], send) + assert stream_id == "stream-1" + assert [e.message.root.params["seq"] for e in received] == [3, 4] + + +@pytest.mark.asyncio +async def test_stream_count_bound_evicts_least_recently_used_stream(): + store = InMemoryMCPEventStore(max_streams=2) + old_id = await store.store_event("stream-old", _message(0)) + mid_id = await store.store_event("stream-mid", _message(1)) + # Touch stream-old so stream-mid becomes the LRU candidate. + kept_id = await store.store_event("stream-old", _message(2)) + await store.store_event("stream-new", _message(3)) + + received, send = _collector() + assert await store.replay_events_after(mid_id, send) is None + assert received == [] + + old_received, send_old = _collector() + assert await store.replay_events_after(old_id, send_old) == "stream-old" + assert [e.event_id for e in old_received] == [kept_id]