From 0091d9f0714763eaad3c8450e5eadfd7555cba11 Mon Sep 17 00:00:00 2001 From: lhd <1037135398@qq.com> Date: Tue, 17 Mar 2026 20:43:55 +0800 Subject: [PATCH] feat(tools): add tool_search for deferred MCP tool loading (#1176) * feat(tools): add tool_search for deferred MCP tool loading When multiple MCP servers are enabled, total tool count can exceed 30-50, causing context bloat and degraded tool selection accuracy. This adds a deferred tool loading mechanism controlled by `tool_search.enabled` config. - Add ToolSearchConfig with single `enabled` field - Add DeferredToolRegistry with regex search (select:, +keyword, keyword) - Add tool_search tool returning OpenAI-compatible function JSON - Add DeferredToolFilterMiddleware to hide deferred schemas from bind_tools - Add section to system prompt - Enable MCP tool_name_prefix to prevent cross-server name collisions - Add 34 unit tests covering registry, tool, prompt, and middleware * fix: reset stale deferred registry and bump config_version - Reset deferred registry upfront in get_available_tools() to prevent stale tool entries when MCP servers are disabled between calls - Bump config_version to 2 for new tool_search config field Co-Authored-By: Claude Opus 4.6 * fix(tests): mock get_app_config in prompt section tests for CI CI has no config.yaml, causing TestDeferredToolsPromptSection to fail with FileNotFoundError. Add autouse fixture to mock get_app_config. Co-Authored-By: Claude Opus 4.6 --------- Co-authored-by: Claude Opus 4.6 --- .../deerflow/agents/lead_agent/agent.py | 9 +- .../deerflow/agents/lead_agent/prompt.py | 31 ++ .../deferred_tool_filter_middleware.py | 60 +++ .../harness/deerflow/config/app_config.py | 6 + .../deerflow/config/tool_search_config.py | 35 ++ .../packages/harness/deerflow/mcp/tools.py | 2 +- .../deerflow/tools/builtins/tool_search.py | 168 ++++++++ .../packages/harness/deerflow/tools/tools.py | 59 ++- backend/tests/test_tool_search.py | 363 ++++++++++++++++++ config.example.yaml | 14 +- 10 files changed, 721 insertions(+), 26 deletions(-) create mode 100644 backend/packages/harness/deerflow/agents/middlewares/deferred_tool_filter_middleware.py create mode 100644 backend/packages/harness/deerflow/config/tool_search_config.py create mode 100644 backend/packages/harness/deerflow/tools/builtins/tool_search.py create mode 100644 backend/tests/test_tool_search.py diff --git a/backend/packages/harness/deerflow/agents/lead_agent/agent.py b/backend/packages/harness/deerflow/agents/lead_agent/agent.py index 2fa588c..122011c 100644 --- a/backend/packages/harness/deerflow/agents/lead_agent/agent.py +++ b/backend/packages/harness/deerflow/agents/lead_agent/agent.py @@ -240,6 +240,11 @@ def _build_middlewares(config: RunnableConfig, model_name: str | None, agent_nam if model_config is not None and model_config.supports_vision: middlewares.append(ViewImageMiddleware()) + # Add DeferredToolFilterMiddleware to hide deferred tool schemas from model binding + if app_config.tool_search.enabled: + from deerflow.agents.middlewares.deferred_tool_filter_middleware import DeferredToolFilterMiddleware + middlewares.append(DeferredToolFilterMiddleware()) + # Add SubagentLimitMiddleware to truncate excess parallel task calls subagent_enabled = config.get("configurable", {}).get("subagent_enabled", False) if subagent_enabled: @@ -314,13 +319,11 @@ def make_lead_agent(config: RunnableConfig): if is_bootstrap: # Special bootstrap agent with minimal prompt for initial custom agent creation flow - system_prompt = apply_prompt_template(subagent_enabled=subagent_enabled, max_concurrent_subagents=max_concurrent_subagents, available_skills=set(["bootstrap"])) - return create_agent( model=create_chat_model(name=model_name, thinking_enabled=thinking_enabled), tools=get_available_tools(model_name=model_name, subagent_enabled=subagent_enabled) + [setup_agent], middleware=_build_middlewares(config, model_name=model_name), - system_prompt=system_prompt, + system_prompt=apply_prompt_template(subagent_enabled=subagent_enabled, max_concurrent_subagents=max_concurrent_subagents, available_skills=set(["bootstrap"])), state_schema=ThreadState, ) diff --git a/backend/packages/harness/deerflow/agents/lead_agent/prompt.py b/backend/packages/harness/deerflow/agents/lead_agent/prompt.py index 3417bbf..47630af 100644 --- a/backend/packages/harness/deerflow/agents/lead_agent/prompt.py +++ b/backend/packages/harness/deerflow/agents/lead_agent/prompt.py @@ -235,6 +235,8 @@ You: "Deploying to staging..." [proceed] {skills_section} +{deferred_tools_section} + {subagent_section} @@ -417,6 +419,31 @@ def get_agent_soul(agent_name: str | None) -> str: return "" +def get_deferred_tools_prompt_section() -> str: + """Generate block for the system prompt. + + Lists only deferred tool names so the agent knows what exists + and can use tool_search to load them. + Returns empty string when tool_search is disabled or no tools are deferred. + """ + from deerflow.tools.builtins.tool_search import get_deferred_registry + + try: + from deerflow.config import get_app_config + + if not get_app_config().tool_search.enabled: + return "" + except FileNotFoundError: + return "" + + registry = get_deferred_registry() + if not registry: + return "" + + names = "\n".join(e.name for e in registry.entries) + return f"\n{names}\n" + + def apply_prompt_template(subagent_enabled: bool = False, max_concurrent_subagents: int = 3, *, agent_name: str | None = None, available_skills: set[str] | None = None) -> str: # Get memory context memory_context = _get_memory_context(agent_name) @@ -446,11 +473,15 @@ def apply_prompt_template(subagent_enabled: bool = False, max_concurrent_subagen # Get skills section skills_section = get_skills_prompt_section(available_skills) + # Get deferred tools section (tool_search) + deferred_tools_section = get_deferred_tools_prompt_section() + # Format the prompt with dynamic skills and memory prompt = SYSTEM_PROMPT_TEMPLATE.format( agent_name=agent_name or "DeerFlow 2.0", soul=get_agent_soul(agent_name), skills_section=skills_section, + deferred_tools_section=deferred_tools_section, memory_context=memory_context, subagent_section=subagent_section, subagent_reminder=subagent_reminder, diff --git a/backend/packages/harness/deerflow/agents/middlewares/deferred_tool_filter_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/deferred_tool_filter_middleware.py new file mode 100644 index 0000000..604cdf3 --- /dev/null +++ b/backend/packages/harness/deerflow/agents/middlewares/deferred_tool_filter_middleware.py @@ -0,0 +1,60 @@ +"""Middleware to filter deferred tool schemas from model binding. + +When tool_search is enabled, MCP tools are registered in the DeferredToolRegistry +and passed to ToolNode for execution, but their schemas should NOT be sent to the +LLM via bind_tools (that's the whole point of deferral — saving context tokens). + +This middleware intercepts wrap_model_call and removes deferred tools from +request.tools so that model.bind_tools only receives active tool schemas. +The agent discovers deferred tools at runtime via the tool_search tool. +""" + +import logging +from collections.abc import Awaitable, Callable +from typing import override + +from langchain.agents import AgentState +from langchain.agents.middleware import AgentMiddleware +from langchain.agents.middleware.types import ModelCallResult, ModelRequest, ModelResponse + +logger = logging.getLogger(__name__) + + +class DeferredToolFilterMiddleware(AgentMiddleware[AgentState]): + """Remove deferred tools from request.tools before model binding. + + ToolNode still holds all tools (including deferred) for execution routing, + but the LLM only sees active tool schemas — deferred tools are discoverable + via tool_search at runtime. + """ + + def _filter_tools(self, request: ModelRequest) -> ModelRequest: + from deerflow.tools.builtins.tool_search import get_deferred_registry + + registry = get_deferred_registry() + if not registry: + return request + + deferred_names = {e.name for e in registry.entries} + active_tools = [t for t in request.tools if getattr(t, "name", None) not in deferred_names] + + if len(active_tools) < len(request.tools): + logger.debug(f"Filtered {len(request.tools) - len(active_tools)} deferred tool schema(s) from model binding") + + return request.override(tools=active_tools) + + @override + def wrap_model_call( + self, + request: ModelRequest, + handler: Callable[[ModelRequest], ModelResponse], + ) -> ModelCallResult: + return handler(self._filter_tools(request)) + + @override + async def awrap_model_call( + self, + request: ModelRequest, + handler: Callable[[ModelRequest], Awaitable[ModelResponse]], + ) -> ModelCallResult: + return await handler(self._filter_tools(request)) diff --git a/backend/packages/harness/deerflow/config/app_config.py b/backend/packages/harness/deerflow/config/app_config.py index 80ef5cc..0f73c35 100644 --- a/backend/packages/harness/deerflow/config/app_config.py +++ b/backend/packages/harness/deerflow/config/app_config.py @@ -17,6 +17,7 @@ from deerflow.config.subagents_config import load_subagents_config_from_dict from deerflow.config.summarization_config import load_summarization_config_from_dict from deerflow.config.title_config import load_title_config_from_dict from deerflow.config.tool_config import ToolConfig, ToolGroupConfig +from deerflow.config.tool_search_config import ToolSearchConfig, load_tool_search_config_from_dict load_dotenv() @@ -32,6 +33,7 @@ class AppConfig(BaseModel): tool_groups: list[ToolGroupConfig] = Field(default_factory=list, description="Available tool groups") skills: SkillsConfig = Field(default_factory=SkillsConfig, description="Skills configuration") extensions: ExtensionsConfig = Field(default_factory=ExtensionsConfig, description="Extensions configuration (MCP servers and skills state)") + tool_search: ToolSearchConfig = Field(default_factory=ToolSearchConfig, description="Tool search / deferred loading configuration") model_config = ConfigDict(extra="allow", frozen=False) checkpointer: CheckpointerConfig | None = Field(default=None, description="Checkpointer configuration") @@ -101,6 +103,10 @@ class AppConfig(BaseModel): if "subagents" in config_data: load_subagents_config_from_dict(config_data["subagents"]) + # Load tool_search config if present + if "tool_search" in config_data: + load_tool_search_config_from_dict(config_data["tool_search"]) + # Load checkpointer config if present if "checkpointer" in config_data: load_checkpointer_config_from_dict(config_data["checkpointer"]) diff --git a/backend/packages/harness/deerflow/config/tool_search_config.py b/backend/packages/harness/deerflow/config/tool_search_config.py new file mode 100644 index 0000000..cdeddab --- /dev/null +++ b/backend/packages/harness/deerflow/config/tool_search_config.py @@ -0,0 +1,35 @@ +"""Configuration for deferred tool loading via tool_search.""" + +from pydantic import BaseModel, Field + + +class ToolSearchConfig(BaseModel): + """Configuration for deferred tool loading via tool_search. + + When enabled, MCP tools are not loaded into the agent's context directly. + Instead, they are listed by name in the system prompt and discoverable + via the tool_search tool at runtime. + """ + + enabled: bool = Field( + default=False, + description="Defer tools and enable tool_search", + ) + + +_tool_search_config: ToolSearchConfig | None = None + + +def get_tool_search_config() -> ToolSearchConfig: + """Get the tool search config, loading from AppConfig if needed.""" + global _tool_search_config + if _tool_search_config is None: + _tool_search_config = ToolSearchConfig() + return _tool_search_config + + +def load_tool_search_config_from_dict(data: dict) -> ToolSearchConfig: + """Load tool search config from a dict (called during AppConfig loading).""" + global _tool_search_config + _tool_search_config = ToolSearchConfig.model_validate(data) + return _tool_search_config diff --git a/backend/packages/harness/deerflow/mcp/tools.py b/backend/packages/harness/deerflow/mcp/tools.py index 78fea42..02d7cb1 100644 --- a/backend/packages/harness/deerflow/mcp/tools.py +++ b/backend/packages/harness/deerflow/mcp/tools.py @@ -53,7 +53,7 @@ async def get_mcp_tools() -> list[BaseTool]: if oauth_interceptor is not None: tool_interceptors.append(oauth_interceptor) - client = MultiServerMCPClient(servers_config, tool_interceptors=tool_interceptors) + client = MultiServerMCPClient(servers_config, tool_interceptors=tool_interceptors, tool_name_prefix=True) # Get all tools from all servers tools = await client.get_tools() diff --git a/backend/packages/harness/deerflow/tools/builtins/tool_search.py b/backend/packages/harness/deerflow/tools/builtins/tool_search.py new file mode 100644 index 0000000..c20f799 --- /dev/null +++ b/backend/packages/harness/deerflow/tools/builtins/tool_search.py @@ -0,0 +1,168 @@ +"""Tool search — deferred tool discovery at runtime. + +Contains: +- DeferredToolRegistry: stores deferred tools and handles regex search +- tool_search: the LangChain tool the agent calls to discover deferred tools + +The agent sees deferred tool names in but cannot +call them until it fetches their full schema via the tool_search tool. +Source-agnostic: no mention of MCP or tool origin. +""" + +import json +import logging +import re +from dataclasses import dataclass + +from langchain.tools import BaseTool +from langchain_core.tools import tool +from langchain_core.utils.function_calling import convert_to_openai_function + +logger = logging.getLogger(__name__) + +MAX_RESULTS = 5 # Max tools returned per search + + +# ── Registry ── + + +@dataclass +class DeferredToolEntry: + """Lightweight metadata for a deferred tool (no full schema in context).""" + + name: str + description: str + tool: BaseTool # Full tool object, returned only on search match + + +class DeferredToolRegistry: + """Registry of deferred tools, searchable by regex pattern.""" + + def __init__(self): + self._entries: list[DeferredToolEntry] = [] + + def register(self, tool: BaseTool) -> None: + self._entries.append( + DeferredToolEntry( + name=tool.name, + description=tool.description or "", + tool=tool, + ) + ) + + def search(self, query: str) -> list[BaseTool]: + """Search deferred tools by regex pattern against name + description. + + Supports three query forms (aligned with Claude Code): + - "select:name1,name2" — exact name match + - "+keyword rest" — name must contain keyword, rank by rest + - "keyword query" — regex match against name + description + + Returns: + List of matched BaseTool objects (up to MAX_RESULTS). + """ + if query.startswith("select:"): + names = {n.strip() for n in query[7:].split(",")} + return [e.tool for e in self._entries if e.name in names][:MAX_RESULTS] + + if query.startswith("+"): + parts = query[1:].split(None, 1) + required = parts[0].lower() + candidates = [e for e in self._entries if required in e.name.lower()] + if len(parts) > 1: + candidates.sort( + key=lambda e: _regex_score(parts[1], e), + reverse=True, + ) + return [e.tool for e in candidates][:MAX_RESULTS] + + # General regex search + try: + regex = re.compile(query, re.IGNORECASE) + except re.error: + regex = re.compile(re.escape(query), re.IGNORECASE) + + scored = [] + for entry in self._entries: + searchable = f"{entry.name} {entry.description}" + if regex.search(searchable): + score = 2 if regex.search(entry.name) else 1 + scored.append((score, entry)) + + scored.sort(key=lambda x: x[0], reverse=True) + return [entry.tool for _, entry in scored][:MAX_RESULTS] + + @property + def entries(self) -> list[DeferredToolEntry]: + return list(self._entries) + + def __len__(self) -> int: + return len(self._entries) + + +def _regex_score(pattern: str, entry: DeferredToolEntry) -> int: + try: + regex = re.compile(pattern, re.IGNORECASE) + except re.error: + regex = re.compile(re.escape(pattern), re.IGNORECASE) + return len(regex.findall(f"{entry.name} {entry.description}")) + + +# ── Singleton ── + +_registry: DeferredToolRegistry | None = None + + +def get_deferred_registry() -> DeferredToolRegistry | None: + return _registry + + +def set_deferred_registry(registry: DeferredToolRegistry) -> None: + global _registry + _registry = registry + + +def reset_deferred_registry() -> None: + """Reset the deferred registry singleton. Useful for testing.""" + global _registry + _registry = None + + +# ── Tool ── + + +@tool +def tool_search(query: str) -> str: + """Fetches full schema definitions for deferred tools so they can be called. + + Deferred tools appear by name in in the system + prompt. Until fetched, only the name is known — there is no parameter + schema, so the tool cannot be invoked. This tool takes a query, matches + it against the deferred tool list, and returns the matched tools' complete + definitions. Once a tool's schema appears in that result, it is callable. + + Query forms: + - "select:Read,Edit,Grep" — fetch these exact tools by name + - "notebook jupyter" — keyword search, up to max_results best matches + - "+slack send" — require "slack" in the name, rank by remaining terms + + Args: + query: Query to find deferred tools. Use "select:" for + direct selection, or keywords to search. + + Returns: + Matched tool definitions as JSON array. + """ + registry = get_deferred_registry() + if registry is None: + return "No deferred tools available." + + matched_tools = registry.search(query) + if not matched_tools: + return f"No tools found matching: {query}" + + # Use LangChain's built-in serialization to produce OpenAI function format. + # This is model-agnostic: all LLMs understand this standard schema. + tool_defs = [convert_to_openai_function(t) for t in matched_tools[:MAX_RESULTS]] + + return json.dumps(tool_defs, indent=2, ensure_ascii=False) diff --git a/backend/packages/harness/deerflow/tools/tools.py b/backend/packages/harness/deerflow/tools/tools.py index 5e7b560..73a3c64 100644 --- a/backend/packages/harness/deerflow/tools/tools.py +++ b/backend/packages/harness/deerflow/tools/tools.py @@ -5,6 +5,7 @@ from langchain.tools import BaseTool from deerflow.config import get_app_config from deerflow.reflection import resolve_variable from deerflow.tools.builtins import ask_clarification_tool, present_file_tool, task_tool, view_image_tool +from deerflow.tools.builtins.tool_search import reset_deferred_registry logger = logging.getLogger(__name__) @@ -42,27 +43,6 @@ def get_available_tools( config = get_app_config() loaded_tools = [resolve_variable(tool.use, BaseTool) for tool in config.tools if groups is None or tool.group in groups] - # Get cached MCP tools if enabled - # NOTE: We use ExtensionsConfig.from_file() instead of config.extensions - # to always read the latest configuration from disk. This ensures that changes - # made through the Gateway API (which runs in a separate process) are immediately - # reflected when loading MCP tools. - mcp_tools = [] - if include_mcp: - try: - from deerflow.config.extensions_config import ExtensionsConfig - from deerflow.mcp.cache import get_cached_mcp_tools - - extensions_config = ExtensionsConfig.from_file() - if extensions_config.get_enabled_mcp_servers(): - mcp_tools = get_cached_mcp_tools() - if mcp_tools: - logger.info(f"Using {len(mcp_tools)} cached MCP tool(s)") - except ImportError: - logger.warning("MCP module not available. Install 'langchain-mcp-adapters' package to enable MCP tools.") - except Exception as e: - logger.error(f"Failed to get cached MCP tools: {e}") - # Conditionally add tools based on config builtin_tools = BUILTIN_TOOLS.copy() @@ -81,4 +61,41 @@ def get_available_tools( builtin_tools.append(view_image_tool) logger.info(f"Including view_image_tool for model '{model_name}' (supports_vision=True)") + # Get cached MCP tools if enabled + # NOTE: We use ExtensionsConfig.from_file() instead of config.extensions + # to always read the latest configuration from disk. This ensures that changes + # made through the Gateway API (which runs in a separate process) are immediately + # reflected when loading MCP tools. + mcp_tools = [] + # Reset deferred registry upfront to prevent stale state from previous calls + reset_deferred_registry() + if include_mcp: + try: + from deerflow.config.extensions_config import ExtensionsConfig + from deerflow.mcp.cache import get_cached_mcp_tools + + extensions_config = ExtensionsConfig.from_file() + if extensions_config.get_enabled_mcp_servers(): + mcp_tools = get_cached_mcp_tools() + if mcp_tools: + logger.info(f"Using {len(mcp_tools)} cached MCP tool(s)") + + # When tool_search is enabled, register MCP tools in the + # deferred registry and add tool_search to builtin tools. + if config.tool_search.enabled: + from deerflow.tools.builtins.tool_search import DeferredToolRegistry, set_deferred_registry + from deerflow.tools.builtins.tool_search import tool_search as tool_search_tool + + registry = DeferredToolRegistry() + for t in mcp_tools: + registry.register(t) + set_deferred_registry(registry) + builtin_tools.append(tool_search_tool) + logger.info(f"Tool search active: {len(mcp_tools)} tools deferred") + except ImportError: + logger.warning("MCP module not available. Install 'langchain-mcp-adapters' package to enable MCP tools.") + except Exception as e: + logger.error(f"Failed to get cached MCP tools: {e}") + + logger.info(f"Total tools loaded: {len(loaded_tools)}, built-in tools: {len(builtin_tools)}, MCP tools: {len(mcp_tools)}") return loaded_tools + builtin_tools + mcp_tools diff --git a/backend/tests/test_tool_search.py b/backend/tests/test_tool_search.py new file mode 100644 index 0000000..b813bc7 --- /dev/null +++ b/backend/tests/test_tool_search.py @@ -0,0 +1,363 @@ +"""Tests for the tool_search (deferred tool loading) feature.""" + +import json +import sys + +import pytest +from langchain_core.tools import tool as langchain_tool + +from deerflow.config.tool_search_config import ToolSearchConfig, load_tool_search_config_from_dict +from deerflow.tools.builtins.tool_search import ( + DeferredToolRegistry, + get_deferred_registry, + reset_deferred_registry, + set_deferred_registry, +) + +# ── Fixtures ── + + +def _make_mock_tool(name: str, description: str): + """Create a minimal LangChain tool for testing.""" + + @langchain_tool(name) + def mock_tool(arg: str) -> str: + """Mock tool.""" + return f"{name}: {arg}" + + mock_tool.description = description + return mock_tool + + +@pytest.fixture +def registry(): + """Create a fresh DeferredToolRegistry with test tools.""" + reg = DeferredToolRegistry() + reg.register(_make_mock_tool("github_create_issue", "Create a new issue in a GitHub repository")) + reg.register(_make_mock_tool("github_list_repos", "List repositories for a GitHub user")) + reg.register(_make_mock_tool("slack_send_message", "Send a message to a Slack channel")) + reg.register(_make_mock_tool("slack_list_channels", "List available Slack channels")) + reg.register(_make_mock_tool("sentry_list_issues", "List issues from Sentry error tracking")) + reg.register(_make_mock_tool("database_query", "Execute a SQL query against the database")) + return reg + + +@pytest.fixture(autouse=True) +def _reset_singleton(): + """Reset the module-level singleton before/after each test.""" + reset_deferred_registry() + yield + reset_deferred_registry() + + +# ── ToolSearchConfig Tests ── + + +class TestToolSearchConfig: + def test_default_disabled(self): + config = ToolSearchConfig() + assert config.enabled is False + + def test_enabled(self): + config = ToolSearchConfig(enabled=True) + assert config.enabled is True + + def test_load_from_dict(self): + config = load_tool_search_config_from_dict({"enabled": True}) + assert config.enabled is True + + def test_load_from_empty_dict(self): + config = load_tool_search_config_from_dict({}) + assert config.enabled is False + + +# ── DeferredToolRegistry Tests ── + + +class TestDeferredToolRegistry: + def test_register_and_len(self, registry): + assert len(registry) == 6 + + def test_entries(self, registry): + names = [e.name for e in registry.entries] + assert "github_create_issue" in names + assert "slack_send_message" in names + + def test_search_select_single(self, registry): + results = registry.search("select:github_create_issue") + assert len(results) == 1 + assert results[0].name == "github_create_issue" + + def test_search_select_multiple(self, registry): + results = registry.search("select:github_create_issue,slack_send_message") + names = {t.name for t in results} + assert names == {"github_create_issue", "slack_send_message"} + + def test_search_select_nonexistent(self, registry): + results = registry.search("select:nonexistent_tool") + assert results == [] + + def test_search_plus_keyword(self, registry): + results = registry.search("+github") + names = {t.name for t in results} + assert names == {"github_create_issue", "github_list_repos"} + + def test_search_plus_keyword_with_ranking(self, registry): + results = registry.search("+github issue") + assert len(results) == 2 + # "github_create_issue" should rank higher (has "issue" in name) + assert results[0].name == "github_create_issue" + + def test_search_regex_keyword(self, registry): + results = registry.search("slack") + names = {t.name for t in results} + assert "slack_send_message" in names + assert "slack_list_channels" in names + + def test_search_regex_description(self, registry): + results = registry.search("SQL") + assert len(results) == 1 + assert results[0].name == "database_query" + + def test_search_regex_case_insensitive(self, registry): + results = registry.search("GITHUB") + assert len(results) == 2 + + def test_search_invalid_regex_falls_back_to_literal(self, registry): + # "[" is invalid regex, should be escaped and used as literal + results = registry.search("[") + assert results == [] + + def test_search_name_match_ranks_higher(self, registry): + # "issue" appears in both github_create_issue (name) and sentry_list_issues (name+desc) + results = registry.search("issue") + names = [t.name for t in results] + # Both should be found (both have "issue" in name) + assert "github_create_issue" in names + assert "sentry_list_issues" in names + + def test_search_max_results(self): + reg = DeferredToolRegistry() + for i in range(10): + reg.register(_make_mock_tool(f"tool_{i}", f"Tool number {i}")) + results = reg.search("tool") + assert len(results) <= 5 # MAX_RESULTS = 5 + + def test_search_empty_registry(self): + reg = DeferredToolRegistry() + assert reg.search("anything") == [] + + def test_empty_registry_len(self): + reg = DeferredToolRegistry() + assert len(reg) == 0 + + +# ── Singleton Tests ── + + +class TestSingleton: + def test_default_none(self): + assert get_deferred_registry() is None + + def test_set_and_get(self, registry): + set_deferred_registry(registry) + assert get_deferred_registry() is registry + + def test_reset(self, registry): + set_deferred_registry(registry) + reset_deferred_registry() + assert get_deferred_registry() is None + + +# ── tool_search Tool Tests ── + + +class TestToolSearchTool: + def test_no_registry(self): + from deerflow.tools.builtins.tool_search import tool_search + + result = tool_search.invoke({"query": "github"}) + assert result == "No deferred tools available." + + def test_no_match(self, registry): + from deerflow.tools.builtins.tool_search import tool_search + + set_deferred_registry(registry) + result = tool_search.invoke({"query": "nonexistent_xyz_tool"}) + assert "No tools found matching" in result + + def test_returns_valid_json(self, registry): + from deerflow.tools.builtins.tool_search import tool_search + + set_deferred_registry(registry) + result = tool_search.invoke({"query": "select:github_create_issue"}) + parsed = json.loads(result) + assert isinstance(parsed, list) + assert len(parsed) == 1 + assert parsed[0]["name"] == "github_create_issue" + + def test_returns_openai_function_format(self, registry): + from deerflow.tools.builtins.tool_search import tool_search + + set_deferred_registry(registry) + result = tool_search.invoke({"query": "select:slack_send_message"}) + parsed = json.loads(result) + func_def = parsed[0] + # OpenAI function format should have these keys + assert "name" in func_def + assert "description" in func_def + assert "parameters" in func_def + + def test_keyword_search_returns_json(self, registry): + from deerflow.tools.builtins.tool_search import tool_search + + set_deferred_registry(registry) + result = tool_search.invoke({"query": "github"}) + parsed = json.loads(result) + assert len(parsed) == 2 + names = {d["name"] for d in parsed} + assert names == {"github_create_issue", "github_list_repos"} + + +# ── Prompt Section Tests ── + + +class TestDeferredToolsPromptSection: + @pytest.fixture(autouse=True) + def _mock_app_config(self, monkeypatch): + """Provide a minimal AppConfig mock so tests don't need config.yaml.""" + from unittest.mock import MagicMock + + from deerflow.config.tool_search_config import ToolSearchConfig + + mock_config = MagicMock() + mock_config.tool_search = ToolSearchConfig() # disabled by default + monkeypatch.setattr("deerflow.config.get_app_config", lambda: mock_config) + + def test_empty_when_disabled(self): + from deerflow.agents.lead_agent.prompt import get_deferred_tools_prompt_section + + # tool_search.enabled defaults to False + section = get_deferred_tools_prompt_section() + assert section == "" + + def test_empty_when_enabled_but_no_registry(self, monkeypatch): + from deerflow.agents.lead_agent.prompt import get_deferred_tools_prompt_section + from deerflow.config import get_app_config + + monkeypatch.setattr(get_app_config().tool_search, "enabled", True) + section = get_deferred_tools_prompt_section() + assert section == "" + + def test_empty_when_enabled_but_empty_registry(self, monkeypatch): + from deerflow.agents.lead_agent.prompt import get_deferred_tools_prompt_section + from deerflow.config import get_app_config + + monkeypatch.setattr(get_app_config().tool_search, "enabled", True) + set_deferred_registry(DeferredToolRegistry()) + section = get_deferred_tools_prompt_section() + assert section == "" + + def test_lists_tool_names(self, registry, monkeypatch): + from deerflow.agents.lead_agent.prompt import get_deferred_tools_prompt_section + from deerflow.config import get_app_config + + monkeypatch.setattr(get_app_config().tool_search, "enabled", True) + set_deferred_registry(registry) + section = get_deferred_tools_prompt_section() + assert "" in section + assert "" in section + assert "github_create_issue" in section + assert "slack_send_message" in section + assert "sentry_list_issues" in section + # Should only have names, no descriptions + assert "Create a new issue" not in section + + +# ── DeferredToolFilterMiddleware Tests ── + + +class TestDeferredToolFilterMiddleware: + @pytest.fixture(autouse=True) + def _ensure_middlewares_package(self): + """Remove mock entries injected by test_subagent_executor.py. + + That file replaces deerflow.agents and deerflow.agents.middlewares with + MagicMock objects in sys.modules (session-scoped) to break circular imports. + We must clear those mocks so real submodule imports work. + """ + from unittest.mock import MagicMock + + mock_keys = [ + "deerflow.agents", + "deerflow.agents.middlewares", + "deerflow.agents.middlewares.deferred_tool_filter_middleware", + ] + for key in mock_keys: + if isinstance(sys.modules.get(key), MagicMock): + del sys.modules[key] + + def test_filters_deferred_tools(self, registry): + from deerflow.agents.middlewares.deferred_tool_filter_middleware import DeferredToolFilterMiddleware + + set_deferred_registry(registry) + middleware = DeferredToolFilterMiddleware() + + # Build a mock tools list: 2 active + 1 deferred + active_tool = _make_mock_tool("my_active_tool", "An active tool") + deferred_tool = registry.entries[0].tool # github_create_issue + + class FakeRequest: + def __init__(self, tools): + self.tools = tools + + def override(self, **kwargs): + return FakeRequest(kwargs.get("tools", self.tools)) + + request = FakeRequest(tools=[active_tool, deferred_tool]) + filtered = middleware._filter_tools(request) + + assert len(filtered.tools) == 1 + assert filtered.tools[0].name == "my_active_tool" + + def test_no_op_when_no_registry(self): + from deerflow.agents.middlewares.deferred_tool_filter_middleware import DeferredToolFilterMiddleware + + middleware = DeferredToolFilterMiddleware() + active_tool = _make_mock_tool("my_tool", "A tool") + + class FakeRequest: + def __init__(self, tools): + self.tools = tools + + def override(self, **kwargs): + return FakeRequest(kwargs.get("tools", self.tools)) + + request = FakeRequest(tools=[active_tool]) + filtered = middleware._filter_tools(request) + + assert len(filtered.tools) == 1 + assert filtered.tools[0].name == "my_tool" + + def test_preserves_dict_tools(self, registry): + """Dict tools (provider built-ins) should not be filtered.""" + from deerflow.agents.middlewares.deferred_tool_filter_middleware import DeferredToolFilterMiddleware + + set_deferred_registry(registry) + middleware = DeferredToolFilterMiddleware() + + dict_tool = {"type": "function", "function": {"name": "some_builtin"}} + active_tool = _make_mock_tool("my_active_tool", "Active") + + class FakeRequest: + def __init__(self, tools): + self.tools = tools + + def override(self, **kwargs): + return FakeRequest(kwargs.get("tools", self.tools)) + + request = FakeRequest(tools=[dict_tool, active_tool]) + filtered = middleware._filter_tools(request) + + # dict_tool has no .name attr → getattr returns None → not in deferred_names → kept + assert len(filtered.tools) == 2 diff --git a/config.example.yaml b/config.example.yaml index 6dcbf33..afdbded 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -12,7 +12,7 @@ # ============================================================================ # Bump this number when the config schema changes. # Run `make config-upgrade` to merge new fields into your local config.yaml. -config_version: 1 +config_version: 2 # ============================================================================ # Models Configuration @@ -224,6 +224,18 @@ tools: group: bash use: deerflow.sandbox.tools:bash_tool +# ============================================================================ +# Tool Search Configuration (Deferred Tool Loading) +# ============================================================================ +# When enabled, MCP tools are not loaded into the agent's context directly. +# Instead, they are listed by name in the system prompt and discoverable +# via the tool_search tool at runtime. +# This reduces context usage and improves tool selection accuracy when +# multiple MCP servers expose a large number of tools. + +tool_search: + enabled: false + # ============================================================================ # Sandbox Configuration # ============================================================================