feat: Search transforms for tool discovery#3154
Conversation
jlowin
added
the
feature
marvin-context-protocol
bot
added
the
server
jlowin
added
the
DON'T MERGE
WalkthroughThis pull request introduces a new Tool Search feature for FastMCP. The implementation adds a search transform system that replaces large tool catalogs with on-demand search via two concrete implementations: BM25SearchTransform (relevance-based ranking with in-memory indexing) and RegexSearchTransform (pattern matching with zero overhead). The search transforms cause list_tools to return two synthetic tools—search_tools and call_tool—enabling LLMs to discover and interact with tools dynamically. The change also removes four Protocol exports (GetPromptNext, GetResourceNext, GetResourceTemplateNext, GetToolNext) from the transforms module's public API. Comprehensive documentation is added covering usage patterns, customization options, and interaction with authentication and visibility middleware. Possibly related PRs
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: cb88a50213
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
|
||
| Use this to execute tools discovered via search_tools. | ||
| """ | ||
| return await ctx.fastmcp.call_tool(name, arguments) |
There was a problem hiding this comment.
call_tool from invoking itself
The proxy forwards any requested tool name directly to ctx.fastmcp.call_tool, so a request like call_tool(name="call_tool") resolves the same synthetic tool again and recurses until timeout/recursion failure. This is an easy request-level DoS path (and can be triggered by LLM mis-selection), so the proxy should explicitly reject self-references to its own synthetic name before dispatching.
Useful? React with 👍 / 👎.
| current_hash = _catalog_hash(tools) | ||
| if current_hash != self._last_hash: |
There was a problem hiding this comment.
The rebuild gate is keyed only by _catalog_hash(tools), and that hash is based on names, so catalogs with unchanged names but updated descriptions/parameter schemas are treated as unchanged. In dynamic providers or list-tools middleware that mutates tool metadata per request, BM25 will keep ranking against stale documents and return stale Tool objects from the previous snapshot; the staleness key should include searchable metadata, not just names.
Useful? React with 👍 / 👎.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
docs/docs.json (1)
572-583:⚠️ Potential issue | 🟠 MajorAdd search transform SDK pages to docs.json navigation.
The four new search transform documentation files exist in
docs/python-sdk/but are not referenced indocs.jsonunder the transforms group (lines 574–582):
fastmcp-server-transforms-search-__init__.mdxfastmcp-server-transforms-search-base.mdxfastmcp-server-transforms-search-bm25.mdxfastmcp-server-transforms-search-regex.mdxPer the documentation guidelines, these files must be included in
docs.jsonto be published. Add the missing entries to the transforms navigation group, or confirm whether the documentation bot auto-updates this file.
🧹 Nitpick comments (6)
src/fastmcp/server/transforms/search/base.py (1)
162-176: Bypass + filter logic is sound, but consider the multi-transform stacking scenario.If two
BaseSearchTransformsubclasses are stacked,_search_bypassis shared (module-levelContextVar). When the inner transform's search tool calls_get_visible_tools, the bypass causes both transforms to pass through, which is correct — the inner search sees the full catalog. Worth a brief note in the docstring if this stacking pattern is expected to be supported.src/fastmcp/server/transforms/search/bm25.py (1)
93-97: Consider explicit keyword arguments instead of**kwargs.
BM25SearchTransform.__init__accepts**kwargsand forwards them tosuper().__init__(). This obscures the accepted parameters from type checkers and IDE autocompletion. Explicitly declaringmax_results,always_visible,search_tool_name, andcall_tool_namewould improve discoverability.Proposed fix
- def __init__(self, **kwargs: Any) -> None: - super().__init__(**kwargs) + def __init__( + self, + *, + max_results: int = 5, + always_visible: list[str] | None = None, + search_tool_name: str = "search_tools", + call_tool_name: str = "call_tool", + ) -> None: + super().__init__( + max_results=max_results, + always_visible=always_visible, + search_tool_name=search_tool_name, + call_tool_name=call_tool_name, + )src/fastmcp/server/transforms/search/regex.py (1)
43-55: Consider ReDoS mitigation for untrusted regex patterns.The
querystring is provided by the LLM/client and compiled directly as a regex. Malicious or pathological patterns (e.g.,(a+)+$) can cause catastrophic backtracking in Python'sreengine. While the searchable text is short (server-defined tool metadata), this is still a potential denial-of-service vector if the server is exposed to untrusted clients.A lightweight mitigation would be to apply a timeout or use
re2(if available), or simply set a maximum pattern length. Even a length cap goes a long way:🛡️ Optional: add a pattern length cap
async def _search(self, tools: Sequence[Tool], query: str) -> Sequence[Tool]: + if len(query) > 200: + return [] try: compiled = re.compile(query, re.IGNORECASE) except re.error: return []Otherwise, the search logic is clean: the
re.errorcatch for invalid patterns is good, and the earlybreakon_max_resultsavoids unnecessary iteration.docs/servers/transforms/tool-search.mdx (3)
36-59: Code example is clear but uses...for function bodies.The guideline calls for "complete, runnable code examples that users can copy and execute." The
...ellipsis bodies are a reasonable shorthand for tools whose implementation doesn't matter to the example, but worth noting that a user copy-pasting this won't get a working demo. Consider adding minimal return values (e.g.,return [],return True) so the snippet is directly executable.
63-73: Client-side example lacks context forclient.This snippet uses
await client.call_tool(...)without showing howclientis created or that it runs inside anasyncfunction. A reader unfamiliar with the client setup may not be able to run this. Consider adding a brief note (e.g., "Assuming an existingClientsession") or linking to client documentation. As per coding guidelines, code blocks should be "fully runnable with all necessary imports."
115-136: Add a closing section with next steps or related information.The page ends at line 136 without a conclusion. Per the MDX documentation guidelines, sections should end with next steps or related information. Consider adding a brief "Next Steps" or "Related Topics" section that links to the transforms overview, authorization/visibility documentation, or the Python SDK reference for
BaseSearchTransform.
| def _catalog_hash(tools: Sequence[Tool]) -> str: | ||
| """SHA256 hash of sorted tool names for staleness detection.""" | ||
| key = "|".join(sorted(t.name for t in tools)) | ||
| return hashlib.sha256(key.encode()).hexdigest() |
There was a problem hiding this comment.
Catalog hash uses only tool names — description/parameter changes won't trigger reindex.
_catalog_hash hashes sorted tool names, so if a tool's description or parameters change (without adding/removing tools), the BM25 index will serve stale results. This is documented, but worth calling out since tool descriptions can be dynamically generated.
If this is intentional to keep the check cheap, consider adding a one-line comment in the hash function body noting the tradeoff.
|
Sorry, wrong link earlier — please ignore 😅 |
|
The review flagged removing |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a0c239ee6a
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| tools = await ctx.fastmcp.list_tools() | ||
| finally: | ||
| _search_bypass.reset(token) | ||
| return [t for t in tools if t.name not in self._always_visible] |
There was a problem hiding this comment.
_get_visible_tools forwards ctx.fastmcp.list_tools() directly into search/serialization, but FastMCP.list_tools() returns all versions while MCP tools/list deduplicates by name in the wire handler. With versioned tools, search can return multiple entries for one name (including older schemas), but the call_tool proxy only accepts a name and will execute the highest version, so clients can select arguments from a stale schema and hit validation failures.
Useful? React with 👍 / 👎.
| if name == self._search_tool_name: | ||
| return self._make_search_tool() | ||
| if name == self._call_tool_name: | ||
| return self._make_call_tool() |
There was a problem hiding this comment.
get_tool always intercepts search_tool_name and call_tool_name, so any existing real tool with either of those names (or with colliding custom names) becomes unreachable once this transform is enabled. Because the proxy also blocks calls to synthetic names, the collision silently removes valid tools from use instead of failing fast, which is a breaking behavior for affected servers.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f82a1fd43e
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
|
||
| Use this to execute tools discovered via search_tools. | ||
| """ | ||
| if name in {transform._call_tool_name, transform._search_tool_name}: |
There was a problem hiding this comment.
The self-call guard only rejects transform._call_tool_name and transform._search_tool_name, so it misses aliases introduced by downstream name-rewriting transforms. For example, when Namespace is added after this transform, clients see ns_call_tool; calling ns_call_tool with name="ns_call_tool" bypasses this check and ctx.fastmcp.call_tool(...) resolves the same proxy again, causing unbounded recursion/timeouts. Fresh evidence: Namespace.list_tools() prefixes every tool name, including synthetic ones, so transformed aliases are reachable in normal configuration.
Useful? React with 👍 / 👎.
|
|
||
| def _catalog_hash(tools: Sequence[Tool]) -> str: | ||
| """SHA256 hash of sorted tool searchable text for staleness detection.""" | ||
| key = "|".join(sorted(_extract_searchable_text(t) for t in tools)) |
There was a problem hiding this comment.
The BM25 rebuild key hashes only _extract_searchable_text, which includes names/descriptions and parameter descriptions but omits other schema fields (for example parameter types, required fields, defaults, and output schema). When a dynamic provider updates those non-text fields without changing searchable text, _search() reuses stale _indexed_tools objects and returns outdated tool schemas, so clients can generate arguments that no longer validate against the current tool. Fresh evidence: _catalog_hash now uses searchable text (not just names), but still does not fingerprint the full tool schema.
Useful? React with 👍 / 👎.
jlowin
removed
the
DON'T MERGE
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: fb6236872f
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| raise ValueError( | ||
| f"'{name}' is a synthetic search tool and cannot be called via the call_tool proxy" | ||
| ) | ||
| return await ctx.fastmcp.call_tool(name, arguments) |
There was a problem hiding this comment.
call_tool
The proxy always forwards discovered tools with ctx.fastmcp.call_tool(name, arguments) and never propagates task execution metadata, so any discovered tool configured with task mode required will fail at runtime (the task router rejects synchronous calls without task metadata). This only appears when servers use search transform with task-augmented tools, but in that setup those tools become effectively unusable through the documented search_tools → call_tool flow.
Useful? React with 👍 / 👎.
RegexSearchTransform and BM25SearchTransform collapse large tool catalogs into a search interface so LLMs discover tools on demand instead of receiving the full listing.
Transforms that replace list_tools() with synthetic components (like search) need to read the real catalog at call time without triggering their own replacement logic. CatalogTransform handles the re-entrant bypass via per-instance ContextVar, exposing transform_tools() as the subclass hook and get_tool_catalog() for catalog access.
jlowin
force-pushed
the
feat/tool-search-transforms
branch
from
7beb80e to
727e430
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 727e430660
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
2 participants
This may hold for 3.1.
When a server exposes hundreds of tools, sending the full catalog to an LLM wastes tokens and hurts selection accuracy. Search transforms solve this by replacing
list_tools()with a search interface — the LLM discovers tools on demand instead of seeing everything upfront.Two strategies, both zero-dependency:
RegexSearchTransformfor pattern matching andBM25SearchTransformfor natural-language relevance ranking. Adding one collapses the entire catalog into two synthetic tools:Search results respect the full auth pipeline — middleware, visibility transforms, session-level
disable_components, and component auth checks all filter what's discoverable. The search tool querieslist_tools()through the complete pipeline at search time using a contextvar bypass that only skips its own hiding behavior.