diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 6d54c3871e5..d35ea26c23a 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -1896,6 +1896,32 @@ if MCP_AVAILABLE: # Give it a moment to start up await asyncio.sleep(0.1) + # Handle stale mcp-session-id headers (Fixes #20292) + # When clients like VSCode reconnect after a reload, they may send a + # stale mcp-session-id that no longer exists in the session manager. + # This causes a 404 "Session not found" error loop. Strip the header + # so the session manager creates a fresh session instead. + _mcp_session_header = b"mcp-session-id" + _stale_session_id: Optional[str] = None + for header_name, header_value in scope.get("headers", []): + if header_name == _mcp_session_header: + _stale_session_id = header_value.decode("utf-8", errors="replace") + break + + if _stale_session_id is not None: + # Check if this session ID exists in the session manager + _known_sessions = getattr(session_manager, "_server_instances", None) + if _known_sessions is not None and _stale_session_id not in _known_sessions: + verbose_logger.warning( + "MCP session ID '%s' not found in active sessions. " + "Stripping stale mcp-session-id header to force new session creation.", + _stale_session_id, + ) + scope["headers"] = [ + (k, v) for k, v in scope["headers"] + if k != _mcp_session_header + ] + await session_manager.handle_request(scope, receive, send) except Exception as e: raise e 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 new file mode 100644 index 00000000000..6efe711ff63 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py @@ -0,0 +1,199 @@ +""" +Tests for MCP stale session ID handling (Fixes #20292). + +When VSCode reconnects to LiteLLM's MCP endpoint after a reload, it sends a stale +`mcp-session-id` header. The session manager returns a 404 because the old session +was cleaned up. This test verifies that stale session IDs are detected and stripped +so a new session is created automatically. +""" + +import pytest +from unittest.mock import AsyncMock, MagicMock, patch + + +@pytest.mark.asyncio +async def test_stale_mcp_session_id_is_stripped(): + """ + When the mcp-session-id header references a session that no longer exists, + handle_streamable_http_mcp should strip the header before forwarding the + request to the session manager so a fresh session is created. + """ + try: + from litellm.proxy._experimental.mcp_server.server import ( + handle_streamable_http_mcp, + session_manager, + ) + except ImportError: + pytest.skip("MCP server not available") + + stale_session_id = "stale-session-id-12345" + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp", + "headers": [ + (b"content-type", b"application/json"), + (b"mcp-session-id", stale_session_id.encode()), + (b"authorization", b"Bearer test-key"), + ], + } + + receive = AsyncMock() + send = AsyncMock() + + # Simulate: session manager has NO sessions (the stale one was cleaned up) + captured_scope = {} + + async def mock_handle_request(s, r, se): + # Capture the scope that was actually passed + captured_scope.update(s) + + with patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=(MagicMock(), None, None, 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, + "handle_request", + side_effect=mock_handle_request, + ), patch.object( + session_manager, + "_server_instances", + {}, # Empty dict = no active sessions + ): + await handle_streamable_http_mcp(scope, receive, send) + + # Verify the mcp-session-id header was stripped + header_names = [k for k, v in captured_scope.get("headers", [])] + assert b"mcp-session-id" not in header_names, ( + "Stale mcp-session-id header should have been stripped from the scope" + ) + + +@pytest.mark.asyncio +async def test_valid_mcp_session_id_is_preserved(): + """ + When the mcp-session-id header references a session that still exists, + handle_streamable_http_mcp should NOT strip the header. + """ + try: + from litellm.proxy._experimental.mcp_server.server import ( + handle_streamable_http_mcp, + session_manager, + ) + except ImportError: + pytest.skip("MCP server not available") + + valid_session_id = "valid-session-id-67890" + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp", + "headers": [ + (b"content-type", b"application/json"), + (b"mcp-session-id", valid_session_id.encode()), + (b"authorization", b"Bearer test-key"), + ], + } + + receive = AsyncMock() + send = AsyncMock() + + captured_scope = {} + + async def mock_handle_request(s, r, se): + captured_scope.update(s) + + # Session manager HAS this session + mock_instances = {valid_session_id: MagicMock()} + + with patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=(MagicMock(), None, None, 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, + "handle_request", + side_effect=mock_handle_request, + ), patch.object( + session_manager, + "_server_instances", + mock_instances, + ): + await handle_streamable_http_mcp(scope, receive, send) + + # Verify the mcp-session-id header was preserved + header_names = [k for k, v in captured_scope.get("headers", [])] + assert b"mcp-session-id" in header_names, ( + "Valid mcp-session-id header should have been preserved" + ) + + +@pytest.mark.asyncio +async def test_no_mcp_session_id_header_works_normally(): + """ + When no mcp-session-id header is present (initial connection), + handle_streamable_http_mcp should work without any issues. + """ + try: + from litellm.proxy._experimental.mcp_server.server import ( + handle_streamable_http_mcp, + session_manager, + ) + except ImportError: + pytest.skip("MCP server not available") + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp", + "headers": [ + (b"content-type", b"application/json"), + (b"authorization", b"Bearer test-key"), + ], + } + + receive = AsyncMock() + send = AsyncMock() + + captured_scope = {} + + async def mock_handle_request(s, r, se): + captured_scope.update(s) + + with patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=(MagicMock(), None, None, 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, + "handle_request", + side_effect=mock_handle_request, + ), patch.object( + session_manager, + "_server_instances", + {}, + ): + await handle_streamable_http_mcp(scope, receive, send) + + # Verify headers are unchanged (no mcp-session-id was added or anything weird) + header_names = [k for k, v in captured_scope.get("headers", [])] + assert b"mcp-session-id" not in header_names + assert b"content-type" in header_names