From 1098fb4c60ec2b4fec36783c480e86282d503095 Mon Sep 17 00:00:00 2001 From: RongJie G <111257566+CorgiBoyG@users.noreply.github.com> Date: Sun, 18 Oct 2026 11:24:41 +0800 Subject: [PATCH] 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 --- README.md | 2 + .../harness/deerflow/agents/AGENTS.md | 2 +- .../deerflow/agents/middlewares/AGENTS.md | 4 +- .../deferred_tool_filter_middleware.py | 63 +++++-- .../middlewares/mcp_routing_middleware.py | 27 +-- .../agents/middlewares/message_utils.py | 27 ++- .../tool_promotion_audit_middleware.py | 4 + .../harness/deerflow/agents/thread_state.py | 32 ++-- .../tests/test_deferred_filter_middleware.py | 104 ++++++++++- .../test_deferred_promotion_integration.py | 170 +++++++++++++++++- .../tests/test_deferred_tool_crosscontext.py | 9 +- .../tests/test_mcp_routing_auto_promote.py | 109 ++++++++++- .../test_skill_policy_deferred_discovery.py | 22 ++- ...subagent_deferred_promotion_integration.py | 16 +- backend/tests/test_thread_state_promoted.py | 32 ++++ 15 files changed, 564 insertions(+), 59 deletions(-) diff --git a/README.md b/README.md index c29e0a4b41..ccbeca064a 100644 --- a/README.md +++ b/README.md @@ -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` diff --git a/backend/packages/harness/deerflow/agents/AGENTS.md b/backend/packages/harness/deerflow/agents/AGENTS.md index ffe83eb07a..ee5edcbcbd 100644 --- a/backend/packages/harness/deerflow/agents/AGENTS.md +++ b/backend/packages/harness/deerflow/agents/AGENTS.md @@ -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, diff --git a/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md b/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md index 66473a31d5..a646c87576 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md +++ b/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md @@ -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. diff --git a/backend/packages/harness/deerflow/agents/middlewares/deferred_tool_filter_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/deferred_tool_filter_middleware.py index 27976cb697..f0fdb4587f 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/deferred_tool_filter_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/deferred_tool_filter_middleware.py @@ -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, diff --git a/backend/packages/harness/deerflow/agents/middlewares/mcp_routing_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/mcp_routing_middleware.py index 92ba8e9833..34f67c0133 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/mcp_routing_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/mcp_routing_middleware.py @@ -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: diff --git a/backend/packages/harness/deerflow/agents/middlewares/message_utils.py b/backend/packages/harness/deerflow/agents/middlewares/message_utils.py index 53857fb614..9f6c0b30f6 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/message_utils.py +++ b/backend/packages/harness/deerflow/agents/middlewares/message_utils.py @@ -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. diff --git a/backend/packages/harness/deerflow/agents/middlewares/tool_promotion_audit_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/tool_promotion_audit_middleware.py index 980c5a769e..00d1c8d22c 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/tool_promotion_audit_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/tool_promotion_audit_middleware.py @@ -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() diff --git a/backend/packages/harness/deerflow/agents/thread_state.py b/backend/packages/harness/deerflow/agents/thread_state.py index 883723ec10..a75e837a7e 100644 --- a/backend/packages/harness/deerflow/agents/thread_state.py +++ b/backend/packages/harness/deerflow/agents/thread_state.py @@ -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) diff --git a/backend/tests/test_deferred_filter_middleware.py b/backend/tests/test_deferred_filter_middleware.py index 9a6200a8a8..b2bd7f8fe9 100644 --- a/backend/tests/test_deferred_filter_middleware.py +++ b/backend/tests/test_deferred_filter_middleware.py @@ -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"} diff --git a/backend/tests/test_deferred_promotion_integration.py b/backend/tests/test_deferred_promotion_integration.py index 38c10234d9..9647029700 100644 --- a/backend/tests/test_deferred_promotion_integration.py +++ b/backend/tests/test_deferred_promotion_integration.py @@ -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) diff --git a/backend/tests/test_deferred_tool_crosscontext.py b/backend/tests/test_deferred_tool_crosscontext.py index e4a1a26aec..2e42699a64 100644 --- a/backend/tests/test_deferred_tool_crosscontext.py +++ b/backend/tests/test_deferred_tool_crosscontext.py @@ -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 diff --git a/backend/tests/test_mcp_routing_auto_promote.py b/backend/tests/test_mcp_routing_auto_promote.py index 83f54b7311..e55902eadd 100644 --- a/backend/tests/test_mcp_routing_auto_promote.py +++ b/backend/tests/test_mcp_routing_auto_promote.py @@ -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 diff --git a/backend/tests/test_skill_policy_deferred_discovery.py b/backend/tests/test_skill_policy_deferred_discovery.py index 5150f48e4c..f386fce2c9 100644 --- a/backend/tests/test_skill_policy_deferred_discovery.py +++ b/backend/tests/test_skill_policy_deferred_discovery.py @@ -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 diff --git a/backend/tests/test_subagent_deferred_promotion_integration.py b/backend/tests/test_subagent_deferred_promotion_integration.py index d4a56f91da..a97083adf8 100644 --- a/backend/tests/test_subagent_deferred_promotion_integration.py +++ b/backend/tests/test_subagent_deferred_promotion_integration.py @@ -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", + } diff --git a/backend/tests/test_thread_state_promoted.py b/backend/tests/test_thread_state_promoted.py index 6cd98b711a..0ff2ea4690 100644 --- a/backend/tests/test_thread_state_promoted.py +++ b/backend/tests/test_thread_state_promoted.py @@ -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"}