fix(mcp): keep the MCP reference intact when the semantic filter narrows tools

litellm/proxy/hooks/mcp_semantic_filter/hook.py

@@ -155,6 +155,40 @@ async def _filter_expanded_tools(
 
         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 @@ async def async_pre_call_hook(
         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 "

tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py

@@ -865,24 +865,264 @@ async def mock_embedding_async(*args, **kwargs):
     )
 
     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"]
     ), "Filter stats must be emitted for the litellm_proxy expansion path"
     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(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"]
 
-    print(f"✅ Expanded litellm_proxy tools filtered: {len(expanded_tools)} -> {len(filtered)}, stats={stats}")
+    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 mock_embedding_async(*args, **kwargs):
 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