fix(mcp): validate MCP tool prefix against registry instead of heuristic - #25088
fix(mcp): validate MCP tool prefix against registry instead of heuristic#25088klhq wants to merge 7 commits into
Conversation
When the semantic tool filter finds no matches for a user query, it previously returned all available tools. With many MCP servers this can exceed provider tool limits. Now it drops MCP tools (prefixed) and returns only non-MCP tools on zero matches.
The semantic filter treated all tools equally — it didn't distinguish between MCP tools (which it should filter) and non-MCP built-in tools (which should pass through untouched). This caused two failures: - On match: non-MCP tools were dropped, losing built-in functionality - On zero match: all tools returned unfiltered, exceeding provider limits Move MCP/non-MCP separation into the hook. The hook now passes only MCP tools (identified by server-name prefix) to filter_tools() and recombines with non-MCP tools after. filter_tools() returns empty list on zero matches instead of all tools.
…CP tools When all tools are MCP and zero semantic matches are found, return the first top_k MCP tools instead of an empty list. Avoids sending an empty tool list to the LLM.
is_tool_name_prefixed() just checks if "-" exists in the name, which misclassifies non-MCP tools with hyphens (e.g., text-to-speech). Replace with _is_mcp_tool() that splits on the first "-" and checks if the prefix is a known registered MCP server name from the registry.
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
klhq
changed the title
Greptile SummaryThis PR fixes a misclassification bug where non-MCP tools containing hyphens (e.g., Key changes:
One minor concern: Confidence Score: 5/5Safe to merge — the core bug fix is correct, tests are thorough and mock-only, and the only remaining issue is a minor per-request allocation inefficiency. All findings are P2 (style/performance). The registry-queried-per-tool concern does not cause incorrect behavior — it is purely an efficiency issue on an already-hot path that could be addressed in a follow-up. No correctness, security, or data-integrity issues were found. The stale-cache concern from the previous review thread has been resolved (registry is now queried fresh). Tests are well-structured and cover the new logic branches. hook.py — the partition loop calls _get_registered_server_prefixes() once per tool; trivial to hoist the call outside the loop.
|
| Filename | Overview |
|---|---|
| litellm/proxy/hooks/mcp_semantic_filter/hook.py | Core hook rewritten to separate MCP vs non-MCP tools via registry lookup; non-MCP tools now bypass semantic filtering entirely. Registry is queried once per tool in the loop (O(N) allocations per request) — a minor performance concern. |
| litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py | Changed filter_tools() to return [] (instead of available_tools) when no semantic matches are found; hook is now responsible for adding non-MCP tools back, making the contract cleaner. |
| tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py | Existing test correctly updated with registry mock and prefixed tool names; 6 new tests cover hyphenated non-MCP pass-through, fallback-to-top-k, zero-match semantics, and _is_mcp_tool unit cases. All tests are mock-only (no real network calls). |
Flowchart
%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[async_pre_call_hook] --> B{tools present?}
B -- No --> Z[return None]
B -- Yes --> C{MCP refs to expand?}
C -- Yes --> D[expand MCP references]
D --> E[continue with expanded tools]
C -- No --> E
E --> F{messages present?}
F -- No --> Z
F -- Yes --> G[extract user query]
G --> H{query found?}
H -- No --> Z
H -- Yes --> I[partition tools via _is_mcp_tool]
I --> I1[_get_registered_server_prefixes called once per tool ⚠️]
I --> J{any MCP tools?}
J -- No --> Z
J -- Yes --> K[filter_tools on MCP tools only]
K --> L{zero MCP matches?}
L -- No --> M[filtered_mcp_tools + non_mcp_tools]
L -- Yes --> N{non_mcp_tools empty?}
N -- Yes --> O[fallback: first top_k MCP tools]
O --> M
N -- No --> P[filtered_mcp_tools = empty]
P --> M
M --> Q[update data and return]
Reviews (2): Last reviewed commit: "fix: remove stale prefix cache and use s..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| def _get_registered_server_prefixes(self) -> set: | ||
| """Get the set of known MCP server prefixes from the registry.""" | ||
| if self._registered_server_prefixes is None: | ||
| from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( | ||
| global_mcp_server_manager, | ||
| ) | ||
|
|
||
| registry = global_mcp_server_manager.get_registry() | ||
| self._registered_server_prefixes = { | ||
| normalize_server_name(get_server_prefix(server)) | ||
| for server in registry.values() | ||
| if get_server_prefix(server) | ||
| } | ||
| return self._registered_server_prefixes |
There was a problem hiding this comment.
_registered_server_prefixes is set once on the first request and never cleared. If a new MCP server is added dynamically via the admin API after that first call, its prefix will be absent from the cached set. The hook will then classify that server's tools as non_mcp_tools — they will bypass semantic filtering entirely and always reach the LLM.
The cache is populated lazily with:
if self._registered_server_prefixes is None:
...But self._registered_server_prefixes is never reset to None after that. Any runtime POST /mcp/server call that adds a new server is invisible to this logic.
Consider one of:
- Clearing
_registered_server_prefixes(reset toNone) wheneverglobal_mcp_server_manager's registry changes — e.g. hook into the registration/removal events. - Using a TTL-based cache (e.g. re-query every N seconds).
- Querying the registry fresh on each call if the registry is cheap to read (it returns an in-memory dict).
There was a problem hiding this comment.
Fixed in 45dbc8c. Removed the cache entirely. get_registry() returns an in-memory dict, so querying it fresh each call is negligible compared to the embedding call that follows. No need to cache a copy of something already in memory.
| mcp_tools = [t for t in tools if self._is_mcp_tool(_tool_name(t))] | ||
| non_mcp_tools = [t for t in tools if not self._is_mcp_tool(_tool_name(t))] |
There was a problem hiding this comment.
Each tool has _is_mcp_tool(_tool_name(t)) called twice — once for each list comprehension — and both comprehensions iterate the full tools list. For large tool lists this is 2× the work needed.
A single pass that partitions tools would be cleaner and more efficient:
mcp_tools, non_mcp_tools = [], []
for t in tools:
(mcp_tools if self._is_mcp_tool(_tool_name(t)) else non_mcp_tools).append(t)This also makes the intent explicit: every tool goes into exactly one bucket.
There was a problem hiding this comment.
Fixed in 45dbc8c, replaced with a single-pass partition loop.
Query the MCP registry fresh each call so dynamically added servers are recognized immediately. Replace double list comprehension with a single-pass partition loop.
|
This PR builds on #24986 which separates MCP from non-MCP tools in the hook. #24986 used This PR replaces that heuristic with a registry lookup: split on the first |
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
|
Closing — no longer needed. #25085 (which fixes the same underlying #24986 picks up the new |
1 participant
Relevant issues
Fixes #25081
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
tests/test_litellm/directory, Adding at least 1 test is a hard requirement - see detailsmake test-unit@greptileaiand received a Confidence Score of at least 4/5 before requesting a maintainer reviewDelays in PR merge?
If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).
CI (LiteLLM team)
Branch creation CI run
Link:
CI run for the last commit
Link:
Merge / cherry-pick CI run
Links:
Type
🐛 Bug Fix
Changes
is_tool_name_prefixed()checks"-" in tool_name, which misclassifies non-MCP tools with hyphens (e.g.,text-to-speech,code-review) as MCP tools. These tools get sent through the semantic filter and silently dropped.is_tool_name_prefixed()in the hook with_is_mcp_tool()that splits on the first-and validates the prefix against known MCP server names from the registryglobal_mcp_server_manager.get_registry())_is_mcp_toolunit test