mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-21 20:16:18 +00:00
0f2813f6b9
* fix(agents): keep queued guard warnings when a model call is retried LoopDetectionMiddleware, TokenBudgetMiddleware and ToolProgressMiddleware pop their queued warning/hint before calling the model. When the call raises, LLMErrorHandlingMiddleware (outside them) retries by running their wrap_model_call again, and by then the queue is empty, so the retried request goes out without the warning. Loop detection and the token budget have already marked it as sent, so it is never queued again, and a loop runs on to the hard stop unwarned. Put the drained items back in front of the queue when the handler raises, then re-raise. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * fix(agents): trim restored loop warnings from the tail and drop a dead helper _restore_pending_warnings put the restored warnings at the front and then trimmed the front, so if the cap ever fired it would drop exactly what it restored. Trim the tail, as tool progress does. _augment_request had no callers after the wrap_model_call change. Add the sync twin of the tool progress retry test. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * refactor(agents): drop tool progress's unused _augment_request Its only remaining reference was a test name; the dedup that test checks lives in _inject_hints. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
775 lines
33 KiB
Python
775 lines
33 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 _restore_pending(self, runtime: Runtime, hints: list[str]) -> None:
|
|
"""Requeue hints taken for a model call that raised.
|
|
|
|
LLMErrorHandlingMiddleware sits outside this middleware and retries a
|
|
failed call by running this wrap again, so the retry must still find
|
|
the hints.
|
|
"""
|
|
if not hints:
|
|
return
|
|
key = self._pending_key(runtime)
|
|
with self._lock:
|
|
if key[0] not in self._phase_states:
|
|
return
|
|
queue = self._pending[key]
|
|
queue[:0] = hints
|
|
del queue[_MAX_PENDING_PER_RUN:]
|
|
|
|
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 _inject_hints(self, request: ModelRequest, hints: list[str]) -> ModelRequest:
|
|
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:
|
|
hints = self._drain_pending(request.runtime)
|
|
try:
|
|
return handler(self._inject_hints(request, hints))
|
|
except Exception:
|
|
self._restore_pending(request.runtime, hints)
|
|
raise
|
|
|
|
@override
|
|
async def awrap_model_call(
|
|
self,
|
|
request: ModelRequest,
|
|
handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
|
|
) -> ModelCallResult:
|
|
hints = self._drain_pending(request.runtime)
|
|
try:
|
|
return await handler(self._inject_hints(request, hints))
|
|
except Exception:
|
|
self._restore_pending(request.runtime, hints)
|
|
raise
|
|
|
|
# ------------------------------------------------------------------
|
|
# 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
|