Fix: litellm user auth not passing issue

This commit is contained in:
Sameer Kankute 2025-12-02 11:25:32 +05:30
parent 4c7a988454
commit 082c8af37f
5 changed files with 452 additions and 33 deletions

View file

@ -6717,6 +6717,33 @@
"supports_vision": true,
"tool_use_system_prompt_tokens": 159
},
"claude-opus-4-5": {
"cache_creation_input_token_cost": 6.25e-06,
"cache_creation_input_token_cost_above_1hr": 1e-05,
"cache_read_input_token_cost": 5e-07,
"input_cost_per_token": 5e-06,
"litellm_provider": "anthropic",
"max_input_tokens": 200000,
"max_output_tokens": 64000,
"max_tokens": 64000,
"mode": "chat",
"output_cost_per_token": 2.5e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.01,
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
"supports_assistant_prefill": true,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 159
},
"claude-sonnet-4-20250514": {
"deprecation_date": "2026-05-14",
"cache_creation_input_token_cost": 3.75e-06,
@ -7824,26 +7851,298 @@
"source": "https://www.databricks.com/product/pricing/foundation-model-serving"
},
"databricks/databricks-claude-3-7-sonnet": {
"input_cost_per_token": 2.5e-06,
"input_dbu_cost_per_token": 3.571e-05,
"input_cost_per_token": 2.9999900000000002e-06,
"input_dbu_cost_per_token": 4.2857e-05,
"litellm_provider": "databricks",
"max_input_tokens": 200000,
"max_output_tokens": 128000,
"max_tokens": 200000,
"metadata": {
"notes": "Input/output cost per token is dbu cost * $0.070, based on databricks Claude 3.7 conversion. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
"notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 1.7857e-05,
"output_db_cost_per_token": 0.000214286,
"source": "https://www.databricks.com/product/pricing/foundation-model-serving",
"output_cost_per_token": 1.5000020000000002e-05,
"output_dbu_cost_per_token": 0.000214286,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_assistant_prefill": true,
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
},
"databricks/databricks-claude-haiku-4-5": {
"input_cost_per_token": 1.00002e-06,
"input_dbu_cost_per_token": 1.4286e-05,
"litellm_provider": "databricks",
"max_input_tokens": 200000,
"max_output_tokens": 64000,
"max_tokens": 200000,
"metadata": {
"notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 5.00003e-06,
"output_dbu_cost_per_token": 7.1429e-05,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_assistant_prefill": true,
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
},
"databricks/databricks-claude-opus-4": {
"input_cost_per_token": 1.5000020000000002e-05,
"input_dbu_cost_per_token": 0.000214286,
"litellm_provider": "databricks",
"max_input_tokens": 200000,
"max_output_tokens": 32000,
"max_tokens": 200000,
"metadata": {
"notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 7.500003000000001e-05,
"output_dbu_cost_per_token": 0.001071429,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_assistant_prefill": true,
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
},
"databricks/databricks-claude-opus-4-1": {
"input_cost_per_token": 1.5000020000000002e-05,
"input_dbu_cost_per_token": 0.000214286,
"litellm_provider": "databricks",
"max_input_tokens": 200000,
"max_output_tokens": 32000,
"max_tokens": 200000,
"metadata": {
"notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 7.500003000000001e-05,
"output_dbu_cost_per_token": 0.001071429,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_assistant_prefill": true,
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
},
"databricks/databricks-claude-opus-4-5": {
"input_cost_per_token": 5.00003e-06,
"input_dbu_cost_per_token": 7.1429e-05,
"litellm_provider": "databricks",
"max_input_tokens": 200000,
"max_output_tokens": 64000,
"max_tokens": 200000,
"metadata": {
"notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 2.5000010000000002e-05,
"output_dbu_cost_per_token": 0.000357143,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_assistant_prefill": true,
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
},
"databricks/databricks-claude-sonnet-4": {
"input_cost_per_token": 2.9999900000000002e-06,
"input_dbu_cost_per_token": 4.2857e-05,
"litellm_provider": "databricks",
"max_input_tokens": 200000,
"max_output_tokens": 64000,
"max_tokens": 200000,
"metadata": {
"notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 1.5000020000000002e-05,
"output_dbu_cost_per_token": 0.000214286,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_assistant_prefill": true,
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
},
"databricks/databricks-claude-sonnet-4-1": {
"input_cost_per_token": 2.9999900000000002e-06,
"input_dbu_cost_per_token": 4.2857e-05,
"litellm_provider": "databricks",
"max_input_tokens": 200000,
"max_output_tokens": 64000,
"max_tokens": 200000,
"metadata": {
"notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 1.5000020000000002e-05,
"output_dbu_cost_per_token": 0.000214286,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_assistant_prefill": true,
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
},
"databricks/databricks-claude-sonnet-4-5": {
"input_cost_per_token": 2.9999900000000002e-06,
"input_dbu_cost_per_token": 4.2857e-05,
"litellm_provider": "databricks",
"max_input_tokens": 200000,
"max_output_tokens": 64000,
"max_tokens": 200000,
"metadata": {
"notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 1.5000020000000002e-05,
"output_dbu_cost_per_token": 0.000214286,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_assistant_prefill": true,
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
},
"databricks/databricks-gemini-2-5-flash": {
"input_cost_per_token": 3.0001999999999996e-07,
"input_dbu_cost_per_token": 4.285999999999999e-06,
"litellm_provider": "databricks",
"max_input_tokens": 1048576,
"max_output_tokens": 65535,
"max_tokens": 1048576,
"metadata": {
"notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 2.49998e-06,
"output_dbu_cost_per_token": 3.5714e-05,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_function_calling": true,
"supports_tool_choice": true
},
"databricks/databricks-gemini-2-5-pro": {
"input_cost_per_token": 1.24999e-06,
"input_dbu_cost_per_token": 1.7857e-05,
"litellm_provider": "databricks",
"max_input_tokens": 1048576,
"max_output_tokens": 65536,
"max_tokens": 1048576,
"metadata": {
"notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 9.999990000000002e-06,
"output_dbu_cost_per_token": 0.000142857,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_function_calling": true,
"supports_tool_choice": true
},
"databricks/databricks-gemma-3-12b": {
"input_cost_per_token": 1.5000999999999998e-07,
"input_dbu_cost_per_token": 2.1429999999999996e-06,
"litellm_provider": "databricks",
"max_input_tokens": 128000,
"max_output_tokens": 32000,
"max_tokens": 128000,
"metadata": {
"notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 5.0001e-07,
"output_dbu_cost_per_token": 7.143e-06,
"source": "https://www.databricks.com/product/pricing/foundation-model-serving"
},
"databricks/databricks-gpt-5": {
"input_cost_per_token": 1.24999e-06,
"input_dbu_cost_per_token": 1.7857e-05,
"litellm_provider": "databricks",
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"max_tokens": 400000,
"metadata": {
"notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 9.999990000000002e-06,
"output_dbu_cost_per_token": 0.000142857,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving"
},
"databricks/databricks-gpt-5-1": {
"input_cost_per_token": 1.24999e-06,
"input_dbu_cost_per_token": 1.7857e-05,
"litellm_provider": "databricks",
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"max_tokens": 400000,
"metadata": {
"notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 9.999990000000002e-06,
"output_dbu_cost_per_token": 0.000142857,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving"
},
"databricks/databricks-gpt-5-mini": {
"input_cost_per_token": 2.4997000000000006e-07,
"input_dbu_cost_per_token": 3.571e-06,
"litellm_provider": "databricks",
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"max_tokens": 400000,
"metadata": {
"notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 1.9999700000000004e-06,
"output_dbu_cost_per_token": 2.8571e-05,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving"
},
"databricks/databricks-gpt-5-nano": {
"input_cost_per_token": 4.998e-08,
"input_dbu_cost_per_token": 7.14e-07,
"litellm_provider": "databricks",
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"max_tokens": 400000,
"metadata": {
"notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 3.9998000000000007e-07,
"output_dbu_cost_per_token": 5.714000000000001e-06,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving"
},
"databricks/databricks-gpt-oss-120b": {
"input_cost_per_token": 1.5000999999999998e-07,
"input_dbu_cost_per_token": 2.1429999999999996e-06,
"litellm_provider": "databricks",
"max_input_tokens": 131072,
"max_output_tokens": 131072,
"max_tokens": 131072,
"metadata": {
"notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 5.9997e-07,
"output_dbu_cost_per_token": 8.571e-06,
"source": "https://www.databricks.com/product/pricing/foundation-model-serving"
},
"databricks/databricks-gpt-oss-20b": {
"input_cost_per_token": 7e-08,
"input_dbu_cost_per_token": 1e-06,
"litellm_provider": "databricks",
"max_input_tokens": 131072,
"max_output_tokens": 131072,
"max_tokens": 131072,
"metadata": {
"notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 3.0001999999999996e-07,
"output_dbu_cost_per_token": 4.285999999999999e-06,
"source": "https://www.databricks.com/product/pricing/foundation-model-serving"
},
"databricks/databricks-gte-large-en": {
"input_cost_per_token": 1.2999e-07,
"input_cost_per_token": 1.2999000000000001e-07,
"input_dbu_cost_per_token": 1.857e-06,
"litellm_provider": "databricks",
"max_input_tokens": 8192,
@ -7868,14 +8167,14 @@
"notes": "Input/output cost per token is dbu cost * $0.070, based on databricks Llama 3.1 70B conversion. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 1.5e-06,
"output_cost_per_token": 1.5000300000000002e-06,
"output_dbu_cost_per_token": 2.1429e-05,
"source": "https://www.databricks.com/product/pricing/foundation-model-serving",
"supports_tool_choice": true
},
"databricks/databricks-llama-4-maverick": {
"input_cost_per_token": 5e-06,
"input_dbu_cost_per_token": 7.143e-05,
"input_cost_per_token": 5.0001e-07,
"input_dbu_cost_per_token": 7.143e-06,
"litellm_provider": "databricks",
"max_input_tokens": 128000,
"max_output_tokens": 128000,
@ -7884,13 +8183,13 @@
"notes": "Databricks documentation now provides both DBU costs (_dbu_cost_per_token) and dollar costs(_cost_per_token)."
},
"mode": "chat",
"output_cost_per_token": 1.5e-05,
"output_dbu_cost_per_token": 0.00021429,
"output_cost_per_token": 1.5000300000000002e-06,
"output_dbu_cost_per_token": 2.1429e-05,
"source": "https://www.databricks.com/product/pricing/foundation-model-serving",
"supports_tool_choice": true
},
"databricks/databricks-meta-llama-3-1-405b-instruct": {
"input_cost_per_token": 5e-06,
"input_cost_per_token": 5.00003e-06,
"input_dbu_cost_per_token": 7.1429e-05,
"litellm_provider": "databricks",
"max_input_tokens": 128000,
@ -7900,14 +8199,29 @@
"notes": "Input/output cost per token is dbu cost * $0.070, based on databricks Llama 3.1 70B conversion. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 1.500002e-05,
"output_db_cost_per_token": 0.000214286,
"output_cost_per_token": 1.5000020000000002e-05,
"output_dbu_cost_per_token": 0.000214286,
"source": "https://www.databricks.com/product/pricing/foundation-model-serving",
"supports_tool_choice": true
},
"databricks/databricks-meta-llama-3-1-8b-instruct": {
"input_cost_per_token": 1.5000999999999998e-07,
"input_dbu_cost_per_token": 2.1429999999999996e-06,
"litellm_provider": "databricks",
"max_input_tokens": 200000,
"max_output_tokens": 128000,
"max_tokens": 200000,
"metadata": {
"notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 4.5003000000000007e-07,
"output_dbu_cost_per_token": 6.429000000000001e-06,
"source": "https://www.databricks.com/product/pricing/foundation-model-serving"
},
"databricks/databricks-meta-llama-3-3-70b-instruct": {
"input_cost_per_token": 1.00002e-06,
"input_dbu_cost_per_token": 1.4286e-05,
"input_cost_per_token": 5.0001e-07,
"input_dbu_cost_per_token": 7.143e-06,
"litellm_provider": "databricks",
"max_input_tokens": 128000,
"max_output_tokens": 128000,
@ -7916,8 +8230,8 @@
"notes": "Input/output cost per token is dbu cost * $0.070, based on databricks Llama 3.1 70B conversion. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 2.99999e-06,
"output_dbu_cost_per_token": 4.2857e-05,
"output_cost_per_token": 1.5000300000000002e-06,
"output_dbu_cost_per_token": 2.1429e-05,
"source": "https://www.databricks.com/product/pricing/foundation-model-serving",
"supports_tool_choice": true
},
@ -7932,7 +8246,7 @@
"notes": "Input/output cost per token is dbu cost * $0.070, based on databricks Llama 3.1 70B conversion. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 2.99999e-06,
"output_cost_per_token": 2.9999900000000002e-06,
"output_dbu_cost_per_token": 4.2857e-05,
"source": "https://www.databricks.com/product/pricing/foundation-model-serving",
"supports_tool_choice": true
@ -7948,13 +8262,13 @@
"notes": "Input/output cost per token is dbu cost * $0.070, based on databricks Llama 3.1 70B conversion. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 9.9902e-07,
"output_cost_per_token": 1.00002e-06,
"output_dbu_cost_per_token": 1.4286e-05,
"source": "https://www.databricks.com/product/pricing/foundation-model-serving",
"supports_tool_choice": true
},
"databricks/databricks-mpt-30b-instruct": {
"input_cost_per_token": 9.9902e-07,
"input_cost_per_token": 1.00002e-06,
"input_dbu_cost_per_token": 1.4286e-05,
"litellm_provider": "databricks",
"max_input_tokens": 8192,
@ -7964,7 +8278,7 @@
"notes": "Input/output cost per token is dbu cost * $0.070, based on databricks Llama 3.1 70B conversion. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 9.9902e-07,
"output_cost_per_token": 1.00002e-06,
"output_dbu_cost_per_token": 1.4286e-05,
"source": "https://www.databricks.com/product/pricing/foundation-model-serving",
"supports_tool_choice": true

View file

@ -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

View file

@ -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

View file

@ -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
(

View file

@ -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__])