mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-09 13:39:26 +00:00
feat(observability): persist deferred tool promotions (#5183)
* feat(observability): persist deferred tool promotions Signed-off-by: PeaceMaker-best <221849497+PeaceMaker-best@users.noreply.github.com> * fix(ci): trim agent guidance chain Signed-off-by: PeaceMaker-best <221849497+PeaceMaker-best@users.noreply.github.com> --------- Signed-off-by: PeaceMaker-best <221849497+PeaceMaker-best@users.noreply.github.com> Co-authored-by: PeaceMaker-best <221849497+PeaceMaker-best@users.noreply.github.com>
This commit is contained in:
PeaceMaker-best
GitHub
parent
0f7d8709d3
commit
3c36217a51
@ -763,6 +763,12 @@ finalization) keep the existing completion-data behavior: they receive the
|
||||
zero-delivery receipt but do not overwrite RunStore completion fields with an
|
||||
empty snapshot.
|
||||
|
||||
The same run event history records loop-detection decisions and deferred MCP
|
||||
tool promotions for both the lead agent and ordinary task subagents. Promotion
|
||||
events identify newly promoted deferred-tool names and whether routing metadata or
|
||||
`tool_search` selected them, without copying the search query, routing keywords,
|
||||
schemas, arguments, results, or catalog hash into the promotion event itself.
|
||||
|
||||
#### LangSmith Tracing
|
||||
|
||||
DeerFlow has built-in [LangSmith](https://smith.langchain.com) integration for observability. When enabled, all LLM calls, agent runs, and tool executions are traced and visible in the LangSmith dashboard.
|
||||
|
||||
@ -77,10 +77,10 @@ through run-event or specialized APIs:
|
||||
| `middleware:{tag}` | `middleware` | `record_middleware()` |
|
||||
|
||||
Current middleware tags are `guardrail`, `loop_detection`,
|
||||
`safety_termination`, `skill_activation`, and `skill_secrets`. The pattern is
|
||||
intentionally open so new middleware tags are additive. Because the full event
|
||||
type is limited to 32 characters and `middleware:` uses 11, a tag must contain
|
||||
1-21 characters.
|
||||
`safety_termination`, `skill_activation`, `skill_secrets`, and
|
||||
`tool_promotion`. The pattern is intentionally open so new middleware tags are
|
||||
additive. Because the full event type is limited to 32 characters and
|
||||
`middleware:` uses 11, a tag must contain 1-21 characters.
|
||||
|
||||
`middleware:loop_detection` records transitions into the warned state (first
|
||||
per call hash or per tool-frequency burst) and each hard stop produced by
|
||||
@ -95,6 +95,22 @@ subagent, and its agent id when applicable. Tool arguments, prompts, message
|
||||
content, tool results, and argument-derived hashes are not persisted in this
|
||||
event.
|
||||
|
||||
`middleware:tool_promotion` records deferred MCP schemas newly promoted for
|
||||
the active catalog. `changes.source` distinguishes automatic `routing_hint`
|
||||
promotion from an explicit `tool_search`; the remaining fields contain sorted
|
||||
new tool names, their count, `is_subagent`, and the optional subagent
|
||||
`agent_id`. Explicit-search events observe the final `Command` after the active
|
||||
skill policy has removed denied schemas. Automatic promotion records the graph
|
||||
state transition, but promotion never bypasses the independent skill or
|
||||
authorization execution policies. Repeated model passes or searches that add
|
||||
no names emit no event. Queries, routing keywords, catalog hashes, tool schemas
|
||||
and descriptions, arguments, and results are not copied into this middleware
|
||||
event. Other event types retain their existing payload contracts.
|
||||
|
||||
Ordinary task-tool subagents forward both loop-detection and tool-promotion
|
||||
appends to the parent run loop through dedicated recorder context keys. The
|
||||
loop-bound `RunJournal` itself never enters the isolated subagent loop.
|
||||
|
||||
### Opaque Run Outputs
|
||||
|
||||
`run.end.content` is the root graph output and is intentionally opaque. Its
|
||||
@ -191,6 +207,6 @@ be used by new producers.
|
||||
representations: memory retains Python values, while JSONL and database
|
||||
stores read them back as strings.
|
||||
- Durable batch subagent loop detection and deferred-tool promotion do not
|
||||
currently emit middleware events.
|
||||
emit middleware events because those runs have no parent run journal.
|
||||
- Journal attribution, token accounting, and external tracing metadata still
|
||||
depend on manual instrumentation at several LLM call sites.
|
||||
|
||||
@ -541,6 +541,14 @@ def build_middlewares(
|
||||
)
|
||||
)
|
||||
|
||||
# Observe the final tool_search Command after every inner policy/result
|
||||
# transformer has run. Tool wrappers are first-in-list outermost, so this
|
||||
# must be registered before SkillToolPolicyMiddleware.
|
||||
if deferred_setup is not None and deferred_setup.deferred_names:
|
||||
from deerflow.agents.middlewares.tool_promotion_audit_middleware import DeferredToolPromotionAuditMiddleware
|
||||
|
||||
middlewares.append(DeferredToolPromotionAuditMiddleware(deferred_setup.deferred_names, deferred_setup.catalog_hash))
|
||||
|
||||
# Enabled skills are only discoverable metadata. Apply allowed-tools at
|
||||
# runtime after explicit slash activation or an actual skill-file load.
|
||||
from deerflow.agents.middlewares.skill_tool_policy_middleware import SkillToolPolicyMiddleware
|
||||
|
||||
@ -91,15 +91,16 @@ Before changing a later authorization phase, read the [authorization RFC](../../
|
||||
21. **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.
|
||||
22. **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
|
||||
23. **ViewImageMiddleware** - *(optional, if the model supports vision)* Appends a hidden HumanMessage with base64 image data, identified by a reserved ID prefix plus a server-owned metadata marker, to `ModelRequest.messages` in `wrap_model_call` / `awrap_model_call`. The payload lives only in that request and is never returned as a state update, so no checkpoint carries it and an interrupted run cannot strand it in history; state keeps only the lightweight `viewed_images` metadata. It owns that context and rebuilds it per call: its own message is swept out of the request first — a thread checkpointed by the earlier `before_model`/`after_model` pair (which wrote the payload into state and took it back out with `RemoveMessage`) can carry one that reached state but was never removed, and leaving it in would resend that base64 in every later request for the life of the thread — then a freshly built one is appended when warranted. The sweep requires both the reserved ID prefix and the server-owned marker, so a client cannot get its own message dropped; unmarked leftovers predating the marker are left in place and merely not duplicated
|
||||
24. **McpRoutingMiddleware** - *(optional, if `tool_search.enabled` and PR1 MCP routing metadata produce a routing index)* Auto-promotes matching deferred MCP tool schemas before the model call by writing a minimal `promoted` state update. It matches only the latest real `HumanMessage`, uses the global `tool_search.auto_promote_top_k` limit (default 3, clamped to 1..5), never executes tools, and must be installed before `DeferredToolFilterMiddleware`
|
||||
25. **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)
|
||||
26. **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
|
||||
27. **SubagentLimitMiddleware** - *(optional, if `subagent_enabled`)* Truncates excess ordinary `task` tool calls to enforce both the per-response concurrency limit (`max_concurrent_subagents`, resolved against startup `subagent_runtime.max_running` and the 1-64 safety range before construction) and the per-run total delegation cap (`max_total_subagents` runtime override or `subagents.max_total_per_run`, default 6, clamped to 1-50). The total cap counts current-run entries in the durable delegation ledger (entries are tagged with `run_id` when captured), so repeated planning checkpoints in one run cannot keep launching legal-sized batches indefinitely, while later user turns in the same thread get a fresh run budget. Explicit durable `batch_task` calls are a separate mode with persisted total/live/running limits and are not rewritten into ordinary ledger entries. If the ordinary cap is exhausted, the middleware strips remaining `task` calls, forces `finish_reason="stop"`, and appends a visible limit note so the run can synthesize existing results instead of ending with an empty tool-call response.
|
||||
28. **LoopDetectionMiddleware** - *(optional, if `loop_detection.enabled`)* Detects repeated tool-call loops; hard-stop clears both structured `tool_calls` and raw provider tool-call metadata before forcing a final text answer; stamps `loop_capped` via `consume_stop_reason` (#3875 Phase 2), symmetric to `TokenBudgetMiddleware`; persists warned-state transitions (first per call hash or per tool-frequency burst) and hard stops as `middleware:loop_detection`, attributed with `is_subagent` and the optional `agent_id`, without tool arguments, message content, tool results, or argument-derived hashes. Ordinary task-tool subagents receive a loop-detection-only recorder proxy that forwards the append to the parent run loop; never pass the loop-bound `RunJournal` itself into their isolated event loop. Durable batch subagents have no parent run journal and do not persist these transitions
|
||||
29. **TokenBudgetMiddleware** - *(optional, if `token_budget.enabled`)* Enforces per-run token limits
|
||||
30. **Custom middlewares** - *(optional)* Any `custom_middlewares` passed to `build_middlewares` are injected here, before config-declared extensions and the terminal-response/safety/clarification tail
|
||||
31. **Configured extension middlewares** - *(optional, if `extensions.middlewares` is set in `config.yaml` or `extensions_config.json`)* Zero-argument `AgentMiddleware` classes loaded from `module.path:ClassName` entries via `deerflow.reflection.resolve_class`. Missing packages, invalid classes, and broken modules fail loudly at agent creation. These run after built-ins/programmatic custom middleware and after the lead/subagent loop/token guards, but before the terminal-response/safety/clarification tail; subagents receive the same configured extension middleware class list before their safety tail. Treat these files as trusted operator config because middleware paths instantiate arbitrary code. Gateway skill/MCP toggle endpoints preserve this field through `to_file_dict()` but must not add a write path for `extensions.middlewares` without an explicit trust-boundary review. Lead-only vs subagent-only middleware lists and per-context constructor parameters are not expressible in this MVP.
|
||||
32. **TerminalResponseMiddleware** - When a provider returns an empty terminal `AIMessage` after tool execution, injects a hidden recovery prompt and retries the model once; a second empty response is replaced in checkpoint state by a visible error fallback marked for the run worker, so the run finishes as an error instead of a silent success
|
||||
33. **ModelLengthFinishReasonMiddleware** - Records `stop_reason=model_length_capped` when provider-specific length detectors match a terminal `AIMessage` without tool-call intent (`finish_reason=length` / `MAX_TOKENS`, or `stop_reason=max_tokens`), preserving the original assistant content and never reparsing textual tool-call-like envelopes
|
||||
34. **SafetyFinishReasonMiddleware** - *(optional, if `safety_finish_reason.enabled`)* Suppresses tool execution when the provider safety-terminated the response (e.g. `finish_reason=content_filter`); registered after terminal-response/custom/configured middlewares so LangChain's reverse-order `after_model` dispatch runs it first
|
||||
35. **ClarificationMiddleware** - Intercepts `ask_clarification`, writes a readable `ToolMessage.content` fallback plus a structured `ToolMessage.artifact.human_input` payload, and interrupts via `Command(goto=END)` (must be last). `after_model` drops same-turn sibling tool calls so they cannot run before the user answers; a malformed `ask_clarification` parked on `invalid_tool_calls` is the same stop signal. `disable_clarification` runs keep the siblings. Payloads are versioned — legacy `free_text`/`choice_with_other` stay `version: 1`; the v2 `form` mode (from `fields`) is `version: 2` so older frontends reject it and fall back to plain text. Field normalization is deterministic and lives in the middleware (it short-circuits before tool execution, so tool-arg typing gives no runtime validation), and it is atomic: any structurally broken entry — non-dict, bad/duplicate name, a name colliding with a JS `Object.prototype` member (`__proto__`/`constructor`), or exceeding the caps (16 fields / 24 options per field / 200 chars per text / `MAX_FORM_SERIALIZED_BYTES` = 16KB UTF-8, the per-item caps alone admitting forms whose IM text fallback overruns channel limits) — degrades the whole form to the legacy option/free-text modes, so a card never renders "complete" while missing a field. Benign issues degrade locally (unknown types — incl. unhashable JSON like `type: []`, which must not raise from the membership probe — and option-less selects become `text`); options are trimmed/deduped with blanks dropped (form- and top-level) since the frontend rejects blank labels. XML-to-dict option payloads are recursively flattened from dict/list containers in source order, scalar leaves kept, residual XML tags stripped before that trimming. Checkboxes are booleans defaulting to "no"; `required` on one means consent semantics. The response protocol is unchanged (v1 `text`/`option`): form cards submit a text summary as `response_kind: "text"`, so journal persistence needs no new allowlist entries. Because this middleware can short-circuit before `on_tool_end`, `RunJournal` does a root-run reconciliation for `ToolMessage`s whose `tool_call_id` came from the current run, so cards survive checkpoint compaction. That reconciliation is **not** `ask_clarification`-only — any middleware that answers a tool call has the same gap, and a result the user saw must not vanish on reload (#4666 — `ReadBeforeWriteMiddleware` blocked-write errors reached the UI but not the event store). It is bounded by three conditions, not a name allowlist: the message is user-visible, the call belongs to this run's **lead agent** (`_remember_current_run_tool_calls` records lead-agent calls only; subagent results stay in `subagent.step`), and it is not already persisted. Human Input Card replies are `hide_from_ui` `HumanMessage`s with `additional_kwargs.human_input_response`; `RunJournal` persists only allowlisted hidden sources (currently `ask_clarification`) as `llm.human.input`.
|
||||
24. **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
|
||||
25. **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, returns the original `Command`, and omits private payloads
|
||||
26. **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)
|
||||
27. **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
|
||||
28. **SubagentLimitMiddleware** - *(optional, if `subagent_enabled`)* Truncates excess ordinary `task` tool calls to enforce both the per-response concurrency limit (`max_concurrent_subagents`, resolved against startup `subagent_runtime.max_running` and the 1-64 safety range before construction) and the per-run total delegation cap (`max_total_subagents` runtime override or `subagents.max_total_per_run`, default 6, clamped to 1-50). The total cap counts current-run entries in the durable delegation ledger (entries are tagged with `run_id` when captured), so repeated planning checkpoints in one run cannot keep launching legal-sized batches indefinitely, while later user turns in the same thread get a fresh run budget. Explicit durable `batch_task` calls are a separate mode with persisted total/live/running limits and are not rewritten into ordinary ledger entries. If the ordinary cap is exhausted, the middleware strips remaining `task` calls, forces `finish_reason="stop"`, and appends a visible limit note so the run can synthesize existing results instead of ending with an empty tool-call response.
|
||||
29. **LoopDetectionMiddleware** - *(optional, if `loop_detection.enabled`)* Detects repeated tool-call loops; hard-stop clears both structured `tool_calls` and raw provider tool-call metadata before forcing a final text answer; stamps `loop_capped` via `consume_stop_reason` (#3875 Phase 2), symmetric to `TokenBudgetMiddleware`; persists warned-state transitions (first per call hash or per tool-frequency burst) and hard stops as `middleware:loop_detection`, attributed with `is_subagent` and the optional `agent_id`, without tool arguments, message content, tool results, or argument-derived hashes. Ordinary task subagents get dedicated recorder keys through a parent-loop proxy; never pass `RunJournal` into their isolated loop. Durable batch subagents have no parent run journal and do not persist these transitions
|
||||
30. **TokenBudgetMiddleware** - *(optional, if `token_budget.enabled`)* Enforces per-run token limits
|
||||
31. **Custom middlewares** - *(optional)* Any `custom_middlewares` passed to `build_middlewares` are injected here, before config-declared extensions and the terminal-response/safety/clarification tail
|
||||
32. **Configured extension middlewares** - *(optional, if `extensions.middlewares` is set in `config.yaml` or `extensions_config.json`)* Zero-argument `AgentMiddleware` classes loaded from `module.path:ClassName` entries via `deerflow.reflection.resolve_class`. Missing packages, invalid classes, and broken modules fail loudly at agent creation. These run after built-ins/programmatic custom middleware and after the lead/subagent loop/token guards, but before the terminal-response/safety/clarification tail; subagents receive the same configured extension middleware class list before their safety tail. Treat these files as trusted operator config because middleware paths instantiate arbitrary code. Gateway skill/MCP toggle endpoints preserve this field through `to_file_dict()` but must not add a write path for `extensions.middlewares` without an explicit trust-boundary review. Lead-only vs subagent-only middleware lists and per-context constructor parameters are not expressible in this MVP.
|
||||
33. **TerminalResponseMiddleware** - When a provider returns an empty terminal `AIMessage` after tool execution, injects a hidden recovery prompt and retries the model once; a second empty response is replaced in checkpoint state by a visible error fallback marked for the run worker, so the run finishes as an error instead of a silent success
|
||||
34. **ModelLengthFinishReasonMiddleware** - Records `stop_reason=model_length_capped` when provider-specific length detectors match a terminal `AIMessage` without tool-call intent (`finish_reason=length` / `MAX_TOKENS`, or `stop_reason=max_tokens`), preserving the original assistant content and never reparsing textual tool-call-like envelopes
|
||||
35. **SafetyFinishReasonMiddleware** - *(optional, if `safety_finish_reason.enabled`)* Suppresses tool execution when the provider safety-terminated the response (e.g. `finish_reason=content_filter`); registered after terminal-response/custom/configured middlewares so LangChain's reverse-order `after_model` dispatch runs it first
|
||||
36. **ClarificationMiddleware** - Intercepts `ask_clarification`, writes a readable `ToolMessage.content` fallback plus a structured `ToolMessage.artifact.human_input` payload, and interrupts via `Command(goto=END)` (must be last). `after_model` drops same-turn sibling tool calls so they cannot run before the user answers; a malformed `ask_clarification` parked on `invalid_tool_calls` is the same stop signal. `disable_clarification` runs keep the siblings. Payloads are versioned — legacy `free_text`/`choice_with_other` stay `version: 1`; the v2 `form` mode (from `fields`) is `version: 2` so older frontends reject it and fall back to plain text. Field normalization is deterministic and lives in the middleware (it short-circuits before tool execution, so tool-arg typing gives no runtime validation), and it is atomic: any structurally broken entry — non-dict, bad/duplicate name, a name colliding with a JS `Object.prototype` member (`__proto__`/`constructor`), or exceeding the caps (16 fields / 24 options per field / 200 chars per text / `MAX_FORM_SERIALIZED_BYTES` = 16KB UTF-8, the per-item caps alone admitting forms whose IM text fallback overruns channel limits) — degrades the whole form to the legacy option/free-text modes, so a card never renders "complete" while missing a field. Benign issues degrade locally (unknown types — incl. unhashable JSON like `type: []`, which must not raise from the membership probe — and option-less selects become `text`); options are trimmed/deduped with blanks dropped (form- and top-level) since the frontend rejects blank labels. XML-to-dict option payloads are recursively flattened from dict/list containers in source order, scalar leaves kept, residual XML tags stripped before that trimming. Checkboxes are booleans defaulting to "no"; `required` on one means consent semantics. The response protocol is unchanged (v1 `text`/`option`): form cards submit a text summary as `response_kind: "text"`, so journal persistence needs no new allowlist entries. Because this middleware can short-circuit before `on_tool_end`, `RunJournal` does a root-run reconciliation for `ToolMessage`s whose `tool_call_id` came from the current run, so cards survive checkpoint compaction. That reconciliation is **not** `ask_clarification`-only — any middleware that answers a tool call has the same gap, and a result the user saw must not vanish on reload (#4666 — `ReadBeforeWriteMiddleware` blocked-write errors reached the UI but not the event store). It is bounded by three conditions, not a name allowlist: the message is user-visible, the call belongs to this run's **lead agent** (`_remember_current_run_tool_calls` records lead-agent calls only; subagent results stay in `subagent.step`), and it is not already persisted. Human Input Card replies are `hide_from_ui` `HumanMessage`s with `additional_kwargs.human_input_response`; `RunJournal` persists only allowlisted hidden sources (currently `ask_clarification`) as `llm.human.input`.
|
||||
|
||||
@ -1,3 +1,4 @@
|
||||
"""Private runtime-context keys for narrowly scoped middleware audit recorders."""
|
||||
|
||||
LOOP_DETECTION_RECORDER_CONTEXT_KEY = "__run_loop_detection_recorder"
|
||||
TOOL_PROMOTION_RECORDER_CONTEXT_KEY = "__run_tool_promotion_recorder"
|
||||
|
||||
@ -11,6 +11,7 @@ from langchain.agents.middleware import AgentMiddleware
|
||||
from langchain_core.messages import HumanMessage
|
||||
from langgraph.runtime import Runtime
|
||||
|
||||
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, is_real_user_message
|
||||
|
||||
@ -101,10 +102,20 @@ class McpRoutingMiddleware(AgentMiddleware[AgentState]):
|
||||
matched.sort(key=lambda item: (-item[0], item[1]))
|
||||
return [name for _, name in matched[: self._top_k]]
|
||||
|
||||
def _state_update(self, state: Mapping[str, Any] | None) -> dict[str, Any] | None:
|
||||
def _state_update(self, state: Mapping[str, Any] | None, runtime: Runtime | None) -> dict[str, Any] | None:
|
||||
names = self._matched_names(state)
|
||||
if not names:
|
||||
return None
|
||||
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
|
||||
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,
|
||||
producer=type(self).__name__,
|
||||
hook="before_model",
|
||||
source="routing_hint",
|
||||
tool_names=set(names) - already_promoted,
|
||||
)
|
||||
logger.debug(
|
||||
"McpRoutingMiddleware auto-promoted %d deferred tool schema(s) catalog=%s names=%s",
|
||||
len(names),
|
||||
@ -120,11 +131,11 @@ class McpRoutingMiddleware(AgentMiddleware[AgentState]):
|
||||
|
||||
@override
|
||||
def before_model(self, state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
|
||||
return self._state_update(state)
|
||||
return self._state_update(state, runtime)
|
||||
|
||||
@override
|
||||
async def abefore_model(self, state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
|
||||
return self._state_update(state)
|
||||
return self._state_update(state, runtime)
|
||||
|
||||
|
||||
def assert_mcp_routing_before_deferred_filter(middlewares: Sequence[AgentMiddleware]) -> None:
|
||||
|
||||
@ -392,6 +392,10 @@ def build_subagent_runtime_middlewares(
|
||||
slash_source_owner_token=slash_source_owner_token,
|
||||
)
|
||||
)
|
||||
if deferred_setup is not None and deferred_setup.deferred_names:
|
||||
from deerflow.agents.middlewares.tool_promotion_audit_middleware import DeferredToolPromotionAuditMiddleware
|
||||
|
||||
middlewares.append(DeferredToolPromotionAuditMiddleware(deferred_setup.deferred_names, deferred_setup.catalog_hash))
|
||||
middlewares.append(
|
||||
SkillToolPolicyMiddleware(
|
||||
available_skills=available_skills,
|
||||
|
||||
@ -0,0 +1,143 @@
|
||||
"""Persist effective deferred-tool promotion decisions without sensitive payloads."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections.abc import Awaitable, Callable, Iterable, Mapping
|
||||
from typing import Any, override
|
||||
|
||||
from langchain.agents import AgentState
|
||||
from langchain.agents.middleware import AgentMiddleware
|
||||
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.audit_context import TOOL_PROMOTION_RECORDER_CONTEXT_KEY
|
||||
from deerflow.runtime.events.catalog import MIDDLEWARE_TOOL_PROMOTION_TAG
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_TOOL_SEARCH_NAME = "tool_search"
|
||||
|
||||
|
||||
def record_tool_promotion(
|
||||
runtime: Runtime | None,
|
||||
*,
|
||||
producer: str,
|
||||
hook: str,
|
||||
source: str,
|
||||
tool_names: Iterable[str],
|
||||
) -> None:
|
||||
"""Record one effective promotion decision, failing open on telemetry errors."""
|
||||
names = sorted(set(tool_names))
|
||||
if not names:
|
||||
return
|
||||
|
||||
context = getattr(runtime, "context", None)
|
||||
if not isinstance(context, dict):
|
||||
return
|
||||
is_subagent = context.get("is_subagent") is True
|
||||
recorder = context.get(TOOL_PROMOTION_RECORDER_CONTEXT_KEY)
|
||||
if recorder is None:
|
||||
# Lead runs own a RunJournal. Ordinary task-tool subagents receive only
|
||||
# the narrow, loop-safe recorder key above.
|
||||
recorder = context.get("__run_journal")
|
||||
if recorder is None:
|
||||
return
|
||||
|
||||
try:
|
||||
claim = getattr(recorder, "claim_tool_promotions", None)
|
||||
if callable(claim):
|
||||
names = claim(names)
|
||||
if not names:
|
||||
return
|
||||
recorder.record_middleware(
|
||||
tag=MIDDLEWARE_TOOL_PROMOTION_TAG,
|
||||
name=producer,
|
||||
hook=hook,
|
||||
action="promote",
|
||||
changes={
|
||||
"source": source,
|
||||
"tool_names": names,
|
||||
"count": len(names),
|
||||
"is_subagent": is_subagent,
|
||||
"agent_id": context.get("agent_id") if is_subagent else None,
|
||||
},
|
||||
)
|
||||
except Exception: # noqa: BLE001
|
||||
# Observation must never alter the agent trajectory it describes.
|
||||
logger.warning("Failed to record middleware:tool_promotion event", exc_info=True)
|
||||
|
||||
|
||||
class DeferredToolPromotionAuditMiddleware(AgentMiddleware[AgentState]):
|
||||
"""Observe final ``tool_search`` Commands after policy filtering.
|
||||
|
||||
This wrapper must remain outer of ``SkillToolPolicyMiddleware``. Tool-call
|
||||
wrappers unwind in reverse registration order, so observing the handler's
|
||||
final return value is what prevents denied schemas from being reported as
|
||||
effective promotions.
|
||||
"""
|
||||
|
||||
def __init__(self, deferred_names: frozenset[str], catalog_hash: str | None) -> None:
|
||||
super().__init__()
|
||||
self._deferred = deferred_names
|
||||
self._catalog_hash = catalog_hash
|
||||
|
||||
def release_policy_parameters(self) -> dict[str, object]:
|
||||
return {
|
||||
"deferred_names": sorted(self._deferred),
|
||||
"catalog_hash": self._catalog_hash,
|
||||
"observation": "final_tool_search_command",
|
||||
}
|
||||
|
||||
def _current_promoted(self, state: Mapping[str, Any] | None) -> set[str]:
|
||||
promoted = (state or {}).get("promoted")
|
||||
if not isinstance(promoted, Mapping) or promoted.get("catalog_hash") != self._catalog_hash:
|
||||
return set()
|
||||
names = promoted.get("names")
|
||||
if not isinstance(names, list) or not all(isinstance(name, str) for name in names):
|
||||
return set()
|
||||
return set(names)
|
||||
|
||||
def _new_promotions(self, request: ToolCallRequest, result: ToolMessage | Command) -> list[str]:
|
||||
if request.tool_call.get("name") != _TOOL_SEARCH_NAME:
|
||||
return []
|
||||
if not isinstance(result, Command) or not isinstance(result.update, dict):
|
||||
return []
|
||||
promoted = result.update.get("promoted")
|
||||
if not isinstance(promoted, dict) or promoted.get("catalog_hash") != self._catalog_hash:
|
||||
return []
|
||||
names = promoted.get("names")
|
||||
if not isinstance(names, list) or not all(isinstance(name, str) for name in names):
|
||||
return []
|
||||
return sorted((set(names) & self._deferred) - self._current_promoted(request.state))
|
||||
|
||||
def _record(self, request: ToolCallRequest, result: ToolMessage | Command) -> None:
|
||||
record_tool_promotion(
|
||||
request.runtime,
|
||||
producer=type(self).__name__,
|
||||
hook="wrap_tool_call",
|
||||
source="tool_search",
|
||||
tool_names=self._new_promotions(request, result),
|
||||
)
|
||||
|
||||
@override
|
||||
def wrap_tool_call(
|
||||
self,
|
||||
request: ToolCallRequest,
|
||||
handler: Callable[[ToolCallRequest], ToolMessage | Command],
|
||||
) -> ToolMessage | Command:
|
||||
result = handler(request)
|
||||
self._record(request, result)
|
||||
return result
|
||||
|
||||
@override
|
||||
async def awrap_tool_call(
|
||||
self,
|
||||
request: ToolCallRequest,
|
||||
handler: Callable[[ToolCallRequest], Awaitable[ToolMessage | Command]],
|
||||
) -> ToolMessage | Command:
|
||||
result = await handler(request)
|
||||
self._record(request, result)
|
||||
return result
|
||||
@ -78,12 +78,19 @@ def core_ordering_constraints() -> tuple[OrderingConstraint, ...]:
|
||||
"""
|
||||
from deerflow.agents.middlewares.read_before_write_middleware import ReadBeforeWriteMiddleware
|
||||
from deerflow.agents.middlewares.sandbox_audit_middleware import SandboxAuditMiddleware
|
||||
from deerflow.agents.middlewares.skill_tool_policy_middleware import SkillToolPolicyMiddleware
|
||||
from deerflow.agents.middlewares.tool_error_handling_middleware import ToolErrorHandlingMiddleware
|
||||
from deerflow.agents.middlewares.tool_progress_middleware import ToolProgressMiddleware
|
||||
from deerflow.agents.middlewares.tool_promotion_audit_middleware import DeferredToolPromotionAuditMiddleware
|
||||
from deerflow.agents.middlewares.tool_receipt_middleware import ToolReceiptMiddleware
|
||||
from deerflow.guardrails.middleware import GuardrailMiddleware
|
||||
|
||||
return (
|
||||
OrderingConstraint(
|
||||
outer=DeferredToolPromotionAuditMiddleware,
|
||||
inner=SkillToolPolicyMiddleware,
|
||||
reason=("DeferredToolPromotionAuditMiddleware must observe the policy-filtered tool_search Command so denied schemas are never reported as effective promotions"),
|
||||
),
|
||||
OrderingConstraint(
|
||||
outer=ToolProgressMiddleware,
|
||||
inner=ToolErrorHandlingMiddleware,
|
||||
|
||||
@ -105,6 +105,17 @@ one accumulated receipt across multiple goal-continuation `_stream_once` calls;
|
||||
journal tests drive LangChain's real async callback dispatcher against a single
|
||||
journal to pin serialized, deduplicated parallel tool callbacks.
|
||||
|
||||
**Deferred-tool promotion event deduplication** (`runtime/journal.py`): one
|
||||
`RunJournal` owns the lead graph's run-scoped atomic promotion claim. Parallel
|
||||
`tool_search` Sends read the same pre-step state, so state diffing alone can
|
||||
label the same schema as new more than once even though the `promoted` reducer
|
||||
unions it once. Producers claim sorted candidate names before appending
|
||||
`middleware:tool_promotion`; later overlapping decisions emit only their
|
||||
unclaimed remainder. Ordinary task-tool subagents use an equivalent claim on
|
||||
their per-execution parent-loop proxy, preserving separate events when two
|
||||
different delegated agents promote the same tool. The active catalog is fixed
|
||||
for one graph execution, so the claim needs no persisted catalog hash.
|
||||
|
||||
**Targeted run-event attribution** (`runtime/events/store/`):
|
||||
`RunEventStore.find_latest_ai_message_run_ids()` has a complete-or-error
|
||||
contract. Its default implementation walks `list_messages()` backward in
|
||||
|
||||
@ -81,12 +81,14 @@ MIDDLEWARE_LOOP_DETECTION_TAG = "loop_detection"
|
||||
MIDDLEWARE_SAFETY_TERMINATION_TAG = "safety_termination"
|
||||
MIDDLEWARE_SKILL_ACTIVATION_TAG = "skill_activation"
|
||||
MIDDLEWARE_SKILL_SECRETS_TAG = "skill_secrets"
|
||||
MIDDLEWARE_TOOL_PROMOTION_TAG = "tool_promotion"
|
||||
MIDDLEWARE_EVENT_TAGS = (
|
||||
MIDDLEWARE_GUARDRAIL_TAG,
|
||||
MIDDLEWARE_LOOP_DETECTION_TAG,
|
||||
MIDDLEWARE_SAFETY_TERMINATION_TAG,
|
||||
MIDDLEWARE_SKILL_ACTIVATION_TAG,
|
||||
MIDDLEWARE_SKILL_SECRETS_TAG,
|
||||
MIDDLEWARE_TOOL_PROMOTION_TAG,
|
||||
)
|
||||
|
||||
JOURNAL_RUN_EVENT_DEFINITIONS = (
|
||||
|
||||
@ -19,8 +19,9 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
|
||||
from datetime import UTC, datetime
|
||||
from typing import TYPE_CHECKING, Any, cast
|
||||
from uuid import UUID
|
||||
@ -267,6 +268,8 @@ class RunJournal(BaseCallbackHandler):
|
||||
self._counted_external_source_ids: set[str] = set()
|
||||
self._counted_message_llm_run_ids: set[str] = set()
|
||||
self._memory_context_recorded = False
|
||||
self._tool_promotion_claim_lock = threading.Lock()
|
||||
self._claimed_tool_promotions: set[str] = set()
|
||||
|
||||
# Convenience fields
|
||||
self._last_ai_msg: str | None = None
|
||||
@ -859,6 +862,14 @@ class RunJournal(BaseCallbackHandler):
|
||||
content={"name": name, "hook": hook, "action": action, "changes": changes},
|
||||
)
|
||||
|
||||
def claim_tool_promotions(self, tool_names: Iterable[str]) -> list[str]:
|
||||
"""Atomically claim names not yet reported by this run's lead agent."""
|
||||
candidates = sorted(set(tool_names))
|
||||
with self._tool_promotion_claim_lock:
|
||||
claimed = [name for name in candidates if name not in self._claimed_tool_promotions]
|
||||
self._claimed_tool_promotions.update(claimed)
|
||||
return claimed
|
||||
|
||||
def record_memory_context(self, *, content_sha256: str) -> None:
|
||||
"""Record the effective hidden memory block for this run.
|
||||
|
||||
|
||||
@ -23,6 +23,6 @@
|
||||
|
||||
**Isolated-loop callback boundary**: sync delegation from an active event loop and `execute_async()` copy the ambient ContextVars into the persistent subagent loop so checkpoint lineage, user identity, tracing context, tags, metadata, and LangGraph's namespaced message-stream handler survive. Before submission, `_copy_isolated_subagent_context()` copies the callback manager/list and removes only handlers marked `deerflow_loop_bound`; `RunJournal` carries that marker because it owns parent-loop tasks and a SQL store/pool. LangGraph merges inherited callbacks with the child run's explicit `SubagentTokenCollector`/tracing callbacks, so letting `RunJournal` cross loops causes duplicate accounting and `Future attached to a different loop` failures, while dropping the whole callback chain silently removes child token frames. Do not replace the boundary with a blank `Context`; the inherited checkpoint namespace and framework stream callback are required by the stream-isolation contract above.
|
||||
|
||||
Loop-detection audit records cross this boundary through a deliberately narrow exception: `task_tool` captures the parent loop and passes `_ParentLoopMiddlewareRecorderProxy` to `SubagentExecutor`, which installs it under the loop-detection-only context key and schedules the real journal append back onto its owner loop with `call_soon_threadsafe`. Other subagent middleware consumers still do not see `__run_journal`. The task tool closes the proxy before returning: close fences later child events and yields once on the owner loop so every previously accepted append reaches the journal before the parent run captures completion data. The proxy, not `RunJournal`, crosses into the isolated loop; do not broaden it into a generic journal facade or call the event store from the subagent loop. Durable batch subagents have no parent run journal and therefore do not use this bridge.
|
||||
Loop-detection and deferred-tool-promotion audit records cross this boundary through a deliberately narrow exception: `task_tool` captures the parent loop and passes one `_ParentLoopMiddlewareRecorderProxy` to `SubagentExecutor`, which installs it under separate loop-detection and tool-promotion context keys and schedules the real journal append back onto its owner loop with `call_soon_threadsafe`. The proxy exposes only `record_middleware` plus an execution-local atomic promotion claim; that claim deduplicates overlapping parallel `tool_search` calls that read the same pre-step state. Other subagent middleware consumers still do not see `__run_journal`. The task tool closes the shared proxy once before returning: close fences later child events and yields once on the owner loop so every previously accepted append reaches the journal before the parent run captures completion data. The proxy, not `RunJournal`, crosses into the isolated loop; do not broaden it into a generic journal facade or call the event store from the subagent loop. Durable batch subagents have no parent run journal and therefore do not use this bridge.
|
||||
|
||||
**Reverse direction of the loop boundary — deferred cleanup & final usage delivery (#5069)**: when a task-tool poller exits unexpectedly, the registry cleanup is pinned **to** the persistent subagent loop via the public `run_on_isolated_subagent_loop()` (executor) so it survives caller-loop teardown — `asyncio.run()` cancels caller-loop tasks on exit, so a caller-loop `asyncio.create_task` would be cancelled before running. The final usage report crosses the boundary the **other way**: `_schedule_deferred_subagent_cleanup` captures the parent run's loop at unwind time (alive in every path that continues the run — the polling-timeout branch returns normally, and a generic poller error becomes an error `ToolMessage`), and `_deliver_final_usage_report` hands the report back onto that loop with `call_soon_threadsafe`. `record_external_llm_usage_records` must never be invoked from the persistent loop or a worker thread (`to_thread`): the journal's accumulators are unlocked read-modify-write fields and `get_completion_data()` iterates `_tokens_by_model`, so a cross-thread write silently loses token updates or breaks iteration mid-run — calling `_report_subagent_usage` directly inside `_deferred_cleanup_subagent_task` would reintroduce exactly this race. The deferred cleaner captures only the resolved usage recorder (plus ids and the captured report loop) — never the whole `runtime`: the strongly-referenced cleanup task lives for up to the full poll budget, and through `runtime` it would pin the parent run's journal and event store for that entire window. When the captured parent loop is already closed (synchronous `asyncio.run` teardown), the report is dropped on purpose — the run has persisted its completion data and nothing reads the counters back — and logged at info with the execution id and unaccounted record count, because the registry entry is removed right after and those records exist nowhere else.
|
||||
|
||||
@ -26,7 +26,7 @@ from langchain_core.runnables import RunnableConfig
|
||||
from langchain_core.runnables.config import var_child_runnable_config
|
||||
from langgraph.errors import GraphRecursionError
|
||||
|
||||
from deerflow.agents.middlewares.audit_context import LOOP_DETECTION_RECORDER_CONTEXT_KEY
|
||||
from deerflow.agents.middlewares.audit_context import LOOP_DETECTION_RECORDER_CONTEXT_KEY, TOOL_PROMOTION_RECORDER_CONTEXT_KEY
|
||||
from deerflow.agents.thread_state import SandboxState, ThreadDataState, ThreadState
|
||||
from deerflow.authz.principal import normalize_authz_attributes
|
||||
from deerflow.config import get_app_config
|
||||
@ -789,6 +789,7 @@ class SubagentExecutor:
|
||||
execution_capacity: SubagentExecutionCapacity | None = None,
|
||||
acceptance_criteria: list[str] | None = None,
|
||||
loop_detection_recorder: Any | None = None,
|
||||
tool_promotion_recorder: Any | None = None,
|
||||
):
|
||||
"""Initialize the executor.
|
||||
|
||||
@ -836,6 +837,8 @@ class SubagentExecutor:
|
||||
parent task tool. Native subagents execute on a separate event
|
||||
loop, so this must be a proxy rather than the parent
|
||||
``RunJournal`` itself.
|
||||
tool_promotion_recorder: Optional loop-safe recorder for deferred-tool
|
||||
promotion events. It follows the same isolated-loop boundary.
|
||||
"""
|
||||
self.config = config
|
||||
self.app_config = app_config
|
||||
@ -884,6 +887,7 @@ class SubagentExecutor:
|
||||
# in report_contract.render_acceptance_criteria_block.
|
||||
self.acceptance_criteria = acceptance_criteria
|
||||
self.loop_detection_recorder = loop_detection_recorder
|
||||
self.tool_promotion_recorder = tool_promotion_recorder
|
||||
|
||||
self._base_tools = _filter_tools(
|
||||
tools,
|
||||
@ -1494,6 +1498,8 @@ class SubagentExecutor:
|
||||
context["agent_id"] = self.config.name
|
||||
if self.loop_detection_recorder is not None:
|
||||
context[LOOP_DETECTION_RECORDER_CONTEXT_KEY] = self.loop_detection_recorder
|
||||
if self.tool_promotion_recorder is not None:
|
||||
context[TOOL_PROMOTION_RECORDER_CONTEXT_KEY] = self.tool_promotion_recorder
|
||||
|
||||
logger.info(f"[trace={self.trace_id}] Subagent {self.config.name} starting async execution with max_turns={self.config.max_turns}")
|
||||
|
||||
|
||||
@ -15,6 +15,8 @@
|
||||
- `batch_task`, `batch_status`, `cancel_batch` - Explicit durable batch submission/progress/cancellation. Added only while the startup SQL-backed batch submitter is installed; large results stay in the owner-scoped API/JSONL export rather than the lead context.
|
||||
- Direct `create_deerflow_agent` integrations receive cloned tools bound to their explicit `SubagentRuntime`. The bound `task` forwards that runtime's exact execution controller and optional caller-owned `AppConfig` into registry/model/tool resolution and `SubagentExecutor`; bound batch tools use the same config snapshot and resolve only that runtime's submitter before falling back to no other application's active worker. Keep the original tool name/schema unchanged so model contracts and user-tool deduplication remain stable.
|
||||
|
||||
The ordinary `task` boundary carries one narrow parent-loop middleware recorder into the isolated subagent runtime under separate loop-detection and tool-promotion keys. It schedules only `record_middleware` calls back onto the loop that owns `RunJournal`, keeps an execution-local atomic promotion claim so parallel searches do not double-report one new schema, is fenced and drained once before `task` returns, and never exposes the journal or event store to the child loop. Durable batch tasks have no parent run journal and do not use this bridge.
|
||||
|
||||
Scheduled-task runtime note:
|
||||
- Scheduled background runs set `context.non_interactive=true` and therefore exclude `ask_clarification` from the lead-agent tool list. This keeps scheduler-triggered runs from stalling on human confirmation mid-execution. `non_interactive` is an internal-only context key: it is merged from `body.context` only when the request authenticated as the process-internal user (the scheduler path), never from arbitrary HTTP/IM clients.
|
||||
|
||||
|
||||
@ -86,7 +86,7 @@ def _record_middleware_on_parent_loop(journal: Any, kwargs: dict[str, Any]) -> N
|
||||
|
||||
|
||||
class _ParentLoopMiddlewareRecorderProxy:
|
||||
"""Forward subagent loop-detection events to the parent run's event loop.
|
||||
"""Forward narrowly scoped subagent middleware events to the parent loop.
|
||||
|
||||
``RunJournal`` owns parent-loop tasks and may wrap an event store backed by
|
||||
a loop-bound SQL pool. Subagents execute on a persistent isolated loop, so
|
||||
@ -98,6 +98,17 @@ class _ParentLoopMiddlewareRecorderProxy:
|
||||
self._loop = loop
|
||||
self._state_lock = threading.Lock()
|
||||
self._closed = False
|
||||
self._claimed_tool_promotions: set[str] = set()
|
||||
|
||||
def claim_tool_promotions(self, tool_names: list[str]) -> list[str]:
|
||||
"""Atomically deduplicate promotions within this one child execution."""
|
||||
candidates = sorted(set(tool_names))
|
||||
with self._state_lock:
|
||||
if self._closed:
|
||||
return []
|
||||
claimed = [name for name in candidates if name not in self._claimed_tool_promotions]
|
||||
self._claimed_tool_promotions.update(claimed)
|
||||
return claimed
|
||||
|
||||
def record_middleware(self, **kwargs: Any) -> None:
|
||||
with self._state_lock:
|
||||
@ -887,17 +898,18 @@ async def task_tool(
|
||||
# system-channel authority over framework instructions.
|
||||
"acceptance_criteria": acceptance_criteria,
|
||||
}
|
||||
loop_detection_recorder = None
|
||||
middleware_recorder = None
|
||||
parent_journal = parent_context.get("__run_journal")
|
||||
if parent_journal is not None:
|
||||
# The task tool runs on the parent run's loop. Pass only a proxy across
|
||||
# the isolated-subagent boundary so middleware persistence is delivered
|
||||
# on the loop that owns the RunJournal and its event store.
|
||||
loop_detection_recorder = _ParentLoopMiddlewareRecorderProxy(
|
||||
middleware_recorder = _ParentLoopMiddlewareRecorderProxy(
|
||||
parent_journal,
|
||||
asyncio.get_running_loop(),
|
||||
)
|
||||
executor_kwargs["loop_detection_recorder"] = loop_detection_recorder
|
||||
executor_kwargs["loop_detection_recorder"] = middleware_recorder
|
||||
executor_kwargs["tool_promotion_recorder"] = middleware_recorder
|
||||
if resolved_app_config is not None:
|
||||
executor_kwargs["app_config"] = resolved_app_config
|
||||
if run_extensions is not None:
|
||||
@ -1189,5 +1201,5 @@ async def task_tool(
|
||||
raise asyncio.CancelledError
|
||||
raise
|
||||
finally:
|
||||
if loop_detection_recorder is not None:
|
||||
await loop_detection_recorder.aclose()
|
||||
if middleware_recorder is not None:
|
||||
await middleware_recorder.aclose()
|
||||
|
||||
@ -91,12 +91,15 @@ def test_every_duplicate_participant_must_satisfy_the_constraint():
|
||||
|
||||
|
||||
def test_core_constraints_are_declared():
|
||||
from deerflow.agents.middlewares.skill_tool_policy_middleware import SkillToolPolicyMiddleware
|
||||
from deerflow.agents.middlewares.tool_error_handling_middleware import ToolErrorHandlingMiddleware
|
||||
from deerflow.agents.middlewares.tool_progress_middleware import ToolProgressMiddleware
|
||||
from deerflow.agents.middlewares.tool_promotion_audit_middleware import DeferredToolPromotionAuditMiddleware
|
||||
from deerflow.extensions.ordering import core_ordering_constraints
|
||||
|
||||
pairs = {(c.outer, c.inner) for c in core_ordering_constraints()}
|
||||
assert (ToolProgressMiddleware, ToolErrorHandlingMiddleware) in pairs
|
||||
assert (DeferredToolPromotionAuditMiddleware, SkillToolPolicyMiddleware) in pairs
|
||||
|
||||
|
||||
def test_core_constraints_are_a_plain_tuple():
|
||||
|
||||
@ -849,6 +849,8 @@ def test_compiled_skill_policy_chain_filters_schema_and_blocks_execution(monkeyp
|
||||
def test_build_middlewares_places_mcp_routing_before_deferred_filter(monkeypatch):
|
||||
from deerflow.agents.middlewares.deferred_tool_filter_middleware import DeferredToolFilterMiddleware
|
||||
from deerflow.agents.middlewares.mcp_routing_middleware import McpRoutingMiddleware
|
||||
from deerflow.agents.middlewares.skill_tool_policy_middleware import SkillToolPolicyMiddleware
|
||||
from deerflow.agents.middlewares.tool_promotion_audit_middleware import DeferredToolPromotionAuditMiddleware
|
||||
from deerflow.tools.builtins.tool_search import DeferredToolSetup
|
||||
|
||||
app_config = _make_app_config([_make_model("safe-model", supports_thinking=False)], loop_detection=LoopDetectionConfig(enabled=False))
|
||||
@ -870,6 +872,9 @@ def test_build_middlewares_places_mcp_routing_before_deferred_filter(monkeypatch
|
||||
|
||||
routing_idx = next(i for i, middleware in enumerate(middlewares) if isinstance(middleware, McpRoutingMiddleware))
|
||||
filter_idx = next(i for i, middleware in enumerate(middlewares) if isinstance(middleware, DeferredToolFilterMiddleware))
|
||||
audit_idx = next(i for i, middleware in enumerate(middlewares) if isinstance(middleware, DeferredToolPromotionAuditMiddleware))
|
||||
policy_idx = next(i for i, middleware in enumerate(middlewares) if isinstance(middleware, SkillToolPolicyMiddleware))
|
||||
assert audit_idx < policy_idx
|
||||
assert routing_idx < filter_idx
|
||||
|
||||
|
||||
|
||||
@ -1,6 +1,7 @@
|
||||
"""Tests for PR2 MCP routing auto-promotion."""
|
||||
|
||||
import asyncio
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from langchain.agents import create_agent
|
||||
@ -16,6 +17,14 @@ from deerflow.tools.mcp_metadata import tag_mcp_routing, tag_mcp_tool
|
||||
from deerflow.utils.messages import ORIGINAL_USER_CONTENT_KEY
|
||||
|
||||
|
||||
class _Recorder:
|
||||
def __init__(self):
|
||||
self.calls: list[dict] = []
|
||||
|
||||
def record_middleware(self, **kwargs):
|
||||
self.calls.append(kwargs)
|
||||
|
||||
|
||||
@as_tool
|
||||
def active_tool(x: str) -> str:
|
||||
"An always-active tool."
|
||||
@ -137,6 +146,110 @@ def test_before_model_returns_minimal_promoted_update_and_reducer_unions():
|
||||
}
|
||||
|
||||
|
||||
def test_auto_promotion_records_only_new_names_without_sensitive_routing_data():
|
||||
recorder = _Recorder()
|
||||
runtime = SimpleNamespace(context={"__run_journal": recorder})
|
||||
middleware = McpRoutingMiddleware(
|
||||
{
|
||||
"postgres_query": {"priority": 100, "keywords": ["secret-orders-keyword"]},
|
||||
"metrics_query": {"priority": 90, "keywords": ["secret-metrics-keyword"]},
|
||||
},
|
||||
"private-catalog-hash",
|
||||
3,
|
||||
)
|
||||
state = {
|
||||
"messages": [HumanMessage(content="secret-orders-keyword secret-metrics-keyword")],
|
||||
"promoted": {"catalog_hash": "private-catalog-hash", "names": ["metrics_query"]},
|
||||
}
|
||||
|
||||
update = middleware.before_model(state, runtime)
|
||||
|
||||
assert update == {"promoted": {"catalog_hash": "private-catalog-hash", "names": ["postgres_query", "metrics_query"]}}
|
||||
assert recorder.calls == [
|
||||
{
|
||||
"tag": "tool_promotion",
|
||||
"name": "McpRoutingMiddleware",
|
||||
"hook": "before_model",
|
||||
"action": "promote",
|
||||
"changes": {
|
||||
"source": "routing_hint",
|
||||
"tool_names": ["postgres_query"],
|
||||
"count": 1,
|
||||
"is_subagent": False,
|
||||
"agent_id": None,
|
||||
},
|
||||
}
|
||||
]
|
||||
persisted = repr(recorder.calls)
|
||||
assert "secret-orders-keyword" not in persisted
|
||||
assert "secret-metrics-keyword" not in persisted
|
||||
assert "private-catalog-hash" not in persisted
|
||||
|
||||
state["promoted"] = update["promoted"]
|
||||
middleware.before_model(state, runtime)
|
||||
assert len(recorder.calls) == 1
|
||||
|
||||
# The same bare name under an old catalog does not prove a promotion in
|
||||
# the active catalog, so catalog drift starts a new effective set.
|
||||
state["promoted"] = {"catalog_hash": "stale-hash", "names": ["postgres_query"]}
|
||||
middleware.before_model(state, runtime)
|
||||
assert len(recorder.calls) == 2
|
||||
assert recorder.calls[-1]["changes"]["tool_names"] == ["metrics_query", "postgres_query"]
|
||||
|
||||
|
||||
def test_auto_promotion_uses_narrow_subagent_recorder_and_is_fail_open(caplog):
|
||||
recorder = _Recorder()
|
||||
|
||||
class BrokenJournal:
|
||||
def record_middleware(self, **kwargs):
|
||||
raise RuntimeError("event store unavailable")
|
||||
|
||||
middleware = McpRoutingMiddleware(
|
||||
{"postgres_query": {"priority": 100, "keywords": ["orders"]}},
|
||||
"hash1",
|
||||
3,
|
||||
)
|
||||
runtime = SimpleNamespace(
|
||||
context={
|
||||
"is_subagent": True,
|
||||
"agent_id": "researcher",
|
||||
"__run_tool_promotion_recorder": recorder,
|
||||
"__run_journal": BrokenJournal(),
|
||||
}
|
||||
)
|
||||
|
||||
update = middleware.before_model({"messages": [HumanMessage(content="orders")]}, runtime)
|
||||
|
||||
assert update == {"promoted": {"catalog_hash": "hash1", "names": ["postgres_query"]}}
|
||||
assert recorder.calls[0]["changes"]["is_subagent"] is True
|
||||
assert recorder.calls[0]["changes"]["agent_id"] == "researcher"
|
||||
|
||||
runtime.context["__run_tool_promotion_recorder"] = BrokenJournal()
|
||||
with caplog.at_level("WARNING"):
|
||||
assert middleware.before_model({"messages": [HumanMessage(content="orders")]}, runtime) == update
|
||||
assert "Failed to record middleware:tool_promotion event" in caplog.text
|
||||
|
||||
|
||||
def test_malformed_existing_promotion_state_cannot_break_routing():
|
||||
recorder = _Recorder()
|
||||
middleware = McpRoutingMiddleware(
|
||||
{"postgres_query": {"priority": 100, "keywords": ["orders"]}},
|
||||
"hash1",
|
||||
3,
|
||||
)
|
||||
|
||||
update = middleware.before_model(
|
||||
{
|
||||
"messages": [HumanMessage(content="orders")],
|
||||
"promoted": {"catalog_hash": "hash1", "names": [{"not": "a tool name"}]},
|
||||
},
|
||||
SimpleNamespace(context={"__run_journal": recorder}),
|
||||
)
|
||||
|
||||
assert update == {"promoted": {"catalog_hash": "hash1", "names": ["postgres_query"]}}
|
||||
assert recorder.calls[0]["changes"]["tool_names"] == ["postgres_query"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_abefore_model_matches_sync_behavior():
|
||||
middleware = McpRoutingMiddleware(
|
||||
|
||||
@ -18,6 +18,7 @@ from deerflow.runtime.events.catalog import (
|
||||
MIDDLEWARE_EVENT_PATTERN,
|
||||
MIDDLEWARE_EVENT_TAG_MAX_LENGTH,
|
||||
MIDDLEWARE_EVENT_TAGS,
|
||||
MIDDLEWARE_TOOL_PROMOTION_TAG,
|
||||
RUN_EVENT_CATEGORY_MAX_LENGTH,
|
||||
RUN_EVENT_TYPE_MAX_LENGTH,
|
||||
SUBAGENT_RUN_EVENT_DEFINITIONS,
|
||||
@ -367,6 +368,15 @@ def test_dynamic_middleware_event_rejects_tags_that_do_not_fit_persistence(tag):
|
||||
)
|
||||
|
||||
|
||||
def test_tool_promotion_tag_is_declared_and_fits_the_persisted_event_type():
|
||||
pattern = _load_contract()["dynamic_event_patterns"][0]
|
||||
|
||||
assert MIDDLEWARE_TOOL_PROMOTION_TAG == "tool_promotion"
|
||||
assert MIDDLEWARE_TOOL_PROMOTION_TAG in MIDDLEWARE_EVENT_TAGS
|
||||
assert MIDDLEWARE_TOOL_PROMOTION_TAG in pattern["known_tags"]
|
||||
assert len(MIDDLEWARE_EVENT_PATTERN.event_type(MIDDLEWARE_TOOL_PROMOTION_TAG)) <= RUN_EVENT_TYPE_MAX_LENGTH
|
||||
|
||||
|
||||
def test_subagent_observed_events_exactly_match_its_catalog_and_payloads():
|
||||
long_result = "r" * (SUBAGENT_STEP_MAX_CHARS + 1)
|
||||
long_error = "e" * (SUBAGENT_STEP_MAX_CHARS + 1)
|
||||
|
||||
@ -5,6 +5,8 @@ Uses MemoryRunEventStore as the backend for direct event inspection.
|
||||
|
||||
import asyncio
|
||||
import weakref
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from threading import Barrier
|
||||
from unittest.mock import MagicMock
|
||||
from uuid import uuid4
|
||||
|
||||
@ -20,6 +22,20 @@ def test_run_journal_is_marked_as_loop_bound():
|
||||
assert RunJournal.deerflow_loop_bound is True
|
||||
|
||||
|
||||
def test_tool_promotion_claim_is_atomic_across_parallel_sync_wrappers():
|
||||
journal = RunJournal("r-claim", "t-claim", MemoryRunEventStore())
|
||||
barrier = Barrier(16)
|
||||
|
||||
def claim():
|
||||
barrier.wait()
|
||||
return journal.claim_tool_promotions(["mcp_a"])
|
||||
|
||||
with ThreadPoolExecutor(max_workers=16) as pool:
|
||||
results = list(pool.map(lambda _: claim(), range(16)))
|
||||
|
||||
assert sum((result for result in results), []) == ["mcp_a"]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_close_flushes_and_detaches_runtime_dependencies():
|
||||
class ProgressReporter:
|
||||
|
||||
@ -11,6 +11,7 @@ from pydantic import Field
|
||||
from deerflow.agents.middlewares.deferred_tool_filter_middleware import DeferredToolFilterMiddleware
|
||||
from deerflow.agents.middlewares.mcp_routing_middleware import McpRoutingMiddleware
|
||||
from deerflow.agents.middlewares.skill_tool_policy_middleware import SkillToolPolicyMiddleware
|
||||
from deerflow.agents.middlewares.tool_promotion_audit_middleware import DeferredToolPromotionAuditMiddleware
|
||||
from deerflow.agents.thread_state import ThreadState
|
||||
from deerflow.runtime.secret_context import SKILL_TOOL_POLICY_DECISION_CONTEXT_KEY, write_slash_skill_source_path
|
||||
from deerflow.runtime.serialization import serialize
|
||||
@ -24,6 +25,14 @@ _CALC_CALLS: list[str] = []
|
||||
_DENIED_CALLS: list[str] = []
|
||||
|
||||
|
||||
class _Recorder:
|
||||
def __init__(self):
|
||||
self.calls: list[dict] = []
|
||||
|
||||
def record_middleware(self, **kwargs):
|
||||
self.calls.append(kwargs)
|
||||
|
||||
|
||||
@tool
|
||||
def calc(expression: str) -> str:
|
||||
"""Evaluate an arithmetic expression."""
|
||||
@ -158,6 +167,8 @@ def test_passive_empty_skill_policy_preserves_deferred_mcp_discovery_and_calling
|
||||
def test_active_skill_can_search_promote_and_call_allowed_deferred_tool():
|
||||
restricted = _skill("restricted", ["calc"])
|
||||
policy, context = _active_policy(restricted)
|
||||
recorder = _Recorder()
|
||||
context["__run_journal"] = recorder
|
||||
setup = _deferred_setup()
|
||||
model = _RecordingModel(
|
||||
[
|
||||
@ -191,6 +202,7 @@ def test_active_skill_can_search_promote_and_call_allowed_deferred_tool():
|
||||
model=model,
|
||||
tools=[calc, denied_lookup, setup.tool_search_tool],
|
||||
middleware=[
|
||||
DeferredToolPromotionAuditMiddleware(setup.deferred_names, setup.catalog_hash),
|
||||
policy,
|
||||
DeferredToolFilterMiddleware(setup.deferred_names, setup.catalog_hash),
|
||||
],
|
||||
@ -213,6 +225,8 @@ def test_active_skill_can_search_promote_and_call_allowed_deferred_tool():
|
||||
assert len(search_result) == 1
|
||||
assert '"name": "calc"' in search_result[0].content
|
||||
assert "denied_lookup" not in search_result[0].content
|
||||
assert recorder.calls[0]["changes"]["tool_names"] == ["calc"]
|
||||
assert "denied_lookup" not in recorder.calls[0]["changes"]["tool_names"]
|
||||
|
||||
|
||||
def test_tool_search_promotion_cannot_expose_or_execute_denied_deferred_tool():
|
||||
|
||||
@ -3840,6 +3840,7 @@ class TestSubagentGuardrailAttribution:
|
||||
oauth_id=None,
|
||||
run_id=None,
|
||||
loop_detection_recorder=None,
|
||||
tool_promotion_recorder=None,
|
||||
name="general-purpose",
|
||||
parent_model="test-model",
|
||||
):
|
||||
@ -3864,6 +3865,7 @@ class TestSubagentGuardrailAttribution:
|
||||
oauth_id=oauth_id,
|
||||
run_id=run_id,
|
||||
loop_detection_recorder=loop_detection_recorder,
|
||||
tool_promotion_recorder=tool_promotion_recorder,
|
||||
)
|
||||
|
||||
@pytest.mark.anyio
|
||||
@ -3930,6 +3932,32 @@ class TestSubagentGuardrailAttribution:
|
||||
assert "__run_journal" not in context
|
||||
assert context.get("agent_id") == "general-purpose"
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_aexecute_propagates_narrow_tool_promotion_recorder(
|
||||
self,
|
||||
classes,
|
||||
executor_module,
|
||||
monkeypatch,
|
||||
):
|
||||
"""Promotion audit crosses the child-loop boundary without the raw journal."""
|
||||
recorder = object()
|
||||
executor = self._make_executor(
|
||||
classes,
|
||||
run_id="run-42",
|
||||
tool_promotion_recorder=recorder,
|
||||
)
|
||||
fake_agent = _FakeStreamAgent()
|
||||
monkeypatch.setattr(executor, "_build_initial_state", self._noop_build_initial_state)
|
||||
monkeypatch.setattr(executor, "_create_agent", lambda *a, **kw: fake_agent)
|
||||
|
||||
await executor._aexecute("do something")
|
||||
|
||||
context = fake_agent.captured_context
|
||||
assert context is not None
|
||||
assert context.get("__run_tool_promotion_recorder") is recorder
|
||||
assert "__run_journal" not in context
|
||||
assert context.get("agent_id") == "general-purpose"
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_aexecute_propagates_channel_user_id_to_subagent_context(
|
||||
self,
|
||||
|
||||
@ -58,6 +58,8 @@ def test_parent_loop_middleware_recorder_proxy_delivers_on_owner_loop():
|
||||
LoopPinnedJournal(),
|
||||
parent_loop,
|
||||
)
|
||||
assert proxy.claim_tool_promotions(["mcp_b", "mcp_a", "mcp_a"]) == ["mcp_a", "mcp_b"]
|
||||
assert proxy.claim_tool_promotions(["mcp_a"]) == []
|
||||
proxy.record_middleware(
|
||||
tag="loop_detection",
|
||||
name="LoopDetectionMiddleware",
|
||||
@ -74,6 +76,7 @@ def test_parent_loop_middleware_recorder_proxy_delivers_on_owner_loop():
|
||||
assert kwargs["action"] == "warn"
|
||||
|
||||
asyncio.run_coroutine_threadsafe(proxy.aclose(), parent_loop).result(timeout=5)
|
||||
assert proxy.claim_tool_promotions(["late_tool"]) == []
|
||||
proxy.record_middleware(tag="loop_detection", name="LoopDetectionMiddleware", hook="after_model", action="hard_stop", changes={})
|
||||
time.sleep(0.05)
|
||||
assert len(calls) == 1, "events emitted after the parent task boundary must be dropped"
|
||||
@ -429,7 +432,7 @@ def test_task_tool_forwards_the_run_extension_snapshot_to_executor(monkeypatch):
|
||||
assert captured["executor_kwargs"]["extensions"] is loaded
|
||||
|
||||
|
||||
def test_task_tool_installs_and_closes_narrow_loop_detection_recorder(monkeypatch):
|
||||
def test_task_tool_installs_and_closes_narrow_middleware_recorder(monkeypatch):
|
||||
journal = MagicMock()
|
||||
runtime = _make_runtime()
|
||||
runtime.context["__run_journal"] = journal
|
||||
@ -458,6 +461,7 @@ def test_task_tool_installs_and_closes_narrow_loop_detection_recorder(monkeypatc
|
||||
|
||||
kwargs = captured["executor_kwargs"]
|
||||
proxy = kwargs["loop_detection_recorder"]
|
||||
assert kwargs["tool_promotion_recorder"] is proxy
|
||||
assert proxy.is_closed is True
|
||||
proxy.record_middleware(tag="loop_detection", name="LoopDetectionMiddleware", hook="after_model", action="warn", changes={})
|
||||
journal.record_middleware.assert_not_called()
|
||||
|
||||
@ -642,6 +642,8 @@ def test_subagent_runtime_middlewares_attach_deferred_filter_when_setup_has_name
|
||||
|
||||
from deerflow.agents.middlewares.deferred_tool_filter_middleware import DeferredToolFilterMiddleware
|
||||
from deerflow.agents.middlewares.safety_finish_reason_middleware import SafetyFinishReasonMiddleware
|
||||
from deerflow.agents.middlewares.skill_tool_policy_middleware import SkillToolPolicyMiddleware
|
||||
from deerflow.agents.middlewares.tool_promotion_audit_middleware import DeferredToolPromotionAuditMiddleware
|
||||
from deerflow.tools.builtins.tool_search import build_deferred_tool_setup
|
||||
from deerflow.tools.mcp_metadata import tag_mcp_tool
|
||||
|
||||
@ -659,9 +661,14 @@ def test_subagent_runtime_middlewares_attach_deferred_filter_when_setup_has_name
|
||||
middlewares = build_subagent_runtime_middlewares(app_config=app_config, deferred_setup=setup)
|
||||
|
||||
filters = [m for m in middlewares if isinstance(m, DeferredToolFilterMiddleware)]
|
||||
audits = [m for m in middlewares if isinstance(m, DeferredToolPromotionAuditMiddleware)]
|
||||
assert len(filters) == 1
|
||||
assert len(audits) == 1
|
||||
audit_idx = next(i for i, m in enumerate(middlewares) if isinstance(m, DeferredToolPromotionAuditMiddleware))
|
||||
policy_idx = next(i for i, m in enumerate(middlewares) if isinstance(m, SkillToolPolicyMiddleware))
|
||||
filter_idx = next(i for i, m in enumerate(middlewares) if isinstance(m, DeferredToolFilterMiddleware))
|
||||
safety_idx = next(i for i, m in enumerate(middlewares) if isinstance(m, SafetyFinishReasonMiddleware))
|
||||
assert audit_idx < policy_idx < filter_idx
|
||||
assert filter_idx < safety_idx
|
||||
|
||||
|
||||
@ -703,6 +710,7 @@ def test_subagent_runtime_middlewares_place_mcp_routing_before_deferred_filter(m
|
||||
def test_subagent_runtime_middlewares_skip_deferred_filter_without_names(monkeypatch):
|
||||
"""No deferred setup (disabled / no MCP tool) -> no DeferredToolFilterMiddleware."""
|
||||
from deerflow.agents.middlewares.deferred_tool_filter_middleware import DeferredToolFilterMiddleware
|
||||
from deerflow.agents.middlewares.tool_promotion_audit_middleware import DeferredToolPromotionAuditMiddleware
|
||||
from deerflow.tools.builtins.tool_search import DeferredToolSetup
|
||||
|
||||
app_config = _make_app_config()
|
||||
@ -711,6 +719,7 @@ def test_subagent_runtime_middlewares_skip_deferred_filter_without_names(monkeyp
|
||||
for setup in (None, DeferredToolSetup(None, frozenset(), None)):
|
||||
middlewares = build_subagent_runtime_middlewares(app_config=app_config, deferred_setup=setup)
|
||||
assert not any(isinstance(m, DeferredToolFilterMiddleware) for m in middlewares)
|
||||
assert not any(isinstance(m, DeferredToolPromotionAuditMiddleware) for m in middlewares)
|
||||
|
||||
|
||||
def test_subagent_runtime_middlewares_attach_loop_detection_when_enabled(monkeypatch):
|
||||
|
||||
189
backend/tests/test_tool_promotion_audit_middleware.py
Normal file
189
backend/tests/test_tool_promotion_audit_middleware.py
Normal file
@ -0,0 +1,189 @@
|
||||
"""Tests for deferred-tool promotion audit events."""
|
||||
|
||||
import asyncio
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from langchain_core.messages import ToolMessage
|
||||
from langgraph.types import Command
|
||||
|
||||
from deerflow.agents.middlewares.tool_promotion_audit_middleware import DeferredToolPromotionAuditMiddleware
|
||||
from deerflow.runtime.events.store.memory import MemoryRunEventStore
|
||||
from deerflow.runtime.journal import RunJournal
|
||||
|
||||
|
||||
class _Recorder:
|
||||
def __init__(self):
|
||||
self.calls: list[dict] = []
|
||||
self.claimed: set[str] = set()
|
||||
|
||||
def claim_tool_promotions(self, tool_names):
|
||||
names = sorted(set(tool_names) - self.claimed)
|
||||
self.claimed.update(names)
|
||||
return names
|
||||
|
||||
def record_middleware(self, **kwargs):
|
||||
self.calls.append(kwargs)
|
||||
|
||||
|
||||
class _ToolRequest:
|
||||
def __init__(self, *, name="tool_search", state=None, context=None, query="private query"):
|
||||
self.tool_call = {"name": name, "id": "tc1", "args": {"query": query}}
|
||||
self.state = state or {}
|
||||
self.runtime = SimpleNamespace(context=context or {})
|
||||
|
||||
|
||||
def _middleware():
|
||||
return DeferredToolPromotionAuditMiddleware(frozenset({"mcp_a", "mcp_b"}), "h1")
|
||||
|
||||
|
||||
def test_records_only_new_names_from_the_final_current_catalog_command():
|
||||
recorder = _Recorder()
|
||||
request = _ToolRequest(
|
||||
state={"promoted": {"catalog_hash": "h1", "names": ["mcp_b"]}},
|
||||
context={"__run_journal": recorder},
|
||||
query="credential-adjacent query",
|
||||
)
|
||||
result = Command(
|
||||
update={
|
||||
"promoted": {"catalog_hash": "h1", "names": ["not_deferred", "mcp_b", "mcp_a", "mcp_a"]},
|
||||
"messages": [ToolMessage(content="private schema", tool_call_id="tc1", name="tool_search")],
|
||||
}
|
||||
)
|
||||
|
||||
assert _middleware().wrap_tool_call(request, lambda _: result) is result
|
||||
|
||||
assert recorder.calls == [
|
||||
{
|
||||
"tag": "tool_promotion",
|
||||
"name": "DeferredToolPromotionAuditMiddleware",
|
||||
"hook": "wrap_tool_call",
|
||||
"action": "promote",
|
||||
"changes": {
|
||||
"source": "tool_search",
|
||||
"tool_names": ["mcp_a"],
|
||||
"count": 1,
|
||||
"is_subagent": False,
|
||||
"agent_id": None,
|
||||
},
|
||||
}
|
||||
]
|
||||
persisted = repr(recorder.calls)
|
||||
assert "credential-adjacent query" not in persisted
|
||||
assert "private schema" not in persisted
|
||||
assert "h1" not in persisted
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_records_promotion_but_repeated_stale_and_non_search_results_do_not():
|
||||
recorder = _Recorder()
|
||||
middleware = _middleware()
|
||||
request = _ToolRequest(context={"__run_journal": recorder})
|
||||
promoted = Command(update={"promoted": {"catalog_hash": "h1", "names": ["mcp_a"]}})
|
||||
|
||||
async def handle_promoted(_):
|
||||
return promoted
|
||||
|
||||
assert await middleware.awrap_tool_call(request, handle_promoted) is promoted
|
||||
assert recorder.calls[0]["changes"]["tool_names"] == ["mcp_a"]
|
||||
|
||||
request.state = {"promoted": {"catalog_hash": "h1", "names": ["mcp_a"]}}
|
||||
|
||||
async def handle_stale(_):
|
||||
return Command(update={"promoted": {"catalog_hash": "stale", "names": ["mcp_b"]}})
|
||||
|
||||
assert await middleware.awrap_tool_call(request, handle_promoted) is promoted
|
||||
assert await middleware.awrap_tool_call(request, handle_stale)
|
||||
request.tool_call["name"] = "another_tool"
|
||||
assert await middleware.awrap_tool_call(request, handle_promoted) is promoted
|
||||
assert len(recorder.calls) == 1
|
||||
|
||||
|
||||
def test_final_handler_result_is_observed_without_rebuilding_it():
|
||||
"""An outer audit wrapper must see names after inner policy filtering."""
|
||||
recorder = _Recorder()
|
||||
request = _ToolRequest(context={"__run_journal": recorder})
|
||||
policy_filtered = Command(update={"promoted": {"catalog_hash": "h1", "names": ["mcp_a"]}})
|
||||
|
||||
observed = _middleware().wrap_tool_call(request, lambda _: policy_filtered)
|
||||
|
||||
assert observed is policy_filtered
|
||||
assert recorder.calls[0]["changes"]["tool_names"] == ["mcp_a"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"result",
|
||||
[
|
||||
ToolMessage(content="no match", tool_call_id="tc1", name="tool_search"),
|
||||
Command(update={}),
|
||||
Command(update={"promoted": {"catalog_hash": "h1", "names": []}}),
|
||||
Command(update={"promoted": {"catalog_hash": "h1", "names": ["mcp_a", 7]}}),
|
||||
],
|
||||
)
|
||||
def test_non_promotion_and_malformed_results_emit_nothing(result):
|
||||
recorder = _Recorder()
|
||||
request = _ToolRequest(context={"__run_journal": recorder})
|
||||
|
||||
assert _middleware().wrap_tool_call(request, lambda _: result) is result
|
||||
assert recorder.calls == []
|
||||
|
||||
|
||||
def test_recorder_failure_does_not_replace_the_tool_result(caplog):
|
||||
class BrokenRecorder:
|
||||
def record_middleware(self, **kwargs):
|
||||
raise RuntimeError("event store unavailable")
|
||||
|
||||
request = _ToolRequest(context={"__run_journal": BrokenRecorder()})
|
||||
result = Command(update={"promoted": {"catalog_hash": "h1", "names": ["mcp_a"]}})
|
||||
|
||||
with caplog.at_level("WARNING"):
|
||||
assert _middleware().wrap_tool_call(request, lambda _: result) is result
|
||||
|
||||
assert "Failed to record middleware:tool_promotion event" in caplog.text
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_event_round_trips_through_run_journal_and_store():
|
||||
store = MemoryRunEventStore()
|
||||
journal = RunJournal("run-1", "thread-1", store, flush_threshold=100)
|
||||
request = _ToolRequest(context={"__run_journal": journal})
|
||||
result = Command(update={"promoted": {"catalog_hash": "h1", "names": ["mcp_a"]}})
|
||||
|
||||
assert _middleware().wrap_tool_call(request, lambda _: result) is result
|
||||
await journal.flush()
|
||||
|
||||
events = await store.list_events("thread-1", "run-1")
|
||||
assert len(events) == 1
|
||||
assert events[0]["event_type"] == "middleware:tool_promotion"
|
||||
assert events[0]["category"] == "middleware"
|
||||
assert events[0]["content"]["changes"] == {
|
||||
"source": "tool_search",
|
||||
"tool_names": ["mcp_a"],
|
||||
"count": 1,
|
||||
"is_subagent": False,
|
||||
"agent_id": None,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_parallel_searches_claim_the_same_new_name_only_once():
|
||||
"""Parallel tool Sends share pre-step state but must not duplicate events."""
|
||||
store = MemoryRunEventStore()
|
||||
journal = RunJournal("run-parallel", "thread-1", store, flush_threshold=100)
|
||||
middleware = _middleware()
|
||||
result = Command(update={"promoted": {"catalog_hash": "h1", "names": ["mcp_a"]}})
|
||||
requests = [
|
||||
_ToolRequest(context={"__run_journal": journal}),
|
||||
_ToolRequest(context={"__run_journal": journal}),
|
||||
]
|
||||
|
||||
async def handler(_):
|
||||
await asyncio.sleep(0)
|
||||
return result
|
||||
|
||||
observed = await asyncio.gather(*(middleware.awrap_tool_call(request, handler) for request in requests))
|
||||
await journal.flush()
|
||||
|
||||
assert observed == [result, result]
|
||||
events = await store.list_events("thread-1", "run-parallel")
|
||||
assert [event["event_type"] for event in events] == ["middleware:tool_promotion"]
|
||||
@ -397,7 +397,7 @@
|
||||
"pattern": "middleware:{tag}",
|
||||
"category": "middleware",
|
||||
"producer": "RunJournal.record_middleware(tag, ...)",
|
||||
"known_tags": ["guardrail", "loop_detection", "safety_termination", "skill_activation", "skill_secrets"],
|
||||
"known_tags": ["guardrail", "loop_detection", "safety_termination", "skill_activation", "skill_secrets", "tool_promotion"],
|
||||
"event_type_schema": {
|
||||
"type": "string",
|
||||
"pattern": "^middleware:",
|
||||
@ -444,7 +444,7 @@
|
||||
{
|
||||
"id": "middleware-coverage",
|
||||
"status": "partial",
|
||||
"notes": "Loop-detection events cover lead-agent and ordinary task-tool subagent runs; durable batch subagent loop detection and deferred-tool promotion do not currently emit middleware events."
|
||||
"notes": "Loop-detection and deferred-tool promotion events cover lead-agent and ordinary task-tool subagent runs. Durable batch subagents have no parent run journal and do not currently emit either event."
|
||||
},
|
||||
{
|
||||
"id": "run-scoped-observation-context",
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user