diff --git a/README.md b/README.md index b80ada9..ea8e609 100644 --- a/README.md +++ b/README.md @@ -21,6 +21,8 @@ Version `0.1.x` contains the converged foundation layer: - loop configuration, lifecycle contexts, observer protocol, intervention merging, and observer dispatch helpers. - streamed tool-call recovery checks for missing required arguments. +- LLM binding, response normalization, streaming assembly/watchdogs, retry + classification, runaway recovery, and physical-call orchestration. The initial extraction is based on the already-merged integration branches: @@ -30,10 +32,10 @@ The initial extraction is based on the already-merged integration branches: Those revisions are provenance, not runtime dependencies. AgentCore tests and builds without either product checkout. -The LLM call runtime and agent loop remain in the products until tool -execution, model profiles, retry classification, execution-context storage, -and runtime hooks have product-neutral boundaries. Moving those files before -that boundary exists would only hide product coupling inside this package. +The agent loop remains in the products until tool execution, model profiles, +tool-call parsing, and execution-context storage have converged boundaries. +The LLM runtime accepts the three remaining product decisions through explicit +hooks: wall-deadline lookup, provider-chain state, and sticky-session policy. ## Repository boundary @@ -109,11 +111,10 @@ edit both products' core copies, that is evidence it belongs here. 1. **Foundation** (this version): messages, token estimation, compaction, context budget, trimming. -2. **Runtime contracts** (in progress): LLM protocols and loop types are now - shared; errors, execution-context storage, retry classification, and - explicit product hooks remain. -3. **LLM runtime:** binding, calls, streaming, response normalization, runaway - recovery, and the public `llm_client` facade. +2. **Runtime contracts** (complete): LLM protocols, loop types, errors, retry + classification, and explicit product hooks are shared. +3. **LLM runtime** (complete): binding, calls, streaming, response + normalization, runaway recovery, and the public `llm_client` facade. 4. **Agent loop:** model/tool parsing, tool execution, and `agent_loop`. 5. Remove product compatibility facades after downstream imports have moved to `agent_core`. diff --git a/agent_core/__init__.py b/agent_core/__init__.py index f64a392..7c1287f 100644 --- a/agent_core/__init__.py +++ b/agent_core/__init__.py @@ -1,5 +1,6 @@ """Product-neutral building blocks for Apodex agent runtimes.""" +from agent_core.errors import AgentCoreError from agent_core.llm import LLMClient, LLMResponse, StreamDelta from agent_core.messages import ( Message, @@ -11,6 +12,7 @@ ) __all__ = [ + "AgentCoreError", "LLMClient", "LLMResponse", "Message", diff --git a/agent_core/errors.py b/agent_core/errors.py new file mode 100644 index 0000000..89512b2 --- /dev/null +++ b/agent_core/errors.py @@ -0,0 +1,164 @@ +"""Exception hierarchy for AgentCore.""" + +from __future__ import annotations + +from typing import Any + + +class AgentCoreError(Exception): + """Base exception for all AgentCore errors.""" + + +# ── Kernel errors ─────────────────────────────────────────────────────────── + + +class KernelError(AgentCoreError): + """Errors originating from the OS kernel layer.""" + + +class TaskNotFoundError(KernelError): + def __init__(self, task_id: str) -> None: + super().__init__(f"Task not found: {task_id}") + self.task_id = task_id + + +class InvalidStateTransition(KernelError): + def __init__(self, task_id: str, current: str, target: str) -> None: + super().__init__(f"Invalid transition for {task_id}: {current} → {target}") + + +class ServiceNotRegistered(KernelError): + def __init__(self, service_type: type) -> None: + super().__init__(f"Service not registered: {service_type.__name__}") + + +class PermissionDenied(KernelError): + def __init__(self, role: str, tool: str) -> None: + super().__init__(f"Role '{role}' has no permission for tool '{tool}'") + +# LLM request errors + +class LLMError(AgentCoreError): + """Errors from the LLM/provider layer.""" + + +class LLMReasoningRunaway(LLMError): + """A live stream spent its semantic budget on reasoning-only output. + + Unlike :class:`LLMStreamStalled`, the provider is healthy and actively + emitting chunks. The failure is semantic: no non-whitespace visible text + or tool-call delta appeared before the configured time/token guard fired. + + ``partial_response`` is intentionally carried separately from provider + usage. Early stream cancellation often happens before the terminal usage + chunk arrives, so its estimated reasoning tokens must never be presented + as authoritative billing data. + """ + + def __init__( + self, + *, + elapsed_s: float, + estimated_tokens: int, + trigger: str, + partial_response: Any, + ) -> None: + self.elapsed_s = float(elapsed_s) + self.estimated_tokens = int(estimated_tokens) + self.trigger = trigger + self.partial_response = partial_response + super().__init__( + "reasoning-only stream exceeded " + f"{trigger} guard (elapsed={self.elapsed_s:.1f}s, " + f"estimated_tokens={self.estimated_tokens})", + ) + + +class LLMStreamStalled(LLMError, TimeoutError): + """A streaming LLM call went silent mid-flight. + + Subclasses ``asyncio.TimeoutError`` so every existing transient- + timeout handler (retry/backoff in ``call_llm``, chain wrappers, + classification) treats it identically without changes; carried + fields make the distinct failure mode visible in logs and traces. + """ + + def __init__( + self, stall_s: float, chunks_seen: int, elapsed_s: float, + ) -> None: + self.stall_s = stall_s + self.chunks_seen = chunks_seen + self.elapsed_s = elapsed_s + super().__init__( + f"stream stalled: no chunks for {stall_s:.0f}s " + f"(chunks_seen={chunks_seen}, elapsed={elapsed_s:.0f}s)", + ) + + +class LLMDeadlineExceeded(LLMError, TimeoutError): + """An LLM attempt was stopped by an enclosing runtime deadline. + + ``reason`` is deliberately carried on the underlying exception as well as + on :class:`LLMCallExhausted`. Some callers unwrap ``last_exc`` before + handing it to a provider-chain policy; a dedicated type prevents that + policy from mistaking an exhausted run budget for an ordinary transient + provider timeout. + """ + + def __init__(self, reason: str, detail: str) -> None: + self.reason = reason + super().__init__(f"{reason}: {detail}") + + +class LLMCallExhausted(LLMError, RuntimeError): + """Raised by ``call_llm`` when retries are exhausted or the error is + structurally unrecoverable (4xx without proxy-wrap, or a chain-aware + fallback signal like ``model_not_found``). + + Wraps the last exception encountered so the caller (typically the + product's agent loop) can surface it to a provider-chain wrapper for + L1→L2→L3 rotation. Carries ``last_exc`` separately because + ``raise from`` is too opaque for chain-aware classification — a chain + wrapper calls ``classify_error(last_exc)`` directly. + + ``last_exc`` must always agree with ``reason``: it is the exception that + *caused this raise*, not merely the most recent failure seen. A deadline + refusal therefore carries :class:`LLMDeadlineExceeded` even when earlier + attempts failed for unrelated reasons. The wrapper's ``reason`` remains + authoritative, while the underlying exception preserves the same reason + if a caller unwraps it before classification. + + ``prior_exc`` is where that earlier, superseded failure goes: diagnostic + context for logs and post-mortems, deliberately outside the field + classification reads. + """ + + def __init__( + self, + last_exc: BaseException, + reason: str, + *, + prior_exc: BaseException | None = None, + ) -> None: + self.last_exc = last_exc + self.reason = reason + self.prior_exc = prior_exc + detail = f"call_llm {reason}: {last_exc!r}" + if prior_exc is not None and prior_exc is not last_exc: + detail += f" (after {prior_exc!r})" + super().__init__(detail) + + +__all__ = [ + "AgentCoreError", + "InvalidStateTransition", + "KernelError", + "LLMCallExhausted", + "LLMDeadlineExceeded", + "LLMError", + "LLMReasoningRunaway", + "LLMStreamStalled", + "PermissionDenied", + "ServiceNotRegistered", + "TaskNotFoundError", +] diff --git a/agent_core/llm.py b/agent_core/llm.py index 750df2d..acd7630 100644 --- a/agent_core/llm.py +++ b/agent_core/llm.py @@ -37,9 +37,9 @@ class StreamDelta: usage: dict[str, int] = field(default_factory=dict[str, int]) finish_reason: str = "" model: str = "" - # Vendor label of the leg serving this stream, stamped by - # ``LLMFallbackChain.stream`` (constant once the chain commits to an - # entry — failover only fires before the first yield). The stream + # Vendor label of the leg serving this stream, stamped by a product's + # provider-chain wrapper (constant once the chain commits to an entry — + # failover only fires before the first yield). The stream # assembler folds it into ``LLMResponse.response_metadata`` so per-call # billing attribution works for streamed calls too — without this the # streaming path had no channel for the provider and every billing diff --git a/agent_core/messages.py b/agent_core/messages.py index 8682bdf..52225db 100644 --- a/agent_core/messages.py +++ b/agent_core/messages.py @@ -20,8 +20,8 @@ is preserved as the ``content`` value when the underlying client returns it; callers that need flat text use :func:`text_of`. -This module replaces ``langchain_core.messages`` (BaseMessage / SystemMessage / -HumanMessage / AIMessage / ToolMessage). It is intentionally dependency-free. +This module replaces the former framework-specific message classes with a +small, dependency-free wire contract. """ from __future__ import annotations @@ -34,8 +34,8 @@ class ToolCall(TypedDict): """OpenAI-style tool_call payload — ``function.arguments`` is JSON-encoded. - Wire key order is fixed ``{type, id, function}`` to match the LangChain - serializer the served checkpoints were aligned against; do not reorder. + Wire key order is fixed ``{type, id, function}`` to match the serializer + byte shape the served checkpoints were aligned against; do not reorder. """ id: str @@ -141,7 +141,7 @@ def for_wire(messages: list[Message]) -> list[Message]: # Key insertion order: ``content`` first, then ``role``. Some served # checkpoints are sensitive to this byte shape (wire byte-equality with -# LangChain's ``_convert_message_to_dict`` — see migration gotcha #2). Do +# the legacy message serializer — see migration gotcha #2). Do # not reorder these dict literals. diff --git a/agent_core/runtime/env.py b/agent_core/runtime/env.py new file mode 100644 index 0000000..e66f41e --- /dev/null +++ b/agent_core/runtime/env.py @@ -0,0 +1,35 @@ +"""Shared env-variable prefix cascade. + +Several modules independently re-walked the same +``AGENT_CORE_ / MIROHARNESS_ / FRONTIER_AGENT_`` prefix order looking for a +configured value — one copy per module, easy to drift if a prefix is ever +added or reordered in only one of them. This is the single implementation +they converge on: import :func:`first_configured` and let it own the order +rather than passing a locally-spelled tuple back in. +""" + +from __future__ import annotations + +import os + +# The portable ``AGENT_CORE_`` spelling wins when multiple aliases are +# configured, followed by the MiroHarness and FrontierAgent compatibility +# names. Order matters: +# callers rely on the first configured prefix winning. +ENV_PREFIXES = ("AGENT_CORE_", "MIROHARNESS_", "FRONTIER_AGENT_") + + +def first_configured(suffix: str, prefixes: tuple[str, ...] = ENV_PREFIXES) -> tuple[str, str] | None: + """Return the ``(name, value)`` of the first non-empty ``{prefix}{suffix}`` env var. + + ``None`` when none of the prefixed names are set (or all are blank). + """ + for prefix in prefixes: + name = f"{prefix}{suffix}" + raw = os.environ.get(name, "").strip() + if raw: + return name, raw + return None + + +__all__ = ["ENV_PREFIXES", "first_configured"] diff --git a/agent_core/runtime/llm_request_overrides.py b/agent_core/runtime/llm_request_overrides.py new file mode 100644 index 0000000..2abddcf --- /dev/null +++ b/agent_core/runtime/llm_request_overrides.py @@ -0,0 +1,68 @@ +"""Task-local overrides for one physical LLM request. + +The runtime occasionally needs to change generation behaviour for one retry +without mutating a cached/shared client. ``ContextVar`` keeps that override +isolated across concurrent tasks and automatically restores the client's +normal profile on exit. + +Provider adapters opt in to the semantic override they understand. Today the +OpenAI-compatible adapter maps it onto SGLang/Qwen +``chat_template_kwargs`` and an explicitly configured ``reasoning_effort``. +Unsupported adapters simply keep their normal request shape; the runtime's +retry prompt and output cap remain the portable fallback. +""" + +from __future__ import annotations + +from collections.abc import Generator +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass + + +@dataclass(frozen=True) +class ThinkingRetryOverride: + """Semantic thinking controls for a single retry attempt.""" + + mode: str = "reduced" + thinking_budget: int | None = None + reasoning_effort: str | None = None + + @property + def disabled(self) -> bool: + return self.mode == "disabled" + + +_THINKING_RETRY_OVERRIDE: ContextVar[ThinkingRetryOverride | None] = ContextVar( + "agent_core_thinking_retry_override", + default=None, +) + + +def current_thinking_retry_override() -> ThinkingRetryOverride | None: + """Return the override active for the current async task, if any.""" + + return _THINKING_RETRY_OVERRIDE.get() + + +@contextmanager +def thinking_retry_override( + override: ThinkingRetryOverride | None, +) -> Generator[None, None, None]: + """Apply ``override`` only inside this context and async task.""" + + if override is None: + yield + return + token = _THINKING_RETRY_OVERRIDE.set(override) + try: + yield + finally: + _THINKING_RETRY_OVERRIDE.reset(token) + + +__all__ = [ + "ThinkingRetryOverride", + "current_thinking_retry_override", + "thinking_retry_override", +] diff --git a/agent_core/runtime/loop/__init__.py b/agent_core/runtime/loop/__init__.py index f6c4def..f7ddd1e 100644 --- a/agent_core/runtime/loop/__init__.py +++ b/agent_core/runtime/loop/__init__.py @@ -4,6 +4,25 @@ DefaultCompactionPolicy, DefaultMessageCompactor, ) +from agent_core.runtime.loop.llm_client import ( + RUNAWAY_STATE_KEY, + TRUNCATION_CONTINUATION_GUIDANCE, + LLMCallExhausted, + LLMDeadlineExceeded, + LLMReasoningRunaway, + LLMStreamStalled, + ThinkTagSplitter, + bind_max_tokens, + bind_session_id, + bind_temperature, + bind_tools, + call_llm, + extract_final_content, + extract_leaked_reasoning, + extract_model_name, + extract_usage, + is_truncated_with_text, +) from agent_core.runtime.loop.message_trimmer import ( MessageTrimmer, NullTrimmer, @@ -11,9 +30,26 @@ ) __all__ = [ + "RUNAWAY_STATE_KEY", + "TRUNCATION_CONTINUATION_GUIDANCE", "DefaultCompactionPolicy", "DefaultMessageCompactor", + "LLMCallExhausted", + "LLMDeadlineExceeded", + "LLMReasoningRunaway", + "LLMStreamStalled", "MessageTrimmer", "NullTrimmer", "TaskBoundaryTrimmer", + "ThinkTagSplitter", + "bind_max_tokens", + "bind_session_id", + "bind_temperature", + "bind_tools", + "call_llm", + "extract_final_content", + "extract_leaked_reasoning", + "extract_model_name", + "extract_usage", + "is_truncated_with_text", ] diff --git a/agent_core/runtime/loop/_bind.py b/agent_core/runtime/loop/_bind.py new file mode 100644 index 0000000..17d6aa9 --- /dev/null +++ b/agent_core/runtime/loop/_bind.py @@ -0,0 +1,201 @@ +# pyright: reportPrivateUsage=false, reportUnusedFunction=false +from __future__ import annotations + +import logging +from collections.abc import Callable +from dataclasses import dataclass, replace +from typing import Any + +from agent_core.llm import LLMResponse +from agent_core.messages import Message + +logger = logging.getLogger(__name__) + +@dataclass +class _BoundLLM: + """Lightweight binding wrapper around a native :class:`LLMClient`. + + The agent loop builds its per-turn LLM by chaining + ``bind_tools(bind_session_id(llm, task_id), tools)``. :class:`LLMClient` + is a plain protocol with no bind-style hook of its own, so the bound + knobs are carried here and threaded into :meth:`LLMClient.chat` / + :meth:`LLMClient.stream` per call. Each ``bind_*`` returns a fresh + wrapper via :func:`dataclasses.replace` — never a mutation of the + shared long-lived client. + """ + + client: Any + tools: list[dict[str, Any]] | None = None + temperature: float | None = None + extra_headers: dict[str, str] | None = None + max_tokens: int | None = None + + @property + def model(self) -> str: + return getattr(self.client, "model", "") or "" + + def _call_kwargs( + self, + timeout: float | None, + *, + tools: list[dict[str, Any]] | None = None, + temperature: float | None = None, + max_tokens: int | None = None, + extra_headers: dict[str, str] | None = None, + ) -> dict[str, Any]: + """Merge the bound fields with any kwarg-style overrides. + + ``_BoundLLM`` is bind-style (knobs live on the dataclass), but it + gets nested *underneath* kwarg-style callers: product proxies and + provider-chain wrappers unconditionally forward + ``tools=/temperature=/max_tokens=/extra_headers=`` to their inner + client. When that inner client is a ``_BoundLLM`` — the shape a + product's sub-agent spawning helper builds — a bind-only signature + raised ``TypeError: stream() got an unexpected keyword argument + 'tools'`` and killed every sub-agent on turn 1. Accepting + merging + these kwargs makes ``_BoundLLM`` a tolerant drop-in. + + Precedence: an explicit non-``None`` kwarg overrides the bound + field (the caller asked for it this call); otherwise the bound + field is used. ``extra_headers`` is merged (bound base, kwarg wins + per key) so a forwarded header never drops the session-affinity + header bound earlier. + """ + kw: dict[str, Any] = {} + eff_tools = tools if tools is not None else self.tools + if eff_tools: + kw["tools"] = eff_tools + eff_temperature = ( + temperature if temperature is not None else self.temperature + ) + if eff_temperature is not None: + kw["temperature"] = eff_temperature + merged_headers = {**(self.extra_headers or {}), **(extra_headers or {})} + if merged_headers: + kw["extra_headers"] = merged_headers + eff_max_tokens = ( + max_tokens if max_tokens is not None else self.max_tokens + ) + if eff_max_tokens is not None: + kw["max_tokens"] = eff_max_tokens + if timeout is not None: + kw["timeout"] = timeout + return kw + + async def chat( + self, + messages: list[Message], + *, + tools: list[dict[str, Any]] | None = None, + temperature: float | None = None, + max_tokens: int | None = None, + extra_headers: dict[str, str] | None = None, + timeout: float | None = None, + ) -> LLMResponse: + return await self.client.chat( + messages, + **self._call_kwargs( + timeout, + tools=tools, + temperature=temperature, + max_tokens=max_tokens, + extra_headers=extra_headers, + ), + ) + + def stream( + self, + messages: list[Message], + *, + tools: list[dict[str, Any]] | None = None, + temperature: float | None = None, + max_tokens: int | None = None, + extra_headers: dict[str, str] | None = None, + timeout: float | None = None, + ) -> Any: + return self.client.stream( + messages, + **self._call_kwargs( + timeout, + tools=tools, + temperature=temperature, + max_tokens=max_tokens, + extra_headers=extra_headers, + ), + ) + + +def _ensure_bound(llm: Any) -> _BoundLLM: + """Wrap a raw ``LLMClient`` in a ``_BoundLLM``; pass an existing one through.""" + return llm if isinstance(llm, _BoundLLM) else _BoundLLM(client=llm) + + +def bind_tools(llm: Any, tools: list[Any]) -> Any: + """Bind native tool objects (or pre-built + OpenAI function-schema dicts) so the model can emit multiple tool_calls + per turn. ``parallel_tool_calls=True`` is applied by the client adapter. + No-op on an empty tool list. + """ + if not tools: + return llm + schemas = [ + t.to_openai_schema() if hasattr(t, "to_openai_schema") else t + for t in tools + ] + return replace(_ensure_bound(llm), tools=schemas) + + +def bind_session_id( + llm: Any, + task_id: str, + *, + sticky_session_enabled: Callable[[], bool] | None = None, +) -> Any: + """Attach ``x-upstream-session-id: `` to every LLM request. + + Pinning it at client-construction time is the obvious approach, but + the LLM is per-profile-cached (one client shared across tasks), so the + header is bound per-call as ``extra_headers`` instead, which the + OpenAI-compatible adapter forwards to the SDK's ``extra_headers`` kwarg. + + Why it matters: EAS-backed gateways use this header for **session + affinity** — the same session-id consistently + routes to the same backend worker, preserving KV-cache across a task's + turns. + + No-op when ``task_id`` is empty (standalone debug) or when + the product session-affinity kill switch is falsey. This header is set + ONLY here, so a client built without going through this bind is + deliberately unpinned. Sticky routing was once + suspected of amplifying a high-concurrency stampede and disabled for it; + that attribution was retracted when the real cause turned out to be silent + mid-stream stalls, now handled by the stall watchdog. + """ + if not task_id or ( + sticky_session_enabled is not None and not sticky_session_enabled() + ): + return llm + bound = _ensure_bound(llm) + headers = dict(bound.extra_headers or {}) + headers["x-upstream-session-id"] = task_id + return replace(bound, extra_headers=headers) + + +def bind_temperature(llm: Any, temperature: float) -> Any: + """Bind ``temperature`` for a single invocation. + + Used by retry observers to escalate sampling on a retry turn without + mutating the long-lived LLM (a fresh ``_BoundLLM`` is returned). + """ + return replace(_ensure_bound(llm), temperature=temperature) + + +def bind_max_tokens(llm: Any, max_tokens: int) -> Any: + """Bind ``max_tokens`` for a single invocation (fresh ``_BoundLLM``). + + Public counterpart to :func:`bind_temperature` for callers outside + this package (e.g. sub-agent spawning) that need to override a child + LLM's ``max_tokens`` without reaching into ``_ensure_bound`` / + ``_BoundLLM`` directly. + """ + return replace(_ensure_bound(llm), max_tokens=max_tokens) diff --git a/agent_core/runtime/loop/_call.py b/agent_core/runtime/loop/_call.py new file mode 100644 index 0000000..702b219 --- /dev/null +++ b/agent_core/runtime/loop/_call.py @@ -0,0 +1,1192 @@ +# pyright: reportMissingTypeArgument=false, reportPrivateUsage=false, reportUnknownArgumentType=false, reportUnknownMemberType=false, reportUnknownParameterType=false, reportUnknownVariableType=false +from __future__ import annotations + +import asyncio +import logging +import random +import time +from collections.abc import Awaitable, Callable +from typing import Any + +from agent_core.errors import ( + LLMCallExhausted, + LLMDeadlineExceeded, + LLMReasoningRunaway, + LLMStreamStalled, +) +from agent_core.llm import LLMResponse +from agent_core.loop_types import ( + ATTEMPT_ACCEPTED, + ATTEMPT_ACCEPTED_DEGRADED, + ATTEMPT_DISCARDED, + ATTEMPT_FAILED, +) +from agent_core.messages import Message, user_msg +from agent_core.runtime.llm_request_overrides import ( + ThinkingRetryOverride, + thinking_retry_override, +) +from agent_core.runtime.retriable import ( + get_status_code as _get_status_code, +) +from agent_core.runtime.retriable import ( + is_context_length_error, + is_transient_network, +) + +from ._bind import _ensure_bound +from ._response import _visible_response_text, extract_usage +from ._runaway import ( + _RUNAWAY_BACKOFF_S, + _RUNAWAY_EXPAND_ENABLED, + _RUNAWAY_MAX_RETRIES, + _bind_expanded_max_tokens, + _bind_reduced_max_tokens, + _env_int, + _is_runaway_response, + _phase_reasoning_guard, + _runaway_retry_policy, +) +from ._streaming import ( + _accepts_tool_call_arg_chunks, + _stream_llm_response, + _stream_stall_max_before_advance, +) +from .tool_call_recovery import stream_tool_calls_missing_required_arguments + +logger = logging.getLogger(__name__) + + +def _get_retry_after(exc: Exception) -> float | None: + """Extract a Retry-After header value (seconds) from a 429 exception.""" + for attr in ("response", "headers"): + obj = getattr(exc, attr, None) + if obj is None: + continue + headers = getattr(obj, "headers", obj) if attr == "response" else obj + if not hasattr(headers, "get"): + continue + val = headers.get("retry-after") or headers.get("Retry-After") + if val: + try: + return float(val) + except (ValueError, TypeError): + pass + return None + + +# ±25% jitter on the exponential schedules. Without it, parallel runs +# that failed together retry together: 5 attempts timing out inside one +# 3-minute window had their retries re-collide 20 minutes later — a +# synchronised stampede against an already struggling gateway. ``retry_wait_fixed`` and literal ``Retry-After`` +# values are intentionally NOT jittered (explicit caller contracts). +_BACKOFF_JITTER = 0.25 + + +def _jittered(base: float) -> float: + return base * random.uniform(1 - _BACKOFF_JITTER, 1 + _BACKOFF_JITTER) + + +def _default_backoff(attempt: int) -> float: + """Exponential schedule for timeouts / transient errors: + 2/4/8/16/32/60s base, ±25% jitter.""" + return _jittered(min(2 * (2 ** attempt), 60)) + + +def _default_rate_limit_backoff(attempt: int) -> float: + """Exponential schedule for 429 fallback (no Retry-After): + 30/60/120/240/300s base, ±25% jitter.""" + return _jittered(min(30 * (2 ** attempt), 300)) + + +# ── Global LLM concurrency gate (opt-in, default OFF) ───────────────── +# Client-side admission control: caps in-flight LLM attempts process- +# wide so a heavy fan-out (8 runs × 2-4 sub-agents) queues at the +# client — observable, cancellable — instead of inside the gateway, +# where the thundering herd was observed to correlate with silent-stream +# black holes. 0 / unset disables; deploys size it to the endpoint's +# decode slots (e.g. 12-16 for a large mixture-of-experts gateway). +# +# The semaphore is loop-affine: rebuilt whenever the running loop +# changes (worker = one loop for life; tests get one per case). +# Suffix resolved through the shared ``AGENT_CORE_`` / compatibility-prefix +# cascade by ``_env_int``. +_LLM_GATE_ENV = "LLM_MAX_CONCURRENT" +_llm_gate_state: tuple[Any, asyncio.Semaphore] | None = None + + +def _llm_gate() -> asyncio.Semaphore | None: + limit = _env_int(_LLM_GATE_ENV, 0) + if limit <= 0: + return None + global _llm_gate_state + loop = asyncio.get_running_loop() + if _llm_gate_state is None or _llm_gate_state[0] is not loop: + _llm_gate_state = (loop, asyncio.Semaphore(limit)) + return _llm_gate_state[1] + + +# ── Wall-deadline budget closure ────────────────────────────────────── +# A product observer stamps the loop's absolute monotonic soft deadline +# into execution-scope metadata, and its sub-agent fan-in already clamps +# its wait to it. ``call_llm`` was the remaining leak: each +# attempt got the full configured timeout regardless of remaining wall +# — an attempt with a fresh 1200 s budget could launch with 90 s of wall +# left. Each attempt now clamps its timeout to the +# remaining budget, refuses to start under the floor, and abandons +# backoff sleeps that would cross the deadline. +# +# Floor: below this many remaining seconds an LLM attempt can't return +# anything useful — fail fast with reason="wall_deadline" so the loop +# stops cleanly and the post-loop salvage (force_final_answer) gets the +# reserve instead. The remaining budget is not read from here: the product +# owns that lookup and injects it as the ``wall_deadline_remaining`` +# callback, so this package never touches context-local storage. +_WALL_DEADLINE_FLOOR_S = 20.0 + +# Floor for the non-streaming replay that recovers a tool call whose streamed +# arguments came back empty. The replay is opportunistic: below this many +# seconds of remaining attempt budget it would almost certainly time out, and +# the streamed response we already hold is a better outcome than burning the +# rest of the turn on a doomed second request. +_STREAM_RECOVERY_MIN_TIMEOUT_S = 10.0 + + +def _stream_recovery_budget_too_small( + remaining_s: float, + attempt_budget_s: float, +) -> bool: + """Whether the empty-arguments replay should be skipped for lack of time. + + Skipped only when the remaining budget is under the absolute floor *and* + under half of what the attempt started with. The second term matters: a + deployment that configures a short ``timeout`` would otherwise never get + the recovery at all, since the post-stream remainder is always slightly + below a floor set at or above ``timeout``. Being wrong here is cheap — + a failed replay falls back to the streamed response. + """ + return ( + remaining_s < _STREAM_RECOVERY_MIN_TIMEOUT_S + and remaining_s < attempt_budget_s / 2 + ) + +async def call_llm( + llm: Any, + messages: list[Message], + timeout: int, + max_retries: int, + turn: int, + # Two accepted shapes, picked apart at runtime by + # ``_accepts_tool_call_arg_chunks``: with or without the keyword-only + # ``tool_call_args_chunks``. Hence Callable[...] rather than a fixed + # parameter list. + on_delta: Callable[..., Awaitable[None]] | None = None, + retry_wait_fixed: int | None = None, + runaway_state: dict[str, Any] | None = None, + first_chunk_s: float | None = None, + on_attempt: Callable[[dict[str, Any]], Awaitable[None]] | None = None, + reasoning_only_timeout_s: float | None = None, + reasoning_only_max_tokens: int | None = None, + logical_call_timeout_s: float | None = None, + max_completion_tokens_hint: int | None = None, + context_token_limit_hint: int | None = None, + wall_deadline_remaining: Callable[[], float | None] | None = None, + chain_fallback_active: Callable[[], bool] | None = None, +) -> LLMResponse | None: + """Call ``llm.chat`` (``llm.stream`` when ``on_delta`` is set) with + exponential backoff on transient errors. + + Default retry schedule (when ``retry_wait_fixed`` is ``None``): + exponential 2/4/8/16/32s capped at 60 for timeouts and generic errors; + 429 honours ``Retry-After`` (clamped at 300s) when present, else + exponential 30/60/120/240s capped at 300. + + Failure modes: + + - **Non-transient HTTP (400 schema, 401/403 auth, 404 route)**: + raises :class:`LLMCallExhausted` immediately (reason=``non_transient``) + — retrying won't help and burning fallback keys/chain legs is just + slower failure. + - **Chain-aware fallback signal** (``model_not_found`` / overload / + credit-exhausted / safety-filter — anything ``is_retriable_with_fallback`` + flags): raises immediately (reason=``chain_advance``). Same-key retry + can't change the outcome; the next chain leg may. + - **Retries exhausted on transient errors** (timeout / stream stall + / 5xx / 429 / proxy-wrap): raises (reason=``exhausted``) carrying + the last transient exception. + + Streaming calls additionally run under the inter-chunk stall + watchdog (:class:`LLMStreamStalled`, + ``LLM_STREAM_STALL_S`` under any supported prefix, default 180 s): a stream that + goes silent mid-flight is aborted and retried like a timeout + instead of pinning the attempt for the full ``timeout``. After + ``LLM_STREAM_STALL_MAX`` under any supported prefix (default + 2) — and only when an outer provider chain is active — the call + raises ``LLMCallExhausted(reason="chain_advance")`` instead of + retrying the same black-holed endpoint. What the caller does with it + is turn-dependent and decided by the caller's loop: a turn-1 + exhaustion is re-raised to the outer chain wrapper to rotate to the next + leg; past turn 1 the loop stops with ``llm_error`` and hands off to + salvage. Either way the retry budget is no longer burned on a dead + gateway. + + When ``retry_wait_fixed`` is set (int seconds), all transient retries + use that fixed wait regardless of error class — used by noisy-endpoint + workflows — a self-hosted endpoint at high concurrency, where stream + timeouts stampede and the server-side worker recovery cycle is longer + than the exponential schedule's early attempts. A fixed + 60-90s wait gives the worker pool one full recovery cycle between + attempts. + + ``runaway_state`` (optional, mutable dict the caller keeps alive + across turns) enables the reasoning-runaway recovery documented at + :data:`RUNAWAY_STATE_KEY`: a successful-but-empty capped completion + is resampled same-key at a reduced ``max_tokens`` with transient guidance + instead of being returned for the loop to nudge blind. The runaway + retries share the ``max_retries`` attempt budget; an unrecovered + runaway is returned as-is (never raised) so existing nudge handling + remains the floor. + + Streaming reasoning can be stopped before it reaches the provider cap by + setting ``reasoning_only_timeout_s`` and/or + ``reasoning_only_max_tokens``. The semantic guard remains armed only while + the stream has reasoning but no non-whitespace visible text or tool-call + delta. It shares the same reduced-cap recovery path as a completed + capped-empty response. + + ``logical_call_timeout_s`` is a separate opt-in deadline spanning gate + wait, all physical attempts, and retry backoff. It never extends an earlier + run-level wall deadline. Products with execution-scope storage inject that + deadline through ``wall_deadline_remaining``; AgentCore never imports a + product's context-local storage. + + Callers wrap with try/except :class:`LLMCallExhausted` and decide + whether to surface (chain advance) or degrade (partial content). + """ + from agent_core.runtime.retriable import ( + is_empty_completion, + is_overloaded_error, + is_retriable_with_fallback, + ) + + def _transient_backoff(attempt: int) -> float: + """Backoff for transient errors: caller-fixed wait, else the + jittered exponential default. (429 has its own schedule.)""" + return ( + retry_wait_fixed + if retry_wait_fixed is not None + else _default_backoff(attempt) + ) + + logical_timeout_s = max(float(logical_call_timeout_s or 0), 0.0) + logical_deadline = ( + time.monotonic() + logical_timeout_s if logical_timeout_s else None + ) + + def _nearest_deadline() -> tuple[float | None, str]: + candidates: list[tuple[float, str]] = [] + wall_remaining = ( + wall_deadline_remaining() + if wall_deadline_remaining is not None + else None + ) + if wall_remaining is not None: + candidates.append((wall_remaining, "wall_deadline")) + if logical_deadline is not None: + candidates.append(( + logical_deadline - time.monotonic(), + "logical_call_deadline", + )) + if not candidates: + return None, "" + return min(candidates, key=lambda item: item[0]) + + def _chain_fallback_active() -> bool: + return bool(chain_fallback_active and chain_fallback_active()) + + def _effective_timeout_or_deadline_exhausted( + *, attempt: int, reason: str, + ) -> tuple[float, str]: + effective_timeout = timeout + timeout_deadline_reason = "" + deadline_remaining, deadline_reason = _nearest_deadline() + if deadline_remaining is not None and deadline_remaining < timeout: + if deadline_remaining < _WALL_DEADLINE_FLOOR_S: + deadline_exc = LLMDeadlineExceeded( + deadline_reason, + f"{deadline_remaining:.0f}s left " + f"< {_WALL_DEADLINE_FLOOR_S:.0f}s attempt floor", + ) + logger.warning( + "LLM call refused: %.0fs to %s (< %.0fs " + "floor; turn=%d, attempt=%d/%d, reason=%s) — surfacing " + "deadline for clean loop exit", + deadline_remaining, deadline_reason, + _WALL_DEADLINE_FLOOR_S, + turn, attempt + 1, max_retries, reason, + ) + # ``last_exc`` must match ``reason``: a chain wrapper + # classifies it directly, so handing it a stale 429 from an + # earlier attempt under reason="wall_deadline" would have it + # back off and retry straight past the deadline we are + # refusing to cross. The superseded failure stays available + # as ``prior_exc`` for the post-mortem. + raise LLMCallExhausted( + deadline_exc, deadline_reason, prior_exc=last_exc, + ) from deadline_exc + effective_timeout = deadline_remaining + timeout_deadline_reason = deadline_reason + logger.info( + "LLM call timeout clamped to %s remaining: %ds → %.0fs " + "(turn=%d, attempt=%d/%d, reason=%s)", + deadline_reason, + timeout, effective_timeout, turn, attempt + 1, max_retries, + reason, + ) + return effective_timeout, timeout_deadline_reason + + def _deadline_allows_retry_after(delay_s: float) -> bool: + deadline_remaining, _deadline_reason = _nearest_deadline() + return ( + deadline_remaining is None + or delay_s + _WALL_DEADLINE_FLOOR_S < deadline_remaining + ) + + async def _emit_attempt(event: dict[str, Any]) -> None: + if on_attempt is None: + return + try: + await on_attempt(event) + except Exception: + # Attempt observability is passive. A broken consumer must not + # turn a valid provider response into an LLM failure. + logger.warning("LLM attempt callback failed", exc_info=True) + + def _response_attempt_fields(response: LLMResponse) -> dict[str, Any]: + return { + "usage": extract_usage(response), + "finish_reason": response.finish_reason or "", + "visible_chars": len(_visible_response_text(response)), + "reasoning_chars": len(response.reasoning_content or ""), + "tool_calls_count": len(response.tool_calls or []), + } + + last_exc: BaseException | None = None + # Phase index and retry spend diverge when expansion is skipped. + runaway_retries = 0 + runaway_attempts = 0 + last_runaway_reason = "" + stream_stall_count = 0 + if runaway_state is not None: + # Per-call diagnostics. A protocol stream observer surfaces these so + # a consumer can distinguish "content was clipped" from earlier + # reasoning-only attempts that were streamed and then resampled. + runaway_state["last_call_runaway_responses"] = 0 + runaway_state["last_call_runaway_reasoning_chars"] = 0 + runaway_state["last_call_recovered"] = False + runaway_state["last_call_reason"] = "" + # ``llm_active`` may be re-bound with a reduced max_tokens after a + # repeated reasoning runaway — error retries then reuse the bound + # variant too, which is fine (the cap only applies post-runaway). + llm_active = _ensure_bound(llm) + messages_active = messages + retry_thinking: ThinkingRetryOverride | None = None + guard_timeout_s = reasoning_only_timeout_s + guard_max_tokens = reasoning_only_max_tokens + # One PHYSICAL request per index, not one retry per index. They only + # diverge when a stream is discarded and replayed inside a single retry + # (empty tool arguments, below): the caller derives ``attempt_id`` + # from this number, so reusing it would emit two ``finished`` events + # under one id and double-count that attempt's usage downstream. + physical_attempt_index = 0 + for attempt in range(max_retries): + attempt_thinking = retry_thinking + # Budget closure: clamp this attempt to the remaining wall (when + # a deadline is stamped) so a retry chain can never outlive the + # loop's own budget. Under the floor, refuse to start at all. + effective_timeout, attempt_deadline_reason = ( + _effective_timeout_or_deadline_exhausted( + attempt=attempt, reason="pre_gate", + ) + ) + physical_attempt_index += 1 + attempt_index = physical_attempt_index + attempt_started = time.monotonic() + attempt_first_delta: float | None = None + active_cap = ( + getattr(llm_active, "max_tokens", None) + or max_completion_tokens_hint + ) + await _emit_attempt({ + "phase": "started", + "attempt_index": attempt_index, + "max_tokens": active_cap, + "thinking_mode": ( + attempt_thinking.mode if attempt_thinking else "profile_default" + ), + "thinking_budget": ( + attempt_thinking.thinking_budget if attempt_thinking else None + ), + }) + + attempt_delta = on_delta + if on_delta is not None: + downstream_accepts_tool_chunks = _accepts_tool_call_arg_chunks( + on_delta, + ) + + async def _attempt_delta( + delta: str, + accumulated: str, + delta_index: int, + thinking_delta: str = "", + *, + tool_call_args_chunks: list[dict] | None = None, + ) -> None: + nonlocal attempt_first_delta + if ( + attempt_first_delta is None + and (delta or thinking_delta or tool_call_args_chunks) + ): + attempt_first_delta = time.monotonic() + if downstream_accepts_tool_chunks: + await on_delta( + delta, + accumulated, + delta_index, + thinking_delta, + tool_call_args_chunks=tool_call_args_chunks or [], + ) + else: + await on_delta( + delta, accumulated, delta_index, thinking_delta, + ) + + attempt_delta = _attempt_delta + + async def _chat_active(read_timeout: float) -> LLMResponse: + with thinking_retry_override(attempt_thinking): + return await llm_active.chat( + messages_active, timeout=read_timeout, + ) + + async def _stream_active(read_timeout: float) -> LLMResponse: + if attempt_delta is None: + raise RuntimeError("streaming retry requires a delta callback") + with thinking_retry_override(attempt_thinking): + return await _stream_llm_response( + llm_active, messages_active, read_timeout, attempt_delta, + first_chunk_s=first_chunk_s, + reasoning_only_timeout_s=guard_timeout_s, + reasoning_only_max_tokens=guard_max_tokens, + ) + + async def _finish_attempt( + *, + outcome: str, + reason: str, + recovery_action: str, + response: LLMResponse | None = None, + error: BaseException | None = None, + ended_at: float | None = None, + ) -> None: + # ``ended_at`` back-dates the close for an attempt that finished + # earlier than this call — the discarded stream below is reported + # only once its replacement is known, and must not be charged for + # the replay's wall time. + now = time.monotonic() if ended_at is None else ended_at + event: dict[str, Any] = { + "phase": "finished", + "attempt_index": attempt_index, + "outcome": outcome, + "reason": reason, + "recovery_action": recovery_action, + "duration_ms": int((now - attempt_started) * 1000), + "ttft_ms": ( + int((attempt_first_delta - attempt_started) * 1000) + if attempt_first_delta is not None + else None + ), + "max_tokens": active_cap, + "thinking_mode": ( + attempt_thinking.mode + if attempt_thinking else "profile_default" + ), + "thinking_budget": ( + attempt_thinking.thinking_budget + if attempt_thinking else None + ), + "error_type": type(error).__name__ if error is not None else "", + } + if response is not None: + event.update(_response_attempt_fields(response)) + await _emit_attempt(event) + + retry_reason = "transient_error" + retry_error: BaseException | None = None + try: + # Opt-in global admission gate — excess attempts wait HERE + # (client-side, visible) rather than queueing blind inside + # the gateway. Backoff sleeps run outside the gate so a + # waiting retry never holds a slot. The wait itself is + # bounded by the wall deadline reserve, and the provider + # timeout is recomputed after the slot is acquired so queue + # time cannot leak past the loop budget. + gate = _llm_gate() + gate_acquired = False + try: + if gate is not None: + deadline_remaining, _deadline_reason = _nearest_deadline() + gate_wait_timeout = None + if deadline_remaining is not None: + gate_wait_timeout = ( + deadline_remaining - _WALL_DEADLINE_FLOOR_S + ) + if gate_wait_timeout <= 0: + _effective_timeout_or_deadline_exhausted( + attempt=attempt, reason="gate_wait", + ) + if gate_wait_timeout is None: + await gate.acquire() + else: + try: + await asyncio.wait_for( + gate.acquire(), timeout=gate_wait_timeout, + ) + except TimeoutError as exc: + deadline_exc = LLMDeadlineExceeded( + _deadline_reason, + "concurrency-gate wait consumed the remaining budget", + ) + raise LLMCallExhausted( + deadline_exc, + _deadline_reason, + prior_exc=last_exc, + ) from exc + gate_acquired = True + effective_timeout, attempt_deadline_reason = ( + _effective_timeout_or_deadline_exhausted( + attempt=attempt, reason="post_gate", + ) + ) + if attempt_delta is None: + response = await asyncio.wait_for( + _chat_active(effective_timeout), + timeout=effective_timeout, + ) + else: + response = await _stream_active(effective_timeout) + empty_arg_tools = stream_tool_calls_missing_required_arguments( + response, llm_active, + ) + if empty_arg_tools: + # The stream has only been observed/assembled here: no + # assistant history or tool execution has happened yet, + # so replacing it with one non-streaming replay cannot + # duplicate a side effect. + streamed_response = response + stream_ended_at = time.monotonic() + logger.warning( + "Streamed tool call(s) %s had blank arguments despite " + "required schema fields (turn=%d, attempt=%d/%d); " + "replaying the same request non-streaming", + empty_arg_tools, turn, attempt + 1, max_retries, + ) + recovered: LLMResponse | None = None + recovery_error: BaseException | None = None + try: + # Clamp to whatever is left of THIS attempt's own + # budget as well as the wall/logical deadline: the + # replay is a second physical request inside one + # attempt, so without the first term a turn could + # quietly cost 2x ``timeout`` whenever no deadline + # is stamped (direct loop use, SDK, tests). + recovery_timeout = min( + _effective_timeout_or_deadline_exhausted( + attempt=attempt, + reason="stream_empty_tool_arguments", + )[0], + max( + effective_timeout + - (stream_ended_at - attempt_started), + 0.0, + ), + ) + if _stream_recovery_budget_too_small( + recovery_timeout, float(effective_timeout), + ): + raise TimeoutError( + f"only {recovery_timeout:.0f}s of the " + f"{effective_timeout:.0f}s attempt budget " + f"left for the replay", + ) + recovered = await asyncio.wait_for( + _chat_active(recovery_timeout), + timeout=recovery_timeout, + ) + except Exception as exc: + # Recovery is opportunistic. Anything it raises — + # an exhausted deadline, a timeout, a provider 5xx + # — must not be worse than not having tried: keep + # the streamed response and let the loop apply its + # normal tool-validation feedback. + recovery_error = exc + + if recovered is None: + logger.warning( + "Non-streaming replay failed (%s: %s); keeping " + "the streamed response with blank tool " + "arguments (turn=%d, attempt=%d/%d)", + type(recovery_error).__name__, recovery_error, + turn, attempt + 1, max_retries, + ) + response.response_metadata = { + **(response.response_metadata or {}), + "stream_empty_args_fallback": False, + "stream_empty_args_tools": empty_arg_tools, + "stream_empty_args_recovery_error": type( + recovery_error, + ).__name__, + } + else: + # Close the discarded stream as its own attempt. + # Every other discard path in this function does + # the same, and downstream depends on it twice: + # a protocol stream observer drains its sentence / + # ```` filters on a non-delivered outcome so + # the abandoned bytes cannot bleed into the replay, + # and attempt-finished is the billing record for a + # request whose payload never reaches the loop. + # That is also why the replay keeps its OWN usage + # untouched: merging the two would bill the stream + # a second time. + await _finish_attempt( + outcome=ATTEMPT_DISCARDED, + reason="stream_empty_tool_arguments", + recovery_action="replay_non_streaming", + response=streamed_response, + ended_at=stream_ended_at, + ) + physical_attempt_index += 1 + attempt_index = physical_attempt_index + attempt_started = stream_ended_at + attempt_first_delta = None + await _emit_attempt({ + "phase": "started", + "attempt_index": attempt_index, + "max_tokens": active_cap, + "thinking_mode": ( + attempt_thinking.mode + if attempt_thinking else "profile_default" + ), + "thinking_budget": ( + attempt_thinking.thinking_budget + if attempt_thinking else None + ), + }) + response = recovered + response.response_metadata = { + **(response.response_metadata or {}), + "stream_empty_args_fallback": True, + "stream_empty_args_tools": empty_arg_tools, + "stream_finish_reason": ( + streamed_response.finish_reason or "" + ), + } + finally: + if gate is not None and gate_acquired: + gate.release() + if _is_runaway_response(response): + last_runaway_reason = "reasoning_runaway" + if runaway_state is not None: + runaway_state["last_call_runaway_responses"] += 1 + runaway_state["last_call_runaway_reasoning_chars"] += len( + getattr(response, "reasoning_content", "") or "", + ) + # Diagnostic only. ``consecutive_turns`` used to gate whether + # this retry reduced the cap; the reduction is unconditional + # now, so the counter survives purely so the log line (and a + # post-mortem reading it) can tell a first-time runaway from a + # model that has been running away turn after turn. + prior_turn_runaway = bool( + runaway_state + and runaway_state.get("consecutive_turns", 0), + ) + if ( + runaway_attempts < _RUNAWAY_MAX_RETRIES + and attempt < max_retries - 1 + and _deadline_allows_retry_after(_RUNAWAY_BACKOFF_S) + ): + runaway_attempts += 1 + next_retry = runaway_retries + 1 + response_usage = extract_usage(response) or {} + expansion_cap = active_cap or response_usage.get( + "completion_tokens", + ) + if next_retry == 1 and _RUNAWAY_EXPAND_ENABLED: + expanded_llm = _bind_expanded_max_tokens( + llm_active, + active_cap=expansion_cap, + messages=messages, + context_token_limit_hint=context_token_limit_hint, + ) + if expanded_llm is not None: + runaway_retries = 1 + llm_active = expanded_llm + else: + runaway_retries = 2 + llm_active = _bind_reduced_max_tokens( + llm_active, response, + ) + else: + runaway_retries = max(next_retry, 2) + llm_active = _bind_reduced_max_tokens( + llm_active, response, + ) + next_cap = getattr(llm_active, "max_tokens", None) + guard_timeout_s, guard_max_tokens = _phase_reasoning_guard( + runaway_retries, + previous_cap=expansion_cap, + next_cap=next_cap, + timeout_s=reasoning_only_timeout_s, + max_tokens=reasoning_only_max_tokens, + ) + retry_thinking, recovery_guidance, recovery_action = ( + _runaway_retry_policy(runaway_retries, next_cap) + ) + messages_active = [*messages, user_msg(recovery_guidance)] + await _finish_attempt( + outcome=ATTEMPT_DISCARDED, + reason="reasoning_runaway", + recovery_action=recovery_action, + response=response, + ) + logger.warning( + "LLM reasoning runaway: capped completion with no " + "visible content (turn=%d, attempt=%d/%d, " + "runaway_retry=%d/%d, runaway_phase=%d, next_cap=%s, " + "next_thinking_mode=%s, next_thinking_budget=%s, " + "prior_turn_runaway=%s); resampling", + turn, attempt + 1, max_retries, + runaway_attempts, _RUNAWAY_MAX_RETRIES, + runaway_retries, next_cap, + retry_thinking.mode, retry_thinking.thinking_budget, + prior_turn_runaway, + ) + await asyncio.sleep(_RUNAWAY_BACKOFF_S) + continue + if runaway_state is not None: + runaway_state["consecutive_turns"] = ( + runaway_state.get("consecutive_turns", 0) + 1 + ) + logger.error( + "LLM reasoning runaway persisted after %d resamples " + "(turn=%d); returning empty response for loop-level " + "nudge handling", + runaway_retries, turn, + ) + if runaway_state is not None: + runaway_state["last_call_reason"] = "reasoning_runaway" + # DELIVERED, not failed: this response is returned below, so + # the loop appends it to history, bills it, and salvages the + # turn with its no-tool nudge. Marking it ``failed`` would + # make consumers drop bytes the loop actually used and would + # flip the enclosing trace call to ``status="failed"`` even + # though it produced a turn. ``reason`` carries the health. + await _finish_attempt( + outcome=ATTEMPT_ACCEPTED_DEGRADED, + reason="reasoning_runaway", + recovery_action="return_to_loop", + response=response, + ) + return response + if runaway_state is not None: + runaway_state["consecutive_turns"] = 0 + runaway_state["last_call_recovered"] = bool(runaway_retries) + if runaway_retries: + runaway_state["last_call_reason"] = ( + last_runaway_reason or "reasoning_runaway" + ) + await _finish_attempt( + outcome=ATTEMPT_ACCEPTED, + reason="", + recovery_action="accepted", + response=response, + ) + return response + except LLMCallExhausted as exc: + await _finish_attempt( + outcome=ATTEMPT_FAILED, + reason=exc.reason, + recovery_action="raise", + error=exc.last_exc, + ) + raise + except LLMReasoningRunaway as exc: + last_runaway_reason = "reasoning_runaway_early" + partial_response = exc.partial_response + if runaway_state is not None: + runaway_state["last_call_runaway_responses"] += 1 + runaway_state["last_call_runaway_reasoning_chars"] += len( + getattr(partial_response, "reasoning_content", "") or "", + ) + prior_turn_runaway = bool( + runaway_state + and runaway_state.get("consecutive_turns", 0) + ) + if ( + runaway_attempts < _RUNAWAY_MAX_RETRIES + and attempt < max_retries - 1 + and _deadline_allows_retry_after(_RUNAWAY_BACKOFF_S) + ): + runaway_attempts += 1 + next_retry = runaway_retries + 1 + if next_retry == 1 and _RUNAWAY_EXPAND_ENABLED: + expanded_llm = _bind_expanded_max_tokens( + llm_active, + active_cap=active_cap, + messages=messages, + context_token_limit_hint=context_token_limit_hint, + ) + if expanded_llm is not None: + runaway_retries = 1 + llm_active = expanded_llm + else: + runaway_retries = 2 + llm_active = _bind_reduced_max_tokens( + llm_active, active_cap=active_cap, + ) + else: + runaway_retries = max(next_retry, 2) + llm_active = _bind_reduced_max_tokens( + llm_active, active_cap=active_cap, + ) + next_cap = getattr(llm_active, "max_tokens", None) + guard_timeout_s, guard_max_tokens = _phase_reasoning_guard( + runaway_retries, + previous_cap=active_cap, + next_cap=next_cap, + timeout_s=reasoning_only_timeout_s, + max_tokens=reasoning_only_max_tokens, + ) + retry_thinking, recovery_guidance, recovery_action = ( + _runaway_retry_policy(runaway_retries, next_cap) + ) + messages_active = [*messages, user_msg(recovery_guidance)] + await _finish_attempt( + outcome=ATTEMPT_DISCARDED, + reason="reasoning_runaway_early", + recovery_action=recovery_action, + response=partial_response, + error=exc, + ) + logger.warning( + "LLM reasoning runaway stopped early: no visible/tool " + "progress (turn=%d, attempt=%d/%d, trigger=%s, " + "elapsed=%.1fs, estimated_tokens=%d, runaway_retry=%d/%d, " + "runaway_phase=%d, next_cap=%s, next_thinking_mode=%s, " + "next_thinking_budget=%s, prior_turn_runaway=%s); resampling", + turn, attempt + 1, max_retries, exc.trigger, + exc.elapsed_s, exc.estimated_tokens, + runaway_attempts, _RUNAWAY_MAX_RETRIES, + runaway_retries, next_cap, + retry_thinking.mode, retry_thinking.thinking_budget, + prior_turn_runaway, + ) + await asyncio.sleep(_RUNAWAY_BACKOFF_S) + continue + if runaway_state is not None: + runaway_state["consecutive_turns"] = ( + runaway_state.get("consecutive_turns", 0) + 1 + ) + runaway_state["last_call_reason"] = ( + "reasoning_runaway_early" + ) + logger.error( + "LLM reasoning runaway stopped early but no resample slot " + "remains (turn=%d, trigger=%s, elapsed=%.1fs, " + "estimated_tokens=%d); returning partial response for " + "loop-level nudge handling", + turn, exc.trigger, exc.elapsed_s, exc.estimated_tokens, + ) + await _finish_attempt( + outcome=ATTEMPT_ACCEPTED_DEGRADED, + reason="reasoning_runaway_early", + recovery_action="return_to_loop", + response=partial_response, + error=exc, + ) + return partial_response + except LLMStreamStalled as exc: + # Mid-stream silence (gateway queue black-hole / dropped + # connection): the stream was already closed by the + # watchdog; retry under the normal transient budget. Logged + # distinctly from the total-timeout so traces show HOW the + # attempt died, not just that it took too long. + last_exc = exc + retry_error = exc + retry_reason = "stream_stalled" + stream_stall_count += 1 + logger.warning( + "LLM stream stalled: no chunks for %.0fs (turn=%d, " + "attempt=%d/%d, chunks_seen=%d, elapsed=%.0fs, " + "stall_count=%d); aborting stream and retrying", + exc.stall_s, turn, attempt + 1, max_retries, + exc.chunks_seen, exc.elapsed_s, stream_stall_count, + ) + # Repeated mid-stream black-holes mean THIS endpoint is dead + # for this call — a same-key retry just re-queues into the same + # saturated gateway: one observed run burnt 56 stalls × + # ~180-330s = 207 minutes on a single endpoint this way. + # When an outer chain is + # active, stop burning the retry budget and surface + # ``chain_advance``. The caller's loop then either rotates the + # leg (turn 1) or stops with ``llm_error`` for salvage (turn > 1) — + # both stop the wall-burn. With no chain configured there's + # nothing to advance to, so fall through to the normal transient + # retry (wall-clamped). + stall_max = _stream_stall_max_before_advance() + if ( + stall_max > 0 + and stream_stall_count >= stall_max + and _chain_fallback_active() + ): + logger.error( + "LLM stream stalled %d× (turn=%d); surfacing for " + "chain advance instead of retrying the same " + "black-holed endpoint", + stream_stall_count, turn, + ) + await _finish_attempt( + outcome=ATTEMPT_FAILED, + reason="stream_stalled", + recovery_action="chain_advance", + error=exc, + ) + raise LLMCallExhausted(exc, "chain_advance") from exc + backoff = _transient_backoff(attempt) + except TimeoutError as exc: + prior_exc = last_exc + last_exc = exc + retry_error = exc + retry_reason = "timeout" + deadline_remaining, deadline_reason = _nearest_deadline() + if ( + attempt_deadline_reason + and deadline_remaining is not None + and deadline_remaining <= 0 + and deadline_reason == attempt_deadline_reason + ): + deadline_exc = LLMDeadlineExceeded( + deadline_reason, + "deadline-clamped provider attempt exhausted its budget", + ) + await _finish_attempt( + outcome=ATTEMPT_FAILED, + reason=deadline_reason, + recovery_action="raise", + error=deadline_exc, + ) + raise LLMCallExhausted( + deadline_exc, deadline_reason, prior_exc=prior_exc, + ) from exc + logger.warning( + "LLM call timed out (turn=%d, attempt=%d/%d)", + turn, attempt + 1, max_retries, + ) + backoff = _transient_backoff(attempt) + except Exception as exc: + last_exc = exc + retry_error = exc + retry_reason = "transient_error" + if ( + attempt_thinking is not None + and attempt_thinking.mode == "expanded" + and is_context_length_error(exc) + and runaway_attempts < _RUNAWAY_MAX_RETRIES + and attempt < max_retries - 1 + ): + runaway_attempts += 1 + runaway_retries = 2 + llm_active = _bind_reduced_max_tokens( + llm_active, active_cap=active_cap, + ) + guard_timeout_s = reasoning_only_timeout_s + guard_max_tokens = reasoning_only_max_tokens + next_cap = getattr(llm_active, "max_tokens", None) + retry_thinking, recovery_guidance, _ = _runaway_retry_policy( + runaway_retries, next_cap, + ) + messages_active = [*messages, user_msg(recovery_guidance)] + await _finish_attempt( + outcome=ATTEMPT_DISCARDED, + reason="context_length", + recovery_action="retry_reduced_after_context_overflow", + error=exc, + ) + logger.warning( + "Expanded-thinking retry exceeded context; degrading " + "without another expanded attempt (turn=%d, attempt=%d/%d, " + "next_cap=%s, next_thinking_budget=%s)", + turn, attempt + 1, max_retries, next_cap, + retry_thinking.thinking_budget, + ) + continue + # Chain-aware shortcut: model_not_found / overload / credit + # / safety_filter is deterministic on this (provider, input). + # Skip the rest of the retry budget and surface so an outer + # chain wrapper can advance the leg right now. + if is_retriable_with_fallback(exc): + # Overload (503) and empty completions frequently clear on a + # same-key resample (temperature>0 re-rolls the sampler). When + # NO outer chain is active to advance a leg + # (``chain_fallback_active()`` is False), + # short-circuiting these would trade a recoverable blip for a + # turn-1 trial loss, so fall through to the transient-backoff + # retry below (503 / no-status both land in the generic retry + # path). Surface immediately when a chain IS active, or for the + # genuinely deterministic failures (auth / model_unavailable / + # credit / safety) where retrying the same key cannot help. + resample_may_recover = ( + is_overloaded_error(exc) or is_empty_completion(exc) + ) + if _chain_fallback_active() or not resample_may_recover: + logger.error( + "LLM call hit chain-fallback signal (turn=%d, " + "attempt=%d/%d, %s); surfacing for layer advance: %s", + turn, attempt + 1, max_retries, + type(exc).__name__, exc, + ) + await _finish_attempt( + outcome=ATTEMPT_FAILED, + reason="chain_advance", + recovery_action="chain_advance", + error=exc, + ) + raise LLMCallExhausted(exc, "chain_advance") from exc + logger.warning( + "Retriable-fallback signal but no active chain (turn=%d, " + "attempt=%d/%d, %s); same-key retry within budget: %s", + turn, attempt + 1, max_retries, + type(exc).__name__, exc, + ) + status = _get_status_code(exc) + if status and status in (400, 401, 403, 404): + # Proxy-wrap escape hatch: OpenAI-compatible gateways + # (new-api, etc.) sometimes package an upstream 5xx / + # timeout as a 400 envelope (body carries + # ``code=bad_response_status_code`` / + # ``type=new_api_error``). The literal status is 400 but + # the semantics are transient — sleeping and retrying + # the same key fixes it. Vanilla 400 (bad JSON, schema + # mismatch) still falls through to the non-transient + # branch. + if status == 400 and is_transient_network(exc): + backoff = _transient_backoff(attempt) + logger.warning( + "Proxy-wrapped transient %d (turn=%d, attempt=%d/%d): %s", + status, turn, attempt + 1, max_retries, exc, + ) + else: + logger.error( + "Non-transient LLM error %d (turn=%d): %s", + status, turn, exc, + ) + await _finish_attempt( + outcome=ATTEMPT_FAILED, + reason="non_transient", + recovery_action="raise", + error=exc, + ) + raise LLMCallExhausted(exc, "non_transient") from exc + elif status == 429: + retry_reason = "rate_limited" + # 429 honours Retry-After (clamped at 300s ceiling so a + # buggy upstream returning ``Retry-After: 86400`` cannot + # silently stall the loop for a day); falls back to the + # exponential rate-limit schedule when no header is set. + # Workflows that opt into ``retry_wait_fixed`` still use + # their fixed schedule — they're tuning for a known + # worker-recovery cycle, not a true rate limit. + if retry_wait_fixed is not None: + backoff = retry_wait_fixed + else: + retry_after = _get_retry_after(exc) + backoff = ( + min(retry_after, 300) + if retry_after + else _default_rate_limit_backoff(attempt) + ) + logger.warning( + "LLM rate-limited 429 (turn=%d, attempt=%d/%d, wait=%ds): %s", + turn, attempt + 1, max_retries, int(backoff), exc, + ) + else: + backoff = _transient_backoff(attempt) + logger.warning( + "LLM call error (turn=%d, attempt=%d/%d): %s", + turn, attempt + 1, max_retries, exc, + ) + + if attempt < max_retries - 1: + # Don't sleep past the nearest logical/run deadline: when the + # backoff plus a useful attempt no longer fit, stop burning it + # and surface now (salvage gets what's left). + deadline_remaining, deadline_reason = _nearest_deadline() + if ( + deadline_remaining is not None + and backoff + _WALL_DEADLINE_FLOOR_S > deadline_remaining + ): + logger.warning( + "Abandoning LLM retries: backoff %ds would cross the " + "%s (%.0fs left, turn=%d, attempt=%d/%d)", + int(backoff), deadline_reason, deadline_remaining, turn, + attempt + 1, max_retries, + ) + await _finish_attempt( + outcome=ATTEMPT_FAILED, + reason=deadline_reason, + recovery_action="abandon_retry", + error=retry_error, + ) + # Both deadline kinds surface the same way. Wall-deadline + # exhaustion used to ``break`` into the generic + # reason="exhausted" raise below, which threw the signal + # away: the caller could not tell "the run is out of wall + # budget, go salvage" from "this key's retries are spent, + # try another leg", even though the attempt event emitted + # just above already reported ``deadline_reason``. + deadline_exc = LLMDeadlineExceeded( + deadline_reason, + "reached before retry " + f"(backoff {int(backoff)}s + {_WALL_DEADLINE_FLOOR_S:.0f}s " + f"floor > {deadline_remaining:.0f}s left)", + ) + raise LLMCallExhausted( + deadline_exc, deadline_reason, prior_exc=retry_error, + ) from deadline_exc + await _finish_attempt( + outcome=ATTEMPT_DISCARDED, + reason=retry_reason, + recovery_action="retry_same_key", + error=retry_error, + ) + await asyncio.sleep(backoff) + else: + await _finish_attempt( + outcome=ATTEMPT_FAILED, + reason=retry_reason, + recovery_action="raise_exhausted", + error=retry_error, + ) + + logger.error("LLM call failed after %d retries (turn=%d)", max_retries, turn) + # Should always have an exception captured here — every except clause + # sets last_exc. Defensive RuntimeError covers a hypothetical + # max_retries=0 invocation, which would skip the body entirely. + if last_exc is None: + last_exc = RuntimeError( + f"call_llm exhausted with no captured exception " + f"(max_retries={max_retries}, turn={turn})", + ) + raise LLMCallExhausted(last_exc, "exhausted") from last_exc diff --git a/agent_core/runtime/loop/_response.py b/agent_core/runtime/loop/_response.py new file mode 100644 index 0000000..133dc48 --- /dev/null +++ b/agent_core/runtime/loop/_response.py @@ -0,0 +1,440 @@ +# pyright: reportMissingTypeArgument=false, reportPrivateUsage=false, reportUnknownArgumentType=false, reportUnknownMemberType=false, reportUnknownParameterType=false, reportUnknownVariableType=false, reportUnnecessaryIsInstance=false, reportUnusedFunction=false +from __future__ import annotations + +import logging +import re +from typing import Any + +from agent_core.llm import LLMResponse +from agent_core.messages import Message + +logger = logging.getLogger(__name__) +LEAKED_REASONING_KEY = "leaked_reasoning" + +# Inlined ``…`` blocks may be carried through history (so +# the model sees its prior reasoning on the next turn) but must never +# surface as a final answer. Stripped at the answer-extraction site. +_THINK_BLOCK_RE = re.compile(r"[\s\S]*?\s*") +_DANGLING_THINK_RE = re.compile(r"[\s\S]*\Z") + + +def _strip_thinking_blocks(text: str) -> str: + """Strip inlined ```` from a model-facing answer. + + Handles closed pairs, unclosed openers, and the SGLang + ``preserve_thinking`` quirk where the closing tag is emitted without + an opener — everything before the last ```` is the thinking + trace and must be stripped, not just the tag character. + """ + text = _THINK_BLOCK_RE.sub("", text) + text = _DANGLING_THINK_RE.sub("", text) + if "" in text: + text = text.rsplit("", 1)[-1] + return text.strip() + + +def _flatten_message_text(content: Any) -> str: + """Collapse an assistant message's ``content`` (str | list of str/text + blocks | other) into plain text. Thinking blocks are NOT stripped here. + + The list form covers providers that return typed content blocks; both + the current ``text`` key and the legacy ``content`` key are accepted.""" + if not content: + return "" + if isinstance(content, str): + return content + if isinstance(content, list): + parts: list[str] = [] + for block in content: + if isinstance(block, str): + parts.append(block) + elif isinstance(block, dict): + val = block.get("text") or block.get("content") or "" + if val: + parts.append(str(val)) + return "\n".join(parts) + return str(content) + + +def _visible_response_text(response: Any) -> str: + """The model-facing answer text of one response, thinking stripped.""" + return _strip_thinking_blocks(_flatten_message_text(getattr(response, "content", ""))) + +def extract_final_content(messages: list[Message]) -> str: + """Find the most recent assistant message with non-empty visible text. + + Walks backward past empty assistant messages so we surface the last + meaningful answer instead of an empty shell (common when the final + turn was a tool-only call or a safety-driven empty reply). Inlined + ``…`` blocks are stripped — they're history-only + reasoning and must never reach a downstream judge as the answer. + """ + for msg in reversed(messages): + if msg.get("role") != "assistant": + continue + cleaned = _strip_thinking_blocks(_flatten_message_text(msg.get("content", ""))) + if cleaned: + return cleaned + return "" + + +def extract_leaked_reasoning(response: Any) -> str: + """Pull reasoning recovered from leaked think/reasoning tags, if any. + + A product's tool-call parser stashes the salvaged inner text under + ``response.response_metadata[LEAKED_REASONING_KEY]`` when it strips out + leaked ```` / ```` / ```` + blocks. Mirroring it onto the turn context lets observers surface it + without grovelling through the raw response. + """ + meta = getattr(response, "response_metadata", None) or {} + value = meta.get(LEAKED_REASONING_KEY, "") + return value if isinstance(value, str) else "" + +def _pick_int(*candidates: Any) -> int: + """Return the first non-zero int-coercible candidate, else 0. + + ``None`` and unparseable values are skipped (so the next candidate + is tried), matching the ``... or ... or 0`` chains this replaces + while tolerating gateway-side ``null`` (seen on aggregating gateways + for ``cached_tokens``: the key is present but the value is JSON null, + which ``dict.get`` returns as ``None``). + + ``0`` is treated as "no signal" so a missing field can fall through + to a real source — same semantics as the chained ``or``. This is + safe because the downstream billing rollup sums per-call ints; the + only effect of picking 0 from candidate A over 0 from candidate B + is which provider name lands in the audit log, not the total. + """ + for v in candidates: + if v is None: + continue + try: + n = int(v) + except (TypeError, ValueError): + continue + if n: + return n + return 0 + + +def extract_usage(response: Any) -> dict[str, int | str] | None: + """Extract token usage from an LLM response, normalized to OpenAI shape. + + Returns a dict with keys ``provider`` / ``model`` / ``prompt_tokens`` / + ``completion_tokens`` / ``cache_read_tokens`` / ``cache_write_tokens`` / + ``cached_tokens`` / ``cache_creation_tokens`` / ``reasoning_tokens`` + (the shape consumed by the protocol stream and worker-trace observers), + or ``None`` when the response carries + no usage info. Besides the native :class:`LLMResponse`, it still + handles the two legacy response shapes products may hand it: the + canonical ``usage_metadata`` map (``input_tokens`` / ``output_tokens``, + no ``model`` key — that lives in ``response_metadata.model_name``) and + the OpenAI raw ``response_metadata.token_usage`` shape. + + ``provider`` is sourced from ``response_metadata.provider_actually_used`` + (stamped per attempt by a product's provider-chain wrapper). Empty + string when the construction path didn't stamp it — a bare client with + no chain wrapper around it. Downstream billing should treat ``""`` as + "vendor unknown — fall back to whatever the model id implies". + + Cache token fields: + + - ``cache_read_tokens`` — cache READ (a.k.a. cache hit). Bills at + ~0.1× base input on Anthropic, free on OpenAI. + - ``cache_write_tokens`` — cache WRITE (cache creation), summing + Anthropic's 5m-TTL and 1h-TTL counts (5m ~1.25×, 1h ~2×). Only + Anthropic exposes write; OpenAI returns 0 here. + - ``cached_tokens`` — backward-compat alias = ``read + write``. + Pre-split this name was cache-read-only; cost boards using it + with a single rate would have under-attributed Anthropic write + spend, which is what motivated this split. + - ``cache_creation_tokens`` — backward-compat alias of + ``cache_write_tokens`` (deprecated; prefer the new name). + + ``reasoning_tokens`` captures OpenAI o-series / Gemini thinking + output, billed as completion tokens but worth surfacing separately. + + Every key above is present on every non-``None`` return, zero-filled + when the provider reported nothing, so consumers can index the shape + without probing which response object they were handed. + """ + # Native path: ``LLMResponse.usage`` is already normalised by the client + # adapter (prompt/completion/total/cached_tokens), so read it directly. + # The ``usage_metadata`` / ``token_usage`` parsing below is retained as a + # fallback for legacy response objects that expose those shapes. + if isinstance(response, LLMResponse): + usage = response.usage or {} + inp = int(usage.get("prompt_tokens", 0) or 0) + out = int(usage.get("completion_tokens", 0) or 0) + if not inp and not out: + return None + # Vendor label stamped by a provider-chain wrapper on the + # non-streaming path; ``_stream_llm_response`` folds the streamed + # ``StreamDelta.provider`` in here. Empty when the client was built + # without a provider-stamp wrapper — downstream billing treats "" + # as "vendor unknown", same as the legacy branches below. + rmd = response.response_metadata or {} + provider = str(rmd.get("provider_actually_used") or "") if isinstance( + rmd, dict, + ) else "" + if "cache_read_tokens" in usage or "cache_write_tokens" in usage: + cache_read = int(usage.get("cache_read_tokens", 0) or 0) + cache_write = int(usage.get("cache_write_tokens", 0) or 0) + else: + # Backward compatibility for native adapters that still expose + # the pre-split cache fields. + cache_read = int(usage.get("cached_tokens", 0) or 0) + cache_write = int(usage.get("cache_creation_tokens", 0) or 0) + out_dict: dict[str, int | str] = { + "provider": provider, + "model": response.model or "", + "prompt_tokens": inp, + "completion_tokens": out, + "cache_read_tokens": cache_read, + "cache_write_tokens": cache_write, + "cached_tokens": cache_read + cache_write, + "cache_creation_tokens": cache_write, + } + # Reasoning/thinking tokens (Anthropic extended thinking / OpenAI + # reasoning models). They are part of completion_tokens but surfaced + # separately for cost / analysis; the client's usage dict carries them. + # Always present, including as 0: this branch and the legacy shapes + # below must return the SAME key set, or a consumer indexing + # ``usage["reasoning_tokens"]`` works on one response object and + # raises KeyError on the other. + out_dict["reasoning_tokens"] = int(usage.get("reasoning_tokens", 0) or 0) + return out_dict + + rmd = getattr(response, "response_metadata", None) or {} + if not isinstance(rmd, dict): + rmd = {} + # ``model_actually_used`` is stamped by a provider-chain wrapper and is + # the only model identifier present on streaming chunks (the provider's + # own ``model_name`` lands on non-streaming responses but not on + # streamed usage chunks). Falling through to it keeps streaming usage + # attribution alive. + model = ( + rmd.get("model_name") + or rmd.get("model") + or rmd.get("model_actually_used") + or "" + ) + provider = str(rmd.get("provider_actually_used") or "") + + def _build( + inp: int, + out: int, + cached: int, + cache_create: int, + reasoning: int, + ) -> dict[str, int | str]: + # ``cached`` carries cache READ; ``cache_create`` carries cache + # WRITE. The legacy ``cached_tokens`` / ``cache_creation_tokens`` + # keys are kept as a derived sum and an alias respectively so + # existing consumers don't break — see module docstring on the + # New consumers should read the explicit + # ``cache_read_tokens`` / ``cache_write_tokens`` keys. + cache_read = int(cached or 0) + cache_write = int(cache_create or 0) + return { + "provider": provider, + "model": model, + "prompt_tokens": int(inp or 0), + "completion_tokens": int(out or 0), + "cache_read_tokens": cache_read, + "cache_write_tokens": cache_write, + # Backward-compat: sum is the intuitive read of "cached" + # for cost boards using a single field. + "cached_tokens": cache_read + cache_write, + # Backward-compat alias for legacy callers; identical to + # ``cache_write_tokens`` (deprecated, will be removed in + # a future cleanup once all consumers migrate). + "cache_creation_tokens": cache_write, + "reasoning_tokens": int(reasoning or 0), + } + + # Legacy canonical shape (input_tokens / output_tokens). + um = getattr(response, "usage_metadata", None) + if um is not None and not isinstance(um, dict): + try: + um = dict(um) + except (TypeError, ValueError): + um = None + if isinstance(um, dict): + idetails = um.get("input_token_details") or {} + odetails = um.get("output_token_details") or {} + if not isinstance(idetails, dict): + idetails = {} + if not isinstance(odetails, dict): + odetails = {} + inp = _pick_int(um.get("input_tokens"), um.get("prompt_tokens")) + out = _pick_int(um.get("output_tokens"), um.get("completion_tokens")) + cached = _pick_int(idetails.get("cache_read")) + cache_create = _pick_int(idetails.get("cache_creation")) + reasoning = _pick_int(odetails.get("reasoning")) + if inp or out: + # Apodex (and some other OpenAI-compatible gateways) on + # non-streaming calls populate ``input_tokens`` / + # ``output_tokens`` on the canonical map but leave + # ``input_token_details`` empty — cache hits only show up + # on the raw ``prompt_tokens_details.cached_tokens`` field. + # Cross-check the raw shape when the canonical pass came back + # zero so streaming and non-streaming don't silently disagree + # on cached token attribution. + if not cached or not cache_create or not reasoning: + tu_raw = rmd.get("token_usage") or rmd.get("usage") + if isinstance(tu_raw, dict): + ptd_raw = tu_raw.get("prompt_tokens_details") or {} + ctd_raw = tu_raw.get("completion_tokens_details") or {} + if not isinstance(ptd_raw, dict): + ptd_raw = {} + if not isinstance(ctd_raw, dict): + ctd_raw = {} + if not cached: + cached = _pick_int( + ptd_raw.get("cached_tokens"), + tu_raw.get("cache_read_input_tokens"), + # Symmetric with the write path below: some + # gateways nest the Anthropic READ key + # under prompt_tokens_details, not at root. + ptd_raw.get("cache_read_input_tokens"), + ptd_raw.get("cache_read_tokens"), + tu_raw.get("cache_read_tokens"), + ) + if not cache_create: + cache_create = _pick_int( + tu_raw.get("cache_creation_input_tokens"), + # Apodex/qwen nest the Anthropic write key + # under prompt_tokens_details — see comments + # in the raw-shape branch below for the + # full alias list. + ptd_raw.get("cache_creation_input_tokens"), + ptd_raw.get("cache_creation_tokens"), + ptd_raw.get("cache_write_tokens"), + tu_raw.get("cache_write_tokens"), + ) + if not reasoning: + reasoning = _pick_int( + ctd_raw.get("reasoning_tokens"), + tu_raw.get("reasoning_tokens"), + ) + return _build(inp, out, cached, cache_create, reasoning) + + # OpenAI raw shape (response_metadata.token_usage / usage). + tu = rmd.get("token_usage") or rmd.get("usage") + if isinstance(tu, dict): + ptd = tu.get("prompt_tokens_details") or {} + ctd = tu.get("completion_tokens_details") or {} + if not isinstance(ptd, dict): + ptd = {} + if not isinstance(ctd, dict): + ctd = {} + inp = _pick_int(tu.get("prompt_tokens"), tu.get("input_tokens")) + out = _pick_int(tu.get("completion_tokens"), tu.get("output_tokens")) + # Cache READ (cache-hit tokens). Field name varies by provider: + # - OpenAI / OpenAI-compatible: ptd.cached_tokens + # - Anthropic direct (Messages API): tu.cache_read_input_tokens + # at the usage root, *not* nested under prompt_tokens_details + # - Apodex / bedrock via OpenAIClient gateway nest the Anthropic + # READ key UNDER prompt_tokens_details — mirror of the write + # path's ptd.cache_creation_input_tokens candidate below. Without + # ptd.cache_read_input_tokens the read count silently dropped to + # 0 on every bedrock-via-gateway call while write was captured + # so reads are not silently lost when writes are present. + # - Some custom gateways flatten: ptd.cache_read_tokens or + # tu.cache_read_tokens + cached = _pick_int( + ptd.get("cached_tokens"), + tu.get("cache_read_input_tokens"), + ptd.get("cache_read_input_tokens"), # bedrock/apodex nested shape + ptd.get("cache_read_tokens"), + tu.get("cache_read_tokens"), + ) + # Cache WRITE (cache-creation tokens). Field name varies: + # - Anthropic direct: tu.cache_creation_input_tokens at root + # - OpenAI-compatible nested (no _input_ infix): + # ptd.cache_creation_tokens + # - Apodex / qwen3.5 / some custom gateways nest the Anthropic + # name UNDER prompt_tokens_details — same key, different + # parent. Observed shape (2026-05): + # usage: { prompt_tokens_details: { + # cached_tokens: ..., cache_creation_input_tokens: ... } + # } + # Without this candidate the write count was silently dropped + # on every apodex non-streaming call (DAG analyzer / synth + # / decision_llm), under-attributing write spend on Claude + # served through the apodex gateway. + # - OpenRouter passthrough alias: ptd.cache_write_tokens + # (some wrappers drop Anthropic's standard names for this alias) + # - Some custom gateways flatten: tu.cache_write_tokens + cache_create = _pick_int( + tu.get("cache_creation_input_tokens"), + ptd.get("cache_creation_input_tokens"), # apodex/qwen shape + ptd.get("cache_creation_tokens"), + ptd.get("cache_write_tokens"), + tu.get("cache_write_tokens"), + ) + # Anthropic 1h-TTL extension (extended prompt-cache, ~2× base + # rate vs 5m's ~1.25×) surfaces under a nested ``cache_creation`` + # dict alongside the 5m count. Both bill as write, just at + # different rates — sum them so the schema field captures the + # full write footprint. Cost boards needing per-TTL breakdown + # should consume the raw provider response directly. + # + # **Provider scope**: the 1h-TTL extension is Anthropic-direct + # only as of 2026-05; Bedrock supports only the 5m TTL and + # omits the nested ``cache_creation`` dict entirely, so this + # branch is a no-op there (gracefully degrades to just the + # 5m count read above). + cc_nested = tu.get("cache_creation") + if isinstance(cc_nested, dict): + cache_create += _pick_int( + cc_nested.get("ephemeral_1h_input_tokens"), + ) + # If the root ``cache_creation_input_tokens`` was absent but + # the 5m count is nested here, pick it up. Guard against + # double-counting when both root + nested are populated. + if not tu.get("cache_creation_input_tokens"): + cache_create += _pick_int( + cc_nested.get("ephemeral_5m_input_tokens"), + ) + # Reasoning tokens (o-series / Gemini thinking / qwen thinking + # via aliyun gateway). Standard location is nested under + # completion_tokens_details, but some gateways flatten to root. + reasoning = _pick_int( + ctd.get("reasoning_tokens"), + tu.get("reasoning_tokens"), + ) + if inp or out: + return _build(inp, out, cached, cache_create, reasoning) + + return None + + +def extract_model_name( + llm: Any, profile: dict[str, Any] | None = None, +) -> str: + """Best-effort model id from a YAML profile or LLM attribute. + + Resolution order: + + 1. ``profile["llm"]["model"]`` if a profile dict was passed (workflow + YAML profiles are the authoritative source — they're what the + benchmark run was configured with). + 2. Common model-id attributes on the bound LLM + (``model_name`` / ``model`` / ``model_id``) — covers OpenAI, + Anthropic, Qwen alike. + + Returns ``""`` when nothing identifies the model — observers treat + empty as "omit the field" rather than recording an empty string. + """ + if profile: + name = (profile.get("llm") or {}).get("model") + if isinstance(name, str) and name: + return name + for attr in ("model_name", "model", "model_id"): + v = getattr(llm, attr, None) + if isinstance(v, str) and v: + return v + return "" diff --git a/agent_core/runtime/loop/_runaway.py b/agent_core/runtime/loop/_runaway.py new file mode 100644 index 0000000..2141c33 --- /dev/null +++ b/agent_core/runtime/loop/_runaway.py @@ -0,0 +1,358 @@ +# pyright: reportPrivateUsage=false, reportUnknownArgumentType=false, reportUnknownMemberType=false, reportUnknownVariableType=false, reportUnusedFunction=false +from __future__ import annotations + +import logging +from dataclasses import replace +from typing import Any + +from agent_core.llm import LLMResponse +from agent_core.messages import Message +from agent_core.runtime.env import first_configured +from agent_core.runtime.llm_request_overrides import ThinkingRetryOverride +from agent_core.tokens import estimate_message_tokens + +from ._bind import _ensure_bound +from ._response import _visible_response_text, extract_usage + +logger = logging.getLogger(__name__) +# ── Reasoning-runaway detection (capped-empty completions) ─────────── +# +# Some reasoning models behind OpenAI-compatible gateways can burn the entire +# ``max_tokens`` budget inside +# the reasoning channel and return a *successful* response with zero +# visible content and no tool calls (``finish_reason="length"``). +# ``call_llm`` treats that signature as retriable-same-key. The very first +# runaway switches the retry to a reduced ``max_tokens`` and appends a +# transient, throwaway reminder asking for concise reasoning plus visible +# output/tool use. The reminder never enters durable message history. +# If the retry budget is exhausted the response is returned as-is so the +# loop's existing no-tool nudge stays the behavioural floor — a runaway must +# never escalate into a fatal ``llm_error`` stop past turn 1. + +_RUNAWAY_MIN_OUTPUT_TOKENS = 1024 +# Ceiling for retry caps after a confirmed runaway. The actual bound cap +# is derived from the observed completion usage so low-cap profiles +# (e.g. CI smoke at 1024) are not accidentally raised to this value. +_RUNAWAY_RETRY_MAX_TOKENS = 8192 +_RUNAWAY_EXPAND_FACTOR = 1.5 +_RUNAWAY_CONTEXT_RESERVE_TOKENS = 1024 +# The reminder rides as a ``user`` turn, NOT a trailing ``system`` one. +# Providers disagree on non-leading system messages: an Anthropic adapter +# only lifts a LEADING system message out of the array and maps anything +# else through the ``{"role": "user"}`` fallthrough, and chat templates on +# SGLang/vLLM commonly render only the first system block. A user turn means +# every provider sees the same instruction in the same position. +_RUNAWAY_EXPANDED_GUIDANCE = ( + "[system reminder] The previous attempt used its full private-reasoning " + "budget without producing visible output. This task may legitimately " + "require extended scientific or mathematical reasoning, so this retry " + "has a larger thinking budget. Complete the analysis, while reserving " + "enough output for either one valid tool call or a visible answer." +) +_RUNAWAY_RECOVERY_GUIDANCE = ( + "[system reminder] The expanded-thinking retry still produced no visible " + "answer or tool call, or could not fit in the available context. Use only " + "a short, bounded reasoning pass now, then promptly emit either one valid " + "tool call or visible answer text. Do not re-derive the full plan." +) +_RUNAWAY_DIRECT_RECOVERY_GUIDANCE = ( + "[system reminder] Three consecutive attempts spent their budgets in " + "private reasoning without producing a visible answer or tool call. " + "Thinking is disabled for this retry. Immediately emit either one valid " + "tool call that advances the task or a visible best-effort answer." +) + + +def _env_value(suffix: str) -> tuple[str, str] | None: + """Return the first configured shared/legacy environment value. + + ``AGENT_CORE_*`` is the portable spelling. The product-prefixed names + remain supported while the shared package is extracted. The order is + fixed and owned by :data:`agent_core.runtime.env.ENV_PREFIXES` — a + second copy here would let the two drift apart the moment a prefix is + added or reordered in only one of them. + """ + return first_configured(suffix) + + +def _env_int(suffix: str, default: int) -> int: + configured = _env_value(suffix) + if configured is None: + return default + name, raw = configured + try: + return int(raw) + except ValueError: + logger.warning("Invalid %s=%r; using default %d", name, raw, default) + return default + + +def _env_float(suffix: str, default: float) -> float: + configured = _env_value(suffix) + if configured is None: + return default + name, raw = configured + try: + return float(raw) + except ValueError: + logger.warning("Invalid %s=%r; using default %s", name, raw, default) + return default + + +# Three semantic retries: expanded thinking → reduced thinking → thinking off. +# Read once at import; deployments may lower it as a cost-control knob. +_RUNAWAY_MAX_RETRIES = _env_int("RUNAWAY_MAX_RETRIES", 3) +_RUNAWAY_BACKOFF_S = 2.0 +_RUNAWAY_DISABLE_THINKING_AFTER = 3 +_RUNAWAY_THINKING_BUDGET_MAX = 4096 +_RUNAWAY_THINKING_BUDGET_MIN = 512 +# A lowered retry budget intentionally skips the expensive expanded phase. +_RUNAWAY_EXPAND_ENABLED = _RUNAWAY_MAX_RETRIES >= _RUNAWAY_DISABLE_THINKING_AFTER +_RUNAWAY_EXPANDED_OUTPUT_RESERVE = 0.25 +# Key under which a caller threads cross-turn runaway state +# (a mutable dict) through its metadata into ``call_llm``. Contents: +# consecutive_turns — diagnostic streak counter (log only) +# last_call_runaway_responses — surfaced on ``llm_finished`` +# last_call_runaway_reasoning_chars — surfaced on ``llm_finished`` +# last_call_recovered / last_call_reason — surfaced on ``llm_finished`` +RUNAWAY_STATE_KEY = "_runaway_state" + +def _is_runaway_response(response: Any) -> bool: + """True for a successful response whose budget went entirely to + reasoning: no visible content, no tool calls, and either + ``finish_reason="length"`` or a completion-token count too large to + be a plain empty reply (gateways that drop ``finish_reason``).""" + if not isinstance(response, LLMResponse): + return False + if response.tool_calls: + return False + if _visible_response_text(response): + return False + if response.finish_reason == "length": + return True + usage = extract_usage(response) or {} + return int(usage.get("completion_tokens") or 0) >= _RUNAWAY_MIN_OUTPUT_TOKENS + +# Continuation asked of a model whose previous reply was cut off mid-sentence. +# Deliberately does NOT reduce ``max_tokens`` the way the runaway path does: this +# model was producing real output when the cap hit, so giving it less room would +# truncate it again sooner. Brevity is requested in words instead. +TRUNCATION_CONTINUATION_GUIDANCE = ( + "[system reminder] Your previous reply hit the output token limit and was " + "cut off mid-sentence. The partial text is above. Continue from exactly " + "where it stopped — do not repeat what you already wrote, and do not start " + "over. Be brief and reach a tool call or a complete answer this time." +) + + +def is_truncated_with_text(response: Any) -> bool: + """True for a reply the token cap cut off *after* it had produced text. + + The other half of :func:`_is_runaway_response`, which handles the same + ``finish_reason="length"`` with the visible text *empty*. Between them they + cover the signal, and the split matters because the two need opposite + treatment: a runaway gets resampled at a smaller cap, while this one already + contains work worth keeping and needs to be continued. + + Nothing detected this case before. It fell through as an ordinary turn, + reached ``if not parsed_calls`` and — under ``no_tool_behavior="stop"`` — + ended the run on a sentence cut mid-token. + + Unlike the runaway detector, there is **no completion-token fallback** for + gateways that drop ``finish_reason``. That heuristic reads "a large + completion with nothing visible cannot be a plain empty reply", which is + sound only while the text is empty. With text present a large completion is + what a long legitimate answer looks like, so the same heuristic would + declare every one of them truncated. An explicit ``finish_reason`` is the + only evidence that can carry this. + """ + if not isinstance(response, LLMResponse): + return False + if response.tool_calls: + return False + if response.finish_reason != "length": + return False + return bool(_visible_response_text(response)) + + +def _runaway_retry_max_tokens(response: Any) -> int | None: + """Return a retry cap that cannot exceed the observed runaway cap. + + Each successive runaway inside one call halves again (the second + reduction is derived from a completion that was ALREADY capped), so the + squeeze is progressive. ``_RUNAWAY_MIN_OUTPUT_TOKENS`` is the floor: below + it a capped-empty completion is no longer even detectable as a runaway + (see :func:`_is_runaway_response`), so shrinking past it would trade a + diagnosable failure for a silent empty reply. + """ + usage = extract_usage(response) or {} + try: + completion_tokens = int(usage.get("completion_tokens") or 0) + except (TypeError, ValueError): + completion_tokens = 0 + if completion_tokens <= 0: + return None + if completion_tokens <= _RUNAWAY_MIN_OUTPUT_TOKENS: + return completion_tokens + return max( + _RUNAWAY_MIN_OUTPUT_TOKENS, + min(_RUNAWAY_RETRY_MAX_TOKENS, completion_tokens // 2), + ) + + +def _runaway_retry_max_tokens_from_cap(active_cap: Any) -> int | None: + """Derive a safe retry cap when an early-cancelled stream has no usage. + + The active request cap is authoritative for the upper bound. This helper + must never raise a low-cap profile toward the normal 8K recovery ceiling. + """ + try: + cap = int(active_cap) + except (TypeError, ValueError): + return None + if cap <= 0: + return None + if cap <= _RUNAWAY_MIN_OUTPUT_TOKENS: + return cap + return max( + _RUNAWAY_MIN_OUTPUT_TOKENS, + min(_RUNAWAY_RETRY_MAX_TOKENS, cap // 2), + ) + + +def _bind_reduced_max_tokens( + llm: Any, + response: Any | None = None, + *, + active_cap: Any = None, +) -> Any: + """Bind the runaway-retry ``max_tokens`` cap for follow-up attempts. + + Mirrors :func:`bind_temperature` — falls back to the original LLM + when no usable cap can be inferred. + """ + retry_max_tokens = ( + _runaway_retry_max_tokens(response) + if response is not None + else _runaway_retry_max_tokens_from_cap(active_cap) + ) + if retry_max_tokens is None: + logger.debug( + "Runaway max_tokens cap could not be inferred from usage; " + "retrying at the existing budget.", + ) + return llm + return replace(_ensure_bound(llm), max_tokens=retry_max_tokens) + + +def _expanded_retry_max_tokens( + active_cap: Any, + messages: list[Message], + context_token_limit_hint: int | None, +) -> int | None: + """Return a safe 1.5× cap, or ``None`` when context cannot hold it.""" + try: + cap = int(active_cap) + except (TypeError, ValueError): + return None + if cap <= 0: + return None + expanded = int(cap * _RUNAWAY_EXPAND_FACTOR) + if context_token_limit_hint: + available = max( + int(context_token_limit_hint) + - sum(estimate_message_tokens(message) for message in messages) + - _RUNAWAY_CONTEXT_RESERVE_TOKENS, + 0, + ) + expanded = min(expanded, available) + return expanded if expanded > cap else None + + +def _bind_expanded_max_tokens( + llm: Any, + *, + active_cap: Any, + messages: list[Message], + context_token_limit_hint: int | None, +) -> Any | None: + expanded = _expanded_retry_max_tokens( + active_cap, messages, context_token_limit_hint, + ) + if expanded is None: + return None + return replace(_ensure_bound(llm), max_tokens=expanded) + + +def _phase_reasoning_guard( + retry_number: int, + *, + previous_cap: Any, + next_cap: Any, + timeout_s: float | None, + max_tokens: int | None, +) -> tuple[float | None, int | None]: + """Scale early-runaway guards with the expanded completion cap.""" + if retry_number != 1: + return timeout_s, max_tokens + try: + ratio = int(next_cap) / int(previous_cap) + except (TypeError, ValueError, ZeroDivisionError): + return timeout_s, max_tokens + if ratio <= 1.0: + return timeout_s, max_tokens + return ( + timeout_s * ratio if timeout_s else timeout_s, + int(max_tokens * ratio) if max_tokens else max_tokens, + ) + + +def _runaway_retry_policy( + retry_number: int, + next_cap: Any, +) -> tuple[ThinkingRetryOverride, str, str]: + """Choose the next retry's task-local thinking policy and guidance.""" + if retry_number == 1: + try: + cap = max(int(next_cap), 1) + except (TypeError, ValueError): + budget = None + else: + budget = max( + int(cap * (1.0 - _RUNAWAY_EXPANDED_OUTPUT_RESERVE)), + 1, + ) + return ( + ThinkingRetryOverride( + mode="expanded", + thinking_budget=budget, + reasoning_effort="high", + ), + _RUNAWAY_EXPANDED_GUIDANCE, + "retry_expanded_cap_and_thinking", + ) + + if retry_number >= _RUNAWAY_DISABLE_THINKING_AFTER: + return ( + ThinkingRetryOverride(mode="disabled", reasoning_effort="low"), + _RUNAWAY_DIRECT_RECOVERY_GUIDANCE, + "retry_thinking_disabled", + ) + + try: + cap = int(next_cap) + except (TypeError, ValueError): + cap = _RUNAWAY_THINKING_BUDGET_MAX * 2 + budget = max( + _RUNAWAY_THINKING_BUDGET_MIN, + min(_RUNAWAY_THINKING_BUDGET_MAX, max(cap, 1) // 2), + ) + return ( + ThinkingRetryOverride( + mode="reduced", + thinking_budget=budget, + reasoning_effort="low", + ), + _RUNAWAY_RECOVERY_GUIDANCE, + "retry_reduced_cap_and_thinking", + ) diff --git a/agent_core/runtime/loop/_streaming.py b/agent_core/runtime/loop/_streaming.py new file mode 100644 index 0000000..4367e3a --- /dev/null +++ b/agent_core/runtime/loop/_streaming.py @@ -0,0 +1,539 @@ +# pyright: reportPrivateUsage=false, reportUnknownArgumentType=false, reportUnknownMemberType=false, reportUnknownVariableType=false, reportUnusedFunction=false +from __future__ import annotations + +import asyncio +import contextlib +import inspect +import logging +import time +from collections.abc import Awaitable, Callable +from typing import Any + +from agent_core.errors import ( + LLMReasoningRunaway, + LLMStreamStalled, +) +from agent_core.llm import LLMResponse +from agent_core.messages import Message, ToolCall + +from ._runaway import _env_float, _env_int + +logger = logging.getLogger(__name__) +# ── Stream-stall watchdog ───────────────────────────────────────────── +# A streaming request can be black-holed without a chunk, error, or connection +# close. Inter-chunk deadlines distinguish that state from slow decoding. +# +# The watchdog bounds the gap between consecutive stream chunks. Any +# chunk — visible text, reasoning delta, tool-call args — resets the +# timer, so reasoning-runaway streams (which decode at full speed) are +# NOT flagged. On a stall the stream is closed and the attempt retried +# under the normal transient-error budget, converting a 1200 s hang +# into a ``stall_timeout`` one. +# +# Default 180 s — generous against buffering proxies that hold chunks +# (a full-response buffer flush arrives as one late burst) and against +# providers that think silently without streaming reasoning. Tunable +# via LLM_STREAM_STALL_S under any supported prefix; <= 0 disables. +_STREAM_STALL_DEFAULT_S = 180.0 +_STREAM_STALL_ENV = "LLM_STREAM_STALL_S" + +# First-chunk (TTFT) bound — a tighter leash on the FIRST chunk only. +# On a healthy gateway the first chunk (text or reasoning delta) lands +# within seconds (measured: TTFT under 3 s on every successful call), +# so a request with ZERO chunks after tens of seconds is almost +# certainly queued into a black hole — retrying immediately beats +# waiting out the generic 180 s stall bound. Inter-chunk gaps keep the +# looser stall bound (mid-generation pauses are legitimate). +# 0 / unset disables: the first chunk is then bounded by the stall +# timeout like any other gap. Deploys facing a large shared gateway +# set ~45. +_FIRST_CHUNK_ENV = "LLM_FIRST_CHUNK_S" + +# After this many mid-stream stalls within a single ``call_llm`` +# invocation, stop same-key retrying and surface ``chain_advance`` — a +# stalled stream is likely gateway black-holing, and same-key retries keep +# hitting the same dead backend. +# +# What "surface" buys depends on which turn stalled (the product's agent +# loop decides): only a turn-1 exhaustion is re-raised to the outer chain +# wrapper to rotate to the next provider leg; past turn 1 the loop stops with +# ``llm_error`` to preserve partial content and hand off to salvage. Either +# way the win is the same — we stop burning the wall budget on a dead +# gateway instead of retrying it to exhaustion. +# +# Only fires when an outer chain is active: with no fallback leg +# configured (single-endpoint benchmark runs) there's nothing to advance +# to, so same-key retry within the wall budget stays the floor. Tunable +# via LLM_STREAM_STALL_MAX under any supported prefix; default 2 (one retry, then +# advance). <= 0 disables the escape (retry to the budget as before). +_STREAM_STALL_MAX_ENV = "LLM_STREAM_STALL_MAX" +_STREAM_STALL_MAX_DEFAULT = 2 + + +def _stream_stall_timeout_s() -> float: + return _env_float(_STREAM_STALL_ENV, _STREAM_STALL_DEFAULT_S) + + +def _first_chunk_timeout_s() -> float: + return _env_float(_FIRST_CHUNK_ENV, 0.0) + + +def _stream_stall_max_before_advance() -> int: + return _env_int(_STREAM_STALL_MAX_ENV, _STREAM_STALL_MAX_DEFAULT) + + +class _ThinkTagSplitter: + """Stateful split of an inline ``...`` text stream. + + Some providers (notably Qwen-style reasoning models behind + OpenAI-compatible endpoints) inline their reasoning channel as + literal ``...`` substrings inside the regular + content stream rather than as a typed content block. To stream the + two channels separately, we need a per-call state machine that + survives across chunks (since a tag may straddle a chunk boundary). + + ``feed(text)`` returns ``(visible_text, thinking_text)`` extracted + from this chunk. ``flush()`` drains buffered bytes at stream end + (treating unmatched leftovers as visible if outside / as thinking + if inside an unclosed ````). The splitter holds back at + most ``len(CLOSE) - 1 == 7`` chars to disambiguate a partial tag + from real content — flush() releases those. + """ + + OPEN = "" + CLOSE = "" + + def __init__(self) -> None: + self._inside = False + self._buffer = "" + + @property + def has_state(self) -> bool: + """True when we're mid-tag or holding a partial-tag tail.""" + return self._inside or bool(self._buffer) + + @staticmethod + def _suffix_overlap(buf: str, tag: str) -> int: + """Return n s.t. ``buf[-n:] == tag[:n]`` (longest match). + + Used to decide how many trailing bytes to hold back as a + possible partial-tag start. Returns 0 when no overlap — those + bytes are safe to emit immediately. + """ + max_n = min(len(tag) - 1, len(buf)) + for n in range(max_n, 0, -1): + if buf[-n:] == tag[:n]: + return n + return 0 + + def feed(self, text: str) -> tuple[str, str]: + visible_parts: list[str] = [] + thinking_parts: list[str] = [] + buf = self._buffer + text + self._buffer = "" + + while buf: + if self._inside: + idx = buf.find(self.CLOSE) + if idx == -1: + hold = self._suffix_overlap(buf, self.CLOSE) + safe = len(buf) - hold + if safe: + thinking_parts.append(buf[:safe]) + self._buffer = buf[safe:] + break + thinking_parts.append(buf[:idx]) + buf = buf[idx + len(self.CLOSE):] + self._inside = False + else: + idx = buf.find(self.OPEN) + if idx == -1: + hold = self._suffix_overlap(buf, self.OPEN) + safe = len(buf) - hold + if safe: + visible_parts.append(buf[:safe]) + self._buffer = buf[safe:] + break + if idx: + visible_parts.append(buf[:idx]) + buf = buf[idx + len(self.OPEN):] + self._inside = True + + return ("".join(visible_parts), "".join(thinking_parts)) + + def flush(self) -> tuple[str, str]: + remainder = self._buffer + self._buffer = "" + if not remainder: + return ("", "") + if self._inside: + return ("", remainder) + return (remainder, "") + + +async def _stream_llm_response( + llm: Any, + messages: list[Message], + timeout: float, + on_delta: Callable[..., Awaitable[None]], + first_chunk_s: float | None = None, + reasoning_only_timeout_s: float | None = None, + reasoning_only_max_tokens: int | None = None, +) -> LLMResponse: + """Stream ``StreamDelta``s through ``on_delta`` and fold them into an + ``LLMResponse``. + + ``LLMClient.stream`` yields normalised ``StreamDelta``s, so assembly is + a plain fold: visible content is concatenated, reasoning is accumulated, + and tool-calls are stitched by ``index`` (id set-once, name/arguments + appended). + + Terminal metadata: OpenAI streaming surfaces usage / finish_reason / + model only on the last chunks — usage on a separate ``choices=[]`` chunk + under ``include_usage``, finish_reason on the final content chunk. + ``StreamDelta`` carries all three (plus the chain-stamped ``provider``), + and the fold below keeps the last non-empty value of each, so the + assembled ``LLMResponse`` reports real usage and finish_reason for + streamed calls exactly like the non-streaming path. Before those fields + existed a streamed run billed 0 tokens and no observer could see + ``finish_reason="length"``. + + Tool-call argument streaming: each chunk's ``tool_call_deltas`` are + re-shaped and forwarded to ``on_delta`` as the ``tool_call_args_chunks`` + keyword arg so observers can decode a specific arg's value + progressively (e.g. a report tool's markdown ``content`` body streamed + out as text). The callback is invoked when any of visible text / + thinking / tool-call chunks is present in the chunk. + + Stall watchdog: the gap between consecutive chunks is bounded by the + inter-chunk stall timeout (see :data:`_STREAM_STALL_ENV`). A stream + that goes silent — gateway queue black-hole, dropped connection + without FIN — raises :class:`LLMStreamStalled` after ``stall_s`` + instead of pinning the attempt for the full call ``timeout``. Any + chunk (text / reasoning / tool args) resets the timer, so + slow-but-alive generations are never flagged. The FIRST chunk gets + an optionally tighter bound (:data:`_FIRST_CHUNK_ENV`) because a + healthy gateway delivers TTFT in seconds — zero chunks after tens + of seconds means black-holed, not thinking. The bound is one + long-lived ``asyncio.timeout`` rescheduled per chunk — a single + timer-handle mutation — rather than a per-chunk ``wait_for`` (which + would allocate a future + timer on a loop that runs ~100k times for + a long generation). + + Semantic reasoning watchdog: when enabled, the timer starts on the first + reasoning delta and is never reset by more reasoning. Non-whitespace + visible output or a non-empty tool-call delta permanently disarms it for + that attempt. The token threshold is an approximate chars/4 liveness + estimate only; provider terminal usage remains authoritative for billing. + """ + accumulated = "" + thinking_accum = "" + delta_index = 0 + # Typed as ToolCall, not dict[str, Any]: the slots below are assembled + # in the wire shape LLMResponse.tool_calls declares, and the literal + # keeps ToolCall's fixed {id, type, function} key order. + tool_call_acc: dict[int, ToolCall] = {} + # Terminal metadata streamed late by the provider — kept so the assembled + # LLMResponse carries usage/finish_reason/model (else streaming runs report + # 0 usage and observers never see finish_reason="length"). + final_usage: dict[str, int] = {} + final_finish_reason = "" + final_model = "" + # Vendor label stamped on the deltas by a product's provider-chain + # wrapper — folded into ``response_metadata`` below so + # streamed calls carry billing attribution like the non-streaming path. + final_provider = "" + think_splitter = _ThinkTagSplitter() + accepts_tool_call_chunks = _accepts_tool_call_arg_chunks(on_delta) + reasoning_timeout_s = max(float(reasoning_only_timeout_s or 0), 0.0) + reasoning_token_limit = max(int(reasoning_only_max_tokens or 0), 0) + reasoning_guard_enabled = bool(reasoning_timeout_s or reasoning_token_limit) + reasoning_only_started: float | None = None + reasoning_token_estimate = 0 + productive_output_seen = False + stall_s = _stream_stall_timeout_s() + # Per-call value (LoopConfig.first_chunk_timeout ← profile + # ``agent.first_chunk_s``) wins over the process-wide env knob; + # an explicit 0 disables even when the env is set. + first_s = ( + first_chunk_s if first_chunk_s is not None else _first_chunk_timeout_s() + ) + # The scope is armed with the (tighter) first-chunk bound when set, + # then rescheduled to the inter-chunk stall cadence once chunks flow. + initial_s = first_s if first_s > 0 else stall_s + chunks_seen = 0 + stream_started = time.monotonic() + loop = asyncio.get_running_loop() + chunk_stream = llm.stream(messages, timeout=timeout) + stall_scope: asyncio.Timeout | None = None + + def _assembled_response() -> LLMResponse: + response_metadata = ( + {"provider_actually_used": final_provider} if final_provider else {} + ) + visible_content = accumulated + if visible_content: + visible_content = ( + visible_content.lstrip() if visible_content.strip() else "" + ) + # Drop slots that never received a function name. A slot is created + # for ANY streamed tool-call delta carrying an index (see the + # ``setdefault`` below), including a content-free one that merely opens + # a tool-call block, and one whose stream was cut before the name + # arrived — both observed against a production endpoint. + # + # Such a call is unexecutable, and keeping it is worse than dropping + # it: the loop records it in DURABLE history as ``name=""``, and some + # chat templates then fail to render that history at all ("can only + # concatenate str (not \"NoneType\") to str", returned as HTTP 400). + # Every later request in the session replays the same + # history and is rejected the same way, so one malformed delta ends the + # run — and it ends it looking like an ordinary empty submission, not + # like the infrastructure fault it is. + # + # NAMED calls to tools that do not exist are deliberately kept: those + # reach the executor and come back as "unknown tool 'x'", which the + # model can read and act on. + complete_tool_calls = [ + tool_call_acc[k] + for k in sorted(tool_call_acc) + if tool_call_acc[k]["function"]["name"] + ] + # Warn rather than drop silently — a provider emitting these + # consistently is a real upstream defect, and this is the only place + # that can still see it. + dropped = len(tool_call_acc) - len(complete_tool_calls) + if dropped: + logger.warning( + "dropped %d streamed tool_call(s) with no function name", dropped, + ) + return LLMResponse( + content=visible_content, + tool_calls=complete_tool_calls, + reasoning_content=thinking_accum, + usage=final_usage, + finish_reason=final_finish_reason, + model=final_model, + response_metadata=response_metadata, + ) + + async def _close_chunk_stream() -> None: + with contextlib.suppress(Exception): + await asyncio.wait_for(chunk_stream.aclose(), timeout=5.0) + + try: + async with asyncio.timeout(timeout): + async with asyncio.timeout( + initial_s if initial_s > 0 else None, + ) as stall_scope: + async for delta in chunk_stream: + chunks_seen += 1 + raw_visible = delta.content or "" + typed_thinking = delta.reasoning_content or "" + tc_chunks = delta.tool_call_deltas or [] + # Capture terminal metadata as it arrives (usage on the + # late ``include_usage`` chunk, finish_reason on the last + # content chunk). Last non-empty wins. + if getattr(delta, "usage", None): + final_usage = delta.usage + if getattr(delta, "finish_reason", ""): + final_finish_reason = delta.finish_reason + if getattr(delta, "model", ""): + final_model = delta.model + if getattr(delta, "provider", ""): + final_provider = delta.provider + # Inline ``...`` tags (Qwen-style) are + # split out so ``delta`` carries answer-only text and + # ``thinking_delta`` collects both inline + typed + # reasoning. When neither the chunk nor the splitter + # has tag state, short-circuit to avoid scanning every + # clean chunk. + if raw_visible and ( + "" in raw_visible + or "" in raw_visible + or think_splitter.has_state + ): + visible, inline_thinking = think_splitter.feed( + raw_visible, + ) + else: + visible, inline_thinking = raw_visible, "" + thinking = (typed_thinking + inline_thinking) if ( + typed_thinking or inline_thinking + ) else "" + # Stitch streamed tool-call deltas by index — id set + # once, name/arguments appended — into wire-shaped slots. + for tcd in tc_chunks: + idx = tcd.get("index") or 0 + slot = tool_call_acc.setdefault(idx, { + "id": "", "type": "function", + "function": {"name": "", "arguments": ""}, + }) + if tcd.get("id"): + slot["id"] = tcd["id"] + if tcd.get("name"): + slot["function"]["name"] += tcd["name"] + if tcd.get("arguments"): + slot["function"]["arguments"] += tcd["arguments"] + if visible or thinking or tc_chunks: + if visible: + accumulated += visible + if thinking: + thinking_accum += thinking + tool_progress = any( + d.get("id") or d.get("name") or d.get("arguments") + for d in tc_chunks + ) + if visible.strip() or tool_progress: + productive_output_seen = True + # Forward arg deltas in the {name, args, id, index} + # shape observers expect (``args`` = partial JSON + # fragment), mirroring the old chunk extractor. + arg_chunks = [ + {"name": d.get("name"), + "args": d.get("arguments") or "", + "id": d.get("id"), "index": d.get("index")} + for d in tc_chunks + ] + if accepts_tool_call_chunks: + await on_delta( + visible, accumulated, delta_index, thinking, + tool_call_args_chunks=arg_chunks, + ) + else: + await on_delta( + visible, accumulated, delta_index, thinking, + ) + if ( + reasoning_guard_enabled + and not productive_output_seen + and thinking + ): + now = loop.time() + if reasoning_only_started is None: + reasoning_only_started = now + # Liveness-only estimate. Keep it separate from + # provider usage: early cancellation commonly + # prevents the terminal billing chunk from arriving. + reasoning_token_estimate = ( + len(thinking_accum) + 3 + ) // 4 + reasoning_elapsed = now - reasoning_only_started + time_exhausted = bool( + reasoning_timeout_s + and reasoning_elapsed >= reasoning_timeout_s + ) + tokens_exhausted = bool( + reasoning_token_limit + and reasoning_token_estimate + >= reasoning_token_limit + ) + if time_exhausted or tokens_exhausted: + trigger = "time" if time_exhausted else "tokens" + partial_response = _assembled_response() + await _close_chunk_stream() + raise LLMReasoningRunaway( + elapsed_s=reasoning_elapsed, + estimated_tokens=reasoning_token_estimate, + trigger=trigger, + partial_response=partial_response, + ) + delta_index += 1 + # One timer scope enforces both the resettable inter-chunk + # stall and the non-resettable semantic deadline. + # Reasoning chunks may move the stall edge forward, but + # min() keeps the first-reasoning deadline fixed. + # Productive output removes only the semantic candidate; + # ordinary stall handling remains. + watchdog_deadlines: list[float] = [] + if stall_s > 0: + watchdog_deadlines.append(loop.time() + stall_s) + if ( + reasoning_timeout_s + and reasoning_only_started is not None + and not productive_output_seen + ): + watchdog_deadlines.append( + reasoning_only_started + reasoning_timeout_s, + ) + stall_scope.reschedule( + min(watchdog_deadlines) + if watchdog_deadlines + else None + ) + except TimeoutError as exc: + if stall_scope is not None and stall_scope.expired(): + # OUR stall bound fired (the outer total-timeout raises with + # its own scope expired and this one fresh; an external + # cancellation re-raises CancelledError instead — neither is + # misclassified). Close the generator now, while we're not + # being cancelled, so the underlying HTTP stream is released + # immediately. + reasoning_elapsed = ( + loop.time() - reasoning_only_started + if reasoning_only_started is not None + else 0.0 + ) + if ( + reasoning_timeout_s + and not productive_output_seen + and reasoning_only_started is not None + and reasoning_elapsed >= reasoning_timeout_s + ): + partial_response = _assembled_response() + await _close_chunk_stream() + raise LLMReasoningRunaway( + elapsed_s=reasoning_elapsed, + estimated_tokens=reasoning_token_estimate, + trigger="time", + partial_response=partial_response, + ) from exc + await _close_chunk_stream() + raise LLMStreamStalled( + initial_s if chunks_seen == 0 else stall_s, chunks_seen, + time.monotonic() - stream_started, + ) from exc + raise + + # Drain any bytes the splitter held back at a partial-tag boundary. + visible_flush, thinking_flush = think_splitter.flush() + if visible_flush or thinking_flush: + if visible_flush: + accumulated += visible_flush + if thinking_flush: + thinking_accum += thinking_flush + if accepts_tool_call_chunks: + await on_delta( + visible_flush, accumulated, delta_index, thinking_flush, + tool_call_args_chunks=[], + ) + else: + await on_delta(visible_flush, accumulated, delta_index, thinking_flush) + + # Qwen chat templates delimit thinking from the visible/tool-call region + # with ``\n\n``. After SGLang's reasoning + tool parsers consume + # both structured regions, those separators can be the only bytes left in + # ``content`` (whitespace-only → drop entirely), or they lead the real + # visible text (``\n\nAnswer…`` → lstrip the remnant). Either way they + # carry no user-visible meaning; keeping the leading remnant doubles the + # separator when ``thinking_in_history`` reconstructs the turn. + return _assembled_response() + + +# Public alias — callers outside this package should depend on this name +# rather than reaching into the underscore-prefixed module directly; +# re-exported from ``llm_client``. +ThinkTagSplitter = _ThinkTagSplitter + + +def _accepts_tool_call_arg_chunks(callback: Callable[..., Awaitable[None]]) -> bool: + try: + sig = inspect.signature(callback) + except (TypeError, ValueError): + return True + for param in sig.parameters.values(): + if param.kind == inspect.Parameter.VAR_KEYWORD: + return True + if param.name == "tool_call_args_chunks": + return True + return False diff --git a/agent_core/runtime/loop/llm_client.py b/agent_core/runtime/loop/llm_client.py new file mode 100644 index 0000000..9b88d1c --- /dev/null +++ b/agent_core/runtime/loop/llm_client.py @@ -0,0 +1,65 @@ +# pyright: reportPrivateUsage=false, reportUnknownVariableType=false +"""Bind, invoke, and normalize LLM clients used by the agent loop. + +Provider adaptation stays here so the loop remains a readable sequence of +turn-level operations. +""" + +from __future__ import annotations + +from agent_core.errors import ( + LLMCallExhausted as LLMCallExhausted, +) +from agent_core.errors import ( + LLMDeadlineExceeded, + LLMReasoningRunaway, + LLMStreamStalled, +) +from agent_core.runtime.loop._bind import ( + _ensure_bound as _ensure_bound, +) +from agent_core.runtime.loop._bind import ( + bind_max_tokens, + bind_session_id, + bind_temperature, + bind_tools, +) +from agent_core.runtime.loop._call import call_llm +from agent_core.runtime.loop._response import ( + extract_final_content, + extract_leaked_reasoning, + extract_model_name, + extract_usage, +) +from agent_core.runtime.loop._runaway import ( + RUNAWAY_STATE_KEY, + TRUNCATION_CONTINUATION_GUIDANCE, + is_truncated_with_text, +) +from agent_core.runtime.loop._streaming import ThinkTagSplitter +from agent_core.tokens import ( + estimate_message_tokens, + estimate_text_tokens, +) + +__all__ = [ + "RUNAWAY_STATE_KEY", + "TRUNCATION_CONTINUATION_GUIDANCE", + "LLMCallExhausted", + "LLMDeadlineExceeded", + "LLMReasoningRunaway", + "LLMStreamStalled", + "ThinkTagSplitter", + "bind_max_tokens", + "bind_session_id", + "bind_temperature", + "bind_tools", + "call_llm", + "estimate_message_tokens", + "estimate_text_tokens", + "extract_final_content", + "extract_leaked_reasoning", + "extract_model_name", + "extract_usage", + "is_truncated_with_text", +] diff --git a/agent_core/runtime/retriable.py b/agent_core/runtime/retriable.py new file mode 100644 index 0000000..7c7a361 --- /dev/null +++ b/agent_core/runtime/retriable.py @@ -0,0 +1,553 @@ +"""Error classification for the LLM-call retry / chain-escalation machinery. + +Every LLM call site — this package's ``call_llm``, a product's +provider-chain wrapper, any workflow-specific retry helper — needs to turn +a raw ``Exception`` into one of a small number of decisions: + +- *Same key, sleep + retry* — transient network glitches, 429 + rate-limits, proxy-wrapped upstream blips. +- *Different key (or provider), no sleep* — overload, credit exhaustion, + safety-filter rejection, model-not-hosted, and auth failure: a wrong or + unauthorised key is fixed by the next leg, not by sleeping. +- *Short-circuit to salvage* — input exceeded the model's context. +- *Surface the enclosing deadline* — the run or logical call budget ended; + neither sleeping nor rotating providers can buy more time. +- *Surface immediately* — everything left over, e.g. a malformed-request + 4xx that no key or provider can satisfy. + +This module is the **single source of truth** for that classification. It +sits in the shared package rather than beside any one caller so the +generic retry loop and a product's chain helper cannot drift apart — +consistency here is what makes a sub-agent going through the generic +engine and a reporter going through the chain helper behave the same way. + +Public predicates (each pure on a single ``Exception``): + +- ``is_overloaded_error(err)`` — Anthropic 529, ``overloaded`` / + ``capacity`` substrings, narrow OpenAI 503 *overload* shape. + Deliberately does **not** match generic ``service_unavailable`` + because proxies / infra often emit that for transport problems + unrelated to upstream capacity. +- ``is_credit_exhausted(err)`` — API key credit exhausted / + insufficient quota; rotating to a different key fixes it. +- ``is_rate_limited(err)`` — per-key 429. Caller should + back off on the same key, NOT escalate the chain layer. +- ``is_context_length_error(err)`` — input exceeds the model's + context window. Caller must short-circuit to L4 salvage; no + retry, no rotation can help. +- ``is_transient_network(err)`` — request-level transient failure + that backoff-same-key fixes: ``timeout`` / ``timed out`` / connection + reset, 5xx **without** an overload signature, and the common + *proxy-wrapped* shape where an OpenAI-compatible gateway packages an + upstream 5xx into a 400 envelope with + ``code=bad_response_status_code`` / ``type=new_api_error``. +- ``is_safety_filter(err)`` — upstream content-moderation + rejection (Aliyun PAI-EAS / DashScope ``DataInspectionFailed``, + Anthropic ``input_filtered`` / ``output_filtered``, OpenAI + ``content_policy_violation``). Same-key retry is hopeless — the + filter is deterministic on the input — so this advances the chain + layer if a different provider is available, otherwise surfaces. +- ``is_model_unavailable(err)`` — provider says it doesn't host + this model (OpenRouter / new-api distributor ``model_not_found`` / + ``no_such_model`` / ``no available channel``). The current + provider can never serve the request; same-key retry just wastes + time. Advance the chain to the next leg (a different distributor + group / provider with the same canonical model) — that is what + the chain machinery exists for. +- ``is_auth_failure(err)`` — 401 / AuthenticationError / + ``invalid_api_key`` / ``Missing Authentication header``. + Same-key retry can never succeed (the key is wrong / revoked / + not whitelisted for this model). The next chain leg uses a + different key (often a different provider entirely), so advancing + is the only way forward. 403 ``forbidden`` is deliberately NOT in + this set because it can mean "scope / region / model-access + denied", which has the same root on a sibling provider. +- ``is_stream_stall(err)`` — the runtime's streaming watchdog + already retried this endpoint/key up to its stall threshold and then + surfaced for chain advance. Do not classify it as transient-network; + same-key backoff would repeat the dead-stream budget. +- :class:`~agent_core.errors.LLMDeadlineExceeded` — carries either + ``wall_deadline`` or ``logical_call_deadline`` and is deliberately excluded + from transient-network retry. + +Composed predicate: + +- ``is_retriable_with_fallback(err)`` — ``overload`` OR + ``credit_exhausted`` OR ``safety_filter`` OR ``model_unavailable`` + OR ``auth_failure`` OR ``empty_completion`` OR ``stream_stall``. The + chain-escalation trigger a chain wrapper uses to advance L1 → L2 → L3. + Intentionally narrower than "anything we might retry" — rate-limit and + transient-network errors retry on the same key, so they are NOT + included here. + +Label helper: + +- ``classify_error(err)`` — returns a deadline reason when given + :class:`~agent_core.errors.LLMDeadlineExceeded`, otherwise one of ``"context_length"`` / + ``"safety_filter"`` / ``"model_unavailable"`` / ``"auth_failure"`` / + ``"empty_completion"`` / ``"overloaded"`` / ``"credit_exhausted"`` / + ``"stream_stall"`` / ``"rate_limited"`` / ``"transient_network"`` / + ``"other"`` for the ``report.fallback`` SSE payload + + ``usage_summary.by_model[].outcome``. + +All functions are pure (no I/O, no LLM). Safe to call from +anywhere — the generic engine retry loop, the chain helper, an +observer, a test. +""" + +from __future__ import annotations + +import re + +from agent_core.errors import LLMDeadlineExceeded + +_OVERLOAD_PATTERNS = ( + re.compile(r"overload", re.IGNORECASE), + re.compile(r"capacity", re.IGNORECASE), + re.compile(r"529", re.IGNORECASE), + # Anthropic explicit error type from the SDK. + re.compile(r"overloaded_error", re.IGNORECASE), +) + +_CREDIT_PATTERNS = ( + re.compile(r"credit", re.IGNORECASE), + re.compile(r"insufficient[_\s]*quota", re.IGNORECASE), + re.compile(r"insufficient[_\s]*balance", re.IGNORECASE), + re.compile(r"billing", re.IGNORECASE), + re.compile(r"payment[_\s]*required", re.IGNORECASE), + re.compile(r"\b402\b"), +) + +_RATE_LIMIT_PATTERNS = ( + re.compile(r"rate[_\s]*limit", re.IGNORECASE), + re.compile(r"\b429\b"), +) + +_CONTEXT_LENGTH_PATTERNS = ( + re.compile(r"context[_\s]*length", re.IGNORECASE), + re.compile(r"context_length_exceeded", re.IGNORECASE), + re.compile(r"longer than the model", re.IGNORECASE), + re.compile(r"maximum context", re.IGNORECASE), +) + +# Transient network / proxy-wrap signatures. ``bad_response_status_code`` +# + ``new_api_error`` are the new-api gateway's way of forwarding an +# upstream 5xx / timeout as a 400 envelope to the client — sleeping and +# retrying the same key is the right response, NOT raising a 400 as +# non-transient. Observed 2026-05-12 in multi-turn smokes against an +# OpenAI-compatible proxy. +_TRANSIENT_NETWORK_PATTERNS = ( + re.compile(r"\btimeout\b", re.IGNORECASE), + re.compile(r"timed[\s_]*out", re.IGNORECASE), + re.compile( + r"connection[\s_]*(?:reset|refused|aborted|error|closed)", + re.IGNORECASE, + ), + re.compile(r"bad_response_status_code", re.IGNORECASE), + re.compile(r"new_api_error", re.IGNORECASE), + # Upstream gateway timeouts that get text-wrapped before our status + # extractor sees them. + re.compile(r"gateway[\s_]*time[\s_]*out", re.IGNORECASE), + re.compile(r"upstream[\s_]*(?:timeout|error)", re.IGNORECASE), +) + +# Runtime stream watchdog. Matched by type name / message text instead of +# importing ``LLMStreamStalled`` from core to keep infra free of core imports. +_STREAM_STALL_PATTERNS = ( + re.compile(r"\bLLMStreamStalled\b"), + re.compile(r"stream[_\s-]*stalled", re.IGNORECASE), + re.compile(r"no chunks for", re.IGNORECASE), +) + +# Upstream content-moderation rejections. Same-key retry is hopeless — +# the filter is deterministic on the same input — so these advance the +# chain to the next provider when one is configured. Patterns cover the +# four providers we have first-hand evidence of: +# +# - Aliyun PAI-EAS / DashScope wraps Qwen behind a ``数据安全检查`` layer; +# rejection shape is 400 + ``code: data_inspection_failed`` + +# ``type: data_inspection_failed`` (observed in a browse-comparison +# smoke against Italian/political prompts). +# - Anthropic returns ``input_filtered`` / ``output_filtered`` blocks on +# policy violations (rare on Claude 4.x but documented). +# - OpenAI ``content_policy_violation`` / ``content_filter`` on Azure + +# o1/gpt-4 deployments with strict moderation enabled. +# - GPT-5.x specifically emits ``Invalid prompt: we've limited access to +# this content for safety reasons. This type of information may be used +# to benefit or to harm people...`` as both a pre-flight 400 and a +# mid-stream error event (observed when reporter_v2 prompted gpt-5.x on +# biomedical / dual-use research questions, 2026-05-29). The +# distinctive substrings are unique enough to anchor on and they +# survive minor wording tweaks. +_SAFETY_FILTER_PATTERNS = ( + re.compile(r"data[_\s]*inspection[_\s]*failed", re.IGNORECASE), + re.compile(r"content[_\s]*policy[_\s]*violation", re.IGNORECASE), + re.compile(r"content[_\s]*filter(?:ed)?", re.IGNORECASE), + re.compile(r"input[_\s]*filtered", re.IGNORECASE), + re.compile(r"output[_\s]*filtered", re.IGNORECASE), + re.compile(r"inappropriate[_\s]*content", re.IGNORECASE), + re.compile(r"prompt[_\s]*blocked", re.IGNORECASE), + # GPT-5.x "Invalid prompt: we've limited access to this content for + # safety reasons..." family. We match short substrings so the check + # survives wording tweaks and translated variants. + re.compile(r"limited[_\s]*access[_\s]*to[_\s]*this[_\s]*content[_\s]*for[_\s]*safety", re.IGNORECASE), + re.compile(r"may[_\s]*be[_\s]*used[_\s]*to[_\s]*benefit[_\s]*or[_\s]*to[_\s]*harm", re.IGNORECASE), + re.compile(r"violates[_\s]*our[_\s]*usage[_\s]*policies", re.IGNORECASE), + re.compile(r"your[_\s]*request[_\s]*was[_\s]*blocked", re.IGNORECASE), + # OpenRouter and general content moderation blocks: + re.compile(r"content[_\s]*moderation", re.IGNORECASE), + re.compile(r"moderation[_\s]*policy", re.IGNORECASE), + re.compile(r"request[_\s]*blocked", re.IGNORECASE), + re.compile(r"safety[_\s]*system", re.IGNORECASE), +) + +# Provider says it doesn't host this model. Same-key retry is futile; +# advance the chain to the next leg. Patterns cover: +# - OpenRouter / new-api distributor: ``code=model_not_found`` body + +# ``"No available channel for model ... under group ..."`` message +# (observed in the 2026-05-16 heavy-mode e2e against +# ``api.miromind.site``). +# - OpenAI-compatible gateways that surface a 404 / 400 with +# ``no_such_model`` or ``model_not_supported``. +# - Anthropic: structured ``not_found_error`` body (snake_case literal in +# ``body.type``) + the ``NotFoundError`` SDK class name surfaced via +# ``type(err).__name__`` in :func:`_stringify` (both openai-python and +# anthropic-python raise a ``NotFoundError`` class on 404). Patterns +# are written precisely — a permissive ``not[_\s]*found[_\s]*error`` +# would false-match Python's builtin ``FileNotFoundError`` and silently +# advance the chain on unrelated file-IO errors. +_MODEL_UNAVAILABLE_PATTERNS = ( + re.compile(r"model[_\s]*not[_\s]*found", re.IGNORECASE), + re.compile(r"no[_\s]*such[_\s]*model", re.IGNORECASE), + re.compile(r"model[_\s]*not[_\s]*supported", re.IGNORECASE), + re.compile(r"no[_\s]*available[_\s]*channel", re.IGNORECASE), + re.compile(r"unsupported[_\s]*model", re.IGNORECASE), + re.compile(r"not_found_error"), + re.compile(r"\bNotFoundError\b"), + # OpenRouter (observed in chaos scenario 02, 2026-05-21): a 400 + # BadRequest with body ``"X is not a valid model ID"``. Same root + # cause as model_not_found — the gateway refuses to route. Same-key + # retry can't fix a typo'd model name; the next chain leg may use a + # different canonical model spec and succeed. + re.compile(r"not[_\s]*a[_\s]*valid[_\s]*model[_\s]*id", re.IGNORECASE), + re.compile(r"\binvalid[_\s]*model[_\s]*id\b", re.IGNORECASE), + re.compile(r"\bunknown[_\s]*model\b", re.IGNORECASE), +) + +# Authentication failures from upstream providers. Discovered by a +# key-clobber chaos scenario (2026-05-21): a bare ``OPENROUTER_API_KEY`` +# clobber surfaced ``openai.AuthenticationError`` with body +# ``{"error": {"message": "Missing Authentication header", "code": 401}}``. +# Without auth here, ``is_retriable_with_fallback`` returned False and the +# product's chain wrapper never +# advanced to L2 / L3 — the reporter crashed with exit 1. Real users +# hit this every time a key is revoked or scoped wrong; rotating to the +# next chain leg (different key OR different provider) is the only +# recovery, hence chain-advance is correct. +# +# Patterns cover the four wire shapes we have first-hand evidence of: +# - OpenAI / OpenRouter raw 401 body shapes (``invalid_api_key`` / +# ``invalid_authentication`` / bare ``unauthorized``). +# - OpenRouter's specific "Missing Authentication header" surface +# (happens when the SDK suppresses an obviously-bogus Bearer value). +# - The openai-python SDK class name surfaced via ``type(err).__name__`` +# in ``_stringify`` — both ``AuthenticationError`` (openai-python) and +# the equivalent anthropic-python shape. +# - Bare ``401`` status code (the bottom-of-the-barrel fallback when an +# upstream wrapper strips structured fields but keeps the status). +_AUTH_FAILURE_PATTERNS = ( + re.compile(r"\bAuthenticationError\b"), + re.compile(r"\bauthentication[_\s]*failed", re.IGNORECASE), + re.compile(r"\binvalid[_\s]*api[_\s]*key", re.IGNORECASE), + re.compile(r"\binvalid[_\s]*authentication", re.IGNORECASE), + re.compile(r"\bmissing[_\s]*authentication", re.IGNORECASE), + re.compile(r"\bunauthorized\b", re.IGNORECASE), + re.compile(r"\bunauthenticated\b", re.IGNORECASE), + re.compile(r"\b401\b"), +) + + +def _stringify(err: BaseException) -> str: + """Concatenate every signal an LLM SDK might surface.""" + parts: list[str] = [type(err).__name__, str(err)] + for attr in ("status_code", "response", "body", "message"): + val = getattr(err, attr, None) + if val is not None: + parts.append(str(val)) + return " | ".join(parts) + + +def get_status_code(err: BaseException) -> int | None: + """Best-effort integer HTTP status extraction. + + Public so :mod:`agent_core.runtime.loop._call` can share this + heuristic instead of keeping its own copy — both call sites need the + exact same status-attribute lookup order. + """ + for attr in ("status_code", "status", "code"): + val = getattr(err, attr, None) + if isinstance(val, int): + return val + return None + + +# Private alias for this module's own call site below. +_get_status_code = get_status_code + + +def is_overloaded_error(err: BaseException) -> bool: + """True if ``err`` indicates the upstream provider is at capacity. + + Triggers fallback key rotation (the next key shares the provider so + capacity is rarely fixed by rotation alone — but it's the cheapest + signal we have and miroflow's prod observed that key-specific + capacity quirks DO exist). + """ + blob = _stringify(err) + return any(p.search(blob) for p in _OVERLOAD_PATTERNS) + + +def is_credit_exhausted(err: BaseException) -> bool: + """True if ``err`` indicates the current API key has run out of + credit / quota. Rotating to a different key usually fixes it.""" + blob = _stringify(err) + return any(p.search(blob) for p in _CREDIT_PATTERNS) + + +def is_rate_limited(err: BaseException) -> bool: + """True if ``err`` is a per-key rate-limit (429). Rotating keys + likely helps; backing off also helps.""" + blob = _stringify(err) + return any(p.search(blob) for p in _RATE_LIMIT_PATTERNS) + + +def is_context_length_error(err: BaseException) -> bool: + """True if the input exceeds the model's context window. + + Per spec §5 decision table: callers must short-circuit to L4 + salvage rather than retry or rotate keys — retrying will just hit + the same wall. + """ + blob = _stringify(err) + return any(p.search(blob) for p in _CONTEXT_LENGTH_PATTERNS) + + +def is_transient_network(err: BaseException) -> bool: + """True if ``err`` is a request-level transient (timeout / connection + reset / upstream 5xx / proxy-wrapped upstream blip). + + Decision per spec §5: sleep + retry on the SAME key. Does NOT + escalate chain layers — the next provider would see the same + transient at roughly the same rate, so burning fallback keys here + is counter-productive. + + The 5xx-without-overload branch catches the common case where a + proxy hands back ``502 / 503 / 504`` without any overload substring + — that's a network problem, not a capacity problem. ``is_overloaded_error`` + keeps priority so a 503 with ``overloaded_error`` in the body still + routes to rotation rather than backoff. + """ + if isinstance(err, LLMDeadlineExceeded): + return False + if is_stream_stall(err): + return False + if is_overloaded_error(err): + return False + # model_unavailable is also a 5xx (typically 503 from distributor + # proxies) but the right response is "advance the chain", not + # "backoff and retry same key" — same-key retries are guaranteed + # to fail with the same model_not_found. Surrender precedence to + # is_retriable_with_fallback here. + if is_model_unavailable(err): + return False + status = _get_status_code(err) + if status is not None and 500 <= status < 600: + return True + blob = _stringify(err) + return any(p.search(blob) for p in _TRANSIENT_NETWORK_PATTERNS) + + +def is_stream_stall(err: BaseException) -> bool: + """True when ``call_llm`` has surfaced a repeated stream watchdog stall. + + The watchdog has already spent the configured same-endpoint stall budget + before this exception reaches an outer chain runner, so the correct chain + decision is immediate key/provider advance rather than same-key backoff. + """ + blob = _stringify(err) + return any(p.search(blob) for p in _STREAM_STALL_PATTERNS) + + +def is_safety_filter(err: BaseException) -> bool: + """True if ``err`` is an upstream content-moderation rejection. + + Retrying the same key is hopeless (filter is deterministic on the + input). Caller should advance the chain to a different provider if + one is configured; if not, the error surfaces to the user. + """ + blob = _stringify(err) + return any(p.search(blob) for p in _SAFETY_FILTER_PATTERNS) + + +def is_model_unavailable(err: BaseException) -> bool: + """True if the provider doesn't host the requested model. + + Distributor proxies (OpenRouter aggregator, new-api / miromind.site + gateway) return ``code=model_not_found`` (often with a 503 status + when the upstream channel pool is empty) when the model name they + received isn't routable to any backend. Same-key retry is pointless — + the next provider leg in the chain may have a different upstream + that DOES host the model, so advance instead. + """ + blob = _stringify(err) + return any(p.search(blob) for p in _MODEL_UNAVAILABLE_PATTERNS) + + +def is_auth_failure(err: BaseException) -> bool: + """True if ``err`` is an upstream authentication failure (401). + + Catches the four observed shapes: ``openai.AuthenticationError`` SDK + class, bare ``unauthorized`` text, ``invalid_api_key`` / + ``invalid_authentication`` structured codes, and OpenRouter's + "Missing Authentication header" wire surface. Status 403 is + deliberately excluded — 403 means "key authenticated but not + authorised for this resource", which often has the same root on a + sibling provider (e.g. account scoped to specific model families). + Surface 403 to the operator instead of silently advancing. + + Caller (chain wrapper) advances to the next leg on True. Same-key + retry can never succeed because the rejection is deterministic on + (current key, current model). See module docstring for the chaos + test that discovered this gap (2026-05-21). + """ + blob = _stringify(err) + return any(p.search(blob) for p in _AUTH_FAILURE_PATTERNS) + + +# A bare ``ValueError("No generation chunks were returned")`` — the shape +# LangChain-based clients raise — arrives when an upstream stream completes +# but yields no usable *content*. That is the dominant failure shape for +# self-hosted reasoning models that run away in the ``reasoning_content`` +# channel and never emit a content token before hitting ``max_tokens``. +# Observed 2026-05-29: a single empty completion among the ~150 LLM calls of +# one run was fatal because nothing classified it as recoverable, so every +# multi-call run eventually died on one. The HTTP call +# *succeeded* — this is not a timeout/network class — so it gets its own +# detector rather than folding into ``is_transient_network``. +_EMPTY_COMPLETION_PATTERNS = ( + re.compile( + r"no[\s_]*generation[\s_]*chunks?[\s_]*(?:were[\s_]*)?returned", + re.IGNORECASE, + ), + re.compile( + r"no[\s_]*completion[\s_]*(?:tokens?|content)[\s_]*returned", + re.IGNORECASE, + ), + re.compile(r"empty[\s_]*completion", re.IGNORECASE), +) + + +def is_empty_completion(err: BaseException) -> bool: + """True if the upstream returned a successful response with no content. + + Distinct from a network/timeout error: the call succeeded at the HTTP + layer but produced zero content tokens (reasoning-runaway, + all-tokens-in-thinking, or an empty stream). Routed through + :func:`is_retriable_with_fallback` so the caller retries the same key + first (a temperature>0 resample frequently recovers) and then advances + the chain to a different provider, which always recovers. + """ + blob = _stringify(err) + return any(p.search(blob) for p in _EMPTY_COMPLETION_PATTERNS) + + +def is_retriable_with_fallback(err: BaseException) -> bool: + """The chain-escalation trigger. + + Per spec §5 decision table: overload, credit_exhausted, + safety_filter, model_unavailable, AND auth_failure advance the + chain layer. rate_limit + transient_network trigger backoff-same-key + instead (handled by the caller, not this predicate). Each of these + is deterministic on (current provider, current input) — only + switching providers / keys can change the outcome. + + ``empty_completion`` also routes here: it isn't deterministic on the + input (a temp>0 resample may recover), but the caller's same-key + retry budget runs first, and advancing the chain afterwards is the + guaranteed recovery — so it belongs to the same predicate. + """ + return ( + is_overloaded_error(err) + or is_credit_exhausted(err) + or is_safety_filter(err) + or is_model_unavailable(err) + or is_auth_failure(err) + or is_empty_completion(err) + or is_stream_stall(err) + ) + + +def classify_error(err: BaseException) -> str: + """Short reason label for the ``report.fallback`` SSE payload. + + Precedence (top wins): + runtime deadline → ``context_length`` → ``safety_filter`` → + ``model_unavailable`` → + ``auth_failure`` → ``empty_completion`` → ``overloaded`` → + ``credit_exhausted`` → ``stream_stall`` → ``rate_limited`` → + ``transient_network`` → ``other``. + + Context-length wins outright because its caller behaviour differs + (short-circuit to salvage). Safety-filter wins next because the + operator dashboard needs to distinguish "model refused" from + capacity issues. Model-unavailable wins over overload because the + operator response is different — overload is "wait or fan out", + model-unavailable is "fix the chain config". Auth-failure sits + above overload/credit because the operator action is also a + config fix (rotate / revoke key) — splitting it out from + ``other`` makes dashboards immediately point at the right knob. + ``empty_completion`` outranks overload/credit for the same reason in + reverse: several gateways answer a capacity problem with a 200 and an + empty body, so labelling it by its own shape keeps a silent-empty + endpoint from being read as ordinary overload. + """ + if isinstance(err, LLMDeadlineExceeded): + return err.reason + if is_context_length_error(err): + return "context_length" + if is_safety_filter(err): + return "safety_filter" + if is_model_unavailable(err): + return "model_unavailable" + if is_auth_failure(err): + return "auth_failure" + if is_empty_completion(err): + return "empty_completion" + if is_overloaded_error(err): + return "overloaded" + if is_credit_exhausted(err): + return "credit_exhausted" + if is_stream_stall(err): + return "stream_stall" + if is_rate_limited(err): + return "rate_limited" + if is_transient_network(err): + return "transient_network" + return "other" + + +__all__ = [ + "classify_error", + "get_status_code", + "is_auth_failure", + "is_context_length_error", + "is_credit_exhausted", + "is_empty_completion", + "is_model_unavailable", + "is_overloaded_error", + "is_rate_limited", + "is_retriable_with_fallback", + "is_safety_filter", + "is_stream_stall", + "is_transient_network", +] diff --git a/agent_core/tokens.py b/agent_core/tokens.py index 51ae815..05cd7e5 100644 --- a/agent_core/tokens.py +++ b/agent_core/tokens.py @@ -49,7 +49,7 @@ def _tool_calls_text(tool_calls: Any) -> str: """Serialise ``tool_calls`` the way their token cost is actually incurred. Reading named keys off each call was the earlier approach and it silently - measured nothing: it looked for ``name`` / ``args``, which is the LangChain + measured nothing: it looked for the legacy flat ``name`` / ``args`` shape, while the canonical shape in ``core.messages.ToolCall`` is OpenAI's ``{"id", "type", "function": {"name", "arguments"}}``. Every OpenAI-shaped tool call therefore contributed zero, and tool arguments are routinely the diff --git a/docs/llm-runtime-boundary.md b/docs/llm-runtime-boundary.md new file mode 100644 index 0000000..724982f --- /dev/null +++ b/docs/llm-runtime-boundary.md @@ -0,0 +1,35 @@ +# LLM runtime boundary + +This extraction moves one complete physical-call layer into AgentCore: + +- client binding and per-call overrides; +- response content, usage, and model-name normalization; +- streaming assembly, first-chunk/stall watchdogs, and reasoning guards; +- retry, backoff, provider-fallback classification, and runaway recovery; +- the public `llm_client` facade. + +The source was converged from ApodexHarness and FrontierAgentInternal as one +batch. Small pre-existing differences were resolved as compatible supersets: + +- `bind_max_tokens` and `ThinkTagSplitter` remain public; +- response blocks accept both `text` and legacy `content` fields; +- HTTP status extraction is shared through the public retry classifier; +- the portable `AGENT_CORE_` environment prefix wins, followed by the two + legacy product prefixes. + +AgentCore does not import product execution context or provider adapters. +Products supply their remaining decisions through three callbacks: + +- `call_llm(..., wall_deadline_remaining=...)` reads the active execution + scope's remaining wall budget; +- `call_llm(..., chain_fallback_active=...)` reports whether another provider + chain leg exists; +- `bind_session_id(..., sticky_session_enabled=...)` applies the product's + session-affinity kill switch. + +Products keep thin wrappers that inject these callbacks and re-export the +shared API. Provider clients consume `current_thinking_retry_override()` to +translate semantic retry intent into provider-specific request fields. + +Model profiles, provider client construction, tool parsing/execution, and the +agent loop remain product-owned in this phase. diff --git a/pyproject.toml b/pyproject.toml index 04730fb..6da2bdd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,6 +34,20 @@ line-length = 100 select = ["E", "F", "I", "UP", "FA", "B", "SIM", "RUF"] ignore = ["E501", "RUF001", "RUF002", "RUF003"] +[tool.ruff.lint.per-file-ignores] +# ``call_llm``'s per-attempt async closures are created and awaited inside the +# same physical-attempt iteration; none escapes to observe a later loop +# binding. B023 cannot infer that lifetime and reports all 21 captured +# attempt-local values. +# +# Do NOT "fix" this by binding them as default arguments (the usual B023 +# remedy): ``_finish_attempt`` MUST observe ``attempt_index`` / +# ``attempt_started`` / ``attempt_first_delta`` as they stand when it runs, +# because the empty-tool-arguments replay rebinds all three mid-iteration to +# open a second physical attempt. Snapshotting them at definition time would +# silently bill the replay against the discarded stream's attempt id. +"agent_core/runtime/loop/_call.py" = ["B023"] + [tool.pyright] pythonVersion = "3.12" include = ["agent_core"] diff --git a/tests/test_llm_runtime_deadline_reason.py b/tests/test_llm_runtime_deadline_reason.py new file mode 100644 index 0000000..eff64fb --- /dev/null +++ b/tests/test_llm_runtime_deadline_reason.py @@ -0,0 +1,299 @@ +"""``LLMCallExhausted.last_exc`` must always agree with ``.reason``. + +A chain wrapper classifies ``last_exc`` directly to decide whether to sleep, +rotate a provider leg, or give up. So a deadline-driven raise that carried an +unrelated earlier failure in ``last_exc`` told the wrapper the opposite of +what ``reason`` said: reason=``wall_deadline`` ("the run is out of budget") +paired with a stale 429 ("sleep and retry this key"), and the wrapper would +retry straight past the deadline the reason had just announced. + +The superseded failure is still reachable as ``prior_exc`` — diagnostics, +deliberately outside the field classification reads. +""" + +from __future__ import annotations + +import asyncio +import time +from types import SimpleNamespace + +import pytest + +import agent_core.runtime.loop._call as call_module +from agent_core.llm import LLMResponse +from agent_core.messages import user_msg +from agent_core.runtime.loop.llm_client import ( + LLMCallExhausted, + LLMDeadlineExceeded, + call_llm, +) + + +class _RateLimit429(Exception): + def __init__(self) -> None: + super().__init__("rate limited") + self.status_code = 429 + + +@pytest.fixture +def _instant_sleep(monkeypatch): + """Keep backoff schedules deterministic without burning wall time.""" + real_sleep = asyncio.sleep + + async def _record(_duration): + await real_sleep(0) + + monkeypatch.setattr("asyncio.sleep", _record) + + +@pytest.mark.asyncio +async def test_pre_attempt_refusal_after_a_failed_attempt_reports_the_deadline( + monkeypatch, +) -> None: + """Attempt 1 fails 429, its backoff fits, then attempt 2 is refused. + + Covers the pre-attempt floor check specifically: the budget is ample when + the retry sleep is decided and under the floor by the time the next + attempt would start, so the refusal happens before the provider is + touched a second time. + """ + monkeypatch.setattr(call_module, "_WALL_DEADLINE_FLOOR_S", 20.0) + calls = 0 + slept = False + real_sleep = asyncio.sleep + + async def _sleep(_duration): + nonlocal slept + slept = True + await real_sleep(0) + + monkeypatch.setattr("asyncio.sleep", _sleep) + + async def _chat(_messages, **_kw): + nonlocal calls + calls += 1 + raise _RateLimit429 + + # Keyed on the retry sleep actually having happened, not on how often the + # deadline is polled: ample while the backoff is being decided, under the + # floor once it has been taken. + def _remaining() -> float: + return 5.0 if slept else 600.0 + + with pytest.raises(LLMCallExhausted) as exc_info: + await call_llm( + SimpleNamespace(chat=_chat, model="fake"), + [user_msg("hello")], + timeout=120, + max_retries=4, + turn=1, + retry_wait_fixed=1, # 1s + 20s floor < 600s, so the sleep is taken + wall_deadline_remaining=_remaining, + ) + + exhausted = exc_info.value + assert slept, "expected the retry backoff to be taken, not abandoned" + assert calls == 1, "the second attempt must be refused before the provider" + assert exhausted.reason == "wall_deadline" + # The reason is a deadline, so last_exc must be the deadline — not the 429 + # a chain wrapper would have read as "sleep and retry". + assert isinstance(exhausted.last_exc, LLMDeadlineExceeded) + assert exhausted.last_exc.reason == exhausted.reason + assert "wall_deadline" in str(exhausted.last_exc) + assert isinstance(exhausted.prior_exc, _RateLimit429) + + +@pytest.mark.asyncio +async def test_backoff_crossing_the_wall_deadline_keeps_the_deadline_reason( + monkeypatch, _instant_sleep, +) -> None: + """Abandoning the retry sleep must not relabel itself ``exhausted``. + + This path used to ``break`` into the generic exhausted-retries raise, + discarding the wall-deadline signal even though the attempt event emitted + alongside it already reported ``wall_deadline`` — so a caller could not + tell "out of run budget, go salvage" from "this key is spent, try + another leg". + """ + monkeypatch.setattr(call_module, "_WALL_DEADLINE_FLOOR_S", 20.0) + calls = 0 + + async def _chat(_messages, **_kw): + nonlocal calls + calls += 1 + raise _RateLimit429 + + # Above the floor (so the attempt runs) but too small for the 429 backoff + # plus the floor, so the retry sleep is abandoned instead of taken. + with pytest.raises(LLMCallExhausted) as exc_info: + await call_llm( + SimpleNamespace(chat=_chat, model="fake"), + [user_msg("hello")], + timeout=30, + max_retries=4, + turn=1, + wall_deadline_remaining=lambda: 25.0, + ) + + exhausted = exc_info.value + assert calls == 1 + assert exhausted.reason == "wall_deadline" + assert isinstance(exhausted.last_exc, LLMDeadlineExceeded) + assert exhausted.last_exc.reason == exhausted.reason + assert isinstance(exhausted.prior_exc, _RateLimit429) + + +@pytest.mark.asyncio +async def test_logical_deadline_abandoning_backoff_reports_its_own_reason( + monkeypatch, _instant_sleep, +) -> None: + """The logical call deadline is labelled distinctly from the run wall.""" + monkeypatch.setattr(call_module, "_WALL_DEADLINE_FLOOR_S", 20.0) + + async def _chat(_messages, **_kw): + raise _RateLimit429 + + with pytest.raises(LLMCallExhausted) as exc_info: + await call_llm( + SimpleNamespace(chat=_chat, model="fake"), + [user_msg("hello")], + timeout=30, + max_retries=4, + turn=1, + logical_call_timeout_s=25.0, + ) + + assert exc_info.value.reason == "logical_call_deadline" + assert isinstance(exc_info.value.last_exc, LLMDeadlineExceeded) + assert exc_info.value.last_exc.reason == exc_info.value.reason + assert isinstance(exc_info.value.prior_exc, _RateLimit429) + + +@pytest.mark.asyncio +async def test_wall_clamped_final_attempt_timeout_reports_wall_deadline( + monkeypatch, +) -> None: + """A provider still running when the wall budget expires is not exhausted.""" + monkeypatch.setattr(call_module, "_WALL_DEADLINE_FLOOR_S", 0.0) + deadline = time.monotonic() + 0.05 + + async def _chat(_messages, **_kw): + await asyncio.sleep(1) + + with pytest.raises(LLMCallExhausted) as exc_info: + await call_llm( + SimpleNamespace(chat=_chat, model="fake"), + [user_msg("hello")], + timeout=1, + max_retries=1, + turn=1, + wall_deadline_remaining=lambda: deadline - time.monotonic(), + ) + + exhausted = exc_info.value + assert exhausted.reason == "wall_deadline" + assert isinstance(exhausted.last_exc, LLMDeadlineExceeded) + assert exhausted.last_exc.reason == exhausted.reason + assert exhausted.prior_exc is None + + +@pytest.mark.asyncio +async def test_in_flight_logical_deadline_preserves_reason(monkeypatch) -> None: + """The logical deadline remains identifiable after callers unwrap it.""" + monkeypatch.setattr(call_module, "_WALL_DEADLINE_FLOOR_S", 0.0) + + async def _chat(_messages, **_kw): + await asyncio.sleep(1) + + with pytest.raises(LLMCallExhausted) as exc_info: + await call_llm( + SimpleNamespace(chat=_chat, model="fake"), + [user_msg("hello")], + timeout=1, + max_retries=1, + turn=1, + logical_call_timeout_s=0.05, + ) + + exhausted = exc_info.value + assert exhausted.reason == "logical_call_deadline" + assert isinstance(exhausted.last_exc, LLMDeadlineExceeded) + assert exhausted.last_exc.reason == exhausted.reason + assert exhausted.prior_exc is None + + +@pytest.mark.asyncio +async def test_real_concurrency_gate_wait_preserves_wall_deadline( + monkeypatch, +) -> None: + """E2E: a real semaphore holder makes a second call consume its wall budget.""" + monkeypatch.setenv("AGENT_CORE_LLM_MAX_CONCURRENT", "1") + monkeypatch.setattr(call_module, "_llm_gate_state", None) + monkeypatch.setattr(call_module, "_WALL_DEADLINE_FLOOR_S", 0.02) + holder_started = asyncio.Event() + release_holder = asyncio.Event() + calls = 0 + + async def _chat(_messages, **_kw): + nonlocal calls + calls += 1 + holder_started.set() + await release_holder.wait() + return LLMResponse(content="ok") + + llm = SimpleNamespace(chat=_chat, model="fake") + holder = asyncio.create_task(call_llm( + llm, + [user_msg("holder")], + timeout=1, + max_retries=1, + turn=1, + )) + await asyncio.wait_for(holder_started.wait(), timeout=1) + deadline = time.monotonic() + 0.15 + + try: + with pytest.raises(LLMCallExhausted) as exc_info: + await call_llm( + llm, + [user_msg("waiting")], + timeout=1, + max_retries=1, + turn=1, + wall_deadline_remaining=lambda: deadline - time.monotonic(), + ) + finally: + release_holder.set() + await holder + + exhausted = exc_info.value + assert calls == 1, "the waiting call must never reach the provider" + assert exhausted.reason == "wall_deadline" + assert isinstance(exhausted.last_exc, LLMDeadlineExceeded) + assert exhausted.last_exc.reason == exhausted.reason + assert "concurrency-gate" in str(exhausted.last_exc) + + +@pytest.mark.asyncio +async def test_ordinary_exhaustion_still_carries_the_real_failure() -> None: + """With no deadline configured, ``exhausted`` keeps pointing at the 429. + + The fix must not push deadline semantics onto the ordinary path: here + ``last_exc`` IS the cause of the raise, and a chain wrapper reading + ``rate_limited`` off it is doing the right thing. + """ + async def _chat(_messages, **_kw): + raise _RateLimit429 + + with pytest.raises(LLMCallExhausted) as exc_info: + await call_llm( + SimpleNamespace(chat=_chat, model="fake"), + [user_msg("hello")], + timeout=30, + max_retries=1, + turn=1, + ) + + assert exc_info.value.reason == "exhausted" + assert isinstance(exc_info.value.last_exc, _RateLimit429) + assert exc_info.value.prior_exc is None diff --git a/tests/test_llm_runtime_extract_usage.py b/tests/test_llm_runtime_extract_usage.py new file mode 100644 index 0000000..4f81620 --- /dev/null +++ b/tests/test_llm_runtime_extract_usage.py @@ -0,0 +1,605 @@ +"""``extract_usage`` normalizes both usage shapes to OpenAI raw form. + +The downstream observers (``ProtocolStreamObserver`` for SSE §4.4, +``WorkerTraceFileObserver`` for the trace file §6.0) read ``ctx.usage`` +with keys ``provider`` / ``model`` / ``prompt_tokens`` / +``completion_tokens`` / ``cache_read_tokens`` / ``cache_write_tokens`` / +``cached_tokens`` / ``cache_creation_tokens`` / ``reasoning_tokens``. +``provider`` comes from ``response_metadata.provider_actually_used`` +(stamped by a provider-chain wrapper per attempt) — empty string when no +chain wrapper stamped it. Both the legacy canonical ``usage_metadata`` +(``input_tokens`` / ``output_tokens``, no ``model``) and OpenAI raw +``response_metadata.token_usage`` must fold to that target shape. + +2026-05-28 cache schema split: + +The schema previously exposed only ``cached_tokens`` (cache READ only) +and ``cache_creation_tokens`` (cache WRITE). Cost boards using a single +``cached_tokens`` field with a single rate would under-attribute +Anthropic cache write spend (~12× rate difference between read and +write). Post-split the schema carries four fields: + +- ``cache_read_tokens`` — explicit cache READ count +- ``cache_write_tokens`` — explicit cache WRITE count (incl. Anthropic + 1h-TTL extension summed into write) +- ``cached_tokens`` — backward-compat sum (``read + write``); + **semantic change** from pre-split (was read-only) +- ``cache_creation_tokens`` — backward-compat alias of + ``cache_write_tokens`` (deprecated; prefer the new name) +""" +from __future__ import annotations + +from types import SimpleNamespace + +from agent_core.llm import LLMResponse +from agent_core.runtime.loop._response import _pick_int +from agent_core.runtime.loop.llm_client import extract_usage + + +def _resp(*, usage_metadata=None, response_metadata=None): + return SimpleNamespace( + usage_metadata=usage_metadata, + response_metadata=response_metadata, + ) + + +def test_returns_none_when_no_usage_anywhere() -> None: + assert extract_usage(_resp()) is None + assert extract_usage(_resp(response_metadata={})) is None + # Both empty / 0 → also None (no signal to record). + assert extract_usage(_resp(usage_metadata={"input_tokens": 0, + "output_tokens": 0})) is None + + +def test_native_response_zero_fills_reasoning_tokens() -> None: + usage = extract_usage(LLMResponse( + model="native-model", + usage={"prompt_tokens": 7, "completion_tokens": 3}, + )) + + assert usage == { + "provider": "", + "model": "native-model", + "prompt_tokens": 7, + "completion_tokens": 3, + "cache_read_tokens": 0, + "cache_write_tokens": 0, + "cached_tokens": 0, + "cache_creation_tokens": 0, + "reasoning_tokens": 0, + } + + +def test_canonical_usage_metadata_shape() -> None: + # Canonical input/output token fields; model lives in response_metadata. + resp = _resp( + usage_metadata={ + "input_tokens": 120, + "output_tokens": 45, + "input_token_details": {"cache_read": 30}, + }, + response_metadata={"model_name": "qwen35_397b_a17b"}, + ) + u = extract_usage(resp) + assert u == { + "model": "qwen35_397b_a17b", + "prompt_tokens": 120, + "completion_tokens": 45, + "cache_read_tokens": 30, + "cache_write_tokens": 0, + "cached_tokens": 30, # = read + write = 30 + 0 + "cache_creation_tokens": 0, # alias of cache_write_tokens + "reasoning_tokens": 0, + "provider": "", + } + + +def test_openai_raw_token_usage_shape() -> None: + # response_metadata.token_usage with prompt_tokens / completion_tokens. + resp = _resp( + usage_metadata=None, + response_metadata={ + "model_name": "gpt-4o", + "token_usage": { + "prompt_tokens": 200, + "completion_tokens": 60, + "prompt_tokens_details": {"cached_tokens": 50}, + }, + }, + ) + u = extract_usage(resp) + assert u == { + "model": "gpt-4o", + "prompt_tokens": 200, + "completion_tokens": 60, + "cache_read_tokens": 50, + "cache_write_tokens": 0, + "cached_tokens": 50, # = read + write + "cache_creation_tokens": 0, + "reasoning_tokens": 0, + "provider": "", + } + + +def test_usage_metadata_takes_precedence_over_token_usage() -> None: + """When both shapes are present, prefer the canonical metadata one.""" + resp = _resp( + usage_metadata={"input_tokens": 1, "output_tokens": 2}, + response_metadata={ + "model_name": "m", + "token_usage": {"prompt_tokens": 999, "completion_tokens": 999}, + }, + ) + u = extract_usage(resp) + assert u == { + "model": "m", + "prompt_tokens": 1, + "completion_tokens": 2, + "cache_read_tokens": 0, + "cache_write_tokens": 0, + "cached_tokens": 0, + "cache_creation_tokens": 0, + "reasoning_tokens": 0, + "provider": "", + } + + +def test_apodex_nonstream_cached_falls_through_to_raw_shape() -> None: + """Apodex on ``ainvoke`` fills the canonical ``input_tokens`` / + ``output_tokens`` but leaves ``input_token_details`` empty — cached + counts only land on the raw ``prompt_tokens_details.cached_tokens``. + Without the raw-shape fallback, cached tokens get silently dropped + on every non-streaming call (DAG analyzer, decision_llm, synth_llm). + """ + resp = _resp( + usage_metadata={"input_tokens": 5000, "output_tokens": 1000}, + response_metadata={ + "model_name": "mirothinker_v20_397b", + "provider_actually_used": "apodex", + "token_usage": { + "prompt_tokens": 5000, + "completion_tokens": 1000, + "prompt_tokens_details": {"cached_tokens": 2000}, + }, + }, + ) + u = extract_usage(resp) + assert u is not None + assert u["prompt_tokens"] == 5000 + assert u["completion_tokens"] == 1000 + assert u["cached_tokens"] == 2000 + assert u["provider"] == "apodex" + + +def test_response_metadata_alt_keys() -> None: + """``response_metadata.usage`` is a valid alias for ``token_usage``; + ``model`` is a valid alias for ``model_name``.""" + resp = _resp( + usage_metadata=None, + response_metadata={ + "model": "alt-name", + "usage": {"prompt_tokens": 10, "completion_tokens": 5}, + }, + ) + assert extract_usage(resp) == { + "model": "alt-name", + "prompt_tokens": 10, + "completion_tokens": 5, + "cache_read_tokens": 0, + "cache_write_tokens": 0, + "cached_tokens": 0, + "cache_creation_tokens": 0, + "reasoning_tokens": 0, + "provider": "", + } + + +def test_canonical_cache_creation_and_reasoning_tokens() -> None: + """Canonical cache creation and reasoning token detail fields.""" + resp = _resp( + usage_metadata={ + "input_tokens": 500, + "output_tokens": 200, + "input_token_details": {"cache_read": 50, "cache_creation": 400}, + "output_token_details": {"reasoning": 120}, + }, + response_metadata={"model_name": "claude-opus-4-7"}, + ) + u = extract_usage(resp) + assert u == { + "model": "claude-opus-4-7", + "prompt_tokens": 500, + "completion_tokens": 200, + "cache_read_tokens": 50, + "cache_write_tokens": 400, + "cached_tokens": 450, # = 50 read + 400 write + "cache_creation_tokens": 400, # alias of cache_write_tokens + "reasoning_tokens": 120, + "provider": "", + } + + +def test_openai_raw_cache_creation_and_reasoning_tokens() -> None: + """Anthropic raw exposes ``cache_creation_input_tokens`` directly on + the usage dict; OpenAI o-series nests reasoning under + ``completion_tokens_details.reasoning_tokens``.""" + resp = _resp( + usage_metadata=None, + response_metadata={ + "model_name": "o1-preview", + "token_usage": { + "prompt_tokens": 1000, + "completion_tokens": 500, + "prompt_tokens_details": {"cached_tokens": 200}, + "completion_tokens_details": {"reasoning_tokens": 300}, + "cache_creation_input_tokens": 600, + }, + }, + ) + u = extract_usage(resp) + assert u == { + "model": "o1-preview", + "prompt_tokens": 1000, + "completion_tokens": 500, + "cache_read_tokens": 200, + "cache_write_tokens": 600, + "cached_tokens": 800, # = 200 read + 600 write + "cache_creation_tokens": 600, + "reasoning_tokens": 300, + "provider": "", + } + + +def test_openrouter_cache_write_tokens_alias() -> None: + """Claude routed through OpenRouter's OpenAI-compatible wrap exposes + cache-creation tokens under ``prompt_tokens_details.cache_write_tokens`` + instead of the OpenAI-standard ``cache_creation_tokens``. Observed live + on ``api.miromind.site/v1 → openrouter`` for ``claude-sonnet-4.6``. + The fallback chain must pick this alias up so cache_create lands in the + rollup when prompt caching is eventually enabled.""" + resp = _resp( + usage_metadata=None, + response_metadata={ + "model_name": "anthropic/claude-sonnet-4.6", + "token_usage": { + "prompt_tokens": 1000, + "completion_tokens": 50, + "prompt_tokens_details": { + "cached_tokens": 120, + "cache_write_tokens": 480, + }, + }, + }, + ) + u = extract_usage(resp) + assert u == { + "model": "anthropic/claude-sonnet-4.6", + "prompt_tokens": 1000, + "completion_tokens": 50, + "cache_read_tokens": 120, + "cache_write_tokens": 480, + "cached_tokens": 600, # = 120 read + 480 write + "cache_creation_tokens": 480, + "reasoning_tokens": 0, + "provider": "", + } + + +def test_qwen_null_cached_tokens_normalises_to_zero() -> None: + """Qwen3.5-397b via apodex gateway returns + ``prompt_tokens_details.cached_tokens = null`` (the gateway does not + surface prefix-cache info even when the SGLang/vLLM backend has it). + ``dict.get`` returns None for that key, the ``or 0`` chain must + normalise it without crashing on ``int(None)``.""" + resp = _resp( + usage_metadata=None, + response_metadata={ + "model_name": "qwen3.5-397b-a17b", + "token_usage": { + "prompt_tokens": 1616, + "completion_tokens": 505, + "prompt_tokens_details": { + "audio_tokens": None, + "cached_tokens": None, + "text_tokens": 1616, + }, + "completion_tokens_details": { + "reasoning_tokens": 499, + }, + }, + }, + ) + u = extract_usage(resp) + assert u == { + "model": "qwen3.5-397b-a17b", + "prompt_tokens": 1616, + "completion_tokens": 505, + "cache_read_tokens": 0, + "cache_write_tokens": 0, + "cached_tokens": 0, + "cache_creation_tokens": 0, + "reasoning_tokens": 499, + "provider": "", + } + + +def test_provider_actually_used_stamped_by_chain() -> None: + """When a provider-chain wrapper writes ``provider_actually_used``, + extract_usage surfaces it as ``provider`` + so downstream billing can attribute the call per vendor.""" + resp = _resp( + usage_metadata={"input_tokens": 100, "output_tokens": 20}, + response_metadata={ + "model_name": "claude-opus-4-7", + "provider_actually_used": "anthropic", + "model_actually_used": "claude-opus-4-7", + "fallback_used": 1, + }, + ) + u = extract_usage(resp) + assert u is not None + assert u["provider"] == "anthropic" + assert u["model"] == "claude-opus-4-7" + + +def test_anthropic_direct_cache_read_input_tokens() -> None: + """Anthropic Messages API exposes ``cache_read_input_tokens`` at the + ``usage`` root (not nested under ``prompt_tokens_details`` the way + OpenAI does). When Claude is hit *directly* (not via openrouter, + which strips the standard names — see + ``test_openrouter_cache_write_tokens_alias``) the read count lands + there. Symmetric with ``cache_creation_input_tokens`` which we + already accept at the same level.""" + resp = _resp( + usage_metadata=None, + response_metadata={ + "model_name": "claude-sonnet-4-6", + "token_usage": { + "prompt_tokens": 1500, + "completion_tokens": 200, + "cache_read_input_tokens": 800, + "cache_creation_input_tokens": 450, + }, + }, + ) + assert extract_usage(resp) == { + "model": "claude-sonnet-4-6", + "prompt_tokens": 1500, + "completion_tokens": 200, + "cache_read_tokens": 800, + "cache_write_tokens": 450, + "cached_tokens": 1250, # = 800 read + 450 write + "cache_creation_tokens": 450, + "reasoning_tokens": 0, + "provider": "", + } + + +def test_apodex_qwen_nests_cache_creation_under_prompt_tokens_details() -> None: + """Apodex / qwen3.5 gateway shape: BOTH cache fields nested under + ``prompt_tokens_details`` instead of split (read nested, write at + root) as Anthropic-direct does. + + Observed shape (2026-05) — note the write key is the Anthropic name + ``cache_creation_input_tokens`` BUT lives under ``prompt_tokens_details``, + not at the usage root:: + + usage: + prompt_tokens: 12000 + completion_tokens: 500 + prompt_tokens_details: + cached_tokens: 8000 # cache read + cache_creation_input_tokens: 3500 # cache write + + Without the ``ptd.cache_creation_input_tokens`` fallback the write + count would silently drop on every apodex non-streaming call — same + class of regression as the apodex cached-tokens cross-check fix + that already exists for ``cached_tokens``. + """ + resp = _resp( + usage_metadata=None, + response_metadata={ + "model_name": "qwen3.5-397b-a17b", + "provider_actually_used": "apodex", + "token_usage": { + "prompt_tokens": 12000, + "completion_tokens": 500, + "prompt_tokens_details": { + "cached_tokens": 8000, + "cache_creation_input_tokens": 3500, + }, + }, + }, + ) + u = extract_usage(resp) + assert u is not None + assert u["cache_read_tokens"] == 8000 + assert u["cache_write_tokens"] == 3500, ( + "apodex/qwen nest cache_creation_input_tokens under " + "prompt_tokens_details — the extractor must accept this alias" + ) + assert u["cached_tokens"] == 11500 # 8000 read + 3500 write + + +def test_bedrock_via_openai_gateway_nests_anthropic_cache_read_under_ptd() -> None: + """Bedrock served through an OpenAI-compatible gateway (provider_class + ``OpenAIClient``) nests BOTH Anthropic cache keys under + ``prompt_tokens_details`` — the WRITE key + ``cache_creation_input_tokens`` AND the READ key + ``cache_read_input_tokens``. + + Regression for 2026-05-29: the read candidate list lacked the nested + ``cache_read_input_tokens`` (while the write list had its nested + counterpart), so every bedrock-via-gateway call recorded + ``cache_read_tokens == 0`` even on a warm second call. The reported + symptom was ``read=0 / write=24163`` in ``usage_summary``. + """ + body = { + "model_name": "global.anthropic.claude-sonnet-4-6", + "provider_actually_used": "bedrock", + "token_usage": { + "prompt_tokens": 31538, + "completion_tokens": 1993, + "prompt_tokens_details": { + "cache_read_input_tokens": 24000, + "cache_creation_input_tokens": 24163, + }, + }, + } + # Raw-only shape (streaming usage chunk: no usage_metadata). + u = extract_usage(_resp(usage_metadata=None, response_metadata=body)) + assert u is not None + assert u["cache_read_tokens"] == 24000, ( + "nested Anthropic cache_read_input_tokens must be picked up — " + "mirror of the write path's nested candidate" + ) + assert u["cache_write_tokens"] == 24163 + + # Canonical-metadata-present shape (ainvoke: usage_metadata populated + # but input_token_details empty — the cross-check must still find the + # nested read key). + u2 = extract_usage(_resp( + usage_metadata={"input_tokens": 31538, "output_tokens": 1993}, + response_metadata=body, + )) + assert u2 is not None + assert u2["cache_read_tokens"] == 24000 + assert u2["cache_write_tokens"] == 24163 + + +def test_apodex_canonical_path_finds_nested_cache_creation() -> None: + """A gateway shape with canonical metadata populated but + empty details — fallback to ptd.cache_creation_input_tokens MUST + still pick up the write count via the cross-check that the canonical + path runs when its own details came back empty. + """ + resp = _resp( + usage_metadata={ + "input_tokens": 12000, + "output_tokens": 500, + # Canonical details intentionally empty — apodex gateway + # populates only the raw shape's nested details. + }, + response_metadata={ + "model_name": "qwen3.5-397b-a17b", + "provider_actually_used": "apodex", + "token_usage": { + "prompt_tokens": 12000, + "completion_tokens": 500, + "prompt_tokens_details": { + "cached_tokens": 8000, + "cache_creation_input_tokens": 3500, + }, + }, + }, + ) + u = extract_usage(resp) + assert u is not None + assert u["cache_read_tokens"] == 8000 + assert u["cache_write_tokens"] == 3500 + + +def test_anthropic_1h_ttl_extension_sums_into_cache_write() -> None: + """Anthropic's 1h-TTL prompt-cache surfaces under a nested + ``cache_creation`` dict alongside the standard 5m count. + + Both bill as cache WRITE (5m ~1.25× base, 1h ~2× base) — sum them + into ``cache_write_tokens`` so the schema field captures the full + write footprint. Boards needing per-TTL breakdown should consume + the raw provider response directly; the schema picks the practical + "all writes" rollup for the same reason ``cached_tokens`` is now + "all cache activity". + """ + resp = _resp( + usage_metadata=None, + response_metadata={ + "model_name": "claude-opus-4-7", + "token_usage": { + "prompt_tokens": 5000, + "completion_tokens": 800, + "cache_read_input_tokens": 1200, + "cache_creation_input_tokens": 600, # 5m TTL + "cache_creation": { + "ephemeral_1h_input_tokens": 400, # 1h TTL + }, + }, + }, + ) + u = extract_usage(resp) + assert u is not None + assert u["cache_read_tokens"] == 1200 + assert u["cache_write_tokens"] == 600 + 400 # 5m + 1h summed + assert u["cached_tokens"] == 1200 + 600 + 400 # read + write sum + assert u["cache_creation_tokens"] == 1000 # alias of cache_write_tokens + + +def test_cached_tokens_backward_compat_is_sum_not_read_only() -> None: + """**Schema semantic change (2026-05-28)**: ``cached_tokens`` is now + the sum of cache read + cache write, not cache-read-only. + + Pre-split this field exposed only cache-read counts; cost boards + using a single rate on ``cached_tokens`` would have under-attributed + Anthropic write spend by ~12× (write bills at ~1.25× base vs + read's ~0.1×). Post-split the same name returns the intuitive + total so boards using the shortcut over-attribute (preferable to + under-attributing). + + This test exists explicitly to lock in the new semantics — if a + future change reverts ``cached_tokens`` to read-only, this test + blocks it. + """ + resp = _resp( + usage_metadata=None, + response_metadata={ + "model_name": "claude-sonnet-4-6", + "token_usage": { + "prompt_tokens": 10_000, + "completion_tokens": 2_000, + "cache_read_input_tokens": 3_000, + "cache_creation_input_tokens": 2_500, + }, + }, + ) + u = extract_usage(resp) + assert u is not None + assert u["cached_tokens"] == 5_500, ( + "cached_tokens must equal cache_read + cache_write post-split" + ) + # Explicit split fields are the same numbers from the raw response. + assert u["cache_read_tokens"] == 3_000 + assert u["cache_write_tokens"] == 2_500 + + +def test_pick_int_skips_none_and_invalid_falls_through_to_real_signal() -> None: + """``_pick_int`` powers the multi-source token extraction. It must + skip ``None`` (gateway-side ``null``), unparseable values, and + zeros — only stopping on the first non-zero int. This is the + contract that lets us list provider field aliases in priority + order without each one having to special-case ``null``.""" + assert _pick_int(None, None, 42) == 42 + assert _pick_int(0, 0, 100) == 100 + assert _pick_int("not-an-int", None, 7) == 7 + assert _pick_int(None) == 0 + assert _pick_int() == 0 + # Numeric string is coerced — providers occasionally JSON-decode + # token counts as strings on edge cases. + assert _pick_int("123") == 123 + + +def test_usage_metadata_object_coerces_to_dict() -> None: + """A compatibility adapter may return an object for ``usage_metadata``; + coerce it via ``dict(...)`` rather than requiring a concrete dict.""" + class UM: + def keys(self): + return ("input_tokens", "output_tokens") + + def __getitem__(self, k): + return {"input_tokens": 7, "output_tokens": 3}[k] + + resp = _resp(usage_metadata=UM(), + response_metadata={"model_name": "x"}) + u = extract_usage(resp) + assert u is not None + assert u["prompt_tokens"] == 7 + assert u["completion_tokens"] == 3 + assert u["model"] == "x" diff --git a/tests/test_llm_runtime_hooks.py b/tests/test_llm_runtime_hooks.py new file mode 100644 index 0000000..84cedbc --- /dev/null +++ b/tests/test_llm_runtime_hooks.py @@ -0,0 +1,85 @@ +from __future__ import annotations + +import asyncio + +import pytest + +import agent_core.runtime.loop._call as call_module +from agent_core.errors import ( + AgentCoreError, + LLMCallExhausted, + LLMStreamStalled, +) +from agent_core.llm import LLMResponse +from agent_core.messages import user_msg +from agent_core.runtime.llm_request_overrides import ( + ThinkingRetryOverride, + current_thinking_retry_override, + thinking_retry_override, +) +from agent_core.runtime.loop._bind import bind_session_id +from agent_core.runtime.loop.llm_client import call_llm + + +class FakeLLM: + model = "fake" + + def __init__(self) -> None: + self.calls = 0 + + async def chat(self, messages: object, **kwargs: object) -> LLMResponse: + self.calls += 1 + return LLMResponse(content="ok") + + +def test_llm_errors_share_the_agent_core_root_and_timeout_contract() -> None: + stalled = LLMStreamStalled(stall_s=3, chunks_seen=2, elapsed_s=7) + + assert isinstance(stalled, AgentCoreError) + assert isinstance(stalled, TimeoutError) + exhausted = LLMCallExhausted(stalled, "exhausted") + assert exhausted.last_exc is stalled + assert exhausted.reason == "exhausted" + + +def test_bind_session_id_uses_explicit_product_kill_switch() -> None: + llm = FakeLLM() + + assert bind_session_id(llm, "task-1", sticky_session_enabled=lambda: False) is llm + bound = bind_session_id(llm, "task-1", sticky_session_enabled=lambda: True) + assert bound.extra_headers == {"x-upstream-session-id": "task-1"} + + +@pytest.mark.asyncio +async def test_wall_deadline_hook_refuses_before_provider_call(monkeypatch) -> None: + monkeypatch.setattr(call_module, "_WALL_DEADLINE_FLOOR_S", 20.0) + llm = FakeLLM() + + with pytest.raises(LLMCallExhausted) as exc_info: + await call_llm( + llm, + [user_msg("hello")], + timeout=120, + max_retries=2, + turn=1, + wall_deadline_remaining=lambda: 5.0, + ) + + assert exc_info.value.reason == "wall_deadline" + assert llm.calls == 0 + + +@pytest.mark.asyncio +async def test_thinking_override_is_task_local_and_restored() -> None: + async def observe(mode: str) -> str: + with thinking_retry_override(ThinkingRetryOverride(mode=mode)): + await asyncio.sleep(0) + current = current_thinking_retry_override() + assert current is not None + return current.mode + + assert await asyncio.gather(observe("reduced"), observe("disabled")) == [ + "reduced", + "disabled", + ] + assert current_thinking_retry_override() is None diff --git a/tests/test_llm_runtime_nameless_tool_calls.py b/tests/test_llm_runtime_nameless_tool_calls.py new file mode 100644 index 0000000..446d8a1 --- /dev/null +++ b/tests/test_llm_runtime_nameless_tool_calls.py @@ -0,0 +1,117 @@ +"""A streamed tool-call that never received a name must not reach history. + +``_stream_llm_response`` creates an accumulator slot for ANY tool-call delta +carrying an index, pre-filled with ``name=""``. Two real cases leave it empty: +a content-free chunk that merely opens a tool-call block, and a stream cut +before the name delta arrives. + +Emitting such a call is worse than dropping it. It cannot be executed, and the +loop records it in durable history as ``name=""`` — after which some chat +templates fail to render that history at all (``can only concatenate str (not +"NoneType") to str``, HTTP 400 from apodex-1.1). Every later request in the +session replays the same history and is rejected identically, so one malformed +delta ends the run, and ends it looking like an ordinary empty submission +rather than the infrastructure fault it is. +""" +from __future__ import annotations + +import logging + +from agent_core.llm import StreamDelta +from agent_core.runtime.loop._streaming import _stream_llm_response + + +class _FakeLLM: + """Yields a fixed StreamDelta script, like a provider adapter would.""" + + def __init__(self, deltas: list[StreamDelta]) -> None: + self._deltas = deltas + + async def stream(self, messages, timeout=None): + for delta in self._deltas: + yield delta + + +async def _noop(*args, **kwargs) -> None: + return None + + +async def _run(deltas: list[StreamDelta]): + return await _stream_llm_response( + _FakeLLM(deltas), messages=[], timeout=30, on_delta=_noop, + ) + + +async def test_index_only_chunk_does_not_become_a_tool_call() -> None: + """The content-free chunk that opens a tool-call block carries no name.""" + resp = await _run([ + StreamDelta(tool_call_deltas=[{"index": 0}]), + StreamDelta(content="thinking about it"), + ]) + assert resp.tool_calls == [] + assert resp.content == "thinking about it" + + +async def test_stream_cut_before_the_name_arrives_is_dropped() -> None: + """Arguments streamed, then the stream ended — the call is unexecutable.""" + resp = await _run([ + StreamDelta(tool_call_deltas=[{"index": 0, "id": "call_1"}]), + StreamDelta(tool_call_deltas=[{"index": 0, "arguments": '{"command":'}]), + ]) + assert resp.tool_calls == [] + + +async def test_a_named_call_still_survives() -> None: + resp = await _run([ + StreamDelta(tool_call_deltas=[ + {"index": 0, "id": "call_1", "name": "bash"}, + ]), + StreamDelta(tool_call_deltas=[ + {"index": 0, "arguments": '{"command": "ls"}'}, + ]), + ]) + assert len(resp.tool_calls) == 1 + assert resp.tool_calls[0]["function"]["name"] == "bash" + assert resp.tool_calls[0]["function"]["arguments"] == '{"command": "ls"}' + assert resp.tool_calls[0]["id"] == "call_1" + + +async def test_a_named_call_to_a_nonexistent_tool_is_kept() -> None: + """Deliberate: it reaches the executor and comes back as "unknown tool", + which the model can read and correct. Only NAMELESS calls are dropped.""" + resp = await _run([ + StreamDelta(tool_call_deltas=[ + {"index": 0, "id": "c1", "name": "no_such_tool", "arguments": "{}"}, + ]), + ]) + assert [tc["function"]["name"] for tc in resp.tool_calls] == ["no_such_tool"] + + +async def test_only_the_nameless_slot_is_dropped_from_a_mixed_turn() -> None: + """A real parallel-tool turn must not lose its good calls to a bad sibling.""" + resp = await _run([ + StreamDelta(tool_call_deltas=[ + {"index": 0, "id": "c0", "name": "bash", "arguments": "{}"}, + {"index": 1}, + {"index": 2, "id": "c2", "name": "read_file", "arguments": "{}"}, + ]), + ]) + assert [tc["function"]["name"] for tc in resp.tool_calls] == ["bash", "read_file"] + + +async def test_the_drop_is_logged_not_silent(caplog) -> None: + """A provider emitting these consistently is an upstream defect, and this + is the only place left that can see it.""" + with caplog.at_level(logging.WARNING): + await _run([StreamDelta(tool_call_deltas=[{"index": 0}])]) + assert any("no function name" in r.getMessage() for r in caplog.records) + assert any("dropped 1 " in r.getMessage() for r in caplog.records) + + +async def test_a_clean_turn_logs_nothing(caplog) -> None: + with caplog.at_level(logging.WARNING): + resp = await _run([ + StreamDelta(tool_call_deltas=[{"index": 0, "id": "c", "name": "bash"}]), + ]) + assert len(resp.tool_calls) == 1 + assert not [r for r in caplog.records if "function name" in r.getMessage()] diff --git a/tests/test_llm_runtime_proxy_wrap.py b/tests/test_llm_runtime_proxy_wrap.py new file mode 100644 index 0000000..ef5c28b --- /dev/null +++ b/tests/test_llm_runtime_proxy_wrap.py @@ -0,0 +1,205 @@ +"""Regression for the proxy-wrapped-400 transient handling. + +OpenAI-compatible proxies (new-api, etc.) sometimes package an upstream +5xx or timeout as a 400 envelope. The literal HTTP status is 400 but the +semantics are transient — sleeping + retrying the same key fixes it. + +Until this fix, ``call_llm`` treated every 400 as non-transient and +returned ``None`` immediately, surfacing the failure to the caller as +"sub-agent's LLM call failed; report is partial." This was observed in a +multi-turn smoke against an OpenAI-compatible proxy. + +The fix delegates the classification to +``agent_core.runtime.retriable.is_transient_network``, the shared predicate +used by this call site and product provider-chain wrappers. +""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +import pytest + +from agent_core.llm import LLMResponse +from agent_core.messages import user_msg +from agent_core.runtime.loop.llm_client import LLMCallExhausted, call_llm + + +class _ProxyWrapped400(Exception): + """Mimics the shape new-api / OpenAI-proxy gateways surface. + + Carries ``status_code=400`` (literal HTTP status from the proxy) but + the body string flags it as the proxy's own wrap of an upstream + failure. + """ + + def __init__(self) -> None: + super().__init__( + "Error code: 400 - {'error': {'message': '(request id: foo)', " + "'type': 'new_api_error', 'code': 'bad_response_status_code'}}", + ) + self.status_code = 400 + + +class _Vanilla400(Exception): + """Genuine bad-request: schema validation, malformed JSON, etc. + + These ARE non-transient — must continue to return ``None`` without + retry.""" + + def __init__(self) -> None: + super().__init__( + "Error code: 400 - invalid_request_error: schema validation failed", + ) + self.status_code = 400 + + +@pytest.mark.asyncio +async def test_proxy_wrapped_400_is_retried(monkeypatch): + """A 400 carrying ``bad_response_status_code`` must trigger backoff + + retry, NOT immediate ``return None``.""" + sleeps: list[float] = [] + real_sleep = asyncio.sleep + + async def _record(duration): + sleeps.append(duration) + await real_sleep(0) + + monkeypatch.setattr("asyncio.sleep", _record) + + call_count = 0 + + async def _chat(_messages, **_kw): + nonlocal call_count + call_count += 1 + if call_count == 1: + raise _ProxyWrapped400() + return LLMResponse(content="recovered") + + fake_llm = SimpleNamespace(chat=_chat) + + result = await call_llm( + fake_llm, + [user_msg("hi")], + timeout=10, + max_retries=3, + turn=0, + ) + + assert result is not None + assert result.content == "recovered" + assert call_count == 2 # initial fail → 1 backoff → success + backoffs = [s for s in sleeps if s > 0] + assert backoffs, f"expected a backoff sleep on proxy 400, got {sleeps!r}" + + +@pytest.mark.asyncio +async def test_vanilla_400_surfaces_as_LLMCallExhausted(monkeypatch): + """Schema-level 400s remain non-transient — must surface immediately + so a chain wrapper can decide whether to advance to the next leg + (different provider may have stricter / looser schema) instead of + burning retries on a deterministic failure.""" + sleeps: list[float] = [] + real_sleep = asyncio.sleep + + async def _record(duration): + sleeps.append(duration) + await real_sleep(0) + + monkeypatch.setattr("asyncio.sleep", _record) + + call_count = 0 + + async def _chat(_messages, **_kw): + nonlocal call_count + call_count += 1 + raise _Vanilla400() + + fake_llm = SimpleNamespace(chat=_chat) + + with pytest.raises(LLMCallExhausted) as exc_info: + await call_llm( + fake_llm, + [user_msg("hi")], + timeout=10, + max_retries=3, + turn=0, + ) + + assert exc_info.value.reason == "non_transient" + assert isinstance(exc_info.value.last_exc, _Vanilla400) + # No retry — call_count is 1 and no real backoff was scheduled. + assert call_count == 1 + backoffs = [s for s in sleeps if s > 0] + assert not backoffs, ( + f"vanilla 400 should not trigger backoff, got {backoffs!r}" + ) + + +@pytest.mark.asyncio +async def test_401_403_404_classification(monkeypatch): + """The 400-transient escape hatch is narrowly scoped to 400 + body + pattern. 401 / 403 / 404 each surface immediately without same-key + retry, but split on whether a chain advance might recover: + + - **401** → ``chain_advance``. The auth_failure predicate matches + (``\\b401\\b``) and ``is_retriable_with_fallback`` includes it + since ``501e2c645`` — the next chain leg uses a different key + (often a different provider), so retrying *there* can succeed + where same-key retry cannot. + - **403** → ``non_transient``. Explicitly excluded from + ``is_auth_failure`` because 403 means "authenticated but not + authorised", typically a scoping / region / model-access + problem that recurs on sibling providers — surface to the + operator instead of silently advancing. + - **404** → ``non_transient``. A bare "Error code: 404" carries + no ``model_not_found`` / ``no_such_model`` signal, so the + chain-advance predicate doesn't match — falls through to the + generic 400-family non-transient branch. + + Either way ``call_count == 1`` — neither classification retries on + the same key. + """ + real_sleep = asyncio.sleep + + async def _no_op(_duration): + await real_sleep(0) + + monkeypatch.setattr("asyncio.sleep", _no_op) + + cases = [ + (401, "chain_advance"), + (403, "non_transient"), + (404, "non_transient"), + ] + for status, expected_reason in cases: + call_count = 0 + + class _AuthError(Exception): + pass + + err = _AuthError(f"Error code: {status}") + err.status_code = status # type: ignore[attr-defined] + + async def _chat(_messages, _e=err, **_kw): + nonlocal call_count + call_count += 1 + raise _e + + fake_llm = SimpleNamespace(chat=_chat) + with pytest.raises(LLMCallExhausted) as exc_info: + await call_llm( + fake_llm, + [user_msg("hi")], + timeout=10, + max_retries=3, + turn=0, + ) + assert exc_info.value.reason == expected_reason, ( + f"status {status} should be classified as {expected_reason}, " + f"got {exc_info.value.reason}" + ) + assert call_count == 1, ( + f"status {status} should not retry; got {call_count} calls" + ) diff --git a/tests/test_llm_runtime_retry_after.py b/tests/test_llm_runtime_retry_after.py new file mode 100644 index 0000000..c7ca5f8 --- /dev/null +++ b/tests/test_llm_runtime_retry_after.py @@ -0,0 +1,193 @@ +"""Regression for P2-2 / Step 8: clamp ``Retry-After`` at the same 300s +ceiling as the exponential fallback so a misbehaving provider returning +``Retry-After: 86400`` cannot stall the loop with no output. +""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +import pytest + +from agent_core.llm import LLMResponse +from agent_core.messages import user_msg +from agent_core.runtime.loop.llm_client import LLMCallExhausted, call_llm + + +class _RateLimit429(Exception): + """Mimics the shape ``_get_retry_after`` looks for: ``response.headers``.""" + + def __init__(self, retry_after: float) -> None: + super().__init__("rate limited") + self.status_code = 429 + self.response = SimpleNamespace( + headers={"retry-after": str(retry_after)}, + ) + + +@pytest.mark.asyncio +async def test_retry_after_huge_value_is_clamped(monkeypatch): + """``Retry-After: 86400`` must not sleep the loop for a day. + + Without the clamp, a single hostile or buggy upstream response would + silently freeze the agent loop for ``Retry-After`` seconds — the + exact "no error, no progress" pattern the swarm hang audit (P2-2) + flagged. + """ + sleeps: list[float] = [] + + real_sleep = asyncio.sleep + + async def _record(duration): + sleeps.append(duration) + # Yield the loop without burning real time so the test stays + # deterministic; ``asyncio.wait_for`` and friends only need a + # zero-duration yield, not the requested duration. + await real_sleep(0) + + monkeypatch.setattr("asyncio.sleep", _record) + + call_count = 0 + + async def _chat(_messages, **_kw): + nonlocal call_count + call_count += 1 + if call_count == 1: + raise _RateLimit429(retry_after=86400) # one full day + return LLMResponse(content="ok") + + fake_llm = SimpleNamespace(chat=_chat) + + result = await call_llm( + fake_llm, + [user_msg("hi")], + timeout=10, + max_retries=3, + turn=0, + ) + + assert result is not None and result.content == "ok" + # The retry path slept once between attempt 1 (429) and attempt 2. + # Filter out 0-duration yields the test injects via ``real_sleep(0)``. + backoffs = [s for s in sleeps if s > 0] + assert backoffs, f"expected a backoff sleep, got {sleeps!r}" + assert all(s <= 300 for s in backoffs), ( + f"backoff exceeded 300s ceiling: {backoffs!r}" + ) + + +@pytest.mark.asyncio +async def test_retry_after_within_ceiling_unchanged(monkeypatch): + """``Retry-After`` values ≤300 are honoured verbatim — clamp is a + ceiling, not a floor. + """ + sleeps: list[float] = [] + real_sleep = asyncio.sleep + + async def _record(duration): + sleeps.append(duration) + await real_sleep(0) + + monkeypatch.setattr("asyncio.sleep", _record) + + call_count = 0 + + async def _chat(_messages, **_kw): + nonlocal call_count + call_count += 1 + if call_count == 1: + raise _RateLimit429(retry_after=42) + return LLMResponse(content="ok") + + fake_llm = SimpleNamespace(chat=_chat) + + result = await call_llm( + fake_llm, + [user_msg("hi")], + timeout=10, + max_retries=3, + turn=0, + ) + assert result is not None + backoffs = [s for s in sleeps if s > 0] + assert 42 in backoffs, ( + f"expected literal Retry-After=42 in backoffs, got {backoffs!r}" + ) + + +@pytest.mark.asyncio +async def test_default_backoff_is_exponential_on_timeout(monkeypatch): + """``retry_wait_fixed=None`` (default) → timeouts back off on the + 2/4/8s exponential base with ±25% jitter (jitter de-synchronises + retry stampedes across parallel runs — partial3 saw timed-out + attempts re-collide 20 minutes later without it). + + Regression: a prior refactor flattened the default to 30s flat which + over-paused short transient errors. Pin the schedule shape so other + workflows keep mirothinker's non-flat behavior. + """ + sleeps: list[float] = [] + real_sleep = asyncio.sleep + + async def _record(duration): + sleeps.append(duration) + await real_sleep(0) + + monkeypatch.setattr("asyncio.sleep", _record) + + async def _chat_timeout(_messages, **_kw): + raise TimeoutError("simulated stream timeout") + + fake_llm = SimpleNamespace(chat=_chat_timeout) + + with pytest.raises(LLMCallExhausted) as exc_info: + await call_llm( + fake_llm, + [user_msg("hi")], + timeout=10, + max_retries=4, + turn=0, + ) + assert exc_info.value.reason == "exhausted" + backoffs = [s for s in sleeps if s > 0] + # Base min(2 * 2**attempt, 60) for attempts 0..2 (last attempt has + # no sleep), ±25% jitter → assert each within its jitter band. + assert len(backoffs) == 3, f"expected 3 backoffs, got {backoffs!r}" + for got, base in zip(backoffs, [2, 4, 8], strict=False): + assert base * 0.75 <= got <= base * 1.25, ( + f"backoff {got:.2f}s outside jitter band of base {base}s " + f"(full schedule: {backoffs!r})" + ) + + +@pytest.mark.asyncio +async def test_retry_wait_fixed_overrides_default(monkeypatch): + """Workflows that opt into ``retry_wait_fixed=N`` get N-second flat waits.""" + sleeps: list[float] = [] + real_sleep = asyncio.sleep + + async def _record(duration): + sleeps.append(duration) + await real_sleep(0) + + monkeypatch.setattr("asyncio.sleep", _record) + + async def _chat_timeout(_messages, **_kw): + raise TimeoutError() + + fake_llm = SimpleNamespace(chat=_chat_timeout) + + with pytest.raises(LLMCallExhausted): + await call_llm( + fake_llm, + [user_msg("hi")], + timeout=10, + max_retries=3, + turn=0, + retry_wait_fixed=90, + ) + backoffs = [int(s) for s in sleeps if s > 0] + assert backoffs == [90, 90], ( + f"expected fixed 90s waits, got {backoffs!r}" + ) diff --git a/tests/test_llm_runtime_stall.py b/tests/test_llm_runtime_stall.py new file mode 100644 index 0000000..e96f084 --- /dev/null +++ b/tests/test_llm_runtime_stall.py @@ -0,0 +1,323 @@ +"""Stream-stall watchdog in ``call_llm`` / ``_stream_llm_response``. + +2026-06-05 heavy-trace forensics (partial3.json): 12 streaming attempts +against the apodex gateway went silent mid-flight — no chunks, no +error — and each pinned its call for the full 1200 s timeout, driving +4/8 runs into their wall deadline. Decode throughput on *successful* +calls was a steady 110-120 tok/s (a full-cap generation finishes in +~155 s), so a silent stream is a dead stream, not a slow one. + +The watchdog bounds the gap between consecutive stream chunks +(``MIROHARNESS_LLM_STREAM_STALL_S``, default 180 s, <= 0 disables) and +raises :class:`LLMStreamStalled` — a ``TimeoutError`` subclass — so the +existing transient retry/backoff budget handles the redo. +""" + +from __future__ import annotations + +import asyncio + +import pytest + +import agent_core.runtime.loop._streaming as llm_client +from agent_core.llm import StreamDelta +from agent_core.messages import user_msg +from agent_core.runtime.loop.llm_client import ( + LLMCallExhausted, + LLMStreamStalled, + call_llm, +) + +MSGS = [user_msg("q")] + + +async def _sink_delta(*_a, **_kw) -> None: + """Minimal on_delta — presence enables the streaming path.""" + + +class _StallingLLM: + """Streams ``preamble`` chunks then hangs forever (silent stream). + + ``hang_attempts`` controls how many attempts hang before a healthy + one — models the partial3 pattern where a retry after the stall + succeeds at full speed. + """ + + def __init__(self, hang_attempts: int = 99, preamble: int = 1) -> None: + self.hang_attempts = hang_attempts + self.preamble = preamble + self.calls = 0 + + async def stream(self, messages, **_kw): + self.calls += 1 + for i in range(self.preamble): + yield StreamDelta(content=f"chunk{i} ") + if self.calls <= self.hang_attempts: + await asyncio.sleep(3600) # gateway black-hole: no chunks, no error + else: + yield StreamDelta(content="recovered answer") + + +class _HealthyGappyLLM: + """Streams with inter-chunk gaps below the stall threshold.""" + + async def stream(self, messages, **_kw): + for i in range(3): + await asyncio.sleep(0.02) + yield StreamDelta(content=f"part{i} ") + + +@pytest.mark.asyncio +async def test_stall_aborts_attempt_and_exhausts(monkeypatch) -> None: + """A permanently silent stream dies at the stall bound (not the full + call timeout) on every attempt, then surfaces as ``exhausted`` with + the stall as ``last_exc`` — the chain wrapper sees WHY it died.""" + monkeypatch.setenv("MIROHARNESS_LLM_STREAM_STALL_S", "0.1") + llm = _StallingLLM() + with pytest.raises(LLMCallExhausted) as exc_info: + await call_llm( + llm, MSGS, timeout=30, max_retries=2, turn=1, + on_delta=_sink_delta, retry_wait_fixed=0, + ) + assert exc_info.value.reason == "exhausted" + assert isinstance(exc_info.value.last_exc, LLMStreamStalled) + assert exc_info.value.last_exc.chunks_seen == 1 + assert llm.calls == 2 # both attempts ran and stalled + + +@pytest.mark.asyncio +async def test_stall_chain_advances_after_threshold(monkeypatch) -> None: + """Under an active provider chain, repeated stalls stop same-key + retrying and surface ``chain_advance`` at the threshold (default 2) + instead of burning the whole retry budget on a black-holed endpoint. + fail1: 56 stalls × ~180-330 s wasted on one dead 397b gateway.""" + monkeypatch.setenv("MIROHARNESS_LLM_STREAM_STALL_S", "0.1") + llm = _StallingLLM() # hangs forever + with pytest.raises(LLMCallExhausted) as exc_info: + await call_llm( + llm, MSGS, timeout=30, max_retries=5, turn=1, + on_delta=_sink_delta, retry_wait_fixed=0, + chain_fallback_active=lambda: True, + ) + assert exc_info.value.reason == "chain_advance" + assert isinstance(exc_info.value.last_exc, LLMStreamStalled) + # Advanced at the 2nd stall — did NOT consume the full 5-attempt budget. + assert llm.calls == 2 + + +@pytest.mark.asyncio +async def test_stall_no_chain_retries_to_budget(monkeypatch) -> None: + """With no outer chain there's nothing to advance to, so stalls keep + same-key retrying to the budget and surface ``exhausted`` (the + single-endpoint swarm/benchmark floor is unchanged).""" + monkeypatch.setenv("MIROHARNESS_LLM_STREAM_STALL_S", "0.1") + llm = _StallingLLM() + with pytest.raises(LLMCallExhausted) as exc_info: + await call_llm( + llm, MSGS, timeout=30, max_retries=4, turn=1, + on_delta=_sink_delta, retry_wait_fixed=0, + ) + assert exc_info.value.reason == "exhausted" + assert llm.calls == 4 # all attempts ran, no early chain advance + + +@pytest.mark.asyncio +async def test_stall_chain_advance_disabled_by_env(monkeypatch) -> None: + """``MIROHARNESS_LLM_STREAM_STALL_MAX=0`` disables the escape even + under a chain: stalls retry to the budget then surface ``exhausted``.""" + monkeypatch.setenv("MIROHARNESS_LLM_STREAM_STALL_S", "0.1") + monkeypatch.setenv("MIROHARNESS_LLM_STREAM_STALL_MAX", "0") + llm = _StallingLLM() + with pytest.raises(LLMCallExhausted) as exc_info: + await call_llm( + llm, MSGS, timeout=30, max_retries=3, turn=1, + on_delta=_sink_delta, retry_wait_fixed=0, + chain_fallback_active=lambda: True, + ) + assert exc_info.value.reason == "exhausted" + assert llm.calls == 3 + + +@pytest.mark.asyncio +async def test_stall_then_retry_recovers(monkeypatch) -> None: + """First attempt stalls, second streams normally — the retry budget + converts a 1200 s black-hole into one stall-bound redo.""" + monkeypatch.setenv("MIROHARNESS_LLM_STREAM_STALL_S", "0.1") + llm = _StallingLLM(hang_attempts=1) + response = await call_llm( + llm, MSGS, timeout=30, max_retries=3, turn=1, + on_delta=_sink_delta, retry_wait_fixed=0, + ) + assert response is not None + assert "recovered answer" in response.content + assert llm.calls == 2 + + +@pytest.mark.asyncio +async def test_slow_but_alive_stream_not_flagged(monkeypatch) -> None: + """Inter-chunk gaps below the bound never trip the watchdog — slow + decode is healthy, only silence is pathological.""" + monkeypatch.setenv("MIROHARNESS_LLM_STREAM_STALL_S", "0.5") + response = await call_llm( + _HealthyGappyLLM(), MSGS, timeout=30, max_retries=1, turn=1, + on_delta=_sink_delta, retry_wait_fixed=0, + ) + assert response is not None + assert response.content == "part0 part1 part2 " + + +@pytest.mark.asyncio +async def test_stall_disabled_falls_back_to_total_timeout(monkeypatch) -> None: + """``MIROHARNESS_LLM_STREAM_STALL_S=0`` disables the watchdog: the + silent stream then dies at the (plain) total call timeout.""" + monkeypatch.setenv("MIROHARNESS_LLM_STREAM_STALL_S", "0") + llm = _StallingLLM() + with pytest.raises(LLMCallExhausted) as exc_info: + await call_llm( + llm, MSGS, timeout=1, max_retries=1, turn=1, + on_delta=_sink_delta, retry_wait_fixed=0, + ) + assert isinstance(exc_info.value.last_exc, asyncio.TimeoutError) + assert not isinstance(exc_info.value.last_exc, LLMStreamStalled) + + +@pytest.mark.asyncio +async def test_stall_default_and_invalid_env(monkeypatch) -> None: + """No env → documented default; junk env → default (with a warning), + never a crash in the hot path.""" + monkeypatch.delenv("MIROHARNESS_LLM_STREAM_STALL_S", raising=False) + assert llm_client._stream_stall_timeout_s() == 180.0 + monkeypatch.setenv("MIROHARNESS_LLM_STREAM_STALL_S", "not-a-number") + assert llm_client._stream_stall_timeout_s() == 180.0 + + +@pytest.mark.asyncio +async def test_first_chunk_bound_fires_before_stall(monkeypatch) -> None: + """With the TTFT knob set, a stream that never produces its FIRST + chunk dies at the tight first-chunk bound — not the loose stall + bound (healthy TTFT is seconds; zero chunks after the bound means + black-holed, not thinking).""" + monkeypatch.setenv("MIROHARNESS_LLM_STREAM_STALL_S", "30") + monkeypatch.setenv("MIROHARNESS_LLM_FIRST_CHUNK_S", "0.1") + llm = _StallingLLM(preamble=0) # hangs before any chunk + loop = asyncio.get_running_loop() + started = loop.time() + with pytest.raises(LLMCallExhausted) as exc_info: + await call_llm( + llm, MSGS, timeout=600, max_retries=1, turn=1, + on_delta=_sink_delta, retry_wait_fixed=0, + ) + assert isinstance(exc_info.value.last_exc, LLMStreamStalled) + assert exc_info.value.last_exc.chunks_seen == 0 + assert exc_info.value.last_exc.stall_s == pytest.approx(0.1) + assert loop.time() - started < 5, "must die at TTFT bound, not stall/timeout" + + +@pytest.mark.asyncio +async def test_first_chunk_bound_relaxes_after_first_chunk(monkeypatch) -> None: + """Once the first chunk arrives, inter-chunk gaps are judged by the + looser stall bound — gaps longer than the TTFT bound but below the + stall bound are healthy (mid-generation pauses are legitimate).""" + monkeypatch.setenv("MIROHARNESS_LLM_STREAM_STALL_S", "0.5") + monkeypatch.setenv("MIROHARNESS_LLM_FIRST_CHUNK_S", "0.05") + + class _SlowAfterFirst: + async def stream(self, messages, **_kw): + yield StreamDelta(content="fast ") # TTFT well under 0.05 + await asyncio.sleep(0.2) # > first_s, < stall_s — must NOT flag + yield StreamDelta(content="slow") + + response = await call_llm( + _SlowAfterFirst(), MSGS, timeout=30, max_retries=1, turn=1, + on_delta=_sink_delta, retry_wait_fixed=0, + ) + assert response is not None + assert response.content == "fast slow" + + +@pytest.mark.asyncio +async def test_first_chunk_only_mode_disarms_after_first(monkeypatch) -> None: + """TTFT bound with the stall watchdog disabled: the scope disarms + after the first chunk, so a later silence falls through to the + plain total call timeout (not LLMStreamStalled).""" + monkeypatch.setenv("MIROHARNESS_LLM_STREAM_STALL_S", "0") + monkeypatch.setenv("MIROHARNESS_LLM_FIRST_CHUNK_S", "0.05") + llm = _StallingLLM() # 1 preamble chunk, then silent forever + with pytest.raises(LLMCallExhausted) as exc_info: + await call_llm( + llm, MSGS, timeout=1, max_retries=1, turn=1, + on_delta=_sink_delta, retry_wait_fixed=0, + ) + assert isinstance(exc_info.value.last_exc, asyncio.TimeoutError) + assert not isinstance(exc_info.value.last_exc, LLMStreamStalled) + + +@pytest.mark.asyncio +async def test_first_chunk_per_call_param_wins_over_env(monkeypatch) -> None: + """``call_llm(first_chunk_s=...)`` (LoopConfig ← profile + ``agent.first_chunk_s``) overrides the process-wide env knob; an + explicit 0 disables even when the env arms it.""" + monkeypatch.setenv("MIROHARNESS_LLM_STREAM_STALL_S", "30") + monkeypatch.setenv("MIROHARNESS_LLM_FIRST_CHUNK_S", "30") + # Param tightens past the env: dies at 0.1, not 30. + llm = _StallingLLM(preamble=0) + loop = asyncio.get_running_loop() + started = loop.time() + with pytest.raises(LLMCallExhausted) as exc_info: + await call_llm( + llm, MSGS, timeout=600, max_retries=1, turn=1, + on_delta=_sink_delta, retry_wait_fixed=0, first_chunk_s=0.1, + ) + assert isinstance(exc_info.value.last_exc, LLMStreamStalled) + assert loop.time() - started < 5 + + # Param 0 disables the TTFT bound even though the env sets 30: the + # first chunk is then judged by the stall bound (0.2 here). + monkeypatch.setenv("MIROHARNESS_LLM_STREAM_STALL_S", "0.2") + started = loop.time() + with pytest.raises(LLMCallExhausted) as exc_info: + await call_llm( + _StallingLLM(preamble=0), MSGS, timeout=600, max_retries=1, + turn=1, on_delta=_sink_delta, retry_wait_fixed=0, + first_chunk_s=0, + ) + assert isinstance(exc_info.value.last_exc, LLMStreamStalled) + assert exc_info.value.last_exc.stall_s == pytest.approx(0.2) + + +@pytest.mark.asyncio +async def test_first_chunk_default_off_and_invalid_env(monkeypatch) -> None: + """No env → disabled (first chunk rides the stall bound); junk env + → disabled, never a crash in the hot path.""" + monkeypatch.delenv("MIROHARNESS_LLM_FIRST_CHUNK_S", raising=False) + assert llm_client._first_chunk_timeout_s() == 0.0 + monkeypatch.setenv("MIROHARNESS_LLM_FIRST_CHUNK_S", "not-a-number") + assert llm_client._first_chunk_timeout_s() == 0.0 + + +@pytest.mark.asyncio +async def test_stall_closes_stream_promptly(monkeypatch) -> None: + """The stall path acloses the chunk generator before raising so the + underlying HTTP stream is released immediately (the httpcore-leak + family from the mm1 incident), and the whole failure takes ~stall_s, + not ~timeout.""" + monkeypatch.setenv("MIROHARNESS_LLM_STREAM_STALL_S", "0.1") + closed = asyncio.Event() + + class _TrackingLLM(_StallingLLM): + async def stream(self, messages, **_kw): + try: + yield StreamDelta(content="x") + await asyncio.sleep(3600) + finally: + closed.set() + + loop = asyncio.get_running_loop() + started = loop.time() + with pytest.raises(LLMCallExhausted): + await call_llm( + _TrackingLLM(), MSGS, timeout=600, max_retries=1, turn=1, + on_delta=_sink_delta, retry_wait_fixed=0, + ) + assert closed.is_set(), "generator finally must run (stream released)" + assert loop.time() - started < 5, "must fail at stall bound, not timeout" diff --git a/tests/test_llm_runtime_stream_metadata.py b/tests/test_llm_runtime_stream_metadata.py new file mode 100644 index 0000000..6b47b03 --- /dev/null +++ b/tests/test_llm_runtime_stream_metadata.py @@ -0,0 +1,306 @@ +"""Streaming assembly invariants in ``_stream_llm_response``. + +History: the legacy chunk merger combined +``response_metadata`` with ``merge_dicts``, which **concatenated** string +values across chunks. Correct for tokenized content (one stream → one +string) but wrong for header-like fields an OpenAI-compatible proxy +repeats on every chunk (``model_name`` / ``system_fingerprint`` / +``finish_reason``) — ``final.usage.model_usage`` keys came out doubled +(``…claude-4.6-sonnet…claude-4.6-sonnet…``) on new-api streaming runs. + +The native streaming path removes the failure mode by construction: the +client adapter normalises each chunk to a :class:`StreamDelta`, which +carries only ``content`` / ``reasoning_content`` / ``tool_call_deltas`` +— no ``response_metadata`` to merge, so there is nothing to double. +``_stream_llm_response`` folds the deltas into an ``LLMResponse`` whose +``content`` is the clean concatenation. ``response_metadata`` stays empty +UNLESS a serving-leg wrapper stamps ``StreamDelta.provider``, which the +assembler folds into +``response_metadata["provider_actually_used"]`` so streamed calls carry +billing attribution (the doubling bug still cannot recur — only one +known key is written, never a per-chunk dict merge). ``usage`` / +``finish_reason`` / ``model`` are likewise carried through from the +provider's terminal chunks when present (``StreamDelta`` gained those +fields) — previously they were dropped, zeroing streaming usage/billing. +These tests pin that contract. +""" + +from __future__ import annotations + +import asyncio + +import pytest + +from agent_core.llm import StreamDelta +from agent_core.messages import user_msg +from agent_core.runtime.loop.llm_client import ( + call_llm, + extract_usage, +) + + +class _StubStreamingLLM: + """Fake client whose ``stream`` yields a fixed list of + ``StreamDelta``s — the native shape the OpenAIClient adapter emits + per SSE chunk.""" + + def __init__(self, *, deltas: list[StreamDelta]) -> None: + self.deltas = deltas + + async def stream(self, _messages, **_kw): + for delta in self.deltas: + yield delta + + +@pytest.mark.asyncio +async def test_streaming_content_concatenated_no_metadata_doubling(monkeypatch): + """N content deltas fold into a single clean ``LLMResponse``: content + is concatenated (the merge behaviour we WANT), and because + ``StreamDelta`` carries no ``response_metadata``, the assembled + response has empty metadata — the model-name doubling bug cannot + recur. ``extract_usage`` returns ``None`` since the streamed response + carries no usage (the non-streaming path supplies usage instead).""" + real_sleep = asyncio.sleep + + async def _noop(_): + await real_sleep(0) + + monkeypatch.setattr("asyncio.sleep", _noop) + + deltas = [ + StreamDelta(content="hello "), + StreamDelta(content="world"), + ] + llm = _StubStreamingLLM(deltas=deltas) + + delta_log: list[tuple[str, str, int]] = [] + + async def _on_delta(delta, accumulated, idx, thinking=""): + delta_log.append((delta, accumulated, idx)) + + result = await call_llm( + llm, [user_msg("hi")], + timeout=10, max_retries=1, turn=0, + on_delta=_on_delta, + ) + + assert result is not None + # Content concatenation (the merge behaviour we WANT for body text) + # is preserved — exactly once, not doubled. + assert "hello world" in str(result.content) + # No response_metadata is carried on the streamed response, so there + # is no ``model_name`` to concatenate / double. + assert result.response_metadata == {} + # These deltas carry no terminal usage chunk, so usage stays empty here + # and extract_usage reports None. (A usage-bearing stream is covered by + # test_streaming_carries_terminal_usage_and_finish_reason below.) + assert extract_usage(result) is None + # on_delta saw the body text deltas in order. + assert [d for d, _a, _i in delta_log] == ["hello ", "world"] + + +@pytest.mark.asyncio +async def test_streaming_carries_terminal_usage_and_finish_reason(monkeypatch): + """Regression (P1): the terminal ``include_usage`` chunk (empty choices, + usage set) and the last content chunk's ``finish_reason`` must reach the + assembled ``LLMResponse``. Previously ``OpenAIClient.stream`` dropped the + empty-choices usage chunk and ``StreamDelta`` had no usage/finish field, + so streaming runs reported 0 usage and observers never saw + ``finish_reason='length'`` (truncation / salvage / rollback).""" + real_sleep = asyncio.sleep + + async def _noop(_): + await real_sleep(0) + + monkeypatch.setattr("asyncio.sleep", _noop) + + deltas = [ + StreamDelta(content="hi", finish_reason="stop", model="m-1"), + # Terminal include_usage chunk: empty content, carries final usage. + StreamDelta( + usage={ + "prompt_tokens": 12, "completion_tokens": 3, + "total_tokens": 15, "cached_tokens": 4, + }, + model="m-1", + ), + ] + llm = _StubStreamingLLM(deltas=deltas) + + async def _on_delta(*_): + pass + + result = await call_llm( + llm, [user_msg("hi")], + timeout=10, max_retries=1, turn=0, + on_delta=_on_delta, + ) + + assert result is not None + assert str(result.content) == "hi" + assert result.finish_reason == "stop" + assert result.model == "m-1" + assert result.usage == { + "prompt_tokens": 12, "completion_tokens": 3, + "total_tokens": 15, "cached_tokens": 4, + } + # extract_usage now surfaces the streamed usage (was None before the fix). + u = extract_usage(result) + assert u is not None + assert u["prompt_tokens"] == 12 + assert u["completion_tokens"] == 3 + assert u["cached_tokens"] == 4 + + +@pytest.mark.asyncio +async def test_streaming_folds_provider_actually_used(monkeypatch): + """A serving-leg-stamped ``StreamDelta.provider`` is folded into the + assembled response's ``response_metadata`` so streamed calls bill against + the right vendor — the gap that split mirothinker into a ``provider=""`` + bucket and an ``@apodex`` bucket. ``extract_usage`` then surfaces it.""" + real_sleep = asyncio.sleep + + async def _noop(_): + await real_sleep(0) + + monkeypatch.setattr("asyncio.sleep", _noop) + + deltas = [ + StreamDelta(content="hi", provider="apodex"), + StreamDelta( + usage={"prompt_tokens": 7, "completion_tokens": 2, "cached_tokens": 0}, + model="mirothinker_v20_397b", + provider="apodex", + ), + ] + llm = _StubStreamingLLM(deltas=deltas) + + async def _on_delta(*_): + pass + + result = await call_llm( + llm, [user_msg("hi")], + timeout=10, max_retries=1, turn=0, + on_delta=_on_delta, + ) + + assert result is not None + assert result.response_metadata == {"provider_actually_used": "apodex"} + u = extract_usage(result) + assert u is not None + assert u["provider"] == "apodex" + + +@pytest.mark.asyncio +async def test_streaming_many_chunks_content_clean(monkeypatch): + """Three content deltas still concatenate cleanly into one body with + no per-chunk metadata bleed — pins that adding chunks never + re-introduces the doubling the legacy merge produced.""" + real_sleep = asyncio.sleep + + async def _noop(_): + await real_sleep(0) + + monkeypatch.setattr("asyncio.sleep", _noop) + + deltas = [ + StreamDelta(content="a"), + StreamDelta(content="b"), + StreamDelta(content="c"), + ] + llm = _StubStreamingLLM(deltas=deltas) + + async def _on_delta(*_): + pass + + result = await call_llm( + llm, [user_msg("hi")], + timeout=10, max_retries=1, turn=0, + on_delta=_on_delta, + ) + + assert result is not None + assert str(result.content) == "abc" + assert result.response_metadata == {} + + +@pytest.mark.asyncio +async def test_streaming_forwards_tool_call_chunks(monkeypatch): + real_sleep = asyncio.sleep + + async def _noop(_): + await real_sleep(0) + + monkeypatch.setattr("asyncio.sleep", _noop) + + # Native StreamDelta tool-call deltas use the OpenAIClient adapter + # shape: {index, id, name, arguments}. + deltas = [ + StreamDelta( + tool_call_deltas=[ + { + "index": 0, + "id": "call_1", + "name": "finalize_answer", + "arguments": '{"content":"hel', + }, + ], + ), + StreamDelta( + tool_call_deltas=[ + { + "index": 0, + "id": "call_1", + "name": None, + "arguments": 'lo"}', + }, + ], + ), + ] + llm = _StubStreamingLLM(deltas=deltas) + seen: list[list[dict]] = [] + + async def _on_delta( + _delta, + _accumulated, + _idx, + _thinking="", + *, + tool_call_args_chunks=None, + ): + seen.append(tool_call_args_chunks or []) + + result = await call_llm( + llm, [user_msg("hi")], + timeout=10, max_retries=1, turn=0, + on_delta=_on_delta, + ) + + assert result is not None + # ``_stream_llm_response`` forwards per-chunk arg deltas in the + # {name, args, id, index} shape observers decode progressively. + assert seen == [ + [{ + "name": "finalize_answer", + "args": '{"content":"hel', + "id": "call_1", + "index": 0, + }], + [{ + "name": None, + "args": 'lo"}', + "id": "call_1", + "index": 0, + }], + ] + # The stitched tool call assembled across both chunks. + assert result.tool_calls == [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "finalize_answer", + "arguments": '{"content":"hello"}', + }, + }, + ] diff --git a/tests/test_llm_runtime_strip_thinking.py b/tests/test_llm_runtime_strip_thinking.py new file mode 100644 index 0000000..145b400 --- /dev/null +++ b/tests/test_llm_runtime_strip_thinking.py @@ -0,0 +1,65 @@ +"""``_strip_thinking_blocks`` removes ```` artifacts from final answers. + +The kernel's ``extract_final_content`` calls this helper before handing +content to downstream judges / scorers. Three artifact shapes appear in +the wild: + +1. **Proper pair** ``...answer`` — vanilla case. +2. **Dangling opener** ``...EOF`` — model truncated mid-thought. +3. **Orphan closer** ``thinking textanswer`` — SGLang ``enable_thinking + + preserve_thinking`` chat-template sometimes drops the opening tag. + Until 2026-05-09 this case was mishandled (only the 8-char tag was + removed; the thinking trace before it leaked into the answer and + confused the judge). Confirmed against trial 65 of an apodex BC-200 + run where the model's "Wait—wait no!" indecision rambling reached the + judge as the final answer. +""" +from __future__ import annotations + +from agent_core.runtime.loop._response import _strip_thinking_blocks + + +def test_strips_proper_pair() -> None: + assert _strip_thinking_blocks( + "reasoning\nFinal answer: 42" + ) == "Final answer: 42" + + +def test_strips_dangling_opener() -> None: + assert _strip_thinking_blocks( + "answer\ncut off mid-thought" + ) == "answer" + + +def test_strips_orphan_closer_sglang_quirk() -> None: + """SGLang preserve_thinking sometimes emits with no opener.""" + raw = "lots of thinking text including doubts\n**Final: 42**" + assert _strip_thinking_blocks(raw) == "**Final: 42**" + + +def test_strips_orphan_closer_keeps_only_tail_after_last() -> None: + """When multiple stray exist, keep only what's after the last.""" + raw = "thought onesecond pass thinking\nanswer" + assert _strip_thinking_blocks(raw) == "answer" + + +def test_no_think_passthrough() -> None: + assert _strip_thinking_blocks("plain answer") == "plain answer" + + +def test_strips_pair_then_orphan_closer() -> None: + """Mixed: a proper pair earlier and a stray closer later.""" + raw = "pair\ntext
more thinking\nfinal" + # _THINK_BLOCK_RE removes the pair → "text
more thinking\nfinal" + # Then orphan-closer logic keeps after last → "final" + assert _strip_thinking_blocks(raw) == "final" + + +def test_empty_string() -> None: + assert _strip_thinking_blocks("") == "" + + +def test_all_thinking_no_answer() -> None: + """Worst case: model emitted only thinking, never an answer.""" + assert _strip_thinking_blocks("just reasoning") == "" + assert _strip_thinking_blocks("just reasoning") == "" diff --git a/tests/test_retriable.py b/tests/test_retriable.py new file mode 100644 index 0000000..11f9bdb --- /dev/null +++ b/tests/test_retriable.py @@ -0,0 +1,599 @@ +"""Unit tests for the reporter-fallback error classifier. + +Pure-function tests on ``Exception`` instances. Mirrors the patterns +provider gateway's prod sees: Anthropic 529 overloaded_error, OpenAI 503 +overloads, ``insufficient_quota`` from billing-exhausted keys, 429 +rate-limits, and structural errors (which must NOT trigger fallback). +""" + +from __future__ import annotations + +import pytest + +from agent_core.errors import LLMDeadlineExceeded +from agent_core.runtime.retriable import ( + classify_error, + is_credit_exhausted, + is_empty_completion, + is_overloaded_error, + is_rate_limited, + is_retriable_with_fallback, + is_stream_stall, + is_transient_network, +) + + +@pytest.mark.parametrize("reason", ["wall_deadline", "logical_call_deadline"]) +def test_runtime_deadline_is_not_a_transient_provider_timeout(reason: str) -> None: + err = LLMDeadlineExceeded(reason, "budget consumed") + + assert classify_error(err) == reason + assert not is_transient_network(err) + assert not is_retriable_with_fallback(err) + +# ── is_overloaded_error ────────────────────────────────────────────── + + +@pytest.mark.parametrize("msg", [ + "Error code: 529 - Anthropic API is overloaded", + "overloaded_error: capacity exceeded", + "anthropic.OverloadedError: Overloaded", + "OpenAI 503: model is overloaded, please retry", + "Provider returned capacity exhausted", +]) +def test_overload_signatures_detected(msg: str) -> None: + assert is_overloaded_error(RuntimeError(msg)) + + +@pytest.mark.parametrize("msg", [ + "invalid_request_error: messages.0.content must be a string", + "rate_limit_error: too many requests", # different signal — 429 + "schema validation failed", + "json.JSONDecodeError: Expecting value", + "service_unavailable", # deliberately NOT matched (proxy noise) +]) +def test_non_overload_signatures_not_matched(msg: str) -> None: + assert not is_overloaded_error(RuntimeError(msg)) + + +def test_overload_detected_from_status_code_attribute() -> None: + """SDK exceptions often carry ``status_code`` separately from str().""" + class _FakeAnthropicError(Exception): + status_code = 529 + + def __str__(self) -> str: + return "request failed" + + assert is_overloaded_error(_FakeAnthropicError()) + + +# ── is_credit_exhausted ────────────────────────────────────────────── + + +@pytest.mark.parametrize("msg", [ + "insufficient_quota: your credit is exhausted", + "billing_error: account balance is insufficient", + "402 Payment Required", + "Insufficient Balance", +]) +def test_credit_signatures_detected(msg: str) -> None: + assert is_credit_exhausted(RuntimeError(msg)) + + +def test_credit_not_confused_with_overload() -> None: + assert not is_credit_exhausted(RuntimeError("overloaded_error")) + + +# ── is_rate_limited ────────────────────────────────────────────────── + + +@pytest.mark.parametrize("msg", [ + "rate_limit_error: too many requests", + "Error code: 429 - rate limit exceeded", + "RateLimitError: tokens per minute exceeded", +]) +def test_rate_limit_signatures_detected(msg: str) -> None: + assert is_rate_limited(RuntimeError(msg)) + + +# ── is_retriable_with_fallback (union) ─────────────────────────────── + + +def test_overload_is_retriable() -> None: + assert is_retriable_with_fallback(RuntimeError("529 overloaded")) + + +def test_credit_is_retriable() -> None: + assert is_retriable_with_fallback(RuntimeError("insufficient_quota")) + + +def test_rate_limit_is_retriable() -> None: + # Per spec §5: rate_limit triggers backoff-same-key, NOT chain escalation. + assert not is_retriable_with_fallback(RuntimeError("429 rate_limit")) + + +def test_stream_stall_is_retriable_with_fallback() -> None: + err = RuntimeError( + "LLMStreamStalled: stream stalled: no chunks for 180s", + ) + assert is_stream_stall(err) + assert is_retriable_with_fallback(err) + assert classify_error(err) == "stream_stall" + + +def test_structural_error_not_retriable() -> None: + """Bad JSON / schema mismatch should never trigger fallback — + retrying with another key won't fix the request shape.""" + assert not is_retriable_with_fallback( + ValueError("Expected str, got int at field 'response'"), + ) + assert not is_retriable_with_fallback( + RuntimeError("invalid_request_error: schema validation"), + ) + + +# ── is_empty_completion (reasoning-runaway / no-content recovery) ──── + + +@pytest.mark.parametrize("msg", [ + "No generation chunks were returned", + "no generation chunk returned", + "No completion tokens returned", + "empty completion", +]) +def test_empty_completion_detected(msg: str) -> None: + assert is_empty_completion(RuntimeError(msg)) + + +def test_empty_completion_not_overmatched() -> None: + # A normal validation error must not look like an empty completion. + assert not is_empty_completion(ValueError("schema validation failed")) + assert not is_empty_completion(RuntimeError("generation succeeded")) + + +def test_empty_completion_is_retriable_with_fallback() -> None: + """A reasoning-runaway empty completion carrying the observed bare + ``ValueError('No generation chunks were returned')`` must advance + the chain instead of being fatal — otherwise one empty among a heavy + run's ~150 calls kills the whole run.""" + assert is_retriable_with_fallback( + ValueError("No generation chunks were returned"), + ) + + +# ── classify_error ─────────────────────────────────────────────────── + + +def test_classify_returns_specific_reason() -> None: + assert ( + classify_error(ValueError("No generation chunks were returned")) + == "empty_completion" + ) + assert classify_error(RuntimeError("529 overloaded")) == "overloaded" + assert classify_error(RuntimeError("insufficient_quota")) == "credit_exhausted" + assert classify_error(RuntimeError("429 too many requests")) == "rate_limited" + assert classify_error(RuntimeError("validation failed")) == "other" + + +def test_classify_precedence_overload_before_credit() -> None: + """If both signals are present, ``overloaded`` wins — it's the + more specific provider-state signal.""" + err = RuntimeError("529 overloaded; also insufficient_quota mentioned") + assert classify_error(err) == "overloaded" + + +# ── is_context_length_error (NEW per spec §5) ──────────────────────── + + +@pytest.mark.parametrize("msg", [ + "context length exceeded", + "context_length_exceeded", + "this model's maximum context length is 200000 tokens", + "input is longer than the model can handle", + "maximum context size reached", +]) +def test_context_length_signatures_detected(msg: str) -> None: + from agent_core.runtime.retriable import is_context_length_error + assert is_context_length_error(RuntimeError(msg)) + + +@pytest.mark.parametrize("msg", [ + "rate_limit_error: too many requests", + "overloaded_error: capacity exceeded", + "json parse failed", +]) +def test_context_length_not_misclassified(msg: str) -> None: + from agent_core.runtime.retriable import is_context_length_error + assert not is_context_length_error(RuntimeError(msg)) + + +# ── is_retriable_with_fallback no longer matches rate_limit ────────── + + +def test_rate_limit_is_NOT_retriable_with_fallback() -> None: + """Per spec §5 decision table: rate_limit means backoff, same key — + do NOT escalate the chain layer.""" + from agent_core.runtime.retriable import is_retriable_with_fallback + err = RuntimeError("rate_limit_error: too many requests") + assert not is_retriable_with_fallback(err) + + +def test_overload_still_retriable_with_fallback() -> None: + from agent_core.runtime.retriable import is_retriable_with_fallback + err = RuntimeError("anthropic.OverloadedError: Overloaded") + assert is_retriable_with_fallback(err) + + +def test_credit_exhausted_still_retriable_with_fallback() -> None: + from agent_core.runtime.retriable import is_retriable_with_fallback + err = RuntimeError("insufficient_quota") + assert is_retriable_with_fallback(err) + + +def test_classify_error_context_length() -> None: + """`classify_error` returns a new `context_length` label.""" + from agent_core.runtime.retriable import classify_error + err = RuntimeError("context_length_exceeded") + assert classify_error(err) == "context_length" + + +# ── is_transient_network ───────────────────────────────────────────── + + +@pytest.mark.parametrize("msg", [ + "Connection reset by peer", + "Connection refused", + "Connection aborted", + "Connection error: ECONNRESET", + "Connection closed unexpectedly", + "asyncio.TimeoutError: timed out after 60s", + "Read timeout", + "504 Gateway Timeout", + "Upstream error: provider unreachable", + "Upstream timeout while reading response", +]) +def test_transient_network_signatures_detected(msg: str) -> None: + from agent_core.runtime.retriable import is_transient_network + assert is_transient_network(RuntimeError(msg)) + + +def test_transient_network_matches_proxy_wrapped_400() -> None: + """The whole reason this predicate exists: an OpenAI-compatible proxy + (e.g. new-api gateway) packages an upstream 5xx or timeout into a 400 + envelope. The literal HTTP status is 400 but the semantics are + transient — sleeping + retrying the same key fixes it.""" + from agent_core.runtime.retriable import is_transient_network + + err = RuntimeError( + "Error code: 400 - {'error': {'message': '(request id: foo)', " + "'type': 'new_api_error', 'code': 'bad_response_status_code'}}", + ) + # Status attribute would mark this as 400, but the body pattern wins. + err.status_code = 400 # type: ignore[attr-defined] + assert is_transient_network(err) + + +def test_transient_network_matches_5xx_without_overload() -> None: + """A raw 502/503/504 with no overload wording is a transport problem, + not capacity exhaustion — classify as transient so the caller sleeps + the same key instead of burning fallback keys.""" + from agent_core.runtime.retriable import is_transient_network + + err = RuntimeError("502 Bad Gateway") + err.status_code = 502 # type: ignore[attr-defined] + assert is_transient_network(err) + + +def test_transient_network_yields_to_overload() -> None: + """When the same 503 has ``overloaded`` in the body, overload wins — + callers rotate keys instead of sleeping.""" + from agent_core.runtime.retriable import ( + is_overloaded_error, + is_transient_network, + ) + + err = RuntimeError("503 Service Unavailable: model is overloaded") + err.status_code = 503 # type: ignore[attr-defined] + assert is_overloaded_error(err) + assert not is_transient_network(err) + + +def test_transient_network_NOT_matched_by_plain_400() -> None: + """A vanilla 400 Bad Request (schema mismatch, bad JSON, etc.) is + structural — must NOT be classified as transient.""" + from agent_core.runtime.retriable import is_transient_network + + err = RuntimeError("Error code: 400 - invalid_request_error: bad JSON") + err.status_code = 400 # type: ignore[attr-defined] + assert not is_transient_network(err) + + +def test_transient_network_NOT_retriable_with_fallback() -> None: + """Per spec §5: transient_network triggers backoff-same-key, NOT + layer escalation. ``is_retriable_with_fallback`` (the chain-advance + trigger) must therefore stay False for transient errors.""" + from agent_core.runtime.retriable import ( + is_retriable_with_fallback, + is_transient_network, + ) + + err = RuntimeError("Connection reset") + assert is_transient_network(err) + assert not is_retriable_with_fallback(err) + + +def test_classify_error_returns_transient_network() -> None: + from agent_core.runtime.retriable import classify_error + err = RuntimeError("Connection reset by peer") + assert classify_error(err) == "transient_network" + + +def test_classify_precedence_overload_before_transient() -> None: + """503 + ``overloaded`` keeps the overloaded label, not transient.""" + from agent_core.runtime.retriable import classify_error + err = RuntimeError("503 model is overloaded") + err.status_code = 503 # type: ignore[attr-defined] + assert classify_error(err) == "overloaded" + + +# ── is_safety_filter (Aliyun DataInspectionFailed + friends) ───────── + + +@pytest.mark.parametrize("msg", [ + "<400> ***.***.DataInspectionFailed: Input text data may contain " + "inappropriate content. (request id: 202605121426525099564225sJgS6e6)", + "data_inspection_failed", + "content_policy_violation: prompt blocked", + "content_filter: rejected", + "content_filtered", + "input_filtered: anthropic policy", + "output_filtered", + "the request contains inappropriate content", + "prompt_blocked by safety system", + # gpt-5.x mid-stream refusal family (added 2026-05-29 for the + # reporter_v2 mid-stream continuation path): + "Invalid prompt: we've limited access to this content for safety " + "reasons. This type of information may be used to benefit or to " + "harm people...", + "Your request was blocked by our safety system", + "This violates our usage policies", + # OpenRouter content moderation blocks: + "Request blocked: content moderation policy", + "content moderation", + "blocked by safety system", + "moderation policy", + "Request blocked by safety system", +]) +def test_safety_filter_signatures_detected(msg: str) -> None: + from agent_core.runtime.retriable import is_safety_filter + assert is_safety_filter(RuntimeError(msg)), f"missed: {msg!r}" + + +@pytest.mark.parametrize("msg", [ + "rate_limit_error: too many requests", + "529 overloaded_error", + "insufficient_quota", + "context_length_exceeded", + "Connection reset by peer", + "validation failed", +]) +def test_safety_filter_not_misclassified(msg: str) -> None: + """Capacity/quota/network errors must NOT trip the safety predicate + — those route through their own branches with different recovery.""" + from agent_core.runtime.retriable import is_safety_filter + assert not is_safety_filter(RuntimeError(msg)) + + +def test_safety_filter_is_retriable_with_fallback() -> None: + """Per design: same-key retry on a deterministic safety rejection is + hopeless, so the chain must advance to a different provider.""" + from agent_core.runtime.retriable import is_retriable_with_fallback + err = RuntimeError( + "data_inspection_failed: Input text data may contain inappropriate content" + ) + assert is_retriable_with_fallback(err) + + +def test_classify_safety_filter_label() -> None: + from agent_core.runtime.retriable import classify_error + err = RuntimeError("data_inspection_failed") + assert classify_error(err) == "safety_filter" + + +def test_classify_precedence_safety_before_overload() -> None: + """If both signals appear in the same envelope, ``safety_filter`` + wins — operators need to distinguish 'model refused' from 'overload' + on the dashboard.""" + from agent_core.runtime.retriable import classify_error + err = RuntimeError( + "529 overloaded_error; also data_inspection_failed in body" + ) + assert classify_error(err) == "safety_filter" + + +# ── is_model_unavailable (distributor "no available channel" / 503 model_not_found) ── + + +@pytest.mark.parametrize("msg", [ + # The exact shape from the 2026-05-16 heavy-mode e2e against + # api.miromind.site: + "Error code: 503 - {'error': {'code': 'model_not_found', 'message': " + "'No available channel for model claude-sonnet-4-6 under group " + "openrouter (distributor) (request id: foo)', 'type': 'new_api_error'}}", + "model_not_found", + "no_such_model: gpt-foo", + "no available channel for model claude-bar under group anthropic", + "model_not_supported by this provider", + "unsupported_model: provider doesn't host this canonical name", + # OpenRouter shape — chaos scenario 02, 2026-05-21: + "openai.BadRequestError: Error code: 400 - {'error': {'message': " + "'does-not-exist-chaos-9999 is not a valid model ID', 'code': 400}}", + "invalid model ID", + "unknown model: gpt-foo-9999", +]) +def test_model_unavailable_signatures_detected(msg: str) -> None: + from agent_core.runtime.retriable import is_model_unavailable + assert is_model_unavailable(RuntimeError(msg)), f"missed: {msg!r}" + + +@pytest.mark.parametrize("msg", [ + "rate_limit_error: too many requests", + "529 overloaded_error", + "insufficient_quota", + "context_length_exceeded", + "Connection reset by peer", + "validation failed", + "Error code: 503 - service unavailable", # bare 503 without model phrase +]) +def test_model_unavailable_not_misclassified(msg: str) -> None: + """Other 5xx / structural errors must not trip model_unavailable — + they have their own recovery paths (backoff, salvage, raise).""" + from agent_core.runtime.retriable import is_model_unavailable + assert not is_model_unavailable(RuntimeError(msg)) + + +def test_model_unavailable_is_retriable_with_fallback() -> None: + """The whole reason this predicate exists: the next leg in the chain + may host the model — advance instead of same-key retry.""" + from agent_core.runtime.retriable import is_retriable_with_fallback + err = RuntimeError( + "Error code: 503 - model_not_found: No available channel for model X" + ) + assert is_retriable_with_fallback(err) + + +def test_model_unavailable_yields_transient_network() -> None: + """A 503 model_not_found is also a 5xx, but we want chain advance + rather than same-key backoff. ``is_transient_network`` must defer.""" + from agent_core.runtime.retriable import ( + is_model_unavailable, + is_transient_network, + ) + err = RuntimeError( + "Error code: 503 - model_not_found: No available channel" + ) + err.status_code = 503 # type: ignore[attr-defined] + assert is_model_unavailable(err) + assert not is_transient_network(err) + + +def test_classify_model_unavailable_label() -> None: + from agent_core.runtime.retriable import classify_error + err = RuntimeError("model_not_found") + assert classify_error(err) == "model_unavailable" + + +def test_classify_precedence_model_unavailable_before_overload() -> None: + """Operator wants 'fix the chain config' vs 'wait it out' distinction.""" + from agent_core.runtime.retriable import classify_error + err = RuntimeError("529 overloaded_error; also model_not_found in body") + assert classify_error(err) == "model_unavailable" + + +# ── is_auth_failure (chaos-discovered gap, 2026-05-21) ──────────────── +# Background: a key-clobber chaos scenario supplied invalid credentials and +# expected the chain to advance. +# Instead the run crashed exit=1 because the resulting +# openai.AuthenticationError ("Missing Authentication header", 401) +# wasn't in is_retriable_with_fallback. These tests pin the four wire +# shapes (SDK class name, OpenRouter body, invalid_api_key body, bare +# 401) and verify chain-advance + classify_error all agree. + + +@pytest.mark.parametrize("msg", [ + "openai.AuthenticationError: Error code: 401", + "Error code: 401 - {'error': {'message': 'Missing Authentication header', 'code': 401}}", + "invalid_api_key: the provided key is malformed", + "invalid_authentication: token rejected", + "401 Unauthorized", + "unauthenticated", + "authentication_failed", +]) +def test_auth_failure_signatures_detected(msg: str) -> None: + from agent_core.runtime.retriable import is_auth_failure + assert is_auth_failure(RuntimeError(msg)) + + +def test_auth_failure_detected_from_sdk_class_name() -> None: + """openai-python raises ``AuthenticationError`` as a class — observed + via ``type(err).__name__`` even when the body is sanitised.""" + from agent_core.runtime.retriable import is_auth_failure + + class AuthenticationError(Exception): + pass + + assert is_auth_failure(AuthenticationError("rejected")) + + +def test_auth_failure_detected_from_status_code() -> None: + """Status attribute fallback — wrappers sometimes strip the body + but keep the status.""" + from agent_core.runtime.retriable import is_auth_failure + + class _Err(Exception): + status_code = 401 + + def __str__(self) -> str: + return "request failed" + + assert is_auth_failure(_Err()) + + +@pytest.mark.parametrize("msg", [ + "rate_limit_error: too many requests", # 429, not 401 + "insufficient_quota", # billing, not auth + "model_not_found", + "internal server error", + "schema validation failed", +]) +def test_auth_failure_not_misclassified(msg: str) -> None: + from agent_core.runtime.retriable import is_auth_failure + assert not is_auth_failure(RuntimeError(msg)) + + +def test_403_forbidden_is_NOT_auth_failure() -> None: + """403 means "key authenticated but not authorised for this resource" + (e.g. account scoped to specific model families). The root often + repeats on sibling providers in the same chain, so advancing isn't + useful — surface to the operator instead.""" + from agent_core.runtime.retriable import is_auth_failure + + class _Err(Exception): + status_code = 403 + + def __str__(self) -> str: + return "403 Forbidden: insufficient permissions" + + assert not is_auth_failure(_Err()) + + +def test_auth_failure_is_retriable_with_fallback() -> None: + """The whole point of the chaos finding: 401 must trigger chain + advance so the next leg's (different key, often different provider) + can take over. Reporter L1 OpenRouter 401 → L2 anthropic-direct → + L3 fallback_models[0] must walk.""" + from agent_core.runtime.retriable import is_retriable_with_fallback + err = RuntimeError( + "Error code: 401 - {'error': {'message': 'Missing Authentication " + "header', 'code': 401}}", + ) + assert is_retriable_with_fallback(err) + + +def test_classify_returns_auth_failure_label() -> None: + from agent_core.runtime.retriable import classify_error + err = RuntimeError( + "Error code: 401 - {'error': {'message': 'Missing Authentication " + "header', 'code': 401}}", + ) + assert classify_error(err) == "auth_failure" + + +def test_classify_precedence_model_unavailable_before_auth_failure() -> None: + """If a chain returns both 401 AND model_not_found (e.g. distributor + proxy rejected the model name before checking the key), + ``model_unavailable`` wins — the actionable fix is the chain config, + not the credentials.""" + from agent_core.runtime.retriable import classify_error + err = RuntimeError("401 unauthorized; also model_not_found in body") + assert classify_error(err) == "model_unavailable"