mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
fix(harness): scope deferred tool promotion to user turn (#6977)
* fix(harness): scope deferred tool promotion to user turn Signed-off-by: CorgiBoyG <111257566+CorgiBoyG@users.noreply.github.com> * fix(harness): preserve legacy promotions during routing resume --------- Signed-off-by: CorgiBoyG <111257566+CorgiBoyG@users.noreply.github.com> Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
RongJie G
GitHub
parent
07f0d74f70
commit
1098fb4c60
@ -939,6 +939,8 @@ When deferred MCP tool discovery is enabled, `tool_search` supports opt-in liter
|
||||
|
||||
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.
|
||||
|
||||
Deferred MCP schema promotions last for the current user turn. They remain active across that turn's model/tool loop and checkpoint resume, then become deferred again when a new user message starts the next turn; the next turn must discover them again through `tool_search` or match their routing hints. Resuming a legacy in-flight checkpoint preserves previously searched tools even when routing hints also match that turn.
|
||||
|
||||
Deferred `tool_search` regex queries share a 100 ms execution budget across the catalog, including `+required` ranking. Patterns are limited to 256 characters and searchable names/descriptions to 65,536 characters each. Queries that exceed a limit return a retry hint without promoting partial results; use a simpler pattern, `keywords:`, or `select:` with exact names. Oversized catalog fields are reported with the tool name, field and length, and logged with the MCP server name when available; use `keywords:` or `select:` because simplifying the pattern cannot shrink a catalog field. Exact selection remains uncapped and bypasses regex limits; the literal `keywords:` mode also bypasses regex limits. Invalid regex syntax still falls back to a literal substring match.
|
||||
|
||||
OpenViking users can register the official Streamable HTTP endpoint at `/mcp`
|
||||
|
||||
@ -28,7 +28,7 @@ same resulting prompt. Empty overlays preserve bytes.
|
||||
|
||||
**ThreadState** (`packages/harness/deerflow/agents/thread_state.py`):
|
||||
- Extends `AgentState` with: `sandbox`, `thread_data`, `title`, `artifacts`, `todos`, `uploaded_files`, `viewed_images`, `goal`, `promoted`, `delegations`, `skill_context`, `summary_text`
|
||||
- Uses custom reducers: `merge_artifacts` (deduplicate), `merge_viewed_images` (merge/clear), `merge_goal` (preserve the active goal across ordinary state updates unless the goal writer replaces it), `merge_promoted` (catalog-hash-scoped deferred tool promotions), `merge_delegations` (append task delegation entries, same (run,id) wins, terminal status never downgraded, capped to the most recent entries), and `merge_skill_context` (dedupe active-skill references by path, keep the most recently read entries; entries store a name/path/description reference, not the SKILL.md body). `summary_text` is a LastValue channel updated by summarization and projected into model requests as durable context data instead of being stored as a `messages` item.
|
||||
- Uses custom reducers: `merge_artifacts` (deduplicate), `merge_viewed_images` (merge/clear), `merge_goal` (preserve the active goal across ordinary state updates unless the goal writer replaces it), `merge_promoted` (union deferred tools within one catalog/user turn; replace on scope change; hash-only tool updates inherit scope), `merge_delegations` (append task delegation entries, same (run,id) wins, terminal status never downgraded, capped to the most recent entries), and `merge_skill_context` (dedupe active-skill references by path, keep the most recently read entries; entries store a name/path/description reference, not the SKILL.md body). `summary_text` is a LastValue channel updated by summarization and projected into model requests as durable context data instead of being stored as a `messages` item.
|
||||
- Delta-mode `merge_message_writes` normalizes the current message state once,
|
||||
then folds normalized writes in order with message-ID position indexes and
|
||||
deferred tombstone compaction. It preserves public `add_messages` behavior,
|
||||
|
||||
@ -120,9 +120,9 @@ Before changing a later authorization phase, read the [authorization RFC](../../
|
||||
24. **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.
|
||||
25. **MemoryMiddleware** - Queues conversations for async memory update (filters to user + final AI responses); captures the runtime-resolved user so standalone LangGraph Server reads and writes stay in the same bucket
|
||||
26. **ViewImageMiddleware** - *(vision models only)* Enforces the shared vision cap on remote rereads (see [download contract](../../sandbox/AGENTS.md)). Rebuilds hidden base64 context in `ModelRequest.messages`, keeping lightweight `viewed_images` in checkpoints. Caps each request at four recent images and 28 MiB base64; constructor limits may tighten caps. Skips over-budget reads. Re-views refresh reducer recency. Validated blob refs precede compatibility paths. `_image_injection_plan` removes stale server-marked context, including old `before_model`/`after_model` payloads, then shares eligibility for sync/async injection. Removal requires both the reserved ID prefix and marker, preserving user and unmarked legacy messages. Both paths recheck `sandbox:execute` before I/O. Host copies require `actual_path` to match the current user/thread's canonical virtual path, even for old checkpoints or replacement sandboxes; absent/mismatched identity skips reads. Only the async wrapper passes `authorization_checked=True` after its gate to avoid duplicate checks in its worker.
|
||||
27. **McpRoutingMiddleware** - Auto-promotes deferred schemas matching the latest real user message before `DeferredToolFilterMiddleware`, without executing tools. New names emit `middleware:tool_promotion` with `source=routing_hint`; repeat passes emit nothing
|
||||
27. **McpRoutingMiddleware** - Auto-promotes deferred schemas matching the latest real user message before `DeferredToolFilterMiddleware`, without executing tools. `message_utils` owns turn IDs and legacy resume detection. Routing preserves same-turn search names during migration. New names emit `middleware:tool_promotion` with `source=routing_hint`; repeat passes emit nothing
|
||||
28. **DeferredToolPromotionAuditMiddleware** - Observes final `tool_search` `Command`s; keep it outer of `SkillToolPolicyMiddleware` so denied names are excluded. It atomically claims new names per lead run or subagent execution to dedupe parallel searches, derives subagent attribution from the server-installed recorder, returns the original `Command`, and omits private payloads
|
||||
29. **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)
|
||||
29. **DeferredToolFilterMiddleware** - *(optional, if `tool_search.enabled`)* Hides MCP schemas until search/routing promotes them. Promotions union within the catalog/user-message scope and reset next turn. Legacy state without `turn_id` survives only when an AI response follows the latest user message (in-flight resume).
|
||||
30. **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. The subagent builder places its date-only context middleware immediately before this coalescer, so the built-in subagent prompt and hidden date reminder still reach providers as one leading system block
|
||||
31. **SubagentLimitMiddleware** - `subagent_enabled`: caps ordinary `task` calls per response (`max_concurrent_subagents`, resolved before construction against startup `subagent_runtime.max_running` and 1-64) and per run (`max_total_subagents` override or `subagents.max_total_per_run`, default 6, clamped 1-50). Count current-run durable ledger entries, tagged with `run_id` on capture: repeated planning shares the cap; later user runs get a fresh budget. Durable `batch_task` has separate persisted total/live/running limits and no ordinary ledger entries. When exhausted, strip remaining `task` calls, force `finish_reason="stop"`, and append a visible limit note so the agent synthesizes existing results.
|
||||
Capture/retain delegations by `(run_id, tool_call_id)`. Later-run replies cannot close older orphans; missing run IDs match conservatively.
|
||||
|
||||
@ -8,21 +8,24 @@ calls to tools that have not been promoted yet.
|
||||
|
||||
The deferred name set and the catalog hash are injected at construction time
|
||||
(no ContextVar). Promotion state is read from graph state (``state["promoted"]``),
|
||||
scoped by catalog hash so a stale persisted promotion cannot expose a renamed
|
||||
or drifted tool.
|
||||
scoped by catalog hash and genuine user-message ID so a stale promotion cannot
|
||||
expose a renamed tool or remain active in a later user turn.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import override
|
||||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from typing import Any, override
|
||||
|
||||
from langchain.agents import AgentState
|
||||
from langchain.agents.middleware import AgentMiddleware
|
||||
from langchain.agents.middleware.types import ModelCallResult, ModelRequest, ModelResponse
|
||||
from langchain_core.messages import ToolMessage
|
||||
from langgraph.prebuilt.tool_node import ToolCallRequest
|
||||
from langgraph.runtime import Runtime
|
||||
from langgraph.types import Command
|
||||
|
||||
from deerflow.agents.middlewares.message_utils import genuine_user_turn_id, has_ai_response_after_latest_user_message
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@ -31,7 +34,7 @@ class DeferredToolFilterMiddleware(AgentMiddleware[AgentState]):
|
||||
|
||||
ToolNode still holds all tools (including deferred) for execution routing,
|
||||
but the LLM only sees active tool schemas plus tools that have already been
|
||||
promoted (recorded in ``state["promoted"]`` under the current catalog hash).
|
||||
promoted under the current catalog hash and genuine user turn.
|
||||
"""
|
||||
|
||||
def __init__(self, deferred_names: frozenset[str], catalog_hash: str | None):
|
||||
@ -43,16 +46,48 @@ class DeferredToolFilterMiddleware(AgentMiddleware[AgentState]):
|
||||
return {
|
||||
"deferred_names": sorted(self._deferred),
|
||||
"catalog_hash": self._catalog_hash,
|
||||
"promotion_scope": "graph_state_catalog_hash",
|
||||
"promotion_scope": "graph_state_catalog_hash_user_turn",
|
||||
}
|
||||
|
||||
def _promoted(self, state) -> set[str]:
|
||||
def _promotion_is_current(self, state: Mapping[str, Any] | None, promoted: Mapping[str, Any]) -> bool:
|
||||
if promoted.get("catalog_hash") != self._catalog_hash:
|
||||
return False
|
||||
promoted_turn_id = promoted.get("turn_id")
|
||||
if promoted_turn_id is None:
|
||||
return has_ai_response_after_latest_user_message(list((state or {}).get("messages") or []))
|
||||
current_turn_id = genuine_user_turn_id(list((state or {}).get("messages") or []))
|
||||
return current_turn_id is not None and promoted_turn_id == current_turn_id
|
||||
|
||||
def _promoted(self, state: Mapping[str, Any] | None) -> set[str]:
|
||||
promoted = (state or {}).get("promoted")
|
||||
if promoted and promoted.get("catalog_hash") == self._catalog_hash:
|
||||
return set(promoted.get("names") or [])
|
||||
if not isinstance(promoted, Mapping) or not self._promotion_is_current(state, promoted):
|
||||
return set()
|
||||
names = promoted.get("names")
|
||||
if isinstance(names, Sequence) and not isinstance(names, (str, bytes)):
|
||||
return {name for name in names if isinstance(name, str)}
|
||||
return set()
|
||||
|
||||
def _hidden(self, state) -> set[str]:
|
||||
def _turn_state_update(self, state: Mapping[str, Any] | None) -> dict[str, Any] | None:
|
||||
turn_id = genuine_user_turn_id(list((state or {}).get("messages") or []))
|
||||
if not self._deferred or not self._catalog_hash or turn_id is None:
|
||||
return None
|
||||
promoted = (state or {}).get("promoted")
|
||||
if isinstance(promoted, Mapping) and self._promotion_is_current(state, promoted):
|
||||
if promoted.get("turn_id") == turn_id:
|
||||
return None
|
||||
names = promoted.get("names")
|
||||
migrated_names = [name for name in names if isinstance(name, str)] if isinstance(names, Sequence) and not isinstance(names, (str, bytes)) else []
|
||||
else:
|
||||
migrated_names = []
|
||||
return {
|
||||
"promoted": {
|
||||
"catalog_hash": self._catalog_hash,
|
||||
"names": migrated_names,
|
||||
"turn_id": turn_id,
|
||||
}
|
||||
}
|
||||
|
||||
def _hidden(self, state: Mapping[str, Any] | None) -> set[str]:
|
||||
return set(self._deferred) - self._promoted(state)
|
||||
|
||||
def _filter_tools(self, request: ModelRequest) -> ModelRequest:
|
||||
@ -80,6 +115,14 @@ class DeferredToolFilterMiddleware(AgentMiddleware[AgentState]):
|
||||
status="error",
|
||||
)
|
||||
|
||||
@override
|
||||
def before_model(self, state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
|
||||
return self._turn_state_update(state)
|
||||
|
||||
@override
|
||||
async def abefore_model(self, state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
|
||||
return self._turn_state_update(state)
|
||||
|
||||
@override
|
||||
def wrap_model_call(
|
||||
self,
|
||||
|
||||
@ -11,7 +11,7 @@ from langchain.agents.middleware import AgentMiddleware
|
||||
from langchain_core.messages import HumanMessage
|
||||
from langgraph.runtime import Runtime
|
||||
|
||||
from deerflow.agents.middlewares.message_utils import is_genuine_user_message
|
||||
from deerflow.agents.middlewares.message_utils import genuine_user_turn_id, has_ai_response_after_latest_user_message, latest_genuine_user_message
|
||||
from deerflow.agents.middlewares.tool_promotion_audit_middleware import record_tool_promotion
|
||||
from deerflow.config.tool_search_config import clamp_auto_promote_top_k
|
||||
from deerflow.utils.messages import get_original_user_content_text
|
||||
@ -80,10 +80,7 @@ class McpRoutingMiddleware(AgentMiddleware[AgentState]):
|
||||
request, so ``is_genuine_user_message`` — not ``is_real_user_message``,
|
||||
which drops every ``hide_from_ui`` message — decides this.
|
||||
"""
|
||||
for message in reversed(messages):
|
||||
if is_genuine_user_message(message):
|
||||
return message
|
||||
return None
|
||||
return latest_genuine_user_message(messages)
|
||||
|
||||
def _matched_names(self, state: Mapping[str, Any] | None) -> list[str]:
|
||||
if not self._catalog_hash or not self._routing_index:
|
||||
@ -113,8 +110,13 @@ class McpRoutingMiddleware(AgentMiddleware[AgentState]):
|
||||
names = self._matched_names(state)
|
||||
if not names:
|
||||
return None
|
||||
messages = list((state or {}).get("messages") or [])
|
||||
turn_id = genuine_user_turn_id(messages)
|
||||
promoted = (state or {}).get("promoted")
|
||||
raw_promoted_names = promoted.get("names") if isinstance(promoted, Mapping) and promoted.get("catalog_hash") == self._catalog_hash else None
|
||||
same_catalog = isinstance(promoted, Mapping) and promoted.get("catalog_hash") == self._catalog_hash
|
||||
migrate_legacy = same_catalog and turn_id is not None and promoted.get("turn_id") is None and has_ai_response_after_latest_user_message(messages)
|
||||
same_scope = same_catalog and (turn_id is None or promoted.get("turn_id") == turn_id or migrate_legacy)
|
||||
raw_promoted_names = promoted.get("names") if same_scope else None
|
||||
already_promoted = {name for name in raw_promoted_names if isinstance(name, str)} if isinstance(raw_promoted_names, Sequence) and not isinstance(raw_promoted_names, (str, bytes)) else set()
|
||||
record_tool_promotion(
|
||||
runtime,
|
||||
@ -129,12 +131,15 @@ class McpRoutingMiddleware(AgentMiddleware[AgentState]):
|
||||
(self._catalog_hash or "")[:8],
|
||||
names,
|
||||
)
|
||||
return {
|
||||
"promoted": {
|
||||
"catalog_hash": self._catalog_hash,
|
||||
"names": names,
|
||||
}
|
||||
# Adding turn_id replaces the legacy envelope in merge_promoted, so
|
||||
# carry its same-turn search names through before the filter runs.
|
||||
promoted_update = {
|
||||
"catalog_hash": self._catalog_hash,
|
||||
"names": sorted(already_promoted | set(names)) if migrate_legacy else names,
|
||||
}
|
||||
if turn_id is not None:
|
||||
promoted_update["turn_id"] = turn_id
|
||||
return {"promoted": promoted_update}
|
||||
|
||||
@override
|
||||
def before_model(self, state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
|
||||
|
||||
@ -2,7 +2,7 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from langchain_core.messages import HumanMessage, SystemMessage
|
||||
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
|
||||
|
||||
from deerflow.agents.human_input import read_human_input_response
|
||||
from deerflow.utils.messages import UNTRUSTED_INPUT_KEY
|
||||
@ -25,6 +25,31 @@ def is_genuine_user_message(message: object) -> bool:
|
||||
return True
|
||||
|
||||
|
||||
def latest_genuine_user_message(messages: list[object]) -> HumanMessage | None:
|
||||
"""Return the latest user-authored message that defines the active turn."""
|
||||
return next((message for message in reversed(messages) if is_genuine_user_message(message)), None)
|
||||
|
||||
|
||||
def genuine_user_turn_id(messages: list[object]) -> str | None:
|
||||
"""Return the non-empty message ID shared by turn-scoped middlewares."""
|
||||
target = latest_genuine_user_message(messages)
|
||||
return target.id if target is not None and isinstance(target.id, str) and target.id else None
|
||||
|
||||
|
||||
def has_ai_response_after_latest_user_message(messages: list[object]) -> bool:
|
||||
"""Recognize legacy in-flight turns without carrying names into a new turn.
|
||||
|
||||
An AI response after the latest genuine user message lets routing and the
|
||||
deferred filter migrate pre-turn-scope checkpoints consistently. A newly
|
||||
appended user message has no later AI response and expires legacy names.
|
||||
"""
|
||||
target = latest_genuine_user_message(messages)
|
||||
if target is None:
|
||||
return False
|
||||
target_index = next(index for index in range(len(messages) - 1, -1, -1) if messages[index] is target)
|
||||
return any(isinstance(message, AIMessage) for message in messages[target_index + 1 :])
|
||||
|
||||
|
||||
def requires_input_sanitization(message: object) -> bool:
|
||||
"""Return True for message content the input guardrail must neutralize.
|
||||
|
||||
|
||||
@ -17,6 +17,7 @@ from deerflow.agents.middlewares.audit_context import (
|
||||
TOOL_PROMOTION_RECORDER_CONTEXT_KEY,
|
||||
resolve_audit_recorder,
|
||||
)
|
||||
from deerflow.agents.middlewares.message_utils import genuine_user_turn_id
|
||||
from deerflow.runtime.events.catalog import MIDDLEWARE_TOOL_PROMOTION_TAG
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@ -93,6 +94,9 @@ class DeferredToolPromotionAuditMiddleware(AgentMiddleware[AgentState]):
|
||||
promoted = (state or {}).get("promoted")
|
||||
if not isinstance(promoted, Mapping) or promoted.get("catalog_hash") != self._catalog_hash:
|
||||
return set()
|
||||
turn_id = genuine_user_turn_id(list((state or {}).get("messages") or []))
|
||||
if turn_id is not None and promoted.get("turn_id") != turn_id:
|
||||
return set()
|
||||
names = promoted.get("names")
|
||||
if not isinstance(names, list) or not all(isinstance(name, str) for name in names):
|
||||
return set()
|
||||
|
||||
@ -159,27 +159,35 @@ def merge_goal(existing: GoalState | None, new: GoalState | None) -> GoalState |
|
||||
class PromotedTools(TypedDict):
|
||||
catalog_hash: str
|
||||
names: list[str]
|
||||
turn_id: NotRequired[str]
|
||||
|
||||
|
||||
def merge_promoted(existing: PromotedTools | None, new: PromotedTools | None) -> PromotedTools | None:
|
||||
"""Reducer for deferred-tool promotions, scoped by catalog hash.
|
||||
"""Reducer for deferred-tool promotions, scoped by catalog and user turn.
|
||||
|
||||
- new None/empty -> preserve existing (node didn't touch promotions).
|
||||
- catalog_hash changed -> replace wholesale, dropping stale names (prevents a
|
||||
persisted bare name from exposing a different tool after catalog drift).
|
||||
- same catalog_hash -> union names, dedupe, preserve order.
|
||||
- catalog_hash or explicit turn_id changed -> replace wholesale, dropping
|
||||
stale names.
|
||||
- same scope -> union names, dedupe, preserve order.
|
||||
- a missing new turn_id inherits the existing scope. ``tool_search`` writes
|
||||
this shape after the filter establishes the current turn before the model.
|
||||
"""
|
||||
if not new:
|
||||
return existing
|
||||
if existing is None or existing.get("catalog_hash") != new["catalog_hash"]:
|
||||
return {
|
||||
"catalog_hash": new["catalog_hash"],
|
||||
"names": list(dict.fromkeys(new["names"])),
|
||||
}
|
||||
return {
|
||||
"catalog_hash": existing["catalog_hash"],
|
||||
"names": list(dict.fromkeys(existing["names"] + new["names"])),
|
||||
new_turn_id = new.get("turn_id")
|
||||
replace = existing is None or existing.get("catalog_hash") != new["catalog_hash"] or (new_turn_id is not None and existing.get("turn_id") != new_turn_id)
|
||||
existing_names = existing.get("names") if existing is not None else None
|
||||
if not isinstance(existing_names, list) or not all(isinstance(name, str) for name in existing_names):
|
||||
existing_names = []
|
||||
names = new["names"] if replace else existing_names + new["names"]
|
||||
merged: PromotedTools = {
|
||||
"catalog_hash": new["catalog_hash"],
|
||||
"names": list(dict.fromkeys(names)),
|
||||
}
|
||||
turn_id = new_turn_id if new_turn_id is not None else (existing or {}).get("turn_id")
|
||||
if turn_id is not None:
|
||||
merged["turn_id"] = turn_id
|
||||
return merged
|
||||
|
||||
|
||||
TERMINAL_STATUSES: frozenset[str] = frozenset(SUBAGENT_STATUS_VALUES)
|
||||
|
||||
@ -1,5 +1,6 @@
|
||||
"""Tests for DeferredToolFilterMiddleware (closure deferred-set + state promotion)."""
|
||||
|
||||
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
|
||||
@ -45,7 +46,13 @@ def test_hides_all_deferred_when_no_promotion():
|
||||
|
||||
|
||||
def test_promoted_under_matching_hash_passes_through():
|
||||
req = _Req([mcp_a, mcp_b, active_c], {"promoted": {"catalog_hash": "h1", "names": ["mcp_a"]}})
|
||||
req = _Req(
|
||||
[mcp_a, mcp_b, active_c],
|
||||
{
|
||||
"messages": [HumanMessage(content="current", id="user-1")],
|
||||
"promoted": {"catalog_hash": "h1", "names": ["mcp_a"], "turn_id": "user-1"},
|
||||
},
|
||||
)
|
||||
out = _mw()._filter_tools(req)
|
||||
assert {t.name for t in out.overridden} == {"mcp_a", "active_c"}
|
||||
|
||||
@ -56,6 +63,83 @@ def test_promotion_ignored_when_hash_mismatch():
|
||||
assert [t.name for t in out.overridden] == ["active_c"]
|
||||
|
||||
|
||||
def test_promotion_ignored_when_user_turn_changes():
|
||||
state = {
|
||||
"messages": [HumanMessage(content="next", id="user-2")],
|
||||
"promoted": {"catalog_hash": "h1", "names": ["mcp_a"], "turn_id": "user-1"},
|
||||
}
|
||||
req = _Req([mcp_a, active_c], state)
|
||||
|
||||
out = _mw()._filter_tools(req)
|
||||
|
||||
assert [t.name for t in out.overridden] == ["active_c"]
|
||||
|
||||
|
||||
def test_legacy_promotion_without_user_turn_fails_closed():
|
||||
req = _Req([mcp_a, active_c], {"promoted": {"catalog_hash": "h1", "names": ["mcp_a"]}})
|
||||
|
||||
out = _mw()._filter_tools(req)
|
||||
|
||||
assert [t.name for t in out.overridden] == ["active_c"]
|
||||
|
||||
|
||||
def test_before_model_establishes_empty_current_turn_scope():
|
||||
state = {"messages": [HumanMessage(content="first", id="user-1")]}
|
||||
|
||||
assert _mw().before_model(state, runtime=None) == {
|
||||
"promoted": {
|
||||
"catalog_hash": "h1",
|
||||
"names": [],
|
||||
"turn_id": "user-1",
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def test_before_model_clears_prior_turn_promotions():
|
||||
state = {
|
||||
"messages": [HumanMessage(content="first", id="user-1"), AIMessage(content="done"), HumanMessage(content="second", id="user-2")],
|
||||
"promoted": {"catalog_hash": "h1", "names": ["mcp_a"], "turn_id": "user-1"},
|
||||
}
|
||||
|
||||
assert _mw().before_model(state, runtime=None) == {
|
||||
"promoted": {
|
||||
"catalog_hash": "h1",
|
||||
"names": [],
|
||||
"turn_id": "user-2",
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def test_before_model_migrates_in_flight_legacy_checkpoint():
|
||||
state = {
|
||||
"messages": [HumanMessage(content="first", id="user-1"), AIMessage(content="", tool_calls=[])],
|
||||
"promoted": {"catalog_hash": "h1", "names": ["mcp_a"]},
|
||||
}
|
||||
|
||||
assert _mw().before_model(state, runtime=None) == {
|
||||
"promoted": {
|
||||
"catalog_hash": "h1",
|
||||
"names": ["mcp_a"],
|
||||
"turn_id": "user-1",
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def test_before_model_expires_legacy_promotion_after_new_user_message():
|
||||
state = {
|
||||
"messages": [HumanMessage(content="first", id="user-1"), AIMessage(content="done"), HumanMessage(content="second", id="user-2")],
|
||||
"promoted": {"catalog_hash": "h1", "names": ["mcp_a"]},
|
||||
}
|
||||
|
||||
assert _mw().before_model(state, runtime=None) == {
|
||||
"promoted": {
|
||||
"catalog_hash": "h1",
|
||||
"names": [],
|
||||
"turn_id": "user-2",
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def test_no_deferred_names_is_noop():
|
||||
req = _Req([active_c], {})
|
||||
out = DeferredToolFilterMiddleware(frozenset(), "h1")._filter_tools(req)
|
||||
@ -74,11 +158,27 @@ def test_blocked_message_for_unpromoted_deferred_call():
|
||||
def test_no_block_for_promoted_call():
|
||||
class _TCReq:
|
||||
tool_call = {"name": "mcp_a", "id": "tc1"}
|
||||
state = {"promoted": {"catalog_hash": "h1", "names": ["mcp_a"]}}
|
||||
state = {
|
||||
"messages": [HumanMessage(content="current", id="user-1")],
|
||||
"promoted": {"catalog_hash": "h1", "names": ["mcp_a"], "turn_id": "user-1"},
|
||||
}
|
||||
|
||||
assert _mw()._blocked_tool_message(_TCReq()) is None
|
||||
|
||||
|
||||
def test_blocks_tool_call_promoted_only_in_prior_user_turn():
|
||||
class _TCReq:
|
||||
tool_call = {"name": "mcp_a", "id": "tc1"}
|
||||
state = {
|
||||
"messages": [HumanMessage(content="next", id="user-2")],
|
||||
"promoted": {"catalog_hash": "h1", "names": ["mcp_a"], "turn_id": "user-1"},
|
||||
}
|
||||
|
||||
msg = _mw()._blocked_tool_message(_TCReq())
|
||||
|
||||
assert msg is not None and msg.status == "error"
|
||||
|
||||
|
||||
def test_no_block_for_non_deferred_call():
|
||||
class _TCReq:
|
||||
tool_call = {"name": "active_c", "id": "tc1"}
|
||||
|
||||
@ -12,12 +12,16 @@ flows through graph state, so it works regardless of build/execute context.
|
||||
|
||||
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
|
||||
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
|
||||
from langchain_core.tools import tool as as_tool
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.types import Overwrite
|
||||
|
||||
from deerflow.agents.middlewares.deferred_tool_filter_middleware import DeferredToolFilterMiddleware
|
||||
from deerflow.agents.middlewares.mcp_routing_middleware import McpRoutingMiddleware
|
||||
from deerflow.agents.thread_state import ThreadState
|
||||
from deerflow.tools.builtins.tool_search import build_deferred_tool_setup
|
||||
from deerflow.tools.mcp_metadata import tag_mcp_tool
|
||||
@ -61,7 +65,7 @@ def test_tool_search_promotes_into_next_turn():
|
||||
state_schema=ThreadState,
|
||||
)
|
||||
|
||||
result = asyncio.run(graph.ainvoke({"messages": [HumanMessage(content="use the deferred calculator")]}))
|
||||
result = asyncio.run(graph.ainvoke({"messages": [HumanMessage(content="use the deferred calculator", id="user-1")]}))
|
||||
|
||||
assert len(bound) >= 2, f"expected >=2 model binds, got {bound}"
|
||||
# Turn 1: both deferred MCP tools hidden.
|
||||
@ -69,5 +73,163 @@ def test_tool_search_promotes_into_next_turn():
|
||||
# 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"]}
|
||||
# Promotion recorded in graph state, scoped by catalog and user turn.
|
||||
assert result["promoted"] == {
|
||||
"catalog_hash": setup.catalog_hash,
|
||||
"names": ["mcp_calc"],
|
||||
"turn_id": "user-1",
|
||||
}
|
||||
|
||||
|
||||
def test_promoted_schema_is_deferred_again_on_next_user_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_tool(mcp_calc)], enabled=True)
|
||||
search = AIMessage(content="", tool_calls=[{"name": "tool_search", "args": {"query": "select:mcp_calc"}, "id": "search-1", "type": "tool_call"}])
|
||||
model = RecordingModel(messages=iter([search, AIMessage(content="first done"), AIMessage(content="second done")]))
|
||||
graph = create_agent(
|
||||
model=model,
|
||||
tools=[active_tool, mcp_calc, setup.tool_search_tool],
|
||||
middleware=[DeferredToolFilterMiddleware(setup.deferred_names, setup.catalog_hash)],
|
||||
state_schema=ThreadState,
|
||||
checkpointer=InMemorySaver(),
|
||||
)
|
||||
config = {"configurable": {"thread_id": "turn-scoped-promotion"}}
|
||||
|
||||
asyncio.run(graph.ainvoke({"messages": [HumanMessage(content="calculate", id="user-1")]}, config))
|
||||
asyncio.run(graph.ainvoke({"messages": [HumanMessage(content="now just say hello", id="user-2")]}, config))
|
||||
|
||||
assert "mcp_calc" not in bound[0]
|
||||
assert "mcp_calc" in bound[1]
|
||||
assert "mcp_calc" not in bound[2]
|
||||
|
||||
|
||||
def test_promoted_schema_survives_checkpoint_resume_within_same_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_tool(mcp_calc)], enabled=True)
|
||||
search = AIMessage(content="", tool_calls=[{"name": "tool_search", "args": {"query": "select:mcp_calc"}, "id": "search-1", "type": "tool_call"}])
|
||||
model = RecordingModel(messages=iter([search, AIMessage(content="done")]))
|
||||
graph = create_agent(
|
||||
model=model,
|
||||
tools=[active_tool, mcp_calc, setup.tool_search_tool],
|
||||
middleware=[DeferredToolFilterMiddleware(setup.deferred_names, setup.catalog_hash)],
|
||||
state_schema=ThreadState,
|
||||
checkpointer=InMemorySaver(),
|
||||
interrupt_after=["tools"],
|
||||
)
|
||||
config = {"configurable": {"thread_id": "resume-turn-scoped-promotion"}}
|
||||
|
||||
interrupted = asyncio.run(graph.ainvoke({"messages": [HumanMessage(content="calculate", id="user-1")]}, config))
|
||||
resumed = asyncio.run(graph.ainvoke(None, config))
|
||||
|
||||
assert interrupted["promoted"] == {
|
||||
"catalog_hash": setup.catalog_hash,
|
||||
"names": ["mcp_calc"],
|
||||
"turn_id": "user-1",
|
||||
}
|
||||
assert resumed["promoted"] == interrupted["promoted"]
|
||||
assert "mcp_calc" not in bound[0]
|
||||
assert "mcp_calc" in bound[1]
|
||||
|
||||
|
||||
def test_in_flight_legacy_checkpoint_promotion_migrates_on_resume():
|
||||
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_tool(mcp_calc)], enabled=True)
|
||||
search = AIMessage(content="", tool_calls=[{"name": "tool_search", "args": {"query": "select:mcp_calc"}, "id": "search-1", "type": "tool_call"}])
|
||||
model = RecordingModel(messages=iter([search, AIMessage(content="done")]))
|
||||
graph = create_agent(
|
||||
model=model,
|
||||
tools=[active_tool, mcp_calc, setup.tool_search_tool],
|
||||
middleware=[DeferredToolFilterMiddleware(setup.deferred_names, setup.catalog_hash)],
|
||||
state_schema=ThreadState,
|
||||
checkpointer=InMemorySaver(),
|
||||
interrupt_after=["tools"],
|
||||
)
|
||||
config = {"configurable": {"thread_id": "resume-legacy-promotion"}}
|
||||
|
||||
asyncio.run(graph.ainvoke({"messages": [HumanMessage(content="calculate", id="user-1")]}, config))
|
||||
graph.update_state(
|
||||
config,
|
||||
{
|
||||
"promoted": Overwrite(
|
||||
{
|
||||
"catalog_hash": setup.catalog_hash,
|
||||
"names": ["mcp_calc"],
|
||||
}
|
||||
)
|
||||
},
|
||||
as_node="tools",
|
||||
)
|
||||
resumed = asyncio.run(graph.ainvoke(None, config))
|
||||
|
||||
assert resumed["promoted"] == {
|
||||
"catalog_hash": setup.catalog_hash,
|
||||
"names": ["mcp_calc"],
|
||||
"turn_id": "user-1",
|
||||
}
|
||||
assert "mcp_calc" not in bound[0]
|
||||
assert "mcp_calc" in bound[1]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_async", [False, True], ids=["sync", "async"])
|
||||
def test_routing_preserves_searched_tools_on_legacy_checkpoint_resume(use_async):
|
||||
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_tool(mcp_calc), tag_mcp_tool(mcp_other)], enabled=True)
|
||||
search = AIMessage(content="", tool_calls=[{"name": "tool_search", "args": {"query": "select:mcp_calc"}, "id": "search-1", "type": "tool_call"}])
|
||||
calculate = AIMessage(content="", tool_calls=[{"name": "mcp_calc", "args": {"expression": "1+1"}, "id": "calc-1", "type": "tool_call"}])
|
||||
model = RecordingModel(messages=iter([search, calculate, AIMessage(content="done")]))
|
||||
graph = create_agent(
|
||||
model=model,
|
||||
tools=[active_tool, mcp_calc, mcp_other, setup.tool_search_tool],
|
||||
middleware=[
|
||||
McpRoutingMiddleware({"mcp_other": {"priority": 1, "keywords": ["calculate"]}}, setup.catalog_hash, 3),
|
||||
DeferredToolFilterMiddleware(setup.deferred_names, setup.catalog_hash),
|
||||
],
|
||||
state_schema=ThreadState,
|
||||
checkpointer=InMemorySaver(),
|
||||
interrupt_after=["tools"],
|
||||
)
|
||||
config = {"configurable": {"thread_id": "resume-legacy-routing-promotion"}}
|
||||
|
||||
def invoke(value):
|
||||
return asyncio.run(graph.ainvoke(value, config)) if use_async else graph.invoke(value, config)
|
||||
|
||||
invoke({"messages": [HumanMessage(content="calculate", id="user-1")]})
|
||||
graph.update_state(
|
||||
config,
|
||||
{"promoted": Overwrite({"catalog_hash": setup.catalog_hash, "names": ["mcp_calc"]})},
|
||||
as_node="tools",
|
||||
)
|
||||
resumed = invoke(None)
|
||||
|
||||
assert "mcp_calc" not in bound[0]
|
||||
assert {"mcp_calc", "mcp_other"} <= set(bound[1])
|
||||
assert resumed["promoted"]["turn_id"] == "user-1"
|
||||
assert set(resumed["promoted"]["names"]) == {"mcp_calc", "mcp_other"}
|
||||
result = next(message for message in resumed["messages"] if isinstance(message, ToolMessage) and message.tool_call_id == "calc-1")
|
||||
assert result.status == "success"
|
||||
assert result.content == "1+1"
|
||||
invoke(None)
|
||||
|
||||
@ -117,7 +117,14 @@ def test_subagent_reentry_does_not_touch_lead_state():
|
||||
class _Req:
|
||||
def __init__(self):
|
||||
self.tools = [active_tool, mcp_secret]
|
||||
self.state = {"promoted": {"catalog_hash": lead_setup.catalog_hash, "names": ["mcp_secret"]}}
|
||||
self.state = {
|
||||
"messages": [HumanMessage(content="use the secret tool", id="user-1")],
|
||||
"promoted": {
|
||||
"catalog_hash": lead_setup.catalog_hash,
|
||||
"names": ["mcp_secret"],
|
||||
"turn_id": "user-1",
|
||||
},
|
||||
}
|
||||
|
||||
def override(self, tools):
|
||||
self.tools = tools
|
||||
|
||||
@ -8,6 +8,7 @@ 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 langgraph.checkpoint.memory import InMemorySaver
|
||||
|
||||
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
|
||||
@ -325,11 +326,15 @@ def test_auto_promote_makes_schema_visible_in_same_model_cycle():
|
||||
state_schema=ThreadState,
|
||||
)
|
||||
|
||||
result = asyncio.run(graph.ainvoke({"messages": [HumanMessage(content="show orders")]}))
|
||||
result = asyncio.run(graph.ainvoke({"messages": [HumanMessage(content="show orders", id="user-1")]}))
|
||||
|
||||
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 result["promoted"] == {
|
||||
"catalog_hash": setup.catalog_hash,
|
||||
"names": ["postgres_query"],
|
||||
"turn_id": "user-1",
|
||||
}
|
||||
assert not any(isinstance(message, ToolMessage) for message in result["messages"])
|
||||
|
||||
|
||||
@ -359,10 +364,14 @@ def test_auto_promoted_tool_can_be_called_without_tool_search():
|
||||
state_schema=ThreadState,
|
||||
)
|
||||
|
||||
result = asyncio.run(graph.ainvoke({"messages": [HumanMessage(content="show orders")]}))
|
||||
result = asyncio.run(graph.ainvoke({"messages": [HumanMessage(content="show orders", id="user-1")]}))
|
||||
|
||||
assert "postgres_query" in bound[0]
|
||||
assert result["promoted"] == {"catalog_hash": setup.catalog_hash, "names": ["postgres_query"]}
|
||||
assert result["promoted"] == {
|
||||
"catalog_hash": setup.catalog_hash,
|
||||
"names": ["postgres_query"],
|
||||
"turn_id": "user-1",
|
||||
}
|
||||
tool_messages = [message for message in result["messages"] if isinstance(message, ToolMessage)]
|
||||
assert tool_messages
|
||||
assert tool_messages[0].name == "postgres_query"
|
||||
@ -393,11 +402,49 @@ def test_explicit_tool_search_merges_with_auto_promoted_names():
|
||||
state_schema=ThreadState,
|
||||
)
|
||||
|
||||
result = asyncio.run(graph.ainvoke({"messages": [HumanMessage(content="show orders")]}))
|
||||
result = asyncio.run(graph.ainvoke({"messages": [HumanMessage(content="show orders", id="user-1")]}))
|
||||
|
||||
assert result["promoted"] == {
|
||||
"catalog_hash": setup.catalog_hash,
|
||||
"names": ["postgres_query", "metrics_query"],
|
||||
"turn_id": "user-1",
|
||||
}
|
||||
|
||||
|
||||
def test_auto_promotion_is_recomputed_on_each_user_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
|
||||
|
||||
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
|
||||
model = RecordingModel(messages=iter([AIMessage(content="first done"), AIMessage(content="second done")]))
|
||||
graph = create_agent(
|
||||
model=model,
|
||||
tools=final_tools,
|
||||
middleware=[
|
||||
routing_middleware,
|
||||
DeferredToolFilterMiddleware(setup.deferred_names, setup.catalog_hash),
|
||||
],
|
||||
state_schema=ThreadState,
|
||||
checkpointer=InMemorySaver(),
|
||||
)
|
||||
config = {"configurable": {"thread_id": "routing-turn-scope"}}
|
||||
|
||||
asyncio.run(graph.ainvoke({"messages": [HumanMessage(content="show orders", id="user-1")]}, config))
|
||||
result = asyncio.run(graph.ainvoke({"messages": [HumanMessage(content="say hello", id="user-2")]}, config))
|
||||
|
||||
assert "postgres_query" in bound[0]
|
||||
assert "postgres_query" not in bound[1]
|
||||
assert result["promoted"] == {
|
||||
"catalog_hash": setup.catalog_hash,
|
||||
"names": [],
|
||||
"turn_id": "user-2",
|
||||
}
|
||||
|
||||
|
||||
@ -432,3 +479,55 @@ def test_privacy_no_trace_metadata_or_info_logs(caplog):
|
||||
assert state["metadata"] == {"trace": "existing"}
|
||||
assert "sensitive-keyword" not in caplog.text
|
||||
assert "secret_tool" not in caplog.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_async", [False, True], ids=["sync", "async"])
|
||||
@pytest.mark.parametrize(
|
||||
("case", "expected_names", "expected_audit"),
|
||||
[
|
||||
("in_flight", {"archive_lookup", "postgres_query", "metrics_query"}, ["metrics_query"]),
|
||||
("new_user", {"postgres_query", "metrics_query"}, ["metrics_query", "postgres_query"]),
|
||||
("new_catalog", {"postgres_query", "metrics_query"}, ["metrics_query", "postgres_query"]),
|
||||
("stale_turn", {"postgres_query", "metrics_query"}, ["metrics_query", "postgres_query"]),
|
||||
("no_ai", {"postgres_query", "metrics_query"}, ["metrics_query", "postgres_query"]),
|
||||
("malformed_names", {"postgres_query", "metrics_query"}, ["metrics_query", "postgres_query"]),
|
||||
],
|
||||
)
|
||||
def test_routing_legacy_migration_preserves_scope_and_audit(case, expected_names, expected_audit, use_async):
|
||||
middleware = McpRoutingMiddleware(
|
||||
{
|
||||
"postgres_query": {"priority": 100, "keywords": ["orders"]},
|
||||
"metrics_query": {"priority": 90, "keywords": ["orders"]},
|
||||
},
|
||||
"hash1",
|
||||
3,
|
||||
)
|
||||
messages = [HumanMessage(content="orders", id="user-1"), AIMessage(content="searching")]
|
||||
promoted = {"catalog_hash": "hash1", "names": ["archive_lookup", "postgres_query"]}
|
||||
if case == "new_user":
|
||||
messages.append(HumanMessage(content="orders again", id="user-2"))
|
||||
elif case == "new_catalog":
|
||||
promoted["catalog_hash"] = "old-hash"
|
||||
elif case == "stale_turn":
|
||||
promoted["turn_id"] = "old-user"
|
||||
elif case == "no_ai":
|
||||
messages.pop()
|
||||
elif case == "malformed_names":
|
||||
promoted["names"] = "archive_lookup"
|
||||
state = {"messages": messages, "promoted": promoted}
|
||||
recorder = _Recorder()
|
||||
runtime = SimpleNamespace(context={"__run_journal": recorder})
|
||||
|
||||
update = asyncio.run(middleware.abefore_model(state, runtime)) if use_async else middleware.before_model(state, runtime)
|
||||
merged = merge_promoted(promoted, update["promoted"])
|
||||
|
||||
assert set(merged["names"]) == expected_names
|
||||
assert merged["turn_id"] == ("user-2" if case == "new_user" else "user-1")
|
||||
assert recorder.calls[0]["changes"]["tool_names"] == expected_audit
|
||||
assert state["promoted"] is promoted
|
||||
assert "turn_id" not in promoted or promoted["turn_id"] == "old-user"
|
||||
|
||||
state["promoted"] = merged
|
||||
repeated = asyncio.run(middleware.abefore_model(state, runtime)) if use_async else middleware.before_model(state, runtime)
|
||||
assert set(merge_promoted(merged, repeated["promoted"])["names"]) == expected_names
|
||||
assert len(recorder.calls) == 1
|
||||
|
||||
@ -150,7 +150,7 @@ def test_passive_empty_skill_policy_preserves_deferred_mcp_discovery_and_calling
|
||||
)
|
||||
|
||||
result = graph.invoke(
|
||||
{"messages": [HumanMessage(content="use the configured MCP calculator")]},
|
||||
{"messages": [HumanMessage(content="use the configured MCP calculator", id="user-passive")]},
|
||||
context={},
|
||||
)
|
||||
|
||||
@ -160,7 +160,11 @@ def test_passive_empty_skill_policy_preserves_deferred_mcp_discovery_and_calling
|
||||
# skill applied a policy restriction.
|
||||
assert "denied_lookup" not in model.bound_tool_names[1]
|
||||
assert _CALC_CALLS == ["2 + 2"]
|
||||
assert result["promoted"] == {"catalog_hash": setup.catalog_hash, "names": ["calc"]}
|
||||
assert result["promoted"] == {
|
||||
"catalog_hash": setup.catalog_hash,
|
||||
"names": ["calc"],
|
||||
"turn_id": "user-passive",
|
||||
}
|
||||
assert storage.load_calls == 0
|
||||
|
||||
|
||||
@ -211,7 +215,7 @@ def test_active_skill_can_search_promote_and_call_allowed_deferred_tool():
|
||||
|
||||
result = asyncio.run(
|
||||
graph.ainvoke(
|
||||
{"messages": [HumanMessage(content="use the allowed calculator")]},
|
||||
{"messages": [HumanMessage(content="use the allowed calculator", id="user-active")]},
|
||||
context=context,
|
||||
)
|
||||
)
|
||||
@ -220,7 +224,11 @@ def test_active_skill_can_search_promote_and_call_allowed_deferred_tool():
|
||||
assert "calc" in model.bound_tool_names[1]
|
||||
assert "denied_lookup" not in model.bound_tool_names[1]
|
||||
assert _CALC_CALLS == ["2 + 2"]
|
||||
assert result["promoted"] == {"catalog_hash": setup.catalog_hash, "names": ["calc"]}
|
||||
assert result["promoted"] == {
|
||||
"catalog_hash": setup.catalog_hash,
|
||||
"names": ["calc"],
|
||||
"turn_id": "user-active",
|
||||
}
|
||||
search_result = [message for message in result["messages"] if isinstance(message, ToolMessage) and message.tool_call_id == "search-call"]
|
||||
assert len(search_result) == 1
|
||||
assert '"name": "calc"' in search_result[0].content
|
||||
@ -272,7 +280,7 @@ def test_tool_search_promotion_cannot_expose_or_execute_denied_deferred_tool():
|
||||
)
|
||||
|
||||
result = graph.invoke(
|
||||
{"messages": [HumanMessage(content="try the denied lookup")]},
|
||||
{"messages": [HumanMessage(content="try the denied lookup", id="user-denied")]},
|
||||
context=context,
|
||||
)
|
||||
|
||||
@ -281,6 +289,7 @@ def test_tool_search_promotion_cannot_expose_or_execute_denied_deferred_tool():
|
||||
assert result["promoted"] == {
|
||||
"catalog_hash": setup.catalog_hash,
|
||||
"names": [],
|
||||
"turn_id": "user-denied",
|
||||
}
|
||||
search_result = [message for message in result["messages"] if isinstance(message, ToolMessage) and message.tool_call_id == "search-denied"]
|
||||
assert len(search_result) == 1
|
||||
@ -329,7 +338,7 @@ def test_auto_promotion_cannot_expose_or_execute_denied_deferred_tool():
|
||||
)
|
||||
|
||||
result = graph.invoke(
|
||||
{"messages": [HumanMessage(content="run the denied lookup")]},
|
||||
{"messages": [HumanMessage(content="run the denied lookup", id="user-auto-denied")]},
|
||||
context=context,
|
||||
)
|
||||
|
||||
@ -338,6 +347,7 @@ def test_auto_promotion_cannot_expose_or_execute_denied_deferred_tool():
|
||||
assert result["promoted"] == {
|
||||
"catalog_hash": setup.catalog_hash,
|
||||
"names": ["denied_lookup"],
|
||||
"turn_id": "user-auto-denied",
|
||||
}
|
||||
blocked = [message for message in result["messages"] if isinstance(message, ToolMessage) and message.tool_call_id == "auto-denied-call"]
|
||||
assert len(blocked) == 1
|
||||
|
||||
@ -88,7 +88,7 @@ def test_subagent_deferral_recipe_hides_then_promotes():
|
||||
state_schema=ThreadState,
|
||||
)
|
||||
|
||||
result = asyncio.run(graph.ainvoke({"messages": [SystemMessage(content=section), HumanMessage(content="use the deferred calculator")]}))
|
||||
result = asyncio.run(graph.ainvoke({"messages": [SystemMessage(content=section), HumanMessage(content="use the deferred calculator", id="user-1")]}))
|
||||
|
||||
assert len(bound) >= 2, f"expected >=2 model binds, got {bound}"
|
||||
# Turn 1: both deferred MCP tools hidden from the subagent's model binding.
|
||||
@ -97,7 +97,11 @@ def test_subagent_deferral_recipe_hides_then_promotes():
|
||||
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"]}
|
||||
assert result["promoted"] == {
|
||||
"catalog_hash": setup.catalog_hash,
|
||||
"names": ["mcp_calc"],
|
||||
"turn_id": "user-1",
|
||||
}
|
||||
|
||||
|
||||
def test_subagent_builder_emits_working_deferred_filter():
|
||||
@ -162,7 +166,7 @@ def test_subagent_builder_emits_working_deferred_filter():
|
||||
system_prompt=None,
|
||||
state_schema=ThreadState,
|
||||
)
|
||||
result = asyncio.run(graph.ainvoke({"messages": [SystemMessage(content=section), HumanMessage(content="use the deferred calculator")]}))
|
||||
result = asyncio.run(graph.ainvoke({"messages": [SystemMessage(content=section), HumanMessage(content="use the deferred calculator", id="user-1")]}))
|
||||
|
||||
assert len(bound) >= 2, f"expected >=2 model binds, got {bound}"
|
||||
# Turn 1: both deferred MCP tools hidden - the builder-produced filter is active.
|
||||
@ -171,4 +175,8 @@ def test_subagent_builder_emits_working_deferred_filter():
|
||||
# hash correctly (a wrong hash would leave mcp_calc hidden here).
|
||||
assert "mcp_calc" in bound[1]
|
||||
assert "mcp_other" not in bound[1]
|
||||
assert result["promoted"] == {"catalog_hash": setup.catalog_hash, "names": ["mcp_calc"]}
|
||||
assert result["promoted"] == {
|
||||
"catalog_hash": setup.catalog_hash,
|
||||
"names": ["mcp_calc"],
|
||||
"turn_id": "user-1",
|
||||
}
|
||||
|
||||
@ -41,3 +41,35 @@ def test_merge_promoted_replaces_malformed_existing_without_crash():
|
||||
result = merge_promoted({"names": ["stale"]}, {"catalog_hash": "abc", "names": ["search"]})
|
||||
|
||||
assert result == {"catalog_hash": "abc", "names": ["search"]}
|
||||
|
||||
|
||||
def test_merge_promoted_unions_names_within_same_turn():
|
||||
existing = {"catalog_hash": "abc", "names": ["search"], "turn_id": "user-1"}
|
||||
|
||||
result = merge_promoted(existing, {"catalog_hash": "abc", "names": ["fetch"], "turn_id": "user-1"})
|
||||
|
||||
assert result == {"catalog_hash": "abc", "names": ["search", "fetch"], "turn_id": "user-1"}
|
||||
|
||||
|
||||
def test_merge_promoted_replaces_names_on_new_turn():
|
||||
existing = {"catalog_hash": "abc", "names": ["search"], "turn_id": "user-1"}
|
||||
|
||||
result = merge_promoted(existing, {"catalog_hash": "abc", "names": ["fetch"], "turn_id": "user-2"})
|
||||
|
||||
assert result == {"catalog_hash": "abc", "names": ["fetch"], "turn_id": "user-2"}
|
||||
|
||||
|
||||
def test_merge_promoted_tool_search_update_inherits_current_turn():
|
||||
existing = {"catalog_hash": "abc", "names": [], "turn_id": "user-1"}
|
||||
|
||||
result = merge_promoted(existing, {"catalog_hash": "abc", "names": ["search"]})
|
||||
|
||||
assert result == {"catalog_hash": "abc", "names": ["search"], "turn_id": "user-1"}
|
||||
|
||||
|
||||
def test_merge_promoted_recovers_from_malformed_existing_names():
|
||||
existing = {"catalog_hash": "abc", "names": "not-a-list", "turn_id": "user-1"}
|
||||
|
||||
result = merge_promoted(existing, {"catalog_hash": "abc", "names": ["search"]})
|
||||
|
||||
assert result == {"catalog_hash": "abc", "names": ["search"], "turn_id": "user-1"}
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user