Fix tests

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:

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

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(