fix(mcp): re-land native tool preservation with typed annotations (#30645)

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

* fix(mcp): tighten _is_mcp_tool Chat Completions shape check
This commit is contained in:
Ayush Shekhar 2026-06-22 17:53:08 +05:30 committed by GitHub
parent 67c0183fcf
commit 0a359ceb20
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 560 additions and 65 deletions

View file

@ -123,11 +123,70 @@ class SemanticToolFilterHook(CustomLogger):
return openai_tools_as_dicts
def _is_mcp_tool(self, tool: object) -> bool:
"""
Check whether *tool* is registered in the MCP semantic router.
Classification strategy (shape-first, lookup-second):
1. Chat Completions format dicts are always native.
2. Responses API function tools are always native.
3. Everything else is looked up by name in the MCP registry.
"""
if (
isinstance(tool, dict)
and tool.get("type") == "function"
and isinstance(tool.get("function"), dict)
):
return False
if (
isinstance(tool, dict)
and tool.get("type") == "function"
and isinstance(tool.get("name"), str)
):
return False
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"
def _emit_filter_metadata(
self,
data: dict,
mcp_tools: list[object],
filtered_mcp_tools: list[object],
native_tools: list[object],
filtered_tools: list[object],
) -> None:
"""
Emit response-header metadata when MCP tools were filtered.
Stats report MCP-only counts so downstream consumers see accurate
semantic filter metrics. Skips metadata entirely for purely-native
requests to avoid spurious 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(
self,
user_api_key_dict: "UserAPIKeyAuth",
@ -140,53 +199,55 @@ class SemanticToolFilterHook(CustomLogger):
This hook is called before the LLM request is made. It filters the
tools list to only include semantically relevant tools.
Args:
user_api_key_dict: User authentication
cache: Cache instance
data: Request data containing messages and tools
call_type: Type of call (completion, acompletion, etc.)
Returns:
Modified data dict with filtered tools, or None if no changes
"""
# Only filter endpoints that support tools
if call_type not in ("completion", "acompletion", "aresponses"):
verbose_proxy_logger.debug(
f"Skipping semantic filter for call_type={call_type}"
)
return None
# Check if tools are present
tools = data.get("tools")
if not tools:
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
# 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.
if self._should_expand_mcp_tools(tools):
verbose_proxy_logger.debug(
"Detected litellm_proxy MCP references, expanding before semantic filtering"
)
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(
"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(
@ -194,7 +255,6 @@ class SemanticToolFilterHook(CustomLogger):
)
return None
# Check if messages are present (try both "messages" and "input" for responses API)
messages = data.get("messages", [])
if not messages:
messages = data.get("input", [])
@ -204,13 +264,11 @@ class SemanticToolFilterHook(CustomLogger):
)
return None
# Check if filter is enabled
if not self.filter.enabled:
verbose_proxy_logger.debug("Semantic filter disabled, skipping")
return None
try:
# Extract user query from messages
user_query = self.filter.extract_user_query(messages)
if not user_query:
verbose_proxy_logger.debug(
@ -218,33 +276,60 @@ class SemanticToolFilterHook(CustomLogger):
)
return None
native_tools: list[object] = []
mcp_tools: list[object] = []
mcp_indices: set[int] = 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
)
if mcp_tools:
filtered_mcp_tools = await self.filter.filter_tools(
query=user_query,
available_tools=mcp_tools, # type: ignore
)
else:
filtered_mcp_tools = []
filtered_mcp_names: set[str] = set()
for t in filtered_mcp_tools:
name, _ = self.filter._extract_tool_info(t)
if name:
filtered_mcp_names.add(name)
filtered_tools: list[object] = []
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 and emit header (even if count unchanged)
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")
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 +351,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,430 @@ 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 +913,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 +925,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 +936,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 +968,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 +981,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 +1002,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