refactor(tool-search): consolidate MCP metadata tag and harden deferred-tool setup (#3370)

Follow-up to #3342 (deferred MCP tool loading). Maintainability cleanup plus
hardening of malformed/empty tool_search queries; no change to the deferral
mechanism or search ranking.

- Add deerflow/tools/mcp_metadata.py as the single source of truth for the
  "deerflow_mcp" tag (MCP_TOOL_METADATA_KEY + tag_mcp_tool + public
  is_mcp_tool). Removes the duplicated magic string and the private,
  cross-module _is_mcp_tool import.
- tool_search.search: never raise on model-generated input. Extract
  _compile_catalog_regex (shared compile-with-literal-fallback); return empty
  for empty/whitespace queries and a bare "+" instead of matching everything
  or raising IndexError.
- DeferredToolSetup: document the empty-vs-populated invariant.
- build_deferred_tool_setup: comment the two distinct empty-return branches.
- _assemble_deferred: add return type, rename local to deferred_setup, build
  the final list with an explicit append.
- Tests: use tag_mcp_tool instead of per-file tag helpers; cover empty and
  bare-"+" queries.
This commit is contained in:
AochenShen99
2026-06-05 15:21:41 +08:00
committed by GitHub
parent 28b1da2172
commit 2bbc7879fa
8 changed files with 123 additions and 47 deletions
@@ -28,11 +28,25 @@ from langchain_core.tools import InjectedToolCallId, tool
from langchain_core.utils.function_calling import convert_to_openai_function
from langgraph.types import Command
from deerflow.tools.mcp_metadata import is_mcp_tool
logger = logging.getLogger(__name__)
MAX_RESULTS = 5 # Max tools returned per search
def _compile_catalog_regex(pattern: str) -> re.Pattern[str]:
"""Compile ``pattern`` case-insensitively, falling back to a literal match.
Search queries come from the model, so an invalid regex (e.g. an unbalanced
paren) must degrade to a literal substring match rather than raise.
"""
try:
return re.compile(pattern, re.IGNORECASE)
except re.error:
return re.compile(re.escape(pattern), re.IGNORECASE)
# ── Catalog ──
@@ -56,22 +70,25 @@ class DeferredToolCatalog:
return hashlib.sha256(blob.encode("utf-8")).hexdigest()[:16]
def search(self, query: str) -> list[BaseTool]:
query = query.strip()
if not query:
return []
if query.startswith("select:"):
wanted = {n.strip() for n in query[7:].split(",")}
return [t for t in self.tools if t.name in wanted][:MAX_RESULTS]
if query.startswith("+"):
parts = query[1:].split(None, 1)
if not parts:
return [] # bare "+" with no required token — nothing to require
required = parts[0].lower()
candidates = [t for t in self.tools if required in t.name.lower()]
if len(parts) > 1:
candidates.sort(key=lambda t: _catalog_regex_score(parts[1], t), reverse=True)
return candidates[:MAX_RESULTS]
try:
regex = re.compile(query, re.IGNORECASE)
except re.error:
regex = re.compile(re.escape(query), re.IGNORECASE)
regex = _compile_catalog_regex(query)
scored: list[tuple[int, BaseTool]] = []
for t in self.tools:
searchable = f"{t.name} {t.description or ''}"
@@ -82,10 +99,7 @@ class DeferredToolCatalog:
def _catalog_regex_score(pattern: str, t: BaseTool) -> int:
try:
regex = re.compile(pattern, re.IGNORECASE)
except re.error:
regex = re.compile(re.escape(pattern), re.IGNORECASE)
regex = _compile_catalog_regex(pattern)
return len(regex.findall(f"{t.name} {t.description or ''}"))
@@ -94,15 +108,25 @@ def _catalog_regex_score(pattern: str, t: BaseTool) -> int:
@dataclass(frozen=True)
class DeferredToolSetup:
"""Result of assembling deferred-tool support for one agent build.
The three fields move as a unit, so callers branch on ``tool_search_tool``:
- **Empty** ``(None, frozenset(), None)``: deferral is disabled, or no MCP
tool survived policy filtering. Nothing is deferred — bind tools as-is.
- **Populated**: ``tool_search_tool`` is appended to the agent's tools,
``deferred_names`` are withheld from the model until promoted, and
``catalog_hash`` scopes those promotions in graph state.
Invariant: ``tool_search_tool is None`` ⟺ ``deferred_names`` is empty ⟺
``catalog_hash is None``.
"""
tool_search_tool: BaseTool | None
deferred_names: frozenset[str]
catalog_hash: str | None
def _is_mcp_tool(t: BaseTool) -> bool:
return (getattr(t, "metadata", None) or {}).get("deerflow_mcp") is True
def build_tool_search_tool(catalog: DeferredToolCatalog) -> BaseTool:
catalog_hash = catalog.hash
@@ -141,11 +165,17 @@ def build_deferred_tool_setup(filtered_tools: list[BaseTool], *, enabled: bool)
Must be called after skill/agent tool-policy filtering so the catalog never
exposes a tool the current agent is not allowed to use.
Returns an empty setup (see :class:`DeferredToolSetup`) in two distinct
cases: deferral is disabled, or it is enabled but no MCP tool survived
filtering.
"""
if not enabled:
# Deferral disabled: defer nothing; the model binds every tool as before.
return DeferredToolSetup(None, frozenset(), None)
deferred = [t for t in filtered_tools if _is_mcp_tool(t)]
deferred = [t for t in filtered_tools if is_mcp_tool(t)]
if not deferred:
# Enabled, but no MCP tool to defer: same empty result, different reason.
return DeferredToolSetup(None, frozenset(), None)
catalog = DeferredToolCatalog(tuple(deferred))
return DeferredToolSetup(build_tool_search_tool(catalog), catalog.names, catalog.hash)