diff --git a/litellm/proxy/_experimental/mcp_server/operations.py b/litellm/proxy/_experimental/mcp_server/operations.py index 891c37f4fe7..33b61ad310e 100644 --- a/litellm/proxy/_experimental/mcp_server/operations.py +++ b/litellm/proxy/_experimental/mcp_server/operations.py @@ -141,7 +141,6 @@ from litellm.types.mcp import ( without_header, ) 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 @@ -1609,17 +1608,10 @@ async def _list_mcp_resource_templates( return managed_resource_templates -def _registered_tool_metadata( - name: str, registered: RegisteredTool, server: MCPServer, caller: ListedToolsCaller -) -> MCPTool: - """The tool as ``tools/list`` served it to this caller (pinned, overridden, guardrail-masked) when a - listing was recorded, else the registry entry with the admin description override applied.""" - listed: Final = global_mcp_server_manager.get_listed_tool(server, name, caller) - if listed is not None: - return listed - overrides: Final = server.tool_name_to_description - description: Final = overrides.get(name, registered.description) if overrides else registered.description - return MCPTool(name=name, description=description, input_schema=registered.input_schema) +def _registered_tool_metadata(name: str, server: MCPServer, caller: ListedToolsCaller) -> MCPTool | None: + """The tool as ``tools/list`` served it to this caller (pinned, overridden, guardrail-masked), or None when + no listing was recorded so the call hands the hooks name and arguments only.""" + return global_mcp_server_manager.get_listed_tool(server, name, caller) def _resolve_display_name_to_original( @@ -2093,7 +2085,6 @@ async def _execute_mcp_tool( guardrail_context=guardrail_context, tool=_registered_tool_metadata( original_tool_name, - local_tool, mcp_server, listed_tools_caller_for( mcp_server, user_api_key_auth, mcp_auth_header, mcp_server_auth_headers, raw_headers, oauth2_headers @@ -2212,7 +2203,6 @@ async def _execute_mcp_tool( guardrail_context=guardrail_context, tool=_registered_tool_metadata( original_tool_name, - registered_local_tool, prefix_server, listed_tools_caller_for( prefix_server, diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py index afe0ef99419..fcf600d7b98 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py @@ -8070,10 +8070,13 @@ async def test_execute_mcp_tool_sets_model_in_model_call_details(): @pytest.mark.asyncio -async def test_execute_mcp_tool_hands_openapi_registered_tool_metadata_to_pre_call_hooks(): - """OpenAPI-generated tools dispatch through the local registry, so the pre-call hooks must get the - registered description and input schema on that path too, even when no tools/list ran first.""" +async def test_execute_mcp_tool_hands_openapi_hooks_the_listed_entry_and_nothing_before_a_listing(): + """A local-registry tools/call with no prior tools/list hands the pre-call hooks name and arguments + only, as before this metadata existed, so a pre_mcp_call policy never scans a description the caller was + not served. Once the caller has listed, the same call hands the entry that listing served.""" + from litellm.caching.caching import DualCache from litellm.proxy._experimental.mcp_server import operations as mcp_module + from litellm.proxy.utils import ProxyLogging petstore = MCPServer( server_id="petstore-id", @@ -8082,6 +8085,7 @@ async def test_execute_mcp_tool_hands_openapi_registered_tool_metadata_to_pre_ca transport=MCPTransport.http, url=None, spec_path="https://example.com/petstore.yaml", + tool_name_to_description={"list_pets": "ADMIN DESC"}, ) schema = {"type": "object", "properties": {"limit": {"type": "integer"}}} mcp_module.global_mcp_tool_registry.register_tool( @@ -8089,71 +8093,42 @@ async def test_execute_mcp_tool_hands_openapi_registered_tool_metadata_to_pre_ca ) manager = mcp_module.global_mcp_server_manager manager._listed_tools_by_server_id.pop(petstore.server_id, None) - pre_call_tool_check = AsyncMock(return_value={}) + alice = UserAPIKeyAuth(api_key="sk-user", user_id="alice") + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging.pre_call_hook = AsyncMock(return_value={}) + pre_call_tool_check = AsyncMock(wraps=manager.pre_call_tool_check) + + async def call() -> tuple[MCPTool | None, dict]: + await mcp_module.execute_mcp_tool( + name="petstore-list_pets", + arguments={"limit": 10}, + allowed_mcp_servers=[petstore], + start_time=datetime.now(), + user_api_key_auth=alice, + ) + return pre_call_tool_check.call_args.kwargs["tool"], proxy_logging.pre_call_hook.call_args.kwargs["data"] 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), + patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging), ): - await mcp_module.execute_mcp_tool( - name="petstore-list_pets", - arguments={"limit": 10}, - allowed_mcp_servers=[petstore], - start_time=datetime.now(), - user_api_key_auth=UserAPIKeyAuth(api_key="sk-user", user_id="alice"), + never_listed_tool, never_listed_data = await call() + manager._record_listed_tools( + petstore, + [MCPTool(name="list_pets", description="ADMIN DESC", inputSchema=schema)], + ListedToolsCaller(user_api_key_auth=alice), ) + listed_tool, listed_data = await call() finally: mcp_module.global_mcp_tool_registry.unregister_tools_with_prefix("petstore-") + manager._listed_tools_by_server_id.pop(petstore.server_id, None) - handed_tool = pre_call_tool_check.call_args.kwargs["tool"] - assert (handed_tool.name, handed_tool.description, handed_tool.input_schema) == ( - "list_pets", - "List the pets", - schema, - ) - - -@pytest.mark.asyncio -async def test_execute_mcp_tool_hands_openapi_hooks_the_admin_description_clients_saw(): - """tools/list shows the admin's tool_name_to_description wording, so the local-registry call path - must hand the pre-call hooks that same wording rather than the generated one.""" - from litellm.proxy._experimental.mcp_server import operations 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", - tool_name_to_description={"getpetbyid": "ADMIN DESC"}, - ) - schema = {"type": "object", "properties": {"petId": {"type": "integer"}}} - mcp_module.global_mcp_tool_registry.register_tool( - name="petstore-getpetbyid", description="Find pet by ID", input_schema=schema, handler=lambda petId: "ok" - ) - manager = mcp_module.global_mcp_server_manager - manager._listed_tools_by_server_id.pop(petstore.server_id, None) - 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), - ): - await mcp_module.execute_mcp_tool( - name="petstore-getpetbyid", - arguments={"petId": 1}, - allowed_mcp_servers=[petstore], - start_time=datetime.now(), - user_api_key_auth=UserAPIKeyAuth(api_key="sk-user", user_id="alice"), - ) - finally: - mcp_module.global_mcp_tool_registry.unregister_tools_with_prefix("petstore-") - - handed_tool = pre_call_tool_check.call_args.kwargs["tool"] - assert (handed_tool.description, handed_tool.input_schema) == ("ADMIN DESC", schema) + assert never_listed_tool is None + assert (never_listed_data.get("mcp_tool_description"), never_listed_data.get("mcp_input_schema")) == (None, None) + assert listed_tool is not None and (listed_tool.description, listed_tool.input_schema) == ("ADMIN DESC", schema) + assert (listed_data["mcp_tool_description"], listed_data["mcp_input_schema"]) == ("ADMIN DESC", schema) @pytest.mark.asyncio @@ -8267,9 +8242,9 @@ async def test_execute_mcp_tool_hands_openapi_hooks_each_callers_own_listed_entr @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.""" +async def test_execute_mcp_tool_runs_the_longer_colliding_operation_and_hands_hooks_no_registry_metadata(): + """An OpenAPI operation whose name starts with its own server prefix runs instead of the shorter one, and + with no prior listing the pre-call hooks get name and arguments only, never either registry entry.""" from litellm.proxy._experimental.mcp_server import operations as mcp_module petstore = MCPServer( @@ -8306,11 +8281,7 @@ async def test_execute_mcp_tool_hands_hooks_the_metadata_of_the_operation_it_run finally: registry.unregister_tools_with_prefix("petstore-") - handed_tool = pre_call_tool_check.call_args.kwargs["tool"] - assert (handed_tool.description, handed_tool.input_schema) == ( - "long", - {"type": "object", "properties": {"petId": {"type": "integer"}}}, - ) + assert pre_call_tool_check.call_args.kwargs["tool"] is None assert result.content[0].text == "long" diff --git a/tests/unit/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/unit/proxy/_experimental/mcp_server/test_rest_endpoints.py index 1cf23e0b49b..759014b54c5 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -3326,8 +3326,7 @@ async def test_request_selected_tool_specific_guardrail_applies_to_virtual_execu default_on=False, custom_code="def apply_guardrail(inputs, request_data, input_type):\n" ' if inputs.get("tools", [{}])[0].get("function", {}).get("name") == "execute":\n' - ' texts = [t.replace("confidential", "redacted") for t in inputs.get("texts", [])]\n' - f' return {{"action": "{action}", "reason": "resolved tool blocked", "texts": texts}}\n' + f' return {{"action": "{action}", "reason": "resolved tool blocked", "texts": ["redacted"]}}\n' " return allow()\n", ) manager: Final = mcp_server_manager.MCPServerManager()