From 29bb73ffcace33e65f39429493a10b1ecb76c909 Mon Sep 17 00:00:00 2001 From: Gaurav Singh <103016722+gavksingh@users.noreply.github.com> Date: Fri, 27 Feb 2026 13:59:08 -0500 Subject: [PATCH] fix(mcp): strip stale mcp-session-id header to prevent 400 in multi-worker deployments (#20992) (#21417) In a multi-worker Uvicorn setup, a client that reconnects to a different worker sends an mcp-session-id that the new worker has never seen. The MCP SDK returns 400 because the session is unknown. Fix: add _handle_stale_mcp_session() which inspects the inbound mcp-session-id header before the request reaches the SDK. If the session is not in this worker's _server_instances: - Non-DELETE: strip the header so the SDK creates a fresh session - DELETE: return 200 immediately (idempotent, session already gone) No new dependencies, no Redis, no latency added to the hot path. Fixes https://github.com/BerriAI/litellm/issues/20992 --- .../proxy/_experimental/mcp_server/server.py | 99 ++++++++++++------- tests/mcp_tests/test_mcp_server.py | 3 +- 2 files changed, 64 insertions(+), 38 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 48c837c1e4f..5b3d5bd60e2 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 traceback import uuid from datetime import datetime @@ -84,6 +85,7 @@ except ImportError as e: _SESSION_MANAGERS_INITIALIZED = False _INITIALIZATION_LOCK = asyncio.Lock() + if MCP_AVAILABLE: from mcp.server import Server @@ -1919,65 +1921,86 @@ if MCP_AVAILABLE: mgr: "StreamableHTTPSessionManager", ) -> bool: """ - Handle stale MCP session IDs to prevent "Session not found" errors. - - When clients reconnect after a server restart or session cleanup, they may - send a session ID that no longer exists. This function handles two scenarios: - - 1. Non-DELETE requests: Strip the stale session ID header so the session - manager creates a fresh session transparently. - - 2. DELETE requests: Return success (200) immediately for idempotent behavior, - since the desired state (session doesn't exist) is already achieved. + Inspect the incoming ``mcp-session-id`` header **before** the + request reaches the MCP SDK. If the session is stale (not known + to this worker), strip the header so the SDK creates a fresh + stateless session instead of returning a 400. Returns: - True if the request was handled (DELETE on non-existent session) - False if the request should continue to the session manager + True if the request was fully handled (e.g. DELETE on + non-existent session). False if the request should continue + to the session manager. - Fixes https://github.com/BerriAI/litellm/issues/20292 + Fixes https://github.com/BerriAI/litellm/issues/20992 """ _mcp_session_header = b"mcp-session-id" + _headers = scope.get("headers", []) + + def _normalize_header_name(header_name: Any) -> Optional[bytes]: + if isinstance(header_name, bytes): + return header_name.lower() + if isinstance(header_name, str): + return header_name.lower().encode("utf-8", errors="replace") + return None + _session_id: Optional[str] = None - for header_name, header_value in scope.get("headers", []): - if header_name == _mcp_session_header: - _session_id = header_value.decode("utf-8", errors="replace") + for header_name, header_value in _headers: + if _normalize_header_name(header_name) == _mcp_session_header: + if isinstance(header_value, bytes): + _session_id = header_value.decode("utf-8", errors="replace") + else: + _session_id = str(header_value) break if _session_id is None: return False + # Check in-memory session tracking known_sessions = getattr(mgr, "_server_instances", None) - if known_sessions is None or _session_id in known_sessions: - # Session exists or we can't check - let the session manager handle it + # If we cannot inspect known_sessions, let the manager handle it + if known_sessions is None: return False - # Session doesn't exist - handle based on request method + # If session exists in this worker's memory, let the manager handle it + try: + if _session_id in known_sessions: + return False + except Exception: + verbose_logger.debug( + "Unable to inspect active MCP sessions for '%s'. " + "Deferring to session manager.", + _session_id, + ) + return False + + # --- Session not in this worker's memory --- method = scope.get("method", "").upper() - + if method == "DELETE": - # Idempotent DELETE: session doesn't exist, return success verbose_logger.info( - f"DELETE request for non-existent MCP session '{_session_id}'. " - "Returning success (idempotent DELETE)." + "DELETE request for non-existent MCP session '%s'. " + "Returning success (idempotent DELETE).", + _session_id, ) success_response = JSONResponse( status_code=200, - content={"message": "Session terminated successfully"} + content={"message": "Session terminated successfully"}, ) await success_response(scope, receive, send) return True - else: - # Non-DELETE: strip stale session ID to allow new session creation - verbose_logger.warning( - "MCP session ID '%s' not found in active sessions. " - "Stripping stale header to force new session creation.", - _session_id, - ) - scope["headers"] = [ - (k, v) for k, v in scope["headers"] - if k != _mcp_session_header - ] - return False + + # Non-DELETE: strip stale session ID to allow new session creation + verbose_logger.warning( + "MCP session ID '%s' not found in this worker's memory. " + "Stripping stale header to force new session creation.", + _session_id, + ) + scope["headers"] = [ + (k, v) + for k, v in _headers + if _normalize_header_name(k) != _mcp_session_header + ] + return False async def handle_streamable_http_mcp( scope: Scope, receive: Receive, send: Send @@ -2055,7 +2078,9 @@ if MCP_AVAILABLE: # 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) + handled = await _handle_stale_mcp_session( + scope, receive, send, session_manager + ) if handled: # Request was fully handled (e.g., DELETE on non-existent session) return diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index 930b0a03042..dc1e2068365 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -427,7 +427,8 @@ async def test_streamable_http_mcp_handler_mock(): # Call the handler await handle_streamable_http_mcp(mock_scope, mock_receive, mock_send) - # Verify session manager handle_request was called + # Verify session manager handle_request was called with correct args + # send is passed directly (no wrapper) mock_session_manager.handle_request.assert_called_once_with( mock_scope, mock_receive, mock_send )