Merge pull request #17342 from BerriAI/litellm_fix_mcp_auth_header_forwarding

Fix: litellm user auth not passing issue

Author: Sameerlite

litellm/proxy/_experimental/mcp_server/mcp_server_manager.py

@@ -17,16 +17,15 @@
 from fastapi import HTTPException
 from httpx import HTTPStatusError
 from mcp import ReadResourceResult, Resource
+from mcp.types import CallToolRequestParams as MCPCallToolRequestParams
 from mcp.types import (
-    CallToolRequestParams as MCPCallToolRequestParams,
+    CallToolResult,
     GetPromptRequestParams,
     GetPromptResult,
     Prompt,
     ResourceTemplate,
 )
-from mcp.types import CallToolResult
 from mcp.types import Tool as MCPTool
-
 from pydantic import AnyUrl
 
 import litellm
@@ -1949,7 +1948,12 @@ def _get_mcp_server_from_tool_name(self, tool_name: str) -> Optional[MCPServer]:
             ) = split_server_prefix_from_name(tool_name)
             if original_tool_name in self.tool_name_to_mcp_server_name_mapping:
                 for server in self.get_registry().values():
-                    if normalize_server_name(server.name) == normalize_server_name(
+                    if server.server_name is None:
+                        if normalize_server_name(server.name) == normalize_server_name(
+                            server_name_from_prefix
+                        ):
+                            return server
+                    elif normalize_server_name(server.server_name) == normalize_server_name(
                         server_name_from_prefix
                     ):
                         return server

litellm/proxy/litellm_pre_call_utils.py

@@ -611,6 +611,8 @@ def add_user_api_key_auth_to_request_metadata(
         data[_metadata_variable_name]["user_api_end_user_max_budget"] = getattr(
             user_api_key_dict, "end_user_max_budget", None
         )
+        # Add the full UserAPIKeyAuth object for MCP server access control
+        data[_metadata_variable_name]["user_api_key_auth"] = user_api_key_dict
         return data
 
     @staticmethod

litellm/responses/main.py

@@ -22,6 +22,9 @@
 from litellm._logging import verbose_logger
 from litellm.constants import request_timeout
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+from litellm.litellm_core_utils.prompt_templates.common_utils import (
+    update_responses_input_with_model_file_ids,
+)
 from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
 from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
 from litellm.responses.litellm_completion_transformation.handler import (
@@ -38,9 +41,6 @@
     ToolChoice,
     ToolParam,
 )
-from litellm.litellm_core_utils.prompt_templates.common_utils import (
-    update_responses_input_with_model_file_ids,
-)
 
 # Handle ResponseText import with fallback
 if TYPE_CHECKING:
@@ -168,7 +168,8 @@ async def aresponses_api_with_mcp(
     ) = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools)
 
     # Process MCP tools through the complete pipeline (fetch + filter + deduplicate + transform)
-    user_api_key_auth = kwargs.get("user_api_key_auth")
+    # Extract user_api_key_auth from litellm_metadata (where it's added by add_user_api_key_auth_to_request_metadata)
+    user_api_key_auth = kwargs.get("user_api_key_auth") or kwargs.get("litellm_metadata", {}).get("user_api_key_auth")
 
     # Get original MCP tools (for events) and OpenAI tools (for LLM) by reusing existing methods
     (

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py

@@ -9,6 +9,8 @@
 sys.path.insert(0, "../../../../../")
 
 import httpx
+from mcp import ReadResourceResult, Resource
+from mcp.types import GetPromptResult, Prompt, ResourceTemplate, TextResourceContents
 
 from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
     MCPServerManager,
@@ -17,8 +19,6 @@
 from litellm.proxy._types import LiteLLM_MCPServerTable, MCPTransport
 from litellm.types.mcp import MCPAuth
 from litellm.types.mcp_server.mcp_server_manager import MCPOAuthMetadata, MCPServer
-from mcp import ReadResourceResult, Resource
-from mcp.types import GetPromptResult, Prompt, ResourceTemplate, TextResourceContents
 
 
 class TestMCPServerManager:
@@ -1606,6 +1606,104 @@ async def mock_call_tool(params):
         # Verify the MCP client call was awaited exactly once
         assert mock_client.call_tool.await_count == 1
 
+    @pytest.mark.asyncio
+    async def test_get_allowed_mcp_servers_with_user_api_key_auth(self):
+        """
+        Test that get_allowed_mcp_servers properly receives and uses user_api_key_auth
+        when called. This verifies the fix where user_api_key_auth is passed through
+        litellm_metadata from responses API.
+        """
+        from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
+            MCPRequestHandler,
+        )
+        from litellm.proxy._types import LiteLLM_ObjectPermissionTable, UserAPIKeyAuth
+
+        manager = MCPServerManager()
+
+        # Create a mock user_api_key_auth with object_permission
+        object_permission = LiteLLM_ObjectPermissionTable(
+            object_permission_id="perm_123",
+            mcp_servers=["test_server_1", "test_server_2"],
+            mcp_access_groups=[],
+        )
+
+        user_api_key_auth = UserAPIKeyAuth(
+            api_key="sk-test",
+            user_id="user-123",
+            object_permission=object_permission,
+            object_permission_id="perm_123",
+        )
+
+        # Mock MCPRequestHandler.get_allowed_mcp_servers to verify it receives user_api_key_auth
+        with patch.object(
+            MCPRequestHandler,
+            "get_allowed_mcp_servers",
+            new_callable=AsyncMock,
+        ) as mock_get_allowed:
+            # Configure mock to return servers from object_permission
+            mock_get_allowed.return_value = ["test_server_1", "test_server_2"]
+
+            # Call get_allowed_mcp_servers with user_api_key_auth
+            result = await manager.get_allowed_mcp_servers(user_api_key_auth)
+
+            # Verify MCPRequestHandler.get_allowed_mcp_servers was called with user_api_key_auth
+            mock_get_allowed.assert_called_once()
+            call_args = mock_get_allowed.call_args
+            assert call_args[0][0] is user_api_key_auth  # First positional arg should be user_api_key_auth
+            assert call_args[0][0].user_id == "user-123"
+            assert call_args[0][0].object_permission_id == "perm_123"
+            assert call_args[0][0].object_permission is not None
+            assert call_args[0][0].object_permission.mcp_servers == ["test_server_1", "test_server_2"]
+
+            # Verify result contains the expected servers
+            assert "test_server_1" in result
+            assert "test_server_2" in result
+
+    def test_get_mcp_server_from_tool_name_uses_server_name_not_name(self):
+        """
+        Test that _get_mcp_server_from_tool_name uses server.server_name instead of server.name
+        when extracting server name from prefixed tool name (second case).
+        This ensures the fix for using server_name instead of name works correctly.
+        """
+        from litellm.proxy._experimental.mcp_server.utils import (
+            add_server_prefix_to_name,
+        )
+
+        manager = MCPServerManager()
+
+        # Create a server where server_name differs from name
+        # This tests the scenario where server.name != server.server_name
+        server = MCPServer(
+            server_id="test-server-id",
+            name="Test Server Name",  # Different from server_name
+            server_name="test_server",  # This is what should be used
+            alias="test_server",
+            transport=MCPTransport.http,
+        )
+
+        # Register the server
+        manager.registry = {server.server_id: server}
+
+        # Create a tool with prefixed name
+        tool_name = "test_tool"
+        prefixed_tool_name = add_server_prefix_to_name(tool_name, "test_server")
+
+        # Populate the mapping with the original tool name
+        manager.tool_name_to_mcp_server_name_mapping[tool_name] = "test_server"
+        manager.tool_name_to_mcp_server_name_mapping[prefixed_tool_name] = "test_server"
+
+        # Test: _get_mcp_server_from_tool_name should find the server using server.server_name
+        # even when server.name is different
+        resolved_server = manager._get_mcp_server_from_tool_name(prefixed_tool_name)
+
+        # Verify the server was found correctly
+        assert resolved_server is not None
+        assert resolved_server.server_id == server.server_id
+        assert resolved_server.server_name == "test_server"
+        # Verify it matched using server_name, not name
+        assert resolved_server.name == "Test Server Name"  # name is different
+        assert resolved_server.server_name == "test_server"  # server_name matches
+
 
 if __name__ == "__main__":
     pytest.main([__file__])