From 31fb2b0acfdce1725044faa15465bbf32778fb9f Mon Sep 17 00:00:00 2001
From: Sameer Kankute <sameer@berri.ai>
Date: Thu, 30 Apr 2026 17:59:53 +0530
Subject: [PATCH] Fix tests

---
 .../proxy/_experimental/mcp_server/server.py  | 10 ++++++--
 tests/mcp_tests/test_mcp_server.py            | 25 +++++++++++--------
 .../auth/test_user_api_key_auth_mcp.py        |  6 ++++-
 3 files changed, 28 insertions(+), 13 deletions(-)

diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py
index 4349499d280..489ea31fd67 100644
--- a/litellm/proxy/_experimental/mcp_server/server.py
+++ b/litellm/proxy/_experimental/mcp_server/server.py
@@ -2527,9 +2527,15 @@ if MCP_AVAILABLE:
         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()
+            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 (
+                    header_value.decode()
+                    if isinstance(header_value, bytes)
+                    else str(header_value)
+                )
         return None
 
     def _is_initialize_request(body: bytes) -> bool:
diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py
index 6af07585796..1a50f60c607 100644
--- a/tests/mcp_tests/test_mcp_server.py
+++ b/tests/mcp_tests/test_mcp_server.py
@@ -395,11 +395,11 @@ async def test_mcp_http_transport_tool_not_found():
 @pytest.mark.asyncio
 async def test_streamable_http_mcp_handler_mock():
     """Test the streamable HTTP MCP handler functionality"""
-    from litellm.proxy._types import UserAPIKeyAuth
-
-    # Mock the session manager and its methods
-    mock_session_manager = AsyncMock()
-    mock_session_manager.handle_request = AsyncMock()
+    # Mock streamable HTTP session managers and their methods
+    mock_session_manager_stateless = AsyncMock()
+    mock_session_manager_stateless.handle_request = AsyncMock()
+    mock_session_manager_stateful = AsyncMock()
+    mock_session_manager_stateful.handle_request = AsyncMock()
 
     # Mock scope, receive, send with proper ASGI scope format
     mock_scope = {
@@ -411,7 +411,7 @@ async def test_streamable_http_mcp_handler_mock():
         "server": ("localhost", 8000),
         "scheme": "http",
     }
-    mock_receive = AsyncMock()
+    mock_receive = AsyncMock(return_value={"body": b"{}", "more_body": False})
     mock_send = AsyncMock()
 
     # Mock extract_mcp_auth_context to bypass auth checks in the handler
@@ -423,8 +423,12 @@ async def test_streamable_http_mcp_handler_mock():
             True,
         ),
         patch(
-            "litellm.proxy._experimental.mcp_server.server.session_manager",
-            mock_session_manager,
+            "litellm.proxy._experimental.mcp_server.server.session_manager_stateless",
+            mock_session_manager_stateless,
+        ),
+        patch(
+            "litellm.proxy._experimental.mcp_server.server.session_manager_stateful",
+            mock_session_manager_stateful,
         ),
         patch(
             "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context",
@@ -441,8 +445,9 @@ 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
-        mock_session_manager.handle_request.assert_called_once()
+        # Verify stateless session manager handle_request was called
+        mock_session_manager_stateless.handle_request.assert_called_once()
+        mock_session_manager_stateful.handle_request.assert_not_called()
 
 
 @pytest.mark.asyncio
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
index dd352d0999a..47ec0a0c133 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
@@ -1520,7 +1520,11 @@ def test_mcp_path_based_server_segregation(monkeypatch):
         )
 
     monkeypatch.setattr(
-        "litellm.proxy._experimental.mcp_server.server.session_manager",
+        "litellm.proxy._experimental.mcp_server.server.session_manager_stateless",
+        MagicMock(handle_request=dummy_handle_request),
+    )
+    monkeypatch.setattr(
+        "litellm.proxy._experimental.mcp_server.server.session_manager_stateful",
         MagicMock(handle_request=dummy_handle_request),
     )
     monkeypatch.setattr(