diff --git a/agent/repetition_guard.py b/agent/repetition_guard.py index c84f5dec37..e2aa371a27 100644 --- a/agent/repetition_guard.py +++ b/agent/repetition_guard.py @@ -1,57 +1,44 @@ """Cheap content-sanity checks for the truncated-response continuation path. -A model in a degenerate repetition loop can spend its ENTIRE output budget -echoing one fragment; the ``finish_reason=length`` continuation path would -then stitch that fragment into the final response with a "continue" nudge -(one incident: a 60k-char turn delivered as 31 Discord messages). These -helpers detect repetition-dominated fragments BEFORE the nudge so the turn -aborts with a clear error (like the ``_thinking_exhausted`` guard). - -Deliberately conservative: only LONG verbatim repeats (60+ chars) covering a -majority of the fragment trip the guard, so ordinary truncations (a sentence -cut mid-word, a repeated heading, similar code lines) are never blocked. +A model in a degenerate repetition loop can spend its ENTIRE output budget echoing one fragment; +the ``finish_reason=length`` continuation would then stitch it into the final response with a +"continue" nudge (one incident: a 60k-char turn delivered as 31 Discord messages). This detects +repetition-dominated fragments BEFORE the nudge so the turn aborts with a clear error. Deliberately +conservative: only LONG verbatim repeats (60+ chars) covering a majority of the fragment trip it. """ from __future__ import annotations import math +from collections import Counter # Below this length the check doesn't run: short truncations trivially # contain repeated tokens and are legitimately continued. MIN_FRAGMENT_LENGTH = 400 - -# Exact-repeat window; a verbatim repeat this long is far beyond ordinary -# phrasing reuse (citations, headings, similar code). +# Exact-repeat window; far beyond ordinary phrasing reuse (citations, headings, similar code). _REPEAT_WINDOW = 60 - # A window repeating at least this often is a signal even for short fragments. _MIN_REPEAT_COUNT = 5 - # "Repetition-dominated" = repeated windows cover at least this fraction. _DOMINANCE_RATIO = 0.5 def is_repetition_dominated(text: str) -> bool: - """True when a single 60+ char substring recurs often enough to cover at - least half of ``text`` — the signature of a repetition loop, where - continuing would only stitch in more repeats. - - Fail-open: False for non-string / empty / short inputs. - """ + """True when a single 60+ char substring recurs often enough to cover at least half + of ``text`` — the signature of a repetition loop. Fail-open for non-string/short input.""" if not isinstance(text, str): return False n = len(text) if n < MIN_FRAGMENT_LENGTH: return False - # Fast path: one normalized line duplicated enough to cover half the - # fragment (the most common echo shape). Cheap, no big allocations. + # Fast path: one normalized line duplicated enough to cover half the fragment (the common echo shape). if _line_repetition_dominated(text, n): return True - # General path: fixed-size windows sliding one char at a time, catching - # loops that don't align to line boundaries. A window must appear - # ``needed`` times to cover >= _DOMINANCE_RATIO (and >= _MIN_REPEAT_COUNT). + # General path: fixed-size windows sliding one char at a time, catching loops that + # don't align to line boundaries. A window must appear ``needed`` times to cover + # >= _DOMINANCE_RATIO (and >= _MIN_REPEAT_COUNT). window = _REPEAT_WINDOW needed = max(_MIN_REPEAT_COUNT, math.ceil(n * _DOMINANCE_RATIO / window)) counts: dict[str, int] = {} @@ -66,12 +53,5 @@ def is_repetition_dominated(text: str) -> bool: def _line_repetition_dominated(text: str, n: int) -> bool: """True when a single normalized line covers half the fragment via repeats.""" - counts: dict[str, int] = {} - for line in text.splitlines(): - norm = line.strip() - if norm: - counts[norm] = counts.get(norm, 0) + 1 - return any( - c >= _MIN_REPEAT_COUNT and c * len(line) >= n * _DOMINANCE_RATIO - for line, c in counts.items() - ) + counts = Counter(norm for norm in (line.strip() for line in text.splitlines()) if norm) + return any(c >= _MIN_REPEAT_COUNT and c * len(line) >= n * _DOMINANCE_RATIO for line, c in counts.items()) diff --git a/agent/status_output.py b/agent/status_output.py index eb1cf6d32d..ba04821e94 100644 --- a/agent/status_output.py +++ b/agent/status_output.py @@ -1,7 +1,7 @@ """User-facing status / warning / notice plumbing for ``AIAgent``. -Safe printing, quiet-mode gating, deduped context-overflow warnings, and the buffered retry chatter that -is shown only when every retry/fallback is exhausted. +Safe printing, quiet-mode gating, deduped context-overflow warnings, and the buffered retry +chatter that is shown only when every retry/fallback is exhausted. Extracted from ``run_agent.py``; every method resolves through ``AIAgent``'s MRO unchanged. """ import logging @@ -17,125 +17,82 @@ class StatusOutputMixin: """Status/warning/notice emission and retry-chatter buffering (see module docstring).""" def _safe_print(self, *args, **kwargs): - """Print that silently handles broken pipes / closed stdout. - - Headless stdout (systemd, Docker, nohup) can vanish mid-session and a raw ``print()`` would crash cron - jobs. - Routes through ``self._print_fn`` so the CLI can inject an ANSI-aware renderer. - """ + """Print that swallows broken pipes / closed stdout (headless stdout can vanish mid-session); + routes through ``self._print_fn`` so the CLI can inject an ANSI-aware renderer.""" try: - fn = self._print_fn or print - fn(*args, **kwargs) + (self._print_fn or print)(*args, **kwargs) except (OSError, ValueError): pass def _vprint(self, *args, force: bool = False, **kwargs): - """Verbose print — suppressed when actively streaming tokens. - - ``force=True`` always shows (errors/warnings). Printing is allowed during tool execution (no tokens - streaming), muted after the main response (``_mute_post_response``), and fully suppressed under - ``suppress_status_output`` (machine-readable ``hermes chat -q``). - """ + """Verbose print — suppressed while tokens are streaming (allowed during tool execution) and after + the main response; ``force=True`` bypasses both. ``suppress_status_output`` (``hermes chat -q``) wins.""" if getattr(self, "suppress_status_output", False): return - if not force and getattr(self, "_mute_post_response", False): - return - if not force and self._has_stream_consumers() and not self._executing_tools: - return - self._safe_print(*args, **kwargs) + if force or not (getattr(self, "_mute_post_response", False) or (self._has_stream_consumers() and not self._executing_tools)): + self._safe_print(*args, **kwargs) def _should_start_quiet_spinner(self) -> bool: - """Return True when quiet-mode spinner output has a safe sink. - - A raw spinner falling back to ``sys.stdout`` can corrupt protocol streams (ACP JSON-RPC); allow it - only when output is rerouted via ``_print_fn`` or stdout is a real TTY. - """ + """True when quiet-mode spinner output has a safe sink (``_print_fn`` or a real TTY); a raw spinner + on a non-TTY stdout can corrupt protocol streams (ACP JSON-RPC).""" if self._print_fn is not None: return True - stream = getattr(sys, "stdout", None) - if stream is None: - return False try: - return bool(stream.isatty()) + return bool(sys.stdout.isatty()) except (AttributeError, ValueError, OSError): return False def _should_emit_quiet_tool_messages(self) -> bool: - """Return True when quiet-mode tool summaries should print directly. - - The CLI wants compact hints when no callback owns rendering; embedded callers expect true silence. - ``suppress_status_output`` always wins so ``[tool]``/``[done]`` lines never land in captured stdout. - """ + """True when quiet-mode tool summaries should print directly (CLI, no callback owns rendering); + ``suppress_status_output`` always wins so ``[tool]``/``[done]`` never land in captured stdout.""" if getattr(self, "suppress_status_output", False): return False - return ( - self.quiet_mode - and not self.tool_progress_callback - and getattr(self, "platform", "") == "cli" - ) + return self.quiet_mode and not self.tool_progress_callback and getattr(self, "platform", "") == "cli" + + def _call_callback(self, name: str, *args, origin: str) -> None: + """Invoke ``self.(*args)`` if set, swallowing errors — a driver callback must never break the loop.""" + cb = getattr(self, name, None) + if cb: + try: + cb(*args) + except Exception: + logger.debug("%s error in %s", name, origin, exc_info=True) def _emit_status_kind(self, kind: str, message: str, *, origin: str) -> None: - """Print to the CLI (``_vprint(force=True)``) and forward to ``status_callback(kind, message)``. - - Never raises — status output must not interrupt the retry/fallback logic. - """ + """Print to the CLI (``_vprint(force=True)``) and forward to ``status_callback(kind, message)``. Never raises.""" try: self._vprint(f"{self.log_prefix}{message}", force=True) except Exception: pass - if self.status_callback: - try: - self.status_callback(kind, message) - except Exception: - logger.debug("status_callback error in %s", origin, exc_info=True) + self._call_callback("status_callback", kind, message, origin=origin) def _emit_status(self, message: str) -> None: """Emit a lifecycle status message (CLI + gateway ``status_callback``).""" self._emit_status_kind("lifecycle", message, origin="_emit_status") def _emit_warning(self, message: str) -> None: - """Emit a user-visible warning for degraded side paths (aux compression, memory flushes) where the - turn continues but the user must know.""" + """Emit a user-visible warning for degraded side paths where the turn continues but the user must know.""" self._emit_status_kind("warn", message, origin="_emit_warning") - def _warn_context_overflow_blocked( - self, reason: str, preflight_tokens: int, threshold_tokens: int - ) -> None: - """Warn (deduped) when context is over the compression threshold but compression is blocked. - - Without this the session grows until the provider hard limit with no explanation. Dedup is on the - block - *kind* (``cooldown`` / ``ineffective``), not the countdown string; cleared by - ``_clear_context_overflow_warn``. - """ + def _warn_context_overflow_blocked(self, reason: str, preflight_tokens: int, threshold_tokens: int) -> None: + """Warn (deduped on the block *kind* — ``cooldown`` / ``ineffective`` — not the countdown string; + cleared by ``_clear_context_overflow_warn``) when context is over the threshold but compression is blocked.""" _warn_kind = (reason or "unknown").split(":", 1)[0] _warn_key = ("ctx_overflow_blocked", _warn_kind) - if getattr(self, "_last_ctx_overflow_warn", None) != _warn_key: - self._last_ctx_overflow_warn = _warn_key - from agent.conversation_compression import ( - CONTEXT_OVERFLOW_BLOCKED_WARNING_TEMPLATE, - ) - # cooldown + anti-thrash (ineffective) are both "compression blocked". - if _warn_kind in ("cooldown", "ineffective"): - self._touch_activity( - f"compression blocked ({reason})", - provenance=ActivityProvenance.AGENT_COMPRESSION_COOLDOWN, - ) - self._emit_warning( - CONTEXT_OVERFLOW_BLOCKED_WARNING_TEMPLATE.format( - tokens=preflight_tokens, - threshold=threshold_tokens, - reason=reason, - ) - ) + if getattr(self, "_last_ctx_overflow_warn", None) == _warn_key: + return + self._last_ctx_overflow_warn = _warn_key + from agent.conversation_compression import CONTEXT_OVERFLOW_BLOCKED_WARNING_TEMPLATE - def _warn_uncompressed_context_overflow( - self, preflight_tokens: int, context_length: int - ) -> None: - """Surface a deduped warning when uncompressed context exceeds the model limit. + # cooldown + anti-thrash (ineffective) are both "compression blocked". + if _warn_kind in ("cooldown", "ineffective"): + self._touch_activity(f"compression blocked ({reason})", provenance=ActivityProvenance.AGENT_COMPRESSION_COOLDOWN) + self._emit_warning(CONTEXT_OVERFLOW_BLOCKED_WARNING_TEMPLATE.format( + tokens=preflight_tokens, threshold=threshold_tokens, reason=reason, + )) - With compression disabled the session can outgrow the window silently; point the user at /compact. - """ + def _warn_uncompressed_context_overflow(self, preflight_tokens: int, context_length: int) -> None: + """Deduped warning when uncompressed context exceeds the model limit; points the user at /compact.""" _warn_key = ("uncompressed_ctx_overflow", context_length) if getattr(self, "_last_ctx_overflow_warn", None) != _warn_key: self._last_ctx_overflow_warn = _warn_key @@ -151,57 +108,32 @@ class StatusOutputMixin: self._last_ctx_overflow_warn = None def _emit_notice(self, notice) -> None: - """Fire a structured ``AgentNotice`` to the active driver (TUI / CLI). - - Driver-agnostic; swallows callback errors — a notice must never break the agent loop. - """ - if self.notice_callback: - try: - self.notice_callback(notice) - except Exception: - logger.debug("notice_callback error in _emit_notice", exc_info=True) + """Fire a structured ``AgentNotice`` to the active driver (TUI / CLI).""" + self._call_callback("notice_callback", notice, origin="_emit_notice") def _emit_notice_clear(self, key: str) -> None: """Clear a previously-fired sticky notice by ``key`` (e.g. on recovery).""" - if self.notice_clear_callback: - try: - self.notice_clear_callback(key) - except Exception: - logger.debug("notice_clear_callback error in _emit_notice_clear", exc_info=True) + self._call_callback("notice_clear_callback", key, origin="_emit_notice_clear") def _emit_wait_notice(self, text: str) -> None: - """Surface a live wait-state explanation on every driver. - - Rewrites the live status line (CLI spinner via ``thinking_callback``, TUI/Desktop ``thinking.delta``, - gateway ``_touch_activity`` description) so long provider waits are not an anonymous spinner. Never - raises. - """ + """Rewrite the live status line (CLI spinner, TUI ``thinking.delta``, gateway activity) + so long provider waits are not an anonymous spinner.""" self._touch_activity(text) - _thinking_cb = getattr(self, "thinking_callback", None) - if _thinking_cb: - try: - _thinking_cb(text) - except Exception: - logger.debug("thinking_callback error in _emit_wait_notice", exc_info=True) + self._call_callback("thinking_callback", text, origin="_emit_wait_notice") - # ── Buffered retry/fallback status ── - # Retry chatter is buffered and shown only when every retry/fallback is exhausted; dropped on - # success. Backend logs are unaffected (every site still logs). + # ── Buffered retry/fallback status: shown only when every retry/fallback is exhausted, dropped on + # success. Backend logs are unaffected (every site still logs). ── def _buffer_retry_message(self, kind: str, message: str) -> None: """Buffer a retry/fallback line as ``(kind, text)`` until we know whether the turn recovered. ``kind`` is ``"status"`` (replays via ``_emit_status``), ``"vprint"`` (``_vprint(force=True)``) or - ``"warn"`` (``_emit_warning``). Never breaks the retry loop on a buffer hiccup. + ``"warn"`` (``_emit_warning``). """ - try: - buf = getattr(self, "_retry_status_buffer", None) - if buf is None: - buf = [] - self._retry_status_buffer = buf - buf.append((kind, message)) - except Exception: - pass + buf = getattr(self, "_retry_status_buffer", None) + if buf is None: + buf = self._retry_status_buffer = [] + buf.append((kind, message)) def _buffer_status(self, message: str) -> None: self._buffer_retry_message("status", message) @@ -211,59 +143,42 @@ class StatusOutputMixin: def _clear_status_buffer(self) -> None: """Drop buffered retry messages — call on successful recovery.""" - try: - buf = getattr(self, "_retry_status_buffer", None) - if buf: - buf.clear() - except Exception: - pass + buf = getattr(self, "_retry_status_buffer", None) + if buf: + buf.clear() def _emit_pending_fallback_notice(self) -> None: - """Surface the one-shot fallback-switch notice on successful recovery. - - A provider switch is durable state operators must see, unlike the transient retry chatter - ``_clear_status_buffer`` drops. Emitted exactly once, then cleared; on terminal failure the buffered - switch line is flushed instead (``_flush_status_buffer``). - """ - try: - notice = getattr(self, "_pending_fallback_notice", None) - if notice: - # Clear before emitting so a (swallowed) callback error can't - # leave the notice set for a stale re-emit on a later turn. - self._pending_fallback_notice = None - notices = notice if isinstance(notice, list) else [notice] - for item in notices: - try: - self._emit_status(str(item)) - except Exception: - # A single surface callback failure must not hide later - # switches from the same fallback chain. - continue - except Exception: - # Never break the conversation loop on a notice hiccup. - pass + """Surface the one-shot fallback-switch notice on successful recovery: a provider switch is durable + state operators must see, unlike the retry chatter ``_clear_status_buffer`` drops. Emitted once, then + cleared; on terminal failure the buffered switch line is flushed instead (``_flush_status_buffer``).""" + notice = getattr(self, "_pending_fallback_notice", None) + if not notice: + return + # Clear before emitting so a (swallowed) callback error can't leave a stale re-emit. + self._pending_fallback_notice = None + for item in notice if isinstance(notice, list) else [notice]: + try: + self._emit_status(str(item)) + except Exception: + # One surface failure must not hide later switches from the same chain. + continue def _flush_status_buffer(self) -> None: """Emit buffered retry messages — call on terminal failure so the user sees what was tried.""" - try: - # The buffered trace already carries the switch line; drop the one-shot notice to avoid a stale - # duplicate. - self._pending_fallback_notice = None - buf = getattr(self, "_retry_status_buffer", None) - if not buf: - return - # Drain first so a callback exception doesn't double-emit. - messages = list(buf) - buf.clear() - for kind, msg in messages: - try: - if kind == "status": - self._emit_status(msg) - elif kind == "warn": - self._emit_warning(msg) - else: - self._vprint(f"{self.log_prefix}{msg}", force=True) - except Exception: - pass - except Exception: - pass + # The buffered trace already carries the switch line; drop the one-shot notice. + self._pending_fallback_notice = None + buf = getattr(self, "_retry_status_buffer", None) + if not buf: + return + # Drain first so a callback exception doesn't double-emit. + messages = list(buf) + buf.clear() + replay = {"status": self._emit_status, "warn": self._emit_warning} + for kind, msg in messages: + try: + if kind in replay: + replay[kind](msg) + else: + self._vprint(f"{self.log_prefix}{msg}", force=True) + except Exception: + pass diff --git a/agent/stream_delivery.py b/agent/stream_delivery.py index f1388afda9..4c2e02ec81 100644 --- a/agent/stream_delivery.py +++ b/agent/stream_delivery.py @@ -19,278 +19,196 @@ logger = logging.getLogger("run_agent") class StreamDeliveryMixin: """Stream ownership, delta/reasoning hook fan-out and interim-text dedup (see module docstring).""" + @staticmethod + def _call_quietly(cb, *args) -> bool: + """Call ``cb(*args)`` if set, swallowing errors; True when it ran without raising.""" + if cb is None: + return False + try: + cb(*args) + return True + except Exception: + return False + + def _deliver_to_stream_callbacks(self, text: str) -> bool: + """Send ``text`` to the display + TTS delta callbacks; True if at least one accepted it.""" + results = [self._call_quietly(cb, text) for cb in (self.stream_delta_callback, self._stream_callback)] + return any(results) + + def _enqueue_stream_hook(self, event: str, *, label: str | None = None, **fields: Any) -> None: + """Best-effort plugin stream hook enqueue; never raises into the stream path.""" + try: + from agent.plugin_stream_hooks import enqueue_plugin_stream_hook + + enqueue_plugin_stream_hook(event, **self._stream_hook_base_payload(), **fields) + except Exception: + logger.debug("%s plugin hook enqueue failed", label or event, exc_info=True) + def _reset_stream_delivery_tracking(self) -> None: - """Reset tracking for text delivered during the current model response.""" - # Flush the think scrubber's benign partial-tag tail first (#17924), then the context scrubber — - # order matters, the think output feeds the context scrubber. + """Reset tracking for text delivered during the current model response. + + Flushes the think scrubber's benign tail first, routed through the context scrubber (a span + straddling the boundary must still be caught), then the context scrubber's own tail. + """ think_scrubber = getattr(self, "_stream_think_scrubber", None) + ctx_scrubber = getattr(self, "_stream_context_scrubber", None) + + def deliver(tail: str) -> None: + if tail: + self._deliver_to_stream_callbacks(tail) + self._record_streamed_assistant_text(tail) + if think_scrubber is not None: think_tail = think_scrubber.flush() - if think_tail: - # Route the tail through the context scrubber so a span straddling the boundary is caught. - ctx_scrubber = getattr(self, "_stream_context_scrubber", None) - if ctx_scrubber is not None: - think_tail = ctx_scrubber.feed(think_tail) - if think_tail: - callbacks = [cb for cb in (self.stream_delta_callback, self._stream_callback) if cb is not None] - for cb in callbacks: - try: - cb(think_tail) - except Exception: - pass - self._record_streamed_assistant_text(think_tail) - # Flush the context scrubber's benign tail before clearing; mid-span, flush() drops orphaned content. - scrubber = getattr(self, "_stream_context_scrubber", None) - if scrubber is not None: - tail = scrubber.flush() - if tail: - callbacks = [cb for cb in (self.stream_delta_callback, self._stream_callback) if cb is not None] - for cb in callbacks: - try: - cb(tail) - except Exception: - pass - self._record_streamed_assistant_text(tail) + deliver(ctx_scrubber.feed(think_tail) if think_tail and ctx_scrubber is not None else think_tail) + if ctx_scrubber is not None: + deliver(ctx_scrubber.flush()) self._current_streamed_assistant_text = "" def _record_streamed_assistant_text(self, text: str) -> None: - """Accumulate visible assistant text emitted through stream callbacks.""" - # Single-writer guard (#65991): a superseded stream must not pollute the accumulated text, even - # when reached directly via the tool-suppressed path. - if self._stream_writer_superseded(): - return - if isinstance(text, str) and text: - self._current_streamed_assistant_text = ( - getattr(self, "_current_streamed_assistant_text", "") + text - ) + """Accumulate visible assistant text emitted through stream callbacks (superseded writers excluded).""" + if isinstance(text, str) and text and not self._stream_writer_superseded(): + self._current_streamed_assistant_text = getattr(self, "_current_streamed_assistant_text", "") + text @staticmethod def _normalize_interim_visible_text(text: str) -> str: - if not isinstance(text, str): - return "" - return re.sub(r"\s+", " ", text).strip() + return re.sub(r"\s+", " ", text).strip() if isinstance(text, str) else "" def _interim_content_was_streamed(self, content: str) -> bool: - visible_content = self._normalize_interim_visible_text( - self._strip_think_blocks(content or "") - ) - if not visible_content: - return False + visible_content = self._normalize_interim_visible_text(self._strip_think_blocks(content or "")) streamed = self._normalize_interim_visible_text( self._strip_think_blocks(getattr(self, "_current_streamed_assistant_text", "") or "") ) - # Prefix match, not equality: the final may be streamed text plus a trailing delta, or the stream - # partial. The reverse (streamed longer) is NOT matched — it could suppress a needed resend (#65919). - return bool(streamed) and visible_content.startswith(streamed) + # Prefix match, not equality: the final may be streamed text plus a trailing delta. The + # reverse (streamed longer) is NOT matched — it could suppress a needed resend. + return bool(visible_content and streamed) and visible_content.startswith(streamed) - def _extract_codex_interim_visible_parts( - self, - assistant_msg: Dict[str, Any], - ) -> List[str]: - """Extract visible Codex commentary as one string per message item. + def _extract_codex_interim_visible_parts(self, assistant_msg: Dict[str, Any]) -> List[str]: + """Visible Codex commentary (``phase=commentary`` items), one string per message item. - Codex keeps mid-turn narration as ``phase=commentary`` items while the final answer stays in - ``content``; - non-streaming gateways need it via the interim callback. ``phase=analysis`` stays hidden (scratchpad). + ``phase=analysis`` stays hidden (scratchpad); with ``display.show_commentary=false`` + commentary stays on the reasoning channel. """ - if not getattr(self, "show_commentary", True): - # display.show_commentary=false — commentary stays on the - # reasoning channel (pre-commentary-channel behavior). - return [] - items = assistant_msg.get("codex_message_items") - if not isinstance(items, list): - return [] - + items = assistant_msg.get("codex_message_items") if getattr(self, "show_commentary", True) else None messages: List[str] = [] - for item in items: - if not isinstance(item, dict): + for item in items if isinstance(items, list) else (): + if not isinstance(item, dict) or item.get("type") != "message": continue - if item.get("type") != "message": + phase, content_parts = item.get("phase"), item.get("content") + if not isinstance(phase, str) or phase.strip().lower() != "commentary" or not isinstance(content_parts, list): continue - phase = item.get("phase") - if not isinstance(phase, str) or phase.strip().lower() != "commentary": - continue - content_parts = item.get("content") - if not isinstance(content_parts, list): - continue - item_parts: List[str] = [] - for part in content_parts: - if not isinstance(part, dict): - continue - if part.get("type") != "output_text": - continue - text = part.get("text") - if isinstance(text, str) and text.strip(): - item_parts.append(text) - visible = "".join(item_parts).strip() - if visible: - visible = self._strip_think_blocks(visible).strip() - visible = redact_sensitive_text(visible) + visible = "".join( + part["text"] for part in content_parts + if isinstance(part, dict) and part.get("type") == "output_text" + and isinstance(part.get("text"), str) and part["text"].strip() + ).strip() + visible = self._visible_commentary(visible) if visible: messages.append(visible) return messages + def _visible_commentary(self, text: str) -> str: + """Think-stripped, redacted commentary text ("" when nothing visible remains).""" + visible = self._strip_think_blocks(text).strip() + return redact_sensitive_text(visible) if visible else visible + def _extract_codex_interim_visible_text(self, assistant_msg: Dict[str, Any]) -> str: - """Extract all visible Codex commentary for comparison/fallback.""" - return "\n\n".join( - self._extract_codex_interim_visible_parts(assistant_msg) - ).strip() + """All visible Codex commentary joined, for comparison/fallback.""" + return "\n\n".join(self._extract_codex_interim_visible_parts(assistant_msg)).strip() def _interim_assistant_visible_text(self, assistant_msg: Dict[str, Any]) -> str: - """Return the exact assistant text eligible for interim delivery. - - Prefers structured Codex commentary over top-level content — a response can hold commentary AND a - partial final answer while tools are pending, and treating content as progress leaks the answer early. - Content may be a parts list, so flatten before stripping reasoning. - """ - visible = self._extract_codex_interim_visible_text(assistant_msg) - if visible: - return visible - content = assistant_msg.get("content") - return self._strip_think_blocks(flatten_message_text(content)).strip() + """Assistant text eligible for interim delivery: structured Codex commentary first — a response can + hold commentary AND a partial final answer while tools are pending, and treating content as + progress leaks the answer early — else top-level content (may be a parts list).""" + return ( + self._extract_codex_interim_visible_text(assistant_msg) + or self._strip_think_blocks(flatten_message_text(assistant_msg.get("content"))).strip() + ) def _interim_text_was_delivered(self, text: str) -> bool: normalized = self._normalize_interim_visible_text(text) - if not normalized: - return False - return normalized in getattr(self, "_delivered_interim_texts", set()) + return bool(normalized) and normalized in getattr(self, "_delivered_interim_texts", set()) def _record_delivered_interim_text(self, text: str) -> None: normalized = self._normalize_interim_visible_text(text) if normalized: - delivered = getattr(self, "_delivered_interim_texts", None) - if not isinstance(delivered, set): - delivered = set() - self._delivered_interim_texts = delivered - delivered.add(normalized) + if not isinstance(getattr(self, "_delivered_interim_texts", None), set): + self._delivered_interim_texts = set() + self._delivered_interim_texts.add(normalized) - def _fire_streamed_codex_commentary(self, text: str) -> None: - """Deliver a completed live Codex commentary message immediately.""" - cb = getattr(self, "interim_assistant_callback", None) - if cb is None or not isinstance(text, str): - return - visible = self._strip_think_blocks(text).strip() - if visible: - visible = redact_sensitive_text(visible) - if not visible or visible == "(empty)" or self._interim_text_was_delivered(visible): - return - try: - cb(visible, already_streamed=False) - self._record_delivered_interim_text(visible) - except Exception: - logger.debug("interim_assistant_callback error", exc_info=True) - - def _emit_interim_assistant_message( - self, assistant_msg: Dict[str, Any] - ) -> None: - """Surface a real mid-turn assistant commentary message to the UI layer. - - Does NOT set ``_response_was_previewed`` — that means "the final response was shown"; setting it for - narration would make the CLI suppress a different final summary (response-loss blocker). - """ - if not isinstance(assistant_msg, dict): - return - commentary_parts = self._extract_codex_interim_visible_parts(assistant_msg) - undelivered_parts: List[str] = [] - pending_keys: set[str] = set() - for part in commentary_parts: - key = self._normalize_interim_visible_text(part) - if ( - not key - or key in pending_keys - or self._interim_text_was_delivered(part) - ): - continue - pending_keys.add(key) - undelivered_parts.append(part) - visible = ( - "\n\n".join(undelivered_parts).strip() - if commentary_parts - else self._interim_assistant_visible_text(assistant_msg) - ) - if ( - not visible - or visible == "(empty)" - or self._interim_text_was_delivered(visible) - ): - return - already_streamed = self._interim_content_was_streamed(visible) - try: - from agent.plugin_stream_hooks import enqueue_plugin_stream_hook - - enqueue_plugin_stream_hook( - "on_interim_message", - turn_id=getattr(self, "_current_turn_id", "") or "", - iteration=int(getattr(self, "_api_call_count", 0) or 0), - session_id=self.session_id or "", - model=self.model or "", - provider=self.provider or "", - surface=self.platform or "cli", - text=visible, - already_streamed=already_streamed, - ) - except Exception: - logger.debug("on_interim_message plugin hook enqueue failed", exc_info=True) + def _deliver_interim(self, visible: str, *, already_streamed: bool, record: List[str]) -> None: + """Hand ``visible`` to ``interim_assistant_callback`` and mark ``record`` delivered; swallows callback errors.""" cb = getattr(self, "interim_assistant_callback", None) if cb is None: return try: cb(visible, already_streamed=already_streamed) - if undelivered_parts: - for part in undelivered_parts: - self._record_delivered_interim_text(part) - else: - self._record_delivered_interim_text(visible) + for part in record: + self._record_delivered_interim_text(part) except Exception: logger.debug("interim_assistant_callback error", exc_info=True) - def _ensure_stream_writer_state(self) -> None: - """Lazily create the single-writer guard fields. + def _fire_streamed_codex_commentary(self, text: str) -> None: + """Deliver a completed live Codex commentary message immediately.""" + if getattr(self, "interim_assistant_callback", None) is None or not isinstance(text, str): + return + visible = self._visible_commentary(text) + if not visible or visible == "(empty)" or self._interim_text_was_delivered(visible): + return + self._deliver_interim(visible, already_streamed=False, record=[visible]) - Normally set in ``agent_init``; ``AIAgent.__new__``-built instances skip that path and must not crash. - """ + def _emit_interim_assistant_message(self, assistant_msg: Dict[str, Any]) -> None: + """Surface a real mid-turn assistant commentary message to the UI layer. Does NOT set + ``_response_was_previewed`` ("the final response was shown") — the CLI would then suppress a + different final summary.""" + if not isinstance(assistant_msg, dict): + return + commentary_parts = self._extract_codex_interim_visible_parts(assistant_msg) + # Dedup within this message and against earlier deliveries, first occurrence wins. + pending: dict[str, str] = {} + for part in commentary_parts: + key = self._normalize_interim_visible_text(part) + if key and key not in pending and not self._interim_text_was_delivered(part): + pending[key] = part + undelivered_parts = list(pending.values()) + visible = "\n\n".join(undelivered_parts).strip() if commentary_parts else self._interim_assistant_visible_text(assistant_msg) + if not visible or visible == "(empty)" or self._interim_text_was_delivered(visible): + return + already_streamed = self._interim_content_was_streamed(visible) + self._enqueue_stream_hook("on_interim_message", text=visible, already_streamed=already_streamed) + self._deliver_interim(visible, already_streamed=already_streamed, record=undelivered_parts or [visible]) + + def _ensure_stream_writer_state(self) -> None: + """Lazily create the single-writer guard fields (``AIAgent.__new__``-built instances skip ``agent_init``).""" if getattr(self, "_stream_writer_lock", None) is None: self._stream_writer_lock = threading.Lock() - if not hasattr(self, "_stream_writer_token"): - self._stream_writer_token = 0 if getattr(self, "_stream_writer_tls", None) is None: self._stream_writer_tls = threading.local() - if not hasattr(self, "_stream_writer_dropped"): - self._stream_writer_dropped = 0 + for attr in ("_stream_writer_token", "_stream_writer_dropped"): + if not hasattr(self, attr): + setattr(self, attr, 0) def _claim_stream_writer(self) -> int: - """Claim exclusive ownership of the streaming delta sink for this stream attempt; returns its writer - token. + """Claim exclusive ownership of the delta sink for this stream attempt; returns its writer token. - Every attempt (each provider path, each retry) claims right before consuming. Claiming bumps the - shared token, so an earlier attempt still alive on another thread is superseded and its late chunks - fenced out. Stored per-thread: a thread that never claimed is never a writer and can never be fenced. + Every attempt (each provider path, each retry) claims right before consuming; claiming bumps + the shared token, so an earlier attempt still alive on another thread is superseded and its + late chunks fenced out. Stored per-thread: a thread that never claimed can never be fenced. """ self._ensure_stream_writer_state() with self._stream_writer_lock: - self._stream_writer_token += 1 - token = self._stream_writer_token + self._stream_writer_token = token = self._stream_writer_token + 1 self._stream_writer_tls.token = token return token def _stream_writer_is_current(self, token: int) -> bool: - """True when ``token`` is still the active writer (no newer attempt claimed since). - - Lets a stream loop bail out the instant it is superseded. - """ + """True when ``token`` is still the active writer, so a stream loop can bail the instant it is superseded.""" return token == getattr(self, "_stream_writer_token", token) def _stream_writer_superseded(self) -> bool: - """True when this thread claimed the sink but a newer attempt has since claimed it (stale writer, drop - chunks). - - A thread that never claimed (``token is None``) is never reported superseded. - """ - tls = getattr(self, "_stream_writer_tls", None) - token = getattr(tls, "token", None) if tls is not None else None - if token is None: - return False - return token != getattr(self, "_stream_writer_token", token) + """True when this thread claimed the sink but a newer attempt has since claimed it (never for a non-claimer).""" + token = getattr(getattr(self, "_stream_writer_tls", None), "token", None) + return token is not None and token != getattr(self, "_stream_writer_token", token) def _note_dropped_stream_writer(self, where: str) -> None: """Record + log that a superseded stream's delta was discarded.""" @@ -319,126 +237,60 @@ class StreamDeliveryMixin: } def _emit_stream_start(self) -> None: - try: - from agent.plugin_stream_hooks import enqueue_plugin_stream_hook - - enqueue_plugin_stream_hook("on_stream_start", **self._stream_hook_base_payload()) - except Exception: - logger.debug("on_stream_start plugin hook enqueue failed", exc_info=True) + self._enqueue_stream_hook("on_stream_start") def _emit_stream_end(self, *, final_text: str, finished: bool, error: str | None) -> None: - try: - from agent.plugin_stream_hooks import enqueue_plugin_stream_hook - - enqueue_plugin_stream_hook( - "on_stream_end", - **self._stream_hook_base_payload(), - final_text=final_text, - finished=finished, - error=error, - ) - except Exception: - logger.debug("on_stream_end plugin hook enqueue failed", exc_info=True) + self._enqueue_stream_hook("on_stream_end", final_text=final_text, finished=finished, error=error) def _fire_stream_delta(self, text: str) -> None: """Fire all registered stream delta callbacks (display + TTS).""" - # Single-writer guard (#65991): a superseded stream must not interleave - # its tokens into the turn alongside the retry that replaced it. + # A superseded stream must not interleave its tokens alongside the retry that replaced it. if self._stream_writer_superseded(): self._note_dropped_stream_writer("_fire_stream_delta") return - # Prepend one paragraph break before the first text delta after a tool iteration, without + # One paragraph break before the first text delta after a tool iteration, without # stacking blank lines across back-to-back tool iterations. - if getattr(self, "_stream_needs_break", False) and text and text.strip(): + prepended_break = bool(getattr(self, "_stream_needs_break", False) and text and text.strip()) + if prepended_break: self._stream_needs_break = False text = "\n\n" + text - prepended_break = True - else: - prepended_break = False if isinstance(text, str): - # Stateful scrubber (#17924): per-delta regex stripping destroyed downstream state machines when a - # tag was split across deltas (MiniMax-M2.7 sends '' separately). + # Stateful scrubbers: per-delta regex stripping destroyed downstream state machines when a + # tag was split across deltas; memory-context spans split across chunks must not leak to + # the UI. Legacy callers lack the scrubber attributes and get the whole-string fallbacks. think_scrubber = getattr(self, "_stream_think_scrubber", None) - if think_scrubber is not None: - text = think_scrubber.feed(text or "") - else: - # Defensive: legacy callers without the scrubber attribute. - text = self._strip_think_blocks(text or "") - # Then feed through the stateful context scrubber so memory-context - # spans split across chunks cannot leak to the UI (#5719). scrubber = getattr(self, "_stream_context_scrubber", None) - if scrubber is not None: - text = scrubber.feed(text) - else: - # Defensive: legacy callers without the scrubber attribute. - text = sanitize_context(text) + text = think_scrubber.feed(text) if think_scrubber is not None else self._strip_think_blocks(text) + text = scrubber.feed(text) if scrubber is not None else sanitize_context(text) # Only strip leading newlines on the first delta — mid-stream "\n" is legitimate markdown. - if not prepended_break and not getattr( - self, "_current_streamed_assistant_text", "" - ): + if not prepended_break and not getattr(self, "_current_streamed_assistant_text", ""): text = text.lstrip("\n") if not text: return - callbacks = [cb for cb in (self.stream_delta_callback, self._stream_callback) if cb is not None] - delivered = False - for cb in callbacks: - try: - cb(text) - delivered = True - except Exception: - pass - try: - from agent.plugin_stream_hooks import enqueue_plugin_stream_hook - - enqueue_plugin_stream_hook( - "on_stream_delta", - **self._stream_hook_base_payload(), - delta=text, - kind="text", - ) - except Exception: - logger.debug("on_stream_delta plugin hook enqueue failed", exc_info=True) + delivered = self._deliver_to_stream_callbacks(text) + self._enqueue_stream_hook("on_stream_delta", delta=text, kind="text") if delivered: self._record_streamed_assistant_text(text) def _fire_reasoning_delta(self, text: str) -> None: - """Fire reasoning callback if registered.""" - # Single-writer guard (#65991): fence out a superseded stream's - # reasoning deltas the same way as content deltas. + """Fire reasoning callback if registered; superseded writers are fenced like content deltas.""" if self._stream_writer_superseded(): self._note_dropped_stream_writer("_fire_reasoning_delta") return - cb = self.reasoning_callback - if cb is not None: - try: - cb(text) - except Exception: - pass + self._call_quietly(self.reasoning_callback, text) try: - from agent.plugin_stream_hooks import enqueue_plugin_stream_hook, stream_reasoning_deltas_enabled + from agent.plugin_stream_hooks import stream_reasoning_deltas_enabled - if stream_reasoning_deltas_enabled(): - enqueue_plugin_stream_hook( - "on_stream_delta", - **self._stream_hook_base_payload(), - delta=text, - kind="reasoning", - ) + enabled = stream_reasoning_deltas_enabled() except Exception: logger.debug("reasoning on_stream_delta plugin hook enqueue failed", exc_info=True) + return + if enabled: + self._enqueue_stream_hook("on_stream_delta", label="reasoning on_stream_delta", delta=text, kind="reasoning") def _fire_tool_gen_started(self, tool_name: str) -> None: - """Notify the display layer that the model is generating tool call arguments. - - Fires once per tool name so the TUI can show a spinner while a large payload (e.g. a 45 KB write_file) - streams. - """ - cb = self.tool_gen_callback - if cb is not None: - try: - cb(tool_name) - except Exception: - pass + """Notify the display layer that the model is generating tool call arguments (spinner for large payloads).""" + self._call_quietly(self.tool_gen_callback, tool_name) def _has_stream_consumers(self) -> bool: """Return True if any streaming consumer is registered.""" @@ -449,7 +301,4 @@ class StreamDeliveryMixin: return True except Exception: logger.debug("plugin stream hook consumer check failed", exc_info=True) - return ( - self.stream_delta_callback is not None - or getattr(self, "_stream_callback", None) is not None - ) + return self.stream_delta_callback is not None or getattr(self, "_stream_callback", None) is not None diff --git a/agent/stream_diag.py b/agent/stream_diag.py index 61733afc22..5500ee4565 100644 --- a/agent/stream_diag.py +++ b/agent/stream_diag.py @@ -17,29 +17,14 @@ logger = logging.getLogger(__name__) # Lowercased upstream headers captured per attempt for post-hoc analysis. STREAM_DIAG_HEADERS = ( - "cf-ray", - "cf-cache-status", - "x-openrouter-provider", - "x-openrouter-model", - "x-openrouter-id", - "x-request-id", - "x-vercel-id", - "via", - "server", - "x-forwarded-for", + "cf-ray", "cf-cache-status", "x-openrouter-provider", "x-openrouter-model", "x-openrouter-id", + "x-request-id", "x-vercel-id", "via", "server", "x-forwarded-for", ) def stream_diag_init() -> Dict[str, Any]: """Fresh per-attempt diagnostic dict; mutated in place by the streaming functions and read by the retry block.""" - return { - "started_at": time.time(), - "first_chunk_at": None, - "chunks": 0, - "bytes": 0, - "headers": {}, - "http_status": None, - } + return {"started_at": time.time(), "first_chunk_at": None, "chunks": 0, "bytes": 0, "headers": {}, "http_status": None} def stream_diag_capture_response(agent: Any, diag: Dict[str, Any], http_response: Any) -> None: @@ -53,11 +38,9 @@ def stream_diag_capture_response(agent: Any, diag: Dict[str, Any], http_response try: headers = getattr(http_response, "headers", None) or {} captured: Dict[str, str] = {} - # Per-agent override of the headers list (back-compat). - for name in getattr(agent, "_STREAM_DIAG_HEADERS", STREAM_DIAG_HEADERS): + for name in getattr(agent, "_STREAM_DIAG_HEADERS", STREAM_DIAG_HEADERS): # per-agent override (back-compat) try: - val = headers.get(name) - if val: + if val := headers.get(name): captured[name] = str(val)[:120] # keep log lines bounded except Exception: continue @@ -67,48 +50,58 @@ def stream_diag_capture_response(agent: Any, diag: Dict[str, Any], http_response def flatten_exception_chain(error: BaseException) -> str: - """Compact ``Outer(msg) <- Inner(msg) <- ...`` rendering. - - The OpenAI SDK wraps httpx errors so only the wrapper class is visible at the - catch site; the inner RemoteProtocolError/ConnectError/ReadError says WHY the - stream died. Walks ``__cause__`` then ``__context__`` (deduped, max 4 deep). - """ + """Compact ``Outer(msg) <- Inner(msg) <- ...`` rendering, walking ``__cause__`` then ``__context__`` + (deduped, max 4 deep): the OpenAI SDK wraps httpx errors so only the wrapper class is visible at + the catch site; the inner RemoteProtocolError/ConnectError/ReadError says WHY the stream died.""" seen: List[BaseException] = [] link: Optional[BaseException] = error - while link is not None and len(seen) < 4: - if link in seen: - break + while link is not None and len(seen) < 4 and link not in seen: seen.append(link) nxt = getattr(link, "__cause__", None) or getattr(link, "__context__", None) if nxt is None or nxt is link: break link = nxt - parts: List[str] = [] - for e in seen: + + def render(e: BaseException) -> str: msg = str(e).strip().replace("\n", " ") - if len(msg) > 140: - msg = msg[:140] + "…" - parts.append(f"{type(e).__name__}({msg})" if msg else type(e).__name__) - return " <- ".join(parts) if parts else type(error).__name__ + msg = msg[:140] + "…" if len(msg) > 140 else msg + return f"{type(e).__name__}({msg})" if msg else type(e).__name__ + + return " <- ".join(render(e) for e in seen) if seen else type(error).__name__ + + +def _diag_fields(diag: Optional[Dict[str, Any]]) -> tuple: + """(http_status, bytes, chunks, elapsed, ttfb, upstream) for the retry log line; ``-`` when unknown.""" + _bytes = _chunks = 0 + _elapsed = 0.0 + _ttfb = _headers_repr = _http_status = "-" + if isinstance(diag, dict): + try: + _now = time.time() + _bytes = int(diag.get("bytes") or 0) + _chunks = int(diag.get("chunks") or 0) + _started = float(diag.get("started_at") or _now) + _elapsed = max(0.0, _now - _started) + _first = diag.get("first_chunk_at") + if _first is not None: + _ttfb = f"{max(0.0, float(_first) - _started):.2f}s" + headers = diag.get("headers") or {} + if isinstance(headers, dict) and headers: + _headers_repr = " ".join(f"{k}={v}" for k, v in headers.items()) + if diag.get("http_status") is not None: + _http_status = str(diag.get("http_status")) + except Exception: + pass + return _http_status, _bytes, _chunks, _elapsed, _ttfb, _headers_repr def log_stream_retry( - agent: Any, - *, - kind: str, - error: BaseException, - attempt: int, - max_attempts: int, - mid_tool_call: bool, - diag: Optional[Dict[str, Any]] = None, + agent: Any, *, kind: str, error: BaseException, attempt: int, max_attempts: int, + mid_tool_call: bool, diag: Optional[Dict[str, Any]] = None, ) -> None: - """Structured WARNING to ``agent.log`` for a transient stream drop + retry. - - Always logged regardless of UI verbosity (subagent retries no longer spam the - parent's terminal but keep full detail here). With *diag*, also records upstream - headers, HTTP status, bytes/chunks streamed, elapsed and TTFB on the dying attempt — - enough to tell "one CF edge / downstream provider" from "random across runs". - """ + """Structured WARNING to ``agent.log`` for a transient stream drop + retry, always logged regardless of + UI verbosity. With *diag*, also records upstream headers, HTTP status, bytes/chunks, elapsed and TTFB on + the dying attempt — enough to tell "one CF edge / downstream provider" from "random across runs".""" try: try: _summary = agent._summarize_api_error(error) @@ -121,53 +114,13 @@ def log_stream_retry( except Exception: _chain = type(error).__name__ - _now = time.time() - _bytes = 0 - _chunks = 0 - _elapsed = 0.0 - _ttfb = None - _headers_repr = "-" - _http_status = "-" - if isinstance(diag, dict): - try: - _bytes = int(diag.get("bytes") or 0) - _chunks = int(diag.get("chunks") or 0) - _started = float(diag.get("started_at") or _now) - _elapsed = max(0.0, _now - _started) - _first = diag.get("first_chunk_at") - if _first is not None: - _ttfb = max(0.0, float(_first) - _started) - headers = diag.get("headers") or {} - if isinstance(headers, dict) and headers: - _headers_repr = " ".join(f"{k}={v}" for k, v in headers.items()) - if diag.get("http_status") is not None: - _http_status = str(diag.get("http_status")) - except Exception: - pass - logger.warning( - "Stream %s on attempt %s/%s — retrying. " - "subagent_id=%s depth=%s provider=%s base_url=%s " - "error_type=%s error=%s " - "chain=%s " - "http_status=%s bytes=%d chunks=%d elapsed=%.2fs ttfb=%s " - "upstream=[%s]", - kind, - attempt, - max_attempts, - getattr(agent, "_subagent_id", None) or "-", - getattr(agent, "_delegate_depth", 0), - agent.provider or "-", - agent.base_url or "-", - type(error).__name__, - _summary, - _chain, - _http_status, - _bytes, - _chunks, - _elapsed, - f"{_ttfb:.2f}s" if _ttfb is not None else "-", - _headers_repr, + "Stream %s on attempt %s/%s — retrying. subagent_id=%s depth=%s provider=%s base_url=%s " + "error_type=%s error=%s chain=%s http_status=%s bytes=%d chunks=%d elapsed=%.2fs ttfb=%s upstream=[%s]", + kind, attempt, max_attempts, + getattr(agent, "_subagent_id", None) or "-", getattr(agent, "_delegate_depth", 0), + agent.provider or "-", agent.base_url or "-", + type(error).__name__, _summary, _chain, *_diag_fields(diag), extra={"mid_tool_call": mid_tool_call}, ) except Exception: @@ -175,32 +128,23 @@ def log_stream_retry( def emit_stream_drop( - agent: Any, - *, - error: BaseException, - attempt: int, - max_attempts: int, - mid_tool_call: bool, - diag: Optional[Dict[str, Any]] = None, + agent: Any, *, error: BaseException, attempt: int, max_attempts: int, + mid_tool_call: bool, diag: Optional[Dict[str, Any]] = None, ) -> None: """One compact user-visible status line for a stream drop+retry, plus the full WARNING via log_stream_retry. - - Subagent lines get a ``[subagent-N]`` prefix from ``log_prefix``. ``after Xs`` - distinguishes "couldn't connect" (0s) from "died mid-stream" (idle-kill / proxy timeout). - """ + ``after Xs`` distinguishes "couldn't connect" (0s) from "died mid-stream" (idle-kill / proxy timeout).""" kind = "drop mid tool-call" if mid_tool_call else "drop" log_stream_retry( agent, kind=kind, error=error, attempt=attempt, max_attempts=max_attempts, mid_tool_call=mid_tool_call, diag=diag ) provider = agent.provider or "provider" _suffix = "" - if isinstance(diag, dict): - try: - started = diag.get("started_at") - if started is not None: - _suffix = f" after {max(0.0, time.time() - float(started)):.1f}s" - except Exception: - pass + try: + started = diag.get("started_at") if isinstance(diag, dict) else None + if started is not None: + _suffix = f" after {max(0.0, time.time() - float(started)):.1f}s" + except Exception: + pass try: agent._buffer_status( f"⚠️ {provider} stream {kind} ({type(error).__name__}){_suffix} " diff --git a/agent/stream_single_writer.py b/agent/stream_single_writer.py index 52907f28a5..fd6ee1dc73 100644 --- a/agent/stream_single_writer.py +++ b/agent/stream_single_writer.py @@ -1,11 +1,10 @@ """Best-effort accessors for the single-writer stream fence. -The fence lives on ``AIAgent`` (``_claim_stream_writer`` / ``_stream_writer_is_current``) -but is used from other streaming modules. Calling it directly would turn an -*additive* safety net into a fatal AttributeError on a partially-updated checkout, -hot-reloaded gateway, duck-typed agent, or test double (a cron job died this way). -The fence may only drop a *provably* superseded stream, never the sole writer, so -when it is unavailable or raises the correct degradation is "no fence": keep streaming. +The fence lives on ``AIAgent`` (``_claim_stream_writer`` / ``_stream_writer_is_current``) but is +used from other streaming modules. Calling it directly would turn an *additive* safety net into a +fatal AttributeError on a partially-updated checkout, hot-reloaded gateway, duck-typed agent, or +test double (a cron job died this way). The fence may only drop a *provably* superseded stream, +never the sole writer, so when it is unavailable or raises the degradation is "no fence". """ from __future__ import annotations @@ -18,23 +17,22 @@ logger = logging.getLogger(__name__) def claim_stream_writer(agent: Any) -> int: """Claim the delta sink for this stream attempt; ``0`` (never fenced) when the agent lacks the fence or the claim raised.""" - claim = getattr(agent, "_claim_stream_writer", None) - if callable(claim): - try: - return int(claim()) - except Exception: - logger.debug("stream single-writer: claim failed; proceeding unfenced", exc_info=True) - return 0 + return _fence_call(agent, "_claim_stream_writer", int, 0, "claim failed; proceeding unfenced") def stream_writer_is_current(agent: Any, token: int) -> bool: """True when ``token`` is still the active writer; a falsy token or a fence-less agent cannot prove supersession, so True.""" if not token: return True - is_current = getattr(agent, "_stream_writer_is_current", None) - if callable(is_current): + return _fence_call(agent, "_stream_writer_is_current", bool, True, "is_current check failed; treating as current", token) + + +def _fence_call(agent: Any, name: str, cast, fallback, failure_note: str, *args): + """Call ``agent.(*args)`` when it exists; ``fallback`` when missing or raising (logged at debug).""" + fn = getattr(agent, name, None) + if callable(fn): try: - return bool(is_current(token)) + return cast(fn(*args)) except Exception: - logger.debug("stream single-writer: is_current check failed; treating as current", exc_info=True) - return True + logger.debug("stream single-writer: %s", failure_note, exc_info=True) + return fallback diff --git a/agent/subscription_view.py b/agent/subscription_view.py index 4a62017a29..9d390963d8 100644 --- a/agent/subscription_view.py +++ b/agent/subscription_view.py @@ -2,12 +2,9 @@ Companion to :mod:`agent.billing_view` — same fail-open philosophy (``logged_in=False`` when not logged in / portal unreachable; never crash) and decimal money end-to-end. - -The TUI ``SubscriptionOverlay`` drives the plan change in-terminal: preview, then -schedule a downgrade / cancellation / resume (chargeless) or apply an upgrade -(charges the subscription card). The portal deep-link (``portal_url`` + ``org_id``) -remains the fallback for an upgrade that needs 3DS / was declined. Until the NAS -``GET /api/billing/subscription`` endpoint ships, 404s take the fail-open path. +The TUI drives plan changes in-terminal (preview, then schedule a downgrade/cancel/resume +or apply an upgrade); the portal deep-link (``portal_url`` + ``org_id``) is the fallback +for an upgrade that needs 3DS / was declined. A 404 from NAS takes the fail-open path. """ from __future__ import annotations @@ -23,17 +20,10 @@ from agent.billing_view import OrgRoleCapability, fetch_portal_state, format_mon logger = logging.getLogger(__name__) -# ── Parsed sub-structures ──────────────────────────────────────────────────── - - @dataclass(frozen=True) class CurrentSubscription: - """The user's active subscription. ``None`` (not this object) = no plan. - - NAS guarantees a present ``current`` is fully populated: ``tier_id`` / - ``tier_name`` / ``monthly_credits`` / ``cycle_ends_at`` are always set; only - ``credits_remaining`` and the cancel/downgrade fields are optional. - """ + """Active subscription (``None``, not this object, = no plan). NAS guarantees ``tier_id`` / + ``tier_name`` / ``monthly_credits`` / ``cycle_ends_at`` are set; the rest are optional.""" tier_id: Optional[str] = None tier_name: Optional[str] = None @@ -48,12 +38,9 @@ class CurrentSubscription: @dataclass(frozen=True) class SubscriptionTier: - """One row of the tier picker (mirrors NAS ``SubscriptionTierOption``). - - ``is_current`` = active plan (shown, not selectable); ``is_enabled=False`` = a - grandfathered tier the user is on but can no longer select. ``tier_order`` sorts - the picker and drives the upgrade-vs-downgrade hint. - """ + """Tier-picker row (NAS ``SubscriptionTierOption``). ``is_current`` = active plan (shown, not + selectable); ``is_enabled=False`` = grandfathered, no longer selectable; ``tier_order`` sorts + the picker and drives the upgrade-vs-downgrade hint.""" tier_id: str name: str @@ -66,12 +53,9 @@ class SubscriptionTier: @dataclass(frozen=True) class SubscriptionChangePreview: - """Parsed ``POST /api/billing/subscription/preview``. - - ``effect``: ``charge_now`` (upgrade; ``amount_due_now_cents`` is the prorated - charge) · ``scheduled`` (downgrade / same-price change at ``effective_at``) · - ``no_op`` (already on target) · ``blocked`` (commit refused; ``reason`` says why). - """ + """Parsed ``POST /api/billing/subscription/preview``. ``effect``: ``charge_now`` (upgrade; prorated + ``amount_due_now_cents``) · ``scheduled`` (downgrade / same-price change at ``effective_at``) · + ``no_op`` (already on target) · ``blocked`` (commit refused; ``reason`` says why).""" effect: str reason: Optional[str] = None @@ -91,7 +75,7 @@ class SubscriptionState(OrgRoleCapability): logged_in: bool org_name: Optional[str] = None - org_id: Optional[str] = None # org.id from the NAS response + org_id: Optional[str] = None role: Optional[str] = None # "OWNER" | "ADMIN" | "FINANCE_ADMIN" | "SECURITY_ADMIN" | "MEMBER" can_change_plan_raw: Optional[bool] = None context: str = "personal" # "personal" | "team" @@ -104,12 +88,13 @@ class SubscriptionState(OrgRoleCapability): # ── Payload parsing ────────────────────────────────────────────────────────── +def _tier_id(raw: Any) -> Optional[str]: + """Real tier id of a NAS dict, else None ("no plan" is ``current: null``; junk is skipped).""" + return (raw.get("tierId") or raw.get("id") or None) if isinstance(raw, dict) else None + + def _parse_current(raw: Any) -> Optional[CurrentSubscription]: - # "No plan" is wire-represented as current:null; a present current is a real - # plan, so guard on a real tier id and return None otherwise. - if not isinstance(raw, dict): - return None - tier_id = raw.get("tierId") or raw.get("id") + tier_id = _tier_id(raw) if not tier_id: return None return CurrentSubscription( @@ -126,19 +111,12 @@ def _parse_current(raw: Any) -> Optional[CurrentSubscription]: def _coalesce(*vals: Any) -> Any: - """First non-``None`` value. NAS sends ``0`` for the free tier's ``tierOrder`` / - ``dollarsPerMonth``, which a plain ``x or default`` would drop.""" - for v in vals: - if v is not None: - return v - return None + """First non-``None`` value (NAS sends ``0`` for the free tier, which ``x or default`` would drop).""" + return next((v for v in vals if v is not None), None) def _parse_tier(raw: Any) -> Optional[SubscriptionTier]: - """Map one NAS ``SubscriptionTierOption`` dict into a :class:`SubscriptionTier`.""" - if not isinstance(raw, dict): - return None - tier_id = raw.get("tierId") or raw.get("id") + tier_id = _tier_id(raw) if not tier_id: return None return SubscriptionTier( @@ -152,9 +130,7 @@ def _parse_tier(raw: Any) -> Optional[SubscriptionTier]: ) -def subscription_change_preview_from_payload( - payload: dict[str, Any], -) -> SubscriptionChangePreview: +def subscription_change_preview_from_payload(payload: dict[str, Any]) -> SubscriptionChangePreview: """Map a raw ``/subscription/preview`` JSON dict into :class:`SubscriptionChangePreview`.""" effect = payload.get("effect") cents = payload.get("amountDueNowCents") @@ -172,18 +148,10 @@ def subscription_change_preview_from_payload( ) -def subscription_state_from_payload( - payload: dict[str, Any], *, portal_url: Optional[str] = None -) -> SubscriptionState: +def subscription_state_from_payload(payload: dict[str, Any], *, portal_url: Optional[str] = None) -> SubscriptionState: """Map a raw ``/api/billing/subscription`` JSON dict into :class:`SubscriptionState`.""" org, can_change_plan_raw = parse_org_fields(payload) - raw_context = payload.get("context") - raw_tiers = payload.get("tiers") - tiers = ( - tuple(t for t in (_parse_tier(x) for x in raw_tiers) if t is not None) - if isinstance(raw_tiers, list) - else () - ) + raw_context, raw_tiers = payload.get("context"), payload.get("tiers") return SubscriptionState( logged_in=True, org_name=org.get("name"), @@ -192,7 +160,7 @@ def subscription_state_from_payload( can_change_plan_raw=can_change_plan_raw, context=raw_context if raw_context in ("personal", "team") else "personal", current=_parse_current(payload.get("current")), - tiers=tiers, + tiers=tuple(filter(None, map(_parse_tier, raw_tiers))) if isinstance(raw_tiers, list) else (), portal_url=portal_url, ) @@ -201,44 +169,31 @@ def subscription_state_from_payload( def build_subscription_state(*, timeout: float = 15.0) -> SubscriptionState: - """Fetch + parse ``GET /api/billing/subscription``. Fail-open (see - :func:`agent.billing_view.fetch_portal_state`). - - ``HERMES_DEV_SUBSCRIPTION_FIXTURE`` short-circuits to a fixture so every - plan/cancel/downgrade/team/not-admin state is testable on CLI and TUI offline. - """ + """Fetch + parse ``GET /api/billing/subscription``; fail-open like ``fetch_portal_state``. + ``HERMES_DEV_SUBSCRIPTION_FIXTURE`` short-circuits to a fixture so every state is testable offline.""" fixture = dev_fixture_subscription_state() if fixture is not None: return fixture return fetch_portal_state( - "get_subscription_state", - "subscription", + "get_subscription_state", "subscription", failed=lambda **kw: SubscriptionState(logged_in=False, **kw), parse=lambda payload, portal_url: subscription_state_from_payload(payload, portal_url=portal_url), - portal_fallback=lambda base: base, - timeout=timeout, - log=logger, + portal_fallback=lambda base: base, timeout=timeout, log=logger, ) -def subscription_manage_url( - state: SubscriptionState, tier_id: Optional[str] = None -) -> Optional[str]: - """Build ``{portal_origin}/manage-subscription?org_id=[&plan=]``. +def subscription_manage_url(state: SubscriptionState, tier_id: Optional[str] = None) -> Optional[str]: + """Build ``{portal_origin}/manage-subscription?org_id=[&plan=]`` (None if unresolvable). - Mirrors the TUI's ``buildManageUrl``: the target is NAS's OWN ``/manage-subscription`` - page (NOT the Stripe Billing Portal), which routes upgrade→Checkout / - downgrade→scheduled internally. ``org_id`` pins the right account in multi-org - situations. ``tier_id`` (the stable ``tiers[]`` id, never a name/slug) preselects - the picked plan; the portal ignores an unknown tier, so it's appended - unconditionally when picked. None when no portal URL is resolvable. + Mirrors the TUI's ``buildManageUrl``: the target is NAS's OWN ``/manage-subscription`` page + (NOT the Stripe Billing Portal). ``org_id`` pins the account in multi-org situations; ``tier_id`` + (the stable ``tiers[]`` id, never a name/slug) preselects the plan — the portal ignores an + unknown tier, so it's appended unconditionally. """ from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit - if not state.portal_url: - return None try: - parts = urlsplit(state.portal_url) + parts = urlsplit(state.portal_url or "") except Exception: return None if parts.scheme not in ("http", "https") or not parts.netloc: @@ -260,8 +215,7 @@ def subscription_manage_url( def selectable_tiers(state: SubscriptionState) -> list[SubscriptionTier]: - """Enabled paid tiers other than the current plan, cheapest first (``tier_order > 0`` - — dropping to free is a cancellation, not a plan pick).""" + """Enabled paid tiers other than the current plan, cheapest first (dropping to free is a cancellation).""" return sorted( (t for t in (state.tiers or ()) if t.is_enabled and not t.is_current and (t.tier_order or 0) > 0), key=lambda t: t.tier_order or 0, @@ -269,8 +223,7 @@ def selectable_tiers(state: SubscriptionState) -> list[SubscriptionTier]: def format_tier_row(tier: SubscriptionTier) -> str: - """``name · $X/mo[ · $Y credits/mo]`` — thousands-grouped money (mirrors the TUI - Free rows); the credits suffix appears ONLY when monthly credits are present and > 0.""" + """``name · $X/mo[ · $Y credits/mo]`` (grouped money, like the TUI); credits suffix only when > 0.""" row = f"{tier.name} · {format_money(tier.dollars_per_month, grouped=True)}/mo" mc = tier.monthly_credits if mc is not None and mc > 0: @@ -281,70 +234,61 @@ def format_tier_row(tier: SubscriptionTier) -> str: def is_upgrade(state: SubscriptionState, tier_id: str) -> bool: """True when ``tier_id`` ranks above the current plan by ``tier_order``. Prefers the active subscription's tier; falls back to the ``tiers[]`` ``is_current`` marker, else 0.""" - orders = {t.tier_id: (t.tier_order or 0) for t in (state.tiers or ())} + tiers = state.tiers or () + orders = {t.tier_id: (t.tier_order or 0) for t in tiers} cur_id = state.current.tier_id if state.current else None - if cur_id is not None and cur_id in orders: - cur_order = orders[cur_id] - else: - cur_order = next((t.tier_order or 0 for t in (state.tiers or ()) if t.is_current), 0) + cur_order = orders[cur_id] if cur_id in orders else next((t.tier_order or 0 for t in tiers if t.is_current), 0) return orders.get(tier_id, 0) > cur_order -# ── Dev fixtures (throwaway scaffolding — env-var driven, no live portal) ──── +# ── Dev fixtures (env-var driven, no live portal) ──────────────────────────── _DEV_FIXTURE_PORTAL = "https://portal.nousresearch.com/billing" - - -def _dev_current(**over: Any) -> CurrentSubscription: - base: dict[str, Any] = dict( - tier_id="plus", tier_name="Plus", monthly_credits=Decimal("1000"), - credits_remaining=Decimal("420"), cycle_ends_at="2026-07-01", - ) - return CurrentSubscription(**{**base, **over}) +_DEV_TIER_SPECS = (("free", "Free", 0, "0", "0"), ("plus", "Plus", 1, "20", "1000"), + ("super", "Super", 2, "40", "3000"), ("ultra", "Ultra", 3, "80", "7000")) +_DEV_FIXTURE_ALIASES = {"logged_out": "logged-out", "loggedout": "logged-out", "mid-tier": "mid", + "top-tier": "top", "member": "not-admin"} def _dev_tiers(current_id: Optional[str]) -> tuple[SubscriptionTier, ...]: - """A sample plan catalog for fixtures (marks ``current_id`` as the active tier).""" - specs = (("free", "Free", 0, "0", "0"), ("plus", "Plus", 1, "20", "1000"), - ("super", "Super", 2, "40", "3000"), ("ultra", "Ultra", 3, "80", "7000")) + """Sample plan catalog for fixtures (marks ``current_id`` as the active tier).""" return tuple( SubscriptionTier( tier_id=tid, name=name, tier_order=order, dollars_per_month=parse_money(dpm), monthly_credits=parse_money(mc), is_current=(tid == current_id), is_enabled=True, ) - for tid, name, order, dpm, mc in specs + for tid, name, order, dpm, mc in _DEV_TIER_SPECS ) -def dev_fixture_subscription_state() -> Optional[SubscriptionState]: - """``HERMES_DEV_SUBSCRIPTION_FIXTURE`` -> fixture :class:`SubscriptionState`; None when unset. +def _dev_plan(tier_id: str, remaining: str, **over: Any) -> dict[str, Any]: + """``current`` + ``tiers`` fixture fields for being on ``tier_id`` with ``remaining`` credits left.""" + tid, name, _order, _dpm, mc = next(spec for spec in _DEV_TIER_SPECS if spec[0] == tier_id) + current = CurrentSubscription( + tier_id=tid, tier_name=name, monthly_credits=Decimal(mc), credits_remaining=Decimal(remaining), + cycle_ends_at="2026-07-01", **over, + ) + return dict(current=current, tiers=_dev_tiers(tid)) - ``free | mid | top | not-admin | downgrade | cancel | team | logged-out``. Unknown - name → logged-out with ``error`` so the misconfiguration is visible. - """ + +def dev_fixture_subscription_state() -> Optional[SubscriptionState]: + """``HERMES_DEV_SUBSCRIPTION_FIXTURE`` (``free | mid | top | not-admin | downgrade | cancel | team | + logged-out``) -> fixture state; None when unset; unknown name → logged-out with ``error`` set.""" name = (os.getenv("HERMES_DEV_SUBSCRIPTION_FIXTURE") or "").strip().lower() if not name: return None - name = {"logged_out": "logged-out", "loggedout": "logged-out", "mid-tier": "mid", - "top-tier": "top", "member": "not-admin"}.get(name, name) + name = _DEV_FIXTURE_ALIASES.get(name, name) if name == "logged-out": return SubscriptionState(logged_in=False) common = dict(logged_in=True, org_name="Acme Inc", org_id="org_acme", role="OWNER", portal_url=_DEV_FIXTURE_PORTAL) - plus = dict(current=_dev_current(), tiers=_dev_tiers("plus")) states: dict[str, dict[str, Any]] = { "free": dict(current=None, tiers=_dev_tiers(None)), - "mid": plus, - "top": dict( - current=_dev_current(tier_id="ultra", tier_name="Ultra", monthly_credits=Decimal("7000"), credits_remaining=Decimal("5000")), - tiers=_dev_tiers("ultra"), - ), - "not-admin": {**plus, "role": "MEMBER"}, - "downgrade": dict( - current=_dev_current(tier_id="super", tier_name="Super", monthly_credits=Decimal("3000"), credits_remaining=Decimal("1500"), pending_downgrade_tier_name="Plus", pending_downgrade_at="2026-07-15"), - tiers=_dev_tiers("super"), - ), - "cancel": dict(current=_dev_current(cancel_at_period_end=True, cancellation_effective_at="2026-07-01"), tiers=_dev_tiers("plus")), + "mid": _dev_plan("plus", "420"), + "top": _dev_plan("ultra", "5000"), + "not-admin": {**_dev_plan("plus", "420"), "role": "MEMBER"}, + "downgrade": _dev_plan("super", "1500", pending_downgrade_tier_name="Plus", pending_downgrade_at="2026-07-15"), + "cancel": _dev_plan("plus", "420", cancel_at_period_end=True, cancellation_effective_at="2026-07-01"), "team": dict(context="team", current=None, org_name="Acme Engineering", org_id="org_eng"), } if name not in states: diff --git a/agent/think_scrubber.py b/agent/think_scrubber.py index 82a27e3b2a..5ded655d3e 100644 --- a/agent/think_scrubber.py +++ b/agent/think_scrubber.py @@ -1,32 +1,16 @@ """Stateful scrubber for reasoning/thinking blocks in streamed assistant text. -The regex ``run_agent._strip_think_blocks`` is correct for a complete string but, -run per-delta, erases an opening ```` that arrives alone in one delta, so -downstream state machines never see the open tag and leak reasoning. This class -centralises tag suppression upstream: partial tags at delta boundaries are held -back until resolved, and ``flush()`` releases held-back prose that was not a tag. - -Usage:: - - scrubber = StreamingThinkScrubber() - for delta in stream: - visible = scrubber.feed(delta) - if visible: - emit(visible) - tail = scrubber.flush() # at end of stream - -Call ``reset()`` at the top of each turn so an interrupted block cannot taint -the next turn. Tags handled (case-insensitive): ````, ````, -````, ````, ````. - -Boundary rule: an opening tag only starts a block at a block boundary (stream -start, after a newline, or with only whitespace emitted on the current line), so -prose that *mentions* ```` is not suppressed. Closed pairs -(``X``) are always suppressed — a closed pair is intentional. +The regex ``_strip_think_blocks`` is correct for a complete string but, run per-delta, erases an +opening ```` that arrives alone, so downstream state machines leak reasoning. This class +holds partial tags at delta boundaries until resolved; ``flush()`` releases held-back prose that +was not a tag; ``reset()`` at the top of each turn. An open tag only starts a block at a block +boundary (stream start / after a newline / whitespace-only line so far), so prose that *mentions* +```` is not suppressed; closed pairs are always suppressed (intentional). """ from __future__ import annotations +import re from typing import Tuple __all__ = ["StreamingThinkScrubber"] @@ -35,24 +19,22 @@ __all__ = ["StreamingThinkScrubber"] class StreamingThinkScrubber: """Stateful scrubber for streaming reasoning/thinking blocks. - State: ``_in_block`` (inside an open block; text discarded), ``_buf`` - (held-back partial-tag tail), ``_last_emitted_ended_newline`` (True iff the - last emission ended with ``\\n`` or nothing has been emitted yet — decides - whether an open tag at buffer position 0 sits at a block boundary). + State: ``_in_block`` (inside an open block; text discarded), ``_buf`` (held-back partial-tag + tail), ``_last_emitted_ended_newline`` (True iff the last emission ended with ``\\n`` or nothing + was emitted yet — decides whether an open tag at buffer position 0 sits at a block boundary). """ - _OPEN_TAG_NAMES: Tuple[str, ...] = ( - "think", - "thinking", - "reasoning", - "thought", - "REASONING_SCRATCHPAD", - ) + _OPEN_TAG_NAMES: Tuple[str, ...] = ("think", "thinking", "reasoning", "thought", "REASONING_SCRATCHPAD") - # Literal tag strings so the hot path does string ops, not regex per feed(). - _OPEN_TAGS: Tuple[str, ...] = tuple(f"<{name}>" for name in _OPEN_TAG_NAMES) - _CLOSE_TAGS: Tuple[str, ...] = tuple(f"" for name in _OPEN_TAG_NAMES) - _MAX_TAG_LEN: int = max(len(tag) for tag in _OPEN_TAGS + _CLOSE_TAGS) + # Lowercased literal tags so the hot path does string ops, not regex per feed(). + _OPEN_TAGS: Tuple[str, ...] = tuple(f"<{name.lower()}>" for name in _OPEN_TAG_NAMES) + _CLOSE_TAGS: Tuple[str, ...] = tuple(f"" for name in _OPEN_TAG_NAMES) + _ALL_TAGS: Tuple[str, ...] = _OPEN_TAGS + _CLOSE_TAGS + _MAX_TAG_LEN: int = max(len(tag) for tag in _ALL_TAGS) + # Orphan close tag plus trailing whitespace (matches _strip_think_blocks case 3). + _ORPHAN_CLOSE_RE = re.compile( + "(?:" + "|".join(re.escape(t) for t in _CLOSE_TAGS) + r")[ \t\n\r]*", re.IGNORECASE + ) def __init__(self) -> None: self.reset() @@ -65,18 +47,13 @@ class StreamingThinkScrubber: def _emit(self, out: list[str], text: str) -> None: """Append visible prose to *out* (orphan close tags stripped) and track the newline flag.""" + text = self._strip_orphan_close_tags(text) if text: - text = self._strip_orphan_close_tags(text) - if text: - out.append(text) - self._last_emitted_ended_newline = text.endswith("\n") + out.append(text) + self._last_emitted_ended_newline = text.endswith("\n") def feed(self, text: str) -> str: - """Feed one delta; return the scrubbed visible portion. - - Returns "" when the whole delta is reasoning content or is held back - pending resolution of a partial tag at the boundary. - """ + """Feed one delta; return the scrubbed visible portion ("" when it is all reasoning or held back).""" if not text: return "" buf = self._buf + text @@ -88,17 +65,15 @@ class StreamingThinkScrubber: close_idx, close_len = self._find_first_tag(buf, self._CLOSE_TAGS) if close_idx == -1: # No close yet: hold back a possible partial close-tag prefix, drop the rest. - held = self._max_partial_suffix(buf, self._CLOSE_TAGS) - self._buf = buf[-held:] if held else "" - return "".join(out) + self._hold_partial(buf, self._CLOSE_TAGS) + break buf = buf[close_idx + close_len:] self._in_block = False continue - # Priority 1: closed X pair anywhere (no boundary gating — - # even inline pairs are almost certainly leaked reasoning). - # Priority 2: unterminated open tag at a block boundary (gated so - # prose that mentions '' isn't over-stripped). Earliest wins. + # Priority 1: closed X pair anywhere (even inline pairs are almost + # certainly leaked reasoning). Priority 2: unterminated open tag at a block + # boundary (gated so prose mentioning '' isn't over-stripped). Earliest wins. pair = self._find_earliest_closed_pair(buf) open_idx, open_len = self._find_open_at_boundary(buf, out) if pair is not None and (open_idx == -1 or pair[0] <= open_idx): @@ -113,28 +88,22 @@ class StreamingThinkScrubber: # No resolvable tag: hold back any partial-tag prefix at the tail # so a tag split across deltas isn't missed, then emit the rest. - held = max( - self._max_partial_suffix(buf, self._OPEN_TAGS), - self._max_partial_suffix(buf, self._CLOSE_TAGS), - ) - if held: - self._emit(out, buf[:-held]) - self._buf = buf[-held:] - else: - self._emit(out, buf) - return "".join(out) + self._emit(out, self._hold_partial(buf, self._ALL_TAGS)) + break return "".join(out) - def flush(self) -> str: - """End-of-stream flush. + def _hold_partial(self, buf: str, tags: Tuple[str, ...]) -> str: + """Move a trailing partial-tag prefix of *buf* into ``_buf``; return the remainder.""" + held = self._max_partial_suffix(buf, tags) + self._buf = buf[-held:] if held else "" + return buf[:-held] if held else buf - Inside an unterminated block the held-back content is discarded (leaking - partial reasoning is worse than a truncated answer); otherwise the - held-back tail is emitted verbatim. Always resets the boundary flag: - intra-turn retries flush then stream again without ``reset()``, and a - stale False flag made the new stream's opening ```` look mid-line. - """ + def flush(self) -> str: + """End-of-stream flush: inside an unterminated block the held-back content is discarded (leaking + partial reasoning is worse than a truncated answer), otherwise the tail is emitted verbatim. + Always resets the boundary flag — intra-turn retries flush then stream again without ``reset()``, + and a stale False flag made the new stream's opening ```` look mid-line.""" tail = "" if self._in_block else self._buf self._buf = "" self._in_block = False @@ -147,116 +116,54 @@ class StreamingThinkScrubber: def _find_first_tag(buf: str, tags: Tuple[str, ...]) -> Tuple[int, int]: """Return (earliest_index, tag_length) over *tags* (case-insensitive), or (-1, 0).""" buf_lower = buf.lower() - best_idx = -1 - best_len = 0 - for tag in tags: - idx = buf_lower.find(tag.lower()) - if idx != -1 and (best_idx == -1 or idx < best_idx): - best_idx = idx - best_len = len(tag) - return best_idx, best_len + hits = [(idx, len(tag)) for tag in tags if (idx := buf_lower.find(tag)) != -1] + return min(hits) if hits else (-1, 0) def _find_earliest_closed_pair(self, buf: str): - """Return (start_idx, end_idx) of the earliest ``...`` pair, else None. - - Case-insensitive and non-greedy (closest close after the open wins), - matching ``_strip_think_blocks`` case 1; the earliest open tag wins. - """ + """(start_idx, end_idx) of the earliest ``...`` pair (non-greedy, case-insensitive), else None.""" buf_lower = buf.lower() - best: "tuple[int, int] | None" = None + pairs = [] for open_tag, close_tag in zip(self._OPEN_TAGS, self._CLOSE_TAGS): - open_lower = open_tag.lower() - close_lower = close_tag.lower() - open_idx = buf_lower.find(open_lower) - if open_idx == -1: - continue - close_idx = buf_lower.find(close_lower, open_idx + len(open_lower)) - if close_idx == -1: - continue - if best is None or open_idx < best[0]: - best = (open_idx, close_idx + len(close_lower)) - return best + open_idx = buf_lower.find(open_tag) + close_idx = buf_lower.find(close_tag, open_idx + len(open_tag)) if open_idx != -1 else -1 + if close_idx != -1: + pairs.append((open_idx, close_idx + len(close_tag))) + return min(pairs) if pairs else None def _find_open_at_boundary(self, buf: str, already_emitted: list[str]) -> Tuple[int, int]: """Return the earliest block-boundary open-tag (idx, len), or (-1, 0).""" buf_lower = buf.lower() - best_idx = -1 - best_len = 0 + hits = [] for tag in self._OPEN_TAGS: - tag_lower = tag.lower() - search_start = 0 - while True: - idx = buf_lower.find(tag_lower, search_start) - if idx == -1: - break - if self._is_block_boundary(buf, idx, already_emitted): - if best_idx == -1 or idx < best_idx: - best_idx = idx - best_len = len(tag) - break # first boundary hit for this tag is enough - search_start = idx + 1 - return best_idx, best_len + idx = buf_lower.find(tag) + while idx != -1 and not self._is_block_boundary(buf, idx, already_emitted): + idx = buf_lower.find(tag, idx + 1) + if idx != -1: + hits.append((idx, len(tag))) + return min(hits) if hits else (-1, 0) def _is_block_boundary(self, buf: str, idx: int, already_emitted: list[str]) -> bool: - """True iff position *idx* in *buf* is a block boundary. - - Boundary = position 0 with the prior emission ending in a newline (or - nothing emitted yet), or any position whose preceding text on the current - line is whitespace-only (when no newline precedes it in *buf*, the prior - emission must also have ended with a newline). - """ - prior_newline = ( - already_emitted[-1].endswith("\n") if already_emitted else self._last_emitted_ended_newline - ) + """True iff *idx* is a block boundary: position 0 after a newline-terminated (or no) prior emission, + or any position whose preceding text on the current line is whitespace-only (when no newline + precedes it in *buf*, the prior emission must also have ended with a newline).""" + prior_newline = already_emitted[-1].endswith("\n") if already_emitted else self._last_emitted_ended_newline if idx == 0: return prior_newline preceding = buf[:idx] last_nl = preceding.rfind("\n") - if last_nl == -1: - return prior_newline and preceding.strip() == "" - return preceding[last_nl + 1:].strip() == "" + return (prior_newline if last_nl == -1 else True) and preceding[last_nl + 1:].strip() == "" @classmethod def _max_partial_suffix(cls, buf: str, tags: Tuple[str, ...]) -> int: - """Longest buf-suffix that is a strict prefix of any tag (case-insensitive). - - Full-length matches are real tags handled elsewhere, not held-back partials. - """ - if not buf: - return 0 + """Longest buf-suffix that is a strict prefix of any tag (full matches are real tags, handled elsewhere).""" buf_lower = buf.lower() - max_check = min(len(buf_lower), cls._MAX_TAG_LEN - 1) - for i in range(max_check, 0, -1): + for i in range(min(len(buf_lower), cls._MAX_TAG_LEN - 1), 0, -1): suffix = buf_lower[-i:] - for tag in tags: - tag_lower = tag.lower() - if len(tag_lower) > i and tag_lower.startswith(suffix): - return i + if any(len(tag) > i and tag.startswith(suffix) for tag in tags): + return i return 0 @classmethod def _strip_orphan_close_tags(cls, text: str) -> str: """Remove close tags with no matching open (always noise) plus trailing whitespace.""" - if " TextIO: - if self._state.silenced.get(threading.get_ident(), 0) > 0: - return self._state.sink - return self._passthrough + return self._state.sink if self._state.silenced.get(threading.get_ident(), 0) > 0 else self._passthrough def silence(self, ident: int) -> None: with self._state.lock: @@ -61,24 +58,21 @@ class _ThreadRoutingStream: else: self._state.silenced.pop(ident, None) - def write(self, data): # type: ignore[no-untyped-def] + def _forward(self, name: str, fallback, *args): # type: ignore[no-untyped-def] + """Call ``name`` on the current target; a dead target yields ``fallback(*args)`` instead of raising.""" try: - return self._target().write(data) + return getattr(self._target(), name)(*args) except Exception: - return len(data) if isinstance(data, str) else 0 + return fallback(*args) + + def write(self, data): # type: ignore[no-untyped-def] + return self._forward("write", lambda d: len(d) if isinstance(d, str) else 0, data) def flush(self): # type: ignore[no-untyped-def] - try: - return self._target().flush() - except Exception: - return None + return self._forward("flush", lambda: None) def writelines(self, lines): # type: ignore[no-untyped-def] - target = self._target() - try: - return target.writelines(lines) - except Exception: - return None + return self._forward("writelines", lambda _l: None, lines) def isatty(self) -> bool: try: @@ -111,15 +105,12 @@ def _ensure_installed(attr: str, passthrough: TextIO) -> "_ThreadRoutingStream": passthrough = current if current is not None else passthrough sink = _sinks.get(attr) if sink is None or sink.closed: - sink = open(os.devnull, "w", encoding="utf-8") - _sinks[attr] = sink + sink = _sinks[attr] = open(os.devnull, "w", encoding="utf-8") state = _routing_states.get(attr) if state is None or state.sink is not sink: - state = _RoutingState(sink) - _routing_states[attr] = state - proxy = _ThreadRoutingStream(passthrough, state) + state = _routing_states[attr] = _RoutingState(sink) + proxy = _installed[attr] = _ThreadRoutingStream(passthrough, state) setattr(sys, attr, proxy) - _installed[attr] = proxy return proxy @@ -127,12 +118,11 @@ def _ensure_installed(attr: str, passthrough: TextIO) -> "_ThreadRoutingStream": def thread_scoped_silence() -> Iterator[None]: """Silence ``stdout``/``stderr`` for the *current thread only*.""" ident = threading.get_ident() - out_proxy = _ensure_installed("stdout", sys.__stdout__ or sys.stdout) - err_proxy = _ensure_installed("stderr", sys.__stderr__ or sys.stderr) - out_proxy.silence(ident) - err_proxy.silence(ident) + proxies = (_ensure_installed("stdout", sys.__stdout__ or sys.stdout), _ensure_installed("stderr", sys.__stderr__ or sys.stderr)) + for proxy in proxies: + proxy.silence(ident) try: yield finally: - out_proxy.unsilence(ident) - err_proxy.unsilence(ident) + for proxy in proxies: + proxy.unsilence(ident) diff --git a/agent/tool_guardrails.py b/agent/tool_guardrails.py index 6ec3b289ae..f3d49605e0 100644 --- a/agent/tool_guardrails.py +++ b/agent/tool_guardrails.py @@ -9,7 +9,7 @@ from __future__ import annotations import hashlib import json -from dataclasses import dataclass, field +from dataclasses import asdict, dataclass, field, fields from typing import Any, Mapping from utils import safe_json_loads @@ -17,12 +17,11 @@ from agent.tool_result_classification import file_mutation_result_landed IDEMPOTENT_TOOL_NAMES = frozenset({ - "read_file", "search_files", "web_search", "web_extract", "session_search", "skill_view", - "skills_list", "browser_snapshot", "browser_console", "browser_get_images", - "mcp_filesystem_read_file", "mcp_filesystem_read_text_file", - "mcp_filesystem_read_multiple_files", "mcp_filesystem_list_directory", - "mcp_filesystem_list_directory_with_sizes", "mcp_filesystem_directory_tree", - "mcp_filesystem_get_file_info", "mcp_filesystem_search_files", + "read_file", "search_files", "web_search", "web_extract", "session_search", "skill_view", "skills_list", + "browser_snapshot", "browser_console", "browser_get_images", "mcp_filesystem_read_file", + "mcp_filesystem_read_text_file", "mcp_filesystem_read_multiple_files", "mcp_filesystem_list_directory", + "mcp_filesystem_list_directory_with_sizes", "mcp_filesystem_directory_tree", "mcp_filesystem_get_file_info", + "mcp_filesystem_search_files", }) MUTATING_TOOL_NAMES = frozenset({ @@ -31,55 +30,83 @@ MUTATING_TOOL_NAMES = frozenset({ "send_message", "cronjob_manage", "delegate_task", "process_manage", }) -# Tools legitimately re-invoked with identical args while waiting on external -# progress (pollers). The identical-call NOTICE never fires for these. +# Pollers: legitimately re-invoked with identical args; the identical-call NOTICE never fires. STALL_GUARD_REPEATABLE_TOOLS = frozenset({"process_manage"}) -# Poller naming conventions on generated / MCP surfaces (``_get_result``). -_STALL_GUARD_REPEATABLE_SUFFIXES = ("_get_result", "_poll") - -# Notice fires on the Nth consecutive identical (tool, args, result) call; 3 -# tolerates one legitimate double-check while catching observed re-issue loops. +_STALL_GUARD_REPEATABLE_SUFFIXES = ("_get_result", "_poll") # generated / MCP poller conventions +# Nth consecutive identical (tool, args, result) call that fires the notice; 3 tolerates one double-check. STALL_GUARD_IDENTICAL_CALL_THRESHOLD = 3 - -# Result-reference stubbing: from the 2nd consecutive identical call whose fresh -# result is byte-identical, the duplicate payload is replaced by a reference -# stub. Results under this size aren't worth stubbing; errors are never stubbed. +# From the 2nd byte-identical repeat the duplicate payload becomes a reference stub; smaller results +# aren't worth it, errors never are. The args preview keeps WHAT was called if compression evicts the original. IDENTICAL_RESULT_STUB_MIN_CHARS = 512 -# Canonical-args preview kept in the stub so the model still knows WHAT the call -# was if compression later evicts the referenced result. _RESULT_STUB_ARGS_PREVIEW_CHARS = 120 -# Tools whose "failure" is normal work output (red test run, empty grep, page -# timeout). same_tool_failure (DIFFERENT commands) never halts these; only an -# exact-args replay with no intervening change, or an identical-result streak, can. +# Tools whose "failure" is normal work output (red test run, empty grep, page timeout). +# same_tool_failure (DIFFERENT commands) never halts these; only an exact-args replay with +# no intervening change, or an identical-result streak, can. FAILURE_TOLERANT_TOOL_NAMES = frozenset({ "terminal", "execute_code", "process_manage", "process", "browser_navigate", "web_extract", }) -# A successful call to one of these marks progress for every failing signature -# still counted this turn: the next retry is a new experiment (edit -> re-run), not a replay. +# A successful call to one of these marks progress for every failing signature still counted +# this turn: the next retry is a new experiment (edit -> re-run), not a replay. PROGRESS_RESET_TOOL_NAMES = frozenset({ - "write_file", "patch", "terminal", "execute_code", "browser_click", "browser_type", - "browser_press", "browser_navigate", "process_manage", "process", "delegate_task", - "send_message", "cronjob", "cronjob_manage", "todo", "todo_list", "memory", "skill_manage", + "write_file", "patch", "terminal", "execute_code", "browser_click", "browser_type", "browser_press", + "browser_navigate", "process_manage", "process", "delegate_task", "send_message", "cronjob", + "cronjob_manage", "todo", "todo_list", "memory", "skill_manage", }) +_BOOL_FIELDS = ("warnings_enabled", "hard_stop_enabled", "non_interactive_hard_stop_enabled") +# Threshold field -> (nested section, nested key). The flat legacy key is the field name itself. +_THRESHOLD_SOURCES: dict[str, tuple[str, str]] = { + "exact_failure_warn_after": ("warn_after", "exact_failure"), + "same_tool_failure_warn_after": ("warn_after", "same_tool_failure"), + "no_progress_warn_after": ("warn_after", "idempotent_no_progress"), + "exact_failure_block_after": ("hard_stop_after", "exact_failure"), + "same_tool_failure_halt_after": ("hard_stop_after", "same_tool_failure"), + "no_progress_block_after": ("hard_stop_after", "idempotent_no_progress"), +} + +# Per-turn caps on runaway-prone tools (counters reset in reset_for_turn). +_DEFAULT_MAX_WEB_SEARCHES_PER_TURN = 50 +_DEFAULT_MAX_SUBAGENTS_PER_TURN = 50 + +# Interactive surfaces plus bounded supervised task loops (subagent stopped by its parent; +# api_server has a live client) doing real edit -> re-run work keep the warn-only default. +_ATTENDED_PLATFORMS = frozenset({"cli", "tui", "desktop", "acp", "subagent", "api_server"}) + def is_stall_guard_repeatable(tool_name: str) -> bool: """Whether a tool is exempt from the identical-call loop notice.""" - return tool_name in STALL_GUARD_REPEATABLE_TOOLS or tool_name.endswith( - _STALL_GUARD_REPEATABLE_SUFFIXES - ) + return tool_name in STALL_GUARD_REPEATABLE_TOOLS or tool_name.endswith(_STALL_GUARD_REPEATABLE_SUFFIXES) + + +def _is_non_interactive_platform(platform: str | None) -> bool: + """True for gateway/cron sessions where tool loops are unattended.""" + if not isinstance(platform, str) or not platform.strip(): + return False + return platform.strip().lower() not in _ATTENDED_PLATFORMS + + +@dataclass(frozen=True) +class LoopCapConfig: + """Per-turn hard ceilings on web_search calls / subagent spawns; count total calls (not + repeats), fire regardless of ``hard_stop_enabled``; ``0`` disables a cap.""" + + max_web_searches: int = _DEFAULT_MAX_WEB_SEARCHES_PER_TURN + max_subagents: int = _DEFAULT_MAX_SUBAGENTS_PER_TURN + + @classmethod + def from_mapping(cls, data: Mapping[str, Any] | None) -> "LoopCapConfig": + """Build config from the ``tool_loop_guardrails.loop_caps`` section.""" + if not isinstance(data, Mapping): + return cls() + return cls(**{f.name: _int_at_least(data.get(f.name), f.default, 0) for f in fields(cls)}) @dataclass(frozen=True) class ToolCallGuardrailConfig: - """Thresholds for per-turn tool-call loop detection. - - Warnings never prevent execution. Hard stops are opt-in for interactive - platforms but default on for unattended gateway/cron platforms where nobody - can interrupt a model that ignores loop warnings. - """ + """Thresholds for per-turn tool-call loop detection. Warnings never prevent execution; hard + stops are opt-in on interactive platforms, default on for unattended gateway/cron platforms.""" warnings_enabled: bool = True hard_stop_enabled: bool = False @@ -92,108 +119,32 @@ class ToolCallGuardrailConfig: no_progress_block_after: int = 5 idempotent_tools: frozenset[str] = field(default_factory=lambda: IDEMPOTENT_TOOL_NAMES) mutating_tools: frozenset[str] = field(default_factory=lambda: MUTATING_TOOL_NAMES) - loop_caps: "LoopCapConfig" = field(default_factory=lambda: LoopCapConfig()) + loop_caps: LoopCapConfig = field(default_factory=LoopCapConfig) @classmethod def from_mapping( - cls, - data: Mapping[str, Any] | None, - *, - platform: str | None = None, + cls, data: Mapping[str, Any] | None, *, platform: str | None = None, ) -> "ToolCallGuardrailConfig": - """Build config from the `tool_loop_guardrails` config.yaml section. - - Nested ``warn_after`` / ``hard_stop_after`` keys win over the flat legacy keys. - """ + """Build config from `tool_loop_guardrails`; nested ``warn_after`` / ``hard_stop_after`` win over flat legacy keys.""" if not isinstance(data, Mapping): data = {} d = cls() - hard_stop_enabled = _as_bool(data.get("hard_stop_enabled"), d.hard_stop_enabled) - non_interactive_hard_stop_enabled = _as_bool( - data.get("non_interactive_hard_stop_enabled"), d.non_interactive_hard_stop_enabled, - ) - if _is_non_interactive_platform(platform) and non_interactive_hard_stop_enabled: - hard_stop_enabled = True + flags = {name: _as_bool(data.get(name), getattr(d, name)) for name in _BOOL_FIELDS} + if flags["non_interactive_hard_stop_enabled"] and _is_non_interactive_platform(platform): + flags["hard_stop_enabled"] = True - thresholds: dict[str, int] = {} - for field_name, (section_name, key) in _THRESHOLD_SOURCES.items(): + def threshold(name: str, section_name: str, key: str) -> int: section = data.get(section_name) - if not isinstance(section, Mapping): - section = {} - thresholds[field_name] = _int_at_least( - section.get(key, data.get(field_name)), getattr(d, field_name), 1, - ) + nested = section.get(key, data.get(name)) if isinstance(section, Mapping) else data.get(name) + return _int_at_least(nested, getattr(d, name), 1) - return cls( - warnings_enabled=_as_bool(data.get("warnings_enabled"), d.warnings_enabled), - hard_stop_enabled=hard_stop_enabled, - non_interactive_hard_stop_enabled=non_interactive_hard_stop_enabled, - loop_caps=LoopCapConfig.from_mapping(data.get("loop_caps")), - **thresholds, - ) - - -# Threshold field -> (nested section, nested key). The flat legacy key is the field name itself. -_THRESHOLD_SOURCES: dict[str, tuple[str, str]] = { - "exact_failure_warn_after": ("warn_after", "exact_failure"), - "same_tool_failure_warn_after": ("warn_after", "same_tool_failure"), - "no_progress_warn_after": ("warn_after", "idempotent_no_progress"), - "exact_failure_block_after": ("hard_stop_after", "exact_failure"), - "same_tool_failure_halt_after": ("hard_stop_after", "same_tool_failure"), - "no_progress_block_after": ("hard_stop_after", "idempotent_no_progress"), -} - - -# Per-turn caps on runaway-prone tools; counters reset in reset_for_turn at the -# start of every agent loop, so the limit is per turn, not per session. Dozens -# of searches / subagent spawns in one loop is already pathological. -_DEFAULT_MAX_WEB_SEARCHES_PER_TURN = 50 -_DEFAULT_MAX_SUBAGENTS_PER_TURN = 50 - - -@dataclass(frozen=True) -class LoopCapConfig: - """Per-turn hard ceilings on web_search calls / subagent spawns. - - Unlike the loop detector (keyed on repeated identical/failing calls) these - count total calls within the turn and fire regardless of - ``hard_stop_enabled``. ``0`` disables a cap. - """ - - max_web_searches: int = _DEFAULT_MAX_WEB_SEARCHES_PER_TURN - max_subagents: int = _DEFAULT_MAX_SUBAGENTS_PER_TURN - - @classmethod - def from_mapping(cls, data: Mapping[str, Any] | None) -> "LoopCapConfig": - """Build config from the ``tool_loop_guardrails.loop_caps`` section.""" - if not isinstance(data, Mapping): - return cls() - defaults = cls() - return cls( - max_web_searches=_int_at_least(data.get("max_web_searches"), defaults.max_web_searches, 0), - max_subagents=_int_at_least(data.get("max_subagents"), defaults.max_subagents, 0), - ) - - -_INTERACTIVE_PLATFORMS = frozenset({"cli", "tui", "desktop", "acp"}) -# Not chat gateways, but bounded supervised task loops (a subagent is stopped by -# its parent; api_server has a live client). Both do real edit -> re-run work, -# so they keep the interactive warn-only default. -_SUPERVISED_TASK_PLATFORMS = frozenset({"subagent", "api_server"}) - - -def _is_non_interactive_platform(platform: str | None) -> bool: - """True for gateway/cron sessions where tool loops are unattended.""" - if not isinstance(platform, str) or not platform.strip(): - return False - key = platform.strip().lower() - return key not in _INTERACTIVE_PLATFORMS and key not in _SUPERVISED_TASK_PLATFORMS + thresholds = {name: threshold(name, *src) for name, src in _THRESHOLD_SOURCES.items()} + return cls(loop_caps=LoopCapConfig.from_mapping(data.get("loop_caps")), **flags, **thresholds) @dataclass(frozen=True) class IdenticalCallObservation: - """Outcome of observing one completed call: ``notice`` is appended after the - result, ``stub`` replaces a byte-identical duplicate result. Both may be set.""" + """``notice`` is appended after the result, ``stub`` replaces a byte-identical duplicate result.""" notice: str | None = None stub: str | None = None @@ -212,7 +163,7 @@ class ToolCallSignature: def to_metadata(self) -> dict[str, str]: """Return public metadata without raw argument values.""" - return {"tool_name": self.tool_name, "args_hash": self.args_hash} + return asdict(self) @dataclass(frozen=True) @@ -235,15 +186,9 @@ class ToolGuardrailDecision: return self.action in {"block", "halt"} def to_metadata(self) -> dict[str, Any]: - data: dict[str, Any] = { - "action": self.action, - "code": self.code, - "message": self.message, - "tool_name": self.tool_name, - "count": self.count, - } - if self.signature is not None: - data["signature"] = self.signature.to_metadata() + data = asdict(self) + if data["signature"] is None: + del data["signature"] return data @@ -259,32 +204,73 @@ def canonical_tool_args(args: Mapping[str, Any]) -> str: def classify_tool_failure(tool_name: str, result: str | None) -> tuple[bool, str]: - """Fallback classifier used only when callers don't pass ``failed``. - - Mirrors ``agent.display._detect_tool_failure`` exactly so the guardrail - never disagrees with the CLI's user-visible ``[error]`` tag. - """ + """Fallback classifier used only when callers don't pass ``failed``; mirrors + ``agent.display._detect_tool_failure`` so the guardrail never disagrees with the CLI's ``[error]`` tag.""" if result is None or file_mutation_result_landed(tool_name, result): return False, "" if tool_name == "terminal": data = safe_json_loads(result) - if isinstance(data, dict): - exit_code = data.get("exit_code") - if exit_code is not None and exit_code != 0: - return True, f" [exit {exit_code}]" - return False, "" + exit_code = data.get("exit_code") if isinstance(data, dict) else None + return (True, f" [exit {exit_code}]") if exit_code is not None and exit_code != 0 else (False, "") if tool_name == "memory": data = safe_json_loads(result) if isinstance(data, dict) and data.get("success") is False and "exceed the limit" in data.get("error", ""): return True, " [full]" - lower = result[:500].lower() - if '"error"' in lower or '"failed"' in lower or result.startswith("Error"): - return True, " [error]" + return (True, " [error]") if '"error"' in lower or '"failed"' in lower or result.startswith("Error") else (False, "") - return False, "" + +# Guardrail verdict text injected into the conversation, keyed by decision code. +# ``same_tool_failure_warning`` is built by _tool_failure_recovery_hint (tool-specific). +_DECISION_MESSAGES: dict[str, str] = { + "repeated_exact_failure_block": ( + "Blocked {tool_name}: the same tool call failed {count} times with identical arguments. " + "Stop retrying it unchanged; change strategy or explain the blocker." + ), + "idempotent_no_progress_block": ( + "Blocked {tool_name}: this read-only call returned the same result {count} times. " + "Stop repeating it unchanged; use the result already provided or try a different query." + ), + "same_tool_failure_halt": ( + "Stopped {tool_name}: it failed {count} times this turn. " + "Stop retrying the same failing tool path and choose a different approach." + ), + "repeated_exact_failure_warning": ( + "{tool_name} has failed {count} times with identical arguments. This looks like a loop; " + "inspect the error and change strategy instead of retrying it unchanged." + ), + "idempotent_no_progress_warning": ( + "{tool_name} returned the same result {count} times. Use the result already provided " + "or change the query instead of repeating it unchanged." + ), + "identical_call_streak_halt": ( + "Stopped {tool_name}: the same call with identical arguments returned the same result " + "{count} times in a row. Stop repeating it unchanged; use the result already provided or change strategy." + ), + "loop_web_search_cap": ( + "Blocked web_search: this turn has already made {cap} web searches, the per-turn limit. " + "This looks like a runaway search loop. Work with the results you already have and give the user your answer." + ), + "loop_subagent_cap": ( + "Blocked delegate_task: this turn has already spawned {count} subagents (limit {cap}). " + "This looks like a runaway delegation loop. Finish the work with the results you have and answer the user." + ), +} + +_IDENTICAL_CALL_NOTICE = ( + "[hermes note: this is the {ordinal} consecutive identical call to " + "{tool_name} with identical arguments returning the same result. " + "Do not repeat it — change arguments, use a different tool, or " + "proceed with what you have.]" +) + +# tool -> (LoopCapConfig field, controller counter attribute, decision code) +_LOOP_CAPS: dict[str, tuple[str, str, str]] = { + "web_search": ("max_web_searches", "_turn_web_search_count", "loop_web_search_cap"), + "delegate_task": ("max_subagents", "_turn_subagent_count", "loop_subagent_cap"), +} class ToolCallGuardrailController: @@ -301,16 +287,13 @@ class ToolCallGuardrailController: self._progress_since_failure: dict[ToolCallSignature, bool] = {} self._no_progress: dict[ToolCallSignature, tuple[str, int]] = {} self._halt_decision: ToolGuardrailDecision | None = None - # Identical-call streak: CONSECUTIVE identical (tool, args) calls with - # identical results. Any different call or result resets it, so re-reads - # after edits and varied polling are never flagged. + # Identical-call streak: CONSECUTIVE identical (tool, args, result) calls; any different call or + # result resets it, so re-reads after edits and varied polling are never flagged. self._identical_streak_sig: ToolCallSignature | None = None self._identical_streak_result_hash: str = "" self._identical_streak_count: int = 0 - # tool_call_id of the streak's FIRST call, so a stub can point at the full payload. self._identical_streak_first_call_id: str = "" - # tool_call_id -> spillover path, so a stub referencing a result that - # entered context only as a persisted-output preview can't dangle. + # tool_call_id -> spillover path, so a stub referencing a persisted-output preview can't dangle. self._persisted_result_paths: dict[str, str] = {} self._turn_web_search_count = 0 self._turn_subagent_count = 0 @@ -319,16 +302,17 @@ class ToolCallGuardrailController: def halt_decision(self) -> ToolGuardrailDecision | None: return self._halt_decision - def _halt( - self, action: str, code: str, message: str, - tool_name: str, count: int, signature: ToolCallSignature, + def _decide( + self, action: str, code: str, tool_name: str, count: int, signature: ToolCallSignature, + *, message: str | None = None, **fmt: Any, ) -> ToolGuardrailDecision: - """Build a block/halt decision and record it as the turn's halt decision.""" - self._halt_decision = ToolGuardrailDecision( - action=action, code=code, message=message, - tool_name=tool_name, count=count, signature=signature, - ) - return self._halt_decision + """Build a warn/block/halt decision; block/halt is also recorded as the turn's halt decision.""" + if message is None: + message = _DECISION_MESSAGES[code].format(tool_name=tool_name, count=count, **fmt) + decision = ToolGuardrailDecision(action, code, message, tool_name, count, signature) + if decision.should_halt: + self._halt_decision = decision + return decision def before_call(self, tool_name: str, args: Mapping[str, Any] | None) -> ToolGuardrailDecision: args = _coerce_args(args) @@ -337,116 +321,59 @@ class ToolCallGuardrailController: # Loop caps apply regardless of hard_stop_enabled (which only governs the detector). cap_block = self._check_loop_cap(tool_name, args, signature) - if cap_block is not None: - return cap_block - - if not self.config.hard_stop_enabled: - return allow - - exact_count = self._exact_failure_counts.get(signature, 0) - if self._progress_since_failure.get(signature): - # Something landed since this call last failed — let it run; the - # streak restarts in after_call if it fails again. - exact_count = 0 + if cap_block is not None or not self.config.hard_stop_enabled: + return cap_block or allow + # A mutation since this call last failed makes the retry a new experiment. + exact_count = 0 if self._progress_since_failure.get(signature) else self._exact_failure_counts.get(signature, 0) if exact_count >= self.config.exact_failure_block_after: - return self._halt( - "block", "repeated_exact_failure_block", - f"Blocked {tool_name}: the same tool call failed {exact_count} " - "times with identical arguments. Stop retrying it unchanged; " - "change strategy or explain the blocker.", - tool_name, exact_count, signature, - ) - - if self._is_idempotent(tool_name): - record = self._no_progress.get(signature) - if record is not None and record[1] >= self.config.no_progress_block_after: - repeat_count = record[1] - return self._halt( - "block", "idempotent_no_progress_block", - f"Blocked {tool_name}: this read-only call returned the same " - f"result {repeat_count} times. Stop repeating it unchanged; " - "use the result already provided or try a different query.", - tool_name, repeat_count, signature, - ) - + return self._decide("block", "repeated_exact_failure_block", tool_name, exact_count, signature) + record = self._no_progress.get(signature) if self._is_idempotent(tool_name) else None + if record is not None and record[1] >= self.config.no_progress_block_after: + return self._decide("block", "idempotent_no_progress_block", tool_name, record[1], signature) return allow def after_call( - self, - tool_name: str, - args: Mapping[str, Any] | None, - result: str | None, - *, - failed: bool | None = None, + self, tool_name: str, args: Mapping[str, Any] | None, result: str | None, + *, failed: bool | None = None, ) -> ToolGuardrailDecision: args = _coerce_args(args) signature = ToolCallSignature.from_call(tool_name, args) if failed is None: failed, _ = classify_tool_failure(tool_name, result) - - def warn(code: str, message: str, count: int) -> ToolGuardrailDecision: - return ToolGuardrailDecision( - action="warn", code=code, message=message, - tool_name=tool_name, count=count, signature=signature, - ) + warnings = self.config.warnings_enabled if failed: - # An identical failing call is only a REPLAY if nothing landed in - # between; a mutation since the last identical failure makes the - # retry a new experiment, so the exact-args streak restarts. + # An identical failing call is only a REPLAY if nothing landed in between; + # a mutation since the last identical failure restarts the exact-args streak. if self._progress_since_failure.pop(signature, False): self._exact_failure_counts.pop(signature, None) - exact_count = self._exact_failure_counts.get(signature, 0) + 1 - self._exact_failure_counts[signature] = exact_count + exact_count = self._exact_failure_counts[signature] = self._exact_failure_counts.get(signature, 0) + 1 + same_count = self._same_tool_failure_counts[tool_name] = self._same_tool_failure_counts.get(tool_name, 0) + 1 self._no_progress.pop(signature, None) - - same_count = self._same_tool_failure_counts.get(tool_name, 0) + 1 - self._same_tool_failure_counts[tool_name] = same_count - - # same_tool_failure counts DIFFERENT args on one tool; for - # failure-tolerant tools a run of distinct red commands is diagnosis, - # not a loop — warn, never halt (exact-args replay still applies). + # same_tool_failure counts DIFFERENT args on one tool; for failure-tolerant + # tools a run of distinct red commands is diagnosis, not a loop — warn, never halt. if ( self.config.hard_stop_enabled and tool_name not in FAILURE_TOLERANT_TOOL_NAMES and same_count >= self.config.same_tool_failure_halt_after ): - return self._halt( - "halt", "same_tool_failure_halt", - f"Stopped {tool_name}: it failed {same_count} times this turn. " - "Stop retrying the same failing tool path and choose a different approach.", - tool_name, same_count, signature, + return self._decide("halt", "same_tool_failure_halt", tool_name, same_count, signature) + if warnings and exact_count >= self.config.exact_failure_warn_after: + return self._decide("warn", "repeated_exact_failure_warning", tool_name, exact_count, signature) + if warnings and same_count >= self.config.same_tool_failure_warn_after: + return self._decide( + "warn", "same_tool_failure_warning", tool_name, same_count, signature, + message=_tool_failure_recovery_hint(tool_name, same_count), ) - - if self.config.warnings_enabled and exact_count >= self.config.exact_failure_warn_after: - return warn( - "repeated_exact_failure_warning", - f"{tool_name} has failed {exact_count} times with identical arguments. " - "This looks like a loop; inspect the error and change strategy " - "instead of retrying it unchanged.", - exact_count, - ) - - if self.config.warnings_enabled and same_count >= self.config.same_tool_failure_warn_after: - return warn( - "same_tool_failure_warning", - _tool_failure_recovery_hint(tool_name, same_count), - same_count, - ) - return ToolGuardrailDecision(tool_name=tool_name, count=exact_count, signature=signature) self._exact_failure_counts.pop(signature, None) self._same_tool_failure_counts.pop(tool_name, None) - - # A successful mutation is progress for every failing signature still - # counted this turn (next identical retry runs against changed state). - # Pure loops never mutate between attempts, so the replay detector keeps its teeth. + # A successful mutation is progress for every failing signature still counted + # this turn. Pure loops never mutate between attempts, so the replay detector keeps its teeth. if tool_name in PROGRESS_RESET_TOOL_NAMES or file_mutation_result_landed(tool_name, result): - for sig in list(self._exact_failure_counts): - self._progress_since_failure[sig] = True + self._progress_since_failure.update(dict.fromkeys(self._exact_failure_counts, True)) self._same_tool_failure_counts.clear() - if not self._is_idempotent(tool_name): self._no_progress.pop(signature, None) return ToolGuardrailDecision(tool_name=tool_name, signature=signature) @@ -455,51 +382,29 @@ class ToolCallGuardrailController: previous = self._no_progress.get(signature) repeat_count = previous[1] + 1 if previous is not None and previous[0] == result_hash else 1 self._no_progress[signature] = (result_hash, repeat_count) - - if self.config.warnings_enabled and repeat_count >= self.config.no_progress_warn_after: - return warn( - "idempotent_no_progress_warning", - f"{tool_name} returned the same result {repeat_count} times. " - "Use the result already provided or change the query instead of " - "repeating it unchanged.", - repeat_count, - ) - + if warnings and repeat_count >= self.config.no_progress_warn_after: + return self._decide("warn", "idempotent_no_progress_warning", tool_name, repeat_count, signature) return ToolGuardrailDecision(tool_name=tool_name, count=repeat_count, signature=signature) def _is_idempotent(self, tool_name: str) -> bool: return tool_name not in self.config.mutating_tools and tool_name in self.config.idempotent_tools def observe_call( - self, - tool_name: str, - args: Mapping[str, Any] | None, - result: str | None, - *, - tool_call_id: str = "", - failed: bool = False, - ) -> "IdenticalCallObservation": + self, tool_name: str, args: Mapping[str, Any] | None, result: str | None, + *, tool_call_id: str = "", failed: bool = False, + ) -> IdenticalCallObservation: """Track consecutive identical calls; return notice + dedupe stub info. - ``notice`` fires from the ``STALL_GUARD_IDENTICAL_CALL_THRESHOLD``-th - consecutive identical (tool, args, result) call; purely observational, - pollers exempt. ``stub`` replaces the CURRENT result from the 2nd - byte-identical repeat — the tool still executed, only the context - representation is deduplicated, so polling semantics survive (a changed - result flows through whole and resets the streak). Pollers are NOT - exempt from stubbing: an unchanged poll is exactly when the stub saves - the most. Short results, failed results and non-string results are never - stubbed. Callers substitute at result construction time, which is cache-safe. + ``notice`` fires from the threshold-th consecutive identical (tool, args, result) call + (observational, pollers exempt). ``stub`` replaces the CURRENT result from the 2nd byte-identical + repeat — the tool still executed, only the context representation is deduplicated, so polling + semantics survive; pollers are NOT exempt here since an unchanged poll is where it saves most. """ is_plain_str = isinstance(result, str) signature = ToolCallSignature.from_call(tool_name, _coerce_args(args)) result_hash = _result_hash(result) if is_plain_str else "" - if ( - is_plain_str - and self._identical_streak_sig == signature - and self._identical_streak_result_hash == result_hash - ): + if is_plain_str and (signature, result_hash) == (self._identical_streak_sig, self._identical_streak_result_hash): self._identical_streak_count += 1 else: # New streak; non-string (multimodal) results never form one. @@ -507,39 +412,20 @@ class ToolCallGuardrailController: self._identical_streak_result_hash = result_hash self._identical_streak_count = 1 if is_plain_str else 0 self._identical_streak_first_call_id = tool_call_id or "" - count = self._identical_streak_count notice = None if not is_stall_guard_repeatable(tool_name) and count >= STALL_GUARD_IDENTICAL_CALL_THRESHOLD: - ordinal = f"{count}{'th' if 11 <= count % 100 <= 13 else {1: 'st', 2: 'nd', 3: 'rd'}.get(count % 10, 'th')}" - notice = ( - f"[hermes note: this is the {ordinal} consecutive identical call to " - f"{tool_name} with identical arguments returning the same result. " - "Do not repeat it — change arguments, use a different tool, or " - "proceed with what you have.]" - ) - # The no-progress BLOCK in before_call only covers idempotent_tools; this - # streak is tool-agnostic, so with hard stops on, halt at the same threshold - # (a model replaying a successful `terminal` call otherwise runs to the budget). - if ( - self.config.hard_stop_enabled - and count >= self.config.no_progress_block_after - and self._halt_decision is None - ): - self._halt( - "halt", "identical_call_streak_halt", - f"Stopped {tool_name}: the same call with identical arguments " - f"returned the same result {count} times in a row. Stop " - "repeating it unchanged; use the result already provided or " - "change strategy.", - tool_name, count, signature, - ) + notice = _IDENTICAL_CALL_NOTICE.format(ordinal=_ordinal(count), tool_name=tool_name) + # The no-progress BLOCK in before_call only covers idempotent_tools; this streak + # is tool-agnostic, so with hard stops on, halt at the same threshold (a model + # replaying a successful `terminal` call otherwise runs to the budget). + if self.config.hard_stop_enabled and count >= self.config.no_progress_block_after and self._halt_decision is None: + self._decide("halt", "identical_call_streak_halt", tool_name, count, signature) stub = None if is_plain_str and count >= 2 and not failed and len(result) >= IDENTICAL_RESULT_STUB_MIN_CHARS: stub = self._build_result_reference_stub(tool_name, args) - return IdenticalCallObservation(notice=notice, stub=stub) def record_persisted_result(self, tool_call_id: str, file_path: str) -> None: @@ -549,10 +435,7 @@ class ToolCallGuardrailController: def _build_result_reference_stub(self, tool_name: str, args: Mapping[str, Any] | None) -> str: """Reference stub for a byte-identical duplicate result (tool + args preview).""" - try: - args_preview = canonical_tool_args(_coerce_args(args)) - except TypeError: - args_preview = "{}" + args_preview = canonical_tool_args(_coerce_args(args)) if len(args_preview) > _RESULT_STUB_ARGS_PREVIEW_CHARS: args_preview = args_preview[:_RESULT_STUB_ARGS_PREVIEW_CHARS] + "…" first_id = self._identical_streak_first_call_id @@ -564,62 +447,29 @@ class ToolCallGuardrailController: ) spill_path = self._persisted_result_paths.get(first_id) if first_id else None if spill_path: - stub += ( - f"\n[The referenced result was persisted to: {spill_path} — " - "page through it with read_file if you need the full content.]" - ) + stub += f"\n[The referenced result was persisted to: {spill_path} — page through it with read_file if you need the full content.]" return stub def _check_loop_cap( - self, - tool_name: str, - args: Mapping[str, Any], - signature: ToolCallSignature, + self, tool_name: str, args: Mapping[str, Any], signature: ToolCallSignature, ) -> ToolGuardrailDecision | None: - """Block once a per-turn cap is reached, else advance the counter and return None. - - A cap of 0 disables that limit. Blocking happens BEFORE the call when - the count is already at the cap, so the (cap+1)-th call is refused. - """ - caps = self.config.loop_caps - if tool_name == "web_search": - cap, count, increment = caps.max_web_searches, self._turn_web_search_count, 1 - code, message = "loop_web_search_cap", ( - f"Blocked web_search: this turn has already made {cap} " - "web searches, the per-turn limit. This looks like a " - "runaway search loop. Work with the results you already " - "have and give the user your answer." - ) - elif tool_name == "delegate_task": - cap, count = caps.max_subagents, self._turn_subagent_count - # Control actions (list/steer/stop) spawn nothing and must keep working after the cap is hit. - increment = _subagent_spawn_count(args) if cap else 0 - if increment == 0: - return None - code, message = "loop_subagent_cap", ( - f"Blocked delegate_task: this turn has already spawned " - f"{count} subagents (limit {cap}). " - "This looks like a runaway delegation loop. Finish the " - "work with the results you have and answer the user." - ) - else: + """Block once a per-turn cap is reached (BEFORE the call, so the (cap+1)-th is refused), else advance + the counter and return None. delegate_task control actions spawn nothing and keep working after the cap.""" + spec = _LOOP_CAPS.get(tool_name) + if spec is None: return None - - if cap and count >= cap: - return self._halt("block", code, message, tool_name, count, signature) - if tool_name == "web_search": - self._turn_web_search_count += increment - else: - self._turn_subagent_count += increment + cap_field, count_attr, code = spec + cap, count = getattr(self.config.loop_caps, cap_field), getattr(self, count_attr) + increment = 1 if tool_name == "web_search" else (_subagent_spawn_count(args) if cap else 0) + if increment and cap and count >= cap: + return self._decide("block", code, tool_name, count, signature, cap=cap) + setattr(self, count_attr, count + increment) return None def toolguard_synthetic_result(decision: ToolGuardrailDecision) -> str: """Build a synthetic role=tool content string for a blocked tool call.""" - return json.dumps( - {"error": decision.message, "guardrail": decision.to_metadata()}, - ensure_ascii=False, - ) + return json.dumps({"error": decision.message, "guardrail": decision.to_metadata()}, ensure_ascii=False) def append_toolguard_guidance(result: str, decision: ToolGuardrailDecision) -> str: @@ -650,42 +500,32 @@ def _tool_failure_recovery_hint(tool_name: str, count: int) -> str: ) +def _ordinal(count: int) -> str: + return f"{count}{'th' if 11 <= count % 100 <= 13 else {1: 'st', 2: 'nd', 3: 'rd'}.get(count % 10, 'th')}" + + def _coerce_args(args: Mapping[str, Any] | None) -> Mapping[str, Any]: return args if isinstance(args, Mapping) else {} def _result_hash(result: str | None) -> str: parsed = safe_json_loads(result or "") - if parsed is not None: - try: - canonical = _canonical_json(parsed) - except TypeError: - canonical = str(parsed) - else: - canonical = result or "" - return _sha256(canonical) + return _sha256(_canonical_json(parsed) if parsed is not None else (result or "")) + + +_BOOL_WORDS = {w: True for w in ("1", "true", "yes", "on", "enabled")} | {w: False for w in ("0", "false", "no", "off", "disabled")} def _as_bool(value: Any, default: bool) -> bool: - if value is None: - return default - if isinstance(value, bool): - return value - if isinstance(value, (int, float)): + if isinstance(value, (bool, int, float)): return bool(value) if isinstance(value, str): - lowered = value.strip().lower() - if lowered in {"1", "true", "yes", "on", "enabled"}: - return True - if lowered in {"0", "false", "no", "off", "disabled"}: - return False + return _BOOL_WORDS.get(value.strip().lower(), default) return default def _int_at_least(value: Any, default: int, minimum: int) -> int: - """Int parser: junk/None/below-minimum fall back to default (caps use minimum 0 so 0 = disabled).""" - if value is None: - return default + """junk/None/below-minimum fall back to default (caps use minimum 0 so 0 = disabled).""" try: parsed = int(value) except (TypeError, ValueError): @@ -694,8 +534,7 @@ def _int_at_least(value: Any, default: int, minimum: int) -> int: def _subagent_spawn_count(args: Mapping[str, Any]) -> int: - """Subagents one delegate_task call spawns: batch size when ``tasks`` is a - non-empty list, else 1; control actions (list/steer/stop) spawn 0.""" + """Subagents one delegate_task call spawns: ``len(tasks)`` for a non-empty batch, else 1; control actions 0.""" if str(args.get("action") or "").strip().lower() in ("list", "steer", "stop"): return 0 tasks = args.get("tasks")