From ebc09ce130fcea64cd2aebb459a614346ff2630a Mon Sep 17 00:00:00 2001 From: Ryker_Feng <90562015+18062706139fcz@users.noreply.github.com> Date: Fri, 10 Jul 2026 07:54:36 +0800 Subject: [PATCH] feat(mcp): auto-promote deferred MCP tools from routing hints (#4019) * feat(mcp): auto-promote deferred MCP tools from routing hints When tool_search.enabled=true defers MCP tool schemas, PR1 routing hints still require the model to spend a tool_search discovery round trip before it can call the tool the routing metadata already points at. This adds a McpRoutingMiddleware that matches the latest user message against PR1 routing keywords and promotes the matching deferred schemas before the model call, removing that round trip. Design (soft routing, opt-in, additive): - Matches only the latest real HumanMessage (shared is_real_user_message helper, reused by SkillActivationMiddleware so the two cannot drift); case-insensitive substring match, no tokenizer dependency. - Ordering: priority desc, then tool name asc; capped by the new global tool_search.auto_promote_top_k (default 3, clamped 1..5). Does not add or consume a per-tool auto_promote_top_k (PR1 schema unchanged); a per-tool value is ignored with a DEBUG note. - Returns a plain {"promoted": ...} state update (not a Command) and relies on ThreadState.merge_promoted for union/dedupe, so auto-promote and a model-triggered tool_search converge on the same catalog hash. - Installed before DeferredToolFilterMiddleware on every deferred-tool path (lead agent, subagent, embedded client, webhook via shared builders); a construction-time assert rejects the reversed order. catalog_hash is None / no routing index is a complete no-op, so bootstrap and ACP skip it. - Privacy: never executes tools, never promotes policy-filtered tools, adds no routing keywords or matched tool names to trace metadata or INFO/WARN logs. No behavior change when tool_search.enabled=false. Tests: index construction, matching semantics, middleware state updates, same-cycle deferred-filter interaction, lead/subagent/embedded-client builder wiring + order invariant, config clamping, config.example.yaml parseability, and privacy assertions. * refactor(mcp): address auto-promote review nits - executor: access app_config.tool_search.auto_promote_top_k directly to match the lead-agent and embedded-client paths (drop the over-defensive getattr that masked missing config); update the subagent test mock to carry tool_search. - tool_search / mcp_routing_middleware: cross-reference the duplicated routing priority/keyword normalization between the builder and the middleware's defensive _normalize_index so they cannot silently drift. - MCP_SERVER.md: document that auto-promote keyword matching is a case-insensitive substring test (not word-boundary), advising distinctive keywords. --- README.md | 2 +- backend/AGENTS.md | 22 +- backend/README.md | 8 +- backend/docs/MCP_SERVER.md | 15 +- .../deerflow/agents/lead_agent/agent.py | 25 +- .../middlewares/mcp_routing_middleware.py | 137 ++++++++ .../skill_activation_middleware.py | 11 +- .../tool_error_handling_middleware.py | 7 + backend/packages/harness/deerflow/client.py | 8 +- .../deerflow/config/tool_search_config.py | 19 +- .../harness/deerflow/subagents/executor.py | 26 +- .../deerflow/tools/builtins/tool_search.py | 65 +++- .../harness/deerflow/utils/messages.py | 18 ++ backend/tests/test_client.py | 62 ++++ .../tests/test_lead_agent_model_resolution.py | 27 ++ .../tests/test_mcp_routing_auto_promote.py | 296 ++++++++++++++++++ backend/tests/test_subagent_executor.py | 2 +- .../test_tool_error_handling_middleware.py | 17 + backend/tests/test_tool_search.py | 47 ++- config.example.yaml | 4 + 20 files changed, 778 insertions(+), 40 deletions(-) create mode 100644 backend/packages/harness/deerflow/agents/middlewares/mcp_routing_middleware.py create mode 100644 backend/tests/test_mcp_routing_auto_promote.py diff --git a/README.md b/README.md index d13cdea87..89b076757 100644 --- a/README.md +++ b/README.md @@ -354,7 +354,7 @@ See the [Sandbox Configuration Guide](backend/docs/CONFIGURATION.md#sandbox) to DeerFlow supports configurable MCP servers and skills to extend its capabilities. For HTTP/SSE MCP servers, OAuth token flows are supported (`client_credentials`, `refresh_token`). For stdio MCP servers, per-tool call timeouts can be configured with `tool_call_timeout`. -MCP routing hints can also prefer a specific MCP tool for matching requests without changing tool binding; when `tool_search` defers MCP schemas, the hint directs the agent to fetch the tool first. +MCP routing hints can also prefer a specific MCP tool for matching requests without forbidding other tools. When `tool_search` defers MCP schemas, matching routing metadata can auto-promote up to `tool_search.auto_promote_top_k` deferred schemas before the model call. See the [MCP Server Guide](backend/docs/MCP_SERVER.md) for detailed instructions. #### IM Channels diff --git a/backend/AGENTS.md b/backend/AGENTS.md index e33760434..a9142d43f 100644 --- a/backend/AGENTS.md +++ b/backend/AGENTS.md @@ -243,14 +243,15 @@ Lead-agent middlewares are assembled in strict order across three functions: the 20. **TitleMiddleware** - Auto-generates the thread title after the first complete exchange and normalizes structured message content before prompting the title model. If a first-turn run is interrupted before this middleware can write a title, `runtime/runs/worker.py` keeps the run in a finalizing state, persists a local fallback title from the latest checkpoint or original run input, and then syncs it to `threads_meta.display_name`. Replacement runs admitted by `multitask_strategy="interrupt"` / `"rollback"` wait for older same-thread finalization before entering the graph; the interrupted run only skips the fallback title write once a later run has started and may have advanced the checkpoint. 21. **MemoryMiddleware** - Queues conversations for async memory update (filters to user + final AI responses) 22. **ViewImageMiddleware** - *(optional, if the model supports vision)* Injects base64 image data before the LLM call -23. **DeferredToolFilterMiddleware** - *(optional, if `tool_search.enabled`)* Hides deferred (MCP) tool schemas from the bound model until `tool_search` promotes them (reads per-thread promotions from `ThreadState.promoted`, hash-scoped) -24. **SystemMessageCoalescingMiddleware** - Merges every SystemMessage into a single leading SystemMessage per request; provider-agnostic fix for strict backends (vLLM/SGLang/Qwen/Anthropic) that reject non-leading system messages. Touches the per-request payload only (checkpoint state unchanged); on midnight crossings only the latest `dynamic_context_reminder` SystemMessage survives -25. **SubagentLimitMiddleware** - *(optional, if `subagent_enabled`)* Truncates excess `task` tool calls to enforce the `MAX_CONCURRENT_SUBAGENTS` limit -26. **LoopDetectionMiddleware** - *(optional, if `loop_detection.enabled`)* Detects repeated tool-call loops; hard-stop clears both structured `tool_calls` and raw provider tool-call metadata before forcing a final text answer; stamps `loop_capped` via `consume_stop_reason` (#3875 Phase 2), symmetric to `TokenBudgetMiddleware` -27. **TokenBudgetMiddleware** - *(optional, if `token_budget.enabled`)* Enforces per-run token limits -28. **Custom middlewares** - *(optional)* Any `custom_middlewares` passed to `build_middlewares` are injected here, before the safety/clarification tail -29. **SafetyFinishReasonMiddleware** - *(optional, if `safety_finish_reason.enabled`)* Suppresses tool execution when the provider safety-terminated the response (e.g. `finish_reason=content_filter`); registered after custom middlewares so LangChain's reverse-order `after_model` dispatch runs it first -30. **ClarificationMiddleware** - Intercepts `ask_clarification` tool calls, writes a readable `ToolMessage.content` fallback plus structured `ToolMessage.artifact.human_input` request payload, and interrupts via `Command(goto=END)` (must be last). Because this middleware can short-circuit tool execution before LangChain emits `on_tool_end`, `RunJournal` performs a root-run final reconciliation for allowlisted clarification `ToolMessage`s whose `tool_call_id` was produced by the current run, so human-input request cards remain recoverable from `run_events` after checkpoint compaction. Human Input Card replies are submitted as `hide_from_ui` `HumanMessage`s with `additional_kwargs.human_input_response`; `RunJournal` persists only allowlisted hidden response sources (currently `ask_clarification`) as `llm.human.input`, which preserves answered-card state after compaction without exposing generic internal hidden context. +23. **McpRoutingMiddleware** - *(optional, if `tool_search.enabled` and PR1 MCP routing metadata produce a routing index)* Auto-promotes matching deferred MCP tool schemas before the model call by writing a minimal `promoted` state update. It matches only the latest real `HumanMessage`, uses the global `tool_search.auto_promote_top_k` limit (default 3, clamped to 1..5), never executes tools, and must be installed before `DeferredToolFilterMiddleware` +24. **DeferredToolFilterMiddleware** - *(optional, if `tool_search.enabled`)* Hides deferred (MCP) tool schemas from the bound model until `tool_search` or `McpRoutingMiddleware` promotes them (reads per-thread promotions from `ThreadState.promoted`, hash-scoped) +25. **SystemMessageCoalescingMiddleware** - Merges every SystemMessage into a single leading SystemMessage per request; provider-agnostic fix for strict backends (vLLM/SGLang/Qwen/Anthropic) that reject non-leading system messages. Touches the per-request payload only (checkpoint state unchanged); on midnight crossings only the latest `dynamic_context_reminder` SystemMessage survives +26. **SubagentLimitMiddleware** - *(optional, if `subagent_enabled`)* Truncates excess `task` tool calls to enforce the `MAX_CONCURRENT_SUBAGENTS` limit +27. **LoopDetectionMiddleware** - *(optional, if `loop_detection.enabled`)* Detects repeated tool-call loops; hard-stop clears both structured `tool_calls` and raw provider tool-call metadata before forcing a final text answer; stamps `loop_capped` via `consume_stop_reason` (#3875 Phase 2), symmetric to `TokenBudgetMiddleware` +28. **TokenBudgetMiddleware** - *(optional, if `token_budget.enabled`)* Enforces per-run token limits +29. **Custom middlewares** - *(optional)* Any `custom_middlewares` passed to `build_middlewares` are injected here, before the safety/clarification tail +30. **SafetyFinishReasonMiddleware** - *(optional, if `safety_finish_reason.enabled`)* Suppresses tool execution when the provider safety-terminated the response (e.g. `finish_reason=content_filter`); registered after custom middlewares so LangChain's reverse-order `after_model` dispatch runs it first +31. **ClarificationMiddleware** - Intercepts `ask_clarification` tool calls, writes a readable `ToolMessage.content` fallback plus structured `ToolMessage.artifact.human_input` request payload, and interrupts via `Command(goto=END)` (must be last). Because this middleware can short-circuit tool execution before LangChain emits `on_tool_end`, `RunJournal` performs a root-run final reconciliation for allowlisted clarification `ToolMessage`s whose `tool_call_id` was produced by the current run, so human-input request cards remain recoverable from `run_events` after checkpoint compaction. Human Input Card replies are submitted as `hide_from_ui` `HumanMessage`s with `additional_kwargs.human_input_response`; `RunJournal` persists only allowlisted hidden response sources (currently `ask_clarification`) as `llm.human.input`, which preserves answered-card state after compaction without exposing generic internal hidden context. ### Configuration System @@ -380,7 +381,7 @@ Proxied through nginx: `/api/langgraph/*` → Gateway LangGraph-compatible runti **Events**: `task_started`, `task_running`, `task_completed`/`task_failed`/`task_timed_out` **Guardrail caps & `stop_reason` (#3875 Phase 2)**: three independent axes can end a subagent run early, and all now surface *why* through one additive field rather than a new status enum. **Turn axis**: `recursion_limit` on the subagent `run_config` equals `max_turns`, so exhausting the turn budget raises `GraphRecursionError` from `agent.astream`; `executor.py::_aexecute` catches it specifically (before the generic `except Exception`). **Token axis**: `TokenBudgetMiddleware` is attached per-agent via `build_subagent_runtime_middlewares` from `subagents.token_budget` (default 2,000,000 tokens, warn at 0.7, hard-stop at 1.0 — a backstop against a subagent that burns tokens on trivial work). It does *not* raise: at the hard-stop threshold it strips the in-flight turn's tool calls, forces `finish_reason="stop"`, and lets the run complete naturally with a final answer. **Loop axis**: `LoopDetectionMiddleware` (attached at the same point) catches repeated identical tool-call sets — or one tool *type* called many times with varying args — and its hard-stop likewise strips `tool_calls` and forces a final answer without raising, recording `loop_capped`. Each guard exposes its cap on a per-`run_id` `consume_stop_reason(run_id)` accessor; `_aexecute` collects **every** middleware with that method (duck-typed via `hasattr`, so the executor has no import coupling to the guard classes) and surfaces the first non-`None` reason — adding a future guard needs no executor change. **Surfacing**: whichever axis fired, `_aexecute` stamps a normal status plus an additive reason — `completed` + `stop_reason=token_capped|turn_capped|loop_capped` when a usable final answer (or partial recovered from the last streamed chunk via `_extract_final_result` → `utils/messages.py::message_content_to_text`, returning a `"No response Generated"` sentinel when no text survived) was produced; `failed` + `stop_reason=turn_capped` when nothing usable survived. `SubagentResult.stop_reason` flows through `task_tool.py::_task_result_command` → `format_subagent_result_message` (renders `Task Succeeded (capped: ...)` / `Task failed (capped: ...)`) and `make_subagent_additional_kwargs`, which stamps the additive `subagent_stop_reason` key alongside the normal `subagent_status`. **Why additive, not an enum**: a new status value would break v1 consumers; an optional field is ignored by older frontends and ledger readers, so the cross-language contract (`contracts/subagent_status_contract.json` v2 + `subagents/status_contract.py` + `frontend/.../subtask-result.ts`, pinned by `test_status_values_match_contract` / `test_stop_reason_values_match_contract`) stays backward-compatible. The durable delegation ledger captures `stop_reason` onto the entry and renders model-facing guidance ("hit a guardrail cap with a partial result; reuse it, retry tighter, or raise the per-agent budget (`max_turns` / `token_budget`)") so the lead reuses a capped completion knowingly instead of mistaking it for a clean one. (Phase 1 shipped this surfacing as a `MAX_TURNS_REACHED` status enum in #3949; Phase 2 replaced that enum with the additive `stop_reason` field per the agreed design — the `max_turns_reached` status value and `SubagentStatus.MAX_TURNS_REACHED` are gone.) **Step capture & persistence (#3779)**: `executor.py` captures both assistant turns (`AIMessage`) **and** tool outputs (`ToolMessage`) via `subagents/step_events.py::capture_new_step_messages`, which walks the *newly-appended tail* of each `stream_mode="values"` chunk (not just `messages[-1]`) so a multi-tool-call turn — where LangGraph's `ToolNode` appends several `ToolMessage`s in one super-step — keeps every tool output instead of dropping all but the last. `runtime/runs/worker.py::_SubagentEventBuffer` additionally persists these `task_*` custom events to the `RunEventStore` as `subagent.start`/`subagent.step`/`subagent.end` (`category="subagent"`, `task_id` in `metadata`). It **batches** writes via `put_batch` (flushing on a terminal `subagent.end`, at `FLUSH_THRESHOLD` events, and in the worker's `finally`) rather than one `put()` per step, since `put()` is a documented low-frequency path (per-thread advisory lock per call) and a deep subagent (`max_turns=150`) emits hundreds of steps on the hot stream loop. `build_subagent_step` caps both the per-step `text` and each tool call's serialized `args` at `SUBAGENT_STEP_MAX_CHARS` (flagged `truncated` / `args_truncated`) so a large `write_file`/`bash` payload can't produce an unbounded row. The dedicated category keeps them out of `list_messages` (the thread feed) while `list_events` returns them for the frontend's fetch-on-expand backfill. `list_events` accepts `task_id` (filters on `metadata["task_id"]` — SQL-side in `DbRunEventStore` via `event_metadata["task_id"].as_string()`, in-memory in the JSONL/memory stores) plus an `after_seq` forward cursor, so the card pages through one subagent's steps without the run-wide `limit` truncating the tail (no schema migration: the filter rides the existing run-scoped index). `step_events.py` is a pure, unit-tested layer (`build_subagent_step` / `subagent_run_event`). -**Deferred MCP tools** (if `tool_search.enabled`): `SubagentExecutor._build_initial_state` assembles deferral after policy filtering via the shared `assemble_deferred_tools` (fail-closed), appends the `tool_search` tool, injects the `` section into the subagent's `SystemMessage`, and threads the setup to `_create_agent`, which attaches `DeferredToolFilterMiddleware` through `build_subagent_runtime_middlewares(deferred_setup=...)`. Subagents thus withhold full MCP schemas until promotion, same as the lead agent; each task run gets a fresh `ThreadState` so promotion is isolated per run +**Deferred MCP tools** (if `tool_search.enabled`): `SubagentExecutor._build_initial_state` assembles deferral after policy filtering via the shared `assemble_deferred_tools` (fail-closed), appends the `tool_search` tool, injects the `` section into the subagent's `SystemMessage`, and threads the setup to `_create_agent`, which attaches `McpRoutingMiddleware` (when PR1 routing metadata matches deferred tools) before `DeferredToolFilterMiddleware` through `build_subagent_runtime_middlewares(...)`. Subagents thus withhold full MCP schemas until promotion, same as the lead agent; each task run gets a fresh `ThreadState` so promotion is isolated per run **Checkpointer isolation**: Subagent graphs are compiled with `checkpointer=False` to avoid inheriting the parent run's checkpointer, since subagents are one-shot and never resume. ### Tool System (`packages/harness/deerflow/tools/`) @@ -715,7 +716,8 @@ Returns `{}` when Langfuse is not in the enabled providers — LangSmith-only de - `memory` - Memory system (enabled, storage_path, debounce_seconds, model_name, max_facts, fact_confidence_threshold, injection_enabled, max_injection_tokens, staleness_review_enabled, staleness_age_days, staleness_min_candidates, staleness_max_removals_per_cycle, staleness_protected_categories) **`extensions_config.json`**: -- `mcpServers` - Map of server name → config (enabled, type, command, args, env, url, headers, oauth, description, `routing`, `tools`, `tool_call_timeout`). `routing.mode="prefer"` emits `` prompt guidance; if `tool_search` defers the hinted tool, the guidance points at promotion first. It does not hard-disable other tools. +- `mcpServers` - Map of server name → config (enabled, type, command, args, env, url, headers, oauth, description, `routing`, `tools`, `tool_call_timeout`). `routing.mode="prefer"` emits `` prompt guidance; if `tool_search` defers the hinted tool, `McpRoutingMiddleware` can also auto-promote matching deferred schemas before the model call. It does not hard-disable other tools. +- `tool_search.auto_promote_top_k` - Global MCP routing auto-promote breadth. Default `3`, clamped to `1..5`; applies only when `tool_search.enabled=true` and only to policy-filtered deferred MCP tools with `routing.mode="prefer"` and non-empty keywords. - `skills` - Map of skill name → state (enabled) Both can be modified at runtime via Gateway API endpoints or `DeerFlowClient` methods. diff --git a/backend/README.md b/backend/README.md index 85ddb0702..e02559f83 100644 --- a/backend/README.md +++ b/backend/README.md @@ -346,10 +346,10 @@ MCP servers and skill states in a single file: ``` `routing` adds soft MCP preference hints to the agent prompt. It helps the -model prefer a configured MCP tool for matching requests without changing the -bound tool schemas or forbidding other tools. When `tool_search.enabled=true` -defers MCP schemas, the hint tells the model to fetch the deferred tool with -`tool_search` before preferring it. +model prefer a configured MCP tool for matching requests without forbidding +other tools. When `tool_search.enabled=true` defers MCP schemas, matching +routing metadata can auto-promote up to `tool_search.auto_promote_top_k` +deferred schemas before the model call. ### Environment Variables diff --git a/backend/docs/MCP_SERVER.md b/backend/docs/MCP_SERVER.md index 3e5854cb6..db2b18e38 100644 --- a/backend/docs/MCP_SERVER.md +++ b/backend/docs/MCP_SERVER.md @@ -21,8 +21,9 @@ as internal database questions that should use a PostgreSQL MCP tool before web search. Routing hints are soft model guidance: they add a `` prompt section, but they do not forbid other tools. Use agent-level allow/deny policy for hard restrictions. If `tool_search.enabled` -defers MCP tool schemas, the hint references `tool_search` so the model fetches -the deferred tool before preferring it. +defers MCP tool schemas, matching routing metadata can also auto-promote the +deferred schema before the model call. Auto-promotion is controlled by the +top-level `config.yaml -> tool_search.auto_promote_top_k` setting. ```json { @@ -53,13 +54,21 @@ the deferred tool before preferring it. - `routing.mode`: `off` disables hints; `prefer` emits hints. - `routing.priority`: `0` to `100`; higher-priority hints are rendered first. + When `tool_search.enabled=true`, priority also orders auto-promote matches. - `routing.keywords`: operator-authored terms that describe when to prefer the - MCP tool. Empty keywords are allowed but do not emit a hint line. + MCP tool. Empty keywords are allowed but do not emit a hint line and do not + trigger auto-promotion. Auto-promote matching is a case-insensitive substring + test against the latest user message (not token/word-boundary matching), so + prefer distinctive keywords — a short term like `api` also matches `rapid`. + Over-matching only exposes an extra tool schema (soft/additive), never + disables other tools. - `tools..routing`: overrides only the fields explicitly set for that tool. The key is the MCP server's original tool name, before the `_` prefix added for model binding. If the server-level `routing.mode` is `off`, a tool override must set `mode: "prefer"`; setting only `priority` or `keywords` still inherits `off` and emits no hint. +- `tool_search.auto_promote_top_k`: global limit for auto-promoted deferred MCP + schemas per model call. Default `3`; valid range `1..5`. ## Per-Tool Timeout (Stdio MCP Servers) diff --git a/backend/packages/harness/deerflow/agents/lead_agent/agent.py b/backend/packages/harness/deerflow/agents/lead_agent/agent.py index ae618299a..ce538d86c 100644 --- a/backend/packages/harness/deerflow/agents/lead_agent/agent.py +++ b/backend/packages/harness/deerflow/agents/lead_agent/agent.py @@ -223,6 +223,7 @@ def build_middlewares( available_skills: set[str] | None = None, app_config: AppConfig | None = None, deferred_setup=None, + mcp_routing_middleware: AgentMiddleware | None = None, user_id: str | None = None, ): """Build the lead-agent middleware chain based on runtime configuration. @@ -240,6 +241,8 @@ def build_middlewares( app_config: Explicit AppConfig; falls back to ``get_app_config()`` when omitted. deferred_setup: Optional deferred-MCP-tool setup that attaches ``DeferredToolFilterMiddleware`` when ``tool_search`` is enabled. + mcp_routing_middleware: Optional PR2 middleware that auto-promotes + deferred MCP schemas before the deferred filter runs. user_id: Effective user ID for user-scoped skill loading. Passed through to ``SkillActivationMiddleware`` so it can resolve per-user custom skills. @@ -302,6 +305,11 @@ def build_middlewares( if model_config is not None and model_config.supports_vision: middlewares.append(ViewImageMiddleware()) + # Auto-promote deferred MCP schemas from PR1 routing metadata before the + # deferred filter decides which schemas to hide for this model call. + if mcp_routing_middleware is not None: + middlewares.append(mcp_routing_middleware) + # Hide deferred tool schemas from model binding until tool_search promotes them. # The deferred set + catalog hash come from the build-time setup (assembled # after tool-policy filtering); promotion is read from graph state. @@ -309,6 +317,9 @@ def build_middlewares( from deerflow.agents.middlewares.deferred_tool_filter_middleware import DeferredToolFilterMiddleware middlewares.append(DeferredToolFilterMiddleware(deferred_setup.deferred_names, deferred_setup.catalog_hash)) + from deerflow.agents.middlewares.mcp_routing_middleware import assert_mcp_routing_before_deferred_filter + + assert_mcp_routing_before_deferred_filter(middlewares) # Coalesce every SystemMessage into a single leading one before the request # reaches the provider. Strict backends (vLLM, SGLang, Qwen, Anthropic) @@ -386,7 +397,7 @@ def _make_lead_agent(config: RunnableConfig, *, app_config: AppConfig): # Lazy import to avoid circular dependency from deerflow.tools import get_available_tools from deerflow.tools.builtins import setup_agent, update_agent - from deerflow.tools.builtins.tool_search import assemble_deferred_tools, get_mcp_routing_hints_prompt_section + from deerflow.tools.builtins.tool_search import assemble_deferred_tools, build_mcp_routing_middleware, get_mcp_routing_hints_prompt_section cfg = _get_runtime_config(config) resolved_app_config = app_config @@ -490,6 +501,11 @@ def _make_lead_agent(config: RunnableConfig, *, app_config: AppConfig): if non_interactive: filtered = [tool for tool in filtered if tool.name not in _NON_INTERACTIVE_DISABLED_TOOL_NAMES] final_tools, setup = assemble_deferred_tools(filtered, enabled=resolved_app_config.tool_search.enabled) + mcp_routing_middleware = build_mcp_routing_middleware( + final_tools, + setup, + top_k=resolved_app_config.tool_search.auto_promote_top_k, + ) if skill_setup.describe_skill_tool: final_tools.append(skill_setup.describe_skill_tool) return create_agent( @@ -501,6 +517,7 @@ def _make_lead_agent(config: RunnableConfig, *, app_config: AppConfig): available_skills=set(_BOOTSTRAP_SKILL_NAMES), app_config=resolved_app_config, deferred_setup=setup, + mcp_routing_middleware=mcp_routing_middleware, user_id=resolved_user_id, ), system_prompt=apply_prompt_template( @@ -546,6 +563,11 @@ def _make_lead_agent(config: RunnableConfig, *, app_config: AppConfig): if non_interactive: filtered = [tool for tool in filtered if tool.name not in _NON_INTERACTIVE_DISABLED_TOOL_NAMES] final_tools, setup = assemble_deferred_tools(filtered, enabled=resolved_app_config.tool_search.enabled) + mcp_routing_middleware = build_mcp_routing_middleware( + final_tools, + setup, + top_k=resolved_app_config.tool_search.auto_promote_top_k, + ) mcp_routing_hints_section = get_mcp_routing_hints_prompt_section(filtered, deferred_names=setup.deferred_names) if skill_setup.describe_skill_tool: final_tools.append(skill_setup.describe_skill_tool) @@ -559,6 +581,7 @@ def _make_lead_agent(config: RunnableConfig, *, app_config: AppConfig): available_skills=available_skills, app_config=resolved_app_config, deferred_setup=setup, + mcp_routing_middleware=mcp_routing_middleware, user_id=resolved_user_id, ), system_prompt=apply_prompt_template( diff --git a/backend/packages/harness/deerflow/agents/middlewares/mcp_routing_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/mcp_routing_middleware.py new file mode 100644 index 000000000..938e4025e --- /dev/null +++ b/backend/packages/harness/deerflow/agents/middlewares/mcp_routing_middleware.py @@ -0,0 +1,137 @@ +"""Auto-promote deferred MCP tools from routing metadata before model calls.""" + +from __future__ import annotations + +import logging +from collections.abc import Mapping, Sequence +from typing import Any, TypedDict, override + +from langchain.agents import AgentState +from langchain.agents.middleware import AgentMiddleware +from langchain_core.messages import HumanMessage +from langgraph.runtime import Runtime + +from deerflow.config.tool_search_config import clamp_auto_promote_top_k +from deerflow.utils.messages import get_original_user_content_text, is_real_user_message + +logger = logging.getLogger(__name__) + + +class McpRoutingIndexEntry(TypedDict): + priority: int + keywords: list[str] + + +McpRoutingIndex = Mapping[str, McpRoutingIndexEntry] + + +class McpRoutingMiddleware(AgentMiddleware[AgentState]): + """Write minimal deferred-tool promotion state from latest user text. + + The middleware intentionally receives only serialized routing data. It does + not hold ``BaseTool`` objects, does not execute tools, and does not filter + tool calls. ``DeferredToolFilterMiddleware`` remains responsible for hiding + unpromoted schemas and blocking unpromoted deferred tool calls. + """ + + def __init__( + self, + routing_index: McpRoutingIndex, + catalog_hash: str | None, + top_k: int, + ) -> None: + super().__init__() + self._catalog_hash = catalog_hash + self._top_k = clamp_auto_promote_top_k(top_k) + self._routing_index = self._normalize_index(routing_index) + + @staticmethod + def _normalize_index(routing_index: McpRoutingIndex) -> dict[str, tuple[int, tuple[str, ...]]]: + # Defensive re-normalization: this middleware is built to accept arbitrary + # serialized routing data, not only the output of + # tool_search._routing_priority / _routing_keywords. In practice it is a + # no-op over the builder's output; keep the coercion rules aligned with + # those two helpers if either side changes. + normalized: dict[str, tuple[int, tuple[str, ...]]] = {} + for raw_name, raw_entry in routing_index.items(): + name = str(raw_name) + if not name: + continue + try: + priority = int(raw_entry.get("priority", 0)) + except (TypeError, ValueError): + priority = 0 + raw_keywords = raw_entry.get("keywords") or [] + if not isinstance(raw_keywords, Sequence) or isinstance(raw_keywords, (str, bytes)): + raw_keywords = [] + keywords = tuple(keyword for keyword in (str(item).strip() for item in raw_keywords) if keyword) + if not keywords: + continue + normalized[name] = (priority, keywords) + return normalized + + @staticmethod + def _latest_user_message(messages: list[Any]) -> HumanMessage | None: + for message in reversed(messages): + if is_real_user_message(message): + return message + return None + + def _matched_names(self, state: Mapping[str, Any] | None) -> list[str]: + if not self._catalog_hash or not self._routing_index: + return [] + messages = list((state or {}).get("messages") or []) + target = self._latest_user_message(messages) + if target is None: + return [] + + text = get_original_user_content_text(target.content, target.additional_kwargs) + if not text: + return [] + + haystack = text.casefold() + matched: list[tuple[int, str]] = [] + for name, (priority, keywords) in self._routing_index.items(): + if any(keyword.casefold() in haystack for keyword in keywords): + matched.append((priority, name)) + + if not matched: + return [] + + matched.sort(key=lambda item: (-item[0], item[1])) + return [name for _, name in matched[: self._top_k]] + + def _state_update(self, state: Mapping[str, Any] | None) -> dict[str, Any] | None: + names = self._matched_names(state) + if not names: + return None + logger.debug( + "McpRoutingMiddleware auto-promoted %d deferred tool schema(s) catalog=%s names=%s", + len(names), + (self._catalog_hash or "")[:8], + names, + ) + return { + "promoted": { + "catalog_hash": self._catalog_hash, + "names": names, + } + } + + @override + def before_model(self, state: AgentState, runtime: Runtime) -> dict[str, Any] | None: + return self._state_update(state) + + @override + async def abefore_model(self, state: AgentState, runtime: Runtime) -> dict[str, Any] | None: + return self._state_update(state) + + +def assert_mcp_routing_before_deferred_filter(middlewares: Sequence[AgentMiddleware]) -> None: + """Fail fast if auto-promote would run after deferred schema filtering.""" + from deerflow.agents.middlewares.deferred_tool_filter_middleware import DeferredToolFilterMiddleware + + routing_idx = next((idx for idx, middleware in enumerate(middlewares) if isinstance(middleware, McpRoutingMiddleware)), None) + filter_idx = next((idx for idx, middleware in enumerate(middlewares) if isinstance(middleware, DeferredToolFilterMiddleware)), None) + if routing_idx is not None and filter_idx is not None and routing_idx > filter_idx: + raise RuntimeError(f"McpRoutingMiddleware must be installed before DeferredToolFilterMiddleware (routing index {routing_idx}, deferred filter index {filter_idx})") diff --git a/backend/packages/harness/deerflow/agents/middlewares/skill_activation_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/skill_activation_middleware.py index 7a8405715..4f86be6d8 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/skill_activation_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/skill_activation_middleware.py @@ -27,7 +27,7 @@ from deerflow.skills.slash import parse_slash_skill_reference, resolve_slash_ski from deerflow.skills.storage import get_or_new_skill_storage, get_or_new_user_skill_storage from deerflow.skills.storage.skill_storage import SkillStorage from deerflow.skills.types import SKILL_MD_FILE, SecretRequirement, Skill, SkillCategory -from deerflow.utils.messages import get_original_user_content_text +from deerflow.utils.messages import get_original_user_content_text, is_real_user_message if TYPE_CHECKING: from deerflow.config.app_config import AppConfig @@ -36,7 +36,6 @@ logger = logging.getLogger(__name__) _SLASH_SKILL_ACTIVATION_KEY = "slash_skill_activation" _SLASH_SKILL_ACTIVATION_TARGET_ID_KEY = "slash_skill_activation_target_id" -_SUMMARY_MESSAGE_NAME = "summary" # _SECRETS_BINDING_AUDIT_KEY: last audited binding (skill and secret names only, # never values) so unchanged bindings are not re-recorded each call. @@ -73,13 +72,7 @@ def is_slash_skill_activation_reminder(message: object) -> bool: def _is_user_activation_target(message: object) -> bool: - if not isinstance(message, HumanMessage): - return False - if message.name == _SUMMARY_MESSAGE_NAME: - return False - if message.additional_kwargs.get("hide_from_ui"): - return False - return True + return is_real_user_message(message) class SkillActivationMiddleware(AgentMiddleware): diff --git a/backend/packages/harness/deerflow/agents/middlewares/tool_error_handling_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/tool_error_handling_middleware.py index 0a69ecf6f..03c7c9d6f 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/tool_error_handling_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/tool_error_handling_middleware.py @@ -280,6 +280,7 @@ def build_subagent_runtime_middlewares( model_name: str | None = None, lazy_init: bool = True, deferred_setup: "DeferredToolSetup | None" = None, + mcp_routing_middleware: AgentMiddleware | None = None, agent_name: str | None = None, ) -> list[AgentMiddleware]: """Middlewares shared by subagent runtime before subagent-only middlewares.""" @@ -304,6 +305,9 @@ def build_subagent_runtime_middlewares( middlewares.append(ViewImageMiddleware()) + if mcp_routing_middleware is not None: + middlewares.append(mcp_routing_middleware) + # Hide deferred (MCP) tool schemas from the subagent's model binding until # tool_search promotes them. This is the same wiring the lead agent gets. The deferred # set + catalog hash come from the build-time setup (assembled after @@ -313,6 +317,9 @@ def build_subagent_runtime_middlewares( from deerflow.agents.middlewares.deferred_tool_filter_middleware import DeferredToolFilterMiddleware middlewares.append(DeferredToolFilterMiddleware(deferred_setup.deferred_names, deferred_setup.catalog_hash)) + from deerflow.agents.middlewares.mcp_routing_middleware import assert_mcp_routing_before_deferred_filter + + assert_mcp_routing_before_deferred_filter(middlewares) # LoopDetectionMiddleware — subagents inherit none of the lead's runaway # guards today (see #3875): with no loop detection a degenerate subagent tool diff --git a/backend/packages/harness/deerflow/client.py b/backend/packages/harness/deerflow/client.py index 383ad1a18..641835445 100644 --- a/backend/packages/harness/deerflow/client.py +++ b/backend/packages/harness/deerflow/client.py @@ -46,7 +46,7 @@ from deerflow.runtime.goal import DEFAULT_MAX_GOAL_CONTINUATIONS, build_goal_sta from deerflow.runtime.user_context import get_effective_user_id from deerflow.skills.describe import build_skill_search_setup from deerflow.skills.storage import get_or_new_user_skill_storage -from deerflow.tools.builtins.tool_search import assemble_deferred_tools, get_mcp_routing_hints_prompt_section +from deerflow.tools.builtins.tool_search import assemble_deferred_tools, build_mcp_routing_middleware, get_mcp_routing_hints_prompt_section from deerflow.trace_context import DEERFLOW_TRACE_METADATA_KEY, generate_trace_id, get_current_trace_id, reset_current_trace_id, set_current_trace_id from deerflow.tracing import build_tracing_callbacks, inject_langfuse_metadata from deerflow.uploads.manager import ( @@ -257,6 +257,11 @@ class DeerFlowClient: tools = self._get_tools(model_name=model_name, subagent_enabled=subagent_enabled) final_tools, deferred_setup = assemble_deferred_tools(tools, enabled=self._app_config.tool_search.enabled) + mcp_routing_middleware = build_mcp_routing_middleware( + final_tools, + deferred_setup, + top_k=self._app_config.tool_search.auto_promote_top_k, + ) mcp_routing_hints_section = get_mcp_routing_hints_prompt_section(tools, deferred_names=deferred_setup.deferred_names) # Wire deferred skill discovery — mirrors agent.py so config flag works on both paths. @@ -286,6 +291,7 @@ class DeerFlowClient: custom_middlewares=self._middlewares, app_config=self._app_config, deferred_setup=deferred_setup, + mcp_routing_middleware=mcp_routing_middleware, user_id=get_effective_user_id(), ), "system_prompt": apply_prompt_template( diff --git a/backend/packages/harness/deerflow/config/tool_search_config.py b/backend/packages/harness/deerflow/config/tool_search_config.py index cdeddabf2..0a1a8938b 100644 --- a/backend/packages/harness/deerflow/config/tool_search_config.py +++ b/backend/packages/harness/deerflow/config/tool_search_config.py @@ -1,6 +1,14 @@ """Configuration for deferred tool loading via tool_search.""" -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, field_validator + +AUTO_PROMOTE_TOP_K_MIN = 1 +AUTO_PROMOTE_TOP_K_MAX = 5 + + +def clamp_auto_promote_top_k(value: int) -> int: + """Clamp the global MCP routing auto-promote breadth to PR2's range.""" + return max(AUTO_PROMOTE_TOP_K_MIN, min(AUTO_PROMOTE_TOP_K_MAX, int(value))) class ToolSearchConfig(BaseModel): @@ -15,6 +23,15 @@ class ToolSearchConfig(BaseModel): default=False, description="Defer tools and enable tool_search", ) + auto_promote_top_k: int = Field( + default=3, + description="Maximum number of deferred MCP tool schemas auto-promoted from routing metadata per model call", + ) + + @field_validator("auto_promote_top_k") + @classmethod + def _clamp_auto_promote_top_k(cls, value: int) -> int: + return clamp_auto_promote_top_k(value) _tool_search_config: ToolSearchConfig | None = None diff --git a/backend/packages/harness/deerflow/subagents/executor.py b/backend/packages/harness/deerflow/subagents/executor.py index dfb655716..139067308 100644 --- a/backend/packages/harness/deerflow/subagents/executor.py +++ b/backend/packages/harness/deerflow/subagents/executor.py @@ -437,13 +437,25 @@ class SubagentExecutor: # Reuse shared middleware composition with lead agent. ``agent_name`` # lets the builder resolve the per-agent token_budget override. - middlewares = build_subagent_runtime_middlewares( - app_config=app_config, - model_name=self.model_name, - lazy_init=True, - deferred_setup=deferred_setup, - agent_name=self.config.name, - ) + mcp_routing_middleware = None + if deferred_setup is not None and deferred_setup.deferred_names: + from deerflow.tools.builtins.tool_search import build_mcp_routing_middleware + + mcp_routing_middleware = build_mcp_routing_middleware( + tools if tools is not None else self.tools, + deferred_setup, + top_k=app_config.tool_search.auto_promote_top_k, + ) + middleware_kwargs = { + "app_config": app_config, + "model_name": self.model_name, + "lazy_init": True, + "deferred_setup": deferred_setup, + "agent_name": self.config.name, + } + if mcp_routing_middleware is not None: + middleware_kwargs["mcp_routing_middleware"] = mcp_routing_middleware + middlewares = build_subagent_runtime_middlewares(**middleware_kwargs) # Collect every guard middleware that exposes ``consume_stop_reason`` # (TokenBudgetMiddleware, LoopDetectionMiddleware) so _aexecute can read # each after the run and surface whichever cap fired. Duck-typed diff --git a/backend/packages/harness/deerflow/tools/builtins/tool_search.py b/backend/packages/harness/deerflow/tools/builtins/tool_search.py index f39e3d364..57b5ad026 100644 --- a/backend/packages/harness/deerflow/tools/builtins/tool_search.py +++ b/backend/packages/harness/deerflow/tools/builtins/tool_search.py @@ -6,6 +6,8 @@ Contains: catalog; it records promotions into graph state via ``Command``. - build_deferred_tool_setup: assembles the catalog + tool from a policy-filtered tool list (call AFTER tool-policy filtering). +- build_mcp_routing_middleware: builds the PR2 auto-promote middleware from + serialized routing metadata on policy-filtered deferred tools. The agent sees deferred tool names in but cannot call them until it fetches their full schema via the tool_search tool. The @@ -21,7 +23,7 @@ import re from collections.abc import Iterable from dataclasses import dataclass from functools import cached_property -from typing import Annotated +from typing import TYPE_CHECKING, Annotated, Any from langchain.tools import BaseTool from langchain_core.messages import ToolMessage @@ -31,6 +33,9 @@ from langgraph.types import Command from deerflow.tools.mcp_metadata import get_mcp_routing, is_mcp_tool +if TYPE_CHECKING: + from langchain.agents.middleware import AgentMiddleware + logger = logging.getLogger(__name__) MAX_RESULTS = 5 # Max tools returned per search @@ -202,6 +207,64 @@ def assemble_deferred_tools(filtered_tools: list[BaseTool], *, enabled: bool) -> return final_tools, deferred_setup +def _routing_priority(value: Any) -> int: + # Produces the typed priority stored in the routing index. McpRoutingMiddleware + # ._normalize_index re-parses this defensively (it is built to accept arbitrary + # serialized data), so keep the two coercion rules in sync if either changes. + try: + return int(value) + except (TypeError, ValueError): + return 0 + + +def _routing_keywords(value: Any) -> list[str]: + # See _routing_priority: McpRoutingMiddleware._normalize_index re-normalizes + # keywords defensively; keep both coercion rules aligned. + if not isinstance(value, list): + return [] + return [keyword for keyword in (str(item).strip() for item in value) if keyword] + + +def build_mcp_routing_middleware( + tools: Iterable[BaseTool], + deferred_setup: DeferredToolSetup, + *, + top_k: int, +) -> "AgentMiddleware | None": + """Build PR2 auto-promotion middleware from policy-filtered deferred tools. + + The builder may inspect ``BaseTool.metadata`` at construction time, but the + returned middleware receives only a flat serializable routing index. + """ + if deferred_setup.catalog_hash is None or not deferred_setup.deferred_names: + return None + + routing_index: dict[str, dict[str, Any]] = {} + for candidate in tools: + tool_name = getattr(candidate, "name", "") + if tool_name not in deferred_setup.deferred_names: + continue + routing = get_mcp_routing(candidate) + if routing is None or routing.get("mode") != "prefer": + continue + keywords = _routing_keywords(routing.get("keywords")) + if not keywords: + continue + if routing.get("auto_promote_top_k") is not None: + logger.debug("Ignoring per-tool MCP routing auto_promote_top_k for %s in PR2", tool_name) + routing_index[str(tool_name)] = { + "priority": _routing_priority(routing.get("priority", 0)), + "keywords": keywords, + } + + if not routing_index: + return None + + from deerflow.agents.middlewares.mcp_routing_middleware import McpRoutingMiddleware + + return McpRoutingMiddleware(routing_index, deferred_setup.catalog_hash, top_k) + + # Prompt rendering diff --git a/backend/packages/harness/deerflow/utils/messages.py b/backend/packages/harness/deerflow/utils/messages.py index be1160721..e54807467 100644 --- a/backend/packages/harness/deerflow/utils/messages.py +++ b/backend/packages/harness/deerflow/utils/messages.py @@ -3,7 +3,10 @@ from __future__ import annotations from collections.abc import Mapping from typing import Any +from langchain_core.messages import HumanMessage + ORIGINAL_USER_CONTENT_KEY = "original_user_content" +SUMMARY_MESSAGE_NAME = "summary" def message_content_to_text(content: Any) -> str: @@ -72,3 +75,18 @@ def get_original_user_content_text(content: Any, additional_kwargs: Mapping[str, if isinstance(original_content, str): return original_content return message_content_to_text(content) + + +def is_real_user_message(message: object) -> bool: + """Return whether ``message`` is a real user-authored HumanMessage. + + Middleware-injected hidden HumanMessages and summarization markers should not + drive user-intent features such as slash-skill activation or MCP routing. + """ + if not isinstance(message, HumanMessage): + return False + if message.name == SUMMARY_MESSAGE_NAME: + return False + if message.additional_kwargs.get("hide_from_ui"): + return False + return True diff --git a/backend/tests/test_client.py b/backend/tests/test_client.py index 823baaaab..a32f767d2 100644 --- a/backend/tests/test_client.py +++ b/backend/tests/test_client.py @@ -1087,6 +1087,68 @@ class TestEnsureAgent: skill_names_arg = mock_apply_prompt.call_args.kwargs.get("skill_names") assert skill_names_arg is None, "skill_names must be None when deferred_discovery=False" + def test_mcp_routing_middleware_wired_when_tool_search_enabled(self, client, mock_app_config): + """Embedded client builds McpRoutingMiddleware from routed deferred MCP tools. + + RFC §10.3/§12.5 requires verifying the actual embedded-client builder path + rather than assuming it inherits lead-agent behavior. Exercises the real + assemble_deferred_tools + build_mcp_routing_middleware wiring and asserts a + genuine McpRoutingMiddleware reaches build_middlewares. + """ + from langchain_core.tools import tool as as_tool + + from deerflow.agents.middlewares.mcp_routing_middleware import McpRoutingMiddleware + from deerflow.tools.mcp_metadata import tag_mcp_routing, tag_mcp_tool + + @as_tool + def postgres_query(sql: str) -> str: + "Query Postgres." + return sql + + tag_mcp_tool(postgres_query) + tag_mcp_routing(postgres_query, {"mode": "prefer", "priority": 100, "keywords": ["orders"]}) + + mock_app_config.tool_search.enabled = True + mock_app_config.tool_search.auto_promote_top_k = 3 + mock_app_config.skills.deferred_discovery = False + client._app_config = mock_app_config + config = client._get_runnable_config("t1") + + with ( + patch("deerflow.client.create_chat_model"), + patch("deerflow.client.create_agent", return_value=MagicMock()), + patch("deerflow.client.build_middlewares", return_value=[]) as mock_build_middlewares, + patch("deerflow.client.apply_prompt_template", return_value="prompt"), + patch.object(client, "_get_tools", return_value=[postgres_query]), + patch("deerflow.runtime.checkpointer.get_checkpointer", return_value=None), + patch("deerflow.client.get_enabled_skills_for_config", return_value=[]), + ): + client._ensure_agent(config) + + routing_arg = mock_build_middlewares.call_args.kwargs.get("mcp_routing_middleware") + assert isinstance(routing_arg, McpRoutingMiddleware) + assert routing_arg._matched_names({"messages": [HumanMessage(content="show orders")]}) == ["postgres_query"] + + def test_mcp_routing_middleware_absent_when_tool_search_disabled(self, client, mock_app_config): + """No routing middleware is built on the embedded path when tool_search is off.""" + mock_app_config.tool_search.enabled = False + mock_app_config.skills.deferred_discovery = False + client._app_config = mock_app_config + config = client._get_runnable_config("t1") + + with ( + patch("deerflow.client.create_chat_model"), + patch("deerflow.client.create_agent", return_value=MagicMock()), + patch("deerflow.client.build_middlewares", return_value=[]) as mock_build_middlewares, + patch("deerflow.client.apply_prompt_template", return_value="prompt"), + patch.object(client, "_get_tools", return_value=[]), + patch("deerflow.runtime.checkpointer.get_checkpointer", return_value=None), + patch("deerflow.client.get_enabled_skills_for_config", return_value=[]), + ): + client._ensure_agent(config) + + assert mock_build_middlewares.call_args.kwargs.get("mcp_routing_middleware") is None + # --------------------------------------------------------------------------- # get_model diff --git a/backend/tests/test_lead_agent_model_resolution.py b/backend/tests/test_lead_agent_model_resolution.py index a9d669518..4083cf576 100644 --- a/backend/tests/test_lead_agent_model_resolution.py +++ b/backend/tests/test_lead_agent_model_resolution.py @@ -424,6 +424,33 @@ def test_build_middlewares_passes_explicit_app_config_to_shared_factory(monkeypa assert middlewares[0] == "base-middleware" +def test_build_middlewares_places_mcp_routing_before_deferred_filter(monkeypatch): + from deerflow.agents.middlewares.deferred_tool_filter_middleware import DeferredToolFilterMiddleware + from deerflow.agents.middlewares.mcp_routing_middleware import McpRoutingMiddleware + from deerflow.tools.builtins.tool_search import DeferredToolSetup + + app_config = _make_app_config([_make_model("safe-model", supports_thinking=False)], loop_detection=LoopDetectionConfig(enabled=False)) + routing = McpRoutingMiddleware({"mcp_thing": {"priority": 100, "keywords": ["orders"]}}, "hash123", 3) + setup = DeferredToolSetup(object(), frozenset({"mcp_thing"}), "hash123") + + monkeypatch.setattr(lead_agent_module, "get_app_config", lambda: app_config) + monkeypatch.setattr(lead_agent_module, "build_lead_runtime_middlewares", lambda *, app_config, lazy_init=True: []) + monkeypatch.setattr(lead_agent_module, "_create_summarization_middleware", lambda *, app_config=None: None) + monkeypatch.setattr(lead_agent_module, "_create_todo_list_middleware", lambda is_plan_mode: None) + + middlewares = lead_agent_module.build_middlewares( + {"configurable": {"is_plan_mode": False, "subagent_enabled": False}}, + model_name="safe-model", + app_config=app_config, + deferred_setup=setup, + mcp_routing_middleware=routing, + ) + + routing_idx = next(i for i, middleware in enumerate(middlewares) if isinstance(middleware, McpRoutingMiddleware)) + filter_idx = next(i for i, middleware in enumerate(middlewares) if isinstance(middleware, DeferredToolFilterMiddleware)) + assert routing_idx < filter_idx + + def test_build_middlewares_uses_loop_detection_config(monkeypatch): app_config = _make_app_config( [_make_model("safe-model", supports_thinking=False)], diff --git a/backend/tests/test_mcp_routing_auto_promote.py b/backend/tests/test_mcp_routing_auto_promote.py new file mode 100644 index 000000000..59027a194 --- /dev/null +++ b/backend/tests/test_mcp_routing_auto_promote.py @@ -0,0 +1,296 @@ +"""Tests for PR2 MCP routing auto-promotion.""" + +import asyncio + +import pytest +from langchain.agents import create_agent +from langchain_core.language_models.fake_chat_models import GenericFakeChatModel +from langchain_core.messages import AIMessage, HumanMessage, ToolMessage +from langchain_core.tools import tool as as_tool + +from deerflow.agents.middlewares.deferred_tool_filter_middleware import DeferredToolFilterMiddleware +from deerflow.agents.middlewares.mcp_routing_middleware import McpRoutingMiddleware, assert_mcp_routing_before_deferred_filter +from deerflow.agents.thread_state import ThreadState, merge_promoted +from deerflow.tools.builtins.tool_search import assemble_deferred_tools, build_mcp_routing_middleware +from deerflow.tools.mcp_metadata import tag_mcp_routing, tag_mcp_tool +from deerflow.utils.messages import ORIGINAL_USER_CONTENT_KEY + + +@as_tool +def active_tool(x: str) -> str: + "An always-active tool." + return x + + +@as_tool +def postgres_query(sql: str) -> str: + "Query Postgres." + return sql + + +@as_tool +def metrics_query(query: str) -> str: + "Query metrics." + return query + + +@as_tool +def archive_lookup(query: str) -> str: + "Search archived records." + return query + + +def _routed(tool, *, keywords: list[str], priority: int = 0, mode: str = "prefer"): + tag_mcp_tool(tool) + tag_mcp_routing( + tool, + { + "mode": mode, + "priority": priority, + "keywords": keywords, + }, + ) + return tool + + +def test_builder_indexes_only_deferred_prefer_tools(): + routed = _routed(postgres_query, keywords=["orders"], priority=100) + off = _routed(metrics_query, keywords=["metrics"], priority=50, mode="off") + empty_keywords = _routed(archive_lookup, keywords=[], priority=90) + final_tools, setup = assemble_deferred_tools([active_tool, routed, off, empty_keywords], enabled=True) + + middleware = build_mcp_routing_middleware(final_tools, setup, top_k=3) + + assert isinstance(middleware, McpRoutingMiddleware) + assert middleware._matched_names({"messages": [HumanMessage(content="show ORDERS")]}) == ["postgres_query"] + assert middleware._matched_names({"messages": [HumanMessage(content="metrics archive")]}) == [] + + +def test_builder_skips_when_tool_search_disabled_or_no_index(): + routed = _routed(postgres_query, keywords=["orders"], priority=100) + final_tools, setup = assemble_deferred_tools([routed], enabled=False) + + assert build_mcp_routing_middleware(final_tools, setup, top_k=3) is None + + _, setup = assemble_deferred_tools([_routed(metrics_query, keywords=[], priority=50)], enabled=True) + assert build_mcp_routing_middleware([metrics_query], setup, top_k=3) is None + + +def test_matching_uses_latest_real_human_message_only(): + middleware = McpRoutingMiddleware( + { + "postgres_query": {"priority": 100, "keywords": ["orders"]}, + "metrics_query": {"priority": 90, "keywords": ["metrics"]}, + }, + "hash1", + 3, + ) + + assert middleware._matched_names({"messages": [HumanMessage(content="orders"), HumanMessage(content="no match now")]}) == [] + assert middleware._matched_names({"messages": [HumanMessage(content="metrics", name="summary"), HumanMessage(content="orders", additional_kwargs={"hide_from_ui": True})]}) == [] + + +def test_matching_supports_casefold_chinese_priority_tiebreak_and_top_k(): + middleware = McpRoutingMiddleware( + { + "z_tool": {"priority": 50, "keywords": ["订单"]}, + "a_tool": {"priority": 50, "keywords": ["orders"]}, + "top_tool": {"priority": 100, "keywords": ["ORDERS"]}, + }, + "hash1", + 2, + ) + + assert middleware._matched_names({"messages": [HumanMessage(content="查订单 and orders")]}) == ["top_tool", "a_tool"] + + +def test_structured_original_user_text_is_used(): + middleware = McpRoutingMiddleware( + {"postgres_query": {"priority": 100, "keywords": ["orders"]}}, + "hash1", + 3, + ) + message = HumanMessage( + content=[{"type": "text", "text": "sanitized replacement"}], + additional_kwargs={ORIGINAL_USER_CONTENT_KEY: "show orders"}, + ) + + assert middleware._matched_names({"messages": [message]}) == ["postgres_query"] + + +def test_before_model_returns_minimal_promoted_update_and_reducer_unions(): + middleware = McpRoutingMiddleware( + {"postgres_query": {"priority": 100, "keywords": ["orders"]}}, + "hash1", + 3, + ) + + update = middleware.before_model( + {"messages": [HumanMessage(content="orders")], "promoted": {"catalog_hash": "hash1", "names": ["metrics_query"]}}, + runtime=None, + ) + + assert update == {"promoted": {"catalog_hash": "hash1", "names": ["postgres_query"]}} + assert merge_promoted({"catalog_hash": "hash1", "names": ["metrics_query"]}, update["promoted"]) == { + "catalog_hash": "hash1", + "names": ["metrics_query", "postgres_query"], + } + + +@pytest.mark.asyncio +async def test_abefore_model_matches_sync_behavior(): + middleware = McpRoutingMiddleware( + {"postgres_query": {"priority": 100, "keywords": ["orders"]}}, + "hash1", + 3, + ) + + assert await middleware.abefore_model({"messages": [HumanMessage(content="orders")]}, runtime=None) == {"promoted": {"catalog_hash": "hash1", "names": ["postgres_query"]}} + + +def test_no_match_and_missing_catalog_hash_return_no_update(): + assert McpRoutingMiddleware({"postgres_query": {"priority": 100, "keywords": ["orders"]}}, None, 3).before_model({"messages": [HumanMessage(content="orders")]}, runtime=None) is None + assert McpRoutingMiddleware({"postgres_query": {"priority": 100, "keywords": ["orders"]}}, "hash1", 3).before_model({"messages": [HumanMessage(content="nothing")]}, runtime=None) is None + + +def test_order_invariant_rejects_reversed_middlewares(): + routing = McpRoutingMiddleware({"postgres_query": {"priority": 100, "keywords": ["orders"]}}, "hash1", 3) + deferred = DeferredToolFilterMiddleware(frozenset({"postgres_query"}), "hash1") + + assert_mcp_routing_before_deferred_filter([routing, deferred]) + with pytest.raises(RuntimeError, match="McpRoutingMiddleware must be installed before DeferredToolFilterMiddleware"): + assert_mcp_routing_before_deferred_filter([deferred, routing]) + + +def test_auto_promote_makes_schema_visible_in_same_model_cycle(): + bound: list[list[str]] = [] + + class RecordingModel(GenericFakeChatModel): + def bind_tools(self, tools, **kwargs): + bound.append([getattr(t, "name", None) for t in tools]) + return self + + routed = _routed(postgres_query, keywords=["orders"], priority=100) + other = _routed(metrics_query, keywords=["metrics"], priority=90) + final_tools, setup = assemble_deferred_tools([active_tool, routed, other], enabled=True) + routing_middleware = build_mcp_routing_middleware(final_tools, setup, top_k=3) + assert routing_middleware is not None + + model = RecordingModel(messages=iter([AIMessage(content="done")])) + graph = create_agent( + model=model, + tools=final_tools, + middleware=[ + routing_middleware, + DeferredToolFilterMiddleware(setup.deferred_names, setup.catalog_hash), + ], + state_schema=ThreadState, + ) + + result = asyncio.run(graph.ainvoke({"messages": [HumanMessage(content="show orders")]})) + + assert "postgres_query" in bound[0] + assert "metrics_query" not in bound[0] + assert result["promoted"] == {"catalog_hash": setup.catalog_hash, "names": ["postgres_query"]} + assert not any(isinstance(message, ToolMessage) for message in result["messages"]) + + +def test_auto_promoted_tool_can_be_called_without_tool_search(): + bound: list[list[str]] = [] + + class RecordingModel(GenericFakeChatModel): + def bind_tools(self, tools, **kwargs): + bound.append([getattr(t, "name", None) for t in tools]) + return self + + routed = _routed(postgres_query, keywords=["orders"], priority=100) + final_tools, setup = assemble_deferred_tools([active_tool, routed], enabled=True) + routing_middleware = build_mcp_routing_middleware(final_tools, setup, top_k=3) + assert routing_middleware is not None + + turn1 = AIMessage(content="", tool_calls=[{"name": "postgres_query", "args": {"sql": "select * from orders"}, "id": "c1", "type": "tool_call"}]) + turn2 = AIMessage(content="done") + model = RecordingModel(messages=iter([turn1, turn2])) + graph = create_agent( + model=model, + tools=final_tools, + middleware=[ + routing_middleware, + DeferredToolFilterMiddleware(setup.deferred_names, setup.catalog_hash), + ], + state_schema=ThreadState, + ) + + result = asyncio.run(graph.ainvoke({"messages": [HumanMessage(content="show orders")]})) + + assert "postgres_query" in bound[0] + assert result["promoted"] == {"catalog_hash": setup.catalog_hash, "names": ["postgres_query"]} + tool_messages = [message for message in result["messages"] if isinstance(message, ToolMessage)] + assert tool_messages + assert tool_messages[0].name == "postgres_query" + assert tool_messages[0].status == "success" + + +def test_explicit_tool_search_merges_with_auto_promoted_names(): + class RecordingModel(GenericFakeChatModel): + def bind_tools(self, tools, **kwargs): + return self + + routed = _routed(postgres_query, keywords=["orders"], priority=100) + other = _routed(metrics_query, keywords=["metrics"], priority=90) + final_tools, setup = assemble_deferred_tools([active_tool, routed, other], enabled=True) + routing_middleware = build_mcp_routing_middleware(final_tools, setup, top_k=3) + assert routing_middleware is not None + + turn1 = AIMessage(content="", tool_calls=[{"name": "tool_search", "args": {"query": "select:metrics_query"}, "id": "c1", "type": "tool_call"}]) + turn2 = AIMessage(content="done") + model = RecordingModel(messages=iter([turn1, turn2])) + graph = create_agent( + model=model, + tools=final_tools, + middleware=[ + routing_middleware, + DeferredToolFilterMiddleware(setup.deferred_names, setup.catalog_hash), + ], + state_schema=ThreadState, + ) + + result = asyncio.run(graph.ainvoke({"messages": [HumanMessage(content="show orders")]})) + + assert result["promoted"] == { + "catalog_hash": setup.catalog_hash, + "names": ["postgres_query", "metrics_query"], + } + + +def test_bootstrap_like_no_mcp_tools_skips_middleware(): + final_tools, setup = assemble_deferred_tools([active_tool], enabled=True) + + assert build_mcp_routing_middleware(final_tools, setup, top_k=3) is None + + +def test_acp_tool_without_mcp_metadata_is_not_indexed(): + final_tools, setup = assemble_deferred_tools([active_tool], enabled=True) + + assert setup.deferred_names == frozenset() + assert build_mcp_routing_middleware(final_tools, setup, top_k=3) is None + + +def test_privacy_no_trace_metadata_or_info_logs(caplog): + caplog.set_level("INFO") + middleware = McpRoutingMiddleware( + {"secret_tool": {"priority": 100, "keywords": ["sensitive-keyword"]}}, + "hash1", + 3, + ) + state = { + "messages": [HumanMessage(content="contains sensitive-keyword")], + "metadata": {"trace": "existing"}, + } + + update = middleware.before_model(state, runtime=None) + + assert update == {"promoted": {"catalog_hash": "hash1", "names": ["secret_tool"]}} + assert state["metadata"] == {"trace": "existing"} + assert "sensitive-keyword" not in caplog.text + assert "secret_tool" not in caplog.text diff --git a/backend/tests/test_subagent_executor.py b/backend/tests/test_subagent_executor.py index 699cf8ca2..2cc00c00f 100644 --- a/backend/tests/test_subagent_executor.py +++ b/backend/tests/test_subagent_executor.py @@ -649,7 +649,7 @@ class TestAgentConstruction: from deerflow.tools.builtins.tool_search import DeferredToolSetup SubagentExecutor = classes["SubagentExecutor"] - app_config = SimpleNamespace(models=[SimpleNamespace(name="default-model")]) + app_config = SimpleNamespace(models=[SimpleNamespace(name="default-model")], tool_search=SimpleNamespace(enabled=True, auto_promote_top_k=3)) captured: dict[str, object] = {} def fake_build_subagent_runtime_middlewares(**kwargs): diff --git a/backend/tests/test_tool_error_handling_middleware.py b/backend/tests/test_tool_error_handling_middleware.py index 2fb15efe8..d1675cf1f 100644 --- a/backend/tests/test_tool_error_handling_middleware.py +++ b/backend/tests/test_tool_error_handling_middleware.py @@ -548,6 +548,23 @@ def test_subagent_runtime_middlewares_attach_deferred_filter_when_setup_has_name assert filter_idx < safety_idx +def test_subagent_runtime_middlewares_place_mcp_routing_before_deferred_filter(monkeypatch): + from deerflow.agents.middlewares.deferred_tool_filter_middleware import DeferredToolFilterMiddleware + from deerflow.agents.middlewares.mcp_routing_middleware import McpRoutingMiddleware + from deerflow.tools.builtins.tool_search import DeferredToolSetup + + app_config = _make_app_config() + _stub_runtime_middleware_imports(monkeypatch) + routing = McpRoutingMiddleware({"mcp_thing": {"priority": 100, "keywords": ["orders"]}}, "hash123", 3) + setup = DeferredToolSetup(object(), frozenset({"mcp_thing"}), "hash123") + + middlewares = build_subagent_runtime_middlewares(app_config=app_config, deferred_setup=setup, mcp_routing_middleware=routing) + + routing_idx = next(i for i, middleware in enumerate(middlewares) if isinstance(middleware, McpRoutingMiddleware)) + filter_idx = next(i for i, middleware in enumerate(middlewares) if isinstance(middleware, DeferredToolFilterMiddleware)) + assert routing_idx < filter_idx + + def test_subagent_runtime_middlewares_skip_deferred_filter_without_names(monkeypatch): """No deferred setup (disabled / no MCP tool) -> no DeferredToolFilterMiddleware.""" from deerflow.agents.middlewares.deferred_tool_filter_middleware import DeferredToolFilterMiddleware diff --git a/backend/tests/test_tool_search.py b/backend/tests/test_tool_search.py index 3722cf3f9..91898e48b 100644 --- a/backend/tests/test_tool_search.py +++ b/backend/tests/test_tool_search.py @@ -15,15 +15,60 @@ from deerflow.tools.builtins.tool_search import get_deferred_tools_prompt_sectio class TestToolSearchConfig: def test_default_disabled(self): assert ToolSearchConfig().enabled is False + assert ToolSearchConfig().auto_promote_top_k == 3 def test_enabled(self): assert ToolSearchConfig(enabled=True).enabled is True + def test_auto_promote_top_k_is_clamped(self): + assert ToolSearchConfig(auto_promote_top_k=0).auto_promote_top_k == 1 + assert ToolSearchConfig(auto_promote_top_k=99).auto_promote_top_k == 5 + def test_load_from_dict(self): - assert load_tool_search_config_from_dict({"enabled": True}).enabled is True + loaded = load_tool_search_config_from_dict({"enabled": True, "auto_promote_top_k": 4}) + assert loaded.enabled is True + assert loaded.auto_promote_top_k == 4 def test_load_from_empty_dict(self): assert load_tool_search_config_from_dict({}).enabled is False + assert load_tool_search_config_from_dict({}).auto_promote_top_k == 3 + + +class TestConfigExampleToolSearchSection: + """Guard the documented ``tool_search`` block in config.example.yaml. + + The example file is the first-run template (``cp config.example.yaml + config.yaml``); a malformed indentation there breaks the whole file for + every downstream consumer, so pin that it parses and carries the PR2 field. + """ + + def _load_example(self): + import os + + import yaml + + example_path = os.path.join(os.path.dirname(__file__), "..", "..", "config.example.yaml") + if not os.path.exists(example_path): + return None + with open(example_path, encoding="utf-8") as f: + return yaml.safe_load(f) + + def test_config_example_parses(self): + # A raw yaml.safe_load raises on malformed indentation; asserting a + # dict result pins that the whole template stays parseable. + data = self._load_example() + if data is None: + return + assert isinstance(data, dict) + + def test_config_example_tool_search_block(self): + data = self._load_example() + if data is None: + return + tool_search = data.get("tool_search") + assert isinstance(tool_search, dict) + assert tool_search.get("enabled") is False + assert tool_search.get("auto_promote_top_k") == 3 class TestDeferredToolsPromptSection: diff --git a/config.example.yaml b/config.example.yaml index 9b86608d2..20f05a590 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -872,6 +872,10 @@ tools: tool_search: enabled: false + # When tool_search is enabled, PR1 MCP routing metadata can auto-promote + # matching deferred MCP tool schemas before a model call. This is the maximum + # number of matched schemas promoted per model call. Valid range: 1..5. + auto_promote_top_k: 3 # ============================================================================ # Tool Output Budget Protection