diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 3ba7880f4df..aab7f93164f 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -127,7 +127,6 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( from litellm.proxy._experimental.mcp_server.sampling_handler import ( MCP_SAMPLING_AVAILABLE, ) -from litellm.proxy._experimental.mcp_server.tool_registry import global_mcp_tool_registry from litellm.proxy._experimental.mcp_server.utils import ( MCP_TOOL_PREFIX_SEPARATOR, MCPMissingUserEnvVarsError, @@ -5140,14 +5139,6 @@ class MCPServerManager: return prefixed_tools def get_listed_tool(self, server: MCPServer, name: str) -> MCPTool | None: - if server.spec_path: - bare_name: Final = strip_known_server_prefix(name, server) - registered: Final = global_mcp_tool_registry.get_tool( - f"{get_server_prefix(server)}{MCP_TOOL_PREFIX_SEPARATOR}{bare_name}" - ) or global_mcp_tool_registry.get_tool(bare_name) - if registered is None: - return None - return MCPTool(name=bare_name, description=registered.description, inputSchema=registered.input_schema) listed: Final = self._listed_tools_by_server_id.get(server.server_id) if not listed: return None diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 88a61f955da..e4bc9bdd952 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -86,6 +86,7 @@ from litellm.proxy.litellm_pre_call_utils import ( ) from litellm.types.mcp import MCPAuth, MCPSpecVersion from litellm.types.mcp_server.mcp_server_manager import MCPInfo, MCPServer +from litellm.types.mcp_server.tool_registry import MCPTool as RegisteredTool from litellm.types.utils import CallTypes, StandardLoggingMCPToolCall from litellm.utils import Rules, client, function_setup @@ -2777,6 +2778,9 @@ if MCP_AVAILABLE: return managed_resource_templates + def _registered_tool_metadata(name: str, registered: RegisteredTool) -> MCPTool: + return MCPTool(name=name, description=registered.description, inputSchema=registered.input_schema) + def _resolve_display_name_to_original( name: str, allowed_mcp_servers: list[MCPServer], @@ -3119,7 +3123,7 @@ if MCP_AVAILABLE: server=mcp_server, raw_headers=raw_headers, litellm_logging_obj=litellm_logging_obj, - tool=global_mcp_server_manager.get_listed_tool(mcp_server, original_tool_name), + tool=_registered_tool_metadata(original_tool_name, local_tool), ) # `pre_call_tool_check` may return guardrail-modified # arguments; honor them on the local path too. @@ -3185,7 +3189,8 @@ if MCP_AVAILABLE: # not in the registry either, `_handle_local_mcp_tool` below reports # 404 and nothing runs, so demanding a server here would turn every # unknown tool name into a misleading 503. - if global_mcp_tool_registry.get_tool(original_tool_name) is not None: + registered_local_tool: Final = global_mcp_tool_registry.get_tool(original_tool_name) + if registered_local_tool is not None: # `mcp_server` is None here because the tool name is not in the # tool -> server mapping, but the name still carries a prefix # that the server-level check above compared against the @@ -3226,7 +3231,7 @@ if MCP_AVAILABLE: server=prefix_server, raw_headers=raw_headers, litellm_logging_obj=litellm_logging_obj, - tool=global_mcp_server_manager.get_listed_tool(prefix_server, original_tool_name), + tool=_registered_tool_metadata(original_tool_name, registered_local_tool), ) if "arguments" in hook_result: arguments = hook_result["arguments"] # pyright: ignore[reportAny] # hook returns untyped args diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index d2cdab12fb0..7d4159aade8 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -7109,6 +7109,8 @@ async def test_execute_mcp_tool_sets_model_in_model_call_details(): fake_tool = MagicMock() fake_tool.name = "list_pets" + fake_tool.description = "test tool" + fake_tool.input_schema = {"type": "object"} start_time = datetime.now(timezone.utc) litellm_logging_obj, _ = function_setup( @@ -7175,7 +7177,7 @@ async def test_execute_mcp_tool_hands_openapi_registered_tool_metadata_to_pre_ca ) schema = {"type": "object", "properties": {"limit": {"type": "integer"}}} mcp_module.global_mcp_tool_registry.register_tool( - name="petstore-list_pets", description="List the pets", input_schema=schema, handler=lambda: None + name="petstore-list_pets", description="List the pets", input_schema=schema, handler=lambda limit: "ok" ) manager = mcp_module.global_mcp_server_manager manager._listed_tools_by_server_id.pop(petstore.server_id, None) @@ -7185,14 +7187,6 @@ async def test_execute_mcp_tool_hands_openapi_registered_tool_metadata_to_pre_ca with ( patch.object(manager, "_get_mcp_server_from_tool_name", return_value=petstore), patch.object(manager, "pre_call_tool_check", new=pre_call_tool_check), - patch( - "litellm.proxy._experimental.mcp_server.server._handle_local_mcp_tool", - new=AsyncMock(return_value=[]), - ), - patch( - "litellm.proxy._experimental.mcp_server.server.MCPRequestHandler.is_tool_allowed", - return_value=True, - ), ): await mcp_module.execute_mcp_tool( name="petstore-list_pets", @@ -7212,6 +7206,54 @@ async def test_execute_mcp_tool_hands_openapi_registered_tool_metadata_to_pre_ca ) +@pytest.mark.asyncio +async def test_execute_mcp_tool_hands_hooks_the_metadata_of_the_operation_it_runs_when_names_collide(): + """An OpenAPI operation whose name starts with its own server prefix must not be reported to the + pre-call hooks with the metadata of the shorter operation, since that is not the one that runs.""" + from litellm.proxy._experimental.mcp_server import server as mcp_module + + petstore = MCPServer( + server_id="petstore-id", + name="petstore", + server_name="petstore", + transport=MCPTransport.http, + url=None, + spec_path="https://example.com/petstore.yaml", + ) + registry = mcp_module.global_mcp_tool_registry + registry.register_tool(name="petstore-get_pet", description="short", input_schema={}, handler=lambda: "short") + registry.register_tool( + name="petstore-petstore-get_pet", + description="long", + input_schema={"type": "object", "properties": {"petId": {"type": "integer"}}}, + handler=lambda: "long", + ) + manager = mcp_module.global_mcp_server_manager + pre_call_tool_check = AsyncMock(return_value={}) + + try: + with ( + patch.object(manager, "_get_mcp_server_from_tool_name", return_value=petstore), + patch.object(manager, "pre_call_tool_check", new=pre_call_tool_check), + ): + result = await mcp_module.execute_mcp_tool( + name="petstore-petstore-get_pet", + arguments={}, + allowed_mcp_servers=[petstore], + start_time=datetime.now(), + user_api_key_auth=UserAPIKeyAuth(api_key="sk-user", user_id="alice"), + ) + finally: + registry.unregister_tools_with_prefix("petstore-") + + handed_tool = pre_call_tool_check.call_args.kwargs["tool"] + assert (handed_tool.description, handed_tool.inputSchema) == ( + "long", + {"type": "object", "properties": {"petId": {"type": "integer"}}}, + ) + assert result.content[0].text == "long" + + @pytest.mark.asyncio async def test_execute_mcp_tool_rest_unresolved_prefixed_name_routes_to_requested_server(): """A prefixed REST name that resolves to no tool must still dispatch to the server_id. 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 cc4d730a096..f1675b72f77 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 @@ -6504,37 +6504,6 @@ 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_reads_openapi_registry_without_a_prior_listing(self): - """OpenAPI tools live in the local registry from registration on, so their metadata must resolve - before any tools/list has run and must disappear with the registration.""" - from litellm.proxy._experimental.mcp_server.tool_registry import global_mcp_tool_registry - - server = MCPServer( - server_id="petstore-id", - name="petstore", - server_name="petstore", - transport=MCPTransport.http, - url=None, - spec_path="https://example.com/petstore.yaml", - ) - schema = {"type": "object", "properties": {"petId": {"type": "integer"}}} - manager = MCPServerManager() - global_mcp_tool_registry.register_tool( - name="petstore-get_pet", description="Fetch a pet", input_schema=schema, handler=lambda: None - ) - try: - for spelling in ("get_pet", "petstore-get_pet"): - tool = manager.get_listed_tool(server, spelling) - assert tool is not None and (tool.name, tool.description, tool.inputSchema) == ( - "get_pet", - "Fetch a pet", - schema, - ) - finally: - global_mcp_tool_registry.unregister_tools_with_prefix("petstore-") - - assert manager.get_listed_tool(server, "petstore-get_pet") is None - @pytest.mark.asyncio async def test_get_allowed_mcp_servers_with_user_api_key_auth(self): """ diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py index 64614c094ba..c534cf73219 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py @@ -43,6 +43,8 @@ async def test_openapi_local_tool_runs_pre_call_tool_check(): fake_tool = MagicMock() fake_tool.name = "list_pets" + fake_tool.description = "test tool" + fake_tool.input_schema = {"type": "object"} pre_call = AsyncMock(return_value={}) handle_local = AsyncMock(return_value=[]) @@ -124,6 +126,8 @@ async def test_openapi_local_tool_blocked_when_pre_call_check_raises(): fake_tool = MagicMock() fake_tool.name = "delete_pet" + fake_tool.description = "test tool" + fake_tool.input_schema = {"type": "object"} pre_call = AsyncMock( side_effect=HTTPException(status_code=403, detail="not allowed") @@ -186,6 +190,8 @@ async def test_openapi_local_tool_denied_when_server_not_resolvable(): fake_tool = MagicMock() fake_tool.name = "list_pets" + fake_tool.description = "test tool" + fake_tool.input_schema = {"type": "object"} pre_call = AsyncMock(return_value={}) handle_local = AsyncMock(return_value=[]) @@ -270,6 +276,8 @@ async def test_openapi_local_tool_injects_resolved_oauth_token(): fake_tool = MagicMock() fake_tool.name = "get_values" + fake_tool.description = "test tool" + fake_tool.input_schema = {"type": "object"} captured: dict = {} async def handle_local(_name, _arguments): @@ -616,6 +624,8 @@ async def test_per_server_auth_header_reaches_both_openapi_dispatch_arms(dispatc if dispatch_arm == "local_registry": fake_tool = MagicMock() fake_tool.name = "list_reports" + fake_tool.description = "test tool" + fake_tool.input_schema = {"type": "object"} with ( patch.object(manager, "_get_mcp_server_from_tool_name", return_value=server), patch.object(mcp_module.global_mcp_tool_registry, "get_tool", return_value=fake_tool), @@ -687,6 +697,8 @@ async def test_local_dispatch_reports_the_outcome_instead_of_success(failure: st fake_tool = MagicMock() fake_tool.name = "list_reports" + fake_tool.description = "test tool" + fake_tool.input_schema = {"type": "object"} fake_tool.handler = raising_handler server = MCPServer( server_id="srv-openapi",