fix(mcp): preserve native tools in semantic filter hook (#26650)

* fix(mcp): preserve native tools in semantic filter hook

The SemanticToolFilterHook.async_pre_call_hook passed ALL tools (MCP +
native) to filter_tools(), which only knows MCP-registered tool names.
Native tools silently failed the name match in _get_tools_by_names()
and were dropped from the request.

Fix: partition tools into native and MCP-registered before filtering.
Run the semantic filter only on MCP tools, then merge native tools
back unconditionally.

Changes:
- Robust _is_mcp_tool() using shape-based detection for OpenAI-format
  dicts, safe regardless of future _extract_tool_info changes
- Single-pass partition loop (no double _is_mcp_tool calls)
- Preserve native tools in MCP expansion path (mixed requests)
- Track MCP expansion to prevent expanded tools bypassing filtering
- filter_stats reports MCP-only counts for accurate metrics
- Extracted _emit_filter_metadata() helper
- Skip spurious filter headers for all-native tool requests

Closes #26212

* remove stale docstring note referencing tools_expanded_from_mcp

* fix: handle Responses API name collision and preserve tool ordering

- Classify Responses API tools ({type: 'function', name: '...'}) as
  native to prevent name collisions with MCP canonical names
- Preserve original request tool ordering using id()-based merge
  instead of naive native+mcp concatenation
- Add 2 regression tests: name collision and ordering preservation

* style: apply black formatting

* fix(mcp): harden semantic filter — preserve all native tool formats, safe metadata access, graceful expansion failure, name-based merge

* lint: suppress PLR0915 on async_pre_call_hook (matches codebase convention)

* ci: retrigger checks after rebase onto litellm_internal_staging
This commit is contained in:
Ayush Shekhar 2026-06-17 17:42:45 +05:30 committed by GitHub
parent 82aa40b7b7
commit 438c825bd4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 588 additions and 52 deletions

View file

@ -123,12 +123,78 @@ class SemanticToolFilterHook(CustomLogger):
return openai_tools_as_dicts
def _is_mcp_tool(self, tool: Any) -> bool:
"""
Check whether *tool* is registered in the MCP semantic router.
Classification strategy:
1. Standard OpenAI Chat Completions-format dicts
(``{"type": "function", "function": {"name": ...}}``) are
**always** classified as native detected by shape (both
``"function"`` and ``"type"`` keys present).
2. Responses API function tools
(``{"type": "function", "name": ..., ...}``) are **always**
classified as native detected by ``type == "function"``
with a top-level string ``name`` but no nested ``"function"``
dict. This prevents name collisions with MCP canonicals.
3. Everything else (MCP ``Tool`` objects, flat dicts, etc.) is
looked up by name in ``self.filter._tool_map``.
"""
# Chat Completions format: {"type": "function", "function": {...}}
if isinstance(tool, dict) and "function" in tool and "type" in tool:
return False # Standard OpenAI function tool — always native
# Responses API format: {"type": "function", "name": "...", ...}
if (
isinstance(tool, dict)
and tool.get("type") == "function"
and isinstance(tool.get("name"), str)
):
return False # Responses API function tool — always native
name, _ = self.filter._extract_tool_info(tool)
return bool(name) and name in self.filter._tool_map
def _get_metadata_variable_name(self, data: dict) -> str:
if "litellm_metadata" in data:
return "litellm_metadata"
return "metadata"
async def async_pre_call_hook(
def _emit_filter_metadata(
self,
data: dict,
mcp_tools: list,
filtered_mcp_tools: list,
native_tools: list,
filtered_tools: list,
) -> None:
"""
Emit response-header metadata when MCP tools were filtered.
Stats report MCP-only counts so downstream consumers see accurate
semantic filter metrics (not inflated by native tools that always
pass through). Skips metadata entirely for purely-native requests
to avoid spurious ``N->N`` headers.
"""
if mcp_tools:
filter_stats = f"{len(mcp_tools)}->{len(filtered_mcp_tools)}"
tool_names_csv = self._get_tool_names_csv(filtered_mcp_tools)
_metadata_variable_name = self._get_metadata_variable_name(data)
metadata = data.setdefault(_metadata_variable_name, {})
metadata["litellm_semantic_filter_stats"] = filter_stats
metadata["litellm_semantic_filter_tools"] = tool_names_csv
verbose_proxy_logger.info(
f"Semantic tool filter: {filter_stats} MCP tools "
f"({len(native_tools)} native preserved, "
f"{len(filtered_tools)} total)"
)
else:
verbose_proxy_logger.info(
f"Semantic tool filter: all {len(native_tools)} tools "
f"are native, no MCP filtering applied"
)
async def async_pre_call_hook( # noqa: PLR0915
self,
user_api_key_dict: "UserAPIKeyAuth",
cache: "DualCache",
@ -163,30 +229,54 @@ class SemanticToolFilterHook(CustomLogger):
verbose_proxy_logger.debug("No tools in request, skipping semantic filter")
return None
original_tool_count = len(tools)
# Check for MCP references (server_url="litellm_proxy") and expand them
# Check for MCP references (server_url="litellm_proxy") and expand them.
# Expanded MCP tools are in OpenAI nested format which
# filter_tools/_extract_tool_info cannot name-match, so we skip
# semantic filtering and return early — expanded tools already
# went through access-control filtering in
# _process_mcp_tools_to_openai_format.
if self._should_expand_mcp_tools(tools):
verbose_proxy_logger.debug(
"Detected litellm_proxy MCP references, expanding before semantic filtering"
)
try:
# Preserve all non-MCP-reference tools before expansion.
# _expand_mcp_tools only returns expanded MCP tools;
# non-MCP tools in the original list are discarded by
# _parse_mcp_tools. Without this, native tools in
# mixed requests (MCP refs + native) are silently lost.
# We detect MCP references by type=="mcp" (the only
# shape _should_expand_mcp_tools triggers on), keeping
# both Chat Completions and Responses API native formats.
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(
"No MCP tools expanded, preserving "
f"{len(native_tools_before_expand)} native tools"
)
return data
verbose_proxy_logger.warning(
"No tools expanded from MCP references"
)
return None
data["tools"] = native_tools_before_expand + expanded_tools
verbose_proxy_logger.info(
f"Expanded {len(tools)} MCP reference(s) to {len(expanded_tools)} tools"
f"Expanded MCP references to {len(expanded_tools)} tools "
f"({len(native_tools_before_expand)} native preserved), "
f"skipping semantic filter (OpenAI nested format)"
)
# Update tools for filtering
tools = expanded_tools
original_tool_count = len(tools)
return data
except Exception as e:
verbose_proxy_logger.error(
@ -218,33 +308,71 @@ class SemanticToolFilterHook(CustomLogger):
)
return None
# ── Partition tools into MCP-registered and native ──
native_tools: List[Any] = []
mcp_tools: List[Any] = []
mcp_indices: set = set()
for i, t in enumerate(tools):
if self._is_mcp_tool(t):
mcp_tools.append(t)
mcp_indices.add(i)
else:
native_tools.append(t)
verbose_proxy_logger.debug(
f"Applying semantic filter to {len(tools)} tools "
f"with query: '{user_query[:50]}...'"
f"Applying semantic filter: {len(mcp_tools)} MCP tools, "
f"{len(native_tools)} native tools, "
f"query: '{user_query[:50]}...'"
)
# Filter tools semantically
filtered_tools = await self.filter.filter_tools(
query=user_query,
available_tools=tools, # type: ignore
)
# ── Filter only MCP tools semantically ──
if mcp_tools:
filtered_mcp_tools = await self.filter.filter_tools(
query=user_query,
available_tools=mcp_tools, # type: ignore
)
else:
filtered_mcp_tools = []
# Always update tools and emit header (even if count unchanged)
# Build set of surviving MCP tool names for order-preserving
# merge. Uses names (not id()) so the merge stays correct
# even if filter_tools internally copies tool objects.
filtered_mcp_names: set = set()
for t in filtered_mcp_tools:
name, _ = self.filter._extract_tool_info(t)
if name:
filtered_mcp_names.add(name)
# Merge preserving original request order — some LLMs are
# sensitive to the position tools appear in.
filtered_tools: List[Any] = []
for i, t in enumerate(tools):
if i in mcp_indices:
name, _ = self.filter._extract_tool_info(t)
if name in filtered_mcp_names:
filtered_tools.append(t)
else:
filtered_tools.append(t)
# Always update tools
data["tools"] = filtered_tools
# Store filter stats and tool names for response header
filter_stats = f"{original_tool_count}->{len(filtered_tools)}"
tool_names_csv = self._get_tool_names_csv(filtered_tools)
_metadata_variable_name = self._get_metadata_variable_name(data)
data[_metadata_variable_name][
"litellm_semantic_filter_stats"
] = filter_stats
data[_metadata_variable_name][
"litellm_semantic_filter_tools"
] = tool_names_csv
verbose_proxy_logger.info(f"Semantic tool filter: {filter_stats} tools")
# Emit metadata — wrapped separately so a metadata error
# does not cause the outer except to return None after
# data["tools"] has already been mutated.
try:
self._emit_filter_metadata(
data=data,
mcp_tools=mcp_tools,
filtered_mcp_tools=filtered_mcp_tools,
native_tools=native_tools,
filtered_tools=filtered_tools,
)
except Exception as e:
verbose_proxy_logger.warning(
f"Failed to emit semantic filter metadata: {e}",
exc_info=True,
)
return data
@ -266,7 +394,7 @@ class SemanticToolFilterHook(CustomLogger):
from litellm.constants import MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH
_metadata_variable_name = self._get_metadata_variable_name(data)
metadata = data[_metadata_variable_name]
metadata = data.get(_metadata_variable_name, {})
filter_stats = metadata.get("litellm_semantic_filter_stats")
if not filter_stats:

View file

@ -452,6 +452,428 @@ async def test_semantic_filter_hook_skips_no_tools():
print("✅ Hook correctly skips requests without tools")
@pytest.mark.asyncio
async def test_semantic_filter_hook_preserves_native_tools():
"""
Regression test: mixed MCP + native tools.
Given: 5 MCP tools (registered in _tool_map) + 2 native OpenAI-format
function tools (not in _tool_map)
When: The hook filters tools
Then: The native tools must survive unconditionally, and only MCP
tools go through the semantic filter.
"""
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,
)
# --- MCP tools (registered in the semantic router) ---
mcp_tools = [
MCPTool(
name=f"mcp_tool_{i}",
description=f"MCP tool {i}",
inputSchema={"type": "object"},
)
for i in range(5)
]
filter_instance._build_router(mcp_tools)
# --- Native OpenAI-format function tools (NOT in _tool_map) ---
native_tools = [
{
"type": "function",
"function": {
"name": "get_current_weather",
"description": "Get the current weather",
"parameters": {"type": "object", "properties": {}},
},
},
{
"type": "function",
"function": {
"name": "search_web",
"description": "Search the web",
"parameters": {"type": "object", "properties": {}},
},
},
]
# Combine: MCP tools + native tools
all_tools = list(mcp_tools) + native_tools
hook = SemanticToolFilterHook(filter_instance)
data = {
"model": "gpt-4",
"messages": [{"role": "user", "content": "What is the weather?"}],
"tools": all_tools,
"metadata": {},
}
result = await hook.async_pre_call_hook(
user_api_key_dict=Mock(),
cache=Mock(),
data=data,
call_type="completion",
)
assert result is not None, "Hook should return modified data"
filtered = result["tools"]
# Native tools must survive
native_in_result = [
t for t in filtered if isinstance(t, dict) and t.get("type") == "function"
]
assert (
len(native_in_result) == 2
), f"Both native tools must survive, got {len(native_in_result)}"
# MCP tools should be filtered (top_k=2)
mcp_in_result = [t for t in filtered if not isinstance(t, dict)]
assert (
len(mcp_in_result) <= 2
), f"MCP tools should be filtered to top_k=2, got {len(mcp_in_result)}"
# Total should be native + filtered MCP
assert len(filtered) <= 4, f"Expected at most 4 tools, got {len(filtered)}"
# Filter stats should be emitted (MCP tools were present)
assert "litellm_semantic_filter_stats" in result["metadata"]
# Stats should report MCP-only counts, not inflated with native tools
stats = result["metadata"]["litellm_semantic_filter_stats"]
mcp_before, mcp_after = stats.split("->")
assert (
int(mcp_before) == 5
), f"Stats 'from' should be MCP count (5), got {mcp_before}"
assert int(mcp_after) == len(
mcp_in_result
), f"Stats 'to' should match filtered MCP count, got {mcp_after}"
print(
f"✅ Hook preserves native tools: {len(all_tools)} -> {len(filtered)} "
f"({len(native_in_result)} native + {len(mcp_in_result)} MCP), "
f"stats={stats}"
)
@pytest.mark.asyncio
async def test_semantic_filter_hook_all_native_tools():
"""
Regression test: all-native request.
Given: Only native OpenAI-format function tools (none registered in
the MCP semantic router)
When: The hook processes the request
Then: All tools pass through, and NO spurious semantic filter response
headers are emitted (no litellm_semantic_filter_stats in metadata).
"""
from litellm.proxy._experimental.mcp_server.semantic_tool_filter import (
SemanticMCPToolFilter,
)
from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook
mock_router = Mock()
filter_instance = SemanticMCPToolFilter(
embedding_model="text-embedding-3-small",
litellm_router_instance=mock_router,
top_k=3,
similarity_threshold=0.3,
enabled=True,
)
# Build router with some MCP tools (so tool_router is not None)
mcp_tools = [
MCPTool(
name="some_mcp_tool",
description="An MCP tool",
inputSchema={"type": "object"},
)
]
from litellm.types.utils import Embedding, EmbeddingResponse
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._build_router(mcp_tools)
# --- Only native tools in the request ---
native_tools = [
{
"type": "function",
"function": {
"name": f"native_func_{i}",
"description": f"Native function {i}",
"parameters": {"type": "object", "properties": {}},
},
}
for i in range(3)
]
hook = SemanticToolFilterHook(filter_instance)
data = {
"model": "gpt-4",
"messages": [{"role": "user", "content": "Hello"}],
"tools": native_tools,
"metadata": {},
}
result = await hook.async_pre_call_hook(
user_api_key_dict=Mock(),
cache=Mock(),
data=data,
call_type="completion",
)
assert result is not None, "Hook should return modified data"
filtered = result["tools"]
# All native tools must pass through
assert (
len(filtered) == 3
), f"All 3 native tools must pass through, got {len(filtered)}"
# No spurious semantic filter stats (P2 fix)
assert (
"litellm_semantic_filter_stats" not in result["metadata"]
), "Should NOT emit semantic filter stats for all-native-tool requests"
print(
f"✅ Hook passes through all {len(filtered)} native tools, "
f"no spurious filter headers emitted"
)
@pytest.mark.asyncio
async def test_semantic_filter_hook_responses_api_name_collision():
"""
Regression test: Responses API native tool with MCP-matching name.
Given: A Responses-API native tool whose top-level ``name`` collides
with an MCP canonical name in ``_tool_map``
When: The hook classifies tools
Then: The native tool must NOT be sent to the semantic filter, even
though its name matches an MCP canonical.
"""
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,
)
# Register an MCP tool with name "github-search"
mcp_tools = [
MCPTool(
name="github-search",
description="Search GitHub repos",
inputSchema={"type": "object"},
)
]
filter_instance._build_router(mcp_tools)
# Responses API native tool with SAME name as MCP canonical
responses_api_tool = {
"type": "function",
"name": "github-search",
"description": "Caller-owned search tool",
"parameters": {"type": "object"},
}
hook = SemanticToolFilterHook(filter_instance)
# Verify classification: should be native, not MCP
assert not hook._is_mcp_tool(responses_api_tool), (
"Responses API tool with type=function + top-level name "
"should be classified as native, not MCP"
)
# Full hook test: all-native request should preserve tools
data = {
"model": "gpt-4",
"messages": [{"role": "user", "content": "Search GitHub"}],
"tools": [responses_api_tool],
"metadata": {},
}
result = await hook.async_pre_call_hook(
user_api_key_dict=Mock(),
cache=Mock(),
data=data,
call_type="completion",
)
# All tools are native → hook returns data with all tools preserved
filtered = (result or data)["tools"]
assert len(filtered) == 1, f"Native tool must survive, got {len(filtered)}"
assert filtered[0]["name"] == "github-search"
print("✅ Responses API tool with MCP-matching name correctly classified as native")
@pytest.mark.asyncio
async def test_semantic_filter_hook_preserves_tool_order():
"""
Regression test: tool ordering preservation.
Given: An interleaved request [mcp_tool_A, native_tool, mcp_tool_B]
When: The hook filters tools (all MCP tools survive)
Then: The output order must match the original request order,
NOT native-first.
"""
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=5,
similarity_threshold=0.3,
enabled=True,
)
# Register MCP tools
mcp_tool_a = MCPTool(
name="github-search",
description="Search GitHub",
inputSchema={"type": "object"},
)
mcp_tool_b = MCPTool(
name="github-issue",
description="Create GitHub issue",
inputSchema={"type": "object"},
)
filter_instance._build_router([mcp_tool_a, mcp_tool_b])
# Mock filter_tools to return both MCP tools (deterministic)
filter_instance.filter_tools = AsyncMock( # type: ignore[method-assign]
return_value=[mcp_tool_a, mcp_tool_b]
)
# Native tool (interleaved between MCP tools)
native_tool = {
"type": "function",
"function": {
"name": "weather_lookup",
"description": "Look up weather",
"parameters": {"type": "object", "properties": {}},
},
}
# Original order: [mcp_A, native, mcp_B]
original_tools = [mcp_tool_a, native_tool, mcp_tool_b]
hook = SemanticToolFilterHook(filter_instance)
data = {
"model": "gpt-4",
"messages": [{"role": "user", "content": "Search GitHub and check weather"}],
"tools": original_tools,
"metadata": {},
}
result = await hook.async_pre_call_hook(
user_api_key_dict=Mock(),
cache=Mock(),
data=data,
call_type="completion",
)
assert result is not None, "Hook should return modified data"
filtered = result["tools"]
# All tools should survive
assert len(filtered) == 3, f"Expected 3 tools, got {len(filtered)}"
# Order must be preserved: [mcp_A, native, mcp_B]
assert filtered[0] is mcp_tool_a, "First tool should be mcp_tool_a"
assert filtered[1] is native_tool, "Second tool should be native_tool"
assert filtered[2] is mcp_tool_b, "Third tool should be mcp_tool_b"
print("✅ Tool ordering preserved: [mcp_A, native, mcp_B] maintained after filtering")
class TestGetToolsByNames:
"""
Regression coverage for SemanticMCPToolFilter._get_tools_by_names
@ -489,9 +911,7 @@ class TestGetToolsByNames:
{"name": "send_email", "description": "send mail"},
]
matched = filter_instance._get_tools_by_names(
["send_email"], available_tools
)
matched = filter_instance._get_tools_by_names(["send_email"], available_tools)
assert len(matched) == 1
assert matched[0]["name"] == "send_email"
@ -503,9 +923,7 @@ class TestGetToolsByNames:
client_name = "litellm_" + canonical
available_tools = [{"name": client_name, "description": "scrape"}]
matched = filter_instance._get_tools_by_names(
[canonical], available_tools
)
matched = filter_instance._get_tools_by_names([canonical], available_tools)
assert len(matched) == 1
# Must return the incoming tool unchanged so the client-facing
@ -516,13 +934,9 @@ class TestGetToolsByNames:
"""Some clients use dash as alias separator; accept that too."""
filter_instance = self._make_filter()
canonical = "weather_svc-get_weather"
available_tools = [
{"name": "mcp-" + canonical, "description": "weather"}
]
available_tools = [{"name": "mcp-" + canonical, "description": "weather"}]
matched = filter_instance._get_tools_by_names(
[canonical], available_tools
)
matched = filter_instance._get_tools_by_names([canonical], available_tools)
assert len(matched) == 1
assert matched[0]["name"] == "mcp-" + canonical
@ -552,9 +966,7 @@ class TestGetToolsByNames:
{"name": "litellm_" + canonical, "description": "wrapped"},
]
matched = filter_instance._get_tools_by_names(
[canonical], available_tools
)
matched = filter_instance._get_tools_by_names([canonical], available_tools)
assert len(matched) == 1
assert matched[0]["name"] == canonical
@ -567,9 +979,7 @@ class TestGetToolsByNames:
separator-anchored suffixes of ``litellm_api-fs-read_file``.
"""
filter_instance = self._make_filter()
available_tools = [
{"name": "litellm_api-fs-read_file", "description": "read"}
]
available_tools = [{"name": "litellm_api-fs-read_file", "description": "read"}]
matched = filter_instance._get_tools_by_names(
["fs-read_file", "api-fs-read_file"], available_tools
@ -590,9 +1000,7 @@ class TestGetToolsByNames:
{"name": "my_" + canonical, "description": "plain search"},
]
matched = filter_instance._get_tools_by_names(
[canonical], available_tools
)
matched = filter_instance._get_tools_by_names([canonical], available_tools)
assert len(matched) == 1
assert matched[0]["name"] == "my_" + canonical