mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-06-10 17:35:57 +00:00
AochenShen99
d9f4724950
* feat(tool-search): add hash-scoped promoted state to ThreadState * feat(tool-search): add immutable DeferredToolCatalog with stable hash * feat(tool-search): add build_deferred_tool_setup + Command-writing tool_search * refactor(tool-search): replace deferred-tool ContextVar with closures + graph state (#3272) Build the deferred catalog + tool_search tool per agent from the policy-filtered tool list (after skill allowed-tools), pass deferred_names + catalog_hash explicitly to DeferredToolFilterMiddleware and the prompt, and record promotions in ThreadState.promoted (scoped by catalog_hash) via a Command-returning tool_search. Removes DeferredToolRegistry and the _registry_var ContextVar so deferral no longer depends on build/execute sharing an async context. MCP tools are tagged with metadata[deerflow_mcp]; client.py assembles deferral the same way. Catalog is built AFTER tool-policy filtering (no policy-excluded tool can leak via tool_search) and assembly is fail-closed. Migrate tests off the deleted registry APIs; delete the obsolete ContextVar-based #2884 regression (re-covered by state-based tests in a follow-up). * test(tool-search): lock tool_search promotion into next model turn via graph state * test(tool-search): cross-context, policy-leak, fail-closed, #2884 isolation regressions * test(tool-search): align real-LLM e2e with closure-based deferred setup * docs: update DeferredToolFilterMiddleware description for closure+state design * style(tests): drop unused import in test_deferred_setup (ruff) * test(tool-search): harden merge_promoted + replace tautological catalog test From independent code review: - merge_promoted: use existing.get("catalog_hash") so a forward-incompatible or externally-injected persisted promoted dict triggers a replace instead of a KeyError crash; add regression test for the malformed-existing case. - test_deferred_catalog: replace the `== [] or True` tautology (a test that could never fail) with a deterministic invalid-regex->literal-fallback check (positive match on calc + negative empty match). - DeferredToolCatalog: comment why frozen-without-slots is required for the cached_property hash/names fields (adding slots=True would break them). * fix(tool-search): read tool_search.enabled from self._app_config in client DeerFlowClient._ensure_agent called get_app_config() directly to read tool_search.enabled, but the client already resolves and stores its config as self._app_config at construction (and uses it everywhere else). The bare call re-resolves config from disk at agent-build time, which raises FileNotFoundError in environments without a config.yaml (CI) — test_client.py's fixture only patches get_app_config during __init__, so the later call hit the real loader. Use self._app_config, matching the rest of the client. * test(tool-search): lock tool_search post-policy append ordering tool_search is appended after skill-allowlist filtering, so the allowlist can no longer deny it by name. Lock the intended contract: it only appears when allowed MCP tools survive the filter, and its catalog (derived from the already policy-filtered list) can never expose a denied tool. Addresses the ordering observation from the Copilot review on #3342.
78 lines
2.9 KiB
Python
78 lines
2.9 KiB
Python
"""End-to-end: tool_search promotes a deferred tool into the next model turn.
|
|
|
|
Locks the full loop through a real ``create_agent`` graph:
|
|
turn 1 -> deferred MCP tools hidden from bind_tools; model calls tool_search
|
|
ToolNode-> tool_search returns Command(update={"promoted": {...}}) -> state
|
|
turn 2 -> middleware reads state["promoted"] (hash-scoped) -> the searched
|
|
tool's schema is now bound; un-searched deferred tools stay hidden
|
|
|
|
This is the behavior #3272's redesign depends on (no ContextVar): promotion
|
|
flows through graph state, so it works regardless of build/execute context.
|
|
"""
|
|
|
|
import asyncio
|
|
|
|
from langchain.agents import create_agent
|
|
from langchain_core.language_models.fake_chat_models import GenericFakeChatModel
|
|
from langchain_core.messages import AIMessage, HumanMessage
|
|
from langchain_core.tools import tool as as_tool
|
|
|
|
from deerflow.agents.middlewares.deferred_tool_filter_middleware import DeferredToolFilterMiddleware
|
|
from deerflow.agents.thread_state import ThreadState
|
|
from deerflow.tools.builtins.tool_search import build_deferred_tool_setup
|
|
|
|
|
|
@as_tool
|
|
def active_tool(x: str) -> str:
|
|
"An always-active tool."
|
|
return x
|
|
|
|
|
|
@as_tool
|
|
def mcp_calc(expression: str) -> str:
|
|
"Evaluate arithmetic."
|
|
return expression
|
|
|
|
|
|
@as_tool
|
|
def mcp_other(x: str) -> str:
|
|
"Another deferred MCP tool."
|
|
return x
|
|
|
|
|
|
def _tag(t):
|
|
t.metadata = {**(t.metadata or {}), "deerflow_mcp": True}
|
|
return t
|
|
|
|
|
|
def test_tool_search_promotes_into_next_turn():
|
|
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
|
|
|
|
setup = build_deferred_tool_setup([active_tool, _tag(mcp_calc), _tag(mcp_other)], enabled=True)
|
|
turn1 = AIMessage(content="", tool_calls=[{"name": "tool_search", "args": {"query": "select:mcp_calc"}, "id": "c1", "type": "tool_call"}])
|
|
turn2 = AIMessage(content="done")
|
|
model = RecordingModel(messages=iter([turn1, turn2]))
|
|
|
|
graph = create_agent(
|
|
model=model,
|
|
tools=[active_tool, mcp_calc, mcp_other, setup.tool_search_tool],
|
|
middleware=[DeferredToolFilterMiddleware(setup.deferred_names, setup.catalog_hash)],
|
|
state_schema=ThreadState,
|
|
)
|
|
|
|
result = asyncio.run(graph.ainvoke({"messages": [HumanMessage(content="use the deferred calculator")]}))
|
|
|
|
assert len(bound) >= 2, f"expected >=2 model binds, got {bound}"
|
|
# Turn 1: both deferred MCP tools hidden.
|
|
assert "mcp_calc" not in bound[0] and "mcp_other" not in bound[0]
|
|
# Turn 2: the searched tool is promoted (visible); the un-searched one stays hidden.
|
|
assert "mcp_calc" in bound[1]
|
|
assert "mcp_other" not in bound[1]
|
|
# Promotion recorded in graph state, scoped by catalog hash.
|
|
assert result["promoted"] == {"catalog_hash": setup.catalog_hash, "names": ["mcp_calc"]}
|