mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-19 02:56:17 +00:00
14c9d44440
* feat(runtime): persist tool-progress phase transitions Record bounded warn, block, and recover decisions for lead and task subagent runs while preserving event-loop isolation, fail-open behavior, and concurrent transition order. * fix(runtime): trust server-owned tool progress attribution * fix(runtime): centralize trusted audit attribution * fix(runtime): preserve complete tool progress audit state * docs: trim tool progress guidance to pass size check * fix(runtime): fence subagent audit recorder loop * docs(readme): sync tool-progress event coverage 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> Co-authored-by: 嗜鵼 <hy2010hy2010@qq.com> Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
749 lines
32 KiB
Python
749 lines
32 KiB
Python
"""Middleware for task-level tool call progress tracking with a state machine.
|
|
|
|
Implements RFC #3177: structured tool result signals drive a per-(thread, tool)
|
|
state machine that detects stagnation and repetition, injects hints early
|
|
(WARNED), and hard-blocks the tool when it has stopped producing value (BLOCKED).
|
|
|
|
Architecture:
|
|
ToolProgressMiddleware (outer)
|
|
└── handler → ToolErrorHandlingMiddleware (inner) → actual tool
|
|
↓
|
|
ToolProgressMiddleware reads deerflow_tool_meta from the normalized result
|
|
|
|
State machine transitions per (thread_id, tool_name):
|
|
ACTIVE → WARNED (at stagnation_threshold problems)
|
|
Any problem-free call resets consecutive_problems=0 and reverts to ACTIVE.
|
|
|
|
Whether WARNED can escalate to BLOCKED depends on recoverable_by_model:
|
|
- recoverable_by_model=True (no_results, not_found, permission, Jaccard-duplicate success):
|
|
WARNED is terminal. The model received a hint and is expected to change strategy;
|
|
blocking would prevent a legitimate retry with different parameters.
|
|
- recoverable_by_model=False, action≠stop (transient, rate_limited):
|
|
WARNED → BLOCKED after warn_escalation_count more problems. The model cannot fix
|
|
these by retrying the same tool, so hard-blocking conserves API calls.
|
|
- recoverable_by_model=False, action=stop (auth, config, internal):
|
|
Immediately BLOCKED on the first occurrence — no retry can help.
|
|
|
|
Division of labor with LoopDetectionMiddleware (middleware position 23):
|
|
ToolProgressMiddleware (position 10) is a result-quality guard — it fires
|
|
after a tool executes, inspects what came back, and blocks *specific tools*
|
|
that have stopped producing new information.
|
|
|
|
LoopDetectionMiddleware is a call-pattern guard — it fires after the model
|
|
responds (before tools execute), inspects the tool_calls signature in the
|
|
AIMessage, and forces the *whole turn* to stop when the model keeps issuing
|
|
the same calls regardless of results.
|
|
|
|
They are complementary, not competing:
|
|
- ToolProgressMiddleware is fine-grained (per-tool BLOCK, other tools normal).
|
|
- LoopDetectionMiddleware is coarse-grained (strips all tool_calls, ends turn).
|
|
- Both can inject HumanMessage hints in the same model call without conflict;
|
|
the model sees both sets of hints and can reason about them.
|
|
- If LoopDetectionMiddleware hard-stops (strips tool_calls), no wrap_tool_call
|
|
is issued so ToolProgressMiddleware never fires — there is no double-stop.
|
|
- If ToolProgressMiddleware BLOCKs a tool (returns an error ToolMessage),
|
|
the model still makes a tool call that LoopDetectionMiddleware tracks; both
|
|
continue to operate on their own independent state.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import re
|
|
import threading
|
|
from collections import OrderedDict, defaultdict
|
|
from collections.abc import Awaitable, Callable, Sequence
|
|
from dataclasses import dataclass, field, replace
|
|
from typing import TYPE_CHECKING, Literal, 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 HumanMessage, 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_PROGRESS_RECORDER_CONTEXT_KEY,
|
|
resolve_audit_recorder,
|
|
)
|
|
from deerflow.agents.middlewares.tool_result_meta import (
|
|
PROGRESS_GUARD_ERROR_TYPE,
|
|
TOOL_META_KEY,
|
|
TOOL_RESULT_ERROR_TYPES,
|
|
TOOL_RESULT_NEXT_ACTIONS,
|
|
TOOL_RESULT_STATUSES,
|
|
ToolResultMeta,
|
|
)
|
|
from deerflow.runtime.events.catalog import MIDDLEWARE_TOOL_PROGRESS_TAG
|
|
|
|
if TYPE_CHECKING:
|
|
from deerflow.config.tool_progress_config import ToolProgressConfig
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_MAX_PENDING_PER_RUN = 3
|
|
# Jaccard word-set computation is capped to avoid O(n) regex work on very large tool results.
|
|
_MAX_CONTENT_FOR_WORDSET = 8192
|
|
|
|
|
|
def _audit_error_type(value: object) -> str | None:
|
|
"""Project an untrusted tool stamp onto the bounded audit vocabulary."""
|
|
if value is None:
|
|
return None
|
|
return value if type(value) is str and value in TOOL_RESULT_ERROR_TYPES else "unknown"
|
|
|
|
|
|
def _audit_next_action(value: object) -> str:
|
|
"""Return only framework-defined recovery actions to persistence."""
|
|
return value if type(value) is str and value in TOOL_RESULT_NEXT_ACTIONS else "unknown"
|
|
|
|
|
|
def _audit_status(value: object) -> str:
|
|
"""Project the producer status onto the canonical result vocabulary."""
|
|
return value if type(value) is str and value in TOOL_RESULT_STATUSES else "unknown"
|
|
|
|
|
|
def _audit_recoverable(value: object) -> bool | None:
|
|
"""Reject truthy non-booleans from producer-supplied tool metadata."""
|
|
return value if type(value) is bool else None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# State data structures
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class ToolPhaseState:
|
|
"""Per (thread_id, tool_name) tracking state."""
|
|
|
|
phase: Literal["active", "warned", "blocked"] = "active"
|
|
consecutive_problems: int = 0
|
|
block_reason: str | None = None
|
|
# Immutable tuple so that dataclasses.replace() calls that omit recent_word_sets
|
|
# (problem paths) cannot accidentally share a mutable list between the old and new
|
|
# state objects and cause silent cross-state corruption via .append().
|
|
recent_word_sets: tuple[frozenset[str], ...] = field(default_factory=tuple)
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class ToolPhaseTransition:
|
|
"""The exact state-machine rule that produced a durable phase change."""
|
|
|
|
action: Literal["warn", "block", "recover", "reset"]
|
|
threshold: int | None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Content helpers
|
|
|
|
|
|
def word_set(content: str) -> frozenset[str]:
|
|
"""Extract lowercase words of length >= 3 for Jaccard similarity.
|
|
|
|
Content is capped at _MAX_CONTENT_FOR_WORDSET chars to bound memory and CPU cost on
|
|
large tool results (e.g. web pages). Tail content beyond the cap is omitted from the
|
|
set, which is acceptable because duplicate-detection is a heuristic, not a guarantee.
|
|
"""
|
|
return frozenset(re.findall(r"\b\w{3,}\b", content[:_MAX_CONTENT_FOR_WORDSET].lower()))
|
|
|
|
|
|
def is_near_duplicate(
|
|
current: frozenset[str],
|
|
recent: Sequence[frozenset[str]],
|
|
threshold: float,
|
|
min_words: int,
|
|
) -> bool:
|
|
"""Return True if current is similar to any of the last 3 recent word sets."""
|
|
if len(current) < min_words:
|
|
return False
|
|
for prev in recent[-3:]:
|
|
if len(prev) < min_words:
|
|
continue
|
|
union = len(current | prev)
|
|
if union == 0:
|
|
continue
|
|
if len(current & prev) / union >= threshold:
|
|
return True
|
|
return False
|
|
|
|
|
|
def _message_content_str(msg: ToolMessage) -> str:
|
|
return msg.content if isinstance(msg.content, str) else ""
|
|
|
|
|
|
def _result_tool_message(result: ToolMessage | Command, tool_call_id: str) -> ToolMessage | None:
|
|
"""Return the ToolMessage for this tool call, including Command-wrapped results."""
|
|
if isinstance(result, ToolMessage):
|
|
return result
|
|
update = result.update
|
|
if not isinstance(update, dict):
|
|
return None
|
|
messages = update.get("messages", [])
|
|
if isinstance(messages, ToolMessage):
|
|
messages = [messages]
|
|
if not isinstance(messages, (list, tuple)):
|
|
return None
|
|
for message in messages:
|
|
if isinstance(message, ToolMessage) and str(message.tool_call_id) == tool_call_id:
|
|
return message
|
|
return None
|
|
|
|
|
|
def _parse_tool_meta(meta_dict: object) -> ToolResultMeta | None:
|
|
"""Safely deserialize a ToolResultMeta from a raw dict; returns None on schema mismatch."""
|
|
if not isinstance(meta_dict, dict):
|
|
return None
|
|
try:
|
|
return ToolResultMeta(**meta_dict)
|
|
except TypeError:
|
|
logger.warning("Unexpected tool meta schema, skipping progress tracking: %s", meta_dict)
|
|
return None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Hint / block reason formatting
|
|
|
|
|
|
def _format_hint(meta: ToolResultMeta) -> str:
|
|
action_map = {
|
|
"rewrite_query": "Try rephrasing your search query with different keywords or approach.",
|
|
"try_alternative": "Consider using a different tool or strategy.",
|
|
"summarize": "Consider summarizing your current findings and moving forward.",
|
|
"stop": "Do not retry this operation — it is not recoverable.",
|
|
# Near-duplicate success results: recommended_next_action is "continue" by default,
|
|
# but the model should still change strategy to avoid re-fetching the same content.
|
|
"continue": "Try rephrasing your query or using a different search term.",
|
|
}
|
|
base = {
|
|
"no_results": "[PROGRESS HINT] Your search returned no results.",
|
|
"not_found": "[PROGRESS HINT] The resource was not found repeatedly.",
|
|
"rate_limited": "[PROGRESS HINT] The tool is being rate-limited.",
|
|
"transient": "[PROGRESS HINT] The tool encountered repeated transient failures.",
|
|
"partial_success": "[PROGRESS HINT] The tool has returned incomplete results multiple times.",
|
|
# Jaccard near-duplicate success: the tool is returning the same content repeatedly.
|
|
"success": "[PROGRESS HINT] The tool is returning duplicate results.",
|
|
}.get(
|
|
meta.error_type or meta.status,
|
|
"[PROGRESS HINT] The tool is not producing new information.",
|
|
)
|
|
suffix = action_map.get(meta.recommended_next_action, "")
|
|
return f"{base} {suffix}".strip()
|
|
|
|
|
|
def _block_reason(meta: ToolResultMeta) -> str:
|
|
return {
|
|
"no_results": "Repeated no-results — rewrite your query or try a different tool.",
|
|
"not_found": "Repeated not-found — rewrite your query or try a different resource.",
|
|
"rate_limited": "Repeated rate-limiting — summarize current findings and proceed.",
|
|
"transient": "Repeated transient failures — try a different approach.",
|
|
"auth": "Authentication failure — this tool cannot be used.",
|
|
"config": "Tool is not configured — this tool cannot be used.",
|
|
"internal": "Repeated internal errors — this tool is unavailable.",
|
|
}.get(
|
|
meta.error_type or "",
|
|
"Tool has not produced new information after multiple attempts — summarize and move on.",
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Middleware
|
|
|
|
|
|
class ToolProgressMiddleware(AgentMiddleware[AgentState]):
|
|
"""State-machine-based tool stagnation guard (RFC #3177)."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
stagnation_threshold: int = 3,
|
|
warn_escalation_count: int = 2,
|
|
inject_assessment: bool = True,
|
|
jaccard_threshold: float = 0.8,
|
|
min_words: int = 10,
|
|
exempt_tools: set[str] | None = None,
|
|
max_tracked_threads: int = 100,
|
|
) -> None:
|
|
self._stagnation_threshold = stagnation_threshold
|
|
self._warn_escalation = warn_escalation_count
|
|
self._inject_assessment = inject_assessment
|
|
self._jaccard_threshold = jaccard_threshold
|
|
self._min_words = min_words
|
|
self._exempt_tools: set[str] = exempt_tools if exempt_tools is not None else {"ask_clarification", "write_todos", "present_files", "task"}
|
|
self._max_tracked_threads = max_tracked_threads
|
|
|
|
# threading.Lock (not asyncio.Lock) also protects embedded callers that use the
|
|
# synchronous wrapper from multiple threads. Recorder callbacks are deliberately
|
|
# invoked after this state lock is released; observability must not stall tool
|
|
# state updates. This matches LoopDetectionMiddleware's publication convention.
|
|
self._lock = threading.Lock()
|
|
# LRU-evicting store: thread_id → {tool_name → ToolPhaseState}
|
|
self._phase_states: OrderedDict[str, dict[str, ToolPhaseState]] = OrderedDict()
|
|
# Pending hint queue: (thread_id, run_id) → [hint texts]
|
|
self._pending: dict[tuple[str, str], list[str]] = defaultdict(list)
|
|
|
|
@classmethod
|
|
def from_config(cls, config: ToolProgressConfig) -> ToolProgressMiddleware:
|
|
return cls(
|
|
stagnation_threshold=config.stagnation_threshold,
|
|
warn_escalation_count=config.warn_escalation_count,
|
|
inject_assessment=config.inject_assessment,
|
|
jaccard_threshold=config.jaccard_similarity_threshold,
|
|
min_words=config.min_word_count_for_similarity,
|
|
exempt_tools=set(config.exempt_tools),
|
|
max_tracked_threads=config.max_tracked_threads,
|
|
)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Runtime helpers
|
|
|
|
@staticmethod
|
|
def _thread_id(runtime: Runtime) -> str:
|
|
tid = runtime.context.get("thread_id") if runtime.context else None
|
|
return str(tid) if tid else "default"
|
|
|
|
@staticmethod
|
|
def _run_id(runtime: Runtime) -> str:
|
|
rid = runtime.context.get("run_id") if runtime.context else None
|
|
return str(rid) if rid else "default"
|
|
|
|
def _pending_key(self, runtime: Runtime) -> tuple[str, str]:
|
|
return self._thread_id(runtime), self._run_id(runtime)
|
|
|
|
def _record_phase_transition(
|
|
self,
|
|
*,
|
|
runtime: Runtime,
|
|
tool_name: str,
|
|
state: ToolPhaseState,
|
|
new_state: ToolPhaseState,
|
|
meta: ToolResultMeta | None,
|
|
hook: Literal["wrap_tool_call", "awrap_tool_call", "before_agent", "abefore_agent"],
|
|
transition: ToolPhaseTransition,
|
|
) -> None:
|
|
"""Persist one effective transition without copying tool content."""
|
|
recorder, is_subagent, agent_id = resolve_audit_recorder(
|
|
getattr(runtime, "context", None),
|
|
recorder_key=TOOL_PROGRESS_RECORDER_CONTEXT_KEY,
|
|
)
|
|
if recorder is None:
|
|
return
|
|
|
|
try:
|
|
recorder.record_middleware(
|
|
tag=MIDDLEWARE_TOOL_PROGRESS_TAG,
|
|
name=type(self).__name__,
|
|
hook=hook,
|
|
action=transition.action,
|
|
changes={
|
|
"is_subagent": is_subagent,
|
|
"agent_id": agent_id,
|
|
"tool_name": tool_name,
|
|
"from_phase": state.phase,
|
|
"to_phase": new_state.phase,
|
|
"consecutive_problems": new_state.consecutive_problems,
|
|
"status": _audit_status(meta.status) if meta is not None else None,
|
|
"error_type": _audit_error_type(meta.error_type) if meta is not None else None,
|
|
"recoverable_by_model": _audit_recoverable(meta.recoverable_by_model) if meta is not None else None,
|
|
"recommended_next_action": _audit_next_action(meta.recommended_next_action) if meta is not None else None,
|
|
"threshold": transition.threshold,
|
|
},
|
|
)
|
|
except Exception: # noqa: BLE001
|
|
# Observability must never alter the progress guard or agent run.
|
|
logger.warning(
|
|
"Failed to record middleware:tool_progress event",
|
|
exc_info=True,
|
|
)
|
|
|
|
# ------------------------------------------------------------------
|
|
# State store (caller holds lock)
|
|
|
|
def _get_state(self, thread_id: str, tool_name: str) -> ToolPhaseState:
|
|
if thread_id not in self._phase_states:
|
|
self._phase_states[thread_id] = {}
|
|
while len(self._phase_states) > self._max_tracked_threads:
|
|
evicted_thread, _ = self._phase_states.popitem(last=False)
|
|
# Evict pending hints for the evicted thread to prevent unbounded growth.
|
|
for key in [k for k in self._pending if k[0] == evicted_thread]:
|
|
del self._pending[key]
|
|
self._phase_states.move_to_end(thread_id)
|
|
return self._phase_states[thread_id].get(tool_name, ToolPhaseState())
|
|
|
|
def _set_state(self, thread_id: str, tool_name: str, state: ToolPhaseState) -> None:
|
|
self._phase_states[thread_id][tool_name] = state
|
|
|
|
def _get_block_reason(self, runtime: Runtime, tool_name: str) -> str | None:
|
|
thread_id = self._thread_id(runtime)
|
|
with self._lock:
|
|
thread_tools = self._phase_states.get(thread_id)
|
|
if thread_tools is None:
|
|
return None
|
|
# Read-only check: do NOT call move_to_end here. Bumping recency on the read path
|
|
# would keep blocked threads permanently warm in the LRU, preventing healthy active
|
|
# threads from occupying those slots. Recency is updated only on _get_state writes.
|
|
tool_state = thread_tools.get(tool_name)
|
|
return tool_state.block_reason if tool_state is not None and tool_state.phase == "blocked" else None
|
|
|
|
def _make_blocked_message(self, request: ToolCallRequest, tool_name: str, block_reason: str) -> ToolMessage:
|
|
return ToolMessage(
|
|
content=f"[TOOL_BLOCKED] {block_reason}",
|
|
tool_call_id=str(request.tool_call.get("id", "")),
|
|
name=tool_name,
|
|
status="error",
|
|
additional_kwargs={
|
|
TOOL_META_KEY: {
|
|
"status": "error",
|
|
"error_type": PROGRESS_GUARD_ERROR_TYPE,
|
|
"recoverable_by_model": True,
|
|
"recommended_next_action": "summarize",
|
|
"source": "progress_middleware",
|
|
}
|
|
},
|
|
)
|
|
|
|
def _update_state_from_result(
|
|
self,
|
|
result: ToolMessage | Command,
|
|
tool_name: str,
|
|
runtime: Runtime,
|
|
tool_call_id: str,
|
|
*,
|
|
hook: Literal["wrap_tool_call", "awrap_tool_call"],
|
|
) -> ToolMessage | Command:
|
|
"""Update the state machine from a tool result; queue hints if warranted."""
|
|
message = _result_tool_message(result, tool_call_id)
|
|
if message is None:
|
|
return result
|
|
meta = _parse_tool_meta((message.additional_kwargs or {}).get(TOOL_META_KEY))
|
|
if meta is None:
|
|
if tool_name not in self._exempt_tools:
|
|
logger.warning(
|
|
"tool_progress: deerflow_tool_meta missing for non-exempt tool %s — verify ToolProgressMiddleware is outer of ToolErrorHandlingMiddleware",
|
|
tool_name,
|
|
)
|
|
return result
|
|
content = _message_content_str(message)
|
|
thread_id = self._thread_id(runtime)
|
|
phase_change: tuple[ToolPhaseState, ToolPhaseState, ToolPhaseTransition] | None = None
|
|
with self._lock:
|
|
state = self._get_state(thread_id, tool_name)
|
|
new_state, hint, transition = self._assess_and_transition(state, meta, content)
|
|
self._set_state(thread_id, tool_name, new_state)
|
|
if transition is not None:
|
|
phase_change = (state, new_state, transition)
|
|
if phase_change is not None:
|
|
old_state, changed_state, phase_transition = phase_change
|
|
self._record_phase_transition(
|
|
runtime=runtime,
|
|
tool_name=tool_name,
|
|
state=old_state,
|
|
new_state=changed_state,
|
|
meta=meta,
|
|
hook=hook,
|
|
transition=phase_transition,
|
|
)
|
|
if new_state.phase != state.phase:
|
|
if new_state.phase == "blocked":
|
|
logger.warning(
|
|
"tool_progress: %s/%s -> BLOCKED: %s",
|
|
thread_id,
|
|
tool_name,
|
|
new_state.block_reason,
|
|
)
|
|
elif new_state.phase == "warned":
|
|
logger.info(
|
|
"tool_progress: %s/%s -> WARNED (consecutive_problems=%d)",
|
|
thread_id,
|
|
tool_name,
|
|
new_state.consecutive_problems,
|
|
)
|
|
elif new_state.phase == "active":
|
|
logger.info(
|
|
"tool_progress: %s/%s -> ACTIVE (reset after good result)",
|
|
thread_id,
|
|
tool_name,
|
|
)
|
|
if hint and self._inject_assessment:
|
|
self._queue_assessment(runtime, hint)
|
|
return result
|
|
|
|
# ------------------------------------------------------------------
|
|
# State machine
|
|
|
|
def _assess_and_transition(
|
|
self,
|
|
state: ToolPhaseState,
|
|
meta: ToolResultMeta,
|
|
content: str,
|
|
) -> tuple[ToolPhaseState, str | None, ToolPhaseTransition | None]:
|
|
"""Return the new state, optional hint, and rule that changed phase.
|
|
|
|
The outer wrap_tool_call gate intercepts already-blocked states before
|
|
the handler is called, so this function is normally reached only for
|
|
active/warned states. If a blocked state arrives (e.g., concurrent
|
|
transition), the function returns it unchanged — no counter inflation,
|
|
no phase regression.
|
|
"""
|
|
# Guard: blocked is a terminal state; nothing should change it here.
|
|
# (In normal flow this branch is unreachable because wrap_tool_call
|
|
# intercepts blocked tools before calling the handler. The check exists
|
|
# to make concurrent-race semantics well-defined and prevent a
|
|
# recoverable-error result from silently demoting the phase back to warned.)
|
|
if state.phase == "blocked":
|
|
return state, None, None
|
|
|
|
# Count this call as a problem before branching so all exit paths leave
|
|
# consecutive_problems in a consistent state (never 0 when the tool has failed).
|
|
new_count = state.consecutive_problems + 1
|
|
|
|
# Immediately block on unrecoverable stop signals (auth, config, internal).
|
|
if not meta.recoverable_by_model and meta.recommended_next_action == "stop":
|
|
return (
|
|
replace(
|
|
state,
|
|
phase="blocked",
|
|
consecutive_problems=new_count,
|
|
block_reason=_block_reason(meta),
|
|
),
|
|
None,
|
|
ToolPhaseTransition(action="block", threshold=None),
|
|
)
|
|
|
|
# Compute word_set only for success results: error/partial_success are problems by
|
|
# definition and never reach the Jaccard check, so the O(n) regex is wasted on them.
|
|
ws = word_set(content) if meta.status == "success" else frozenset()
|
|
is_problem = meta.status in ("error", "partial_success") or (meta.status == "success" and is_near_duplicate(ws, state.recent_word_sets, self._jaccard_threshold, self._min_words))
|
|
|
|
if not is_problem:
|
|
# Good result: reset consecutive count, return to active.
|
|
new_recent = (*state.recent_word_sets, ws)[-3:]
|
|
transition = ToolPhaseTransition(action="recover", threshold=None) if state.phase == "warned" else None
|
|
return replace(state, consecutive_problems=0, phase="active", recent_word_sets=new_recent), None, transition
|
|
|
|
hint: str | None = None
|
|
transition: ToolPhaseTransition | None = None
|
|
|
|
if new_count >= self._stagnation_threshold + self._warn_escalation:
|
|
if meta.recoverable_by_model:
|
|
# Model can fix this by changing strategy — keep warned, re-inject hint.
|
|
# BLOCKED would prevent a legitimate retry with different parameters.
|
|
hint = _format_hint(meta)
|
|
new_state = replace(state, consecutive_problems=new_count, phase="warned")
|
|
if state.phase != "warned":
|
|
transition = ToolPhaseTransition(
|
|
action="warn",
|
|
threshold=self._stagnation_threshold + self._warn_escalation,
|
|
)
|
|
else:
|
|
# Model cannot fix this by retrying — block the tool.
|
|
reason = _block_reason(meta)
|
|
new_state = replace(state, consecutive_problems=new_count, phase="blocked", block_reason=reason)
|
|
transition = ToolPhaseTransition(
|
|
action="block",
|
|
threshold=self._stagnation_threshold + self._warn_escalation,
|
|
)
|
|
elif new_count >= self._stagnation_threshold:
|
|
hint = _format_hint(meta)
|
|
new_state = replace(state, consecutive_problems=new_count, phase="warned")
|
|
if state.phase != "warned":
|
|
transition = ToolPhaseTransition(action="warn", threshold=self._stagnation_threshold)
|
|
else:
|
|
new_state = replace(state, consecutive_problems=new_count)
|
|
|
|
return new_state, hint, transition
|
|
|
|
# ------------------------------------------------------------------
|
|
# Pending queue helpers
|
|
|
|
def _queue_assessment(self, runtime: Runtime, text: str) -> None:
|
|
key = self._pending_key(runtime)
|
|
thread_id = key[0]
|
|
with self._lock:
|
|
# Guard against creating a phantom _pending entry for a thread that was just
|
|
# evicted from _phase_states by the LRU. Such entries can never be cleaned up
|
|
# by the eviction loop (which only walks _phase_states) and accumulate silently.
|
|
if thread_id not in self._phase_states:
|
|
return
|
|
queue = self._pending[key]
|
|
if len(queue) < _MAX_PENDING_PER_RUN:
|
|
queue.append(text)
|
|
|
|
def _drain_pending(self, runtime: Runtime) -> list[str]:
|
|
key = self._pending_key(runtime)
|
|
with self._lock:
|
|
return self._pending.pop(key, [])
|
|
|
|
def _clear_stale_pending(self, runtime: Runtime) -> None:
|
|
thread_id, current_run = self._pending_key(runtime)
|
|
with self._lock:
|
|
for key in list(self._pending):
|
|
if key[0] == thread_id and key[1] != current_run:
|
|
del self._pending[key]
|
|
|
|
def _reset_run_states(
|
|
self,
|
|
runtime: Runtime,
|
|
*,
|
|
hook: Literal["before_agent", "abefore_agent"],
|
|
) -> None:
|
|
"""Reset all per-run tool state for the thread at the start of a new agent run.
|
|
|
|
Every tool's consecutive_problems counter and recent_word_sets Jaccard window are
|
|
cleared unconditionally so state from a previous run never bleeds into the next:
|
|
- BLOCKED/WARNED tools are reset to ACTIVE (they re-block immediately if the root
|
|
cause persists, and the model has no memory of the prior-run hint).
|
|
- ACTIVE tools with non-zero consecutive_problems or non-empty recent_word_sets from
|
|
the previous run are also cleared so a single first-call problem in the new run
|
|
cannot falsely trip WARNED against stale context from a run the model no longer sees.
|
|
|
|
**Graph-entry scoping vs LoopDetectionMiddleware**: this reset at every
|
|
``before_agent`` is an intentional policy choice, not an oversight. Errors like
|
|
``rate_limited`` and ``transient`` are time-bound, so carrying a stale counter into a
|
|
later graph entry risks a false-positive BLOCKED on calls that would now succeed.
|
|
LoopDetectionMiddleware instead keys call-pattern state by ``(thread_id, run_id)``:
|
|
separate user runs are isolated even on a cached graph, while repeated graph entries
|
|
in one Gateway run (including hidden goal continuations) share a loop budget. The two
|
|
middlewares therefore guard different failure modes and use different lifetimes.
|
|
"""
|
|
thread_id = self._thread_id(runtime)
|
|
transitions: list[tuple[str, ToolPhaseState, ToolPhaseState]] = []
|
|
with self._lock:
|
|
thread_tools = self._phase_states.get(thread_id)
|
|
if thread_tools is None:
|
|
return
|
|
for tool_name, tool_state in list(thread_tools.items()):
|
|
new_state = replace(
|
|
tool_state,
|
|
phase="active",
|
|
consecutive_problems=0,
|
|
block_reason=None,
|
|
recent_word_sets=(),
|
|
)
|
|
thread_tools[tool_name] = new_state
|
|
if tool_state.phase != new_state.phase:
|
|
transitions.append((tool_name, tool_state, new_state))
|
|
|
|
for tool_name, tool_state, new_state in transitions:
|
|
self._record_phase_transition(
|
|
runtime=runtime,
|
|
tool_name=tool_name,
|
|
state=tool_state,
|
|
new_state=new_state,
|
|
meta=None,
|
|
hook=hook,
|
|
transition=ToolPhaseTransition(action="reset", threshold=None),
|
|
)
|
|
|
|
# ------------------------------------------------------------------
|
|
# wrap_tool_call
|
|
|
|
@override
|
|
def wrap_tool_call(
|
|
self,
|
|
request: ToolCallRequest,
|
|
handler: Callable[[ToolCallRequest], ToolMessage | Command],
|
|
) -> ToolMessage | Command:
|
|
tool_name = str(request.tool_call.get("name", ""))
|
|
if not tool_name or tool_name in self._exempt_tools:
|
|
return handler(request)
|
|
runtime = getattr(request, "runtime", None)
|
|
if runtime is None:
|
|
return handler(request)
|
|
block_reason = self._get_block_reason(runtime, tool_name)
|
|
if block_reason:
|
|
logger.info(
|
|
"tool_progress: %s/%s call intercepted (blocked): %s",
|
|
self._thread_id(runtime),
|
|
tool_name,
|
|
block_reason,
|
|
)
|
|
return self._make_blocked_message(request, tool_name, block_reason)
|
|
return self._update_state_from_result(
|
|
handler(request),
|
|
tool_name,
|
|
runtime,
|
|
str(request.tool_call.get("id") or ""),
|
|
hook="wrap_tool_call",
|
|
)
|
|
|
|
@override
|
|
async def awrap_tool_call(
|
|
self,
|
|
request: ToolCallRequest,
|
|
handler: Callable[[ToolCallRequest], Awaitable[ToolMessage | Command]],
|
|
) -> ToolMessage | Command:
|
|
tool_name = str(request.tool_call.get("name", ""))
|
|
if not tool_name or tool_name in self._exempt_tools:
|
|
return await handler(request)
|
|
runtime = getattr(request, "runtime", None)
|
|
if runtime is None:
|
|
return await handler(request)
|
|
block_reason = self._get_block_reason(runtime, tool_name)
|
|
if block_reason:
|
|
logger.info(
|
|
"tool_progress: %s/%s call intercepted (blocked): %s",
|
|
self._thread_id(runtime),
|
|
tool_name,
|
|
block_reason,
|
|
)
|
|
return self._make_blocked_message(request, tool_name, block_reason)
|
|
return self._update_state_from_result(
|
|
await handler(request),
|
|
tool_name,
|
|
runtime,
|
|
str(request.tool_call.get("id") or ""),
|
|
hook="awrap_tool_call",
|
|
)
|
|
|
|
# ------------------------------------------------------------------
|
|
# wrap_model_call: drain pending hints and inject before model sees messages
|
|
|
|
def _augment_request(self, request: ModelRequest) -> ModelRequest:
|
|
hints = self._drain_pending(request.runtime)
|
|
if not hints:
|
|
return request
|
|
deduped = list(dict.fromkeys(hints))
|
|
logger.debug(
|
|
"tool_progress: injecting %d hint(s) for %s/%s",
|
|
len(deduped),
|
|
*self._pending_key(request.runtime),
|
|
)
|
|
new_messages = [
|
|
*request.messages,
|
|
HumanMessage(content="\n\n".join(deduped), name="progress_hint"),
|
|
]
|
|
return request.override(messages=new_messages)
|
|
|
|
@override
|
|
def wrap_model_call(
|
|
self,
|
|
request: ModelRequest,
|
|
handler: Callable[[ModelRequest], ModelResponse],
|
|
) -> ModelCallResult:
|
|
return handler(self._augment_request(request))
|
|
|
|
@override
|
|
async def awrap_model_call(
|
|
self,
|
|
request: ModelRequest,
|
|
handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
|
|
) -> ModelCallResult:
|
|
return await handler(self._augment_request(request))
|
|
|
|
# ------------------------------------------------------------------
|
|
# before_agent: clean up stale pending hints from previous runs
|
|
|
|
@override
|
|
def before_agent(self, state: AgentState, runtime: Runtime) -> dict | None:
|
|
self._clear_stale_pending(runtime)
|
|
self._reset_run_states(runtime, hook="before_agent")
|
|
return None
|
|
|
|
@override
|
|
async def abefore_agent(self, state: AgentState, runtime: Runtime) -> dict | None:
|
|
self._clear_stale_pending(runtime)
|
|
self._reset_run_states(runtime, hook="abefore_agent")
|
|
return None
|