From d7623112fa817c4c7e949d0ef49c10e8851352e2 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sun, 13 Sep 2026 05:15:33 +0000 Subject: [PATCH] fix(mcp): hand Agent 365 the admin description override and clear the LIT001 and schema gates get_listed_tool now applies tool_name_to_description so the guardrail evaluates the description clients saw, with a regression test. The listed-tools cache and tool_input_schema fields use Mapping annotations so the change adds no LIT001 findings, and schema.d.ts is regenerated with npm run gen:api Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/mcp_server_manager.py | 13 ++++++---- .../guardrail_hooks/agent_365/agent_365.py | 4 ++-- litellm/types/mcp.py | 4 ++-- .../mcp_server/test_mcp_server_manager.py | 24 +++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 -- 5 files changed, 36 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 696d496da54..560d3d51177 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -242,7 +242,8 @@ _user_env_vars_cache: Final[dict[tuple[str, str], tuple[dict[str, str], float]]] _USER_ENV_VARS_CACHE_TTL: Final = 60 # seconds _USER_ENV_VARS_CACHE_MAX_SIZE: Final = 4096 # cap to prevent unbounded growth -_NO_LISTED_TOOLS: Final[Mapping[str | None, Mapping[str, MCPTool]]] = MappingProxyType({}) +_ListedToolsByCaller: TypeAlias = Mapping[str | None, Mapping[str, MCPTool]] +_NO_LISTED_TOOLS: Final[_ListedToolsByCaller] = MappingProxyType({}) _LISTED_TOOLS_CALLERS_PER_SERVER: Final = 256 # Auth types whose upstream OAuth endpoints (protected-resource + authorization-server metadata) the @@ -1948,9 +1949,7 @@ class MCPServerManager: "gmail_send_email": "zapier_mcp_server", } """ - self._listed_tools_by_server_id: dict[ - str, Mapping[str | None, Mapping[str, MCPTool]] - ] = {} # mutable-ok: refreshed per tools/list + self._listed_tools_by_server_id: dict[str, _ListedToolsByCaller] = {} # mutable-ok: refreshed per tools/list self._upstream_initialize_instructions_by_server_id: dict[str, str] = {} # Per-server monotonic timestamp of last upstream prefetch attempt (success, # empty result, or failure). Used to throttle re-probes for servers that do @@ -5392,7 +5391,11 @@ class MCPServerManager: listed: Final = self._listed_tools_by_server_id.get(server.server_id, _NO_LISTED_TOOLS).get(identity) if not listed: return None - return listed.get(name) or listed.get(strip_known_server_prefix(name, server)) + tool: Final = listed.get(name) or listed.get(strip_known_server_prefix(name, server)) + if tool is None: + return None + description: Final = (server.tool_name_to_description or {}).get(tool.name) + return tool if description is None else tool.model_copy(update={"description": description}) def _create_prefixed_prompts( self, prompts: Sequence[Prompt], server: MCPServer, add_prefix: bool = True diff --git a/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py b/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py index 58d8ba5f865..d81fa2f095b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py +++ b/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py @@ -77,7 +77,7 @@ def _parse_expires_in(raw: object) -> float: return _DEFAULT_TOKEN_TTL_SECONDS -def _parse_tool_input_schema(raw: object) -> dict[str, object] | None: +def _parse_tool_input_schema(raw: object) -> Mapping[str, object] | None: try: return _TOOL_INPUT_SCHEMA_ADAPTER.validate_python(raw) except ValidationError: @@ -101,7 +101,7 @@ class _ToolReference(BaseModel): name: str description: str | None = None - input_schema: dict[str, object] | None = Field(default=None, serialization_alias="inputSchema") + input_schema: Mapping[str, object] | None = Field(default=None, serialization_alias="inputSchema") class _UnavailableDetail(TypedDict): diff --git a/litellm/types/mcp.py b/litellm/types/mcp.py index 13a46d8b84f..87ffaa7cd9e 100644 --- a/litellm/types/mcp.py +++ b/litellm/types/mcp.py @@ -392,7 +392,7 @@ class MCPPreCallRequestObject(BaseModel): arguments: dict[str, Any] server_name: str | None = None tool_description: str | None = None - tool_input_schema: dict[str, object] | None = None + tool_input_schema: Mapping[str, object] | None = None user_api_key_auth: dict[str, Any] | None = None hidden_params: HiddenParams = HiddenParams() @@ -417,7 +417,7 @@ class MCPDuringCallRequestObject(BaseModel): arguments: dict[str, Any] server_name: str | None = None tool_description: str | None = None - tool_input_schema: dict[str, object] | None = None + tool_input_schema: Mapping[str, object] | None = None start_time: float | None = None hidden_params: HiddenParams = HiddenParams() 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 1bdb1d0a427..56e693e3523 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 @@ -6672,6 +6672,30 @@ class TestMCPServerManager: assert by_prefixed_name is not None and by_prefixed_name.description == "v2" assert manager.get_listed_tool(server, "missing") is None + def test_get_listed_tool_uses_admin_description_override_clients_saw(self): + manager = MCPServerManager() + server = MCPServer( + server_id="srv", + name="srv", + transport=MCPTransport.http, + url="http://srv", + tool_name_to_description={"echo": "Admin wording"}, + ) + schema = {"type": "object", "properties": {"text": {"type": "string"}}} + manager._create_prefixed_tools( + [ + MCPTool(name="echo", description="Upstream wording", inputSchema=schema), + MCPTool(name="ping", description="Untouched", inputSchema={}), + ], + server, + ) + + overridden = manager.get_listed_tool(server, "srv-echo") + assert overridden is not None + assert (overridden.name, overridden.description, overridden.inputSchema) == ("echo", "Admin wording", schema) + untouched = manager.get_listed_tool(server, "ping") + assert untouched is not None and untouched.description == "Untouched" + def test_server_definition_change_drops_listed_tools(self): manager = MCPServerManager() server = MCPServer(server_id="srv", name="srv", transport=MCPTransport.http, url="http://srv") diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 3b172b8d35f..7df6fbc5863 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -16781,7 +16781,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) @@ -16887,7 +16886,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys)