diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 2c03cbdae3..d42fd80cbd 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -17,16 +17,15 @@ from urllib.parse import urlparse 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 @@ class MCPServerManager: ) = 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 diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 0a7fc62a42..9dc255bd79 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -611,6 +611,8 @@ class LiteLLMProxyRequestSetup: 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 diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 950ea7063f..e837346df2 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -22,6 +22,9 @@ import litellm 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 @@ from litellm.types.llms.openai import ( 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 ( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 940dcf1a7a..86e25c46c8 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -9,6 +9,8 @@ from fastapi import HTTPException 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._experimental.mcp_server.mcp_server_manager import ( 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 @@ class TestMCPServerManager: # 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__])