scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

fix(mcp): scope semantic filter matching and indexing errors to the requesting call

Browse source 
Match candidates are restricted to the request's own tool names via
route_filter, so routes learned from other principals' listings cannot
displace the caller's tools from top_k. Lazy indexing now uses the
async aadd flow exclusively; an embedding failure, including a
context-window overflow on an oversized description, raises for the
requesting call only and never writes the shared context_window_error,
so one request cannot poison the filter for every user on the worker.
The router is also sized to the configured top_k, which the
semantic-router index layer otherwise silently caps at its default
of 5.
This commit is contained in:
Tin Chi Lo 2026-07-14 19:34:17 -07:00
parent a6390abefa
commit 8ca809d426
2 changed files with 135 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

51
litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py
View file

@ -206,9 +206,19 @@ class SemanticMCPToolFilter:
the filter came through an authenticated expansion; without indexing
them here they can never be selected, so requests either bypass
filtering entirely (N->N) or lose every tool to unrelated matches.
Runs async-only (no synchronous embedding on the request path) and
never writes shared error state: an embedding failure here raises and
is scoped to the requesting call, so one request's oversized tool
description cannot poison the filter for other users on the worker.
"""
from semantic_router.routers import SemanticRouter
from semantic_router.routers.base import Route
from litellm.router_strategy.auto_router.litellm_encoder import (
LiteLLMRouterEncoder,
)
if not self._has_tools_missing_from_index(available_tools):
return
@ -217,14 +227,6 @@ class SemanticMCPToolFilter:
if not missing:
return
if self.tool_router is None:
self._build_router(list(missing.values()))
if self.tool_router is not None:
verbose_logger.info(
f"Semantic tool filter indexed {len(missing)} request-time tools missing from the startup index"
)
return
descriptions = {name: self._extract_tool_info(tool)[1] for name, tool in missing.items()}
routes = [
Route(
@ -235,7 +237,23 @@ class SemanticMCPToolFilter:
)
for name, description in descriptions.items()
]
await self.tool_router.aadd(routes)
if self.tool_router is None:
router = SemanticRouter(
routes=[],
encoder=LiteLLMRouterEncoder(
litellm_router_instance=self.router_instance,
model_name=self.embedding_model,
score_threshold=self.similarity_threshold,
),
auto_sync="local",
top_k=self.top_k,
)
await router.aadd(routes)
self.tool_router = router
else:
await self.tool_router.aadd(routes)
self._tool_map.update(missing)
verbose_logger.info(
f"Semantic tool filter indexed {len(routes)} request-time tools missing from the startup index"
@ -279,19 +297,18 @@ class SemanticMCPToolFilter:
try:
await self._ensure_tools_indexed(available_tools)
if self.context_window_error is not None:
raise SemanticToolFilterContextWindowError(
embedding_model=self.embedding_model,
stage="the MCP tool descriptions during semantic router build",
original_error=self.context_window_error,
)
if self.tool_router is None:
verbose_logger.warning("Semantic router could not be built from the request's tools")
return available_tools
available_names = [name for name in (self._extract_tool_info(t)[0] for t in available_tools) if name]
if not available_names:
return available_tools
limit = top_k or self.top_k
matches = self.tool_router(text=query, limit=limit)
if self.tool_router.top_k < limit:
self.tool_router.top_k = limit
matches = self.tool_router(text=query, limit=limit, route_filter=available_names)
matched_tool_names = self._extract_tool_names_from_matches(matches)
if not matched_tool_names:

101
tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py
View file

@ -1678,6 +1678,8 @@ def _make_keyword_embedding_router(recorded_inputs):
def _vector(text):
lowered = text.lower()
if "kanban" in lowered:
return [0.6, 0.8]
if "linear" in lowered or "issue" in lowered or "ticket" in lowered:
return [1.0, 0.0]
return [0.0, 1.0]
@ -1812,3 +1814,102 @@ async def test_filter_fails_open_when_matches_are_not_in_available_tools():
assert [t.name for t in filtered] == ["linear_stub-get_issue", "linear_stub-list_issues"]
print("✅ Matches outside available_tools fail open instead of dropping every tool")
@pytest.mark.asyncio
async def test_request_time_context_window_error_is_request_scoped():
"""
Regression test: an oversized tool description hitting the embedding
context window while lazily indexing request-time tools must fail only
the requesting call. Previously the lazy path reused the startup build
and recorded the overflow in the shared context_window_error, after
which EVERY user's MCP requests on the worker were blocked with a 400
until restart (index poisoning via a single request).
"""
from litellm.proxy._experimental.mcp_server.semantic_tool_filter import (
SemanticToolFilterContextWindowError,
)
state = {"raise_context_error": True}
filter_instance = _make_context_window_filter(state)
tools = [
MCPTool(name="tool_a", description="Tool A", inputSchema={"type": "object"}),
MCPTool(name="tool_b", description="Tool B", inputSchema={"type": "object"}),
]
with pytest.raises(SemanticToolFilterContextWindowError):
await filter_instance.filter_tools(query="send an email", available_tools=tools)
assert filter_instance.context_window_error is None
assert filter_instance.tool_router is None
state["raise_context_error"] = False
filtered = await filter_instance.filter_tools(query="send an email", available_tools=tools)
assert len(filtered) > 0
assert filter_instance.context_window_error is None
assert filter_instance.tool_router is not None
print("✅ Request-time context window overflow is scoped to the request, not the worker")
@pytest.mark.asyncio
async def test_foreign_index_routes_cannot_displace_available_tools():
"""
Regression test: routes indexed from OTHER principals' tool listings must
not occupy the match candidate set for this request. Previously the
router matched over the whole shared index, so foreign routes that
embedded closer to the query displaced the caller's own tools from
top_k, degrading results to the fail-open list (or, before the
empty-result guard, stripping every tool). Matching is now scoped to the
request's own tool names via route_filter.
"""
filter_instance = _make_keyword_filter([], top_k=1)
foreign_tools = [
MCPTool(
name=f"other_user-linear_tool_{i}",
description=f"Get a Linear issue variant {i}",
inputSchema={"type": "object"},
)
for i in range(6)
]
filter_instance._build_router(foreign_tools)
my_kanban = MCPTool(
name="mine-kanban_board",
description="Manage kanban board cards",
inputSchema={"type": "object"},
)
filtered = await filter_instance.filter_tools(
query="what is Linear ticket LIT-3794 about",
available_tools=[my_kanban, _weather_tool()],
)
assert [t.name for t in filtered] == ["mine-kanban_board"]
print("✅ Foreign index routes cannot displace the caller's own tools")
@pytest.mark.asyncio
async def test_top_k_above_router_default_is_respected():
"""
Regression test: semantic-router's SemanticRouter defaults to top_k=5 at
the index-query layer, silently capping any configured filter top_k
above 5 regardless of the limit passed to __call__. The router must be
sized (and resized) to honor the configured top_k.
"""
filter_instance = _make_keyword_filter([], top_k=6)
tools = [
MCPTool(
name=f"linear_stub-tool_{i}",
description=f"Work with Linear issues part {i}",
inputSchema={"type": "object"},
)
for i in range(6)
]
filtered = await filter_instance.filter_tools(
query="Linear ticket work",
available_tools=tools,
)
assert len(filtered) == 6
print("✅ Configured top_k above the semantic-router default of 5 is honored")