From 1a2aa62c0fc5d2bcd9dfdde389e60aa96285c5d0 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 30 Apr 2026 16:42:23 +0530 Subject: [PATCH] fix(mcp): respect stateful routing and harden initialize detection Ensure streamable MCP requests are dispatched via the computed target session manager, and guard initialize detection against non-object JSON bodies. Update stale-session test patches to target the stateful manager so routing assertions remain correct. Made-with: Cursor --- litellm/proxy/_experimental/mcp_server/server.py | 4 ++-- .../proxy/_experimental/mcp_server/test_mcp_stale_session.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index d43b964e85b..4349499d280 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -2544,7 +2544,7 @@ if MCP_AVAILABLE: return False try: data = json.loads(body) - return data.get("method") == "initialize" + return isinstance(data, dict) and data.get("method") == "initialize" except (json.JSONDecodeError, TypeError): return False @@ -2848,7 +2848,7 @@ if MCP_AVAILABLE: mcp_servers, _client_ip, ): - 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_stale_session.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py index 8b0c27f6ee4..f4a9b118095 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 @@ -384,12 +384,12 @@ async def test_valid_mcp_session_id_is_preserved(): 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, ),