From ca996e18e34fa665a9782c2a8f5db3ffcabb34fa Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Fri, 13 Mar 2026 16:53:06 +0530 Subject: [PATCH] feat(mcp): support stateless and stateful clients via session-id routing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add session_manager_stateful (stateless=False) alongside stateless - Route by mcp-session-id: has ID → stateful, initialize (no ID) → stateful, else → stateless - Peek POST body to detect initialize for routing; replay via wrapped receive - Handle stale session IDs for both managers - Add test_mcp_routing_initialize_to_stateful_no_session_to_stateless - Update test_valid_mcp_session_id_is_preserved, test_concurrent_initialize_session_managers Made-with: Cursor --- .../proxy/_experimental/mcp_server/server.py | 100 +++++++++++-- .../mcp_server/test_mcp_server.py | 133 ++++++++++++++++-- .../mcp_server/test_mcp_stale_session.py | 11 +- 3 files changed, 216 insertions(+), 28 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index cd06de2a2df..e89f791fa1b 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -5,6 +5,7 @@ LiteLLM MCP Server Routes import asyncio import contextlib +import json import time import traceback import uuid @@ -180,13 +181,23 @@ if MCP_AVAILABLE: sse: SseServerTransport = SseServerTransport("/mcp/sse/messages") # Create session managers - session_manager = StreamableHTTPSessionManager( + session_manager_stateless = StreamableHTTPSessionManager( app=server, event_store=None, json_response=False, # enables SSE streaming stateless=True, ) + session_manager_stateful = StreamableHTTPSessionManager( + app=server, + event_store=None, # TODO: Add EventStore for reconnection/event replay if needed + json_response=False, # enables SSE streaming + stateless=False, + ) + + # Keep this alias so existing references to session_manager still work + session_manager = session_manager_stateless + # Create SSE session manager sse_session_manager = StreamableHTTPSessionManager( app=server, @@ -197,11 +208,12 @@ if MCP_AVAILABLE: # Context managers for proper lifecycle management _session_manager_cm = None + _session_manager_stateful_cm = None _sse_session_manager_cm = None async def initialize_session_managers(): """Initialize the session managers. Can be called from main app lifespan.""" - global _SESSION_MANAGERS_INITIALIZED, _session_manager_cm, _sse_session_manager_cm + global _SESSION_MANAGERS_INITIALIZED, _session_manager_cm, _session_manager_stateful_cm, _sse_session_manager_cm # Use async lock to prevent concurrent initialization async with _INITIALIZATION_LOCK: @@ -211,11 +223,13 @@ if MCP_AVAILABLE: verbose_logger.info("Initializing MCP session managers...") # Start the session managers with context managers - _session_manager_cm = session_manager.run() + _session_manager_cm = session_manager_stateless.run() + _session_manager_stateful_cm = session_manager_stateful.run() _sse_session_manager_cm = sse_session_manager.run() # Enter the context managers await _session_manager_cm.__aenter__() + await _session_manager_stateful_cm.__aenter__() await _sse_session_manager_cm.__aenter__() _SESSION_MANAGERS_INITIALIZED = True @@ -225,7 +239,7 @@ if MCP_AVAILABLE: async def shutdown_session_managers(): """Shutdown the session managers.""" - global _SESSION_MANAGERS_INITIALIZED, _session_manager_cm, _sse_session_manager_cm + global _SESSION_MANAGERS_INITIALIZED, _session_manager_cm, _session_manager_stateful_cm, _sse_session_manager_cm if _SESSION_MANAGERS_INITIALIZED: verbose_logger.info("Shutting down MCP session managers...") @@ -233,12 +247,15 @@ if MCP_AVAILABLE: try: if _session_manager_cm: await _session_manager_cm.__aexit__(None, None, None) + if _session_manager_stateful_cm: + await _session_manager_stateful_cm.__aexit__(None, None, None) if _sse_session_manager_cm: await _sse_session_manager_cm.__aexit__(None, None, None) except Exception as e: verbose_logger.exception(f"Error during session manager shutdown: {e}") _session_manager_cm = None + _session_manager_stateful_cm = None _sse_session_manager_cm = None _SESSION_MANAGERS_INITIALIZED = False @@ -304,7 +321,7 @@ if MCP_AVAILABLE: @server.call_tool() async def mcp_server_tool_call( - name: str, arguments: Dict[str, Any] | None + name: str, arguments: Optional[Dict[str, Any]] ) -> CallToolResult: """ Call a specific tool with the provided arguments @@ -347,7 +364,7 @@ if MCP_AVAILABLE: if host_token and hasattr(host_ctx, "session") and host_ctx.session: host_session = host_ctx.session - async def forward_progress(progress: float, total: float | None): + async def forward_progress(progress: float, total: Optional[float]): """Forward progress notifications from external MCP to Host""" try: await host_session.send_progress_notification( @@ -489,7 +506,7 @@ if MCP_AVAILABLE: @server.get_prompt() async def get_prompt( - name: str, arguments: dict[str, str] | None + name: str, arguments: Optional[Dict[str, str]] ) -> GetPromptResult: """ Get a specific prompt with the provided arguments @@ -2270,6 +2287,33 @@ if MCP_AVAILABLE: raw_headers, ) + def _get_session_id_from_scope(scope: Scope) -> Optional[str]: + """ + Extract mcp-session-id from ASGI scope headers. + Returns None if not present. + """ + for header_name, header_value in scope.get("headers", []): + name = header_name if isinstance(header_name, bytes) else header_name.encode() + if name.lower() == b"mcp-session-id": + return header_value.decode() if isinstance(header_value, bytes) else str(header_value) + return None + + def _is_initialize_request(body: bytes) -> bool: + """ + Check if the request body is a JSON-RPC initialize method. + Returns True if method is "initialize", False otherwise or on parse error. + + Note: Assumes the body fits in a single ASGI receive() chunk. MCP initialize + requests are typically small; large chunked bodies may not be fully parsed. + """ + if not body: + return False + try: + data = json.loads(body) + return data.get("method") == "initialize" + except (json.JSONDecodeError, TypeError): + return False + async def _handle_stale_mcp_session( scope: Scope, receive: Receive, @@ -2432,16 +2476,54 @@ if MCP_AVAILABLE: # Give it a moment to start up await asyncio.sleep(0.1) + # Route based on mcp-session-id and request method: + # - Has session ID → stateful (Claude Code, Cursor, VSCode) + # - No session ID + initialize → stateful (so client gets mcp-session-id) + # - No session ID + other → stateless (curl, Inspector, Notion) + session_id = _get_session_id_from_scope(scope) + is_initialize = False + first_msg = None + + if scope.get("method") == "POST" and not session_id: + # Peek at first chunk to detect initialize (assumes body fits in one chunk) + first_msg = await receive() + body = first_msg.get("body", b"") or b"" + is_initialize = _is_initialize_request(body) + + use_stateful = bool(session_id or is_initialize) + target_manager = ( + session_manager_stateful if use_stateful else session_manager_stateless + ) + + verbose_logger.debug( + f"MCP routing to {'stateful' if use_stateful else 'stateless'} manager" + + (f" (session={session_id[:8]}...)" if session_id else "") + + (" (initialize)" if is_initialize else "") + ) + + # Replay first message if we consumed it for peeking + original_receive = receive + if first_msg is not None: + + async def wrapped_receive(): + nonlocal first_msg + if first_msg is not None: + msg, first_msg = first_msg, None + return msg + return await original_receive() + + receive = wrapped_receive + # Handle stale session IDs - either strip them for reconnection # or return success for idempotent DELETE operations handled = await _handle_stale_mcp_session( - scope, receive, send, session_manager + scope, receive, send, target_manager ) if handled: # Request was fully handled (e.g., DELETE on non-existent session) return - await session_manager.handle_request(scope, receive, send) + await target_manager.handle_request(scope, receive, send) except HTTPException: # Re-raise HTTP exceptions to preserve status codes and details raise diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index edc69ad6a4f..d9ef2fb1443 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -720,17 +720,21 @@ async def test_concurrent_initialize_session_managers(): # Reset state before test original_initialized = mcp_server._SESSION_MANAGERS_INITIALIZED original_session_cm = mcp_server._session_manager_cm + original_session_stateful_cm = mcp_server._session_manager_stateful_cm original_sse_session_cm = mcp_server._sse_session_manager_cm try: mcp_server._SESSION_MANAGERS_INITIALIZED = False mcp_server._session_manager_cm = None + mcp_server._session_manager_stateful_cm = None mcp_server._sse_session_manager_cm = None # Mock the session managers to avoid actual MCP initialization with patch( - "litellm.proxy._experimental.mcp_server.server.session_manager" - ) as mock_session_manager, patch( + "litellm.proxy._experimental.mcp_server.server.session_manager_stateless" + ) as mock_session_manager_stateless, patch( + "litellm.proxy._experimental.mcp_server.server.session_manager_stateful" + ) as mock_session_manager_stateful, patch( "litellm.proxy._experimental.mcp_server.server.sse_session_manager" ) as mock_sse_session_manager, patch( "litellm.proxy._experimental.mcp_server.server.verbose_logger" @@ -740,7 +744,8 @@ async def test_concurrent_initialize_session_managers(): mock_cm.__aenter__ = AsyncMock() mock_cm.__aexit__ = AsyncMock() - mock_session_manager.run.return_value = mock_cm + mock_session_manager_stateless.run.return_value = mock_cm + mock_session_manager_stateful.run.return_value = mock_cm mock_sse_session_manager.run.return_value = mock_cm # Create multiple concurrent tasks that call initialize_session_managers @@ -757,18 +762,21 @@ async def test_concurrent_initialize_session_managers(): result == "success" for result in results ), f"Some tasks failed: {results}" - # session_manager.run() should only be called once due to the lock + # Each session manager.run() should only be called once due to the lock assert ( - mock_session_manager.run.call_count == 1 - ), f"Expected 1 call to session_manager.run(), got {mock_session_manager.run.call_count}" + mock_session_manager_stateless.run.call_count == 1 + ), f"Expected 1 call to session_manager_stateless.run(), got {mock_session_manager_stateless.run.call_count}" + assert ( + mock_session_manager_stateful.run.call_count == 1 + ), f"Expected 1 call to session_manager_stateful.run(), got {mock_session_manager_stateful.run.call_count}" assert ( mock_sse_session_manager.run.call_count == 1 ), f"Expected 1 call to sse_session_manager.run(), got {mock_sse_session_manager.run.call_count}" - # The context managers should only be entered once each + # The context managers should only be entered once each (3 managers) assert ( - mock_cm.__aenter__.call_count == 2 - ), f"Expected 2 calls to __aenter__ (one for each session manager), got {mock_cm.__aenter__.call_count}" + mock_cm.__aenter__.call_count == 3 + ), f"Expected 3 calls to __aenter__ (one per session manager), got {mock_cm.__aenter__.call_count}" # State should be properly set assert mcp_server._SESSION_MANAGERS_INITIALIZED is True @@ -777,32 +785,129 @@ async def test_concurrent_initialize_session_managers(): # Restore original state mcp_server._SESSION_MANAGERS_INITIALIZED = original_initialized mcp_server._session_manager_cm = original_session_cm + mcp_server._session_manager_stateful_cm = original_session_stateful_cm mcp_server._sse_session_manager_cm = original_sse_session_cm @pytest.mark.asyncio async def test_streamable_http_session_manager_is_stateless(): """ - Test that the StreamableHTTPSessionManager is initialized with stateless=True. + Test that the StreamableHTTPSessionManager is initialized with both stateless and stateful managers. Regression test for GitHub issue #20242 / PR #19809. When stateless=False, the mcp library rejects non-initialize requests that lack an mcp-session-id header, breaking clients like MCP Inspector, curl, and any HTTP client without automatic session management. + + Now we support both: + - stateless manager for clients without session IDs (curl, Inspector) + - stateful manager for clients with session IDs (Claude Code, Cursor, VSCode) """ try: - from litellm.proxy._experimental.mcp_server.server import session_manager + from litellm.proxy._experimental.mcp_server.server import ( + session_manager_stateful, + session_manager_stateless, + ) except ImportError: pytest.skip("MCP server not available") - # The session manager must be stateless to avoid requiring mcp-session-id + # The stateless session manager must be stateless to avoid requiring mcp-session-id # on every request. This was regressed by PR #19809 (stateless=True -> False). - assert session_manager.stateless is True, ( - "StreamableHTTPSessionManager must be initialized with stateless=True. " + assert session_manager_stateless.stateless is True, ( + "session_manager_stateless must be initialized with stateless=True. " "stateless=False breaks MCP clients that don't manage session IDs. " "See: https://github.com/BerriAI/litellm/issues/20242" ) + # The stateful session manager must be stateful to support progress notifications + assert session_manager_stateful.stateless is False, ( + "session_manager_stateful must be initialized with stateless=False. " + "stateless=True breaks progress notifications for clients that manage session IDs." + ) + + +@pytest.mark.asyncio +async def test_mcp_routing_initialize_to_stateful_no_session_to_stateless(): + """ + Test that routing correctly sends: + - initialize (no mcp-session-id) → stateful manager (so client gets mcp-session-id) + - tools/list (no mcp-session-id) → stateless manager (curl, Inspector) + """ + try: + from litellm.proxy._experimental.mcp_server.server import ( + handle_streamable_http_mcp, + session_manager_stateful, + session_manager_stateless, + ) + except ImportError: + pytest.skip("MCP server not available") + + async def make_request(method_body: bytes, path: str = "/mcp/progress_test"): + scope = { + "type": "http", + "method": "POST", + "path": path, + "headers": [ + (b"content-type", b"application/json"), + (b"authorization", b"Bearer test-key"), + ], + } + receive = AsyncMock(return_value={"type": "http.request", "body": method_body, "more_body": False}) + send = AsyncMock() + + stateless_called = [] + stateful_called = [] + + async def stateless_handle(s, r, se): + stateless_called.append(1) + + async def stateful_handle(s, r, se): + stateful_called.append(1) + + with patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=(MagicMock(), None, ["progress_test"], None, None, None), + ), patch( + "litellm.proxy._experimental.mcp_server.server.set_auth_context", + ), patch( + "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", + True, + ), patch.object( + session_manager_stateless, + "handle_request", + side_effect=stateless_handle, + ), patch.object( + session_manager_stateful, + "handle_request", + side_effect=stateful_handle, + ), patch.object( + session_manager_stateless, + "_server_instances", + {}, + ), patch.object( + session_manager_stateful, + "_server_instances", + {}, + ): + await handle_streamable_http_mcp(scope, receive, send) + + return bool(stateless_called), bool(stateful_called) + + # initialize → stateful + init_body = b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2024-11-05"}}' + stateless_called, stateful_called = await make_request(init_body) + assert stateful_called and not stateless_called, ( + "initialize (no session) should route to stateful, not stateless" + ) + + # tools/list → stateless + tools_body = b'{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}' + stateless_called, stateful_called = await make_request(tools_body) + assert stateless_called and not stateful_called, ( + "tools/list (no session) should route to stateless, not stateful" + ) + @pytest.mark.asyncio @pytest.mark.no_parallel diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py index 5eb8c1e51ac..bf00358db36 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py @@ -7,9 +7,10 @@ they may send a stale `mcp-session-id` header. This test verifies that: 2. For DELETE requests: idempotent behavior returns success even if session doesn't exist """ -import pytest from unittest.mock import AsyncMock, MagicMock, patch +import pytest + class TestHandleStaleMcpSession: """Unit tests for the _handle_stale_mcp_session helper.""" @@ -326,7 +327,7 @@ async def test_valid_mcp_session_id_is_preserved(): try: from litellm.proxy._experimental.mcp_server.server import ( handle_streamable_http_mcp, - session_manager, + session_manager_stateful, ) except ImportError: pytest.skip("MCP server not available") @@ -352,7 +353,7 @@ async def test_valid_mcp_session_id_is_preserved(): async def mock_handle_request(s, r, se): captured_scope.update(s) - # Session manager HAS this session + # Stateful session manager HAS this session (requests with mcp-session-id route there) mock_instances = {valid_session_id: MagicMock()} with patch( @@ -365,11 +366,11 @@ async def test_valid_mcp_session_id_is_preserved(): "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", True, ), patch.object( - session_manager, + session_manager_stateful, "handle_request", side_effect=mock_handle_request, ), patch.object( - session_manager, + session_manager_stateful, "_server_instances", mock_instances, ):