diff --git a/litellm/proxy/hooks/mcp_semantic_filter/hook.py b/litellm/proxy/hooks/mcp_semantic_filter/hook.py index bad6ef44ccd..5f1d061c7cb 100644 --- a/litellm/proxy/hooks/mcp_semantic_filter/hook.py +++ b/litellm/proxy/hooks/mcp_semantic_filter/hook.py @@ -155,6 +155,40 @@ class SemanticToolFilterHook(CustomLogger): return await self.filter.filter_tools(query=user_query, available_tools=expanded_tools) + def _selected_tool_names(self, filtered_tools: list[dict[str, Any]]) -> list[str]: + """Names of the semantically selected tools, as produced by the MCP expansion.""" + names = (self.filter._extract_tool_info(tool)[0] for tool in filtered_tools) + return [name for name in names if name] + + @staticmethod + def _narrow_mcp_references(tools: list[Any], selected_tool_names: list[str]) -> list[Any]: + """ + Restrict each litellm_proxy MCP reference to the semantically selected tools. + + The reference block is preserved rather than replaced with expanded tools, so the + MCP gateway still performs the expansion. That keeps the per-endpoint tool shape + and tool auto-execution intact. Expansion already applied any caller-supplied + allowed_tools, so this selection can only narrow a block further. + + Whether an undecidable selection exposes every tool or none is owned by + SemanticMCPToolFilter.filter_tools, which returns the full set when nothing + matches; the same policy therefore governs references and plain tools. Passing an + empty selection through is safe rather than a hidden allow-all: the gateway reads + the union of every reference's allowed_tools and treats an empty union as unset. + """ + from litellm.responses.mcp.litellm_proxy_mcp_handler import ( + LiteLLM_Proxy_MCP_Handler, + ) + + return [ + ( + {**tool, "allowed_tools": selected_tool_names} + if isinstance(tool, dict) and LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway([tool]) + else tool + ) + for tool in tools + ] + def _is_mcp_tool(self, tool: object) -> bool: """ Check whether *tool* is registered in the MCP semantic router. @@ -261,36 +295,30 @@ class SemanticToolFilterHook(CustomLogger): if self._should_expand_mcp_tools(tools): verbose_proxy_logger.debug("Detected litellm_proxy MCP references, expanding before semantic filtering") + if not self.filter.enabled: + verbose_proxy_logger.debug("Semantic filter disabled, leaving MCP references untouched") + return None + try: native_tools_before_expand = [t for t in tools if not (isinstance(t, dict) and t.get("type") == "mcp")] expanded_tools = await self._expand_mcp_tools(tools, user_api_key_dict) if not expanded_tools: - if native_tools_before_expand: - data["tools"] = native_tools_before_expand - verbose_proxy_logger.warning( - f"No MCP tools expanded, preserving {len(native_tools_before_expand)} native tools" - ) - return data verbose_proxy_logger.warning("No tools expanded from MCP references") return None - if not self.filter.enabled: - data["tools"] = native_tools_before_expand + expanded_tools - verbose_proxy_logger.debug("Semantic filter disabled, forwarding expanded MCP tools unfiltered") - return data - filtered_expanded_tools = await self._filter_expanded_tools(data=data, expanded_tools=expanded_tools) - combined_tools = native_tools_before_expand + filtered_expanded_tools - data["tools"] = combined_tools + selected_tool_names = self._selected_tool_names(filtered_expanded_tools) + narrowed_tools = self._narrow_mcp_references(tools, selected_tool_names) + data["tools"] = narrowed_tools self._emit_filter_metadata_safe( data=data, mcp_tools=expanded_tools, filtered_mcp_tools=filtered_expanded_tools, native_tools=native_tools_before_expand, - filtered_tools=combined_tools, + filtered_tools=narrowed_tools, ) verbose_proxy_logger.info( f"Expanded MCP references to {len(expanded_tools)} tools " diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py index 0f392a54b4c..d864b442bd3 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py @@ -865,14 +865,17 @@ async def test_semantic_filter_hook_filters_expanded_litellm_proxy_tools(): ) assert result is not None, "Hook should return modified data" - filtered = result["tools"] + mcp_references = [tool for tool in result["tools"] if tool.get("type") == "mcp"] + assert len(mcp_references) == 1, "The litellm_proxy MCP reference must be preserved for the MCP gateway to expand" - assert len(filtered) <= 2, f"Expanded tools should be filtered to top_k=2, got {len(filtered)}" - assert len(filtered) < len(expanded_tools), ( - f"Hook must not forward all {len(expanded_tools)} expanded tools unfiltered, got {len(filtered)}" + allowed_tools = mcp_references[0]["allowed_tools"] + assert len(allowed_tools) <= 2, f"Expanded tools should be filtered to top_k=2, got {len(allowed_tools)}" + assert len(allowed_tools) < len(expanded_tools), ( + f"Hook must not forward all {len(expanded_tools)} expanded tools unfiltered, got {len(allowed_tools)}" ) - for tool in filtered: - assert tool in expanded_tools, "Filtered tools must be the original expanded tool dicts" + expanded_names = {tool["name"] for tool in expanded_tools} + for name in allowed_tools: + assert name in expanded_names, "Selected tool names must come from the expanded tools" assert ( "litellm_semantic_filter_stats" in result["metadata"] @@ -880,9 +883,246 @@ async def test_semantic_filter_hook_filters_expanded_litellm_proxy_tools(): stats = result["metadata"]["litellm_semantic_filter_stats"] total, selected = stats.split("->") assert int(total) == 5, f"Stats 'from' should be pre-filter expanded count (5), got {total}" - assert int(selected) == len(filtered), f"Stats 'to' should match post-filter count, got {selected}" + assert int(selected) == len(allowed_tools), f"Stats 'to' should match post-filter count, got {selected}" - print(f"✅ Expanded litellm_proxy tools filtered: {len(expanded_tools)} -> {len(filtered)}, stats={stats}") + print(f"✅ Expanded litellm_proxy tools filtered: {len(expanded_tools)} -> {len(allowed_tools)}, stats={stats}") + + +@pytest.mark.asyncio +async def test_semantic_filter_hook_narrows_mcp_reference_for_chat_completions(): + """ + Regression test (LIT-4451): the hook must narrow the litellm_proxy MCP + reference instead of replacing it with expanded tool definitions. + + Given: A /chat/completions request whose tools are a single + {"type": "mcp", "server_url": "litellm_proxy"} reference that + expands to 5 tools, with the semantic filter selecting top_k=2 + When: The hook processes the request with call_type="acompletion" + Then: The MCP reference survives in data["tools"], carrying the selected + tools in allowed_tools, and no expanded function definitions are + written into the request. + + Replacing the reference made the hook write Responses-API-shaped tools + ({"type": "function", "name": ...}) into /chat/completions, which expects + {"type": "function", "function": {...}}. The provider transformation then + rejected every MCP tool (Anthropic raised KeyError: 'function') or dropped + it silently (Bedrock), so the model saw no MCP tools at all. Replacing the + reference also removed the marker the MCP gateway matches on, which + disabled tool auto-execution for require_approval="never". + """ + from litellm.proxy._experimental.mcp_server.semantic_tool_filter import ( + SemanticMCPToolFilter, + ) + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + from litellm.types.utils import Embedding, EmbeddingResponse + + mock_router = Mock() + + def mock_embedding_sync(*args, **kwargs): + return EmbeddingResponse( + data=[Embedding(embedding=[0.1] * 1536, index=0, object="embedding")], + model="text-embedding-3-small", + object="list", + usage={"prompt_tokens": 10, "total_tokens": 10}, + ) + + async def mock_embedding_async(*args, **kwargs): + return mock_embedding_sync() + + mock_router.embedding = mock_embedding_sync + mock_router.aembedding = mock_embedding_async + + filter_instance = SemanticMCPToolFilter( + embedding_model="text-embedding-3-small", + litellm_router_instance=mock_router, + top_k=2, + similarity_threshold=0.3, + enabled=True, + ) + + registry_tools = [ + MCPTool( + name=f"srv-tool_{i}", + description=f"Registry tool {i}", + inputSchema={"type": "object"}, + ) + for i in range(5) + ] + filter_instance._build_router(registry_tools) + + expanded_tools = [ + { + "type": "function", + "name": f"srv-tool_{i}", + "description": f"Registry tool {i}", + "parameters": {"type": "object", "properties": {}}, + } + for i in range(5) + ] + + hook = SemanticToolFilterHook(filter_instance) + hook._expand_mcp_tools = AsyncMock( # type: ignore[method-assign] + return_value=expanded_tools + ) + + mcp_reference = { + "type": "mcp", + "server_url": "litellm_proxy", + "require_approval": "never", + } + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Send an email"}], + "tools": [mcp_reference], + "metadata": {}, + } + + result = await hook.async_pre_call_hook( + user_api_key_dict=Mock(), + cache=Mock(), + data=data, + call_type="acompletion", + ) + + assert result is not None, "Hook should return modified data" + forwarded = result["tools"] + + assert [tool.get("type") for tool in forwarded] == ["mcp"], ( + "The MCP reference must be the only forwarded tool; writing expanded function " + f"definitions into a chat completion loses every MCP tool. Got: {forwarded}" + ) + assert forwarded[0]["server_url"] == "litellm_proxy", "The MCP reference must keep routing to the gateway" + assert forwarded[0]["require_approval"] == "never", "The MCP reference must keep its auto-execute marker" + + allowed_tools = forwarded[0]["allowed_tools"] + assert allowed_tools, "The narrowed reference must still carry the selected tools" + assert len(allowed_tools) <= 2, f"Selection must narrow the reference to top_k=2, got {allowed_tools}" + assert len(allowed_tools) < len(expanded_tools), ( + f"Hook must not forward all {len(expanded_tools)} expanded tools unfiltered, got {allowed_tools}" + ) + assert set(allowed_tools) <= {tool["name"] for tool in expanded_tools}, ( + f"Selected names must come from the expanded tools, got {allowed_tools}" + ) + + print(f"✅ chat completions: MCP reference preserved, narrowed to {allowed_tools}") + + +@pytest.mark.asyncio +async def test_semantic_filter_hook_zero_matches_exposes_all_tools_on_both_paths(): + """ + A query that matches nothing must expose every MCP tool, whether the request + carries a litellm_proxy MCP reference or plain MCP tool objects. + + Given: A router that returns no matches for the query + When: The hook processes an MCP reference request and a plain MCP tool request + Then: Both expose all 3 tools, because filter_tools owns the undecidable-selection + policy and returns the full set rather than an empty one + + The two paths narrow through different mechanisms (allowed_tools on the reference + versus dropping unmatched entries), so they could drift into opposite fail + behaviours. Pinning both here keeps that single policy honest: flipping + filter_tools to fail closed must fail this test on both paths at once, instead of + silently hard-limiting one surface and not the other. + """ + from litellm.proxy._experimental.mcp_server.semantic_tool_filter import ( + SemanticMCPToolFilter, + ) + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + from litellm.types.utils import Embedding, EmbeddingResponse + + mock_router = Mock() + + def mock_embedding_sync(*args, **kwargs): + return EmbeddingResponse( + data=[Embedding(embedding=[0.1] * 1536, index=0, object="embedding")], + model="text-embedding-3-small", + object="list", + usage={"prompt_tokens": 10, "total_tokens": 10}, + ) + + async def mock_embedding_async(*args, **kwargs): + return mock_embedding_sync() + + mock_router.embedding = mock_embedding_sync + mock_router.aembedding = mock_embedding_async + + registry_tools = [ + MCPTool( + name=f"srv-tool_{i}", + description=f"Registry tool {i}", + inputSchema={"type": "object"}, + ) + for i in range(3) + ] + + def build_hook(): + filter_instance = SemanticMCPToolFilter( + embedding_model="text-embedding-3-small", + litellm_router_instance=mock_router, + top_k=2, + similarity_threshold=0.3, + enabled=True, + ) + filter_instance._build_router(registry_tools) + zero_match_router = Mock(return_value=[]) + zero_match_router.top_k = 2 + filter_instance.tool_router = zero_match_router + return SemanticToolFilterHook(filter_instance) + + expanded_tools = [ + { + "type": "function", + "name": f"srv-tool_{i}", + "description": f"Registry tool {i}", + "parameters": {"type": "object", "properties": {}}, + } + for i in range(3) + ] + + reference_hook = build_hook() + reference_hook._expand_mcp_tools = AsyncMock( # type: ignore[method-assign] + return_value=expanded_tools + ) + reference_data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "something entirely unrelated"}], + "tools": [{"type": "mcp", "server_url": "litellm_proxy", "require_approval": "never"}], + "metadata": {}, + } + reference_result = await reference_hook.async_pre_call_hook( + user_api_key_dict=Mock(), + cache=Mock(), + data=reference_data, + call_type="acompletion", + ) + + reference_tools = (reference_result or reference_data)["tools"] + mcp_references = [tool for tool in reference_tools if tool.get("type") == "mcp"] + assert len(mcp_references) == 1, "The MCP reference must survive a zero-match query" + assert set(mcp_references[0].get("allowed_tools") or []) == {tool["name"] for tool in expanded_tools}, ( + "A zero-match query must leave every expanded tool reachable through the reference" + ) + + plain_hook = build_hook() + plain_data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "something entirely unrelated"}], + "tools": list(registry_tools), + "metadata": {}, + } + plain_result = await plain_hook.async_pre_call_hook( + user_api_key_dict=Mock(), + cache=Mock(), + data=plain_data, + call_type="acompletion", + ) + + plain_tools = (plain_result or plain_data)["tools"] + assert len(plain_tools) == len(registry_tools), ( + f"A zero-match query must not drop plain MCP tools, got {len(plain_tools)} of {len(registry_tools)}" + ) + + print("✅ zero matches: both the MCP reference path and the plain tool path expose every tool") @pytest.mark.asyncio @@ -958,8 +1198,9 @@ async def test_semantic_filter_hook_filters_expanded_tools_with_string_input(): async def test_semantic_filter_hook_expansion_skips_filter_when_disabled(): """ When the filter is disabled at runtime (e.g. via the UI toggle), the - expansion path must forward all expanded tools and emit NO filter - stats, mirroring the generic path's enabled guard. + expansion path must leave the MCP reference untouched and emit NO filter + stats, mirroring the generic path's enabled guard. The MCP gateway then + expands the reference itself, so no tool is narrowed away. """ from litellm.proxy._experimental.mcp_server.semantic_tool_filter import ( SemanticMCPToolFilter, @@ -1009,13 +1250,19 @@ async def test_semantic_filter_hook_expansion_skips_filter_when_disabled(): call_type="aresponses", ) - assert result is not None, "Hook should still expand MCP references when the filter is disabled" - assert len(result["tools"]) == 5, f"All expanded tools must be forwarded when disabled, got {len(result['tools'])}" + assert result is None, "Hook must not modify the request when the filter is disabled" + assert data["tools"] == [ + { + "type": "mcp", + "server_url": "litellm_proxy", + "require_approval": "never", + } + ], "The MCP reference must be left intact for the MCP gateway to expand" assert ( - "litellm_semantic_filter_stats" not in result["metadata"] + "litellm_semantic_filter_stats" not in data["metadata"] ), "No filter stats may be emitted when the filter is disabled" - print("✅ Disabled filter: expansion preserved, no spurious stats") + print("✅ Disabled filter: MCP reference untouched, no spurious stats") @pytest.mark.asyncio