From 97c10cbb425f4ccb6ceae8cd564fb1b50678427b Mon Sep 17 00:00:00 2001 From: ATCharpentier Date: Thu, 30 Jul 2026 13:24:43 -0400 Subject: [PATCH 01/18] feat: typed LLM errors + SSE capacity classification MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit New src/llm/errors.py: LLMError(RuntimeError) hierarchy — Capacity (model-scoped, retryable), RateLimit (account-scoped, fast-fail after internal rotation), Transport, Auth, Request. retry_after stays a plain attribute for the two duck-typed consumers. openai_codex: stream readers parse the SSE error object structurally (type/code/retry_after; error events and response.failed both handled); capacity markers (service_unavailable_error / server_is_overloaded / server_error) escape _send_with_retries immediately as LLMCapacityError — no inner retry burn, no client-breaker count, never account rotation. Terminal raises typed with identical message strings (401→Auth, 429-exhausted→RateLimit, 5xx/conn→Transport, other→Request); 429 rotation via _mark_limited unchanged. ollama/kimi exhaustion raises wrapped the same way (kimi 429 carries the parsed Retry-After). --- src/llm/__init__.py | 10 +++ src/llm/errors.py | 90 +++++++++++++++++++++++ src/llm/kimi.py | 33 +++++++-- src/llm/ollama.py | 19 ++++- src/llm/openai_codex.py | 142 +++++++++++++++++++++++++++++++++---- tests/test_llm_errors.py | 149 +++++++++++++++++++++++++++++++++++++++ 6 files changed, 421 insertions(+), 22 deletions(-) create mode 100644 src/llm/errors.py create mode 100644 tests/test_llm_errors.py diff --git a/src/llm/__init__.py b/src/llm/__init__.py index 5cf7f0f7..ed216949 100644 --- a/src/llm/__init__.py +++ b/src/llm/__init__.py @@ -3,6 +3,14 @@ from .circuit_breaker import CircuitOpenError from .codex_auth import CodexAuth, CodexAuthPool from .cost_tracker import CostTracker +from .errors import ( + LLMAuthError, + LLMCapacityError, + LLMError, + LLMRateLimitError, + LLMRequestError, + LLMTransportError, +) from .kimi import KimiClient from .ollama import OllamaClient from .openai_codex import CodexChatClient @@ -13,5 +21,7 @@ "AuxiliaryLLMClient", "CircuitOpenError", "CodexAuth", "CodexAuthPool", "CodexChatClient", "CostTracker", "KimiClient", "LLMProvider", "LLMResponse", "OllamaClient", "ToolCall", + "LLMAuthError", "LLMCapacityError", "LLMError", "LLMRateLimitError", + "LLMRequestError", "LLMTransportError", "compute_backoff", "compute_backoff_no_jitter", ] diff --git a/src/llm/errors.py b/src/llm/errors.py new file mode 100644 index 00000000..98999843 --- /dev/null +++ b/src/llm/errors.py @@ -0,0 +1,90 @@ +"""Typed LLM provider errors. + +Classified from structured provider signals (SSE error events, HTTP status +exhaustion) instead of leaving every failure a bare ``RuntimeError`` string. +The recovery layer keys its retry/fast-fail decision on these types; the +subsystem guard and circuit breakers key their counting on them. + +Compatibility constraints (both load-bearing): + +- Every class subclasses ``RuntimeError``: legacy consumers catching + ``RuntimeError``/``Exception`` keep working, and message text keeps the + provider's bounded shape. +- ``retry_after`` is a plain attribute (``float | None``) because two + consumers discover it by duck typing (``hasattr(err, "retry_after")``): + ``src/agents/manager.py`` and ``src/tools/autonomous_loop.py``. + +Only whitelisted, safe fields ride on the exception: ``provider``, +``model``, ``retry_after``. Raise sites are responsible for bounding any +response-body text they put in the message (the existing ``[:200]`` / +``[:500]`` discipline); user-facing presentation still goes through +``format_user_facing_error`` at the boundary. +""" + +from __future__ import annotations + + +class LLMError(RuntimeError): + """Base for typed LLM provider failures. + + ``retryable`` documents the default recovery treatment; the recovery + policy makes the actual decision by isinstance checks so subclasses + stay the single source of truth. + """ + + retryable: bool = False + + def __init__( + self, + message: str, + *, + provider: str | None = None, + model: str | None = None, + retry_after: float | None = None, + ) -> None: + super().__init__(message) + self.provider = provider + self.model = model + self.retry_after = retry_after + + +class LLMCapacityError(LLMError): + """The model tier is out of capacity (e.g. ``server_is_overloaded``). + + Model-scoped: every account sees the same failure, so account rotation + must never be triggered by this class. Retryable — the deadline-based + recovery policy owns the wait; it counts once per failed logical + generation toward the model-scoped breaker. + """ + + retryable = True + + +class LLMRateLimitError(LLMError): + """Account quota exhausted (HTTP 429) after internal rotation ran out. + + Account-scoped. The provider client has already marked accounts limited + and rotated through the pool by the time this is raised, so the outer + recovery policy FAST-FAILS it — spending the recovery budget cycling + accounts already marked limited would change quota semantics. + """ + + retryable = False + + +class LLMTransportError(LLMError): + """Connection/stream transport failure after the client's own retries.""" + + retryable = True + + +class LLMAuthError(LLMError): + """Authentication failed with no healthy account left. Fast-fail.""" + + retryable = False + + +class LLMRequestError(LLMError): + """The request itself is invalid (bad model, malformed input). Fast-fail.""" + + retryable = False diff --git a/src/llm/kimi.py b/src/llm/kimi.py index 37bc593c..dae2752d 100644 --- a/src/llm/kimi.py +++ b/src/llm/kimi.py @@ -15,6 +15,7 @@ from ..odin_log import get_logger from .backoff import DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY, DEFAULT_MAX_RETRIES, compute_backoff from .circuit_breaker import CircuitBreaker +from .errors import LLMRateLimitError, LLMRequestError, LLMTransportError from .provider import LLMProvider from .types import LLMResponse, ToolCall @@ -232,8 +233,18 @@ async def _request_with_retry(self, body: dict) -> dict: if resp.status == 429: if attempt >= self.max_retries: self.breaker.record_failure() - raise RuntimeError(f"Kimi rate limited after {self.max_retries + 1} " - f"attempts: {text[:300]}") + hdr = resp.headers.get("Retry-After") + try: + hdr_delay = float(hdr) if hdr else None + except (ValueError, TypeError): + hdr_delay = None + raise LLMRateLimitError( + f"Kimi rate limited after {self.max_retries + 1} " + f"attempts: {text[:300]}", + provider="kimi", + model=self.model, + retry_after=hdr_delay, + ) retry_after = resp.headers.get("Retry-After") try: delay = ( @@ -267,7 +278,16 @@ async def _request_with_retry(self, body: dict) -> dict: continue self.breaker.record_failure() - raise RuntimeError(f"Kimi {resp.status}: {text[:500]}") + exc_cls = ( + LLMTransportError + if resp.status in (500, 502, 503, 504) + else LLMRequestError + ) + raise exc_cls( + f"Kimi {resp.status}: {text[:500]}", + provider="kimi", + model=self.model, + ) except (TimeoutError, aiohttp.ClientError) as e: last_error = e self.breaker.record_failure() @@ -277,8 +297,11 @@ async def _request_with_retry(self, body: dict) -> dict: attempt + 1, self.max_retries + 1, e, delay) await asyncio.sleep(delay) continue - raise RuntimeError(f"Kimi connection error after {self.max_retries + 1} attempts: " - f"{e}") from e + raise LLMTransportError( + f"Kimi connection error after {self.max_retries + 1} attempts: {e}", + provider="kimi", + model=self.model, + ) from e raise RuntimeError(f"Kimi request failed after {self.max_retries + 1} attempts: " f"{last_error}") diff --git a/src/llm/ollama.py b/src/llm/ollama.py index e354e404..0f6bd9f3 100644 --- a/src/llm/ollama.py +++ b/src/llm/ollama.py @@ -16,6 +16,7 @@ from .backoff import DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY, DEFAULT_MAX_RETRIES, compute_backoff from .circuit_breaker import CircuitBreaker from .cost_tracker import estimate_tokens +from .errors import LLMRequestError, LLMTransportError from .provider import LLMProvider from .types import LLMResponse, ToolCall @@ -206,7 +207,16 @@ async def _request_with_retry(self, body: dict) -> dict: await asyncio.sleep(delay) continue self.breaker.record_failure() - raise RuntimeError(f"Ollama {resp.status}: {text[:500]}") + exc_cls = ( + LLMTransportError + if resp.status in (500, 502, 503, 504) + else LLMRequestError + ) + raise exc_cls( + f"Ollama {resp.status}: {text[:500]}", + provider="ollama", + model=self.model, + ) except (TimeoutError, aiohttp.ClientError) as e: last_error = e self.breaker.record_failure() @@ -216,8 +226,11 @@ async def _request_with_retry(self, body: dict) -> dict: attempt + 1, self.max_retries + 1, e, delay) await asyncio.sleep(delay) continue - raise RuntimeError(f"Ollama connection error after {self.max_retries + 1} " - f"attempts: {e}") from e + raise LLMTransportError( + f"Ollama connection error after {self.max_retries + 1} attempts: {e}", + provider="ollama", + model=self.model, + ) from e raise RuntimeError(f"Ollama request failed after {self.max_retries + 1} attempts: " f"{last_error}") diff --git a/src/llm/openai_codex.py b/src/llm/openai_codex.py index 502c784a..eb1e5d78 100644 --- a/src/llm/openai_codex.py +++ b/src/llm/openai_codex.py @@ -10,6 +10,13 @@ from .circuit_breaker import CircuitBreaker from .codex_auth import CodexAuth, CodexAuthPool from .cost_tracker import estimate_tokens +from .errors import ( + LLMAuthError, + LLMCapacityError, + LLMRateLimitError, + LLMRequestError, + LLMTransportError, +) from .types import LLMResponse, ToolCall log = get_logger("codex") @@ -29,8 +36,70 @@ CONNECT_TIMEOUT = 30 +# Markers the backend uses for model-tier capacity exhaustion. These arrive +# INSIDE an HTTP 200 as SSE error events, so no status-code branch ever sees +# them; matched against both error.type and error.code. Observed live +# (2026-07-29/30 sol degradation): type=service_unavailable_error with +# code=server_is_overloaded, plus bare server_error. +_CAPACITY_ERROR_MARKERS = frozenset({ + "service_unavailable_error", + "server_is_overloaded", + "server_error", +}) + + class CodexStreamError(RuntimeError): - """The SSE stream reported a terminal failure event (response.failed / error).""" + """The SSE stream reported a terminal failure event (response.failed / error). + + Carries the structured fields parsed from the event's error object so the + retry engine can classify capacity failures without substring matching. + The message keeps the historical ``{event_type}: {json[:500]}`` shape. + """ + + def __init__( + self, + message: str, + *, + error_type: str | None = None, + error_code: str | None = None, + retry_after: float | None = None, + ) -> None: + super().__init__(message) + self.error_type = error_type + self.error_code = error_code + self.retry_after = retry_after + + @property + def is_capacity(self) -> bool: + return ( + self.error_type in _CAPACITY_ERROR_MARKERS + or self.error_code in _CAPACITY_ERROR_MARKERS + ) + + +def _stream_error_from_event(event_type: str, event: dict) -> CodexStreamError: + """Build a classified CodexStreamError from a terminal SSE event. + + The error object lives at ``event["error"]`` for bare ``error`` events + and at ``event["response"]["error"]`` for ``response.failed``; tolerate + both plus absence (classification simply stays empty and the failure is + treated as transport, today's behavior). + """ + err = event.get("error") + if not isinstance(err, dict): + resp_obj = event.get("response") + err = resp_obj.get("error") if isinstance(resp_obj, dict) else None + if not isinstance(err, dict): + err = {} + error_type = err.get("type") + error_code = err.get("code") + retry_after = err.get("retry_after") + return CodexStreamError( + f"{event_type}: {json.dumps(event)[:500]}", + error_type=error_type if isinstance(error_type, str) else None, + error_code=error_code if isinstance(error_code, str) else None, + retry_after=float(retry_after) if isinstance(retry_after, (int, float)) else None, + ) class CodexChatClient: @@ -537,6 +606,22 @@ async def _send_with_retries(self, body: dict, reader, result_is_empty): try: result = await reader(resp) except CodexStreamError as e: + if e.is_capacity: + # Model-tier capacity exhaustion (e.g. + # server_is_overloaded) inside a 200. Every + # account shares it, so: no account rotation, + # no client-breaker count, no inner retry + # burn. Escape immediately — the deadline- + # based recovery layer owns the wait and + # counts one model-breaker failure per failed + # logical generation. + raise LLMCapacityError( + "Codex capacity: " + f"{e.error_code or e.error_type or 'unknown'}", + provider="codex", + model=str(body.get("model") or self.model), + retry_after=e.retry_after, + ) from e # response.failed / error event: the "200" turned # out to be a failure mid-stream — retryable. self.breaker.record_failure() @@ -553,7 +638,11 @@ async def _send_with_retries(self, body: dict, reader, result_is_empty): ) await asyncio.sleep(wait) continue - raise RuntimeError(f"Codex stream failed: {last_error}") from e + raise LLMTransportError( + f"Codex stream failed: {last_error}", + provider="codex", + model=str(body.get("model") or self.model), + ) from e if not result_is_empty(result): self.breaker.record_success() return result @@ -601,8 +690,10 @@ async def _send_with_retries(self, body: dict, reader, result_is_empty): token, account_id, acct_idx = await self._acquire_auth() continue self.breaker.record_failure() - raise RuntimeError( - f"Codex 401 (auth failed, no healthy account): {error_body[:200]}" + raise LLMAuthError( + f"Codex 401 (auth failed, no healthy account): {error_body[:200]}", + provider="codex", + model=str(body.get("model") or self.model), ) if resp.status == 429: @@ -623,7 +714,16 @@ async def _send_with_retries(self, body: dict, reader, result_is_empty): await asyncio.sleep(wait) token, account_id, acct_idx = await self._acquire_auth() continue - raise RuntimeError(f"Codex API error (429): {error_body[:500]}") + # Internal rotation is exhausted at this point (every + # attempt marked its account limited and re-acquired). + # Typed so the outer recovery FAST-FAILS instead of + # spending its budget cycling limited accounts — + # quota semantics stay exactly as before. + raise LLMRateLimitError( + f"Codex API error (429): {error_body[:500]}", + provider="codex", + model=str(body.get("model") or self.model), + ) if resp.status in (500, 502, 503, 504): self.breaker.record_failure() @@ -640,10 +740,20 @@ async def _send_with_retries(self, body: dict, reader, result_is_empty): ) await asyncio.sleep(wait) continue - raise RuntimeError(f"Codex API error ({resp.status}): {error_body[:500]}") + raise LLMTransportError( + f"Codex API error ({resp.status}): {error_body[:500]}", + provider="codex", + model=str(body.get("model") or self.model), + ) + # Any other status (400 bad model/malformed request, ...): + # the request itself is wrong — fast-fail, never retried. self.breaker.record_failure() - raise RuntimeError(f"Codex API error ({resp.status}): {error_body[:500]}") + raise LLMRequestError( + f"Codex API error ({resp.status}): {error_body[:500]}", + provider="codex", + model=str(body.get("model") or self.model), + ) except (TimeoutError, aiohttp.ClientError) as e: # asyncio.TimeoutError: the total/sock_read timeouts can fire @@ -659,7 +769,11 @@ async def _send_with_retries(self, body: dict, reader, result_is_empty): ) await asyncio.sleep(wait) else: - raise RuntimeError(f"Codex API connection failed: {last_error}") from e + raise LLMTransportError( + f"Codex API connection failed: {last_error}", + provider="codex", + model=str(body.get("model") or self.model), + ) from e raise RuntimeError(f"Codex API failed after {self.max_retries} retries: {last_error}") @@ -781,9 +895,9 @@ async def _read_tool_stream(self, resp: aiohttp.ClientResponse) -> LLMResponse: # generation — surface it so the retry engine treats it as an # error instead of returning partial output as a completed turn. elif event_type in ("response.failed", "error"): - detail = json.dumps(event)[:500] - log.warning("Codex stream terminal failure %s: %s", event_type, detail) - raise CodexStreamError(f"{event_type}: {detail}") + exc = _stream_error_from_event(event_type, event) + log.warning("Codex stream terminal failure %s", exc) + raise exc # Incomplete (length-capped / filtered): keep the partial output # but mark it so callers can tell it isn't a normal completion. @@ -876,9 +990,9 @@ async def _read_stream(self, resp: aiohttp.ClientResponse) -> str: # Terminal failure events — surface to the retry engine instead of # returning partial output as a normal completion. elif event_type in ("response.failed", "error"): - detail = json.dumps(event)[:500] - log.warning("Codex stream terminal failure %s: %s", event_type, detail) - raise CodexStreamError(f"{event_type}: {detail}") + exc = _stream_error_from_event(event_type, event) + log.warning("Codex stream terminal failure %s", exc) + raise exc elif event_type == "response.incomplete": reason = ((event.get("response") or {}).get("incomplete_details") diff --git a/tests/test_llm_errors.py b/tests/test_llm_errors.py new file mode 100644 index 00000000..f7da7d0f --- /dev/null +++ b/tests/test_llm_errors.py @@ -0,0 +1,149 @@ +"""Pins for the typed LLM error hierarchy (src/llm/errors.py). + +The compatibility contracts here are load-bearing: +- RuntimeError subclassing keeps every legacy `except RuntimeError` working. +- `.retry_after` as a plain attribute keeps the two duck-typed consumers + (agents/manager.py, tools/autonomous_loop.py) working unchanged. +""" + +from __future__ import annotations + +import pytest + +from src.llm import ( + LLMAuthError, + LLMCapacityError, + LLMError, + LLMRateLimitError, + LLMRequestError, + LLMTransportError, +) + +ALL_TYPES = [ + LLMCapacityError, + LLMRateLimitError, + LLMTransportError, + LLMAuthError, + LLMRequestError, +] + + +@pytest.mark.parametrize("exc_type", ALL_TYPES) +def test_every_type_is_runtimeerror_and_llmerror(exc_type): + err = exc_type("boom") + assert isinstance(err, RuntimeError) + assert isinstance(err, LLMError) + assert str(err) == "boom" + + +@pytest.mark.parametrize("exc_type", ALL_TYPES) +def test_retry_after_duck_typing_contract(exc_type): + # hasattr must be True even when unset — consumers read it after hasattr(). + bare = exc_type("x") + assert hasattr(bare, "retry_after") + assert bare.retry_after is None + + carrying = exc_type("x", retry_after=12.5) + assert carrying.retry_after == 12.5 + + +@pytest.mark.parametrize("exc_type", ALL_TYPES) +def test_whitelisted_fields_default_none(exc_type): + err = exc_type("x") + assert err.provider is None + assert err.model is None + + stamped = exc_type("x", provider="codex", model="gpt-5.6-sol") + assert stamped.provider == "codex" + assert stamped.model == "gpt-5.6-sol" + + +def test_retryable_classification(): + assert LLMCapacityError.retryable is True + assert LLMTransportError.retryable is True + # Rate-limit is deliberately NOT retryable at the recovery layer: the + # client has already rotated accounts by the time it raises (round-3 + # clarification — quota semantics must remain exactly today's). + assert LLMRateLimitError.retryable is False + assert LLMAuthError.retryable is False + assert LLMRequestError.retryable is False + assert LLMError.retryable is False + + +def test_base_llmerror_catches_all_subtypes(): + for exc_type in ALL_TYPES: + with pytest.raises(LLMError): + raise exc_type("x") + + +class TestCodexStreamErrorClassification: + """Pins for _stream_error_from_event / CodexStreamError.is_capacity.""" + + # Byte-for-byte the event shape observed in the live journal during the + # 2026-07-29/30 sol degradation (110 occurrences on 07-30 alone). + LIVE_OVERLOAD_EVENT = { + "type": "error", + "error": { + "type": "service_unavailable_error", + "code": "server_is_overloaded", + "message": "Our servers are currently overloaded. Please try again later.", + "param": None, + }, + "sequence_number": 2, + } + + def _build(self, event_type, event): + from src.llm.openai_codex import _stream_error_from_event + + return _stream_error_from_event(event_type, event) + + def test_live_overload_event_classifies_as_capacity(self): + exc = self._build("error", self.LIVE_OVERLOAD_EVENT) + assert exc.is_capacity is True + assert exc.error_type == "service_unavailable_error" + assert exc.error_code == "server_is_overloaded" + assert exc.retry_after is None + # Historical message shape preserved: "{event_type}: {json[:500]}". + assert str(exc).startswith("error: ") + assert "server_is_overloaded" in str(exc) + + def test_server_error_code_is_capacity(self): + exc = self._build("error", {"type": "error", "error": {"code": "server_error"}}) + assert exc.is_capacity is True + + def test_response_failed_nested_error_object(self): + event = { + "type": "response.failed", + "response": {"error": {"type": "service_unavailable_error"}}, + } + exc = self._build("response.failed", event) + assert exc.is_capacity is True + assert str(exc).startswith("response.failed: ") + + def test_non_capacity_error_stays_transport_classified(self): + event = {"type": "error", "error": {"type": "invalid_request_error", "code": "nope"}} + exc = self._build("error", event) + assert exc.is_capacity is False + + def test_absent_error_object_tolerated(self): + exc = self._build("response.failed", {"type": "response.failed"}) + assert exc.is_capacity is False + assert exc.error_type is None + assert exc.error_code is None + + def test_retry_after_numeric_extraction(self): + event = {"type": "error", "error": {"code": "server_is_overloaded", "retry_after": 30}} + exc = self._build("error", event) + assert exc.retry_after == 30.0 + + junk = {"type": "error", "error": {"code": "server_is_overloaded", "retry_after": "soon"}} + assert self._build("error", junk).retry_after is None + + def test_capacity_markers_are_the_settled_set(self): + from src.llm.openai_codex import _CAPACITY_ERROR_MARKERS + + assert _CAPACITY_ERROR_MARKERS == { + "service_unavailable_error", + "server_is_overloaded", + "server_error", + } From 4fe5f89e40d908a4ee516ec5bf91ad87bac04a2a Mon Sep 17 00:00:00 2001 From: ATCharpentier Date: Thu, 30 Jul 2026 13:31:33 -0400 Subject: [PATCH 02/18] feat: model-scoped capacity breaker + shared deadline recovery MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit src/llm/model_breaker.py: ModelCapacityBreaker keyed provider:effective- model, owned by a registry (BotServices-scoped — survives client rebuilds and live reloads). Counted once per failed LOGICAL GENERATION, true single-probe half-open with token-attributed resolution (a stray non-probe failure can neither escalate nor release someone else's probe), adaptive cooldown doubling to a cap and resetting on success. src/llm/recovery.py: generate_with_recovery — THE shared policy for chat/ agents/loops. Monotonic deadline (default 300s; resume passes remaining budget so restarts never refresh it), full-jitter backoff capped 45s, retry_after honoured as a clamped floor, waits bounded by remaining budget and raced against the caller's cancel event. Retryable: Capacity/ Transport/CircuitOpen (waits THROUGH the client breaker). Fast-fail: Auth/Request/RateLimit (rotation already exhausted — quota semantics unchanged). Unclassified exceptions escape immediately (the agents bare-except defect is fixed, not spread). CircuitOpenError gains an optional model stamp. --- src/llm/circuit_breaker.py | 5 +- src/llm/model_breaker.py | 241 ++++++++++++++++++++++++++++++++++ src/llm/recovery.py | 209 +++++++++++++++++++++++++++++ tests/test_model_breaker.py | 203 ++++++++++++++++++++++++++++ tests/test_recovery_policy.py | 203 ++++++++++++++++++++++++++++ 5 files changed, 860 insertions(+), 1 deletion(-) create mode 100644 src/llm/model_breaker.py create mode 100644 src/llm/recovery.py create mode 100644 tests/test_model_breaker.py create mode 100644 tests/test_recovery_policy.py diff --git a/src/llm/circuit_breaker.py b/src/llm/circuit_breaker.py index 08ba5b40..33431586 100644 --- a/src/llm/circuit_breaker.py +++ b/src/llm/circuit_breaker.py @@ -22,9 +22,12 @@ class CircuitOpenError(Exception): """Raised when a circuit breaker is open and requests should not be attempted.""" - def __init__(self, provider: str, retry_after: float) -> None: + def __init__(self, provider: str, retry_after: float, model: str | None = None) -> None: self.provider = provider self.retry_after = retry_after + # Optional effective-model stamp for model-scoped diagnostics; + # existing two-arg raise sites are unchanged. + self.model = model super().__init__( f"{provider} is temporarily unavailable (retry in {retry_after:.0f}s)" ) diff --git a/src/llm/model_breaker.py b/src/llm/model_breaker.py new file mode 100644 index 00000000..ee6ca6fd --- /dev/null +++ b/src/llm/model_breaker.py @@ -0,0 +1,241 @@ +"""Model-scoped capacity breakers. + +Capacity exhaustion (``server_is_overloaded``) is a property of a MODEL +TIER, not of an account or a client instance: every account sees the same +failure, and rebuilding a client on live reload does not change upstream +capacity. So capacity admission gets its own breaker, keyed by +``provider:effective-model`` and owned by a registry that lives in +``BotServices`` — client rebuilds and live reloads never reset it. + +Differences from ``CircuitBreaker`` (which keeps its per-client role for +transport/auth/429 counting): + +- **Counted per failed logical generation**, not per HTTP attempt. The + recovery layer calls :meth:`ModelCapacityBreaker.record_generation_failure` + exactly once when a whole generation's recovery budget is exhausted. +- **True single-probe half-open.** When the cooldown elapses, exactly one + caller is admitted as the probe; everyone else keeps waiting until the + probe resolves. (``CircuitBreaker.check()`` lets every concurrent caller + through after the timeout — documented gap.) +- **Adaptive cooldown.** Consecutive opens double the cooldown up to a cap, + so a long outage probes progressively less often; any success resets it. + +Thread-safe; methods are cheap and callable from async code without +awaiting. Waiters poll — the registry deliberately has no cross-event-loop +signalling. +""" + +from __future__ import annotations + +import threading +import time + +from ..odin_log import get_logger + +log = get_logger("model_breaker") + +DEFAULT_GENERATION_THRESHOLD = 1 +DEFAULT_COOLDOWN_BASE = 30.0 +DEFAULT_COOLDOWN_CAP = 300.0 + +# How soon a waiter should poll again while another caller holds the probe. +_PROBE_PENDING_WAIT = 5.0 + + +class AdmissionToken: + """Opaque admission handle returned by :meth:`acquire_attempt`. + + Resolution methods take the token so only the caller that actually holds + the probe slot can escalate or release it — a concurrent non-probe + failure (admitted earlier, while the breaker was still closed) must + never resolve someone else's probe. + """ + + __slots__ = () + + +class ModelCapacityBreaker: + """Capacity breaker for one ``provider:model`` pair.""" + + def __init__( + self, + name: str, + *, + generation_threshold: int = DEFAULT_GENERATION_THRESHOLD, + cooldown_base: float = DEFAULT_COOLDOWN_BASE, + cooldown_cap: float = DEFAULT_COOLDOWN_CAP, + ) -> None: + self.name = name + self.generation_threshold = max(1, generation_threshold) + self.cooldown_base = max(1.0, cooldown_base) + self.cooldown_cap = max(self.cooldown_base, cooldown_cap) + self._open = False + self._probe_token: AdmissionToken | None = None + self._failed_generations = 0 + self._consecutive_opens = 0 + self._opened_at = 0.0 + self._lock = threading.Lock() + + # -- introspection ------------------------------------------------- + + @property + def state(self) -> str: + with self._lock: + if not self._open: + return "closed" + return "probing" if self._probe_token is not None else "open" + + @property + def is_closed(self) -> bool: + with self._lock: + return not self._open + + def _current_cooldown(self) -> float: + # Lock held by caller. First open waits cooldown_base, each + # consecutive re-open doubles it up to the cap. + exponent = max(0, self._consecutive_opens - 1) + return min(self.cooldown_cap, self.cooldown_base * (2.0**exponent)) + + def snapshot(self) -> dict: + with self._lock: + return { + "name": self.name, + "state": ( + "closed" + if not self._open + else ("probing" if self._probe_token is not None else "open") + ), + "failed_generations": self._failed_generations, + "consecutive_opens": self._consecutive_opens, + "cooldown_seconds": self._current_cooldown() if self._open else 0.0, + } + + # -- attempt admission --------------------------------------------- + + def acquire_attempt(self) -> AdmissionToken | float: + """Ask to make one attempt against the model. + + Returns an :class:`AdmissionToken` when admitted — either the + breaker is closed, or the cooldown has elapsed and THIS caller now + holds the single probe slot. The caller must resolve every admitted + attempt with exactly one of :meth:`attempt_succeeded`, + :meth:`attempt_failed_capacity`, or :meth:`abandon`, passing the + token back. Returns a positive number of seconds to wait otherwise. + """ + with self._lock: + if not self._open: + return AdmissionToken() + elapsed = time.monotonic() - self._opened_at + cooldown = self._current_cooldown() + if elapsed < cooldown: + return max(0.1, cooldown - elapsed) + if self._probe_token is not None: + return _PROBE_PENDING_WAIT + token = AdmissionToken() + self._probe_token = token + log.info("Model breaker %s: admitting half-open probe", self.name) + return token + + def attempt_succeeded(self, token: AdmissionToken) -> None: + """Any successful generation attempt: close and reset everything.""" + with self._lock: + was_open = self._open + self._open = False + if self._probe_token is token: + self._probe_token = None + self._failed_generations = 0 + self._consecutive_opens = 0 + if was_open: + log.info("Model breaker %s: attempt succeeded — closed", self.name) + + def attempt_failed_capacity(self, token: AdmissionToken) -> None: + """A capacity failure on one admitted attempt. + + Does NOT count toward generation failures (the recovery layer owns + that). Its job is probe resolution: a failed PROBE re-opens with an + escalated cooldown. A non-probe admission (granted while the breaker + was still closed) resolves as a no-op here. + """ + with self._lock: + if self._open and self._probe_token is token: + self._probe_token = None + self._consecutive_opens += 1 + self._opened_at = time.monotonic() + log.info( + "Model breaker %s: probe failed — cooldown now %.0fs", + self.name, + self._current_cooldown(), + ) + + def abandon(self, token: AdmissionToken) -> None: + """The admitted attempt ended without a capacity verdict. + + Covers cancellation and non-capacity failures (auth, transport, + request errors say nothing about model capacity). Releases the probe + slot without escalating — the cooldown that admitted the probe has + already elapsed, so the next caller may probe immediately. + """ + with self._lock: + if self._probe_token is token: + self._probe_token = None + + # -- generation-level counting ------------------------------------- + + def record_generation_failure(self) -> None: + """One whole logical generation exhausted its recovery budget.""" + with self._lock: + self._failed_generations += 1 + if not self._open and self._failed_generations >= self.generation_threshold: + self._open = True + self._probe_token = None + self._consecutive_opens += 1 + self._opened_at = time.monotonic() + log.warning( + "Model breaker %s: OPEN after %d failed generation(s) " + "(cooldown %.0fs)", + self.name, + self._failed_generations, + self._current_cooldown(), + ) + + +class ModelBreakerRegistry: + """Get-or-create registry of :class:`ModelCapacityBreaker` by key. + + Owned by ``BotServices`` so breaker state survives provider client + rebuilds and live config reloads. Keys are ``f"{provider}:{model}"`` — + always the EFFECTIVE model of the request (per-request overrides + included), per the round-3 clarification. + """ + + def __init__( + self, + *, + generation_threshold: int = DEFAULT_GENERATION_THRESHOLD, + cooldown_base: float = DEFAULT_COOLDOWN_BASE, + cooldown_cap: float = DEFAULT_COOLDOWN_CAP, + ) -> None: + self._generation_threshold = generation_threshold + self._cooldown_base = cooldown_base + self._cooldown_cap = cooldown_cap + self._breakers: dict[str, ModelCapacityBreaker] = {} + self._lock = threading.Lock() + + def for_model(self, provider: str, model: str) -> ModelCapacityBreaker: + key = f"{provider or 'unknown'}:{model or 'unknown'}" + with self._lock: + breaker = self._breakers.get(key) + if breaker is None: + breaker = ModelCapacityBreaker( + key, + generation_threshold=self._generation_threshold, + cooldown_base=self._cooldown_base, + cooldown_cap=self._cooldown_cap, + ) + self._breakers[key] = breaker + return breaker + + def snapshot(self) -> dict[str, dict]: + with self._lock: + breakers = dict(self._breakers) + return {key: breaker.snapshot() for key, breaker in breakers.items()} diff --git a/src/llm/recovery.py b/src/llm/recovery.py new file mode 100644 index 00000000..04aa514d --- /dev/null +++ b/src/llm/recovery.py @@ -0,0 +1,209 @@ +"""Deadline-based recovery for logical LLM generations. + +THE shared retry policy across chat, agents, and autonomous loops — one +implementation replacing the three divergent ones (tool_loop's +CircuitOpenError-only sleep-and-retry, the agents' bare-``except`` double +attempt, the loop manager's consecutive-error counter feeding). + +Design rules (settled with Odin, 2026-07-30): + +- **A monotonic DEADLINE, never an attempt count.** Default five minutes per + logical generation. Callers resuming a persisted turn pass the remaining + budget (computed from the checkpoint's UTC deadline) via + ``deadline_seconds`` — a restart must not grant a fresh five minutes. +- **Retryable**: ``LLMCapacityError``, ``LLMTransportError``, and + ``CircuitOpenError`` (the client breaker — recovery waits THROUGH it + rather than treating it as terminal). Full-jitter backoff capped at + ``backoff_cap``; a server-suggested ``retry_after`` is honoured as a floor + (clamped to ``retry_after_cap``), and every wait is bounded by the + remaining budget. +- **Fast-fail**: ``LLMAuthError``, ``LLMRequestError``, and + ``LLMRateLimitError`` — rate-limit is deliberate: by the time the client + raises it, internal account rotation is already exhausted, and cycling + limited accounts for five minutes would change quota semantics. +- **No bare ``except``**: an unclassified exception is a programming defect + and escapes immediately (retrying defects was the agents-path bug — fixed + here, not spread). +- **Model-breaker discipline**: one admission per attempt via the + token-resolution protocol; exactly ONE ``record_generation_failure`` when + the whole generation exhausts on capacity. Non-capacity outcomes abandon + the probe without escalating. +- **Cancellable**: waits race the caller's cancel event; cancellation + propagates as ``asyncio.CancelledError`` and always releases a held probe. +""" + +from __future__ import annotations + +import asyncio +import time +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from typing import TypeVar + +from ..odin_log import get_logger +from .backoff import compute_backoff +from .circuit_breaker import CircuitOpenError +from .errors import ( + LLMAuthError, + LLMCapacityError, + LLMRateLimitError, + LLMRequestError, + LLMTransportError, +) +from .model_breaker import AdmissionToken, ModelCapacityBreaker + +log = get_logger("llm_recovery") + +T = TypeVar("T") + +DEFAULT_GENERATION_DEADLINE = 300.0 +DEFAULT_BACKOFF_BASE = 1.0 +DEFAULT_BACKOFF_CAP = 45.0 +# Precedent: every existing consumer clamps a server-suggested retry_after +# to 90s (tool_loop, agents/manager, autonomous_loop). +DEFAULT_RETRY_AFTER_CAP = 90.0 + + +@dataclass(frozen=True) +class RecoveryPolicy: + """Knobs for one recovery run; config-backed via ``llm_recovery:``.""" + + deadline_seconds: float = DEFAULT_GENERATION_DEADLINE + backoff_base: float = DEFAULT_BACKOFF_BASE + backoff_cap: float = DEFAULT_BACKOFF_CAP + retry_after_cap: float = DEFAULT_RETRY_AFTER_CAP + + +def _check_cancel(cancel_event: asyncio.Event | None) -> None: + if cancel_event is not None and cancel_event.is_set(): + raise asyncio.CancelledError("cancelled during LLM recovery") + + +async def _sleep_cancellable(seconds: float, cancel_event: asyncio.Event | None) -> None: + """Sleep, waking immediately (as CancelledError) if the event fires.""" + if seconds <= 0: + return + if cancel_event is None: + await asyncio.sleep(seconds) + return + try: + await asyncio.wait_for(cancel_event.wait(), timeout=seconds) + except TimeoutError: + return + raise asyncio.CancelledError("cancelled during LLM recovery wait") + + +def _notify_wait( + on_wait: Callable[[float, float, BaseException], None] | None, + wait: float, + remaining: float, + error: BaseException, +) -> None: + if on_wait is None: + return + try: + on_wait(wait, remaining, error) + except Exception: # pragma: no cover - observability hook must not kill recovery + log.debug("recovery on_wait hook failed", exc_info=True) + + +async def generate_with_recovery( + attempt: Callable[[], Awaitable[T]], + *, + policy: RecoveryPolicy, + breaker: ModelCapacityBreaker | None = None, + deadline_seconds: float | None = None, + cancel_event: asyncio.Event | None = None, + on_wait: Callable[[float, float, BaseException], None] | None = None, +) -> T: + """Run one logical LLM generation with deadline-based recovery. + + ``attempt`` is called repeatedly until it succeeds, a non-retryable + error escapes, cancellation fires, or the budget is exhausted — in which + case the last error is re-raised (after counting one generation failure + on the model breaker when that error is capacity-class). + """ + budget = policy.deadline_seconds if deadline_seconds is None else deadline_seconds + deadline = time.monotonic() + max(0.0, budget) + retry_index = 0 + last_error: BaseException | None = None + + def _exhausted() -> BaseException: + error = last_error + if error is None: + # Budget consumed entirely by breaker-paced waits: capacity is + # the story even though this generation never got an attempt. + error = LLMCapacityError( + "LLM recovery budget exhausted waiting for capacity" + + (f" ({breaker.name})" if breaker is not None else ""), + ) + if isinstance(error, LLMCapacityError) and breaker is not None: + breaker.record_generation_failure() + return error + + while True: + _check_cancel(cancel_event) + + # -- model-breaker admission (single-probe pacing while open) -- + token: AdmissionToken | None = None + if breaker is not None: + admission = breaker.acquire_attempt() + while not isinstance(admission, AdmissionToken): + remaining = deadline - time.monotonic() + if remaining <= 0: + raise _exhausted() + wait = min(float(admission), remaining) + _notify_wait( + on_wait, + wait, + remaining, + last_error + or LLMCapacityError(f"model breaker open ({breaker.name})"), + ) + await _sleep_cancellable(wait, cancel_event) + _check_cancel(cancel_event) + admission = breaker.acquire_attempt() + token = admission + + # -- one attempt --------------------------------------------- + try: + result = await attempt() + except asyncio.CancelledError: + if breaker is not None and token is not None: + breaker.abandon(token) + raise + except LLMCapacityError as exc: + if breaker is not None and token is not None: + breaker.attempt_failed_capacity(token) + last_error = exc + except (LLMAuthError, LLMRequestError, LLMRateLimitError): + if breaker is not None and token is not None: + breaker.abandon(token) + raise + except (LLMTransportError, CircuitOpenError) as exc: + if breaker is not None and token is not None: + breaker.abandon(token) + last_error = exc + except Exception: + # Unclassified = programming defect. Never retried here. + if breaker is not None and token is not None: + breaker.abandon(token) + raise + else: + if breaker is not None and token is not None: + breaker.attempt_succeeded(token) + return result + + # -- retryable failure: wait within the remaining budget ------ + remaining = deadline - time.monotonic() + if remaining <= 0: + raise _exhausted() + jitter = compute_backoff(retry_index, policy.backoff_base, policy.backoff_cap) + suggested = getattr(last_error, "retry_after", None) + floor = 0.0 + if isinstance(suggested, (int, float)) and suggested > 0: + floor = min(float(suggested), policy.retry_after_cap) + wait = min(remaining, max(jitter, floor)) + retry_index += 1 + _notify_wait(on_wait, wait, remaining, last_error) + await _sleep_cancellable(wait, cancel_event) diff --git a/tests/test_model_breaker.py b/tests/test_model_breaker.py new file mode 100644 index 00000000..3e73c9f5 --- /dev/null +++ b/tests/test_model_breaker.py @@ -0,0 +1,203 @@ +"""Pins for the model-scoped capacity breaker (src/llm/model_breaker.py). + +The three properties Odin's design requires, each pinned here: +- counted once per failed LOGICAL GENERATION (not per HTTP attempt); +- true single-probe half-open (exactly one caller admitted, others wait); +- adaptive cooldown that escalates on failed probes and resets on success. + +Plus the token-attribution property: only the probe holder's failure can +escalate or release the probe slot. +""" + +from __future__ import annotations + +import src.llm.model_breaker as mb +from src.llm.model_breaker import AdmissionToken + + +class FakeClock: + def __init__(self) -> None: + self.now = 1000.0 + + def monotonic(self) -> float: + return self.now + + +def make_breaker(monkeypatch, **kwargs): + clock = FakeClock() + monkeypatch.setattr(mb, "time", clock) + breaker = mb.ModelCapacityBreaker("codex:gpt-5.6-sol", **kwargs) + return breaker, clock + + +def admitted(result) -> bool: + return isinstance(result, AdmissionToken) + + +def test_closed_admits_everyone(monkeypatch): + breaker, _ = make_breaker(monkeypatch) + assert breaker.state == "closed" + assert admitted(breaker.acquire_attempt()) + assert admitted(breaker.acquire_attempt()) + + +def test_opens_only_at_generation_threshold(monkeypatch): + breaker, _ = make_breaker(monkeypatch, generation_threshold=2) + breaker.record_generation_failure() + assert breaker.state == "closed" + breaker.record_generation_failure() + assert breaker.state == "open" + + +def test_attempt_failures_never_open_a_closed_breaker(monkeypatch): + # Per-attempt capacity failures inside a generation's recovery window + # must NOT open the breaker — only generation-level counting does. + breaker, _ = make_breaker(monkeypatch) + for _ in range(20): + token = breaker.acquire_attempt() + assert admitted(token) + breaker.attempt_failed_capacity(token) + assert breaker.state == "closed" + assert admitted(breaker.acquire_attempt()) + + +def test_open_breaker_paces_callers_until_cooldown(monkeypatch): + breaker, clock = make_breaker(monkeypatch, cooldown_base=30.0) + breaker.record_generation_failure() + wait = breaker.acquire_attempt() + assert isinstance(wait, float) and 0 < wait <= 30.0 + clock.now += 10.0 + wait2 = breaker.acquire_attempt() + assert isinstance(wait2, float) and wait2 < wait + + +def test_single_probe_half_open(monkeypatch): + breaker, clock = make_breaker(monkeypatch, cooldown_base=30.0) + breaker.record_generation_failure() + clock.now += 31.0 + # First caller past the cooldown claims THE probe slot. + token = breaker.acquire_attempt() + assert admitted(token) + assert breaker.state == "probing" + # Every other concurrent caller keeps waiting. + assert breaker.acquire_attempt() == mb._PROBE_PENDING_WAIT + assert breaker.acquire_attempt() == mb._PROBE_PENDING_WAIT + + +def test_probe_success_closes_and_resets(monkeypatch): + breaker, clock = make_breaker(monkeypatch, cooldown_base=30.0) + breaker.record_generation_failure() + clock.now += 31.0 + token = breaker.acquire_attempt() + assert admitted(token) + breaker.attempt_succeeded(token) + assert breaker.state == "closed" + assert admitted(breaker.acquire_attempt()) + snap = breaker.snapshot() + assert snap["failed_generations"] == 0 + assert snap["consecutive_opens"] == 0 + + +def test_probe_failure_escalates_cooldown(monkeypatch): + breaker, clock = make_breaker(monkeypatch, cooldown_base=30.0, cooldown_cap=300.0) + breaker.record_generation_failure() # open #1: cooldown 30 + clock.now += 31.0 + token = breaker.acquire_attempt() + assert admitted(token) + breaker.attempt_failed_capacity(token) # open #2: cooldown 60 + wait = breaker.acquire_attempt() + assert isinstance(wait, float) and 55.0 < wait <= 60.0 + clock.now += 61.0 + token = breaker.acquire_attempt() + assert admitted(token) + breaker.attempt_failed_capacity(token) # open #3: cooldown 120 + wait = breaker.acquire_attempt() + assert isinstance(wait, float) and 115.0 < wait <= 120.0 + + +def test_cooldown_caps(monkeypatch): + breaker, clock = make_breaker(monkeypatch, cooldown_base=30.0, cooldown_cap=100.0) + breaker.record_generation_failure() + for _ in range(6): # keep failing probes well past the cap + clock.now += 101.0 + token = breaker.acquire_attempt() + assert admitted(token) + breaker.attempt_failed_capacity(token) + wait = breaker.acquire_attempt() + assert isinstance(wait, float) and wait <= 100.0 + + +def test_abandon_releases_probe_without_escalation(monkeypatch): + breaker, clock = make_breaker(monkeypatch, cooldown_base=30.0) + breaker.record_generation_failure() + clock.now += 31.0 + token = breaker.acquire_attempt() + assert admitted(token) + before = breaker.snapshot()["consecutive_opens"] + breaker.abandon(token) + # Slot free again immediately — next caller becomes the probe. + assert admitted(breaker.acquire_attempt()) + assert breaker.snapshot()["consecutive_opens"] == before + + +def test_stray_token_cannot_resolve_someone_elses_probe(monkeypatch): + # A caller admitted while the breaker was CLOSED fails after the breaker + # opened and another caller claimed the probe. Its resolution must not + # escalate or release the probe slot. + breaker, clock = make_breaker(monkeypatch, cooldown_base=30.0) + stray = breaker.acquire_attempt() # admitted while closed + assert admitted(stray) + breaker.record_generation_failure() # opens + clock.now += 31.0 + probe = breaker.acquire_attempt() + assert admitted(probe) + assert breaker.state == "probing" + + opens_before = breaker.snapshot()["consecutive_opens"] + breaker.attempt_failed_capacity(stray) # stray failure arrives late + assert breaker.state == "probing" # probe still held + assert breaker.snapshot()["consecutive_opens"] == opens_before + breaker.abandon(stray) # stray abandon can't release it either + assert breaker.state == "probing" + + breaker.attempt_succeeded(probe) + assert breaker.state == "closed" + + +def test_success_resets_cooldown_escalation(monkeypatch): + breaker, clock = make_breaker(monkeypatch, cooldown_base=30.0, cooldown_cap=300.0) + breaker.record_generation_failure() + clock.now += 31.0 + token = breaker.acquire_attempt() + assert admitted(token) + breaker.attempt_failed_capacity(token) # escalated to 60 + clock.now += 61.0 + token = breaker.acquire_attempt() + assert admitted(token) + breaker.attempt_succeeded(token) # closed, escalation reset + breaker.record_generation_failure() # re-open: cooldown back at base + wait = breaker.acquire_attempt() + assert isinstance(wait, float) and wait <= 30.0 + + +def test_registry_get_or_create_and_keying(monkeypatch): + registry = mb.ModelBreakerRegistry() + a = registry.for_model("codex", "gpt-5.6-sol") + b = registry.for_model("codex", "gpt-5.6-sol") + c = registry.for_model("codex", "gpt-5.6-terra") + assert a is b + assert a is not c + assert a.name == "codex:gpt-5.6-sol" + # Effective-model keying: sol capacity trouble never blocks terra. + a.record_generation_failure() + assert a.state == "open" + assert c.state == "closed" + snap = registry.snapshot() + assert set(snap) == {"codex:gpt-5.6-sol", "codex:gpt-5.6-terra"} + assert snap["codex:gpt-5.6-sol"]["state"] == "open" + + +def test_registry_tolerates_missing_identity(): + registry = mb.ModelBreakerRegistry() + breaker = registry.for_model("", "") + assert breaker.name == "unknown:unknown" diff --git a/tests/test_recovery_policy.py b/tests/test_recovery_policy.py new file mode 100644 index 00000000..fa7d9fae --- /dev/null +++ b/tests/test_recovery_policy.py @@ -0,0 +1,203 @@ +"""Pins for the shared deadline-based recovery (src/llm/recovery.py).""" + +from __future__ import annotations + +import asyncio +import time + +import pytest + +from src.llm.circuit_breaker import CircuitOpenError +from src.llm.errors import ( + LLMAuthError, + LLMCapacityError, + LLMRateLimitError, + LLMRequestError, + LLMTransportError, +) +from src.llm.model_breaker import ModelBreakerRegistry +from src.llm.recovery import RecoveryPolicy, generate_with_recovery + +FAST = RecoveryPolicy( + deadline_seconds=0.5, backoff_base=0.01, backoff_cap=0.05, retry_after_cap=0.2 +) + + +def scripted(*outcomes): + """Attempt callable yielding each outcome in order (exception → raise).""" + calls = {"n": 0} + + async def attempt(): + idx = min(calls["n"], len(outcomes) - 1) + calls["n"] += 1 + outcome = outcomes[idx] + if isinstance(outcome, BaseException): + raise outcome + return outcome + + return attempt, calls + + +async def test_success_first_try(): + attempt, calls = scripted("ok") + result = await generate_with_recovery(attempt, policy=FAST) + assert result == "ok" + assert calls["n"] == 1 + + +async def test_capacity_then_success_is_recovered(): + attempt, calls = scripted(LLMCapacityError("overloaded"), "ok") + result = await generate_with_recovery(attempt, policy=FAST) + assert result == "ok" + assert calls["n"] == 2 + + +@pytest.mark.parametrize( + "exc", [LLMTransportError("stream died"), CircuitOpenError("codex_api", 0.01)] +) +async def test_transport_and_client_breaker_are_retryable(exc): + attempt, calls = scripted(exc, "ok") + result = await generate_with_recovery(attempt, policy=FAST) + assert result == "ok" + assert calls["n"] == 2 + + +async def test_deadline_exhaustion_raises_last_capacity_error(): + attempt, calls = scripted(LLMCapacityError("overloaded")) + with pytest.raises(LLMCapacityError): + await generate_with_recovery(attempt, policy=FAST) + assert calls["n"] >= 2 # kept trying until the budget ran out + + +async def test_exactly_one_generation_failure_recorded_on_exhaustion(): + registry = ModelBreakerRegistry(cooldown_base=100.0) # opens, waits are long + breaker = registry.for_model("codex", "gpt-5.6-sol") + attempt, calls = scripted(LLMCapacityError("overloaded")) + with pytest.raises(LLMCapacityError): + await generate_with_recovery(attempt, policy=FAST, breaker=breaker) + # Many failed ATTEMPTS, exactly ONE generation failure counted. + assert calls["n"] >= 2 + assert breaker.snapshot()["failed_generations"] == 1 + + +@pytest.mark.parametrize( + "exc", + [ + LLMAuthError("401 no healthy account"), + LLMRequestError("400 bad model"), + LLMRateLimitError("429 all accounts limited"), + ], +) +async def test_fast_fail_classes_escape_immediately(exc): + registry = ModelBreakerRegistry() + breaker = registry.for_model("codex", "gpt-5.6-sol") + attempt, calls = scripted(exc) + started = time.monotonic() + with pytest.raises(type(exc)): + await generate_with_recovery(attempt, policy=FAST, breaker=breaker) + assert calls["n"] == 1 + assert time.monotonic() - started < 0.2 # no budget spent + assert breaker.snapshot()["failed_generations"] == 0 + + +async def test_unclassified_exception_is_never_retried(): + # The agents-path bug (bare except retrying programming defects) must + # not be spread into the shared policy. + attempt, calls = scripted(ValueError("defect")) + with pytest.raises(ValueError): + await generate_with_recovery(attempt, policy=FAST) + assert calls["n"] == 1 + + +async def test_retry_after_is_honoured_as_wait_floor(): + attempt, _ = scripted(LLMCapacityError("overloaded", retry_after=0.15), "ok") + started = time.monotonic() + result = await generate_with_recovery(attempt, policy=FAST) + assert result == "ok" + assert time.monotonic() - started >= 0.14 + + +async def test_retry_after_is_capped(): + # A pathological server suggestion must not exceed retry_after_cap. + attempt, _ = scripted(LLMCapacityError("overloaded", retry_after=500.0), "ok") + started = time.monotonic() + result = await generate_with_recovery(attempt, policy=FAST) + assert result == "ok" + assert time.monotonic() - started < 0.45 # capped at 0.2, not 500 + + +async def test_zero_budget_gets_one_attempt_then_raises(): + # Restart-with-expired-deadline semantics: the budget bounds WAITING; + # a single attempt is still made, then the failure surfaces. + attempt, calls = scripted(LLMCapacityError("overloaded")) + with pytest.raises(LLMCapacityError): + await generate_with_recovery(attempt, policy=FAST, deadline_seconds=0.0) + assert calls["n"] == 1 + + +async def test_cancellation_interrupts_a_long_wait_promptly(): + cancel = asyncio.Event() + attempt, _ = scripted(LLMCapacityError("overloaded", retry_after=10.0)) + policy = RecoveryPolicy( + deadline_seconds=30.0, backoff_base=5.0, backoff_cap=10.0, retry_after_cap=10.0 + ) + + async def fire_cancel(): + await asyncio.sleep(0.05) + cancel.set() + + started = time.monotonic() + canceller = asyncio.create_task(fire_cancel()) + with pytest.raises(asyncio.CancelledError): + await generate_with_recovery( + attempt, policy=policy, cancel_event=cancel + ) + await canceller + assert time.monotonic() - started < 1.0 # did not sit out the 10s wait + + +async def test_preset_cancel_prevents_any_attempt(): + cancel = asyncio.Event() + cancel.set() + attempt, calls = scripted("ok") + with pytest.raises(asyncio.CancelledError): + await generate_with_recovery(attempt, policy=FAST, cancel_event=cancel) + assert calls["n"] == 0 + + +async def test_open_breaker_is_waited_through_then_probe_succeeds(): + registry = ModelBreakerRegistry(cooldown_base=0.05, cooldown_cap=0.1) + breaker = registry.for_model("codex", "gpt-5.6-sol") + breaker.record_generation_failure() # open + assert breaker.state == "open" + attempt, calls = scripted("ok") + policy = RecoveryPolicy(deadline_seconds=2.0, backoff_base=0.01, backoff_cap=0.05) + result = await generate_with_recovery(attempt, policy=policy, breaker=breaker) + assert result == "ok" + assert calls["n"] == 1 + assert breaker.state == "closed" # probe success closed it + + +async def test_budget_consumed_by_breaker_waits_counts_generation_failure(): + registry = ModelBreakerRegistry(cooldown_base=60.0) + breaker = registry.for_model("codex", "gpt-5.6-sol") + breaker.record_generation_failure() # open, cooldown far exceeds budget + attempt, calls = scripted("ok") + policy = RecoveryPolicy(deadline_seconds=0.1, backoff_base=0.01, backoff_cap=0.05) + with pytest.raises(LLMCapacityError): + await generate_with_recovery(attempt, policy=policy, breaker=breaker) + assert calls["n"] == 0 # never admitted + assert breaker.snapshot()["failed_generations"] == 2 # the wait-exhaust counted + + +async def test_on_wait_hook_is_called_and_fault_tolerant(): + seen = [] + + def hook(wait, remaining, error): + seen.append((wait, remaining, type(error).__name__)) + raise RuntimeError("hook bug must not break recovery") + + attempt, _ = scripted(LLMCapacityError("overloaded"), "ok") + result = await generate_with_recovery(attempt, policy=FAST, on_wait=hook) + assert result == "ok" + assert seen and seen[0][2] == "LLMCapacityError" From 62e053145626629d86e352acb7aaaa8c749a3a52 Mon Sep 17 00:00:00 2001 From: ATCharpentier Date: Thu, 30 Jul 2026 13:45:50 -0400 Subject: [PATCH 03/18] feat: shared recovery at all three call sites + guard capacity fix MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Chat (_call_llm): the CircuitOpenError-only sleep-and-retry pair is replaced by generate_with_recovery — capacity/transport/breaker-open retried up to the generation deadline, cancellable by /stop (graceful stop path, not an error turn), fast-fail classes unchanged. Loops (_call_loop_llm): same recovery in-iteration with retry_circuit_open= False — CircuitOpenError still escapes to the loop manager (pinned asymmetry preserved; gateway bypass untouched per RFC-001 §4.3). Agents: recovery moved INSIDE the iteration callbacks (_agent_generate, breaker keyed on the agent's effective model); the manager's bare-except single retry ladder is REMOVED — lifetime exhaustion still wins as TIMEOUT, everything else fails fast with readable errors. Gateway: capacity/CircuitOpen no longer feed the sticky guard counter (the double penalty that latched llm_* UNAVAILABLE until restart); capacity marks a self-expiring transient DEGRADED instead. New notify_generation_success(provider) clears a latched guard from bypass paths, driven by immutable response provenance only (missing provenance = no-op, never a guess). SubsystemGuard grows transient-DEGRADED with lazy expiry; real failures/successes still own the counter. Characterization pins amended DELIBERATELY (agent ladder tests, chat breaker-retry pins) — called out for review. New tests/test_guard_capacity.py. --- src/agents/manager.py | 106 +++----- src/discord/llm_gateway.py | 58 +++- src/discord/native_tools/agents_tasks.py | 63 ++++- src/discord/tool_loop.py | 111 +++++--- src/health/subsystem_guard.py | 86 +++++- src/llm/recovery.py | 14 +- tests/characterization/test_chat_tool_loop.py | 13 +- tests/test_agent_lifecycle.py | 253 +++++------------- tests/test_guard_capacity.py | 180 +++++++++++++ tests/test_native_agents_tasks.py | 46 +++- tests/test_typing_resilience.py | 10 +- 11 files changed, 612 insertions(+), 328 deletions(-) create mode 100644 tests/test_guard_capacity.py diff --git a/src/agents/manager.py b/src/agents/manager.py index 9bf0de92..82a797dc 100644 --- a/src/agents/manager.py +++ b/src/agents/manager.py @@ -31,7 +31,10 @@ WAIT_POLL_INTERVAL = 2 # poll interval for wait_for_agents ITERATION_CB_TIMEOUT = 120 # 2 min timeout per LLM call TOOL_EXEC_TIMEOUT = 300 # 5 min timeout per tool execution -MAX_RECOVERY_ATTEMPTS = 1 # retries before transitioning to FAILED +# (The manager-level MAX_RECOVERY_ATTEMPTS retry ladder was removed +# 2026-07-30: transient-failure recovery now lives inside the iteration +# callback via src/llm/recovery.py. AgentInfo.recovery_attempts remains for +# API/trajectory shape compatibility and stays 0.) MAX_NESTING_DEPTH = 2 # default max sub-agent depth (root=0) MAX_CHILDREN_PER_AGENT = 3 # max direct children one agent can spawn @@ -922,7 +925,6 @@ def _check_lifetime() -> bool: agent.transition(AgentState.EXECUTING, f"iteration {iteration + 1}") agent.last_activity = time.time() agent.iteration_count = iteration + 1 - agent.recovery_attempts = 0 # per-iteration recovery budget iter_start = time.time() # Call LLM with recovery support @@ -1118,15 +1120,16 @@ async def _call_llm_with_recovery( system_prompt: str, tools: list[dict], ) -> dict | None: - """Call LLM with single-retry recovery on transient errors. + """Call the LLM for one agent iteration. - Each wait is bounded by the agent's snapshotted iteration_timeout, capped - at the remaining lifetime — the hard deadline must hold even during a - long LLM await (the between-iteration lifetime check alone would let a - quiet agent overrun it). + Transient-failure recovery (capacity/transport/breaker waits) lives + INSIDE the iteration callback via the shared deadline-based policy + (``src/llm/recovery.py``) — the old manager-level bare-``except`` single + retry ladder retried programming defects and is deliberately gone + (design settled with Odin, 2026-07-30). What remains here is the wall: + the agent's snapshotted iteration_timeout capped at remaining lifetime + hard-bounds the callback INCLUDING any recovery waits. - On first failure: EXECUTING → RECOVERING → EXECUTING (retry). - On second failure: EXECUTING → FAILED. Returns the LLM response dict, or None if agent reached terminal state. """ remaining = _remaining_lifetime(agent) @@ -1139,71 +1142,36 @@ async def _call_llm_with_recovery( iteration_callback(agent.messages, system_prompt, tools), timeout=call_timeout, ) - except (TimeoutError, Exception) as first_err: - is_timeout = isinstance(first_err, asyncio.TimeoutError) - if is_timeout and _remaining_lifetime(agent) <= 0: + except TimeoutError: + if _remaining_lifetime(agent) <= 0: # The wait was lifetime-capped and the deadline has passed: # this is lifetime exhaustion, not a stuck LLM call. _lifetime_timeout(agent) return None - err_desc = (f"LLM timeout after {int(call_timeout)}s" if is_timeout - else f"LLM error: {first_err}") - - if agent.recovery_attempts < MAX_RECOVERY_ATTEMPTS: - agent.recovery_attempts += 1 - agent.transition(AgentState.RECOVERING, err_desc) - log.warning( - "Agent %s recovering (attempt %d): %s", - agent.id, - agent.recovery_attempts, - err_desc, - ) - - retry_delay = 1 - if hasattr(first_err, "retry_after"): - retry_delay = min(first_err.retry_after, 90.0) - log.info("Agent %s: circuit breaker wait %.0fs", agent.id, retry_delay) - # The recovery sleep must not outlive the deadline either — a - # 90s breaker wait with 5s of lifetime left sleeps 5s, and the - # retry-entry check below then settles it. - remaining = _remaining_lifetime(agent) - if remaining <= 0: - _lifetime_timeout(agent) - return None - await asyncio.sleep(min(retry_delay, remaining)) - - agent.transition(AgentState.EXECUTING, "retry after recovery") - - remaining = _remaining_lifetime(agent) - if remaining <= 0: - _lifetime_timeout(agent) - return None - retry_timeout = min(agent.iteration_timeout, remaining) - try: - return await asyncio.wait_for( - iteration_callback(agent.messages, system_prompt, tools), - timeout=retry_timeout, - ) - except (TimeoutError, Exception) as retry_err: - retry_is_timeout = isinstance(retry_err, asyncio.TimeoutError) - if retry_is_timeout and _remaining_lifetime(agent) <= 0: - _lifetime_timeout(agent) - return None - # str(asyncio.TimeoutError()) is EMPTY — always store the - # formatted description, never the bare exception string. - retry_desc = (f"retry timed out after {int(retry_timeout)}s" - if retry_is_timeout else f"retry failed: {retry_err}") - log.error("Agent %s recovery failed: %s", agent.id, retry_desc) - agent.transition(AgentState.FAILED, retry_desc) - agent.error = retry_desc - agent.ended_at = time.time() - return None - else: - log.error("Agent %s LLM call failed (no retries left): %s", agent.id, err_desc) - agent.transition(AgentState.FAILED, err_desc) - agent.error = err_desc - agent.ended_at = time.time() + # str(asyncio.TimeoutError()) is EMPTY — always store the formatted + # description, never the bare exception string. + err_desc = f"LLM timeout after {int(call_timeout)}s" + log.error("Agent %s LLM call failed: %s", agent.id, err_desc) + agent.transition(AgentState.FAILED, err_desc) + agent.error = err_desc + agent.ended_at = time.time() + return None + except Exception as exc: + if _remaining_lifetime(agent) <= 0: + # Lifetime exhaustion wins over failure classification (the + # v3.59.0 rule: exhaustion is TIMEOUT, never FAILED). + _lifetime_timeout(agent) return None + # Typed fast-fail (auth / malformed request / quota-exhausted after + # rotation) or a programming defect: neither earns a manager-level + # retry — transient classes were already retried inside the callback + # for up to the generation deadline. + err_desc = f"LLM error: {exc}" if str(exc) else f"LLM error: {type(exc).__name__}" + log.error("Agent %s LLM call failed (no retry): %s", agent.id, err_desc) + agent.transition(AgentState.FAILED, err_desc) + agent.error = err_desc + agent.ended_at = time.time() + return None def _get_last_progress(agent: AgentInfo) -> str: diff --git a/src/discord/llm_gateway.py b/src/discord/llm_gateway.py index 4817b1d9..416da7d1 100644 --- a/src/discord/llm_gateway.py +++ b/src/discord/llm_gateway.py @@ -26,7 +26,11 @@ from collections.abc import Callable from ..llm import CodexChatClient, KimiClient, OllamaClient +from ..llm.circuit_breaker import CircuitOpenError from ..llm.codex_auth import CodexAuthPool +from ..llm.errors import LLMCapacityError +from ..llm.model_breaker import ModelBreakerRegistry, ModelCapacityBreaker +from ..llm.recovery import RecoveryPolicy from ..odin_log import get_logger log = get_logger("discord") @@ -45,6 +49,8 @@ def __init__( cost_tracker, sessions, reflector, + model_breakers: ModelBreakerRegistry | None = None, + recovery_policy_source: Callable[[], RecoveryPolicy] | None = None, ) -> None: self.get_config = get_config self.codex_client = codex_client @@ -55,6 +61,11 @@ def __init__( self.cost_tracker = cost_tracker self.sessions = sessions self.reflector = reflector + # Capacity-breaker registry: BotServices-owned in production so state + # survives client rebuilds/live reloads; a private default keeps + # existing constructions (tests) working. + self.model_breakers = model_breakers or ModelBreakerRegistry() + self._recovery_policy_source = recovery_policy_source or RecoveryPolicy self.provider_lock = asyncio.Lock() # Auxiliary live-reload state: a monotonic generation guards against a # candidate built under a config a concurrent reload has since @@ -81,6 +92,40 @@ def active_client(self): return self.kimi_client return self.codex_client + def recovery_policy(self) -> RecoveryPolicy: + """The live recovery policy (config-backed via wiring).""" + return self._recovery_policy_source() + + def capacity_breaker_for(self, model: str | None = None) -> ModelCapacityBreaker: + """Model-scoped capacity breaker for the active provider. + + ``model`` must be the EFFECTIVE model of the request when the caller + overrides it (agents); defaults to the active client's model. + """ + provider_cfg = getattr(self.get_config(), "llm_provider", None) + active = provider_cfg.active_provider if provider_cfg else "codex" + effective = model + if not effective: + client = self.active_client + effective = getattr(client, "model", None) if client is not None else None + return self.model_breakers.for_model(active, str(effective or "unknown")) + + def notify_generation_success(self, provider: str | None) -> None: + """Success signal from a path that bypasses ``call_with_tools`` + (agents, autonomous loops). + + This is the production ``mark_available`` wiring: a latched + ``llm_*`` guard key can never see a gateway success (check() blocks + the call), but bypass-path successes prove the subsystem is fine. + + ``provider`` MUST come from the response's immutable provenance + (``provenance_provider``) — never from whichever provider is active + after the await. Missing provenance is a no-op, never a guess. + """ + if self.subsystem_guard is None or not provider: + return + self.subsystem_guard.record_success(f"llm_{provider}") + def wire_callbacks(self) -> None: """Attach LLM-backed compaction and reflection callbacks using the active provider.""" @@ -628,7 +673,18 @@ async def call_with_tools( ) except Exception as exc: if self.subsystem_guard is not None: - self.subsystem_guard.record_failure(guard_key, str(exc)) + if isinstance(exc, (LLMCapacityError, CircuitOpenError)): + # Capacity (and the client breaker's echoes of it) + # never feeds the sticky failure counter — the + # model-scoped breaker owns capacity admission, and + # counting both was the double penalty that let an + # outage latch the guard UNAVAILABLE until restart. + # Visibility only: transient DEGRADED, self-expiring. + self.subsystem_guard.mark_degraded_transient( + guard_key, str(exc)[:200], expires_in=120.0 + ) + else: + self.subsystem_guard.record_failure(guard_key, str(exc)) raise if self.subsystem_guard is not None: self.subsystem_guard.record_success(guard_key) diff --git a/src/discord/native_tools/agents_tasks.py b/src/discord/native_tools/agents_tasks.py index 7ed78f1a..aad59def 100644 --- a/src/discord/native_tools/agents_tasks.py +++ b/src/discord/native_tools/agents_tasks.py @@ -22,6 +22,7 @@ from ...agents.manager import AGENT_BLOCKED_TOOLS, filter_agent_tools from ...async_utils import fire_and_forget +from ...llm.recovery import generate_with_recovery from ...odin_log import get_logger from ..background_task import ( MAX_STEPS, @@ -463,6 +464,46 @@ def _handle_list_loops(self) -> str: # --- Agent tool handlers --- + async def _agent_generate( + self, + client, + *, + messages: list[dict], + sys_prompt: str, + tool_defs: list[dict], + agent_effort, + resolved_model, + ): + """One agent LLM generation through the shared recovery policy. + + Replaces the old manager-level bare-``except`` single retry: transient + classes (capacity/transport/open breaker) recover here for up to the + generation deadline; auth/malformed/quota-exhausted fail fast. The + model-scoped breaker is keyed on the agent's EFFECTIVE model, so a + fleet mixing models coordinates capacity per model. The manager's + iteration wall (wait_for) hard-bounds this call INCLUDING recovery + waits; cancellation propagates and releases any held probe. + """ + breaker = self._llm_gateway.capacity_breaker_for(resolved_model) + policy = self._llm_gateway.recovery_policy() + + async def _attempt(): + return await client.chat_with_tools( + messages=messages, + system=sys_prompt, + tools=tool_defs, + reasoning_effort=agent_effort, + model=resolved_model, + ) + + resp = await generate_with_recovery(_attempt, policy=policy, breaker=breaker) + # Bypass-path success clears a latched llm_* guard key — provenance + # only, never the post-await active provider. + self._llm_gateway.notify_generation_success( + getattr(resp, "provenance_provider", None) + ) + return resp + async def _handle_spawn_agent(self, message: object, inp: dict) -> str: """Spawn an autonomous agent for a sub-task. @@ -530,12 +571,13 @@ async def _iteration_cb( self._get_config(), client, model_override=model_override, effort_override=effort_override, ) - resp = await client.chat_with_tools( + resp = await self._agent_generate( + client, messages=messages, - system=sys_prompt, - tools=tool_defs, - reasoning_effort=agent_effort, - model=resolved_model, + sys_prompt=sys_prompt, + tool_defs=tool_defs, + agent_effort=agent_effort, + resolved_model=resolved_model, ) return { "text": resp.text, @@ -842,12 +884,13 @@ async def _iteration_cb(messages, sys, tool_defs): self._get_config(), client, model_override=model_override, effort_override=effort_override, ) - resp = await client.chat_with_tools( + resp = await self._agent_generate( + client, messages=messages, - system=sys, - tools=tool_defs, - reasoning_effort=agent_effort, - model=resolved_model, + sys_prompt=sys, + tool_defs=tool_defs, + agent_effort=agent_effort, + resolved_model=resolved_model, ) return { "text": resp.text or "", diff --git a/src/discord/tool_loop.py b/src/discord/tool_loop.py index 2bc0ef56..bda9b379 100644 --- a/src/discord/tool_loop.py +++ b/src/discord/tool_loop.py @@ -37,6 +37,7 @@ import discord from ..llm import CircuitOpenError +from ..llm.recovery import generate_with_recovery from ..llm.secret_scrubber import scrub_output_secrets from ..observability.correlation import get_turn, set_turn from ..odin_log import get_logger @@ -701,55 +702,56 @@ def _maybe_compress(self, st: _ChatTurn) -> None: ) async def _call_llm(self, st: _ChatTurn): - """Guarded LLM call with typing indicator and circuit-breaker recovery. + """Guarded LLM call with typing indicator and deadline-based recovery. Returns ("ok", llm_resp) or ("done", ). + + Transient failures — capacity (SSE overload inside a 200), transport, + and an open client breaker — are retried by the shared recovery + policy for up to the configured generation deadline (default 5 min), + waiting through breakers and honouring retry_after. Auth failures, + malformed requests, and quota exhaustion (429 after the client's own + account rotation) still fail fast, exactly as before. /stop + interrupts any recovery wait immediately. """ _channel_id = str(st.message.channel.id) + breaker = self._llm_gateway.capacity_breaker_for() + policy = self._llm_gateway.recovery_policy() + + def _on_wait(wait: float, remaining: float, error: BaseException) -> None: + log.info( + "LLM recovery (%s): waiting %.1fs, %.0fs of generation budget left", + type(error).__name__, wait, remaining, + ) + + async def _attempt(): + return await self._llm_gateway.call_with_tools( + messages=st.messages, + system=st.system_prompt, + tools=st.tools or [], + user_id=st.user_id, + channel_id=_channel_id, + tools_used=st.tools_used_in_loop, + ) + # Typing is best-effort (shared helper): a typing failure — setup or # cleanup — must never fail the call or misclassify provider errors. async with _best_effort_typing(st.message.channel): try: - llm_resp = await self._llm_gateway.call_with_tools( - messages=st.messages, - system=st.system_prompt, - tools=st.tools or [], - user_id=st.user_id, - channel_id=_channel_id, - tools_used=st.tools_used_in_loop, + llm_resp = await generate_with_recovery( + _attempt, + policy=policy, + breaker=breaker, + cancel_event=st._cancel, + on_wait=_on_wait, ) - except CircuitOpenError as coe: - wait_secs = min(coe.retry_after, 90.0) - log.info( - "Circuit breaker open for %s, waiting %.0fs for recovery", - coe.provider, - wait_secs, - ) - await asyncio.sleep(wait_secs) - try: - llm_resp = await self._llm_gateway.call_with_tools( - messages=st.messages, - system=st.system_prompt, - tools=st.tools or [], - user_id=st.user_id, - channel_id=_channel_id, - tools_used=st.tools_used_in_loop, - ) - except Exception as retry_err: - await self._turn_recorder._save_turn_trajectory( - st._trajectory, error=str(retry_err), trace=st.trace - ) - self._clear_active(st) - return ( - "done", - ( - f"LLM API error (circuit breaker recovery failed): {retry_err}", - False, - True, - st.tools_used_in_loop, - False, - ), - ) + except asyncio.CancelledError: + if st._cancel.is_set(): + # /stop fired during a recovery wait — the graceful stop + # path, not an error turn (same contract as the + # loop-head cancel checks). + return ("done", self._stopped(st, "llm_recovery")) + raise except Exception as api_err: err_msg = str(api_err) or f"{type(api_err).__name__} (no message)" log.error("LLM API call failed: %s", err_msg, exc_info=True) @@ -1491,17 +1493,32 @@ async def _finish_loop( return outcome_text async def _call_loop_llm(self, st: _LoopTurn): - """LLM call for one loop iteration. CircuitOpenError re-raises to the - loop manager (policy asymmetry — the manager owns backoff). + """LLM call for one loop iteration with deadline-based recovery. + + Typed capacity/transport failures are retried in-iteration by the + shared recovery policy; CircuitOpenError still re-raises to the loop + manager (policy asymmetry — the manager owns backoff between + iterations). The gateway bypass itself is unchanged (RFC-001 §4.3). Returns ("ok", response) or ("done", ). """ - try: - response = await self._llm_gateway.active_client.chat_with_tools( + breaker = self._llm_gateway.capacity_breaker_for() + policy = self._llm_gateway.recovery_policy() + + async def _attempt(): + return await self._llm_gateway.active_client.chat_with_tools( messages=st.messages, system=st.system_prompt, tools=st.tools or [], ) + + try: + response = await generate_with_recovery( + _attempt, + policy=policy, + breaker=breaker, + retry_circuit_open=False, + ) except CircuitOpenError: raise except Exception as e: @@ -1516,6 +1533,12 @@ async def _call_loop_llm(self, st: _LoopTurn): error_text=str(e), ), ) + # Bypass-path success: clear a latched llm_* guard key using the + # response's immutable provenance (never the post-await active + # provider) — the production mark_available wiring. + self._llm_gateway.notify_generation_success( + getattr(response, "provenance_provider", None) + ) return ("ok", response) def _record_loop_iteration(self, st: _LoopTurn, response, _iteration: int) -> bool: diff --git a/src/health/subsystem_guard.py b/src/health/subsystem_guard.py index 198d613c..4686a228 100644 --- a/src/health/subsystem_guard.py +++ b/src/health/subsystem_guard.py @@ -45,6 +45,11 @@ class SubsystemInfo: last_failure_at: float = 0.0 last_success_at: float = 0.0 registered_at: float = field(default_factory=time.monotonic) + # Set only by mark_degraded_transient: a monotonic expiry after which a + # visibility-only DEGRADED lapses back to AVAILABLE on the next read. + # None means any DEGRADED state is counter-driven and recovers only via + # record_success/mark_available. + transient_until: float | None = None def to_dict(self) -> dict[str, Any]: d: dict[str, Any] = { @@ -60,6 +65,8 @@ def to_dict(self) -> dict[str, Any]: d["last_failure_at"] = self.last_failure_at if self.last_success_at: d["last_success_at"] = self.last_success_at + if self.transient_until is not None: + d["transient"] = True return d @@ -205,21 +212,44 @@ def registered(self) -> list[str]: # ── State queries ──────────────────────────────────────────────── + def _resolve(self, name: str) -> SubsystemInfo | None: + """Look up *name*, lapsing an expired transient DEGRADED first. + + Transient degradation is visibility-only (capacity storms); it must + never require an explicit success to clear, so expiry is applied + lazily on every read path. + """ + info = self._subsystems.get(name) + if ( + info is not None + and info.state == SubsystemState.DEGRADED + and info.transient_until is not None + and time.monotonic() >= info.transient_until + ): + info.state = SubsystemState.AVAILABLE + info.transient_until = None + self.stats.record_transition( + name, SubsystemState.DEGRADED, SubsystemState.AVAILABLE, + "transient degradation expired", + ) + log.info("Subsystem %r transient degradation expired — AVAILABLE", name) + return info + def get_state(self, name: str) -> SubsystemState | None: """Return the current state of *name*, or None if unregistered.""" - info = self._subsystems.get(name) + info = self._resolve(name) return info.state if info else None def is_available(self, name: str) -> bool: """True only when the subsystem is fully AVAILABLE.""" - info = self._subsystems.get(name) + info = self._resolve(name) if info is None: return True # unregistered = not tracked = assume available return info.state == SubsystemState.AVAILABLE def is_usable(self, name: str) -> bool: """True when AVAILABLE or DEGRADED (partial functionality OK).""" - info = self._subsystems.get(name) + info = self._resolve(name) if info is None: return True return info.state != SubsystemState.UNAVAILABLE @@ -233,7 +263,7 @@ def check(self, name: str) -> str | None: if err: return err """ - info = self._subsystems.get(name) + info = self._resolve(name) if info is None: self.stats.record_check(blocked=False) return None @@ -253,6 +283,7 @@ def mark_available(self, name: str) -> None: return old = info.state info.consecutive_failures = 0 + info.transient_until = None if old == SubsystemState.AVAILABLE: return info.state = SubsystemState.AVAILABLE @@ -265,6 +296,9 @@ def mark_degraded(self, name: str, reason: str = "") -> None: if info is None: return old = info.state + # An explicit mark supersedes a transient one: from here only a + # success/mark_available clears it, never expiry. + info.transient_until = None if old == SubsystemState.DEGRADED: return info.state = SubsystemState.DEGRADED @@ -282,11 +316,47 @@ def mark_unavailable(self, name: str, reason: str = "") -> None: if old == SubsystemState.UNAVAILABLE: return info.state = SubsystemState.UNAVAILABLE + info.transient_until = None if reason: info.last_failure_reason = reason self.stats.record_transition(name, old, SubsystemState.UNAVAILABLE, reason or "manual") log.warning("Subsystem %r marked UNAVAILABLE: %s", name, reason or "manual") + def mark_degraded_transient( + self, name: str, reason: str = "", *, expires_in: float = 120.0 + ) -> None: + """Visibility-only DEGRADED that lapses on its own after *expires_in*. + + Used for model-capacity storms: capacity is excluded from the sticky + failure counter entirely (the model breaker owns capacity admission), + but the status API should still show the provider as degraded while + the storm lasts. Repeated calls refresh the expiry. + + Never downgrades UNAVAILABLE, and never converts a counter-driven + DEGRADED into a self-clearing one — real degradation still requires + a real success to clear. + """ + info = self._subsystems.get(name) + if info is None: + return + if info.state == SubsystemState.UNAVAILABLE: + return + if info.state == SubsystemState.DEGRADED and info.transient_until is None: + return + old = info.state + info.transient_until = time.monotonic() + max(1.0, expires_in) + if reason: + info.last_failure_reason = reason + if old != SubsystemState.DEGRADED: + info.state = SubsystemState.DEGRADED + self.stats.record_transition( + name, old, SubsystemState.DEGRADED, reason or "transient degradation" + ) + log.info( + "Subsystem %r transiently DEGRADED for %.0fs: %s", + name, expires_in, reason or "capacity", + ) + # ── Automatic threshold-based transitions ──────────────────────── def record_failure(self, name: str, reason: str = "") -> SubsystemState: @@ -302,6 +372,9 @@ def record_failure(self, name: str, reason: str = "") -> SubsystemState: info.consecutive_failures += 1 info.total_failures += 1 info.last_failure_at = time.monotonic() + # A real failure supersedes any transient marker: from here the + # counter owns the state and only a success clears it. + info.transient_until = None if reason: info.last_failure_reason = reason @@ -342,6 +415,7 @@ def record_success(self, name: str) -> SubsystemState: info.consecutive_failures = 0 info.total_successes += 1 info.last_success_at = time.monotonic() + info.transient_until = None old = info.state if old != SubsystemState.AVAILABLE: @@ -359,10 +433,12 @@ def record_success(self, name: str) -> SubsystemState: # ── Observability ──────────────────────────────────────────────── def get_subsystem(self, name: str) -> SubsystemInfo | None: - return self._subsystems.get(name) + return self._resolve(name) def get_status(self) -> dict[str, Any]: """Full status snapshot for the REST API.""" + for name in list(self._subsystems): + self._resolve(name) # lapse expired transient degradations subsystems = [info.to_dict() for info in self._subsystems.values()] available_count = sum(1 for i in self._subsystems.values() if i.state == SubsystemState.AVAILABLE) diff --git a/src/llm/recovery.py b/src/llm/recovery.py index 04aa514d..548fbd28 100644 --- a/src/llm/recovery.py +++ b/src/llm/recovery.py @@ -115,6 +115,7 @@ async def generate_with_recovery( deadline_seconds: float | None = None, cancel_event: asyncio.Event | None = None, on_wait: Callable[[float, float, BaseException], None] | None = None, + retry_circuit_open: bool = True, ) -> T: """Run one logical LLM generation with deadline-based recovery. @@ -122,6 +123,11 @@ async def generate_with_recovery( error escapes, cancellation fires, or the budget is exhausted — in which case the last error is re-raised (after counting one generation failure on the model breaker when that error is capacity-class). + + ``retry_circuit_open=False`` makes ``CircuitOpenError`` fast-fail + instead of being waited through — the autonomous-loop path re-raises it + to the loop manager, which owns pacing between iterations (pinned + policy asymmetry). """ budget = policy.deadline_seconds if deadline_seconds is None else deadline_seconds deadline = time.monotonic() + max(0.0, budget) @@ -180,7 +186,13 @@ def _exhausted() -> BaseException: if breaker is not None and token is not None: breaker.abandon(token) raise - except (LLMTransportError, CircuitOpenError) as exc: + except CircuitOpenError as exc: + if breaker is not None and token is not None: + breaker.abandon(token) + if not retry_circuit_open: + raise + last_error = exc + except LLMTransportError as exc: if breaker is not None and token is not None: breaker.abandon(token) last_error = exc diff --git a/tests/characterization/test_chat_tool_loop.py b/tests/characterization/test_chat_tool_loop.py index bde0f83b..c7057024 100644 --- a/tests/characterization/test_chat_tool_loop.py +++ b/tests/characterization/test_chat_tool_loop.py @@ -257,7 +257,10 @@ async def test_llm_api_error_returns_error_tuple(self): assert is_error is True assert text == "LLM API error: boom" - async def test_circuit_open_waits_and_retries_once(self): + async def test_circuit_open_waits_and_retries(self): + # Amended 2026-07-30: the single hardcoded breaker retry became the + # shared deadline recovery — an open client breaker is waited + # through, then retried; the recovered response is not an error. bot, fake = build( [ CircuitOpenError("codex", 0.0), @@ -269,7 +272,11 @@ async def test_circuit_open_waits_and_retries_once(self): assert text == "recovered" assert len(fake.calls) == 2 - async def test_circuit_open_retry_failure_is_error(self): + async def test_unclassified_error_after_breaker_retry_is_error(self): + # Amended 2026-07-30: a breaker-open retry that then hits an + # UNCLASSIFIED exception fast-fails through the shared recovery — + # the old "(circuit breaker recovery failed)" wording is gone; the + # plain terminal error shape is the contract now. bot, fake = build( [ CircuitOpenError("codex", 0.0), @@ -278,7 +285,7 @@ async def test_circuit_open_retry_failure_is_error(self): ) text, _, is_error, _, _ = await run_loop(bot, FakeMessage("go")) assert is_error is True - assert text.startswith("LLM API error (circuit breaker recovery failed):") + assert text == "LLM API error: still down" # --------------------------------------------------------------------------- diff --git a/tests/test_agent_lifecycle.py b/tests/test_agent_lifecycle.py index f221fa3f..5d47ba66 100644 --- a/tests/test_agent_lifecycle.py +++ b/tests/test_agent_lifecycle.py @@ -16,7 +16,6 @@ ACTIVE_STATES, ITERATION_CB_TIMEOUT, MAX_AGENT_LIFETIME, - MAX_RECOVERY_ATTEMPTS, TERMINAL_STATES, VALID_TRANSITIONS, AgentInfo, @@ -733,7 +732,14 @@ async def flaky_iter(msgs, sys, tools): # wait_for is called with a coroutine, need a different approach pass - async def test_recovery_retry_succeeds(self): + # Deliberate pin amendments (2026-07-30, design settled with Odin): + # transient-failure recovery moved INSIDE the iteration callback via the + # shared deadline policy (src/llm/recovery.py). The manager keeps only + # the wall; the old EXECUTING→RECOVERING→EXECUTING single-retry ladder — + # which retried programming defects via bare except — is gone. These + # tests pin its absence. + + async def test_timeout_fails_without_manager_retry(self): agent = AgentInfo( id="r3", label="test", goal="test", channel_id="c1", requester_id="u1", requester_name="user", @@ -742,92 +748,34 @@ async def test_recovery_retry_succeeds(self): agent.transition(AgentState.EXECUTING) call_count = 0 - original_wait_for = asyncio.wait_for - async def counting_wait_for(coro, *, timeout=None): + async def counting_timeout(coro, *, timeout=None): nonlocal call_count call_count += 1 - if call_count == 1: - try: - coro.close() - except: # noqa: E722 — deliberate maximum-breadth catch; narrowing changes cancellation semantics - pass - raise TimeoutError() - return await original_wait_for(coro, timeout=timeout) - - iter_cb = AsyncMock(return_value={"text": "recovered", "tool_calls": []}) - - with patch("src.agents.manager.asyncio.wait_for", side_effect=counting_wait_for): - with patch("src.agents.manager.asyncio.sleep", new_callable=AsyncMock): - result = await _call_llm_with_recovery(agent, iter_cb, "sys", []) - - assert result is not None - assert result["text"] == "recovered" - assert agent.recovery_attempts == 1 - assert agent.state == AgentState.EXECUTING - # History should show EXECUTING → RECOVERING → EXECUTING - h = agent.state_history - states = [(t.from_state, t.to_state) for t in h] - assert (AgentState.EXECUTING, AgentState.RECOVERING) in states - assert (AgentState.RECOVERING, AgentState.EXECUTING) in states - - async def test_recovery_retry_fails(self): - agent = AgentInfo( - id="r4", label="test", goal="test", - channel_id="c1", requester_id="u1", requester_name="user", - ) - agent.transition(AgentState.READY) - agent.transition(AgentState.EXECUTING) - - asyncio.wait_for - - async def always_timeout(coro, *, timeout=None): try: coro.close() except: # noqa: E722 — deliberate maximum-breadth catch; narrowing changes cancellation semantics pass raise TimeoutError() - iter_cb = AsyncMock(return_value={"text": "x", "tool_calls": []}) - - with patch("src.agents.manager.asyncio.wait_for", side_effect=always_timeout): - with patch("src.agents.manager.asyncio.sleep", new_callable=AsyncMock): - result = await _call_llm_with_recovery(agent, iter_cb, "sys", []) - - assert result is None - assert agent.state == AgentState.FAILED - assert agent.recovery_attempts == 1 - assert agent.ended_at is not None - - async def test_no_recovery_when_attempts_exhausted(self): - agent = AgentInfo( - id="r5", label="test", goal="test", - channel_id="c1", requester_id="u1", requester_name="user", - ) - agent.transition(AgentState.READY) - agent.transition(AgentState.EXECUTING) - agent.recovery_attempts = MAX_RECOVERY_ATTEMPTS # already used up - - async def timeout_coro(coro, *, timeout=None): - try: - coro.close() - except: # noqa: E722 — deliberate maximum-breadth catch; narrowing changes cancellation semantics - pass - raise TimeoutError() - - iter_cb = AsyncMock(return_value={"text": "x", "tool_calls": []}) + iter_cb = AsyncMock(return_value={"text": "never", "tool_calls": []}) - with patch("src.agents.manager.asyncio.wait_for", side_effect=timeout_coro): + with patch("src.agents.manager.asyncio.wait_for", side_effect=counting_timeout): result = await _call_llm_with_recovery(agent, iter_cb, "sys", []) assert result is None + assert call_count == 1 # exactly one call — no ladder assert agent.state == AgentState.FAILED - # No recovery transition in history (directly to FAILED) - h = agent.state_history - recovery_transitions = [t for t in h if t.to_state == AgentState.RECOVERING] - assert len(recovery_transitions) == 0 + assert agent.recovery_attempts == 0 + assert agent.ended_at is not None + recovery_transitions = [ + t for t in agent.state_history if t.to_state == AgentState.RECOVERING + ] + assert recovery_transitions == [] - async def test_exception_triggers_recovery(self): + async def test_exception_fails_fast_no_second_call(self): + # The old ladder would have made a second call and "recovered" — + # that second call must never happen now. agent = AgentInfo( id="r6", label="test", goal="test", channel_id="c1", requester_id="u1", requester_name="user", @@ -852,11 +800,13 @@ async def err_then_ok(coro, *, timeout=None): iter_cb = AsyncMock(return_value={"text": "ok", "tool_calls": []}) with patch("src.agents.manager.asyncio.wait_for", side_effect=err_then_ok): - with patch("src.agents.manager.asyncio.sleep", new_callable=AsyncMock): - result = await _call_llm_with_recovery(agent, iter_cb, "sys", []) + result = await _call_llm_with_recovery(agent, iter_cb, "sys", []) - assert result is not None - assert agent.recovery_attempts == 1 + assert result is None + assert call_count == 1 + assert agent.state == AgentState.FAILED + assert agent.error == "LLM error: transient error" + assert agent.recovery_attempts == 0 # --------------------------------------------------------------------------- @@ -864,7 +814,10 @@ async def err_then_ok(coro, *, timeout=None): # --------------------------------------------------------------------------- class TestRunAgentRecovery: - async def test_full_recovery_lifecycle(self): + async def test_llm_failure_fails_agent_without_manager_retry(self): + # Deliberate amendment (2026-07-30): the old ladder made this agent + # COMPLETE via a manager retry; recovery now lives inside the + # iteration callback, so a failure that escapes it fails the agent. agent = AgentInfo( id="fr1", label="test", goal="test", channel_id="c1", requester_id="u1", requester_name="user", @@ -885,18 +838,17 @@ async def first_timeout(coro, *, timeout=None): raise TimeoutError() return await original_wait_for(coro, timeout=timeout) - iter_cb = AsyncMock(return_value={"text": "recovered", "tool_calls": []}) + iter_cb = AsyncMock(return_value={"text": "would recover", "tool_calls": []}) tool_cb = AsyncMock() with patch("src.agents.manager.asyncio.wait_for", side_effect=first_timeout): - with patch("src.agents.manager.asyncio.sleep", new_callable=AsyncMock): - await _run_agent(agent, "sys", [], iter_cb, tool_cb) + await _run_agent(agent, "sys", [], iter_cb, tool_cb) - assert agent.state == AgentState.COMPLETED - assert agent.recovery_attempts == 1 - h = agent.state_history - states = [t.to_state for t in h] - assert AgentState.RECOVERING in states + assert agent.state == AgentState.FAILED + assert call_count == 1 + assert agent.recovery_attempts == 0 + states = [t.to_state for t in agent.state_history] + assert AgentState.RECOVERING not in states async def test_failed_recovery_lifecycle(self): agent = AgentInfo( @@ -1298,8 +1250,13 @@ def test_terminal_states_match_legacy(self): for state in TERMINAL_STATES: assert state.value in _TERMINAL_STATUSES - def test_max_recovery_attempts_constant(self): - assert MAX_RECOVERY_ATTEMPTS == 1 + def test_manager_retry_ladder_removed(self): + # Deliberate amendment (2026-07-30): transient recovery moved into + # the iteration callback (src/llm/recovery.py); the manager-level + # ladder and its constant must stay gone. + import src.agents.manager as manager_mod + + assert not hasattr(manager_mod, "MAX_RECOVERY_ATTEMPTS") def test_state_enum_is_str(self): for state in AgentState: @@ -1514,7 +1471,7 @@ async def timeout_past_deadline(coro, *, timeout=None): recoveries = [t for t in agent.state_history if t.to_state == AgentState.RECOVERING] assert recoveries == [] - async def test_retry_timeout_stores_readable_error(self): + async def test_timeout_stores_readable_error(self): agent = _exec_agent(iteration_timeout=77.0, max_lifetime=100000.0) async def always_timeout(coro, *, timeout=None): @@ -1526,32 +1483,30 @@ async def always_timeout(coro, *, timeout=None): iter_cb = AsyncMock(return_value={"text": "x", "tool_calls": []}) with patch("src.agents.manager.asyncio.wait_for", side_effect=always_timeout): - with patch("src.agents.manager.asyncio.sleep", new_callable=AsyncMock): - result = await _call_llm_with_recovery(agent, iter_cb, "sys", []) + result = await _call_llm_with_recovery(agent, iter_cb, "sys", []) assert result is None assert agent.state == AgentState.FAILED # str(asyncio.TimeoutError()) is "" — the stored error must never be empty - assert agent.error == "retry timed out after 77s" + assert agent.error == "LLM timeout after 77s" - async def test_exhausted_attempts_store_readable_error(self): + async def test_empty_string_exception_stores_readable_error(self): agent = _exec_agent(iteration_timeout=77.0, max_lifetime=100000.0) - agent.recovery_attempts = MAX_RECOVERY_ATTEMPTS - async def always_timeout(coro, *, timeout=None): + async def raise_bare(coro, *, timeout=None): try: coro.close() except Exception: pass - raise TimeoutError() + raise ConnectionError() # str() == "" iter_cb = AsyncMock(return_value={"text": "x", "tool_calls": []}) - with patch("src.agents.manager.asyncio.wait_for", side_effect=always_timeout): + with patch("src.agents.manager.asyncio.wait_for", side_effect=raise_bare): result = await _call_llm_with_recovery(agent, iter_cb, "sys", []) assert result is None assert agent.state == AgentState.FAILED - assert agent.error == "LLM timeout after 77s" + assert agent.error == "LLM error: ConnectionError" class TestLifetimeEnforcement: @@ -1686,65 +1641,36 @@ async def save(self, turn): class TestRetryPathDeadline: - async def test_lifetime_exhausted_at_retry_entry(self): - """First failure is transient, but the deadline passes during the - recovery sleep — the retry must not start; lifetime TIMEOUT wins.""" + async def test_lifetime_exhaustion_wins_over_exception_class(self): + """A failing call that consumed the lifetime is lifetime exhaustion + (TIMEOUT), never FAILED — the v3.59.0 rule holds on the single-call + path now that the manager retry ladder is gone.""" agent = _exec_agent(iteration_timeout=900.0, max_lifetime=100.0) agent.created_at = time.time() - 50 - async def fail_first(coro, *, timeout=None): + async def fail_and_expire(coro, *, timeout=None): try: coro.close() except Exception: pass + agent.created_at -= 100 # the failed call ate the lifetime raise ConnectionError("transient") - async def sleep_past_deadline(_delay): - agent.created_at -= 100 # recovery sleep consumed the lifetime - - iter_cb = AsyncMock(return_value={"text": "x", "tool_calls": []}) - with patch("src.agents.manager.asyncio.wait_for", side_effect=fail_first): - with patch("src.agents.manager.asyncio.sleep", side_effect=sleep_past_deadline): - result = await _call_llm_with_recovery(agent, iter_cb, "sys", []) - - assert result is None - assert agent.state == AgentState.TIMEOUT - assert "lifetime exceeded" in agent.state_history[-1].reason - - async def test_retry_timeout_at_deadline_is_lifetime(self): - """A retry that times out exactly at the deadline is lifetime - exhaustion (TIMEOUT), not a FAILED recovery.""" - agent = _exec_agent(iteration_timeout=900.0, max_lifetime=100.0) - agent.created_at = time.time() - 50 - calls = 0 - - async def error_then_deadline_timeout(coro, *, timeout=None): - nonlocal calls - calls += 1 - try: - coro.close() - except Exception: - pass - if calls == 1: - raise ConnectionError("transient") - agent.created_at -= 100 # retry wait ran past the deadline - raise TimeoutError() - iter_cb = AsyncMock(return_value={"text": "x", "tool_calls": []}) - with patch("src.agents.manager.asyncio.wait_for", - side_effect=error_then_deadline_timeout): - with patch("src.agents.manager.asyncio.sleep", new_callable=AsyncMock): - result = await _call_llm_with_recovery(agent, iter_cb, "sys", []) + with patch("src.agents.manager.asyncio.wait_for", side_effect=fail_and_expire): + result = await _call_llm_with_recovery(agent, iter_cb, "sys", []) assert result is None - assert calls == 2 assert agent.state == AgentState.TIMEOUT assert "lifetime exceeded" in agent.state_history[-1].reason class TestHardDeadlineDuringToolsAndSleep: """PR #226 review blockers: the deadline must hold BETWEEN tool calls - (no floored bonus budget per tool) and across the recovery sleep.""" + (no floored bonus budget per tool). The recovery-sleep half of the + original class is gone with the manager retry ladder (2026-07-30) — + waits-bounded-by-remaining-budget now lives in src/llm/recovery.py and + is pinned in tests/test_recovery_policy.py.""" async def test_expired_agent_stops_at_next_tool(self): """Odin's repro shape: tiny lifetime + three slow tools ran ~3s on @@ -1788,64 +1714,27 @@ async def slow_tool(name, tool_input): assert len(turn.iterations[0].tool_calls) == 1 assert "timed out" in turn.iterations[0].tool_results[0]["result"] - async def test_recovery_sleep_capped_at_remaining(self): - """A 90s circuit-breaker wait with ~1s of lifetime left must sleep - the remainder, not the full 90s.""" - agent = _exec_agent(iteration_timeout=900.0, max_lifetime=100.0) - agent.created_at = time.time() - 99 # ~1s remaining - - class _BreakerError(Exception): - retry_after = 90.0 - - calls = 0 - - async def fail_then_ok(coro, *, timeout=None): - nonlocal calls - calls += 1 - try: - coro.close() - except Exception: - pass - if calls == 1: - raise _BreakerError("breaker open") - return {"text": "recovered", "tool_calls": []} - - slept: list[float] = [] - - async def capture_sleep(delay): - slept.append(delay) - - iter_cb = AsyncMock(return_value={"text": "x", "tool_calls": []}) - with patch("src.agents.manager.asyncio.wait_for", side_effect=fail_then_ok): - with patch("src.agents.manager.asyncio.sleep", side_effect=capture_sleep): - result = await _call_llm_with_recovery(agent, iter_cb, "sys", []) - - assert result is not None and result["text"] == "recovered" - assert len(slept) == 1 - assert slept[0] <= 1.01 # capped at the remainder, not 90 - - async def test_deadline_passed_before_recovery_sleep(self): - """If the first call consumed the lifetime, the agent times out - BEFORE the recovery sleep — no sleep at all.""" + async def test_no_manager_sleep_exists_on_failure_path(self): + """The manager never sleeps anymore: a failure either times out the + lifetime or fails the agent immediately — no recovery sleep at all.""" agent = _exec_agent(iteration_timeout=900.0, max_lifetime=100.0) agent.created_at = time.time() - 50 - async def fail_and_expire(coro, *, timeout=None): + async def fail_once(coro, *, timeout=None): try: coro.close() except Exception: pass - agent.created_at -= 100 # the failed call ate the lifetime raise ConnectionError("transient") sleep_mock = AsyncMock() iter_cb = AsyncMock(return_value={"text": "x", "tool_calls": []}) - with patch("src.agents.manager.asyncio.wait_for", side_effect=fail_and_expire): + with patch("src.agents.manager.asyncio.wait_for", side_effect=fail_once): with patch("src.agents.manager.asyncio.sleep", sleep_mock): result = await _call_llm_with_recovery(agent, iter_cb, "sys", []) assert result is None - assert agent.state == AgentState.TIMEOUT + assert agent.state == AgentState.FAILED sleep_mock.assert_not_awaited() async def test_expiry_during_final_tool_of_final_iteration_is_timeout(self): diff --git a/tests/test_guard_capacity.py b/tests/test_guard_capacity.py new file mode 100644 index 00000000..802da1e7 --- /dev/null +++ b/tests/test_guard_capacity.py @@ -0,0 +1,180 @@ +"""Pins for the SubsystemGuard capacity fix (design settled with Odin 2026-07-30). + +The trap being closed: check() short-circuits BEFORE the call, record_success +can then never fire, and mark_available had no production caller — so ten +guard increments latched an llm_* key UNAVAILABLE until restart. With 5-min +recovery, a capacity outage would have reached that latch in 1-2 turns +("fail quicker" — forbidden). The fix, pinned here: + +- typed capacity failures and client-breaker echoes NEVER feed the sticky + counter (the model-scoped breaker owns capacity admission); +- capacity marks a self-expiring transient DEGRADED (visibility only); +- bypass-path successes (agents/loops) clear a latched guard via + notify_generation_success, driven by immutable response provenance only. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +import src.health.subsystem_guard as sg +from src.discord.llm_gateway import LLMGateway +from src.health.subsystem_guard import SubsystemGuard, SubsystemState +from src.llm.circuit_breaker import CircuitOpenError +from src.llm.errors import LLMCapacityError + + +def _cfg(active="codex"): + return SimpleNamespace(llm_provider=SimpleNamespace(active_provider=active)) + + +def _guard(): + guard = SubsystemGuard() + guard.register("llm_codex") + return guard + + +def _gw(client, guard): + return LLMGateway( + get_config=_cfg, + codex_client=client, + ollama_client=None, + kimi_client=None, + subsystem_guard=guard, + auxiliary_llm_client=None, + cost_tracker=None, + sessions=MagicMock(), + reflector=MagicMock(), + ) + + +async def _call(gw): + return await gw.call_with_tools(messages=[], system="s", tools=[]) + + +class TestCapacityExcludedFromGuardCounter: + async def test_capacity_error_never_increments_counter(self): + guard = _guard() + client = SimpleNamespace( + chat_with_tools=AsyncMock(side_effect=LLMCapacityError("overloaded")), + model="gpt-5.6-sol", + ) + gw = _gw(client, guard) + for _ in range(25): # far past both thresholds + with pytest.raises(LLMCapacityError): + await _call(gw) + info = guard.get_subsystem("llm_codex") + assert info.consecutive_failures == 0 + # Visibility: transiently DEGRADED, never UNAVAILABLE, never blocking. + assert guard.get_state("llm_codex") == SubsystemState.DEGRADED + assert guard.check("llm_codex") is None + + async def test_circuit_open_error_never_increments_counter(self): + guard = _guard() + client = SimpleNamespace( + chat_with_tools=AsyncMock(side_effect=CircuitOpenError("codex_api", 30.0)), + model="gpt-5.6-sol", + ) + gw = _gw(client, guard) + for _ in range(25): + with pytest.raises(CircuitOpenError): + await _call(gw) + assert guard.get_subsystem("llm_codex").consecutive_failures == 0 + assert guard.get_state("llm_codex") != SubsystemState.UNAVAILABLE + + async def test_real_failures_still_count_and_latch(self): + guard = _guard() + client = SimpleNamespace( + chat_with_tools=AsyncMock(side_effect=RuntimeError("real breakage")), + model="gpt-5.6-sol", + ) + gw = _gw(client, guard) + for _ in range(10): + with pytest.raises(RuntimeError): + await _call(gw) + assert guard.get_state("llm_codex") == SubsystemState.UNAVAILABLE + assert guard.check("llm_codex") is not None + + +class TestTransientDegraded: + def test_transient_expires_on_read(self, monkeypatch): + clock = {"now": 1000.0} + monkeypatch.setattr( + sg, "time", SimpleNamespace(monotonic=lambda: clock["now"], time=lambda: 0.0) + ) + guard = _guard() + guard.mark_degraded_transient("llm_codex", "capacity", expires_in=120.0) + assert guard.get_state("llm_codex") == SubsystemState.DEGRADED + clock["now"] += 121.0 + assert guard.get_state("llm_codex") == SubsystemState.AVAILABLE + assert guard.get_subsystem("llm_codex").transient_until is None + + def test_transient_never_downgrades_unavailable(self): + guard = _guard() + guard.mark_unavailable("llm_codex", "dead") + guard.mark_degraded_transient("llm_codex", "capacity") + assert guard.get_state("llm_codex") == SubsystemState.UNAVAILABLE + + def test_transient_never_converts_counter_degraded(self, monkeypatch): + clock = {"now": 1000.0} + monkeypatch.setattr( + sg, "time", SimpleNamespace(monotonic=lambda: clock["now"], time=lambda: 0.0) + ) + guard = _guard() + for _ in range(3): # counter-driven DEGRADED + guard.record_failure("llm_codex", "real") + assert guard.get_state("llm_codex") == SubsystemState.DEGRADED + guard.mark_degraded_transient("llm_codex", "capacity", expires_in=1.0) + clock["now"] += 100.0 + # A real degradation must NOT lapse on a transient expiry. + assert guard.get_state("llm_codex") == SubsystemState.DEGRADED + + def test_real_failure_supersedes_transient(self, monkeypatch): + clock = {"now": 1000.0} + monkeypatch.setattr( + sg, "time", SimpleNamespace(monotonic=lambda: clock["now"], time=lambda: 0.0) + ) + guard = _guard() + guard.mark_degraded_transient("llm_codex", "capacity", expires_in=60.0) + guard.record_failure("llm_codex", "real") + clock["now"] += 120.0 + # The counter owns it now; expiry must not clear it... + assert guard.get_subsystem("llm_codex").transient_until is None + # ...and a success does. + guard.record_success("llm_codex") + assert guard.get_state("llm_codex") == SubsystemState.AVAILABLE + + def test_status_snapshot_flags_transient(self): + guard = _guard() + guard.mark_degraded_transient("llm_codex", "capacity") + snap = guard.get_status() + entry = next(s for s in snap["subsystems"] if s["name"] == "llm_codex") + assert entry.get("transient") is True + + +class TestNotifyGenerationSuccess: + def test_bypass_success_clears_latched_guard(self): + guard = _guard() + for _ in range(10): + guard.record_failure("llm_codex", "boom") + assert guard.get_state("llm_codex") == SubsystemState.UNAVAILABLE + gw = _gw(SimpleNamespace(model="gpt-5.6-sol"), guard) + gw.notify_generation_success("codex") + assert guard.get_state("llm_codex") == SubsystemState.AVAILABLE + assert guard.check("llm_codex") is None + + def test_missing_provenance_is_a_noop_never_a_guess(self): + guard = _guard() + for _ in range(10): + guard.record_failure("llm_codex", "boom") + gw = _gw(SimpleNamespace(model="gpt-5.6-sol"), guard) + gw.notify_generation_success(None) + gw.notify_generation_success("") + assert guard.get_state("llm_codex") == SubsystemState.UNAVAILABLE + + def test_no_guard_is_tolerated(self): + gw = _gw(SimpleNamespace(model="gpt-5.6-sol"), None) + gw.notify_generation_success("codex") # must not raise diff --git a/tests/test_native_agents_tasks.py b/tests/test_native_agents_tasks.py index 10e83ac1..80ab4cb9 100644 --- a/tests/test_native_agents_tasks.py +++ b/tests/test_native_agents_tasks.py @@ -21,6 +21,28 @@ AgentTaskTools, _parse_spawn_overrides, ) +from src.llm.model_breaker import ModelBreakerRegistry +from src.llm.recovery import RecoveryPolicy + + +def _fake_gateway(client): + """Gateway fake carrying the recovery surface the callbacks now use. + + The breaker registry is real (per-fake, isolated) so the callbacks' + admission/resolution protocol actually runs; the fast policy keeps any + would-be retry from sleeping meaningfully in tests. + """ + registry = ModelBreakerRegistry() + gw = SimpleNamespace(active_client=client) + gw.capacity_breaker_for = lambda model=None: registry.for_model( + "codex", str(model or getattr(client, "model", None) or "unknown") + ) + gw.recovery_policy = lambda: RecoveryPolicy( + deadline_seconds=0.2, backoff_base=0.01, backoff_cap=0.02, retry_after_cap=0.05 + ) + gw.success_notices = [] + gw.notify_generation_success = gw.success_notices.append + return gw def _cfg(tools_enabled=True): @@ -37,7 +59,7 @@ def _cfg(tools_enabled=True): def _deps(**ov): d: dict[str, Any] = dict( get_config=lambda: _cfg(), - llm_gateway=SimpleNamespace(active_client=object()), + llm_gateway=_fake_gateway(object()), channel_state=SimpleNamespace(background_tasks={}, background_tasks_max=50), tool_executor=MagicMock(), skill_manager=MagicMock(), @@ -223,7 +245,7 @@ def test_list_loops(self): class TestSpawnAgent: async def test_validation(self): assert "required" in await _tools()._handle_spawn_agent(_message(), {"label": "a"}) - t = _tools(llm_gateway=SimpleNamespace(active_client=None)) + t = _tools(llm_gateway=_fake_gateway(None)) assert "not available" in await t._handle_spawn_agent( _message(), {"label": "a", "goal": "g"}) @@ -309,7 +331,7 @@ def _spawned_callback(self, agent_effort, client): cfg = _cfg() cfg.openai_codex = SimpleNamespace(agent_reasoning_effort=agent_effort) t = _tools(get_config=lambda: cfg, - llm_gateway=SimpleNamespace(active_client=client)) + llm_gateway=_fake_gateway(client)) t._agent_manager.spawn.return_value = "agent-e" t._agent_manager._agents = {} return t @@ -340,7 +362,7 @@ async def test_callback_reads_config_at_call_time(self): cfg = _cfg() cfg.openai_codex = SimpleNamespace(agent_reasoning_effort=None) t = _tools(get_config=lambda: cfg, - llm_gateway=SimpleNamespace(active_client=client)) + llm_gateway=_fake_gateway(client)) t._agent_manager.spawn.return_value = "agent-e" t._agent_manager._agents = {} await t._handle_spawn_agent(_message(), {"label": "w", "goal": "g"}) @@ -384,7 +406,7 @@ async def test_loop_spawn_callback_passes_effort(self): cfg = _cfg() cfg.openai_codex = SimpleNamespace(agent_reasoning_effort="medium") t = _tools(get_config=lambda: cfg, - llm_gateway=SimpleNamespace(active_client=client)) + llm_gateway=_fake_gateway(client)) t._loop_manager._loops = {"L1": SimpleNamespace( status="running", requester_id="1", requester_name="u", goal="loop goal", iteration_count=1)} @@ -407,7 +429,7 @@ def _spawned(self, cfg_codex, client): cfg = _cfg() cfg.openai_codex = cfg_codex t = _tools(get_config=lambda: cfg, - llm_gateway=SimpleNamespace(active_client=client)) + llm_gateway=_fake_gateway(client)) t._agent_manager.spawn.return_value = "agent-m" t._agent_manager._agents = {} return t @@ -556,7 +578,7 @@ async def test_loop_spawn_callback_same_treatment(self): agent_model="gpt-5.6-luna", model="gpt-5.6-sol") t = _tools(get_config=lambda: cfg, - llm_gateway=SimpleNamespace(active_client=client)) + llm_gateway=_fake_gateway(client)) t._loop_manager._loops = {"L1": SimpleNamespace( status="running", requester_id="1", requester_name="u", goal="loop goal", iteration_count=1)} @@ -613,7 +635,7 @@ async def test_spawn_override_wins_over_config(self): # Per-spawn overrides are only accepted when the axis is Auto. cfg.openai_codex = self._codex_cfg(agent_model="auto", agent_effort="auto") t = _tools(get_config=lambda: cfg, - llm_gateway=SimpleNamespace(active_client=client)) + llm_gateway=_fake_gateway(client)) t._agent_manager.spawn.return_value = "agent-o" t._agent_manager._agents = {} await t._handle_spawn_agent( @@ -634,7 +656,7 @@ async def test_no_override_inherits_config(self): cfg = _cfg() cfg.openai_codex = self._codex_cfg(agent_model="gpt-5.6-terra", agent_effort="high") t = _tools(get_config=lambda: cfg, - llm_gateway=SimpleNamespace(active_client=client)) + llm_gateway=_fake_gateway(client)) t._agent_manager.spawn.return_value = "agent-i" t._agent_manager._agents = {} await t._handle_spawn_agent(_message(), {"label": "w", "goal": "g"}) @@ -652,7 +674,7 @@ async def test_invalid_effort_rejects_without_spawning(self): # Effort axis Auto so the field is accepted for value-validation. cfg.openai_codex = self._codex_cfg(agent_effort="auto") t = _tools(get_config=lambda: cfg, - llm_gateway=SimpleNamespace(active_client=client)) + llm_gateway=_fake_gateway(client)) t._agent_manager._agents = {} out = await t._handle_spawn_agent( _message(), {"label": "w", "goal": "g", "reasoning_effort": "ultra"}) @@ -665,7 +687,7 @@ async def test_loop_per_task_overrides_and_batch_rejection(self): # Both axes Auto so per-task model + effort overrides are accepted. cfg.openai_codex = self._codex_cfg(agent_model="auto", agent_effort="auto") t = _tools(get_config=lambda: cfg, - llm_gateway=SimpleNamespace(active_client=client)) + llm_gateway=_fake_gateway(client)) t._loop_manager._loops = {"L1": SimpleNamespace( status="running", requester_id="1", requester_name="u", goal="loop goal", iteration_count=1)} @@ -778,7 +800,7 @@ async def test_spawn_loop_agents_validation(self): _message(), {"loop_id": "L1", "tasks": ["t"]}) async def test_spawn_loop_agents_no_client(self): - t = _tools(llm_gateway=SimpleNamespace(active_client=None)) + t = _tools(llm_gateway=_fake_gateway(None)) t._loop_manager._loops = {"L1": self._running_loop()} assert "not available" in await t._handle_spawn_loop_agents( _message(), {"loop_id": "L1", "tasks": ["t"]}) diff --git a/tests/test_typing_resilience.py b/tests/test_typing_resilience.py index 98897fb4..07de3427 100644 --- a/tests/test_typing_resilience.py +++ b/tests/test_typing_resilience.py @@ -361,7 +361,15 @@ async def test_call_llm_with_dead_typing_endpoint(self): async def _cwt(**kwargs): return resp - runner._llm_gateway = SimpleNamespace(call_with_tools=_cwt) + from src.llm.model_breaker import ModelBreakerRegistry + from src.llm.recovery import RecoveryPolicy + + registry = ModelBreakerRegistry() + runner._llm_gateway = SimpleNamespace( + call_with_tools=_cwt, + capacity_breaker_for=lambda model=None: registry.for_model("codex", "m"), + recovery_policy=RecoveryPolicy, + ) kind, val = await runner._call_llm(st) assert (kind, val) == ("ok", resp) assert ch.typing_calls == 1 From 633e6176612dda6bf916c8221a18bf4776857fe6 Mon Sep 17 00:00:00 2001 From: ATCharpentier Date: Thu, 30 Jul 2026 13:51:50 -0400 Subject: [PATCH 04/18] feat: durable turn-state store (checkpoints + side-effect ledger) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit src/turn_state/store.py: one SQLite DB (data/turn_state/turns.sqlite3 in production wiring) holding checkpoints, ledger rows, leases, and terminal tombstones so state transitions are atomic; content-addressed blob dir beside it (tmp+rename). KnowledgeStore conventions: WAL, busy_timeout, one connection, write lock, available-property degradation. Fencing is three-part per Odin's round-3 clarification: every state write is conditional on (turn_generation, revision, lease_token); an expired owner replaying an old revision or a cleared token gets StaleTurnError and must stop. Heartbeats extend the lease WITHOUT bumping the revision or advancing last_progress_at; checkpoint writes advance last_progress_at only when progressed=True (recovery waits never fake progress). Recovery deadlines persist as absolute UTC. Ledger: PREPARED→RUNNING→APPLIED/DEFINITELY_FAILED with OUTCOME_UNKNOWN for interrupted work; intent validation rejects empty/duplicate tool_call_ids BEFORE execution; effect fingerprint stored as secondary reconciliation evidence only. Boot sweep suspends stale-lease ACTIVE turns and marks their in-flight ops OUTCOME_UNKNOWN (never rerun). Three TTL clocks: 24h resumable from last REAL progress, 7d diagnostic payloads then tombstone, 90d ledger — OUTCOME_UNKNOWN/MANUAL rows never auto-expire. Terminal payloads compact immediately; identity + ledger tombstones remain. Fail-closed contract: init failure disables the feature loudly; mid-turn write failures raise TurnStateUnavailable. Also: two stale-ladder agent pins amended deliberately (test_recovery), loop-path fake gateway updated (test_trajectory_completeness). --- src/turn_state/__init__.py | 23 + src/turn_state/store.py | 700 ++++++++++++++++++++++++++ tests/test_recovery.py | 59 +-- tests/test_trajectory_completeness.py | 13 +- tests/test_turn_state_store.py | 395 +++++++++++++++ 5 files changed, 1143 insertions(+), 47 deletions(-) create mode 100644 src/turn_state/__init__.py create mode 100644 src/turn_state/store.py create mode 100644 tests/test_turn_state_store.py diff --git a/src/turn_state/__init__.py b/src/turn_state/__init__.py new file mode 100644 index 00000000..11c4f533 --- /dev/null +++ b/src/turn_state/__init__.py @@ -0,0 +1,23 @@ +from .store import ( + LedgerIntentError, + OpState, + StaleTurnError, + TurnKey, + TurnLease, + TurnStateStore, + TurnStateUnavailableError, + TurnStatus, + effect_fingerprint, +) + +__all__ = [ + "LedgerIntentError", + "OpState", + "StaleTurnError", + "TurnKey", + "TurnLease", + "TurnStateStore", + "TurnStateUnavailableError", + "TurnStatus", + "effect_fingerprint", +] diff --git a/src/turn_state/store.py b/src/turn_state/store.py new file mode 100644 index 00000000..63142dec --- /dev/null +++ b/src/turn_state/store.py @@ -0,0 +1,700 @@ +"""Durable turn state: checkpoints, side-effect ledger, leases, tombstones. + +One SQLite database (design settled with Odin, 2026-07-30) so checkpoint and +ledger transitions can be atomic — deliberately NOT sessions (compacted +user-facing history), trajectories/audit (append-only observability), or the +command workspace (disposable). Large payloads live beside it as +content-addressed blobs. + +Identity and fencing: + +- Primary key ``source + channel_id + message_id``; a content hash is NOT + identity (identical text collides, edits change it). +- A random ``turn_generation`` minted at first admission, a monotonic + ``revision``, and a ``lease_token``. Every state write is conditional on + ALL THREE (round-3 clarification: generation+revision alone lets an + expired owner win a write before the new owner advances the revision). + A fence mismatch raises :class:`StaleTurnError` — the caller lost + ownership and must stop immediately. + +Fail-closed contract: + +- Store-init failure at startup → ``available`` False, checkpointing is off + for the process (loud log; turns run legacy). +- Once a turn runs WITH durability, any persistence failure raises + :class:`TurnStateUnavailableError` and the caller must halt further generation + or mutation — never silently fall back to the legacy discard path. + +Time discipline (round-3): the recovery budget's live arithmetic is +monotonic and lives in ``src/llm/recovery.py``; THIS store persists absolute +UTC timestamps only (``time.time()``), so a reboot reconstructs remaining +budgets without monotonic-epoch nonsense. Lease heartbeats and checkpoint +rewrites do NOT advance ``last_progress_at`` — only real progress +(``progressed=True``) does. + +Three retention clocks (``ttl_sweep``): resumable 24h from last real +progress; diagnostic payloads 7d then tombstone; ledger rows ≥90d — with +``OUTCOME_UNKNOWN``/``MANUAL_RESOLUTION_REQUIRED`` never auto-expiring. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import secrets +import sqlite3 +import threading +import time +from dataclasses import dataclass +from pathlib import Path + +from ..odin_log import get_logger + +log = get_logger("turn_state") + +SCHEMA_VERSION = 1 + +DEFAULT_LEASE_TTL = 120.0 +DEFAULT_RESUME_TTL_HOURS = 24.0 +DEFAULT_PAYLOAD_RETENTION_DAYS = 7.0 +DEFAULT_LEDGER_RETENTION_DAYS = 90.0 + + +class TurnStatus: + ACTIVE = "ACTIVE" + SUSPENDED = "SUSPENDED" + TERMINAL_COMPLETED = "TERMINAL_COMPLETED" + TERMINAL_CANCELLED = "TERMINAL_CANCELLED" + TERMINAL_FAILED = "TERMINAL_FAILED" + TERMINAL_REJECTED = "TERMINAL_REJECTED" + TERMINAL_EXPIRED = "TERMINAL_EXPIRED" + + TERMINAL = frozenset({ + TERMINAL_COMPLETED, TERMINAL_CANCELLED, TERMINAL_FAILED, + TERMINAL_REJECTED, TERMINAL_EXPIRED, + }) + + +class OpState: + PREPARED = "PREPARED" + RUNNING = "RUNNING" + APPLIED = "APPLIED" + DEFINITELY_FAILED = "DEFINITELY_FAILED" + OUTCOME_UNKNOWN = "OUTCOME_UNKNOWN" + RECONCILED_APPLIED = "RECONCILED_APPLIED" + RECONCILED_NOT_APPLIED = "RECONCILED_NOT_APPLIED" + MANUAL_RESOLUTION_REQUIRED = "MANUAL_RESOLUTION_REQUIRED" + + # Rows in these states are evidence of possibly-unreconciled external + # effects and must never expire automatically. + NEVER_EXPIRE = frozenset({OUTCOME_UNKNOWN, MANUAL_RESOLUTION_REQUIRED}) + + +class TurnStateUnavailableError(RuntimeError): + """A durability write failed while durability was enabled — fail closed.""" + + +class StaleTurnError(RuntimeError): + """A fenced write lost: another owner holds this turn now. Stop.""" + + +class LedgerIntentError(ValueError): + """Malformed tool-call intents (empty/duplicate ids) — fail BEFORE execution.""" + + +@dataclass(frozen=True) +class TurnKey: + source: str + channel_id: str + message_id: str + + +@dataclass +class TurnLease: + """Ownership handle for one admitted turn. ``revision`` tracks the last + successfully fenced write and is advanced by the store.""" + + key: TurnKey + generation: str + token: str + revision: int + + +def effect_fingerprint(tool_name: str, tool_input: dict) -> str: + """Handler-derived effect fingerprint (secondary reconciliation evidence, + NEVER the primary dedup key — deliberate identical invocations are + legitimate). v1 fingerprints the effective invocation; per-handler + resolved-resource identity is a declared follow-up.""" + try: + canonical = json.dumps( + {"tool": tool_name, "input": tool_input}, sort_keys=True, default=str + ) + except Exception: + canonical = f"{tool_name}:{tool_input!r}" + return hashlib.sha256(canonical.encode("utf-8", "replace")).hexdigest() + + +_DDL = """ +CREATE TABLE IF NOT EXISTS turns ( + source TEXT NOT NULL, + channel_id TEXT NOT NULL, + message_id TEXT NOT NULL, + turn_generation TEXT NOT NULL, + revision INTEGER NOT NULL DEFAULT 0, + lease_token TEXT, + lease_expires_at REAL, + status TEXT NOT NULL, + recovery_deadline_utc REAL, + last_progress_at REAL NOT NULL, + created_at REAL NOT NULL, + suspended_at REAL, + guild_id TEXT, + user_id TEXT, + content_digest TEXT, + code_version TEXT, + schema_version INTEGER NOT NULL, + prompt_policy_hash TEXT, + tool_catalog_hash TEXT, + session_snapshot TEXT, + payload TEXT, + PRIMARY KEY (source, channel_id, message_id) +); +CREATE INDEX IF NOT EXISTS idx_turns_status ON turns(status); +CREATE TABLE IF NOT EXISTS operations ( + source TEXT NOT NULL, + channel_id TEXT NOT NULL, + message_id TEXT NOT NULL, + turn_generation TEXT NOT NULL, + generation_seq INTEGER NOT NULL, + tool_call_id TEXT NOT NULL, + state TEXT NOT NULL, + tool_name TEXT NOT NULL, + iteration INTEGER, + effect_fingerprint TEXT, + result TEXT, + created_at REAL NOT NULL, + updated_at REAL NOT NULL, + PRIMARY KEY (source, channel_id, message_id, turn_generation, + generation_seq, tool_call_id) +); +CREATE INDEX IF NOT EXISTS idx_ops_state ON operations(state); +""" + + +class TurnStateStore: + """Sync sqlite bodies behind ``asyncio.to_thread``-style callers. + + Follows the KnowledgeStore conventions: one long-lived connection, WAL, + busy_timeout, a write lock, and an ``available`` property that degrades + the whole feature (not the process) when init fails. + """ + + def __init__( + self, + db_path: str | Path, + *, + blob_dir: str | Path | None = None, + lease_ttl: float = DEFAULT_LEASE_TTL, + ) -> None: + self.db_path = str(db_path) + self.lease_ttl = lease_ttl + self._blob_dir = Path(blob_dir) if blob_dir else Path(self.db_path).parent / "blobs" + self._write_lock = threading.Lock() + self._conn: sqlite3.Connection | None = None + try: + Path(self.db_path).parent.mkdir(parents=True, exist_ok=True) + self._blob_dir.mkdir(parents=True, exist_ok=True) + conn = sqlite3.connect(self.db_path, check_same_thread=False) + conn.execute("PRAGMA journal_mode=WAL") + conn.execute("PRAGMA busy_timeout=30000") + conn.executescript(_DDL) + conn.commit() + self._conn = conn + swept = self._boot_sweep_sync() + if swept["turns"] or swept["ops"]: + log.warning( + "Turn-state boot sweep: %d stale ACTIVE turn(s) suspended, " + "%d in-flight op(s) marked OUTCOME_UNKNOWN", + swept["turns"], swept["ops"], + ) + except Exception: + log.exception( + "TurnStateStore init failed — checkpoint durability DISABLED " + "for this process (turns run legacy, work is not preserved)" + ) + self._conn = None + + @property + def available(self) -> bool: + return self._conn is not None + + def close(self) -> None: + if self._conn is not None: + try: + self._conn.close() + except Exception: + pass + self._conn = None + + # ── internals ──────────────────────────────────────────────────── + + def _require(self) -> sqlite3.Connection: + if self._conn is None: + raise TurnStateUnavailableError("turn-state store is not available") + return self._conn + + def _fenced_update( + self, lease: TurnLease, set_sql: str, params: list, *, bump_revision: bool = True + ) -> None: + """Run one UPDATE fenced on (generation, revision, lease token).""" + conn = self._require() + new_revision = lease.revision + 1 if bump_revision else lease.revision + sql = ( + f"UPDATE turns SET {set_sql}, revision=? " + "WHERE source=? AND channel_id=? AND message_id=? " + "AND turn_generation=? AND revision=? AND lease_token=?" + ) + try: + with self._write_lock: + cur = conn.execute( + sql, + [*params, new_revision, lease.key.source, lease.key.channel_id, + lease.key.message_id, lease.generation, lease.revision, + lease.token], + ) + conn.commit() + except sqlite3.Error as exc: + raise TurnStateUnavailableError(f"turn-state write failed: {exc}") from exc + if cur.rowcount != 1: + raise StaleTurnError( + f"turn {lease.key} generation {lease.generation[:8]} " + f"rev {lease.revision}: fence lost" + ) + lease.revision = new_revision + + def _op_where(self, lease: TurnLease) -> tuple[str, list]: + return ( + "source=? AND channel_id=? AND message_id=? AND turn_generation=?", + [lease.key.source, lease.key.channel_id, lease.key.message_id, + lease.generation], + ) + + def _verify_lease(self, lease: TurnLease) -> None: + """Ops-table writes are fenced indirectly: verify the turns row still + carries this lease before touching operations.""" + conn = self._require() + row = conn.execute( + "SELECT lease_token, turn_generation FROM turns " + "WHERE source=? AND channel_id=? AND message_id=?", + [lease.key.source, lease.key.channel_id, lease.key.message_id], + ).fetchone() + if row is None or row[0] != lease.token or row[1] != lease.generation: + raise StaleTurnError(f"turn {lease.key}: lease no longer held") + + # ── admission / lifecycle (sync bodies) ────────────────────────── + + def admit_turn_sync( + self, + key: TurnKey, + *, + guild_id: str | None, + user_id: str | None, + content_digest: str | None, + code_version: str | None, + prompt_policy_hash: str | None, + tool_catalog_hash: str | None, + session_snapshot: dict | None, + ) -> TurnLease | None: + """Admit a fresh turn. Returns a lease, or None when this message + already has a row (redelivery/crash artifact — caller runs legacy, + loudly) or when the store is unavailable (caller runs legacy).""" + if self._conn is None: + return None + now = time.time() + generation = secrets.token_hex(16) + token = secrets.token_hex(16) + try: + with self._write_lock: + try: + self._conn.execute( + "INSERT INTO turns (source, channel_id, message_id, " + "turn_generation, revision, lease_token, lease_expires_at, " + "status, last_progress_at, created_at, guild_id, user_id, " + "content_digest, code_version, schema_version, " + "prompt_policy_hash, tool_catalog_hash, session_snapshot) " + "VALUES (?,?,?,?,0,?,?,?,?,?,?,?,?,?,?,?,?,?)", + [key.source, key.channel_id, key.message_id, generation, + token, now + self.lease_ttl, TurnStatus.ACTIVE, now, now, + guild_id, user_id, content_digest, code_version, + SCHEMA_VERSION, prompt_policy_hash, tool_catalog_hash, + json.dumps(session_snapshot or {})], + ) + self._conn.commit() + except sqlite3.IntegrityError: + log.warning( + "Turn %s already has a state row — running without " + "durability for this turn", key, + ) + return None + except sqlite3.Error: + log.exception("Turn admission failed — running without durability") + return None + return TurnLease(key=key, generation=generation, token=token, revision=0) + + def checkpoint_sync( + self, + lease: TurnLease, + payload: dict, + *, + progressed: bool, + recovery_deadline_utc: float | None = None, + ) -> None: + """Persist the turn payload (fenced). ``progressed`` advances + ``last_progress_at``; recovery waits and rewrites must pass False.""" + sets = ["payload=?", "lease_expires_at=?"] + params: list = [json.dumps(payload), time.time() + self.lease_ttl] + if progressed: + sets.append("last_progress_at=?") + params.append(time.time()) + if recovery_deadline_utc is not None: + sets.append("recovery_deadline_utc=?") + params.append(recovery_deadline_utc) + self._fenced_update(lease, ", ".join(sets), params) + + def heartbeat_sync(self, lease: TurnLease) -> None: + """Extend the lease. Never advances last_progress_at, never bumps the + revision (a heartbeat is not a state change).""" + self._fenced_update( + lease, "lease_expires_at=?", [time.time() + self.lease_ttl], + bump_revision=False, + ) + + def suspend_sync(self, lease: TurnLease, payload: dict) -> None: + now = time.time() + self._fenced_update( + lease, + "payload=?, status=?, suspended_at=?, lease_token=NULL, " + "lease_expires_at=NULL", + [json.dumps(payload), TurnStatus.SUSPENDED, now], + ) + + def finish_sync(self, lease: TurnLease, status: str = TurnStatus.TERMINAL_COMPLETED) -> None: + """Terminal transition. The payload is compacted immediately + (identity + ledger tombstones retained).""" + if status not in TurnStatus.TERMINAL: + raise ValueError(f"finish_sync requires a terminal status, got {status}") + self._fenced_update( + lease, + "status=?, payload=NULL, lease_token=NULL, lease_expires_at=NULL", + [status], + ) + + # ── ledger (sync bodies) ───────────────────────────────────────── + + def record_intents_sync( + self, + lease: TurnLease, + generation_seq: int, + intents: list[dict], + *, + iteration: int | None = None, + ) -> None: + """Persist PREPARED rows for one generation's tool calls. + + Validates Odin's rule up front: tool-call ids nonempty and unique + within the generation — malformed duplicates fail BEFORE execution. + """ + ids = [str(i.get("tool_call_id") or "") for i in intents] + if any(not i for i in ids): + raise LedgerIntentError("empty tool_call_id in intents") + if len(set(ids)) != len(ids): + raise LedgerIntentError("duplicate tool_call_id in generation") + conn = self._require() + self._verify_lease(lease) + now = time.time() + rows = [ + (lease.key.source, lease.key.channel_id, lease.key.message_id, + lease.generation, generation_seq, str(i["tool_call_id"]), + OpState.PREPARED, str(i.get("tool_name") or "unknown"), iteration, + effect_fingerprint(str(i.get("tool_name") or ""), i.get("tool_input") or {}), + None, now, now) + for i in intents + ] + try: + with self._write_lock: + conn.executemany( + "INSERT OR REPLACE INTO operations VALUES " + "(?,?,?,?,?,?,?,?,?,?,?,?,?)", + rows, + ) + conn.commit() + except sqlite3.Error as exc: + raise TurnStateUnavailableError(f"ledger intent write failed: {exc}") from exc + + def mark_running_sync(self, lease: TurnLease, generation_seq: int, tool_call_id: str) -> None: + self._set_op_state(lease, generation_seq, tool_call_id, OpState.RUNNING, None) + + def settle_op_sync( + self, + lease: TurnLease, + generation_seq: int, + tool_call_id: str, + *, + state: str, + result_text: str | None, + ) -> None: + if state not in (OpState.APPLIED, OpState.DEFINITELY_FAILED): + raise ValueError(f"settle_op_sync: invalid terminal state {state}") + self._set_op_state(lease, generation_seq, tool_call_id, state, result_text) + + def _set_op_state( + self, + lease: TurnLease, + generation_seq: int, + tool_call_id: str, + state: str, + result_text: str | None, + ) -> None: + conn = self._require() + self._verify_lease(lease) + where, params = self._op_where(lease) + try: + with self._write_lock: + if result_text is None: + cur = conn.execute( + f"UPDATE operations SET state=?, updated_at=? WHERE {where} " + "AND generation_seq=? AND tool_call_id=?", + [state, time.time(), *params, generation_seq, tool_call_id], + ) + else: + cur = conn.execute( + f"UPDATE operations SET state=?, result=?, updated_at=? " + f"WHERE {where} AND generation_seq=? AND tool_call_id=?", + [state, result_text, time.time(), *params, generation_seq, + tool_call_id], + ) + conn.commit() + except sqlite3.Error as exc: + raise TurnStateUnavailableError(f"ledger write failed: {exc}") from exc + if cur.rowcount != 1: + raise StaleTurnError( + f"op {tool_call_id} gen_seq {generation_seq}: no PREPARED row to update" + ) + + # ── resume (sync bodies) ───────────────────────────────────────── + + def load_resumable_sync(self, key: TurnKey) -> dict | None: + """Return the SUSPENDED row (payload + validation columns + ops) or None.""" + if self._conn is None: + return None + row = self._conn.execute( + "SELECT turn_generation, revision, payload, guild_id, user_id, " + "content_digest, code_version, schema_version, prompt_policy_hash, " + "tool_catalog_hash, session_snapshot, recovery_deadline_utc, " + "last_progress_at, suspended_at FROM turns " + "WHERE source=? AND channel_id=? AND message_id=? AND status=?", + [key.source, key.channel_id, key.message_id, TurnStatus.SUSPENDED], + ).fetchone() + if row is None or row[2] is None: + return None + ops = self._conn.execute( + "SELECT generation_seq, tool_call_id, state, tool_name, result " + "FROM operations WHERE source=? AND channel_id=? AND message_id=? " + "AND turn_generation=? ORDER BY generation_seq, tool_call_id", + [key.source, key.channel_id, key.message_id, row[0]], + ).fetchall() + return { + "generation": row[0], + "revision": row[1], + "payload": json.loads(row[2]), + "guild_id": row[3], + "user_id": row[4], + "content_digest": row[5], + "code_version": row[6], + "schema_version": row[7], + "prompt_policy_hash": row[8], + "tool_catalog_hash": row[9], + "session_snapshot": json.loads(row[10] or "{}"), + "recovery_deadline_utc": row[11], + "last_progress_at": row[12], + "suspended_at": row[13], + "operations": [ + {"generation_seq": o[0], "tool_call_id": o[1], "state": o[2], + "tool_name": o[3], "result": o[4]} + for o in ops + ], + } + + def acquire_resume_lease_sync(self, key: TurnKey, expected_generation: str) -> TurnLease | None: + """Reclaim a SUSPENDED turn: same generation (the logical turn's + lineage), fresh lease token. Conditional on status+generation so two + resumers can't both win.""" + if self._conn is None: + return None + token = secrets.token_hex(16) + now = time.time() + try: + with self._write_lock: + cur = self._conn.execute( + "UPDATE turns SET status=?, lease_token=?, lease_expires_at=? " + "WHERE source=? AND channel_id=? AND message_id=? " + "AND status=? AND turn_generation=?", + [TurnStatus.ACTIVE, token, now + self.lease_ttl, + key.source, key.channel_id, key.message_id, + TurnStatus.SUSPENDED, expected_generation], + ) + self._conn.commit() + except sqlite3.Error as exc: + raise TurnStateUnavailableError(f"resume-lease write failed: {exc}") from exc + if cur.rowcount != 1: + return None + row = self._conn.execute( + "SELECT revision FROM turns WHERE source=? AND channel_id=? AND message_id=?", + [key.source, key.channel_id, key.message_id], + ).fetchone() + return TurnLease( + key=key, generation=expected_generation, token=token, + revision=int(row[0]), + ) + + def reject_resumable_sync(self, key: TurnKey, reason: str) -> None: + """A SUSPENDED turn failed resume validation (deleted/edited message, + author mismatch, policy change): terminal, payload compacted.""" + if self._conn is None: + return + try: + with self._write_lock: + self._conn.execute( + "UPDATE turns SET status=?, payload=NULL, lease_token=NULL, " + "lease_expires_at=NULL WHERE source=? AND channel_id=? " + "AND message_id=? AND status=?", + [TurnStatus.TERMINAL_REJECTED, key.source, key.channel_id, + key.message_id, TurnStatus.SUSPENDED], + ) + self._conn.commit() + log.info("Resumable turn %s rejected: %s", key, reason) + except sqlite3.Error: + log.exception("reject_resumable failed (non-fatal)") + + def list_suspended_sync(self, source: str | None = None) -> list[dict]: + if self._conn is None: + return [] + sql = ( + "SELECT source, channel_id, message_id, turn_generation, " + "last_progress_at, suspended_at FROM turns WHERE status=?" + ) + params: list = [TurnStatus.SUSPENDED] + if source: + sql += " AND source=?" + params.append(source) + rows = self._conn.execute(sql, params).fetchall() + return [ + {"source": r[0], "channel_id": r[1], "message_id": r[2], + "generation": r[3], "last_progress_at": r[4], "suspended_at": r[5]} + for r in rows + ] + + # ── sweeps ─────────────────────────────────────────────────────── + + def _boot_sweep_sync(self) -> dict: + """Crash recovery at construction: expired-lease ACTIVE turns become + SUSPENDED; their PREPARED/RUNNING ops become OUTCOME_UNKNOWN (we + cannot know whether the external effect happened — never rerun).""" + conn = self._require() + now = time.time() + with self._write_lock: + stale = conn.execute( + "SELECT source, channel_id, message_id, turn_generation FROM turns " + "WHERE status=? AND (lease_expires_at IS NULL OR lease_expires_at < ?)", + [TurnStatus.ACTIVE, now], + ).fetchall() + ops = 0 + for source, channel_id, message_id, generation in stale: + cur = conn.execute( + "UPDATE operations SET state=?, updated_at=? " + "WHERE source=? AND channel_id=? AND message_id=? " + "AND turn_generation=? AND state IN (?, ?)", + [OpState.OUTCOME_UNKNOWN, now, source, channel_id, message_id, + generation, OpState.PREPARED, OpState.RUNNING], + ) + ops += cur.rowcount + conn.execute( + "UPDATE turns SET status=?, suspended_at=?, lease_token=NULL, " + "lease_expires_at=NULL WHERE source=? AND channel_id=? " + "AND message_id=?", + [TurnStatus.SUSPENDED, now, source, channel_id, message_id], + ) + conn.commit() + return {"turns": len(stale), "ops": ops} + + def ttl_sweep_sync( + self, + *, + resume_ttl_hours: float = DEFAULT_RESUME_TTL_HOURS, + payload_retention_days: float = DEFAULT_PAYLOAD_RETENTION_DAYS, + ledger_retention_days: float = DEFAULT_LEDGER_RETENTION_DAYS, + ) -> dict: + """The three retention clocks. Returns counts for observability.""" + if self._conn is None: + return {} + conn = self._conn + now = time.time() + expired = payloads = ledger = 0 + try: + with self._write_lock: + # Clock 1: resumable window — 24h from last REAL progress. + cur = conn.execute( + "UPDATE turns SET status=? WHERE status=? AND last_progress_at < ?", + [TurnStatus.TERMINAL_EXPIRED, TurnStatus.SUSPENDED, + now - resume_ttl_hours * 3600.0], + ) + expired = cur.rowcount + # Clock 2: diagnostic payload retention — 7d, then tombstone. + cur = conn.execute( + "UPDATE turns SET payload=NULL WHERE payload IS NOT NULL " + "AND status IN (?, ?, ?, ?, ?) AND last_progress_at < ?", + [*sorted(TurnStatus.TERMINAL), + now - payload_retention_days * 86400.0], + ) + payloads = cur.rowcount + # Clock 3: ledger — ≥90d for rows of terminal turns, EXCEPT + # OUTCOME_UNKNOWN / MANUAL_RESOLUTION_REQUIRED (never expire). + cur = conn.execute( + "DELETE FROM operations WHERE state NOT IN (?, ?) " + "AND updated_at < ? AND (source, channel_id, message_id) IN " + "(SELECT source, channel_id, message_id FROM turns " + " WHERE status IN (?, ?, ?, ?, ?))", + [*sorted(OpState.NEVER_EXPIRE), + now - ledger_retention_days * 86400.0, + *sorted(TurnStatus.TERMINAL)], + ) + ledger = cur.rowcount + conn.commit() + except sqlite3.Error: + log.exception("turn-state TTL sweep failed (non-fatal)") + return {"expired_turns": expired, "compacted_payloads": payloads, + "ledger_rows_deleted": ledger} + + # ── blobs ──────────────────────────────────────────────────────── + + def store_blob_sync(self, data: bytes) -> str: + """Content-addressed blob write (tmp+rename). Returns 'blob:'.""" + digest = hashlib.sha256(data).hexdigest() + path = self._blob_dir / digest + if not path.exists(): + tmp = path.with_suffix(".tmp") + try: + tmp.write_bytes(data) + os.replace(tmp, path) + except OSError as exc: + raise TurnStateUnavailableError(f"blob write failed: {exc}") from exc + return f"blob:{digest}" + + def load_blob_sync(self, ref: str) -> bytes: + digest = ref.split(":", 1)[1] if ref.startswith("blob:") else ref + path = self._blob_dir / digest + try: + return path.read_bytes() + except OSError as exc: + raise TurnStateUnavailableError(f"blob read failed: {ref}") from exc diff --git a/tests/test_recovery.py b/tests/test_recovery.py index f56b5eb1..c2db2487 100644 --- a/tests/test_recovery.py +++ b/tests/test_recovery.py @@ -689,7 +689,8 @@ def test_tools_config_custom_recovery(self): class TestAgentPerIterationRecovery: @pytest.mark.asyncio async def test_recovery_attempts_reset_each_iteration(self): - """recovery_attempts is reset to 0 at the start of each iteration.""" + """recovery_attempts stays 0 on healthy iterations (the field is + API-shape compatibility now; the manager ladder is gone).""" from src.agents.manager import AgentInfo, AgentState, _run_agent agent = AgentInfo( @@ -722,50 +723,16 @@ async def mock_tool_cb(name, inp): assert all(v == 0 for v in recovery_values) @pytest.mark.asyncio - async def test_recovery_works_on_second_iteration_after_first_recovery(self): - """After recovery in iteration 1, iteration 2 can also recover.""" - from src.agents.manager import ( - AgentInfo, - AgentState, - _run_agent, - ) - - agent = AgentInfo( - id="test2", label="test", goal="test goal", - channel_id="ch1", requester_id="u1", requester_name="user", - ) - - call_count = 0 - - async def mock_iteration_cb(messages, system_prompt, tools): - nonlocal call_count - call_count += 1 - if call_count in (1, 3): - raise TimeoutError("LLM timeout") - if call_count in (2, 4): - return {"text": "working", "tool_calls": [{"name": "t1", "input": {}}]} - return {"text": "done", "tool_calls": []} - - async def mock_tool_cb(name, inp): - return "ok" - - await _run_agent( - agent=agent, - system_prompt="test", - tools=[], - iteration_callback=mock_iteration_cb, - tool_executor_callback=mock_tool_cb, - ) - assert agent.state == AgentState.COMPLETED - assert call_count == 5 - - @pytest.mark.asyncio - async def test_old_behavior_recovery_exhausted_without_reset(self): - """Verify the reset actually matters by checking iteration-level budget.""" + async def test_callback_failure_fails_agent_no_manager_ladder(self): + """Deliberate amendment (2026-07-30): the manager-level per-iteration + retry ladder is gone — transient recovery lives INSIDE the iteration + callback (src/llm/recovery.py). A failure that escapes the callback + fails the agent on the spot; the would-have-recovered second call + never happens.""" from src.agents.manager import AgentInfo, AgentState, _run_agent agent = AgentInfo( - id="test3", label="test", goal="test goal", + id="test2", label="test", goal="test goal", channel_id="ch1", requester_id="u1", requester_name="user", ) @@ -775,9 +742,7 @@ async def mock_iteration_cb(messages, system_prompt, tools): nonlocal call_count call_count += 1 if call_count == 1: - raise TimeoutError("timeout") - if call_count == 2: - return {"text": "working", "tool_calls": [{"name": "t1", "input": {}}]} + raise TimeoutError("LLM timeout") return {"text": "done", "tool_calls": []} async def mock_tool_cb(name, inp): @@ -790,7 +755,9 @@ async def mock_tool_cb(name, inp): iteration_callback=mock_iteration_cb, tool_executor_callback=mock_tool_cb, ) - assert agent.state == AgentState.COMPLETED + assert agent.state == AgentState.FAILED + assert call_count == 1 + assert agent.recovery_attempts == 0 # ==================================================================== diff --git a/tests/test_trajectory_completeness.py b/tests/test_trajectory_completeness.py index 345c4918..5fef0cb0 100644 --- a/tests/test_trajectory_completeness.py +++ b/tests/test_trajectory_completeness.py @@ -311,6 +311,17 @@ async def _log_execution(**kwargs): self.llm_client = SimpleNamespace(chat_with_tools=_chat_with_tools) self.audit = SimpleNamespace(log_execution=_log_execution) + from src.llm.model_breaker import ModelBreakerRegistry + from src.llm.recovery import RecoveryPolicy + + _registry = ModelBreakerRegistry() + self._fake_gateway = SimpleNamespace( + active_client=self.llm_client, + capacity_breaker_for=lambda model=None: _registry.for_model("codex", "m"), + recovery_policy=RecoveryPolicy, + notify_generation_success=lambda provider: None, + ) + # P4 migration: the runner takes narrow deps now. The recorder is the # REAL one; its save/reflect hooks are shadowed with this fake's # capture methods so assertions keep observing the loop body. @@ -323,7 +334,7 @@ async def _log_execution(**kwargs): get_config=lambda: self.config, get_default_system_prompt=lambda: "sys", get_context_compressor=lambda: None, - llm_gateway=SimpleNamespace(active_client=self.llm_client), + llm_gateway=self._fake_gateway, prompt_builder=SimpleNamespace(build_full_prompt=lambda **kw: "sys"), tool_catalog=SimpleNamespace( merged_definitions=lambda: [{"name": "run_command"}] diff --git a/tests/test_turn_state_store.py b/tests/test_turn_state_store.py new file mode 100644 index 00000000..98a7b351 --- /dev/null +++ b/tests/test_turn_state_store.py @@ -0,0 +1,395 @@ +"""Pins for the durable turn-state store (src/turn_state/store.py). + +The load-bearing properties, each from the settled design: +- three-part write fencing (generation, revision, lease token); +- heartbeats extend the lease but never advance last_progress_at and never + bump the revision; +- boot sweep: stale ACTIVE turns suspend, in-flight ops become + OUTCOME_UNKNOWN (never rerun); +- three TTL clocks, with OUTCOME_UNKNOWN ledger rows never auto-expiring; +- terminal payload compaction (tombstones keep identity + ledger). +""" + +from __future__ import annotations + +import time + +import pytest + +from src.turn_state import ( + LedgerIntentError, + OpState, + StaleTurnError, + TurnKey, + TurnStateStore, + TurnStatus, + effect_fingerprint, +) + +KEY = TurnKey(source="discord", channel_id="c1", message_id="m1") + + +@pytest.fixture +def store(tmp_path): + s = TurnStateStore(tmp_path / "turns.sqlite3", blob_dir=tmp_path / "blobs") + yield s + s.close() + + +def _admit(store, key=KEY): + lease = store.admit_turn_sync( + key, + guild_id="g1", + user_id="u1", + content_digest="digest", + code_version="test", + prompt_policy_hash="pp", + tool_catalog_hash="tc", + session_snapshot={"messages": 3}, + ) + assert lease is not None + return lease + + +def _row(store, key=KEY, cols="status, revision, payload, last_progress_at"): + return store._conn.execute( + f"SELECT {cols} FROM turns WHERE source=? AND channel_id=? AND message_id=?", + [key.source, key.channel_id, key.message_id], + ).fetchone() + + +class TestAdmissionAndCheckpoint: + def test_admit_and_checkpoint_roundtrip(self, store): + lease = _admit(store) + assert lease.revision == 0 + store.checkpoint_sync(lease, {"iteration": 1}, progressed=True) + assert lease.revision == 1 + store.checkpoint_sync(lease, {"iteration": 2}, progressed=True) + assert lease.revision == 2 + status, revision, payload, _ = _row(store) + assert status == TurnStatus.ACTIVE + assert revision == 2 + assert '"iteration": 2' in payload + + def test_double_admission_returns_none(self, store): + _admit(store) + again = store.admit_turn_sync( + KEY, guild_id=None, user_id=None, content_digest=None, + code_version=None, prompt_policy_hash=None, tool_catalog_hash=None, + session_snapshot=None, + ) + assert again is None + + def test_unavailable_store_admits_none(self, tmp_path): + # A file where the DB directory should be → init fails → degraded. + blocker = tmp_path / "blocked" + blocker.write_text("not a directory") + s = TurnStateStore(blocker / "turns.sqlite3") + assert s.available is False + assert s.admit_turn_sync( + KEY, guild_id=None, user_id=None, content_digest=None, + code_version=None, prompt_policy_hash=None, tool_catalog_hash=None, + session_snapshot=None, + ) is None + + def test_recovery_deadline_persisted_utc(self, store): + lease = _admit(store) + deadline = time.time() + 300.0 + store.checkpoint_sync( + lease, {}, progressed=False, recovery_deadline_utc=deadline + ) + (stored,) = _row(store, cols="recovery_deadline_utc") + assert abs(stored - deadline) < 0.001 + + +class TestFencing: + def test_wrong_token_is_stale(self, store): + lease = _admit(store) + import dataclasses + + thief = dataclasses.replace(lease) + thief.token = "someone-else" + with pytest.raises(StaleTurnError): + store.checkpoint_sync(thief, {}, progressed=False) + + def test_old_revision_is_stale(self, store): + lease = _admit(store) + store.checkpoint_sync(lease, {"n": 1}, progressed=False) + import dataclasses + + old = dataclasses.replace(lease) + old.revision = 0 # an expired owner replaying its last known revision + with pytest.raises(StaleTurnError): + store.checkpoint_sync(old, {"stale": True}, progressed=False) + # The current owner still works. + store.checkpoint_sync(lease, {"n": 2}, progressed=False) + + def test_resumed_turn_fences_out_old_owner(self, store): + lease = _admit(store) + store.checkpoint_sync(lease, {"n": 1}, progressed=True) + store.suspend_sync(lease, {"n": 1}) + new_lease = store.acquire_resume_lease_sync(KEY, lease.generation) + assert new_lease is not None + # The pre-suspension owner wakes up and tries to write: fenced out + # (its token was cleared at suspension; the new owner holds a fresh one). + with pytest.raises(StaleTurnError): + store.checkpoint_sync(lease, {"zombie": True}, progressed=False) + store.checkpoint_sync(new_lease, {"n": 2}, progressed=True) + + +class TestHeartbeat: + def test_heartbeat_extends_lease_only(self, store): + lease = _admit(store) + store.checkpoint_sync(lease, {"n": 1}, progressed=True) + _, rev_before, _, progress_before = _row(store) + (lease_before,) = _row(store, cols="lease_expires_at") + time.sleep(0.01) + store.heartbeat_sync(lease) + status, rev_after, _, progress_after = _row(store) + (lease_after,) = _row(store, cols="lease_expires_at") + assert lease_after > lease_before # extended + assert rev_after == rev_before # NOT a state change + assert progress_after == progress_before # never fake progress + + def test_waits_do_not_advance_progress(self, store): + lease = _admit(store) + store.checkpoint_sync(lease, {"n": 1}, progressed=True) + (progress_before,) = _row(store, cols="last_progress_at") + time.sleep(0.01) + store.checkpoint_sync(lease, {"n": 1, "waiting": True}, progressed=False) + (progress_after,) = _row(store, cols="last_progress_at") + assert progress_after == progress_before + + +class TestSuspendResumeReject: + def test_suspend_and_load_resumable(self, store): + lease = _admit(store) + store.record_intents_sync( + lease, 0, + [{"tool_call_id": "call_1", "tool_name": "run_command", + "tool_input": {"command": "ls"}}], + iteration=0, + ) + store.settle_op_sync(lease, 0, "call_1", state=OpState.APPLIED, + result_text="file.txt") + store.suspend_sync(lease, {"messages": ["m"]}) + resumable = store.load_resumable_sync(KEY) + assert resumable is not None + assert resumable["generation"] == lease.generation + assert resumable["payload"] == {"messages": ["m"]} + assert resumable["content_digest"] == "digest" + assert resumable["session_snapshot"] == {"messages": 3} + ops = resumable["operations"] + assert len(ops) == 1 + assert ops[0]["state"] == OpState.APPLIED + assert ops[0]["result"] == "file.txt" + + def test_resume_lease_single_winner(self, store): + lease = _admit(store) + store.suspend_sync(lease, {}) + first = store.acquire_resume_lease_sync(KEY, lease.generation) + second = store.acquire_resume_lease_sync(KEY, lease.generation) + assert first is not None + assert second is None # already ACTIVE under the first resumer + + def test_resume_lease_wrong_generation_loses(self, store): + lease = _admit(store) + store.suspend_sync(lease, {}) + assert store.acquire_resume_lease_sync(KEY, "not-the-generation") is None + + def test_reject_resumable_is_terminal_and_compacted(self, store): + lease = _admit(store) + store.suspend_sync(lease, {"payload": True}) + store.reject_resumable_sync(KEY, "message edited") + status, _, payload, _ = _row(store) + assert status == TurnStatus.TERMINAL_REJECTED + assert payload is None + assert store.load_resumable_sync(KEY) is None + + def test_finish_compacts_payload_keeps_ledger(self, store): + lease = _admit(store) + store.record_intents_sync( + lease, 0, [{"tool_call_id": "c1", "tool_name": "t", "tool_input": {}}] + ) + store.settle_op_sync(lease, 0, "c1", state=OpState.APPLIED, result_text="r") + store.checkpoint_sync(lease, {"big": "payload"}, progressed=True) + store.finish_sync(lease) + status, _, payload, _ = _row(store) + assert status == TurnStatus.TERMINAL_COMPLETED + assert payload is None # compacted immediately + count = store._conn.execute("SELECT COUNT(*) FROM operations").fetchone()[0] + assert count == 1 # ledger tombstone retained + + def test_finish_requires_terminal_status(self, store): + lease = _admit(store) + with pytest.raises(ValueError): + store.finish_sync(lease, "ACTIVE") + + +class TestLedger: + def test_intent_validation(self, store): + lease = _admit(store) + with pytest.raises(LedgerIntentError): + store.record_intents_sync( + lease, 0, [{"tool_call_id": "", "tool_name": "t", "tool_input": {}}] + ) + with pytest.raises(LedgerIntentError): + store.record_intents_sync( + lease, 0, + [{"tool_call_id": "dup", "tool_name": "t", "tool_input": {}}, + {"tool_call_id": "dup", "tool_name": "t2", "tool_input": {}}], + ) + + def test_state_machine_prepared_running_applied(self, store): + lease = _admit(store) + store.record_intents_sync( + lease, 0, [{"tool_call_id": "c1", "tool_name": "t", "tool_input": {}}] + ) + store.mark_running_sync(lease, 0, "c1") + store.settle_op_sync(lease, 0, "c1", state=OpState.APPLIED, result_text="ok") + row = store._conn.execute( + "SELECT state, result FROM operations WHERE tool_call_id='c1'" + ).fetchone() + assert row == (OpState.APPLIED, "ok") + + def test_settle_rejects_non_terminal_states(self, store): + lease = _admit(store) + store.record_intents_sync( + lease, 0, [{"tool_call_id": "c1", "tool_name": "t", "tool_input": {}}] + ) + with pytest.raises(ValueError): + store.settle_op_sync(lease, 0, "c1", state=OpState.RUNNING, result_text=None) + + def test_settle_missing_row_is_stale(self, store): + lease = _admit(store) + with pytest.raises(StaleTurnError): + store.settle_op_sync(lease, 0, "ghost", state=OpState.APPLIED, + result_text="x") + + def test_ledger_writes_require_live_lease(self, store): + lease = _admit(store) + store.suspend_sync(lease, {}) # releases the lease + with pytest.raises(StaleTurnError): + store.record_intents_sync( + lease, 0, [{"tool_call_id": "c1", "tool_name": "t", "tool_input": {}}] + ) + + +class TestBootSweep: + def test_stale_active_suspends_and_ops_go_unknown(self, tmp_path): + db = tmp_path / "turns.sqlite3" + s = TurnStateStore(db, blob_dir=tmp_path / "blobs") + lease = _admit(s) + s.record_intents_sync( + lease, 0, + [{"tool_call_id": "c1", "tool_name": "run_command", "tool_input": {}}, + {"tool_call_id": "c2", "tool_name": "run_command", "tool_input": {}}], + ) + s.mark_running_sync(lease, 0, "c1") + # Simulate crash: expire the lease on disk, drop the handle, reopen. + s._conn.execute("UPDATE turns SET lease_expires_at = ?", [time.time() - 10]) + s._conn.commit() + s.close() + + reopened = TurnStateStore(db, blob_dir=tmp_path / "blobs") + row = reopened._conn.execute("SELECT status FROM turns").fetchone() + assert row[0] == TurnStatus.SUSPENDED + states = { + r[0] for r in reopened._conn.execute( + "SELECT state FROM operations" + ).fetchall() + } + assert states == {OpState.OUTCOME_UNKNOWN} + reopened.close() + + def test_live_lease_survives_reopen(self, tmp_path): + db = tmp_path / "turns.sqlite3" + s = TurnStateStore(db, blob_dir=tmp_path / "blobs") + _admit(s) + s.close() + # Lease still in the future → NOT swept (a concurrent healthy owner + # in another process must not be hijacked). + reopened = TurnStateStore(db, blob_dir=tmp_path / "blobs") + row = reopened._conn.execute("SELECT status FROM turns").fetchone() + assert row[0] == TurnStatus.ACTIVE + reopened.close() + + +class TestTtlSweep: + def _age(self, store, key, *, progress_age_s): + store._conn.execute( + "UPDATE turns SET last_progress_at=? WHERE message_id=?", + [time.time() - progress_age_s, key.message_id], + ) + store._conn.commit() + + def test_resumable_expires_after_ttl(self, store): + lease = _admit(store) + store.suspend_sync(lease, {"p": 1}) + self._age(store, KEY, progress_age_s=25 * 3600) + out = store.ttl_sweep_sync(resume_ttl_hours=24.0) + assert out["expired_turns"] == 1 + (status,) = _row(store, cols="status") + assert status == TurnStatus.TERMINAL_EXPIRED + # Diagnostic payload retained until the 7d clock. + (payload,) = _row(store, cols="payload") + assert payload is not None + + def test_diagnostic_payload_compacts_after_seven_days(self, store): + lease = _admit(store) + store.suspend_sync(lease, {"p": 1}) + self._age(store, KEY, progress_age_s=8 * 86400) + out = store.ttl_sweep_sync() + assert out["expired_turns"] == 1 + assert out["compacted_payloads"] == 1 + (payload,) = _row(store, cols="payload") + assert payload is None + + def test_ledger_expires_except_unknown(self, store): + lease = _admit(store) + store.record_intents_sync( + lease, 0, + [{"tool_call_id": "ok", "tool_name": "t", "tool_input": {}}, + {"tool_call_id": "mystery", "tool_name": "t", "tool_input": {}}], + ) + store.settle_op_sync(lease, 0, "ok", state=OpState.APPLIED, result_text="r") + store._conn.execute( + "UPDATE operations SET state=? WHERE tool_call_id='mystery'", + [OpState.OUTCOME_UNKNOWN], + ) + store.finish_sync(lease, TurnStatus.TERMINAL_FAILED) + self._age(store, KEY, progress_age_s=91 * 86400) + store._conn.execute( + "UPDATE operations SET updated_at=?", [time.time() - 91 * 86400] + ) + store._conn.commit() + out = store.ttl_sweep_sync() + assert out["ledger_rows_deleted"] == 1 + remaining = store._conn.execute( + "SELECT tool_call_id, state FROM operations" + ).fetchall() + # OUTCOME_UNKNOWN never expires automatically. + assert remaining == [("mystery", OpState.OUTCOME_UNKNOWN)] + + +class TestBlobs: + def test_roundtrip_and_content_addressing(self, store): + ref1 = store.store_blob_sync(b"image-bytes") + ref2 = store.store_blob_sync(b"image-bytes") + assert ref1 == ref2 + assert ref1.startswith("blob:") + assert store.load_blob_sync(ref1) == b"image-bytes" + + def test_missing_blob_raises_unavailable(self, store): + from src.turn_state import TurnStateUnavailableError + + with pytest.raises(TurnStateUnavailableError): + store.load_blob_sync("blob:" + "0" * 64) + + +def test_effect_fingerprint_is_deterministic_and_sensitive(): + a = effect_fingerprint("run_command", {"command": "ls", "host": "web"}) + b = effect_fingerprint("run_command", {"host": "web", "command": "ls"}) + c = effect_fingerprint("run_command", {"command": "rm", "host": "web"}) + assert a == b # key order irrelevant + assert a != c From 18b5486bbeac05831259e0d47ad4d1ee4a6cea60 Mon Sep 17 00:00:00 2001 From: ATCharpentier Date: Thu, 30 Jul 2026 13:55:01 -0400 Subject: [PATCH 05/18] feat: _ChatTurn checkpoint codec + field-census enforcement MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit src/turn_state/codec.py classifies EVERY _ChatTurn field as PERSISTED (restored verbatim — includes all six one-shot guard flags, continuation and validation budgets, post-compression messages, the possibly-rebound system_prompt, op details, pending validations) or RECONSTRUCTED with the reason (message re-fetched; _cancel fresh from channel_state; tools re-derived from CURRENT catalog+permissions — current policy wins; policy by stable name; trace closed one-way + fresh linked resume segment). tests/test_turn_checkpoint_codec.py carries the census pin: a new _ChatTurn field fails the test until classified — permanently enforcing that resume can never silently re-arm guard budgets (hard rule). Also: stuck-tracker export/import (fingerprints + warned), FULL trajectory persistence with an iteration_revision guard (extra rows dropped on restore — no double-append), base64 image blocks externalized to content-addressed blobs and re-inlined on restore. --- src/turn_state/codec.py | 298 ++++++++++++++++++++++++++++ tests/test_turn_checkpoint_codec.py | 259 ++++++++++++++++++++++++ 2 files changed, 557 insertions(+) create mode 100644 src/turn_state/codec.py create mode 100644 tests/test_turn_checkpoint_codec.py diff --git a/src/turn_state/codec.py b/src/turn_state/codec.py new file mode 100644 index 00000000..5723b65b --- /dev/null +++ b/src/turn_state/codec.py @@ -0,0 +1,298 @@ +"""Checkpoint codec for the chat turn (`_ChatTurn` ⇄ JSON payload). + +EVERY `_ChatTurn` dataclass field is explicitly classified below as either +PERSISTED (serialized into the payload and restored verbatim) or +RECONSTRUCTED (rebuilt by the resume flow from live state, with the reason +documented). `tests/test_turn_checkpoint_codec.py` holds a field-census pin: +a new `_ChatTurn` field breaks the test until it is classified here — which +permanently enforces the hard rule that a resume can never silently grant +fresh one-shot guard budgets ("never weaken the anti-hedging guards"). + +This module deliberately does NOT import `tool_loop` (which will import the +turn-state machinery): `snapshot_chat_turn` reads attributes off the live +turn object, and `restore_field_values` returns constructor kwargs for the +resume flow to combine with its reconstructed fields. + +Large base64 image blocks (vision injection, pending_image_blocks) are +externalized to the store's content-addressed blob dir and re-inlined on +restore, keeping the SQLite payload row small. +""" + +from __future__ import annotations + +import hashlib +from dataclasses import asdict +from typing import Any + +from ..odin_log import get_logger + +log = get_logger("turn_state") + +CODEC_VERSION = 1 + +# ── The classification (census-pinned) ─────────────────────────────── + +#: Serialized into the payload and restored verbatim. The six one-shot +#: guard flags and every consumed budget MUST be here — restoring a turn +#: without them would silently re-arm the anti-hedging/anti-fabrication +#: guards (hard-rule violation). +PERSISTED_FIELDS: frozenset[str] = frozenset({ + "system_prompt", # CURRENT value — may have been rebound mid-turn + "messages", # post-compression model transcript (blob-externalized) + "user_id", + "chat_cap", + "iteration", + "tools_used_in_loop", + "continuation_count", + "max_continuations", + "fabrication_retried", + "promise_retried", + "unavail_retried", + "hedging_retried", + "code_hedging_retried", + "premature_failure_retried", + "pending_image_blocks", # blob-externalized + "_op_tool_details", + "_pending_validations", + "_validation_required", + "_validation_retries", + "_max_validation_retries", + "_result_store_cap", + "_ch_id", + "_req_id", + "stuck_tracker", # exported as plain state, re-seeded on restore + "_trajectory", # full dict incl. iterations; rebuilt on restore +}) + +#: Rebuilt by the resume flow from live state. Each entry documents why it +#: is NOT persisted: +#: - message: a live discord.Message — re-fetched from channel+message id; +#: deleted or materially edited requests become terminal, never resumed. +#: - _cancel: process-local asyncio.Event — a fresh one from channel_state; +#: cancellation itself is persisted as a TERMINAL turn status. +#: - tools: re-derived from the CURRENT catalog + permission filter — +#: current security policy always wins over persisted tool definitions. +#: - policy: a stable name ("chat") rides the payload metadata; the object +#: is reconstructed from current code, never unpickled folklore. +#: - trace: ContextTraceCollector.finalize() is one-way; the old segment +#: is closed into the payload for diagnostics and a fresh linked resume +#: segment starts (Odin round-2). +RECONSTRUCTED_FIELDS: frozenset[str] = frozenset({ + "message", + "_cancel", + "tools", + "policy", + "trace", +}) + + +def compute_content_digest(text: str) -> str: + """Full sha256 of the request content — the admission validation digest + (NOT identity; identity is source+channel+message id).""" + return hashlib.sha256((text or "").encode("utf-8", "replace")).hexdigest() + + +# ── image-block externalization ────────────────────────────────────── + + +def _externalize_blocks(obj: Any, store_blob) -> Any: + """Recursively swap base64 image payloads for blob refs.""" + if isinstance(obj, list): + return [_externalize_blocks(x, store_blob) for x in obj] + if isinstance(obj, dict): + source = obj.get("source") + if ( + obj.get("type") == "image" + and isinstance(source, dict) + and source.get("type") == "base64" + and isinstance(source.get("data"), str) + ): + ref = store_blob(source["data"].encode("ascii", "replace")) + new_source = {k: v for k, v in source.items() if k != "data"} + new_source["type"] = "blob_ref" + new_source["ref"] = ref + return {**obj, "source": new_source} + return {k: _externalize_blocks(v, store_blob) for k, v in obj.items()} + return obj + + +def _inline_blocks(obj: Any, load_blob) -> Any: + """Inverse of :func:`_externalize_blocks`.""" + if isinstance(obj, list): + return [_inline_blocks(x, load_blob) for x in obj] + if isinstance(obj, dict): + source = obj.get("source") + if ( + obj.get("type") == "image" + and isinstance(source, dict) + and source.get("type") == "blob_ref" + and isinstance(source.get("ref"), str) + ): + data = load_blob(source["ref"]).decode("ascii", "replace") + new_source = {k: v for k, v in source.items() if k != "ref"} + new_source["type"] = "base64" + new_source["data"] = data + return {**obj, "source": new_source} + return {k: _inline_blocks(v, load_blob) for k, v in obj.items()} + return obj + + +# ── stuck-tracker export/import ────────────────────────────────────── + + +def export_stuck_tracker(tracker) -> dict: + return { + "fingerprints": list(tracker._fingerprints), + "window_size": tracker._window_size, + "min_repeats": tracker._min_repeats, + "max_cycle_length": tracker._max_cycle_length, + "names_only": tracker._names_only, + "warned": bool(tracker.warned), + } + + +def import_stuck_tracker(state: dict, tracker_cls): + tracker = tracker_cls( + window_size=int(state.get("window_size", 12)), + min_repeats=int(state.get("min_repeats", 3)), + max_cycle_length=int(state.get("max_cycle_length", 3)), + names_only=bool(state.get("names_only", False)), + ) + tracker._fingerprints.extend(state.get("fingerprints") or []) + tracker.warned = bool(state.get("warned", False)) + return tracker + + +# ── trajectory persistence ─────────────────────────────────────────── + + +def _trajectory_to_payload(trajectory) -> dict: + """FULL trajectory state (unlike TrajectoryTurn.to_dict, which is lossy + by design for the JSONL files). Iterations ride as asdict() rows; the + saved iteration count is the anti-double-append revision (Odin round-2: + "resumption cannot append the same iteration twice").""" + return { + "message_id": trajectory.message_id, + "channel_id": trajectory.channel_id, + "user_id": trajectory.user_id, + "user_name": trajectory.user_name, + "timestamp": trajectory.timestamp, + "source": trajectory.source, + "user_content": trajectory.user_content, + "system_prompt": trajectory.system_prompt, + "history": list(trajectory.history or []), + "iterations": [asdict(it) for it in trajectory.iterations], + "final_response": trajectory.final_response, + "tools_used": list(trajectory.tools_used or []), + "is_error": trajectory.is_error, + "handoff": trajectory.handoff, + "user_content_truncated": trajectory.user_content_truncated, + "user_content_original_chars": trajectory.user_content_original_chars, + "total_input_tokens": trajectory.total_input_tokens, + "total_output_tokens": trajectory.total_output_tokens, + "total_duration_ms": trajectory.total_duration_ms, + "iteration_revision": len(trajectory.iterations), + } + + +def trajectory_from_payload(data: dict): + from ..trajectories.saver import ToolIteration, TrajectoryTurn + + turn = TrajectoryTurn( + message_id=data.get("message_id", ""), + channel_id=data.get("channel_id", ""), + user_id=data.get("user_id", ""), + user_name=data.get("user_name", ""), + timestamp=data.get("timestamp", ""), + source=data.get("source", "discord"), + user_content=data.get("user_content", ""), + system_prompt=data.get("system_prompt", ""), + history=list(data.get("history") or []), + ) + expected = int(data.get("iteration_revision", 0)) + rows = list(data.get("iterations") or [])[:expected] + for row in rows: + try: + turn.iterations.append(ToolIteration(**row)) + except TypeError: + # Forward/backward field drift: keep what maps, never crash resume. + known = {k: v for k, v in row.items() + if k in ToolIteration.__dataclass_fields__} + turn.iterations.append(ToolIteration(**known)) + turn.final_response = data.get("final_response", "") + turn.tools_used = list(data.get("tools_used") or []) + turn.is_error = bool(data.get("is_error", False)) + turn.handoff = bool(data.get("handoff", False)) + turn.user_content_truncated = bool(data.get("user_content_truncated", False)) + turn.user_content_original_chars = int(data.get("user_content_original_chars", 0) or 0) + turn.total_input_tokens = int(data.get("total_input_tokens", 0) or 0) + turn.total_output_tokens = int(data.get("total_output_tokens", 0) or 0) + turn.total_duration_ms = int(data.get("total_duration_ms", 0) or 0) + return turn + + +# ── the codec proper ───────────────────────────────────────────────── + + +def snapshot_chat_turn(st, *, store_blob, generation_seq: int, extra: dict | None = None) -> dict: + """Serialize the live chat turn into a JSON-safe checkpoint payload. + + ``store_blob(bytes) -> ref`` externalizes image payloads. ``extra`` + carries flow metadata (suspension reason, closed trace segment, ...). + """ + payload = { + "codec_version": CODEC_VERSION, + "policy": "chat", + "generation_seq": generation_seq, + "fields": { + "system_prompt": st.system_prompt, + "messages": _externalize_blocks(st.messages, store_blob), + "user_id": st.user_id, + "chat_cap": st.chat_cap, + "iteration": st.iteration, + "tools_used_in_loop": list(st.tools_used_in_loop), + "continuation_count": st.continuation_count, + "max_continuations": st.max_continuations, + "fabrication_retried": st.fabrication_retried, + "promise_retried": st.promise_retried, + "unavail_retried": st.unavail_retried, + "hedging_retried": st.hedging_retried, + "code_hedging_retried": st.code_hedging_retried, + "premature_failure_retried": st.premature_failure_retried, + "pending_image_blocks": _externalize_blocks( + list(st.pending_image_blocks), store_blob + ), + "_op_tool_details": list(st._op_tool_details), + "_pending_validations": list(st._pending_validations), + "_validation_required": st._validation_required, + "_validation_retries": st._validation_retries, + "_max_validation_retries": st._max_validation_retries, + "_result_store_cap": st._result_store_cap, + "_ch_id": st._ch_id, + "_req_id": st._req_id, + "stuck_tracker": export_stuck_tracker(st.stuck_tracker), + "_trajectory": _trajectory_to_payload(st._trajectory), + }, + } + if extra: + payload["extra"] = extra + return payload + + +def restore_field_values(payload: dict, *, load_blob, stuck_tracker_cls) -> dict: + """Persisted-field constructor kwargs for `_ChatTurn`. + + The resume flow combines these with its RECONSTRUCTED fields (live + message, fresh cancel event, current-policy tools, fresh trace, policy + object) — see the classification at the top of this module. + """ + fields = dict(payload.get("fields") or {}) + fields["messages"] = _inline_blocks(fields.get("messages") or [], load_blob) + fields["pending_image_blocks"] = _inline_blocks( + fields.get("pending_image_blocks") or [], load_blob + ) + fields["stuck_tracker"] = import_stuck_tracker( + fields.get("stuck_tracker") or {}, stuck_tracker_cls + ) + fields["_trajectory"] = trajectory_from_payload(fields.get("_trajectory") or {}) + return fields diff --git a/tests/test_turn_checkpoint_codec.py b/tests/test_turn_checkpoint_codec.py new file mode 100644 index 00000000..f12aa075 --- /dev/null +++ b/tests/test_turn_checkpoint_codec.py @@ -0,0 +1,259 @@ +"""Pins for the `_ChatTurn` checkpoint codec (src/turn_state/codec.py). + +The FIELD CENSUS test is the load-bearing one: every `_ChatTurn` dataclass +field must be explicitly classified PERSISTED or RECONSTRUCTED in the codec. +A new field breaks the census until classified — permanently enforcing that +a resume can never silently re-arm the one-shot anti-hedging / +anti-fabrication guard budgets (hard project rule: never weaken them). +""" + +from __future__ import annotations + +import dataclasses +import json +from types import SimpleNamespace + +from src.discord.response_guards import StuckLoopTracker +from src.discord.tool_loop import CHAT_POLICY, _ChatTurn +from src.trajectories.saver import ToolIteration, TrajectoryTurn +from src.turn_state.codec import ( + PERSISTED_FIELDS, + RECONSTRUCTED_FIELDS, + compute_content_digest, + restore_field_values, + snapshot_chat_turn, +) + +GUARD_FLAGS = [ + "fabrication_retried", + "promise_retried", + "unavail_retried", + "hedging_retried", + "code_hedging_retried", + "premature_failure_retried", +] + + +class TestFieldCensus: + def test_every_chat_turn_field_is_classified(self): + """THE completeness pin. If this fails you added/renamed a _ChatTurn + field: classify it in src/turn_state/codec.py (PERSISTED or + RECONSTRUCTED, with the reason) and extend the codec round-trip.""" + actual = {f.name for f in dataclasses.fields(_ChatTurn)} + classified = PERSISTED_FIELDS | RECONSTRUCTED_FIELDS + assert actual == classified, ( + f"unclassified: {sorted(actual - classified)}; " + f"stale: {sorted(classified - actual)}" + ) + + def test_classification_sets_are_disjoint(self): + assert not (PERSISTED_FIELDS & RECONSTRUCTED_FIELDS) + + def test_every_guard_budget_is_persisted(self): + """Resume must never grant fresh guard budgets.""" + for flag in GUARD_FLAGS: + assert flag in PERSISTED_FIELDS + for budget in ("continuation_count", "max_continuations", + "_validation_retries", "_max_validation_retries"): + assert budget in PERSISTED_FIELDS + + +def _blob_dict(): + blobs: dict[str, bytes] = {} + + def store(data: bytes) -> str: + key = f"blob:{len(blobs)}-{len(data)}" + blobs[key] = data + return key + + def load(ref: str) -> bytes: + return blobs[ref] + + return blobs, store, load + + +IMAGE_BLOCK = { + "type": "image", + "source": {"type": "base64", "media_type": "image/png", "data": "aGVsbG8="}, +} + + +def _full_turn(): + """A _ChatTurn with every persisted field forced off its default.""" + trajectory = TrajectoryTurn( + message_id="m1", channel_id="c1", user_id="u1", user_name="user", + timestamp="2026-07-30T12:00:00Z", source="discord", + user_content="do the thing", system_prompt="sys", + history=[{"role": "user", "content": "earlier"}], + ) + trajectory.iterations.append( + ToolIteration( + iteration=0, + tool_calls=[{"id": "call_1", "name": "run_command", "input": {"command": "ls"}}], + tool_results=[{"tool": "run_command", "result": "ok"}], + llm_text="running", input_tokens=10, output_tokens=5, + duration_ms=1200, provider="codex", model="gpt-5.6-sol", + reasoning_effort="xhigh", + ) + ) + tracker = StuckLoopTracker() + tracker.record([{"name": "run_command", "input": {"command": "ls"}}]) + tracker.warned = True + + return _ChatTurn( + message=SimpleNamespace(), # RECONSTRUCTED — never serialized + policy=CHAT_POLICY, + trace=None, + system_prompt="sys REBOUND mid-turn", + tools=[{"name": "run_command"}], + messages=[ + {"role": "user", "content": "do the thing"}, + {"role": "assistant", "content": [ + {"type": "text", "text": "on it"}, + {"type": "tool_use", "id": "call_1", "name": "run_command", + "input": {"command": "ls"}}, + ]}, + {"role": "user", "content": [ + {"type": "tool_result", "tool_use_id": "call_1", "content": "ok"}, + ]}, + {"role": "user", "content": [IMAGE_BLOCK]}, + ], + user_id="u1", + chat_cap=500, + stuck_tracker=tracker, + _trajectory=trajectory, + _result_store_cap=2000, + _cancel=SimpleNamespace(is_set=lambda: False), # RECONSTRUCTED + _ch_id="c1", + _req_id="abcd1234", + iteration=3, + tools_used_in_loop=["run_command", "read_file"], + continuation_count=2, + max_continuations=3, + fabrication_retried=True, + promise_retried=True, + unavail_retried=True, + hedging_retried=True, + code_hedging_retried=True, + premature_failure_retried=True, + pending_image_blocks=[dict(IMAGE_BLOCK)], + _op_tool_details=[{"tool": "run_command", "detail": "x"}], + _pending_validations=["validate ls"], + _validation_required=True, + _validation_retries=1, + _max_validation_retries=2, + ) + + +class TestRoundTrip: + def test_snapshot_is_json_safe_and_restores_every_persisted_field(self): + blobs, store, load = _blob_dict() + st = _full_turn() + payload = snapshot_chat_turn(st, store_blob=store, generation_seq=4) + encoded = json.dumps(payload) # must be JSON-safe + decoded = json.loads(encoded) + + restored = restore_field_values( + decoded, load_blob=load, stuck_tracker_cls=StuckLoopTracker + ) + # Scalars + containers restore verbatim. + for name in PERSISTED_FIELDS - {"stuck_tracker", "_trajectory", + "messages", "pending_image_blocks"}: + assert restored[name] == getattr(st, name), name + # Blob-externalized structures restore to equality. + assert restored["messages"] == st.messages + assert restored["pending_image_blocks"] == st.pending_image_blocks + + def test_all_guard_flags_survive_the_round_trip(self): + """The hard-rule pin: consumed one-shot budgets stay consumed.""" + blobs, store, load = _blob_dict() + payload = snapshot_chat_turn(_full_turn(), store_blob=store, generation_seq=1) + restored = restore_field_values( + json.loads(json.dumps(payload)), load_blob=load, + stuck_tracker_cls=StuckLoopTracker, + ) + for flag in GUARD_FLAGS: + assert restored[flag] is True, flag + assert restored["continuation_count"] == 2 + assert restored["_validation_retries"] == 1 + + def test_stuck_tracker_state_survives(self): + blobs, store, load = _blob_dict() + st = _full_turn() + payload = snapshot_chat_turn(st, store_blob=store, generation_seq=1) + restored = restore_field_values( + json.loads(json.dumps(payload)), load_blob=load, + stuck_tracker_cls=StuckLoopTracker, + ) + tracker = restored["stuck_tracker"] + assert isinstance(tracker, StuckLoopTracker) + assert tracker.warned is True + assert list(tracker._fingerprints) == list(st.stuck_tracker._fingerprints) + + def test_image_data_leaves_the_payload_row(self): + blobs, store, load = _blob_dict() + payload = snapshot_chat_turn(_full_turn(), store_blob=store, generation_seq=1) + assert "aGVsbG8=" not in json.dumps(payload) # externalized + assert len(blobs) >= 1 + + def test_trajectory_restores_with_iteration_revision_guard(self): + blobs, store, load = _blob_dict() + st = _full_turn() + payload = snapshot_chat_turn(st, store_blob=store, generation_seq=1) + data = json.loads(json.dumps(payload)) + # Simulate a double-append artifact: extra iteration rows beyond the + # saved revision must be dropped on restore (Odin round-2). + traj = data["fields"]["_trajectory"] + traj["iterations"] = traj["iterations"] * 3 + restored = restore_field_values( + data, load_blob=load, stuck_tracker_cls=StuckLoopTracker + ) + rebuilt = restored["_trajectory"] + assert len(rebuilt.iterations) == traj["iteration_revision"] == 1 + it = rebuilt.iterations[0] + assert it.tool_calls[0]["id"] == "call_1" + assert rebuilt.system_prompt == "sys" + assert rebuilt.history == [{"role": "user", "content": "earlier"}] + + def test_trajectory_tolerates_field_drift(self): + blobs, store, load = _blob_dict() + payload = snapshot_chat_turn(_full_turn(), store_blob=store, generation_seq=1) + data = json.loads(json.dumps(payload)) + data["fields"]["_trajectory"]["iterations"][0]["field_from_the_future"] = 1 + restored = restore_field_values( + data, load_blob=load, stuck_tracker_cls=StuckLoopTracker + ) + assert len(restored["_trajectory"].iterations) == 1 + + def test_real_blob_store_integration(self, tmp_path): + from src.turn_state import TurnStateStore + + store_obj = TurnStateStore(tmp_path / "t.sqlite3", blob_dir=tmp_path / "blobs") + st = _full_turn() + payload = snapshot_chat_turn( + st, store_blob=store_obj.store_blob_sync, generation_seq=1 + ) + restored = restore_field_values( + json.loads(json.dumps(payload)), + load_blob=store_obj.load_blob_sync, + stuck_tracker_cls=StuckLoopTracker, + ) + assert restored["messages"] == st.messages + store_obj.close() + + def test_payload_metadata(self): + blobs, store, load = _blob_dict() + payload = snapshot_chat_turn( + _full_turn(), store_blob=store, generation_seq=7, + extra={"suspend_reason": "capacity"}, + ) + assert payload["policy"] == "chat" + assert payload["generation_seq"] == 7 + assert payload["extra"] == {"suspend_reason": "capacity"} + + +def test_content_digest_is_full_sha256(): + digest = compute_content_digest("hello") + assert len(digest) == 64 + assert digest != compute_content_digest("hello ") + assert compute_content_digest("") == compute_content_digest(None or "") From 63b5fe5a188dcc651a5c7f47ec0aa37993e87539 Mon Sep 17 00:00:00 2001 From: ATCharpentier Date: Thu, 30 Jul 2026 14:07:22 -0400 Subject: [PATCH 06/18] feat: write invariant wired through the chat loop + suspension MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit New src/turn_state/durability.py: TurnDurability drives the five-point write invariant; disabled instances no-op so the loop keeps one code path. Admission in _prepare_chat_turn (Discord-source chat turns only, resolved the same way the trajectory source is; web shims excluded — v1 fence). Hook placement: WI-1 in _execute_tool_calls (response transcript + PREPARED intents durable BEFORE any execution; malformed intent batches bounce back as matched error results without executing), WI-2 in _run_one_tool (RUNNING gates the effect, fail-closed), WI-3 settles per tool (APPLIED / DEFINITELY_FAILED / OUTCOME_UNKNOWN for interrupted executions, via a no-downgrade guarded settle in the wait_for-timeout race), WI-4 one fenced checkpoint per settled batch, WI-5 at both guard- injection continue points (consumed one-shot flags durable before the LLM retry). on_generation_start persists the absolute UTC recovery deadline before every LLM call. Capacity budget exhaustion now SUSPENDS instead of discarding: work preserved, clean user message naming auto-resume plus the manual resume window; suspension-persistence failure degrades honestly to the plain error (no false preservation claims). Terminal settlement: COMPLETED / FAILED / CANCELLED (cancellation carried past _clear_active clearing the shared event — a cancelled turn never looks resumable), best-effort where no further external effect can follow; durability failures mid-turn halt the turn through the escape guard (fail-closed, proven by integration test: second batch never starts after the store dies). tests/test_write_invariant_integration.py drives the REAL runner against a real store for all of the above. --- src/discord/tool_loop.py | 204 +++++++++++++- src/turn_state/codec.py | 4 + src/turn_state/durability.py | 328 ++++++++++++++++++++++ src/turn_state/store.py | 49 +++- tests/test_typing_resilience.py | 3 + tests/test_write_invariant_integration.py | 265 +++++++++++++++++ 6 files changed, 835 insertions(+), 18 deletions(-) create mode 100644 src/turn_state/durability.py create mode 100644 tests/test_write_invariant_integration.py diff --git a/src/discord/tool_loop.py b/src/discord/tool_loop.py index bda9b379..763ad178 100644 --- a/src/discord/tool_loop.py +++ b/src/discord/tool_loop.py @@ -37,11 +37,14 @@ import discord from ..llm import CircuitOpenError +from ..llm.errors import LLMCapacityError from ..llm.recovery import generate_with_recovery from ..llm.secret_scrubber import scrub_output_secrets from ..observability.correlation import get_turn, set_turn from ..odin_log import get_logger from ..tools import ToolResult +from ..turn_state import LedgerIntentError +from ..turn_state.durability import TurnDurability if TYPE_CHECKING: from ..audit.logger import AuditLogger @@ -322,6 +325,10 @@ class _ChatTurn: _validation_required: bool = False _validation_retries: int = 0 _max_validation_retries: int = 2 + # Process-local durability handle (write-invariant driver). Classified + # RECONSTRUCTED in the checkpoint codec: a resumed turn gets a fresh + # handle bound to the resume lease, never a deserialized one. + durability: TurnDurability = field(default_factory=TurnDurability.disabled) @dataclass @@ -371,6 +378,12 @@ class ToolLoopDeps: audit: AuditLogger loop_manager: LoopManager stuck_loop_tracker_cls: type[StuckLoopTracker] + # Durable turn-state store (None = checkpointing off). Default keeps + # every existing construction working; wiring passes the real store. + turn_store: object | None = None + # Called with (TurnKey, generation) when a turn suspends — wiring points + # it at the resume manager's auto-resume registration. + on_turn_suspended: Callable | None = None class ToolLoopRunner: @@ -392,6 +405,8 @@ def __init__(self, deps: ToolLoopDeps) -> None: self._audit = deps.audit self._loop_manager = deps.loop_manager self._stuck_loop_tracker_cls = deps.stuck_loop_tracker_cls + self._turn_store = deps.turn_store + self._on_turn_suspended = deps.on_turn_suspended # ------------------------------------------------------------------ # Chat pipeline (old _process_with_tools) — orchestrator + phases @@ -421,11 +436,27 @@ async def run( ) try: - return await self._run_chat_iterations(st) + result = await self._run_chat_iterations(st) + # Terminal bookkeeping (best-effort; a suspension already settled + # itself and this no-ops). Cancellation is terminal by design — + # a cancelled turn never becomes resumable. + await st.durability.settle_terminal( + cancelled=st._cancel.is_set(), is_error=bool(result[2]) + ) + return result except asyncio.CancelledError: # Cancellation is not an error turn: release channel ownership - # and let it propagate untouched (no error trajectory). + # and let it propagate untouched (no error trajectory). The + # durable cancel mark is bounded + best-effort so propagation + # stays prompt. self._clear_active(st) + try: + await asyncio.wait_for( + st.durability.settle_terminal(cancelled=True, is_error=False), + timeout=5.0, + ) + except BaseException: # noqa: BLE001 — never delay cancellation + log.warning("Durable cancel mark failed (non-fatal)") raise except Exception as exc: # An escaping exception used to skip BOTH the trajectory record @@ -441,6 +472,14 @@ async def run( log.exception("Trajectory record failed while handling tool-loop escape") finally: self._clear_active(st) + # Fail-closed epilogue: a durability failure (or any escape) + # marks the turn FAILED so a half-written checkpoint can never + # present itself as resumable. Best-effort — a fence loss here + # just means someone else owns the row now. + try: + await st.durability.settle_terminal(cancelled=False, is_error=True) + except Exception: + log.warning("Durable failure mark failed (non-fatal)") raise async def _run_chat_iterations( @@ -468,6 +507,9 @@ async def _run_chat_iterations( kind, val = outcome if kind == "done": return val + # WI-5: the stuck nudge + warned flag go durable before the + # retry generation consumes them. + await st.durability.on_guard_injection(st) continue # nudge injected — next iteration # Gate on actual parsed tool calls, not is_tool_use (which is also @@ -478,6 +520,9 @@ async def _run_chat_iterations( kind, val = await self._finalize_or_retry(st, llm_resp) if kind == "done": return val + # WI-5: consumed one-shot guard flag / continuation budget + + # the injected retry message go durable before the LLM retry. + await st.durability.on_guard_injection(st) continue # retry/continuation message injected # Build internal-format assistant content from LLMResponse @@ -639,6 +684,25 @@ async def _prepare_chat_turn( _req_id = req_hash self._channel_state.set_active_request(_ch_id, _req_id) + # Durable-turn admission: Discord chat turns only, resolved the same + # way the trajectory source is (web/API turns share this runner via + # message shims carrying _odin_source="web" and have no + # re-fetchable Discord message — v1 checkpoint fence). Any refusal → + # a disabled handle and the turn runs exactly as before. + durability = TurnDurability.disabled() + if ( + policy is CHAT_POLICY + and self._turn_store is not None + and _trajectory.source == "discord" + ): + durability = await TurnDurability.admit( + self._turn_store, # type: ignore[arg-type] + message=message, + system_prompt=system_prompt, + tools=tools, + session_snapshot={"history_len": len(history)}, + ) + return _ChatTurn( message=message, policy=policy, @@ -654,6 +718,7 @@ async def _prepare_chat_turn( _cancel=_cancel, _ch_id=_ch_id, _req_id=_req_id, + durability=durability, ) def _clear_active(self, st: _ChatTurn) -> None: @@ -661,6 +726,10 @@ def _clear_active(self, st: _ChatTurn) -> None: def _stopped(self, st: _ChatTurn, where: str) -> tuple[str, bool, bool, list[str], bool]: log.info("Task stopped by /stop in channel %s at %s", st._ch_id, where) + # Carry the cancellation fact past _clear_active (which clears the + # shared event) so terminal settlement records TERMINAL_CANCELLED — + # a cancelled turn must never look resumable or completed. + st.durability.mark_cancelled() self._clear_active(st) suffix = "" if st._pending_validations or st._validation_required: @@ -734,6 +803,11 @@ async def _attempt(): tools_used=st.tools_used_in_loop, ) + # Persist the absolute recovery deadline BEFORE the call: a restart + # mid-recovery reconstructs only the remaining budget, never a fresh + # five minutes. + await st.durability.on_generation_start(st, policy.deadline_seconds) + # Typing is best-effort (shared helper): a typing failure — setup or # cleanup — must never fail the call or misclassify provider errors. async with _best_effort_typing(st.message.channel): @@ -752,20 +826,65 @@ async def _attempt(): # loop-head cancel checks). return ("done", self._stopped(st, "llm_recovery")) raise + except LLMCapacityError as cap_err: + # The whole recovery budget expired on capacity. With + # durability on, the work is preserved and the turn SUSPENDS + # instead of discarding twenty tool calls because the + # twenty-first couldn't reach the model. + if st.durability.enabled: + return ("done", await self._suspend_turn(st, cap_err)) + return ("done", await self._llm_error_done(st, cap_err)) except Exception as api_err: - err_msg = str(api_err) or f"{type(api_err).__name__} (no message)" - log.error("LLM API call failed: %s", err_msg, exc_info=True) - await self._turn_recorder._save_turn_trajectory( - st._trajectory, error=err_msg, trace=st.trace - ) - self._clear_active(st) - return ( - "done", - (f"LLM API error: {err_msg}", False, True, st.tools_used_in_loop, False), - ) + return ("done", await self._llm_error_done(st, api_err)) return ("ok", llm_resp) + async def _llm_error_done(self, st: _ChatTurn, api_err: BaseException): + """The legacy terminal LLM-failure path (unchanged message shape).""" + err_msg = str(api_err) or f"{type(api_err).__name__} (no message)" + log.error("LLM API call failed: %s", err_msg, exc_info=True) + await self._turn_recorder._save_turn_trajectory( + st._trajectory, error=err_msg, trace=st.trace + ) + self._clear_active(st) + return (f"LLM API error: {err_msg}", False, True, st.tools_used_in_loop, False) + + async def _suspend_turn(self, st: _ChatTurn, cap_err: LLMCapacityError): + """Suspend with preserved work; falls back to the plain error when + suspension persistence itself fails (no false preservation claims).""" + reason = str(cap_err) or "capacity exhausted" + preserved = await st.durability.suspend(st, reason) + if not preserved: + return await self._llm_error_done(st, cap_err) + + minutes = max( + 1, round(self._llm_gateway.recovery_policy().deadline_seconds / 60.0) + ) + model = cap_err.model or "The model" + n_tools = len(st.tools_used_in_loop) + text = ( + f"{model} is out of capacity — I retried for ~{minutes} minute(s) " + f"without getting through. I've preserved everything done so far " + f"({n_tools} tool call(s)) and will pick this up automatically when " + "capacity returns, as long as nothing else happens in this channel. " + "You can also reply `resume` within 24h to continue manually." + ) + await self._turn_recorder._save_turn_trajectory( + st._trajectory, + final_response=text, + tools_used=st.tools_used_in_loop, + trace=st.trace, + ) + self._clear_active(st) + if self._on_turn_suspended is not None and st.durability.lease is not None: + try: + self._on_turn_suspended( + st.durability.lease.key, st.durability.lease.generation + ) + except Exception: + log.exception("Auto-resume registration failed (non-fatal)") + return (text, False, True, st.tools_used_in_loop, False) + async def _check_stuck_and_record(self, st: _ChatTurn, llm_resp): """Record this iteration's tool calls + LLM text into the trajectory and stuck tracker; terminate or nudge on a confirmed repeat cycle. @@ -1002,6 +1121,11 @@ async def _run_one_tool(self, st: _ChatTurn, block) -> dict: _rbac_denial = self._tool_executor.check_permission(tool_name, _uid) if isinstance(_rbac_denial, str) and _rbac_denial: # str = deny, None = allow log.warning("RBAC gate denied tool %s for user %s", tool_name, _uid) + # Denied before any effect: settle the PREPARED intent as + # definitely-not-applied so it can never look interrupted. + await st.durability.after_tool( + block, ok=False, uncertain=False, result_text=_rbac_denial + ) return {"type": "tool_result", "tool_use_id": block.id, "content": _rbac_denial} await self._delivery.set_status( TOOL_STATUS_LABELS.get(tool_name, f"Running: {tool_name}") @@ -1021,8 +1145,14 @@ async def _run_one_tool(self, st: _ChatTurn, block) -> dict: except Exception: pass + # WI-2: PREPARED→RUNNING BEFORE the external effect. A durability + # failure here raises and blocks the effect (fail-closed) — the run() + # escape guard turns it into a bounded error turn. + await st.durability.before_tool(block) + t0 = time.monotonic() error = None + uncertain_outcome = False tool_result = None # Handle Discord-native tools try: @@ -1054,15 +1184,21 @@ async def _run_one_tool(self, st: _ChatTurn, block) -> dict: error = str(e) result = f"Tool {tool_name} timed out: {e}" tool_result = None + # An execution that started and died mid-flight may have applied + # its effect — the ledger records OUTCOME_UNKNOWN, never a + # confident failure (replay must not rerun it). + uncertain_outcome = True log.warning("Tool %s timed out after %.1fs", tool_name, time.monotonic() - t0) except (ValueError, KeyError, TypeError) as e: error = str(e) result = f"Tool {tool_name} input error: {e}" tool_result = None + uncertain_outcome = True except Exception as e: error = str(e) result = f"Error executing {tool_name}: {e}" tool_result = None + uncertain_outcome = True log.warning("Unexpected tool error for %s: %s", tool_name, e) elapsed_ms = int((time.monotonic() - t0) * 1000) @@ -1109,6 +1245,16 @@ async def _run_one_tool(self, st: _ChatTurn, block) -> dict: # Truncate large outputs before sending back to the LLM. tool_content = truncate_tool_output(result) + # WI-3: settle the op right after completion. ok → APPLIED; + # a tool-reported failure → DEFINITELY_FAILED; an exception mid- + # execution → OUTCOME_UNKNOWN (the effect may have applied). + await st.durability.after_tool( + block, + ok=(tool_result.ok if tool_result is not None else error is None), + uncertain=uncertain_outcome, + result_text=tool_content, + ) + return { "type": "tool_result", "tool_use_id": block.id, @@ -1165,6 +1311,13 @@ async def _run_one_tool_with_timeout(self, st: _ChatTurn, block, tool_timeout) - ) except TimeoutError: error_msg = f"Tool '{block.name}' timed out after {t}s" + # WI-3 (interrupted): wait_for cancelled _run_one_tool before its + # own settle — the external effect may or may not have applied. + # OUTCOME_UNKNOWN via the no-downgrade settle, never rerun. + try: + await st.durability.after_tool_interrupted(block, error_msg) + except Exception: + log.exception("Ledger settle failed for timed-out %s", block.name) try: await self._audit.log_execution( user_id=str(st.message.author.id), @@ -1190,6 +1343,28 @@ async def _execute_tool_calls(self, st: _ChatTurn, tool_calls) -> list: result block to the message list (gather preserves call order).""" tool_timeout = self._get_config().tools.tool_timeout_seconds + # WI-1: the LLM response transcript + PREPARED intents are durable + # BEFORE any execution. Malformed intents (empty/duplicate call ids) + # fail here and are bounced back as matched error results without + # executing anything (Odin's rule: fail before execution). + try: + await st.durability.on_llm_response(st, tool_calls) + except LedgerIntentError as intent_err: + log.warning("Tool batch rejected before execution: %s", intent_err) + tool_results = [ + { + "type": "tool_result", + "tool_use_id": b.id, + "content": ( + f"Error: malformed tool-call batch ({intent_err}). " + "NOT executed — re-issue the calls with unique ids." + ), + } + for b in tool_calls + ] + st.messages.append({"role": "user", "content": list(tool_results)}) + return tool_results + # Best-effort typing: the 2026-07-16 Discord incident (typing # endpoint returning HTML 500s) aborted whole turns at exactly this # line and dumped raw DiscordServerError HTML into chat. @@ -1281,6 +1456,11 @@ async def _post_iteration(self, st: _ChatTurn, tool_calls, tool_results): ) st.pending_image_blocks.clear() + # WI-4: the batch is fully settled — tool-result continuation, + # trajectory results, validation/vision state, and the advanced + # iteration go durable in one fenced checkpoint (real progress). + await st.durability.on_batch_settled(st) + return None def _check_skill_handoff(self, st: _ChatTurn, tool_calls, tool_results): diff --git a/src/turn_state/codec.py b/src/turn_state/codec.py index 5723b65b..7ff83ea9 100644 --- a/src/turn_state/codec.py +++ b/src/turn_state/codec.py @@ -77,12 +77,16 @@ #: - trace: ContextTraceCollector.finalize() is one-way; the old segment #: is closed into the payload for diagnostics and a fresh linked resume #: segment starts (Odin round-2). +#: - durability: the write-invariant driver — always process-local; a +#: resumed turn gets a fresh handle bound to the RESUME lease (restoring +#: the old one would carry a fenced-out lease token). RECONSTRUCTED_FIELDS: frozenset[str] = frozenset({ "message", "_cancel", "tools", "policy", "trace", + "durability", }) diff --git a/src/turn_state/durability.py b/src/turn_state/durability.py new file mode 100644 index 00000000..2cc12d11 --- /dev/null +++ b/src/turn_state/durability.py @@ -0,0 +1,328 @@ +"""Write-invariant driver for one chat turn. + +`TurnDurability` is the seam between the tool loop and the store: the loop +calls the five invariant hooks (Odin's write invariant, round-2) and this +module does the snapshot + fenced persistence. A disabled instance (store +missing/unavailable, admission refused, non-Discord source) no-ops every +hook so the loop keeps ONE code path. + +The invariant, mapped to hooks: + +1. `on_llm_response` — the successful LLM response (transcript incl. the + assistant tool_use message) and ALL proposed tool intents are durable + BEFORE any execution. +2. `before_tool` — each intent transitions PREPARED→RUNNING before its + external effect; a persistence failure here BLOCKS the effect + (fail-closed). +3. `after_tool` — each result is persisted immediately after + completion (APPLIED / DEFINITELY_FAILED / OUTCOME_UNKNOWN for + timeouts-and-interruptions, which may have applied). +4. `on_batch_settled` — after the parallel batch settles: the tool-result + continuation message, trajectory/guard/validation/stuck state, and the + advanced iteration go durable atomically (one fenced checkpoint). +5. `on_guard_injection` — a guard nudge and its consumed one-shot flag are + durable BEFORE the LLM retry that consumes them. + +`on_generation_start` additionally persists the absolute UTC recovery +deadline before each LLM call, so a process restart reconstructs only the +REMAINING budget (never a fresh five minutes). + +Failure semantics: hooks raise `TurnStateUnavailableError` / `StaleTurnError` +— the loop halts (fail-closed). Terminal bookkeeping (`settle_terminal`, +`suspend`) is best-effort where no further external effect can follow. +""" + +from __future__ import annotations + +import asyncio +import hashlib +import time +from typing import Any + +from ..odin_log import get_logger +from .codec import compute_content_digest, snapshot_chat_turn +from .store import OpState, TurnKey, TurnLease, TurnStateStore, TurnStatus + +log = get_logger("turn_state") + +# Bound stored per-op result text (the model-visible copy already rides the +# transcript; the ledger copy is for replay/reconciliation). +_OP_RESULT_CAP = 4000 + + +def _hash_text(text: str) -> str: + return hashlib.sha256((text or "").encode("utf-8", "replace")).hexdigest() + + +class TurnDurability: + """Per-turn durability handle. Carried on the turn state object + (classified RECONSTRUCTED in the codec — always process-local).""" + + def __init__(self, store: TurnStateStore | None, lease: TurnLease | None) -> None: + self._store = store + self._lease = lease + self.generation_seq = 0 + self.suspended = False + self.settled = False + # Set by the loop's /stop path: _clear_active clears the shared + # cancel EVENT before terminal settlement can read it, so the fact + # must be carried here (a cancelled turn is terminal by design). + self.cancelled = False + + # -- construction -------------------------------------------------- + + @classmethod + def disabled(cls) -> TurnDurability: + return cls(None, None) + + @classmethod + async def admit( + cls, + store: TurnStateStore | None, + *, + message: Any, + system_prompt: str, + tools: list | None, + session_snapshot: dict | None, + ) -> TurnDurability: + """Admit a fresh Discord chat turn; any refusal → disabled handle.""" + if store is None or not store.available: + return cls.disabled() + try: + key = TurnKey( + source="discord", + channel_id=str(message.channel.id), + message_id=str(message.id), + ) + tool_names = sorted(t.get("name", "") for t in (tools or [])) + lease = await asyncio.to_thread( + store.admit_turn_sync, + key, + guild_id=str(getattr(getattr(message, "guild", None), "id", "") or ""), + user_id=str(message.author.id), + content_digest=compute_content_digest( + getattr(message, "content", "") or "" + ), + code_version=_code_version(), + prompt_policy_hash=_hash_text(system_prompt), + tool_catalog_hash=_hash_text(",".join(tool_names)), + session_snapshot=session_snapshot, + ) + except Exception: + log.exception("Turn admission failed — running without durability") + return cls.disabled() + if lease is None: + return cls.disabled() + return cls(store, lease) + + @classmethod + def resumed( + cls, store: TurnStateStore, lease: TurnLease, generation_seq: int + ) -> TurnDurability: + handle = cls(store, lease) + handle.generation_seq = int(generation_seq) + return handle + + # -- state --------------------------------------------------------- + + @property + def enabled(self) -> bool: + return ( + self._store is not None + and self._lease is not None + and not self.suspended + and not self.settled + ) + + @property + def lease(self) -> TurnLease | None: + return self._lease + + # -- snapshot plumbing (runs in a worker thread: blob writes + sqlite) -- + + def _checkpoint_sync( + self, + st, + *, + progressed: bool, + recovery_deadline_utc: float | None = None, + extra: dict | None = None, + ) -> None: + assert self._store is not None and self._lease is not None + payload = snapshot_chat_turn( + st, + store_blob=self._store.store_blob_sync, + generation_seq=self.generation_seq, + extra=extra, + ) + self._store.checkpoint_sync( + self._lease, + payload, + progressed=progressed, + recovery_deadline_utc=recovery_deadline_utc, + ) + + # -- the invariant hooks ------------------------------------------- + + async def on_generation_start(self, st, deadline_seconds: float) -> None: + """Persist the absolute recovery deadline before the LLM call.""" + if not self.enabled: + return + await asyncio.to_thread( + self._checkpoint_sync, + st, + progressed=False, + recovery_deadline_utc=time.time() + max(0.0, deadline_seconds), + ) + + async def on_llm_response(self, st, tool_calls: list) -> None: + """WI-1: response transcript + PREPARED intents, before any effect. + + Parse-error calls are excluded (they are never executed; on resume a + tool_use block with no ledger row truthfully means "never ran"). + Raises LedgerIntentError on empty/duplicate executable ids — the + loop converts that to matched error results WITHOUT executing. + """ + if not self.enabled: + return + self.generation_seq += 1 + await asyncio.to_thread(self._checkpoint_sync, st, progressed=True) + executable = [tc for tc in tool_calls if not getattr(tc, "parse_error", None)] + if not executable: + return + intents = [ + {"tool_call_id": tc.id, "tool_name": tc.name, "tool_input": tc.input} + for tc in executable + ] + assert self._store is not None and self._lease is not None + await asyncio.to_thread( + self._store.record_intents_sync, + self._lease, + self.generation_seq, + intents, + iteration=st.iteration, + ) + + async def before_tool(self, block) -> None: + """WI-2: PREPARED→RUNNING gates the external effect (fail-closed).""" + if not self.enabled: + return + assert self._store is not None and self._lease is not None + await asyncio.to_thread( + self._store.mark_running_sync, self._lease, self.generation_seq, block.id + ) + + async def after_tool(self, block, *, ok: bool, uncertain: bool, result_text: str) -> None: + """WI-3: settle each op right after it completes.""" + if not self.enabled: + return + if ok: + state = OpState.APPLIED + elif uncertain: + state = OpState.OUTCOME_UNKNOWN + else: + state = OpState.DEFINITELY_FAILED + assert self._store is not None and self._lease is not None + await asyncio.to_thread( + self._store.settle_op_sync, + self._lease, + self.generation_seq, + block.id, + state=state, + result_text=(result_text or "")[:_OP_RESULT_CAP], + ) + + async def after_tool_interrupted(self, block, result_text: str) -> None: + """WI-3 for a wait_for-cancelled execution: OUTCOME_UNKNOWN via the + no-downgrade guarded settle (it may race the tool's own settle at + the cancellation boundary; an already-settled row wins).""" + if not self.enabled: + return + assert self._store is not None and self._lease is not None + await asyncio.to_thread( + self._store.settle_interrupted_sync, + self._lease, + self.generation_seq, + block.id, + result_text=(result_text or "")[:_OP_RESULT_CAP], + ) + + async def on_batch_settled(self, st) -> None: + """WI-4: one fenced checkpoint after the batch + bookkeeping.""" + if not self.enabled: + return + await asyncio.to_thread(self._checkpoint_sync, st, progressed=True) + + async def on_guard_injection(self, st) -> None: + """WI-5: consumed one-shot flags durable before the LLM retry. + + Not "real progress" — guard nudges must not extend the resumable + TTL the way completed tool batches do. + """ + if not self.enabled: + return + await asyncio.to_thread(self._checkpoint_sync, st, progressed=False) + + # -- terminal transitions ------------------------------------------ + + async def suspend(self, st, reason: str) -> bool: + """Suspend with preserved work. Returns False when persistence + failed (caller reports the legacy error instead of promising a + preserved checkpoint that does not exist).""" + if not self.enabled: + return False + extra: dict = {"suspend_reason": reason} + try: + if st.trace is not None: + extra["closed_trace"] = st.trace.finalize() + except Exception: + pass + try: + assert self._store is not None and self._lease is not None + payload = await asyncio.to_thread( + lambda: snapshot_chat_turn( + st, + store_blob=self._store.store_blob_sync, + generation_seq=self.generation_seq, + extra=extra, + ) + ) + await asyncio.to_thread(self._store.suspend_sync, self._lease, payload) + except Exception: + log.exception("Turn suspension failed — work NOT preserved") + return False + self.suspended = True + return True + + def mark_cancelled(self) -> None: + self.cancelled = True + + async def settle_terminal(self, *, cancelled: bool, is_error: bool) -> None: + """Best-effort terminal bookkeeping after the turn's reply exists. + + No further external effect follows, so a failure here must not + destroy an already-computed reply — log and move on. + """ + if not self.enabled: + return + if cancelled or self.cancelled: + status = TurnStatus.TERMINAL_CANCELLED + elif is_error: + status = TurnStatus.TERMINAL_FAILED + else: + status = TurnStatus.TERMINAL_COMPLETED + try: + assert self._store is not None and self._lease is not None + await asyncio.to_thread(self._store.finish_sync, self._lease, status) + except Exception: + log.exception("Turn terminal bookkeeping failed (non-fatal)") + self.settled = True + + +def _code_version() -> str: + try: + from ..version import get_version + + return str(get_version()) + except Exception: + return "unknown" diff --git a/src/turn_state/store.py b/src/turn_state/store.py index 63142dec..c93ebc74 100644 --- a/src/turn_state/store.py +++ b/src/turn_state/store.py @@ -284,11 +284,14 @@ def _verify_lease(self, lease: TurnLease) -> None: """Ops-table writes are fenced indirectly: verify the turns row still carries this lease before touching operations.""" conn = self._require() - row = conn.execute( - "SELECT lease_token, turn_generation FROM turns " - "WHERE source=? AND channel_id=? AND message_id=?", - [lease.key.source, lease.key.channel_id, lease.key.message_id], - ).fetchone() + try: + row = conn.execute( + "SELECT lease_token, turn_generation FROM turns " + "WHERE source=? AND channel_id=? AND message_id=?", + [lease.key.source, lease.key.channel_id, lease.key.message_id], + ).fetchone() + except sqlite3.Error as exc: + raise TurnStateUnavailableError(f"turn-state read failed: {exc}") from exc if row is None or row[0] != lease.token or row[1] != lease.generation: raise StaleTurnError(f"turn {lease.key}: lease no longer held") @@ -444,10 +447,44 @@ def settle_op_sync( state: str, result_text: str | None, ) -> None: - if state not in (OpState.APPLIED, OpState.DEFINITELY_FAILED): + # OUTCOME_UNKNOWN is a legitimate settlement: a timed-out or + # interrupted execution may have applied its external effect — the + # ledger must never claim DEFINITELY_FAILED for it (round-2 replay + # rule: ambiguous outcomes stop, never rerun). + if state not in (OpState.APPLIED, OpState.DEFINITELY_FAILED, OpState.OUTCOME_UNKNOWN): raise ValueError(f"settle_op_sync: invalid terminal state {state}") self._set_op_state(lease, generation_seq, tool_call_id, state, result_text) + def settle_interrupted_sync( + self, + lease: TurnLease, + generation_seq: int, + tool_call_id: str, + *, + result_text: str | None, + ) -> None: + """OUTCOME_UNKNOWN settle for an interrupted execution. + + Guarded so it can race the tool's own settle at a cancellation + boundary: it never downgrades an already-settled row, and a missing + row (e.g. a parse-error call that had no intent) is tolerated. + """ + conn = self._require() + self._verify_lease(lease) + where, params = self._op_where(lease) + try: + with self._write_lock: + conn.execute( + f"UPDATE operations SET state=?, result=?, updated_at=? " + f"WHERE {where} AND generation_seq=? AND tool_call_id=? " + "AND state IN (?, ?)", + [OpState.OUTCOME_UNKNOWN, result_text, time.time(), *params, + generation_seq, tool_call_id, OpState.PREPARED, OpState.RUNNING], + ) + conn.commit() + except sqlite3.Error as exc: + raise TurnStateUnavailableError(f"ledger write failed: {exc}") from exc + def _set_op_state( self, lease: TurnLease, diff --git a/tests/test_typing_resilience.py b/tests/test_typing_resilience.py index 07de3427..dc79f162 100644 --- a/tests/test_typing_resilience.py +++ b/tests/test_typing_resilience.py @@ -241,6 +241,8 @@ async def _default_save(trajectory, **kwargs): def _stub_state(channel=None): + from src.turn_state.durability import TurnDurability + return SimpleNamespace( chat_cap=3, iteration=0, @@ -255,6 +257,7 @@ def _stub_state(channel=None): system_prompt="sys", tools=[], user_id="u1", + durability=TurnDurability.disabled(), ) diff --git a/tests/test_write_invariant_integration.py b/tests/test_write_invariant_integration.py new file mode 100644 index 00000000..82ebb771 --- /dev/null +++ b/tests/test_write_invariant_integration.py @@ -0,0 +1,265 @@ +"""Integration pins: the write invariant driven through the REAL chat loop. + +Each test runs ToolLoopRunner.run against a real TurnStateStore and asserts +the durable state the invariant promises: intents before execution, settles +after completion, one fenced checkpoint per batch, suspension with +preserved work on capacity exhaustion, terminal tombstones, and the +fail-closed halt when durability dies mid-turn. +""" + +from __future__ import annotations + +import pytest + +from src.llm.errors import LLMCapacityError +from src.llm.recovery import RecoveryPolicy +from src.tools.result_validator import ToolResult +from src.turn_state import OpState, TurnStateUnavailableError, TurnStatus +from src.turn_state.store import TurnStateStore +from tests.fakes import FakeLLM, FakeMessage, make_bot, text_response, tool_call_response + +FAST_POLICY = RecoveryPolicy( + deadline_seconds=0.2, backoff_base=0.01, backoff_cap=0.02, retry_after_cap=0.05 +) + + +@pytest.fixture(autouse=True) +def _isolated_cwd(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + + +def build_with_store(script, tmp_path, **overrides): + fake = FakeLLM(script) + bot = make_bot(fake_llm=fake, config_overrides=overrides or None) + store = TurnStateStore(tmp_path / "turn_state" / "turns.sqlite3") + bot.tool_loop._turn_store = store + bot.llm_gateway._recovery_policy_source = lambda: FAST_POLICY + return bot, fake, store + + +def capacity_forever(fake): + """Script entry raising capacity and re-arming itself (persistent outage).""" + + def _raise(): + fake.responses.append(_raise) + raise LLMCapacityError( + "Codex capacity: server_is_overloaded", provider="codex", model="fake-model" + ) + + return _raise + + +async def run_loop(bot, msg): + return await bot.tool_loop.run(msg, [{"role": "user", "content": msg.content}]) + + +def turn_row(store, cols="status, revision, payload"): + return store._conn.execute(f"SELECT {cols} FROM turns").fetchone() + + +def op_rows(store): + return store._conn.execute( + "SELECT tool_call_id, state, result FROM operations ORDER BY tool_call_id" + ).fetchall() + + +class TestHappyPath: + async def test_completed_turn_settles_and_compacts(self, tmp_path): + bot, fake, store = build_with_store( + [ + tool_call_response(("parse_time", {"text": "tomorrow 3pm"})), + text_response("Parsed it."), + ], + tmp_path, + ) + text, _, is_error, tools_used, _ = await run_loop(bot, FakeMessage("parse")) + assert text == "Parsed it." + assert is_error is False + + status, revision, payload = turn_row(store) + assert status == TurnStatus.TERMINAL_COMPLETED + assert payload is None # compacted immediately + assert revision >= 3 # deadline + WI-1 + WI-4 at minimum + + ops = op_rows(store) + assert len(ops) == 1 + call_id, state, result = ops[0] + assert call_id == "call-1" + assert state == OpState.APPLIED + assert result # the model-visible result text rode the ledger + + async def test_intents_are_durable_and_running_before_execution(self, tmp_path): + bot, fake, store = build_with_store( + [ + tool_call_response(("run_command", {"command": "ls"})), + text_response("done"), + ], + tmp_path, + ) + seen_states = [] + + async def probing_execute(tool_name, tool_input, *, user_id=None): + # WI-1/WI-2: at execution time the intent row already exists and + # is RUNNING — durable BEFORE the external effect. + rows = op_rows(store) + seen_states.append(rows[0][1] if rows else None) + return ToolResult(output="probe ok", tool_name=tool_name) + + bot.tool_executor.execute = probing_execute + text, *_ = await run_loop(bot, FakeMessage("go")) + assert text == "done" + assert seen_states == [OpState.RUNNING] + +class TestSuspension: + async def test_capacity_exhaustion_suspends_with_preserved_work(self, tmp_path): + bot, fake, store = build_with_store( + [tool_call_response(("parse_time", {"text": "tomorrow"}))], tmp_path + ) + fake.responses.append(capacity_forever(fake)) + + text, _, is_error, tools_used, _ = await run_loop(bot, FakeMessage("parse")) + assert is_error is True + assert "preserved" in text + assert "resume" in text + assert tools_used == ["parse_time"] + + status, _, payload = turn_row(store) + assert status == TurnStatus.SUSPENDED + assert payload is not None + + import json as _json + + fields = _json.loads(payload)["fields"] + # The transcript survived: assistant tool_use + matched tool_result. + roles = [m.get("role") for m in fields["messages"]] + assert "assistant" in roles + assert fields["tools_used_in_loop"] == ["parse_time"] + # The first generation's op settled APPLIED before the outage. + ops = op_rows(store) + assert ops and ops[0][1] == OpState.APPLIED + # Recovery deadline was persisted as absolute UTC. + (deadline,) = store._conn.execute( + "SELECT recovery_deadline_utc FROM turns" + ).fetchone() + assert deadline is not None + + async def test_consumed_guard_flag_is_durable_in_suspension(self, tmp_path): + # Hedging guard fires (one-shot flag consumed) → capacity kills the + # retry generation → the suspended payload must carry the consumed + # flag (WI-5): a future resume gets NO fresh hedging budget. + bot, fake, store = build_with_store( + [text_response("Shall I proceed with the deployment now?")], tmp_path + ) + fake.responses.append(capacity_forever(fake)) + + text, _, is_error, *_ = await run_loop(bot, FakeMessage("deploy it")) + assert is_error is True + + import json as _json + + status, _, payload = turn_row(store) + assert status == TurnStatus.SUSPENDED + fields = _json.loads(payload)["fields"] + assert fields["hedging_retried"] is True + + async def test_suspension_persistence_failure_reports_plain_error(self, tmp_path): + bot, fake, store = build_with_store( + [tool_call_response(("parse_time", {"text": "x"}))], tmp_path + ) + fake.responses.append(capacity_forever(fake)) + + async def kill_store_then_probe(tool_name, tool_input, *, user_id=None): + return ToolResult(output="ok", tool_name=tool_name) + + # Sever durability AFTER the first batch (mid-turn) but make the + # suspension write the first failing one: monkey-close the conn just + # before capacity exhausts. Simplest deterministic point: close it + # inside the capacity callable's first raise. + original = fake.responses[-1] + + def close_then_capacity(): + store._conn.close() + return original() + + fake.responses[-1] = close_then_capacity + + with pytest.raises((TurnStateUnavailableError, Exception)): + # Depending on which durable write hits the closed connection + # first (deadline checkpoint vs suspension), the turn either + # fail-closes (raise) or reports the plain error tuple. Both are + # acceptable; what is FORBIDDEN is a "work preserved" claim. + result = await run_loop(bot, FakeMessage("x")) + assert "preserved" not in result[0] + raise TurnStateUnavailableError("reached tuple path (acceptable)") + + +class TestTerminalStates: + async def test_cancelled_turn_is_terminal_cancelled(self, tmp_path): + bot, fake, store = build_with_store( + [tool_call_response(("run_command", {"command": "x"}))], tmp_path + ) + + async def cancel_during_tool(tool_name, tool_input, *, user_id=None): + ch_id = str(FakeMessage("x").channel.id) + bot.channel_state.cancel_events[ch_id].set() + return ToolResult(output="ok", tool_name=tool_name) + + bot.tool_executor.execute = cancel_during_tool + text, *_ = await run_loop(bot, FakeMessage("go")) + assert text.startswith("Task stopped by user.") + (status,) = store._conn.execute("SELECT status FROM turns").fetchone() + assert status == TurnStatus.TERMINAL_CANCELLED + + async def test_tool_timeout_settles_outcome_unknown(self, tmp_path): + bot, fake, store = build_with_store( + [ + tool_call_response(("run_command", {"command": "x"})), + text_response("moved on"), + ], + tmp_path, + ) + + async def raise_timeout(tool_name, tool_input, *, user_id=None): + raise TimeoutError("simulated in-execution timeout") + + bot.tool_executor.execute = raise_timeout + text, *_ = await run_loop(bot, FakeMessage("go")) + assert text == "moved on" + ops = op_rows(store) + # The interrupted execution may have applied — UNKNOWN, never a + # confident failure, never rerun. + assert ops and ops[0][1] == OpState.OUTCOME_UNKNOWN + + async def test_llm_error_turn_is_terminal_failed(self, tmp_path): + bot, fake, store = build_with_store([RuntimeError("boom")], tmp_path) + text, _, is_error, *_ = await run_loop(bot, FakeMessage("go")) + assert is_error is True + (status,) = store._conn.execute("SELECT status FROM turns").fetchone() + assert status == TurnStatus.TERMINAL_FAILED + + +class TestFailClosed: + async def test_mid_turn_durability_death_halts_the_turn(self, tmp_path): + # Two batches scripted; durability dies during the FIRST tool's + # execution → the WI-3/WI-4 write raises → the turn halts through + # the escape guard instead of silently continuing without durability. + bot, fake, store = build_with_store( + [ + tool_call_response(("run_command", {"command": "one"})), + tool_call_response(("run_command", {"command": "two"})), + text_response("never reached"), + ], + tmp_path, + ) + executed = [] + + async def kill_durability(tool_name, tool_input, *, user_id=None): + executed.append(tool_input["command"]) + store._conn.close() + return ToolResult(output="ok", tool_name=tool_name) + + bot.tool_executor.execute = kill_durability + with pytest.raises(TurnStateUnavailableError): + await run_loop(bot, FakeMessage("go")) + # The second batch never started: fail closed, not fail quiet. + assert executed == ["one"] From 295ee3a1db4291862f067f771a30b927d036142f Mon Sep 17 00:00:00 2001 From: ATCharpentier Date: Thu, 30 Jul 2026 14:17:20 -0400 Subject: [PATCH 07/18] =?UTF-8?q?feat:=20suspended-turn=20resume=20?= =?UTF-8?q?=E2=80=94=20explicit=20+=20auto,=20admission-gated?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit New src/discord/turn_resume.py (TurnResumeManager). Explicit path: a bare resume/continue from the ORIGINAL requester, running inside the normal intake pipeline (lock/delivery/session inherited); the trigger is consumed as a command and never enters the frozen transcript. Auto path: registered at suspension, in-process only (restart => explicit-only by design); the waiter samples a session baseline, then actively claims the model breaker probe slot when the cooldown elapses — the resumed generation IS the capacity probe, and a re-suspension re-arms the waiter with the breaker escalating cooldown as natural pacing. Auto-resume stands down fail-safe when the session advanced or cannot be read. Admission (both paths): re-fetch the original message; same author + unchanged content digest required; deleted/edited => TERMINAL_REJECTED (payload compacted, never executable folklore); tools re-derived from the CURRENT catalog + permission filter — current security policy wins; single resume winner via the fenced resume lease. Replay: unmatched tool_use blocks are repaired from the ledger (APPLIED replays its stored result; unknown/never-ran stated truthfully as result blocks) — matched blocks guaranteed, NOTHING re-executed. The resumed generation carries only its REMAINING recovery budget from the persisted UTC deadline (one attempt when exhausted — never a fresh five minutes); later generations budget normally. run_resumed() enters the same guard envelope as run(); the iteration loop starts from the restored index. Intake: explicit-resume branch ahead of tool_loop.run; the error-path session marker now says the work is PRESERVED and resumable when a suspension row exists (was: "may ask to retry"). --- src/discord/intake_pipeline.py | 67 ++++-- src/discord/tool_loop.py | 38 +++- src/discord/turn_resume.py | 392 +++++++++++++++++++++++++++++++++ src/turn_state/durability.py | 24 +- tests/test_resume_admission.py | 322 +++++++++++++++++++++++++++ 5 files changed, 826 insertions(+), 17 deletions(-) create mode 100644 src/discord/turn_resume.py create mode 100644 tests/test_resume_admission.py diff --git a/src/discord/intake_pipeline.py b/src/discord/intake_pipeline.py index ae5177e6..5b5454dd 100644 --- a/src/discord/intake_pipeline.py +++ b/src/discord/intake_pipeline.py @@ -406,6 +406,9 @@ class MessagePipelineDeps: tool_loop: ToolLoopRunner # the tools route delivery: ResponseDelivery # status, retries, chunked sends housekeeping: Housekeeping # post-turn cache maintenance + # Suspended-turn resume manager (None = feature off). Default keeps every + # existing construction working; wiring passes the real manager. + turn_resume: object | None = None class MessagePipeline: @@ -419,6 +422,7 @@ def __init__(self, deps: MessagePipelineDeps) -> None: self._tool_loop = deps.tool_loop self._delivery = deps.delivery self._housekeeping = deps.housekeeping + self._turn_resume = deps.turn_resume async def run( self, @@ -603,18 +607,35 @@ async def _run_inner( } log.info("Attached %d image(s) to message for Claude vision", len(image_blocks)) try: - ( - response, - already_sent, - is_error, - tools_used, - handoff, - ) = await self._tool_loop.run( - message, - task_history, - system_prompt_override=_sp, - trace=_trace, - ) + # A bare `resume`/`continue` with preserved work in this + # channel resumes the suspended turn instead of starting + # a fresh one — the trigger is consumed as a command and + # never enters the frozen transcript. Inherits this + # pipeline's lock/delivery/session plumbing wholesale. + _resumed = None + if self._turn_resume is not None: + _resumed = await self._turn_resume.try_explicit_resume(message) + if _resumed is not None: + ( + response, + already_sent, + is_error, + tools_used, + handoff, + ) = _resumed + else: + ( + response, + already_sent, + is_error, + tools_used, + handoff, + ) = await self._tool_loop.run( + message, + task_history, + system_prompt_override=_sp, + trace=_trace, + ) except TimeoutError as codex_err: _err = format_user_facing_error(codex_err) log.warning("Codex tool loop timed out: %s", _err) @@ -736,7 +757,27 @@ async def _run_inner( # Save a sanitized error marker instead of the full error response. # The user sees the full error on Discord, but raw refusals and # fabrications are NOT persisted to prevent context poisoning. - if tools_used: + _preserved = False + if self._turn_resume is not None: + try: + _preserved = await self._turn_resume.is_suspended( + channel_id, str(message.id) + ) + except Exception: + _preserved = False + if _preserved: + _tools_note = ( + f" after using tools ({', '.join(tools_used[:5])})" + if tools_used + else "" + ) + sanitized = ( + "[Previous request was interrupted by a model-capacity " + f"outage{_tools_note}. Its work is PRESERVED and resumable — " + "it auto-resumes when capacity returns, or the user can " + "say 'resume'.]" + ) + elif tools_used: sanitized = ( f"[Previous request used tools ({', '.join(tools_used[:5])}) " f"but encountered an error. The user may ask to retry.]" diff --git a/src/discord/tool_loop.py b/src/discord/tool_loop.py index 763ad178..4f21b983 100644 --- a/src/discord/tool_loop.py +++ b/src/discord/tool_loop.py @@ -434,7 +434,25 @@ async def run( st = await self._prepare_chat_turn( message, history, system_prompt_override, trace, policy ) + return await self._run_with_guards(st) + + async def run_resumed(self, st: _ChatTurn) -> tuple[str, bool, bool, list[str], bool]: + """Continue a restored turn (built by TurnResumeManager) through the + same guard envelope as a fresh one. The iteration loop starts from + ``st.iteration`` — the restored transcript already contains every + earlier generation.""" + self._channel_state.set_active_request(st._ch_id, st._req_id) + set_turn( + turn_id=st._trajectory.message_id or None, + source=st._trajectory.source, + channel_id=st._trajectory.channel_id, + ) + await self._delivery.set_status("Resuming preserved work...", task_start=True) + return await self._run_with_guards(st) + async def _run_with_guards( + self, st: _ChatTurn + ) -> tuple[str, bool, bool, list[str], bool]: try: result = await self._run_chat_iterations(st) # Terminal bookkeeping (best-effort; a suspension already settled @@ -486,8 +504,12 @@ async def _run_chat_iterations( self, st: _ChatTurn ) -> tuple[str, bool, bool, list[str], bool]: """The chat iteration loop — every phase-method exit returns through - here; unexpected escapes are handled by run()'s guard above.""" - for iteration in range(st.chat_cap): + here; unexpected escapes are handled by run()'s guard above. + + Starts from ``st.iteration``: 0 for a fresh turn (unchanged), the + interrupted generation's index for a resumed one — the restored + transcript already carries everything before it.""" + for iteration in range(st.iteration, st.chat_cap): st.iteration = iteration if st._cancel.is_set(): return self._stopped(st, "iteration_start") @@ -803,10 +825,19 @@ async def _attempt(): tools_used=st.tools_used_in_loop, ) + # A resumed generation carries only its REMAINING budget (persisted + # UTC deadline): the generation that already spent its five minutes + # gets one attempt, not a fresh window. Later generations budget + # normally. + resume_budget = st.durability.pop_resume_budget() + deadline_seconds = ( + policy.deadline_seconds if resume_budget is None else resume_budget + ) + # Persist the absolute recovery deadline BEFORE the call: a restart # mid-recovery reconstructs only the remaining budget, never a fresh # five minutes. - await st.durability.on_generation_start(st, policy.deadline_seconds) + await st.durability.on_generation_start(st, deadline_seconds) # Typing is best-effort (shared helper): a typing failure — setup or # cleanup — must never fail the call or misclassify provider errors. @@ -816,6 +847,7 @@ async def _attempt(): _attempt, policy=policy, breaker=breaker, + deadline_seconds=deadline_seconds, cancel_event=st._cancel, on_wait=_on_wait, ) diff --git a/src/discord/turn_resume.py b/src/discord/turn_resume.py new file mode 100644 index 00000000..5f277e6e --- /dev/null +++ b/src/discord/turn_resume.py @@ -0,0 +1,392 @@ +"""Resume machinery for suspended chat turns. + +Two entry points (design settled with Odin, 2026-07-30): + +- **Explicit resume** — the user replies ``resume``/``continue`` in the + channel: ``try_explicit_resume`` runs inside the normal intake pipeline + (channel lock, delivery, session append all inherited), consuming the + trigger message as a command — it is NEVER injected into the frozen + transcript. Allowed even after the session has advanced. +- **Auto-resume** — registered at suspension time, in-process only: a + per-turn waiter polls the model breaker; when capacity returns AND the + session has not advanced since suspension, the turn resumes and replies + against the ORIGINAL message. A process restart drops waiters by design — + after a restart, resume is explicit-only. + +Admission (both paths, per the settled design): re-fetch the original +message; require the same author and unchanged content (digest); re-derive +tools from the CURRENT catalog + permission filter — current security +policy always wins over persisted definitions; a deleted or materially +edited request is terminal (``TERMINAL_REJECTED``), never executable +folklore reconstructed from disk. + +Replay safety: unmatched ``tool_use`` blocks (crash between intent +recording and batch settle) are repaired from the ledger — APPLIED ops +replay their stored result; anything else becomes an explicit +"outcome unknown / never ran" result block. Matched blocks are guaranteed; +NOTHING is ever re-executed automatically. +""" + +from __future__ import annotations + +import asyncio +import time +from collections.abc import Callable +from typing import Any + +from ..odin_log import get_logger +from ..turn_state.codec import compute_content_digest, restore_field_values +from ..turn_state.durability import TurnDurability +from ..turn_state.store import OpState, TurnKey, TurnStateStore +from .tool_loop import CHAT_POLICY, ToolLoopRunner, _ChatTurn + +log = get_logger("turn_resume") + +RESUME_TRIGGERS = frozenset({"resume", "continue"}) + +_AUTO_POLL_SECONDS = 15.0 +# Response text cap for the session append on the auto path (mirrors the +# intake pipeline's CHAT_RESPONSE_MAX_CHARS discipline without importing it). +_SESSION_RESPONSE_CAP = 4000 + + +class TurnResumeManager: + def __init__( + self, + *, + store: TurnStateStore, + tool_loop: ToolLoopRunner, + llm_gateway, + channel_state, + sessions, + delivery, + permissions, + tool_catalog, + get_config: Callable, + fetch_message: Callable, + auto_resume_enabled: bool = True, + resume_ttl_hours: float = 24.0, + ) -> None: + self._store = store + self._tool_loop = tool_loop + self._llm_gateway = llm_gateway + self._channel_state = channel_state + self._sessions = sessions + self._delivery = delivery + self._permissions = permissions + self._tool_catalog = tool_catalog + self._get_config = get_config + self._fetch_message = fetch_message # async (channel_id, message_id) -> msg|None + self._auto_resume_enabled = auto_resume_enabled + self._resume_ttl_hours = resume_ttl_hours + self._waiters: dict[TurnKey, asyncio.Task] = {} + + # ── queries ────────────────────────────────────────────────────── + + async def is_suspended(self, channel_id: str, message_id: str) -> bool: + key = TurnKey(source="discord", channel_id=channel_id, message_id=message_id) + row = await asyncio.to_thread(self._store.load_resumable_sync, key) + return row is not None + + async def _latest_suspended_for_channel(self, channel_id: str) -> dict | None: + rows = await asyncio.to_thread(self._store.list_suspended_sync, "discord") + candidates = [r for r in rows if r["channel_id"] == channel_id] + if not candidates: + return None + return max(candidates, key=lambda r: r.get("suspended_at") or 0.0) + + # ── suspension registration (auto-resume) ──────────────────────── + + def on_turn_suspended(self, key: TurnKey, generation: str) -> None: + """Called by the tool loop when a turn suspends. In-process only.""" + if not self._auto_resume_enabled: + return + existing = self._waiters.pop(key, None) + if existing is not None: + existing.cancel() + task = asyncio.get_running_loop().create_task( + self._auto_resume_waiter(key, generation), + name=f"turn-resume:{key.channel_id}:{key.message_id}", + ) + self._waiters[key] = task + task.add_done_callback(lambda _t: self._waiters.pop(key, None)) + + async def _auto_resume_waiter(self, key: TurnKey, generation: str) -> None: + """Wait for capacity, then resume IF nothing else happened. + + The session baseline is sampled at the FIRST poll (after the + suspending turn's intake bookkeeping settled — the channel lock + serializes us behind it); any later growth means the session + advanced and auto-resume stands down (explicit-only from there). + + Capacity detection is ACTIVE: a quiet breaker is never probed by + anyone, so the waiter claims the probe slot itself when the cooldown + elapses and immediately releases it — the resumed generation's own + attempt is the real probe. If capacity is still gone, that attempt + re-suspends the turn (remaining budget ≈ 0 → single attempt), which + re-registers this waiter — the breaker's escalating cooldown paces + the retry cycle for free. + """ + give_up_at = time.monotonic() + self._resume_ttl_hours * 3600.0 + breaker = self._llm_gateway.capacity_breaker_for() + await asyncio.sleep(_AUTO_POLL_SECONDS) + async with self._channel_lock(key.channel_id): + baseline = self._session_len(key.channel_id) + while time.monotonic() < give_up_at: + row = await asyncio.to_thread(self._store.load_resumable_sync, key) + if row is None or row["generation"] != generation: + return # resumed elsewhere / rejected / expired + admission = breaker.acquire_attempt() + if not isinstance(admission, float): + breaker.abandon(admission) # the resume re-acquires for real + current = self._session_len(key.channel_id) + if baseline < 0 or current < 0 or current != baseline: + # Advanced OR unreadable — either way auto-resume must + # stand down (fail-safe); explicit resume stays available. + log.info( + "Auto-resume for %s stands down: session advanced or " + "unreadable (explicit resume still available)", key, + ) + return + await self._run_auto_resume(key, row) + return + await asyncio.sleep(min(_AUTO_POLL_SECONDS, float(admission))) + log.info("Auto-resume waiter for %s expired", key) + + def _channel_lock(self, channel_id: str) -> asyncio.Lock: + return self._channel_state.channel_locks.setdefault(channel_id, asyncio.Lock()) + + def _session_len(self, channel_id: str) -> int: + """Peek the channel's session length WITHOUT get_or_create (which + bumps last_active on pure reads). SessionManager keeps its dict at + `_sessions`; a lookup failure returns -1 so a broken peek can never + satisfy the baseline-equality check and wrongly auto-resume.""" + try: + session = getattr(self._sessions, "_sessions", {}).get(channel_id) + return len(session.messages) if session is not None else 0 + except Exception: + return -1 + + async def _run_auto_resume(self, key: TurnKey, row: dict) -> None: + async with self._channel_lock(key.channel_id): + if self._channel_state.active_requests.get(key.channel_id): + log.info("Auto-resume for %s stands down: channel busy", key) + return + st, message, reason = await self._validate_and_rebuild(key, row) + if st is None: + log.info("Auto-resume for %s rejected: %s", key, reason) + return + log.info("Auto-resuming turn %s (capacity returned)", key) + try: + result = await self._tool_loop.run_resumed(st) + except Exception: + log.exception("Auto-resumed turn failed") + return + text, already_sent, is_error, tools_used, _handoff = result + self._append_session(key.channel_id, text, is_error, tools_used) + if not already_sent and message is not None: + try: + await self._delivery.send_chunked(message, text) + except Exception: + log.exception("Auto-resume delivery failed") + + def _append_session( + self, channel_id: str, text: str, is_error: bool, tools_used: list + ) -> None: + """Minimal mirror of the intake post-turn session bookkeeping (the + auto path runs outside the intake pipeline; reflection and + housekeeping deliberately do not run here).""" + try: + body = (text or "")[:_SESSION_RESPONSE_CAP] + if is_error: + body = ( + "[Resumed request ended with an error" + + (f" after tools ({', '.join(tools_used[:5])})" if tools_used else "") + + ".]" + ) + self._sessions.add_message(channel_id, "assistant", body) + self._sessions.prune() + except Exception: + log.exception("Auto-resume session append failed") + + # ── explicit resume (runs inside the intake pipeline) ──────────── + + @staticmethod + def is_resume_trigger(content: str) -> bool: + return (content or "").strip().lower().rstrip("!.") in RESUME_TRIGGERS + + async def try_explicit_resume(self, message: Any): + """Resume the channel's suspended turn when *message* is a trigger. + + Returns the run() result tuple, a notice tuple when resume was + attempted but rejected, or None when this message is not a resume + trigger (normal processing continues). + """ + content = getattr(message, "content", "") or "" + if not self.is_resume_trigger(content): + return None + channel_id = str(message.channel.id) + row_summary = await self._latest_suspended_for_channel(channel_id) + if row_summary is None: + return None # nothing to resume — treat as a normal message + key = TurnKey( + source="discord", + channel_id=channel_id, + message_id=row_summary["message_id"], + ) + row = await asyncio.to_thread(self._store.load_resumable_sync, key) + if row is None: + return None + # Only the original requester may resume their turn. + if str(message.author.id) != str(row.get("user_id") or ""): + return ( + "There is preserved work in this channel, but only the person " + "who started it can resume it.", + False, False, [], False, + ) + # Stand the auto-waiter down — the human took over. + waiter = self._waiters.pop(key, None) + if waiter is not None: + waiter.cancel() + st, _original, reason = await self._validate_and_rebuild(key, row) + if st is None: + return ( + f"I couldn't resume the preserved work: {reason}. " + "Ask again from scratch if you still need it.", + False, False, [], False, + ) + log.info("Explicitly resuming turn %s", key) + return await self._tool_loop.run_resumed(st) + + # ── admission + rebuild ────────────────────────────────────────── + + async def _validate_and_rebuild(self, key: TurnKey, row: dict): + """Full resume admission. Returns (st, original_message, None) or + (None, None, reason). Hard rejections mark the row terminal.""" + original = None + try: + original = await self._fetch_message(key.channel_id, key.message_id) + except Exception: + original = None + if original is None: + await asyncio.to_thread( + self._store.reject_resumable_sync, key, "original message unavailable" + ) + return None, None, "the original message is gone" + if str(original.author.id) != str(row.get("user_id") or ""): + await asyncio.to_thread( + self._store.reject_resumable_sync, key, "author mismatch" + ) + return None, None, "the original author no longer matches" + digest = compute_content_digest(getattr(original, "content", "") or "") + if digest != (row.get("content_digest") or ""): + await asyncio.to_thread( + self._store.reject_resumable_sync, key, "content edited" + ) + return None, None, "the original message was edited" + + lease = await asyncio.to_thread( + self._store.acquire_resume_lease_sync, key, row["generation"] + ) + if lease is None: + return None, None, "someone else is already resuming it" + + payload = row["payload"] + try: + fields = restore_field_values( + payload, + load_blob=self._store.load_blob_sync, + stuck_tracker_cls=self._tool_loop._stuck_loop_tracker_cls, + ) + except Exception: + log.exception("Checkpoint restore failed — rejecting") + await asyncio.to_thread( + self._store.reject_resumable_sync, key, "checkpoint unreadable" + ) + return None, None, "the checkpoint could not be restored" + + # Current security policy wins: tools re-derived from the live + # catalog + permission filter, never the persisted definitions. + tools = None + if self._get_config().tools.enabled: + tools = self._tool_catalog.merged_definitions() + tools = self._permissions.filter_tools(str(original.author.id), tools) + + self._repair_unmatched_tool_use(fields["messages"], row.get("operations") or []) + + cancel = self._channel_state.cancel_events.setdefault( + key.channel_id, asyncio.Event() + ) + if cancel.is_set(): + return None, None, "the channel is busy stopping another task" + + remaining_budget = 0.0 + deadline_utc = row.get("recovery_deadline_utc") + if deadline_utc: + remaining_budget = max(0.0, float(deadline_utc) - time.time()) + + durability = TurnDurability.resumed( + self._store, + lease, + payload.get("generation_seq", 0), + first_generation_budget=remaining_budget, + ) + st = _ChatTurn( + message=original, + policy=CHAT_POLICY, + trace=None, # the old segment was closed into the payload + tools=tools, + _cancel=cancel, + durability=durability, + **fields, + ) + return st, original, None + + @staticmethod + def _repair_unmatched_tool_use(messages: list, operations: list[dict]) -> None: + """Guarantee matched tool_use/tool_result blocks after a crash. + + Missing results are synthesized from the ledger: APPLIED replays the + stored result; anything else states the truth (unknown / never ran). + Nothing is re-executed. + """ + seen_results: set[str] = set() + use_blocks: dict[str, str] = {} + for msg in messages: + content = msg.get("content") + if not isinstance(content, list): + continue + for block in content: + if not isinstance(block, dict): + continue + if block.get("type") == "tool_use" and block.get("id"): + use_blocks[block["id"]] = block.get("name", "tool") + elif block.get("type") == "tool_result" and block.get("tool_use_id"): + seen_results.add(block["tool_use_id"]) + missing = [cid for cid in use_blocks if cid not in seen_results] + if not missing: + return + ops_by_id = {op["tool_call_id"]: op for op in operations} + repaired = [] + for cid in missing: + op = ops_by_id.get(cid) + if op is not None and op["state"] in ( + OpState.APPLIED, + OpState.RECONCILED_APPLIED, + ): + content = op.get("result") or "[completed; result recorded]" + elif op is None: + content = ( + "[Interrupted before execution — this call never ran; " + "re-issue it if still needed.]" + ) + else: + content = ( + "[Interrupted: outcome unknown — verify current state " + "before re-running this operation.]" + ) + repaired.append( + {"type": "tool_result", "tool_use_id": cid, "content": content} + ) + messages.append({"role": "user", "content": repaired}) + log.info("Repaired %d unmatched tool_use block(s) on resume", len(repaired)) diff --git a/src/turn_state/durability.py b/src/turn_state/durability.py index 2cc12d11..e06deb45 100644 --- a/src/turn_state/durability.py +++ b/src/turn_state/durability.py @@ -68,6 +68,8 @@ def __init__(self, store: TurnStateStore | None, lease: TurnLease | None) -> Non # cancel EVENT before terminal settlement can read it, so the fact # must be carried here (a cancelled turn is terminal by design). self.cancelled = False + # One-shot remaining budget for a resumed generation (see resumed()). + self._resume_budget: float | None = None # -- construction -------------------------------------------------- @@ -117,12 +119,32 @@ async def admit( @classmethod def resumed( - cls, store: TurnStateStore, lease: TurnLease, generation_seq: int + cls, + store: TurnStateStore, + lease: TurnLease, + generation_seq: int, + *, + first_generation_budget: float | None = None, ) -> TurnDurability: + """Handle for a resumed turn. + + ``first_generation_budget`` is the REMAINING recovery budget of the + interrupted generation (from the persisted UTC deadline) — usually + ~0, which means one attempt then re-suspend. A restart never grants + a fresh five minutes to the generation that already spent its + budget; later generations budget normally. + """ handle = cls(store, lease) handle.generation_seq = int(generation_seq) + handle._resume_budget = first_generation_budget return handle + def pop_resume_budget(self) -> float | None: + """One-shot: the restored generation's remaining budget, then None.""" + budget = self._resume_budget + self._resume_budget = None + return budget + # -- state --------------------------------------------------------- @property diff --git a/tests/test_resume_admission.py b/tests/test_resume_admission.py new file mode 100644 index 00000000..d8aee76a --- /dev/null +++ b/tests/test_resume_admission.py @@ -0,0 +1,322 @@ +"""Resume admission + execution pins (src/discord/turn_resume.py). + +Drives real suspend→resume cycles through the actual runner and store: +explicit resume completes the preserved work with full transcript +continuity; every admission rejection (deleted / edited / wrong author) is +terminal; the unmatched-block repair synthesizes truthful results from the +ledger and never re-executes anything. +""" + +from __future__ import annotations + +import asyncio +import json + +import pytest + +import src.discord.turn_resume as tr +from src.discord.turn_resume import TurnResumeManager +from src.llm.errors import LLMCapacityError +from src.llm.recovery import RecoveryPolicy +from src.turn_state import OpState, TurnStateStore, TurnStatus +from tests.fakes import FakeLLM, FakeMessage, make_bot, text_response, tool_call_response + +FAST_POLICY = RecoveryPolicy( + deadline_seconds=0.15, backoff_base=0.01, backoff_cap=0.02, retry_after_cap=0.05 +) + + +@pytest.fixture(autouse=True) +def _isolated_cwd(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + + +def capacity_forever(fake): + def _raise(): + fake.responses.append(_raise) + raise LLMCapacityError( + "Codex capacity: server_is_overloaded", provider="codex", model="fake-model" + ) + + return _raise + + +class Harness: + """A bot + store + resume manager sharing one fetchable message registry.""" + + def __init__(self, script, tmp_path): + self.fake = FakeLLM(script) + self.bot = make_bot(fake_llm=self.fake) + self.store = TurnStateStore(tmp_path / "ts" / "turns.sqlite3") + self.bot.tool_loop._turn_store = self.store + self.bot.llm_gateway._recovery_policy_source = lambda: FAST_POLICY + self.messages: dict[tuple[str, str], object] = {} + + async def fetch(channel_id: str, message_id: str): + return self.messages.get((channel_id, message_id)) + + self.manager = TurnResumeManager( + store=self.store, + tool_loop=self.bot.tool_loop, + llm_gateway=self.bot.llm_gateway, + channel_state=self.bot.channel_state, + sessions=self.bot.sessions, + delivery=self.bot.delivery, + permissions=self.bot.permissions, + tool_catalog=self.bot.tool_catalog, + get_config=lambda: self.bot.config, + fetch_message=fetch, + ) + self.bot.tool_loop._on_turn_suspended = self.manager.on_turn_suspended + + def register(self, msg): + self.messages[(str(msg.channel.id), str(msg.id))] = msg + + async def run(self, msg): + self.register(msg) + return await self.bot.tool_loop.run( + msg, [{"role": "user", "content": msg.content}] + ) + + def row(self, cols="status, payload"): + return self.store._conn.execute(f"SELECT {cols} FROM turns").fetchone() + + +async def suspend_turn(tmp_path, script=None): + h = Harness( + script + if script is not None + else [tool_call_response(("parse_time", {"text": "tomorrow"}))], + tmp_path, + ) + h.fake.responses.append(capacity_forever(h.fake)) + original = FakeMessage("do the long thing") + text, _, is_error, *_ = await h.run(original) + assert is_error is True + assert h.row()[0] == TurnStatus.SUSPENDED + # Cancel any auto-waiter the suspension registered — these tests drive + # the explicit path deterministically. + for task in list(h.manager._waiters.values()): + task.cancel() + await asyncio.sleep(0) + return h, original + + +def resume_msg(original, content="resume", author=None): + return FakeMessage(content, author=author or original.author, channel=original.channel) + + +def make_breaker_probe_ready(h): + """Model the production timeline where the breaker cooldown has elapsed + by the time a resume happens (suspension→resume is minutes, cooldown is + seconds-to-minutes): admit + succeed one probe so the breaker closes.""" + breaker = h.bot.llm_gateway.capacity_breaker_for() + token = breaker.acquire_attempt() + if not isinstance(token, float): + breaker.attempt_succeeded(token) + else: # still pacing — force the window open for the test + breaker._opened_at = 0.0 + token = breaker.acquire_attempt() + if not isinstance(token, float): + breaker.attempt_succeeded(token) + + +def heal_capacity(h, *responses): + """Capacity is back: replace the self-rearming raiser and close the breaker.""" + h.fake.responses.clear() + h.fake.responses.extend(responses) + make_breaker_probe_ready(h) + + +class TestExplicitResume: + async def test_resume_completes_preserved_work_with_continuity(self, tmp_path): + h, original = await suspend_turn(tmp_path) + heal_capacity(h, text_response("Finished what I started.")) + + result = await h.manager.try_explicit_resume(resume_msg(original)) + assert result is not None + text, _, is_error, tools_used, _ = result + assert text == "Finished what I started." + assert is_error is False + assert tools_used == ["parse_time"] # restored, not re-run + + # Transcript continuity: the resumed LLM call saw the earlier + # tool_use + matched tool_result from before the outage. + resumed_call = h.fake.calls[-1]["messages"] + blocks = [ + b + for m in resumed_call + if isinstance(m.get("content"), list) + for b in m["content"] + if isinstance(b, dict) + ] + assert any(b.get("type") == "tool_use" for b in blocks) + assert any(b.get("type") == "tool_result" for b in blocks) + # The tool was NOT re-executed on resume (ledger untouched, still 1 op). + ops = h.store._conn.execute("SELECT COUNT(*) FROM operations").fetchone() + assert ops[0] == 1 + assert h.row()[0] == TurnStatus.TERMINAL_COMPLETED + + async def test_consumed_guard_budget_survives_resume(self, tmp_path): + # Suspended AFTER the hedging guard consumed its one-shot budget: + # the resumed turn must NOT get a fresh one — hedging again ends the + # turn (guard-terminal), it is not retried a second time. + h, original = await suspend_turn( + tmp_path, script=[text_response("Shall I proceed with the deployment now?")] + ) + payload = json.loads(h.row()[1]) + assert payload["fields"]["hedging_retried"] is True + + heal_capacity(h, text_response("Shall I proceed with the deployment now?")) + calls_before = len(h.fake.calls) + result = await h.manager.try_explicit_resume(resume_msg(original)) + assert result is not None + # One LLM call on resume: the guard flag was restored as consumed, so + # no second hedging retry generation was granted (hedging again is + # guard-terminal, not another free retry). + assert len(h.fake.calls) - calls_before == 1 + + async def test_wrong_author_gets_notice(self, tmp_path): + from tests.fakes.discord_objects import FakeAuthor + + h, original = await suspend_turn(tmp_path) + intruder = FakeAuthor(id=999999, name="intruder") + result = await h.manager.try_explicit_resume( + resume_msg(original, author=intruder) + ) + assert result is not None + assert "only the person" in result[0] + assert h.row()[0] == TurnStatus.SUSPENDED # untouched + + async def test_edited_original_is_terminal_rejected(self, tmp_path): + h, original = await suspend_turn(tmp_path) + original.content = "do the long thing (edited)" + result = await h.manager.try_explicit_resume(resume_msg(original)) + assert result is not None + assert "edited" in result[0] + assert h.row()[0] == TurnStatus.TERMINAL_REJECTED + assert h.row()[1] is None # payload compacted + + async def test_deleted_original_is_terminal_rejected(self, tmp_path): + h, original = await suspend_turn(tmp_path) + h.messages.clear() # fetch returns None + result = await h.manager.try_explicit_resume(resume_msg(original)) + assert result is not None + assert "gone" in result[0] + assert h.row()[0] == TurnStatus.TERMINAL_REJECTED + + async def test_non_trigger_and_no_checkpoint_pass_through(self, tmp_path): + h, original = await suspend_turn(tmp_path) + assert await h.manager.try_explicit_resume( + resume_msg(original, content="what's the weather") + ) is None + # A trigger in a channel WITHOUT preserved work is a normal message. + from tests.fakes.discord_objects import FakeChannel + + other_channel_msg = FakeMessage("resume", channel=FakeChannel(id=999888777)) + assert await h.manager.try_explicit_resume(other_channel_msg) is None + + async def test_second_resume_finds_nothing(self, tmp_path): + h, original = await suspend_turn(tmp_path) + heal_capacity(h, text_response("done")) + assert await h.manager.try_explicit_resume(resume_msg(original)) is not None + # Terminal now — a second `resume` is just a normal message. + assert await h.manager.try_explicit_resume(resume_msg(original)) is None + + async def test_resumed_generation_budget_is_remaining_not_fresh(self, tmp_path): + # Capacity STILL down at resume: the interrupted generation gets its + # REMAINING budget (~0 → one attempt), then re-suspends. No fresh + # five minutes for a generation that already spent its budget. + h, original = await suspend_turn(tmp_path) + make_breaker_probe_ready(h) + attempts_before = len(h.fake.calls) + result = await h.manager.try_explicit_resume(resume_msg(original)) + assert result is not None + text = result[0] + assert "preserved" in text # re-suspended, work still safe + assert h.row()[0] == TurnStatus.SUSPENDED + # Exactly ONE attempt was made (zero-budget semantics). + assert len(h.fake.calls) == attempts_before + 1 + for task in list(h.manager._waiters.values()): + task.cancel() + + +class TestAutoResume: + async def test_auto_resume_fires_when_capacity_returns(self, tmp_path, monkeypatch): + monkeypatch.setattr(tr, "_AUTO_POLL_SECONDS", 0.02) + h, original = await suspend_turn(tmp_path) + # Heal capacity + re-register the waiter (suspend_turn cancelled it). + heal_capacity(h, text_response("Auto-finished.")) + rows = h.store.list_suspended_sync("discord") + from src.turn_state import TurnKey + + key = TurnKey("discord", rows[0]["channel_id"], rows[0]["message_id"]) + h.manager.on_turn_suspended(key, rows[0]["generation"]) + for _ in range(200): + await asyncio.sleep(0.02) + if h.row()[0] != TurnStatus.SUSPENDED: + break + assert h.row()[0] == TurnStatus.TERMINAL_COMPLETED + # The reply landed against the ORIGINAL message. + assert any("Auto-finished." in (r["content"] or "") for r in original.replies) + + async def test_auto_resume_stands_down_when_session_advances( + self, tmp_path, monkeypatch + ): + monkeypatch.setattr(tr, "_AUTO_POLL_SECONDS", 0.02) + h, original = await suspend_turn(tmp_path) + # Capacity text is scripted but the breaker stays OPEN (pacing) so + # the waiter loops without resuming yet. + h.fake.responses.clear() + h.fake.responses.append(text_response("should never send")) + rows = h.store.list_suspended_sync("discord") + from src.turn_state import TurnKey + + key = TurnKey("discord", rows[0]["channel_id"], rows[0]["message_id"]) + h.manager.on_turn_suspended(key, rows[0]["generation"]) + await asyncio.sleep(0.06) # baseline sampled while the breaker paces + h.bot.sessions.add_message(str(original.channel.id), "user", "new topic") + await asyncio.sleep(0.02) + make_breaker_probe_ready(h) # capacity "returns" AFTER the advance + await asyncio.sleep(0.3) + assert h.row()[0] == TurnStatus.SUSPENDED # stood down, still resumable + for task in list(h.manager._waiters.values()): + task.cancel() + + +class TestUnmatchedBlockRepair: + def test_repair_synthesizes_truthful_results(self): + messages = [ + {"role": "user", "content": "go"}, + {"role": "assistant", "content": [ + {"type": "tool_use", "id": "a", "name": "run_command", "input": {}}, + {"type": "tool_use", "id": "b", "name": "run_command", "input": {}}, + {"type": "tool_use", "id": "c", "name": "run_command", "input": {}}, + ]}, + ] + operations = [ + {"tool_call_id": "a", "state": OpState.APPLIED, "result": "real output", + "tool_name": "run_command", "generation_seq": 1}, + {"tool_call_id": "b", "state": OpState.OUTCOME_UNKNOWN, "result": None, + "tool_name": "run_command", "generation_seq": 1}, + # "c" has no ledger row: intent was never recorded → never ran. + ] + TurnResumeManager._repair_unmatched_tool_use(messages, operations) + assert messages[-1]["role"] == "user" + by_id = {b["tool_use_id"]: b["content"] for b in messages[-1]["content"]} + assert by_id["a"] == "real output" + assert "outcome unknown" in by_id["b"] + assert "never ran" in by_id["c"] + + def test_matched_transcript_is_untouched(self): + messages = [ + {"role": "assistant", "content": [ + {"type": "tool_use", "id": "a", "name": "t", "input": {}}, + ]}, + {"role": "user", "content": [ + {"type": "tool_result", "tool_use_id": "a", "content": "ok"}, + ]}, + ] + before = json.loads(json.dumps(messages)) + TurnResumeManager._repair_unmatched_tool_use(messages, []) + assert messages == before From f9332511bdcd8fd1c1d5c1169b205a150fe26a0e Mon Sep 17 00:00:00 2001 From: ATCharpentier Date: Thu, 30 Jul 2026 14:20:19 -0400 Subject: [PATCH 08/18] feat: config sections + wiring + retention sweep + docs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two schema-only top-level sections (deliberately NOT added to the tracked config.yml template — six sections already live schema-only, and this keeps the deploy free of the skip-worktree dance): llm_recovery (generation deadline, backoff cap, model-breaker knobs) and turn_state (enabled, db path, auto_resume, the three retention clocks). Wiring: ModelBreakerRegistry + recovery-policy source + TurnStateStore are BotServices-owned (client rebuilds and live reloads never reset breaker or store state); gateway receives both; ToolLoopDeps carries the store; TurnResumeManager constructed when the store is up, its suspension callback attached to the runner, and the manager handed to the message pipeline for explicit resume. Housekeeping cleanup_stale gains the three-clock TTL sweep (guarded, house per-step style); shutdown_services closes the store. --- docs/configuration.md | 42 +++++++++++++++++++ src/config/schema.py | 30 ++++++++++++++ src/discord/housekeeping.py | 18 +++++++++ src/discord/wiring.py | 81 +++++++++++++++++++++++++++++++++++++ 4 files changed, 171 insertions(+) diff --git a/docs/configuration.md b/docs/configuration.md index 6a9ad004..31094e9d 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -275,3 +275,45 @@ shutdown completes. Recovery therefore does not depend on the service unit's `Restart=` policy, Docker restart policy, or any supervisor at all. `Restart=always` (what the packaged unit ships) is still recommended so the service also recovers from crashes and reboots. + +## LLM Recovery (capacity outages) + +Model-capacity errors (e.g. `server_is_overloaded`, which arrives inside an +HTTP 200 as an SSE error event) are retried with a deadline-based policy +shared by chat, agents, and autonomous loops, coordinated by a per-model +circuit breaker. Quota handling is unchanged: HTTP 429 still rotates +accounts inside the provider client; capacity never does. + +```yaml +llm_recovery: + generation_deadline_seconds: 300 # retry budget per LLM generation (waiting, not the attempt) + backoff_cap_seconds: 45 # full-jitter backoff ceiling between attempts + breaker_generation_threshold: 1 # failed generations before the model breaker opens + breaker_cooldown_base_seconds: 30 # first cooldown; doubles per failed probe + breaker_cooldown_cap_seconds: 300 # cooldown ceiling +``` + +All keys are optional (schema defaults shown); the section does not need to +exist in `config.yml`. + +## Turn State (checkpoints and resume) + +Discord chat turns are checkpointed to a durable store so a capacity outage +suspends the turn with its work preserved instead of discarding it. A +suspended turn auto-resumes when capacity returns (if nothing else has +happened in the channel), or the original requester can reply `resume` +within the resumable window. Interrupted tool executions are recorded as +outcome-unknown and are never re-run automatically. + +```yaml +turn_state: + enabled: true + db_path: "./data/turn_state/turns.sqlite3" + auto_resume: true + resume_ttl_hours: 24 # resumable window from last real progress + payload_retention_days: 7 # diagnostic payloads, then compacted to tombstones + ledger_retention_days: 90 # side-effect ledger (outcome-unknown rows never expire) +``` + +All keys are optional; disabling `turn_state.enabled` restores the previous +behavior (capacity exhaustion ends the turn with an error). diff --git a/src/config/schema.py b/src/config/schema.py index 861f20a1..d6ebb557 100644 --- a/src/config/schema.py +++ b/src/config/schema.py @@ -581,6 +581,34 @@ class GracefulDegradationConfig(BaseModel): unavailable_threshold: int = 10 # consecutive failures before UNAVAILABLE +class LLMRecoveryConfig(BaseModel): + """Deadline-based recovery for logical LLM generations (all three call + paths: chat, agents, autonomous loops) plus the model-scoped capacity + breaker. The deadline bounds WAITING between attempts, never the + attempt itself; capacity never rotates accounts (429 rotation is the + provider client's job and is untouched).""" + + generation_deadline_seconds: float = Field(default=300.0, ge=10.0, le=3600.0) + backoff_cap_seconds: float = Field(default=45.0, ge=1.0, le=300.0) + breaker_generation_threshold: int = Field(default=1, ge=1, le=10) + breaker_cooldown_base_seconds: float = Field(default=30.0, ge=1.0, le=600.0) + breaker_cooldown_cap_seconds: float = Field(default=300.0, ge=30.0, le=3600.0) + + +class TurnStateConfig(BaseModel): + """Durable chat-turn checkpoints, side-effect ledger, and resume. + + Discord chat turns only (v1). Disabled => turns run exactly as before + (capacity exhaustion discards work instead of suspending).""" + + enabled: bool = True + db_path: str = "./data/turn_state/turns.sqlite3" + auto_resume: bool = True + resume_ttl_hours: float = Field(default=24.0, ge=1.0, le=24.0 * 14) + payload_retention_days: float = Field(default=7.0, ge=1.0, le=90.0) + ledger_retention_days: float = Field(default=90.0, ge=30.0, le=365.0) + + class AuditConfig(BaseModel): hmac_key: str = "" # Empty = signing disabled @@ -886,6 +914,8 @@ class Config(BaseModel): grafana_alerts: GrafanaAlertConfig = GrafanaAlertConfig() outbound_webhooks: OutboundWebhooksConfig = OutboundWebhooksConfig() graceful_degradation: GracefulDegradationConfig = GracefulDegradationConfig() + llm_recovery: LLMRecoveryConfig = LLMRecoveryConfig() + turn_state: TurnStateConfig = TurnStateConfig() def _substitute_env_vars(text: str) -> str: diff --git a/src/discord/housekeeping.py b/src/discord/housekeeping.py index fcf03d11..6dcf9698 100644 --- a/src/discord/housekeeping.py +++ b/src/discord/housekeeping.py @@ -32,6 +32,7 @@ def __init__( loop_agent_bridge, channel_logger, fts_index, + turn_store=None, ) -> None: self._get_config = get_config self._sessions = sessions @@ -42,6 +43,7 @@ def __init__( self._loop_agent_bridge = loop_agent_bridge self._channel_logger = channel_logger self._fts_index = fts_index + self._turn_store = turn_store def cleanup_stale(self) -> None: """Remove stale entries from per-channel caches to prevent memory leaks. @@ -94,6 +96,22 @@ def cleanup_stale(self) -> None: except Exception: pass + # Turn-state retention: the three clocks (resumable TTL, diagnostic + # payload compaction, ledger expiry — OUTCOME_UNKNOWN never expires). + if self._turn_store is not None: + try: + config = self._get_config() + ts = getattr(config, "turn_state", None) + swept = self._turn_store.ttl_sweep_sync( + resume_ttl_hours=getattr(ts, "resume_ttl_hours", 24.0), + payload_retention_days=getattr(ts, "payload_retention_days", 7.0), + ledger_retention_days=getattr(ts, "ledger_retention_days", 90.0), + ) + if any(swept.values()): + log.info("Turn-state TTL sweep: %s", swept) + except Exception: + pass + def maybe_cleanup(self) -> None: """Run cache cleanup if enough time has passed since the last run.""" try: diff --git a/src/discord/wiring.py b/src/discord/wiring.py index 9352af49..57ab5d93 100644 --- a/src/discord/wiring.py +++ b/src/discord/wiring.py @@ -119,6 +119,12 @@ class BotServices: stuck_loop_tracker_cls: type classify_command_risk: Callable classify_tool_risk: Callable + # Turn durability (2026-07-30): the checkpoint/ledger store (None = + # disabled) and the model-scoped capacity-breaker registry. Both live + # HERE so client rebuilds and live reloads never reset their state. + turn_store: object | None = None + model_breakers: object | None = None + recovery_policy_source: Callable | None = None def build_services(config: Config) -> BotServices: # noqa: PLR0915 — linear composition root @@ -395,6 +401,34 @@ def build_services(config: Config) -> BotServices: # noqa: PLR0915 — linear c ): subsystem_guard.register(_name) + # Turn durability (2026-07-30): checkpoint/ledger store + the model- + # scoped capacity breakers + the recovery policy source. Services-owned + # so provider-client rebuilds and live reloads never reset their state. + from ..llm.model_breaker import ModelBreakerRegistry + from ..llm.recovery import RecoveryPolicy + + _lr = config.llm_recovery + model_breakers = ModelBreakerRegistry( + generation_threshold=_lr.breaker_generation_threshold, + cooldown_base=_lr.breaker_cooldown_base_seconds, + cooldown_cap=_lr.breaker_cooldown_cap_seconds, + ) + + def recovery_policy_source() -> RecoveryPolicy: + live = config # config object is replaced wholesale on hot reload + return RecoveryPolicy( + deadline_seconds=live.llm_recovery.generation_deadline_seconds, + backoff_cap=live.llm_recovery.backoff_cap_seconds, + ) + + turn_store = None + if config.turn_state.enabled: + from ..turn_state import TurnStateStore + + turn_store = TurnStateStore(config.turn_state.db_path) + if not turn_store.available: + turn_store = None # init failed — feature off, logged loudly + # Action diff tracker — records before→after diffs. Always on. diff_tracker = DiffTracker() @@ -521,6 +555,9 @@ def build_services(config: Config) -> BotServices: # noqa: PLR0915 — linear c stuck_loop_tracker_cls=StuckLoopTracker, classify_command_risk=classify_command, classify_tool_risk=classify_tool, + turn_store=turn_store, + model_breakers=model_breakers, + recovery_policy_source=recovery_policy_source, ) @@ -570,6 +607,8 @@ def build_components(bot, services: BotServices) -> BotComponents: cost_tracker=services.cost_tracker, sessions=services.sessions, reflector=services.reflector, + model_breakers=services.model_breakers, + recovery_policy_source=services.recovery_policy_source, ) # Wire LLM callbacks to whichever provider is active @@ -696,6 +735,7 @@ def build_components(bot, services: BotServices) -> BotComponents: audit=services.audit, loop_manager=services.loop_manager, stuck_loop_tracker_cls=services.stuck_loop_tracker_cls, + turn_store=services.turn_store, ) ) agent_task_tools = AgentTaskTools( @@ -751,7 +791,40 @@ def build_components(bot, services: BotServices) -> BotComponents: loop_agent_bridge=services.loop_agent_bridge, channel_logger=services.channel_logger, fts_index=services.fts_index, + turn_store=services.turn_store, ) + + # Suspended-turn resume (2026-07-30): explicit `resume` rides the intake + # pipeline; auto-resume is registered by the tool loop at suspension. + turn_resume = None + if services.turn_store is not None: + from .turn_resume import TurnResumeManager + + async def _fetch_message(channel_id: str, message_id: str): + channel = bot.get_channel(int(channel_id)) + if channel is None: + channel = await bot.fetch_channel(int(channel_id)) + return await channel.fetch_message(int(message_id)) + + turn_resume = TurnResumeManager( + store=services.turn_store, + tool_loop=tool_loop, + llm_gateway=llm_gateway, + channel_state=services.channel_state, + sessions=services.sessions, + delivery=delivery, + permissions=services.permissions, + tool_catalog=tool_catalog, + get_config=lambda: bot.config, + fetch_message=_fetch_message, + auto_resume_enabled=bot.config.turn_state.auto_resume, + resume_ttl_hours=bot.config.turn_state.resume_ttl_hours, + ) + # Late instance-attr wiring: the runner exists before the manager + # (the manager needs the runner), so the suspension callback is + # attached here rather than through the frozen deps. + tool_loop._on_turn_suspended = turn_resume.on_turn_suspended + pipeline = MessagePipeline( MessagePipelineDeps( channel_state=services.channel_state, @@ -763,6 +836,7 @@ def build_components(bot, services: BotServices) -> BotComponents: tool_loop=tool_loop, delivery=delivery, housekeeping=housekeeping, + turn_resume=turn_resume, ) ) intake = MessageIntake( @@ -850,6 +924,13 @@ async def shutdown_services(bot) -> None: except Exception: log.exception("Error closing knowledge") + turn_store = getattr(getattr(bot, "services", None), "turn_store", None) + if turn_store is not None: + try: + turn_store.close() + except Exception: + log.exception("Error closing turn_store") + sessions = getattr(bot, "sessions", None) if sessions is not None: try: From 552e4b2b6ac24f791f9382b262047a3d9fa3b943 Mon Sep 17 00:00:00 2001 From: ATCharpentier Date: Thu, 30 Jul 2026 14:33:38 -0400 Subject: [PATCH 09/18] test: close coverage gaps + amend last ladder pin + ratchet baseline MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Branch coverage for the new paths: /stop before the first iteration and during a recovery wait (both TERMINAL_CANCELLED, graceful), task-cancel with a failing durable settle (cancellation still propagates promptly), escape-path settle failure (the ORIGINAL fail-closed error wins — the epilogue guard widened to BaseException so a cancellation delivered inside the settle can never replace it), failed suspension (plain error, no false preservation claims, TERMINAL_FAILED), duplicate tool_call ids (bounced matched, nothing executed), the housekeeping three-clock sweep (runs + swallows store failures), kimi Retry-After parsing on 429 exhaustion, and guard no-ops for unregistered names. test_agent_trajectory recovery pin amended deliberately (the manager ladder is gone; recovery_attempts stays 0 and the field remains for shape compatibility). Typed-object deps given real types for mypy (BotServices, MessagePipelineDeps, durability suspend closure). coverage-baseline.json regenerated (--update-baseline): all new files ratcheted, 26 files improved, 0 decreased, total 88.2%. --- coverage-baseline.json | 2514 +++++++++++---------- src/discord/intake_pipeline.py | 3 +- src/discord/tool_loop.py | 6 +- src/discord/wiring.py | 14 +- src/turn_state/durability.py | 7 +- tests/test_agent_trajectory.py | 13 +- tests/test_guard_capacity.py | 6 + tests/test_llm_kimi.py | 25 + tests/test_write_invariant_integration.py | 198 ++ 9 files changed, 1569 insertions(+), 1217 deletions(-) diff --git a/coverage-baseline.json b/coverage-baseline.json index dffb7425..8d43c6bf 100644 --- a/coverage-baseline.json +++ b/coverage-baseline.json @@ -1,1202 +1,1316 @@ { - "src/__init__.py": { - "covered": 3, - "missing": 0, - "percent": 100.0, - "statements": 3 - }, - "src/agents/__init__.py": { - "covered": 5, - "missing": 0, - "percent": 100.0, - "statements": 5 - }, - "src/agents/loop_bridge.py": { - "covered": 87, - "missing": 1, - "percent": 98.86, - "statements": 88 - }, - "src/agents/manager.py": { - "covered": 470, - "missing": 62, - "percent": 88.35, - "statements": 532 - }, - "src/agents/trajectory.py": { - "covered": 129, - "missing": 16, - "percent": 88.97, - "statements": 145 - }, - "src/async_utils.py": { - "covered": 13, - "missing": 2, - "percent": 86.67, - "statements": 15 - }, - "src/audit/__init__.py": { - "covered": 3, - "missing": 0, - "percent": 100.0, - "statements": 3 - }, - "src/audit/diff_tracker.py": { - "covered": 60, - "missing": 0, - "percent": 100.0, - "statements": 60 - }, - "src/audit/logger.py": { - "covered": 274, - "missing": 27, - "percent": 91.03, - "statements": 301 - }, - "src/audit/signer.py": { - "covered": 71, - "missing": 2, - "percent": 97.26, - "statements": 73 - }, - "src/config/__init__.py": { - "covered": 30, - "missing": 0, - "percent": 100.0, - "statements": 30 - }, - "src/config/schema.py": { - "covered": 542, - "missing": 1, - "percent": 99.82, - "statements": 543 - }, - "src/constants.py": { - "covered": 17, - "missing": 0, - "percent": 100.0, - "statements": 17 - }, - "src/context/__init__.py": { - "covered": 2, - "missing": 0, - "percent": 100.0, - "statements": 2 - }, - "src/context/loader.py": { - "covered": 42, - "missing": 11, - "percent": 79.25, - "statements": 53 - }, - "src/database/__init__.py": { - "covered": 0, - "missing": 0, - "percent": 100.0, - "statements": 0 - }, - "src/database/repository.py": { - "covered": 31, - "missing": 0, - "percent": 100.0, - "statements": 31 - }, - "src/discord/__init__.py": { - "covered": 3, - "missing": 4, - "percent": 42.86, - "statements": 7 - }, - "src/discord/attachments.py": { - "covered": 315, - "missing": 14, - "percent": 95.74, - "statements": 329 - }, - "src/discord/background_task.py": { - "covered": 219, - "missing": 130, - "percent": 62.75, - "statements": 349 - }, - "src/discord/channel_config.py": { - "covered": 81, - "missing": 0, - "percent": 100.0, - "statements": 81 - }, - "src/discord/channel_logger.py": { - "covered": 104, - "missing": 9, - "percent": 92.04, - "statements": 113 - }, - "src/discord/channel_state.py": { - "covered": 77, - "missing": 16, - "percent": 82.8, - "statements": 93 - }, - "src/discord/client.py": { - "covered": 129, - "missing": 43, - "percent": 75.0, - "statements": 172 - }, - "src/discord/completion.py": { - "covered": 40, - "missing": 1, - "percent": 97.56, - "statements": 41 - }, - "src/discord/delivery.py": { - "covered": 99, - "missing": 3, - "percent": 97.06, - "statements": 102 - }, - "src/discord/helpers/__init__.py": { - "covered": 0, - "missing": 0, - "percent": 100.0, - "statements": 0 - }, - "src/discord/helpers/converters.py": { - "covered": 23, - "missing": 0, - "percent": 100.0, - "statements": 23 - }, - "src/discord/helpers/cooldowns.py": { - "covered": 25, - "missing": 0, - "percent": 100.0, - "statements": 25 - }, - "src/discord/helpers/embeds.py": { - "covered": 28, - "missing": 0, - "percent": 100.0, - "statements": 28 - }, - "src/discord/helpers/pagination.py": { - "covered": 21, - "missing": 20, - "percent": 51.22, - "statements": 41 - }, - "src/discord/helpers/permissions.py": { - "covered": 12, - "missing": 19, - "percent": 38.71, - "statements": 31 - }, - "src/discord/housekeeping.py": { - "covered": 39, - "missing": 11, - "percent": 78.0, - "statements": 50 - }, - "src/discord/intake_pipeline.py": { - "covered": 298, - "missing": 76, - "percent": 79.68, - "statements": 374 - }, - "src/discord/llm_gateway.py": { - "covered": 173, - "missing": 1, - "percent": 99.43, - "statements": 174 - }, - "src/discord/native_tools/__init__.py": { - "covered": 2, - "missing": 0, - "percent": 100.0, - "statements": 2 - }, - "src/discord/native_tools/agents_tasks.py": { - "covered": 304, - "missing": 17, - "percent": 94.7, - "statements": 321 - }, - "src/discord/native_tools/channel_ops.py": { - "covered": 113, - "missing": 0, - "percent": 100.0, - "statements": 113 - }, - "src/discord/native_tools/knowledge.py": { - "covered": 117, - "missing": 0, - "percent": 100.0, - "statements": 117 - }, - "src/discord/native_tools/media.py": { - "covered": 166, - "missing": 0, - "percent": 100.0, - "statements": 166 - }, - "src/discord/native_tools/registry.py": { - "covered": 83, - "missing": 6, - "percent": 93.26, - "statements": 89 - }, - "src/discord/native_tools/scheduling.py": { - "covered": 124, - "missing": 0, - "percent": 100.0, - "statements": 124 - }, - "src/discord/native_tools/skills_tools.py": { - "covered": 87, - "missing": 8, - "percent": 91.58, - "statements": 95 - }, - "src/discord/prompts.py": { - "covered": 82, - "missing": 28, - "percent": 74.55, - "statements": 110 - }, - "src/discord/response_guards.py": { - "covered": 154, - "missing": 0, - "percent": 100.0, - "statements": 154 - }, - "src/discord/scheduled_events.py": { - "covered": 226, - "missing": 0, - "percent": 100.0, - "statements": 226 - }, - "src/discord/slash_commands.py": { - "covered": 20, - "missing": 49, - "percent": 28.99, - "statements": 69 - }, - "src/discord/tool_catalog.py": { - "covered": 27, - "missing": 0, - "percent": 100.0, - "statements": 27 - }, - "src/discord/tool_loop.py": { - "covered": 568, - "missing": 60, - "percent": 90.45, - "statements": 628 - }, - "src/discord/tool_loop_helpers.py": { - "covered": 30, - "missing": 9, - "percent": 76.92, - "statements": 39 - }, - "src/discord/turn_recorder.py": { - "covered": 92, - "missing": 16, - "percent": 85.19, - "statements": 108 - }, - "src/discord/wiring.py": { - "covered": 284, - "missing": 78, - "percent": 78.45, - "statements": 362 - }, - "src/health/__init__.py": { - "covered": 3, - "missing": 0, - "percent": 100.0, - "statements": 3 - }, - "src/health/checker.py": { - "covered": 239, - "missing": 0, - "percent": 100.0, - "statements": 239 - }, - "src/health/grafana_alerts.py": { - "covered": 235, - "missing": 0, - "percent": 100.0, - "statements": 235 - }, - "src/health/metrics.py": { - "covered": 180, - "missing": 18, - "percent": 90.91, - "statements": 198 - }, - "src/health/server.py": { - "covered": 453, - "missing": 186, - "percent": 70.89, - "statements": 639 - }, - "src/health/startup.py": { - "covered": 231, - "missing": 11, - "percent": 95.45, - "statements": 242 - }, - "src/health/subsystem_guard.py": { - "covered": 178, - "missing": 1, - "percent": 99.44, - "statements": 179 - }, - "src/knowledge/__init__.py": { - "covered": 2, - "missing": 0, - "percent": 100.0, - "statements": 2 - }, - "src/knowledge/importer.py": { - "covered": 170, - "missing": 12, - "percent": 93.41, - "statements": 182 - }, - "src/knowledge/store.py": { - "covered": 322, - "missing": 95, - "percent": 77.22, - "statements": 417 - }, - "src/learning/__init__.py": { - "covered": 2, - "missing": 0, - "percent": 100.0, - "statements": 2 - }, - "src/learning/loop_reflection.py": { - "covered": 49, - "missing": 5, - "percent": 90.74, - "statements": 54 - }, - "src/learning/reflector.py": { - "covered": 501, - "missing": 66, - "percent": 88.36, - "statements": 567 - }, - "src/llm/__init__.py": { - "covered": 11, - "missing": 0, - "percent": 100.0, - "statements": 11 - }, - "src/llm/auxiliary.py": { - "covered": 72, - "missing": 0, - "percent": 100.0, - "statements": 72 - }, - "src/llm/backoff.py": { - "covered": 10, - "missing": 0, - "percent": 100.0, - "statements": 10 - }, - "src/llm/circuit_breaker.py": { - "covered": 42, - "missing": 0, - "percent": 100.0, - "statements": 42 - }, - "src/llm/codex_auth.py": { - "covered": 375, - "missing": 27, - "percent": 93.28, - "statements": 402 - }, - "src/llm/context_compressor.py": { - "covered": 165, - "missing": 18, - "percent": 90.16, - "statements": 183 - }, - "src/llm/cost_tracker.py": { - "covered": 86, - "missing": 0, - "percent": 100.0, - "statements": 86 - }, - "src/llm/kimi.py": { - "covered": 221, - "missing": 1, - "percent": 99.55, - "statements": 222 - }, - "src/llm/ollama.py": { - "covered": 172, - "missing": 5, - "percent": 97.18, - "statements": 177 - }, - "src/llm/openai_codex.py": { - "covered": 379, - "missing": 65, - "percent": 85.36, - "statements": 444 - }, - "src/llm/provider.py": { - "covered": 15, - "missing": 3, - "percent": 83.33, - "statements": 18 - }, - "src/llm/secret_scrubber.py": { - "covered": 7, - "missing": 0, - "percent": 100.0, - "statements": 7 - }, - "src/llm/system_prompt.py": { - "covered": 34, - "missing": 3, - "percent": 91.89, - "statements": 37 - }, - "src/llm/types.py": { - "covered": 18, - "missing": 0, - "percent": 100.0, - "statements": 18 - }, - "src/models/__init__.py": { - "covered": 5, - "missing": 0, - "percent": 100.0, - "statements": 5 - }, - "src/models/guild.py": { - "covered": 18, - "missing": 0, - "percent": 100.0, - "statements": 18 - }, - "src/models/infraction.py": { - "covered": 20, - "missing": 0, - "percent": 100.0, - "statements": 20 - }, - "src/models/reminder.py": { - "covered": 20, - "missing": 0, - "percent": 100.0, - "statements": 20 - }, - "src/models/user.py": { - "covered": 14, - "missing": 0, - "percent": 100.0, - "statements": 14 - }, - "src/monitoring/__init__.py": { - "covered": 2, - "missing": 0, - "percent": 100.0, - "statements": 2 - }, - "src/monitoring/resource_usage.py": { - "covered": 154, - "missing": 11, - "percent": 93.33, - "statements": 165 - }, - "src/notifications/__init__.py": { - "covered": 4, - "missing": 0, - "percent": 100.0, - "statements": 4 - }, - "src/notifications/issue_tracker.py": { - "covered": 261, - "missing": 22, - "percent": 92.23, - "statements": 283 - }, - "src/notifications/outbound_webhooks.py": { - "covered": 258, - "missing": 3, - "percent": 98.85, - "statements": 261 - }, - "src/notifications/slack.py": { - "covered": 128, - "missing": 0, - "percent": 100.0, - "statements": 128 - }, - "src/observability/__init__.py": { - "covered": 3, - "missing": 0, - "percent": 100.0, - "statements": 3 - }, - "src/observability/aggregates.py": { - "covered": 122, - "missing": 5, - "percent": 96.06, - "statements": 127 - }, - "src/observability/context_trace.py": { - "covered": 118, - "missing": 12, - "percent": 90.77, - "statements": 130 - }, - "src/observability/correlation.py": { - "covered": 11, - "missing": 2, - "percent": 84.62, - "statements": 13 - }, - "src/observability/failure_classes.py": { - "covered": 14, - "missing": 2, - "percent": 87.5, - "statements": 16 - }, - "src/odin/__init__.py": { - "covered": 7, - "missing": 0, - "percent": 100.0, - "statements": 7 - }, - "src/odin/cli.py": { - "covered": 46, - "missing": 13, - "percent": 77.97, - "statements": 59 - }, - "src/odin/context.py": { - "covered": 122, - "missing": 7, - "percent": 94.57, - "statements": 129 - }, - "src/odin/executor.py": { - "covered": 42, - "missing": 1, - "percent": 97.67, - "statements": 43 - }, - "src/odin/plan_loader.py": { - "covered": 29, - "missing": 8, - "percent": 78.38, - "statements": 37 - }, - "src/odin/planner.py": { - "covered": 95, - "missing": 3, - "percent": 96.94, - "statements": 98 - }, - "src/odin/registry.py": { - "covered": 27, - "missing": 0, - "percent": 100.0, - "statements": 27 - }, - "src/odin/reporter.py": { - "covered": 25, - "missing": 0, - "percent": 100.0, - "statements": 25 - }, - "src/odin/tools/__init__.py": { - "covered": 0, - "missing": 0, - "percent": 100.0, - "statements": 0 - }, - "src/odin/tools/base.py": { - "covered": 8, - "missing": 1, - "percent": 88.89, - "statements": 9 - }, - "src/odin/tools/file_ops.py": { - "covered": 30, - "missing": 0, - "percent": 100.0, - "statements": 30 - }, - "src/odin/tools/http.py": { - "covered": 8, - "missing": 13, - "percent": 38.1, - "statements": 21 - }, - "src/odin/tools/process.py": { - "covered": 27, - "missing": 0, - "percent": 100.0, - "statements": 27 - }, - "src/odin/tools/shell.py": { - "covered": 19, - "missing": 1, - "percent": 95.0, - "statements": 20 - }, - "src/odin/types.py": { - "covered": 37, - "missing": 0, - "percent": 100.0, - "statements": 37 - }, - "src/odin_log/__init__.py": { - "covered": 2, - "missing": 0, - "percent": 100.0, - "statements": 2 - }, - "src/odin_log/logger.py": { - "covered": 8, - "missing": 11, - "percent": 42.11, - "statements": 19 - }, - "src/packaging/__init__.py": { - "covered": 0, - "missing": 0, - "percent": 100.0, - "statements": 0 - }, - "src/permissions/__init__.py": { - "covered": 3, - "missing": 0, - "percent": 100.0, - "statements": 3 - }, - "src/permissions/host_access.py": { - "covered": 154, - "missing": 0, - "percent": 100.0, - "statements": 154 - }, - "src/permissions/manager.py": { - "covered": 78, - "missing": 0, - "percent": 100.0, - "statements": 78 - }, - "src/permissions/token_manager.py": { - "covered": 132, - "missing": 0, - "percent": 100.0, - "statements": 132 - }, - "src/planning/__init__.py": { - "covered": 2, - "missing": 0, - "percent": 100.0, - "statements": 2 - }, - "src/planning/store.py": { - "covered": 101, - "missing": 0, - "percent": 100.0, - "statements": 101 - }, - "src/relevance.py": { - "covered": 21, - "missing": 1, - "percent": 95.45, - "statements": 22 - }, - "src/scheduler/__init__.py": { - "covered": 3, - "missing": 0, - "percent": 100.0, - "statements": 3 - }, - "src/scheduler/history.py": { - "covered": 102, - "missing": 9, - "percent": 91.89, - "statements": 111 - }, - "src/scheduler/scheduler.py": { - "covered": 566, - "missing": 67, - "percent": 89.42, - "statements": 633 - }, - "src/search/__init__.py": { - "covered": 5, - "missing": 0, - "percent": 100.0, - "statements": 5 - }, - "src/search/embedder.py": { - "covered": 15, - "missing": 14, - "percent": 51.72, - "statements": 29 - }, - "src/search/fts.py": { - "covered": 151, - "missing": 0, - "percent": 100.0, - "statements": 151 - }, - "src/search/hybrid.py": { - "covered": 18, - "missing": 0, - "percent": 100.0, - "statements": 18 - }, - "src/search/sqlite_vec.py": { - "covered": 14, - "missing": 0, - "percent": 100.0, - "statements": 14 - }, - "src/search/vectorstore.py": { - "covered": 52, - "missing": 102, - "percent": 33.77, - "statements": 154 - }, - "src/sessions/__init__.py": { - "covered": 2, - "missing": 0, - "percent": 100.0, - "statements": 2 - }, - "src/sessions/manager.py": { - "covered": 775, - "missing": 70, - "percent": 91.72, - "statements": 845 - }, - "src/setup_wizard.py": { - "covered": 51, - "missing": 8, - "percent": 86.44, - "statements": 59 - }, - "src/tools/__init__.py": { - "covered": 7, - "missing": 0, - "percent": 100.0, - "statements": 7 - }, - "src/tools/affordances.py": { - "covered": 56, - "missing": 0, - "percent": 100.0, - "statements": 56 - }, - "src/tools/autonomous_loop.py": { - "covered": 155, - "missing": 83, - "percent": 65.13, - "statements": 238 - }, - "src/tools/branch_freshness.py": { - "covered": 99, - "missing": 0, - "percent": 100.0, - "statements": 99 - }, - "src/tools/bulkhead.py": { - "covered": 95, - "missing": 3, - "percent": 96.94, - "statements": 98 - }, - "src/tools/defs/__init__.py": { - "covered": 0, - "missing": 0, - "percent": 100.0, - "statements": 0 - }, - "src/tools/defs/agents.py": { - "covered": 1, - "missing": 0, - "percent": 100.0, - "statements": 1 - }, - "src/tools/defs/browser_web.py": { - "covered": 1, - "missing": 0, - "percent": 100.0, - "statements": 1 - }, - "src/tools/defs/channel_process_loops.py": { - "covered": 1, - "missing": 0, - "percent": 100.0, - "statements": 1 - }, - "src/tools/defs/devops.py": { - "covered": 1, - "missing": 0, - "percent": 100.0, - "statements": 1 - }, - "src/tools/defs/integrations_email.py": { - "covered": 1, - "missing": 0, - "percent": 100.0, - "statements": 1 - }, - "src/tools/defs/media_scheduling.py": { - "covered": 1, - "missing": 0, - "percent": 100.0, - "statements": 1 - }, - "src/tools/defs/memory_skills.py": { - "covered": 1, - "missing": 0, - "percent": 100.0, - "statements": 1 - }, - "src/tools/defs/system_files.py": { - "covered": 1, - "missing": 0, - "percent": 100.0, - "statements": 1 - }, - "src/tools/defs/tasks_knowledge.py": { - "covered": 1, - "missing": 0, - "percent": 100.0, - "statements": 1 - }, - "src/tools/docker_ops.py": { - "covered": 267, - "missing": 1, - "percent": 99.63, - "statements": 268 - }, - "src/tools/email_client.py": { - "covered": 210, - "missing": 21, - "percent": 90.91, - "statements": 231 - }, - "src/tools/executor.py": { - "covered": 301, - "missing": 32, - "percent": 90.39, - "statements": 333 - }, - "src/tools/git_ops.py": { - "covered": 168, - "missing": 1, - "percent": 99.41, - "statements": 169 - }, - "src/tools/handlers/__init__.py": { - "covered": 0, - "missing": 0, - "percent": 100.0, - "statements": 0 - }, - "src/tools/handlers/browser_web.py": { - "covered": 39, - "missing": 25, - "percent": 60.94, - "statements": 64 - }, - "src/tools/handlers/coding.py": { - "covered": 138, - "missing": 0, - "percent": 100.0, - "statements": 138 - }, - "src/tools/handlers/comms.py": { - "covered": 52, - "missing": 58, - "percent": 47.27, - "statements": 110 - }, - "src/tools/handlers/deps.py": { - "covered": 59, - "missing": 1, - "percent": 98.33, - "statements": 60 - }, - "src/tools/handlers/devops.py": { - "covered": 105, - "missing": 5, - "percent": 95.45, - "statements": 110 - }, - "src/tools/handlers/files_docs.py": { - "covered": 109, - "missing": 0, - "percent": 100.0, - "statements": 109 - }, - "src/tools/handlers/state.py": { - "covered": 301, - "missing": 0, - "percent": 100.0, - "statements": 301 - }, - "src/tools/handlers/system.py": { - "covered": 113, - "missing": 64, - "percent": 63.84, - "statements": 177 - }, - "src/tools/handlers/validation.py": { - "covered": 8, - "missing": 27, - "percent": 22.86, - "statements": 35 - }, - "src/tools/http_probe_ops.py": { - "covered": 75, - "missing": 0, - "percent": 100.0, - "statements": 75 - }, - "src/tools/kubectl_ops.py": { - "covered": 213, - "missing": 1, - "percent": 99.53, - "statements": 214 - }, - "src/tools/mcp_client.py": { - "covered": 265, - "missing": 81, - "percent": 76.59, - "statements": 346 - }, - "src/tools/output_streamer.py": { - "covered": 91, - "missing": 0, - "percent": 100.0, - "statements": 91 - }, - "src/tools/post_validation.py": { - "covered": 315, - "missing": 36, - "percent": 89.74, - "statements": 351 - }, - "src/tools/process_manager.py": { - "covered": 129, - "missing": 21, - "percent": 86.0, - "statements": 150 - }, - "src/tools/recovery.py": { - "covered": 119, - "missing": 0, - "percent": 100.0, - "statements": 119 - }, - "src/tools/registry.py": { - "covered": 22, - "missing": 0, - "percent": 100.0, - "statements": 22 - }, - "src/tools/result_validator.py": { - "covered": 98, - "missing": 11, - "percent": 89.91, - "statements": 109 - }, - "src/tools/risk_classifier.py": { - "covered": 156, - "missing": 9, - "percent": 94.55, - "statements": 165 - }, - "src/tools/skill_context.py": { - "covered": 208, - "missing": 0, - "percent": 100.0, - "statements": 208 - }, - "src/tools/skill_manager.py": { - "covered": 553, - "missing": 157, - "percent": 77.89, - "statements": 710 - }, - "src/tools/ssh.py": { - "covered": 89, - "missing": 10, - "percent": 89.9, - "statements": 99 - }, - "src/tools/ssh_pool.py": { - "covered": 70, - "missing": 13, - "percent": 84.34, - "statements": 83 - }, - "src/tools/terraform_ops.py": { - "covered": 156, - "missing": 1, - "percent": 99.36, - "statements": 157 - }, - "src/tools/time_parser.py": { - "covered": 108, - "missing": 0, - "percent": 100.0, - "statements": 108 - }, - "src/tools/tool_text.py": { - "covered": 7, - "missing": 3, - "percent": 70.0, - "statements": 10 - }, - "src/tools/url_safety.py": { - "covered": 80, - "missing": 15, - "percent": 84.21, - "statements": 95 - }, - "src/tools/web.py": { - "covered": 98, - "missing": 0, - "percent": 100.0, - "statements": 98 - }, - "src/trajectories/__init__.py": { - "covered": 2, - "missing": 0, - "percent": 100.0, - "statements": 2 - }, - "src/trajectories/saver.py": { - "covered": 181, - "missing": 11, - "percent": 94.27, - "statements": 192 - }, - "src/version.py": { - "covered": 8, - "missing": 12, - "percent": 40.0, - "statements": 20 - }, - "src/web/__init__.py": { - "covered": 0, - "missing": 0, - "percent": 100.0, - "statements": 0 - }, - "src/web/api/__init__.py": { - "covered": 75, - "missing": 0, - "percent": 100.0, - "statements": 75 - }, - "src/web/api/agents_loops.py": { - "covered": 148, - "missing": 2, - "percent": 98.67, - "statements": 150 - }, - "src/web/api/codex_admin.py": { - "covered": 203, - "missing": 16, - "percent": 92.69, - "statements": 219 - }, - "src/web/api/config_admin.py": { - "covered": 288, - "missing": 0, - "percent": 100.0, - "statements": 288 - }, - "src/web/api/integrations.py": { - "covered": 266, - "missing": 0, - "percent": 100.0, - "statements": 266 - }, - "src/web/api/knowledge_mem.py": { - "covered": 247, - "missing": 2, - "percent": 99.2, - "statements": 249 - }, - "src/web/api/llm_admin.py": { - "covered": 393, - "missing": 3, - "percent": 99.24, - "statements": 396 - }, - "src/web/api/observability.py": { - "covered": 222, - "missing": 10, - "percent": 95.69, - "statements": 232 - }, - "src/web/api/schedules_api.py": { - "covered": 110, - "missing": 0, - "percent": 100.0, - "statements": 110 - }, - "src/web/api/security.py": { - "covered": 332, - "missing": 17, - "percent": 95.13, - "statements": 349 - }, - "src/web/api/self_update.py": { - "covered": 79, - "missing": 9, - "percent": 89.77, - "statements": 88 - }, - "src/web/api/sessions_chat.py": { - "covered": 288, - "missing": 15, - "percent": 95.05, - "statements": 303 - }, - "src/web/api/skills_api.py": { - "covered": 79, - "missing": 56, - "percent": 58.52, - "statements": 135 - }, - "src/web/api_common.py": { - "covered": 66, - "missing": 21, - "percent": 75.86, - "statements": 87 - }, - "src/web/chat.py": { - "covered": 119, - "missing": 12, - "percent": 90.84, - "statements": 131 - }, - "src/web/websocket.py": { - "covered": 200, - "missing": 27, - "percent": 88.11, - "statements": 227 - } + "src/__init__.py": { + "covered": 3, + "missing": 0, + "percent": 100.0, + "statements": 3 + }, + "src/agents/__init__.py": { + "covered": 5, + "missing": 0, + "percent": 100.0, + "statements": 5 + }, + "src/agents/loop_bridge.py": { + "covered": 90, + "missing": 1, + "percent": 98.9, + "statements": 91 + }, + "src/agents/manager.py": { + "covered": 485, + "missing": 59, + "percent": 89.15, + "statements": 544 + }, + "src/agents/trajectory.py": { + "covered": 133, + "missing": 16, + "percent": 89.26, + "statements": 149 + }, + "src/async_utils.py": { + "covered": 13, + "missing": 2, + "percent": 86.67, + "statements": 15 + }, + "src/audit/__init__.py": { + "covered": 3, + "missing": 0, + "percent": 100.0, + "statements": 3 + }, + "src/audit/diff_tracker.py": { + "covered": 60, + "missing": 0, + "percent": 100.0, + "statements": 60 + }, + "src/audit/logger.py": { + "covered": 276, + "missing": 27, + "percent": 91.09, + "statements": 303 + }, + "src/audit/signer.py": { + "covered": 71, + "missing": 2, + "percent": 97.26, + "statements": 73 + }, + "src/config/__init__.py": { + "covered": 30, + "missing": 0, + "percent": 100.0, + "statements": 30 + }, + "src/config/schema.py": { + "covered": 628, + "missing": 1, + "percent": 99.84, + "statements": 629 + }, + "src/constants.py": { + "covered": 17, + "missing": 0, + "percent": 100.0, + "statements": 17 + }, + "src/context/__init__.py": { + "covered": 2, + "missing": 0, + "percent": 100.0, + "statements": 2 + }, + "src/context/loader.py": { + "covered": 42, + "missing": 11, + "percent": 79.25, + "statements": 53 + }, + "src/database/__init__.py": { + "covered": 0, + "missing": 0, + "percent": 100.0, + "statements": 0 + }, + "src/database/repository.py": { + "covered": 31, + "missing": 0, + "percent": 100.0, + "statements": 31 + }, + "src/discord/__init__.py": { + "covered": 3, + "missing": 4, + "percent": 42.86, + "statements": 7 + }, + "src/discord/attachments.py": { + "covered": 315, + "missing": 14, + "percent": 95.74, + "statements": 329 + }, + "src/discord/background_task.py": { + "covered": 271, + "missing": 107, + "percent": 71.69, + "statements": 378 + }, + "src/discord/channel_config.py": { + "covered": 81, + "missing": 0, + "percent": 100.0, + "statements": 81 + }, + "src/discord/channel_logger.py": { + "covered": 104, + "missing": 9, + "percent": 92.04, + "statements": 113 + }, + "src/discord/channel_state.py": { + "covered": 77, + "missing": 16, + "percent": 82.8, + "statements": 93 + }, + "src/discord/client.py": { + "covered": 130, + "missing": 43, + "percent": 75.14, + "statements": 173 + }, + "src/discord/completion.py": { + "covered": 40, + "missing": 1, + "percent": 97.56, + "statements": 41 + }, + "src/discord/delivery.py": { + "covered": 99, + "missing": 3, + "percent": 97.06, + "statements": 102 + }, + "src/discord/helpers/__init__.py": { + "covered": 0, + "missing": 0, + "percent": 100.0, + "statements": 0 + }, + "src/discord/helpers/converters.py": { + "covered": 23, + "missing": 0, + "percent": 100.0, + "statements": 23 + }, + "src/discord/helpers/cooldowns.py": { + "covered": 25, + "missing": 0, + "percent": 100.0, + "statements": 25 + }, + "src/discord/helpers/embeds.py": { + "covered": 28, + "missing": 0, + "percent": 100.0, + "statements": 28 + }, + "src/discord/helpers/pagination.py": { + "covered": 21, + "missing": 20, + "percent": 51.22, + "statements": 41 + }, + "src/discord/helpers/permissions.py": { + "covered": 12, + "missing": 19, + "percent": 38.71, + "statements": 31 + }, + "src/discord/housekeeping.py": { + "covered": 49, + "missing": 11, + "percent": 81.67, + "statements": 60 + }, + "src/discord/intake_pipeline.py": { + "covered": 338, + "missing": 57, + "percent": 85.57, + "statements": 395 + }, + "src/discord/llm_gateway.py": { + "covered": 368, + "missing": 0, + "percent": 100.0, + "statements": 368 + }, + "src/discord/native_tools/__init__.py": { + "covered": 2, + "missing": 0, + "percent": 100.0, + "statements": 2 + }, + "src/discord/native_tools/agents_tasks.py": { + "covered": 395, + "missing": 16, + "percent": 96.11, + "statements": 411 + }, + "src/discord/native_tools/channel_ops.py": { + "covered": 113, + "missing": 0, + "percent": 100.0, + "statements": 113 + }, + "src/discord/native_tools/knowledge.py": { + "covered": 122, + "missing": 0, + "percent": 100.0, + "statements": 122 + }, + "src/discord/native_tools/media.py": { + "covered": 166, + "missing": 0, + "percent": 100.0, + "statements": 166 + }, + "src/discord/native_tools/registry.py": { + "covered": 83, + "missing": 6, + "percent": 93.26, + "statements": 89 + }, + "src/discord/native_tools/scheduling.py": { + "covered": 124, + "missing": 0, + "percent": 100.0, + "statements": 124 + }, + "src/discord/native_tools/skills_tools.py": { + "covered": 87, + "missing": 8, + "percent": 91.58, + "statements": 95 + }, + "src/discord/prompts.py": { + "covered": 85, + "missing": 25, + "percent": 77.27, + "statements": 110 + }, + "src/discord/response_guards.py": { + "covered": 154, + "missing": 0, + "percent": 100.0, + "statements": 154 + }, + "src/discord/scheduled_events.py": { + "covered": 225, + "missing": 0, + "percent": 100.0, + "statements": 225 + }, + "src/discord/slash_commands.py": { + "covered": 20, + "missing": 49, + "percent": 28.99, + "statements": 69 + }, + "src/discord/tool_catalog.py": { + "covered": 38, + "missing": 0, + "percent": 100.0, + "statements": 38 + }, + "src/discord/tool_loop.py": { + "covered": 713, + "missing": 57, + "percent": 92.6, + "statements": 770 + }, + "src/discord/tool_loop_helpers.py": { + "covered": 30, + "missing": 9, + "percent": 76.92, + "statements": 39 + }, + "src/discord/turn_recorder.py": { + "covered": 92, + "missing": 16, + "percent": 85.19, + "statements": 108 + }, + "src/discord/turn_resume.py": { + "covered": 174, + "missing": 30, + "percent": 85.29, + "statements": 204 + }, + "src/discord/wiring.py": { + "covered": 332, + "missing": 76, + "percent": 81.37, + "statements": 408 + }, + "src/error_presentation.py": { + "covered": 30, + "missing": 0, + "percent": 100.0, + "statements": 30 + }, + "src/health/__init__.py": { + "covered": 3, + "missing": 0, + "percent": 100.0, + "statements": 3 + }, + "src/health/checker.py": { + "covered": 239, + "missing": 0, + "percent": 100.0, + "statements": 239 + }, + "src/health/grafana_alerts.py": { + "covered": 235, + "missing": 0, + "percent": 100.0, + "statements": 235 + }, + "src/health/metrics.py": { + "covered": 193, + "missing": 18, + "percent": 91.47, + "statements": 211 + }, + "src/health/server.py": { + "covered": 465, + "missing": 178, + "percent": 72.32, + "statements": 643 + }, + "src/health/startup.py": { + "covered": 249, + "missing": 11, + "percent": 95.77, + "statements": 260 + }, + "src/health/subsystem_guard.py": { + "covered": 213, + "missing": 0, + "percent": 100.0, + "statements": 213 + }, + "src/json_store.py": { + "covered": 43, + "missing": 4, + "percent": 91.49, + "statements": 47 + }, + "src/knowledge/__init__.py": { + "covered": 2, + "missing": 0, + "percent": 100.0, + "statements": 2 + }, + "src/knowledge/importer.py": { + "covered": 165, + "missing": 9, + "percent": 94.83, + "statements": 174 + }, + "src/knowledge/store.py": { + "covered": 328, + "missing": 89, + "percent": 78.66, + "statements": 417 + }, + "src/learning/__init__.py": { + "covered": 2, + "missing": 0, + "percent": 100.0, + "statements": 2 + }, + "src/learning/loop_reflection.py": { + "covered": 49, + "missing": 5, + "percent": 90.74, + "statements": 54 + }, + "src/learning/reflector.py": { + "covered": 621, + "missing": 60, + "percent": 91.19, + "statements": 681 + }, + "src/llm/__init__.py": { + "covered": 12, + "missing": 0, + "percent": 100.0, + "statements": 12 + }, + "src/llm/auxiliary.py": { + "covered": 65, + "missing": 0, + "percent": 100.0, + "statements": 65 + }, + "src/llm/backoff.py": { + "covered": 10, + "missing": 0, + "percent": 100.0, + "statements": 10 + }, + "src/llm/circuit_breaker.py": { + "covered": 43, + "missing": 0, + "percent": 100.0, + "statements": 43 + }, + "src/llm/codex_auth.py": { + "covered": 375, + "missing": 27, + "percent": 93.28, + "statements": 402 + }, + "src/llm/context_compressor.py": { + "covered": 165, + "missing": 18, + "percent": 90.16, + "statements": 183 + }, + "src/llm/cost_tracker.py": { + "covered": 86, + "missing": 0, + "percent": 100.0, + "statements": 86 + }, + "src/llm/errors.py": { + "covered": 18, + "missing": 0, + "percent": 100.0, + "statements": 18 + }, + "src/llm/kimi.py": { + "covered": 233, + "missing": 1, + "percent": 99.57, + "statements": 234 + }, + "src/llm/model_breaker.py": { + "covered": 100, + "missing": 2, + "percent": 98.04, + "statements": 102 + }, + "src/llm/ollama.py": { + "covered": 179, + "missing": 5, + "percent": 97.28, + "statements": 184 + }, + "src/llm/openai_codex.py": { + "covered": 431, + "missing": 51, + "percent": 89.42, + "statements": 482 + }, + "src/llm/provider.py": { + "covered": 15, + "missing": 3, + "percent": 83.33, + "statements": 18 + }, + "src/llm/recovery.py": { + "covered": 107, + "missing": 5, + "percent": 95.54, + "statements": 112 + }, + "src/llm/secret_scrubber.py": { + "covered": 7, + "missing": 0, + "percent": 100.0, + "statements": 7 + }, + "src/llm/system_prompt.py": { + "covered": 34, + "missing": 3, + "percent": 91.89, + "statements": 37 + }, + "src/llm/types.py": { + "covered": 21, + "missing": 0, + "percent": 100.0, + "statements": 21 + }, + "src/models/__init__.py": { + "covered": 5, + "missing": 0, + "percent": 100.0, + "statements": 5 + }, + "src/models/guild.py": { + "covered": 18, + "missing": 0, + "percent": 100.0, + "statements": 18 + }, + "src/models/infraction.py": { + "covered": 20, + "missing": 0, + "percent": 100.0, + "statements": 20 + }, + "src/models/reminder.py": { + "covered": 20, + "missing": 0, + "percent": 100.0, + "statements": 20 + }, + "src/models/user.py": { + "covered": 14, + "missing": 0, + "percent": 100.0, + "statements": 14 + }, + "src/monitoring/__init__.py": { + "covered": 2, + "missing": 0, + "percent": 100.0, + "statements": 2 + }, + "src/monitoring/resource_usage.py": { + "covered": 154, + "missing": 11, + "percent": 93.33, + "statements": 165 + }, + "src/notifications/__init__.py": { + "covered": 4, + "missing": 0, + "percent": 100.0, + "statements": 4 + }, + "src/notifications/issue_tracker.py": { + "covered": 261, + "missing": 22, + "percent": 92.23, + "statements": 283 + }, + "src/notifications/outbound_webhooks.py": { + "covered": 258, + "missing": 3, + "percent": 98.85, + "statements": 261 + }, + "src/notifications/slack.py": { + "covered": 128, + "missing": 0, + "percent": 100.0, + "statements": 128 + }, + "src/observability/__init__.py": { + "covered": 3, + "missing": 0, + "percent": 100.0, + "statements": 3 + }, + "src/observability/aggregates.py": { + "covered": 122, + "missing": 5, + "percent": 96.06, + "statements": 127 + }, + "src/observability/context_trace.py": { + "covered": 120, + "missing": 12, + "percent": 90.91, + "statements": 132 + }, + "src/observability/correlation.py": { + "covered": 11, + "missing": 2, + "percent": 84.62, + "statements": 13 + }, + "src/observability/failure_classes.py": { + "covered": 14, + "missing": 2, + "percent": 87.5, + "statements": 16 + }, + "src/odin/__init__.py": { + "covered": 7, + "missing": 0, + "percent": 100.0, + "statements": 7 + }, + "src/odin/cli.py": { + "covered": 46, + "missing": 13, + "percent": 77.97, + "statements": 59 + }, + "src/odin/context.py": { + "covered": 122, + "missing": 7, + "percent": 94.57, + "statements": 129 + }, + "src/odin/executor.py": { + "covered": 42, + "missing": 1, + "percent": 97.67, + "statements": 43 + }, + "src/odin/plan_loader.py": { + "covered": 29, + "missing": 8, + "percent": 78.38, + "statements": 37 + }, + "src/odin/planner.py": { + "covered": 95, + "missing": 3, + "percent": 96.94, + "statements": 98 + }, + "src/odin/registry.py": { + "covered": 27, + "missing": 0, + "percent": 100.0, + "statements": 27 + }, + "src/odin/reporter.py": { + "covered": 25, + "missing": 0, + "percent": 100.0, + "statements": 25 + }, + "src/odin/tools/__init__.py": { + "covered": 0, + "missing": 0, + "percent": 100.0, + "statements": 0 + }, + "src/odin/tools/base.py": { + "covered": 8, + "missing": 1, + "percent": 88.89, + "statements": 9 + }, + "src/odin/tools/file_ops.py": { + "covered": 30, + "missing": 0, + "percent": 100.0, + "statements": 30 + }, + "src/odin/tools/http.py": { + "covered": 8, + "missing": 13, + "percent": 38.1, + "statements": 21 + }, + "src/odin/tools/process.py": { + "covered": 29, + "missing": 0, + "percent": 100.0, + "statements": 29 + }, + "src/odin/tools/shell.py": { + "covered": 19, + "missing": 1, + "percent": 95.0, + "statements": 20 + }, + "src/odin/types.py": { + "covered": 37, + "missing": 0, + "percent": 100.0, + "statements": 37 + }, + "src/odin_log/__init__.py": { + "covered": 2, + "missing": 0, + "percent": 100.0, + "statements": 2 + }, + "src/odin_log/logger.py": { + "covered": 8, + "missing": 11, + "percent": 42.11, + "statements": 19 + }, + "src/packaging/__init__.py": { + "covered": 0, + "missing": 0, + "percent": 100.0, + "statements": 0 + }, + "src/permissions/__init__.py": { + "covered": 3, + "missing": 0, + "percent": 100.0, + "statements": 3 + }, + "src/permissions/host_access.py": { + "covered": 154, + "missing": 0, + "percent": 100.0, + "statements": 154 + }, + "src/permissions/manager.py": { + "covered": 78, + "missing": 0, + "percent": 100.0, + "statements": 78 + }, + "src/permissions/token_manager.py": { + "covered": 132, + "missing": 0, + "percent": 100.0, + "statements": 132 + }, + "src/planning/__init__.py": { + "covered": 2, + "missing": 0, + "percent": 100.0, + "statements": 2 + }, + "src/planning/store.py": { + "covered": 101, + "missing": 0, + "percent": 100.0, + "statements": 101 + }, + "src/relevance.py": { + "covered": 21, + "missing": 1, + "percent": 95.45, + "statements": 22 + }, + "src/restart.py": { + "covered": 21, + "missing": 0, + "percent": 100.0, + "statements": 21 + }, + "src/scheduler/__init__.py": { + "covered": 3, + "missing": 0, + "percent": 100.0, + "statements": 3 + }, + "src/scheduler/history.py": { + "covered": 102, + "missing": 9, + "percent": 91.89, + "statements": 111 + }, + "src/scheduler/scheduler.py": { + "covered": 602, + "missing": 67, + "percent": 89.99, + "statements": 669 + }, + "src/search/__init__.py": { + "covered": 5, + "missing": 0, + "percent": 100.0, + "statements": 5 + }, + "src/search/embedder.py": { + "covered": 15, + "missing": 14, + "percent": 51.72, + "statements": 29 + }, + "src/search/fts.py": { + "covered": 151, + "missing": 0, + "percent": 100.0, + "statements": 151 + }, + "src/search/hybrid.py": { + "covered": 18, + "missing": 0, + "percent": 100.0, + "statements": 18 + }, + "src/search/sqlite_vec.py": { + "covered": 14, + "missing": 0, + "percent": 100.0, + "statements": 14 + }, + "src/search/vectorstore.py": { + "covered": 52, + "missing": 102, + "percent": 33.77, + "statements": 154 + }, + "src/sessions/__init__.py": { + "covered": 2, + "missing": 0, + "percent": 100.0, + "statements": 2 + }, + "src/sessions/manager.py": { + "covered": 775, + "missing": 70, + "percent": 91.72, + "statements": 845 + }, + "src/setup_wizard.py": { + "covered": 51, + "missing": 8, + "percent": 86.44, + "statements": 59 + }, + "src/tools/__init__.py": { + "covered": 7, + "missing": 0, + "percent": 100.0, + "statements": 7 + }, + "src/tools/affordances.py": { + "covered": 56, + "missing": 0, + "percent": 100.0, + "statements": 56 + }, + "src/tools/agent_tool_policy.py": { + "covered": 45, + "missing": 0, + "percent": 100.0, + "statements": 45 + }, + "src/tools/autonomous_loop.py": { + "covered": 155, + "missing": 83, + "percent": 65.13, + "statements": 238 + }, + "src/tools/branch_freshness.py": { + "covered": 99, + "missing": 0, + "percent": 100.0, + "statements": 99 + }, + "src/tools/bulkhead.py": { + "covered": 95, + "missing": 3, + "percent": 96.94, + "statements": 98 + }, + "src/tools/defs/__init__.py": { + "covered": 0, + "missing": 0, + "percent": 100.0, + "statements": 0 + }, + "src/tools/defs/agents.py": { + "covered": 5, + "missing": 0, + "percent": 100.0, + "statements": 5 + }, + "src/tools/defs/browser_web.py": { + "covered": 1, + "missing": 0, + "percent": 100.0, + "statements": 1 + }, + "src/tools/defs/channel_process_loops.py": { + "covered": 1, + "missing": 0, + "percent": 100.0, + "statements": 1 + }, + "src/tools/defs/devops.py": { + "covered": 1, + "missing": 0, + "percent": 100.0, + "statements": 1 + }, + "src/tools/defs/integrations_email.py": { + "covered": 1, + "missing": 0, + "percent": 100.0, + "statements": 1 + }, + "src/tools/defs/media_scheduling.py": { + "covered": 1, + "missing": 0, + "percent": 100.0, + "statements": 1 + }, + "src/tools/defs/memory_skills.py": { + "covered": 1, + "missing": 0, + "percent": 100.0, + "statements": 1 + }, + "src/tools/defs/system_files.py": { + "covered": 1, + "missing": 0, + "percent": 100.0, + "statements": 1 + }, + "src/tools/defs/tasks_knowledge.py": { + "covered": 1, + "missing": 0, + "percent": 100.0, + "statements": 1 + }, + "src/tools/docker_ops.py": { + "covered": 267, + "missing": 1, + "percent": 99.63, + "statements": 268 + }, + "src/tools/email_client.py": { + "covered": 210, + "missing": 21, + "percent": 90.91, + "statements": 231 + }, + "src/tools/executor.py": { + "covered": 402, + "missing": 25, + "percent": 94.15, + "statements": 427 + }, + "src/tools/git_ops.py": { + "covered": 176, + "missing": 1, + "percent": 99.44, + "statements": 177 + }, + "src/tools/handlers/__init__.py": { + "covered": 0, + "missing": 0, + "percent": 100.0, + "statements": 0 + }, + "src/tools/handlers/browser_web.py": { + "covered": 43, + "missing": 24, + "percent": 64.18, + "statements": 67 + }, + "src/tools/handlers/coding.py": { + "covered": 138, + "missing": 0, + "percent": 100.0, + "statements": 138 + }, + "src/tools/handlers/comms.py": { + "covered": 52, + "missing": 58, + "percent": 47.27, + "statements": 110 + }, + "src/tools/handlers/deps.py": { + "covered": 60, + "missing": 0, + "percent": 100.0, + "statements": 60 + }, + "src/tools/handlers/devops.py": { + "covered": 105, + "missing": 5, + "percent": 95.45, + "statements": 110 + }, + "src/tools/handlers/files_docs.py": { + "covered": 117, + "missing": 0, + "percent": 100.0, + "statements": 117 + }, + "src/tools/handlers/state.py": { + "covered": 326, + "missing": 0, + "percent": 100.0, + "statements": 326 + }, + "src/tools/handlers/system.py": { + "covered": 140, + "missing": 49, + "percent": 74.07, + "statements": 189 + }, + "src/tools/handlers/validation.py": { + "covered": 29, + "missing": 6, + "percent": 82.86, + "statements": 35 + }, + "src/tools/http_probe_ops.py": { + "covered": 84, + "missing": 0, + "percent": 100.0, + "statements": 84 + }, + "src/tools/image/__init__.py": { + "covered": 6, + "missing": 0, + "percent": 100.0, + "statements": 6 + }, + "src/tools/image/base.py": { + "covered": 57, + "missing": 2, + "percent": 96.61, + "statements": 59 + }, + "src/tools/image/comfyui_backend.py": { + "covered": 29, + "missing": 2, + "percent": 93.55, + "statements": 31 + }, + "src/tools/image/openai_backend.py": { + "covered": 135, + "missing": 15, + "percent": 90.0, + "statements": 150 + }, + "src/tools/image/selector.py": { + "covered": 88, + "missing": 4, + "percent": 95.65, + "statements": 92 + }, + "src/tools/kubectl_ops.py": { + "covered": 213, + "missing": 1, + "percent": 99.53, + "statements": 214 + }, + "src/tools/mcp_client.py": { + "covered": 265, + "missing": 81, + "percent": 76.59, + "statements": 346 + }, + "src/tools/output_streamer.py": { + "covered": 91, + "missing": 0, + "percent": 100.0, + "statements": 91 + }, + "src/tools/post_validation.py": { + "covered": 326, + "missing": 26, + "percent": 92.61, + "statements": 352 + }, + "src/tools/process_manager.py": { + "covered": 155, + "missing": 20, + "percent": 88.57, + "statements": 175 + }, + "src/tools/recovery.py": { + "covered": 119, + "missing": 0, + "percent": 100.0, + "statements": 119 + }, + "src/tools/registry.py": { + "covered": 22, + "missing": 0, + "percent": 100.0, + "statements": 22 + }, + "src/tools/result_validator.py": { + "covered": 99, + "missing": 11, + "percent": 90.0, + "statements": 110 + }, + "src/tools/risk_classifier.py": { + "covered": 156, + "missing": 9, + "percent": 94.55, + "statements": 165 + }, + "src/tools/safe_fetch.py": { + "covered": 124, + "missing": 15, + "percent": 89.21, + "statements": 139 + }, + "src/tools/skill_context.py": { + "covered": 223, + "missing": 0, + "percent": 100.0, + "statements": 223 + }, + "src/tools/skill_manager.py": { + "covered": 555, + "missing": 151, + "percent": 78.61, + "statements": 706 + }, + "src/tools/ssh.py": { + "covered": 214, + "missing": 8, + "percent": 96.4, + "statements": 222 + }, + "src/tools/ssh_pool.py": { + "covered": 70, + "missing": 13, + "percent": 84.34, + "statements": 83 + }, + "src/tools/terraform_ops.py": { + "covered": 156, + "missing": 1, + "percent": 99.36, + "statements": 157 + }, + "src/tools/time_parser.py": { + "covered": 108, + "missing": 0, + "percent": 100.0, + "statements": 108 + }, + "src/tools/tool_text.py": { + "covered": 7, + "missing": 3, + "percent": 70.0, + "statements": 10 + }, + "src/tools/url_safety.py": { + "covered": 104, + "missing": 15, + "percent": 87.39, + "statements": 119 + }, + "src/tools/web.py": { + "covered": 99, + "missing": 0, + "percent": 100.0, + "statements": 99 + }, + "src/tools/workspace.py": { + "covered": 146, + "missing": 5, + "percent": 96.69, + "statements": 151 + }, + "src/trajectories/__init__.py": { + "covered": 2, + "missing": 0, + "percent": 100.0, + "statements": 2 + }, + "src/trajectories/saver.py": { + "covered": 184, + "missing": 11, + "percent": 94.36, + "statements": 195 + }, + "src/turn_state/__init__.py": { + "covered": 2, + "missing": 0, + "percent": 100.0, + "statements": 2 + }, + "src/turn_state/codec.py": { + "covered": 79, + "missing": 0, + "percent": 100.0, + "statements": 79 + }, + "src/turn_state/durability.py": { + "covered": 126, + "missing": 16, + "percent": 88.73, + "statements": 142 + }, + "src/turn_state/store.py": { + "covered": 256, + "missing": 34, + "percent": 88.28, + "statements": 290 + }, + "src/version.py": { + "covered": 15, + "missing": 8, + "percent": 65.22, + "statements": 23 + }, + "src/web/__init__.py": { + "covered": 0, + "missing": 0, + "percent": 100.0, + "statements": 0 + }, + "src/web/api/__init__.py": { + "covered": 74, + "missing": 0, + "percent": 100.0, + "statements": 74 + }, + "src/web/api/agents_loops.py": { + "covered": 148, + "missing": 2, + "percent": 98.67, + "statements": 150 + }, + "src/web/api/codex_admin.py": { + "covered": 203, + "missing": 16, + "percent": 92.69, + "statements": 219 + }, + "src/web/api/config_admin.py": { + "covered": 292, + "missing": 0, + "percent": 100.0, + "statements": 292 + }, + "src/web/api/integrations.py": { + "covered": 266, + "missing": 0, + "percent": 100.0, + "statements": 266 + }, + "src/web/api/knowledge_mem.py": { + "covered": 279, + "missing": 2, + "percent": 99.29, + "statements": 281 + }, + "src/web/api/llm_admin.py": { + "covered": 510, + "missing": 3, + "percent": 99.42, + "statements": 513 + }, + "src/web/api/observability.py": { + "covered": 216, + "missing": 10, + "percent": 95.58, + "statements": 226 + }, + "src/web/api/schedules_api.py": { + "covered": 110, + "missing": 0, + "percent": 100.0, + "statements": 110 + }, + "src/web/api/security.py": { + "covered": 332, + "missing": 17, + "percent": 95.13, + "statements": 349 + }, + "src/web/api/self_update.py": { + "covered": 127, + "missing": 8, + "percent": 94.07, + "statements": 135 + }, + "src/web/api/sessions_chat.py": { + "covered": 288, + "missing": 15, + "percent": 95.05, + "statements": 303 + }, + "src/web/api/skills_api.py": { + "covered": 79, + "missing": 56, + "percent": 58.52, + "statements": 135 + }, + "src/web/api_common.py": { + "covered": 84, + "missing": 3, + "percent": 96.55, + "statements": 87 + }, + "src/web/chat.py": { + "covered": 119, + "missing": 12, + "percent": 90.84, + "statements": 131 + }, + "src/web/websocket.py": { + "covered": 218, + "missing": 27, + "percent": 88.98, + "statements": 245 + } } diff --git a/src/discord/intake_pipeline.py b/src/discord/intake_pipeline.py index 5b5454dd..0b6c8d87 100644 --- a/src/discord/intake_pipeline.py +++ b/src/discord/intake_pipeline.py @@ -47,6 +47,7 @@ from .prompts import PromptBuilder from .tool_loop import ToolLoopRunner from .turn_recorder import TurnRecorder + from .turn_resume import TurnResumeManager log = get_logger("discord") @@ -408,7 +409,7 @@ class MessagePipelineDeps: housekeeping: Housekeeping # post-turn cache maintenance # Suspended-turn resume manager (None = feature off). Default keeps every # existing construction working; wiring passes the real manager. - turn_resume: object | None = None + turn_resume: TurnResumeManager | None = None class MessagePipeline: diff --git a/src/discord/tool_loop.py b/src/discord/tool_loop.py index 4f21b983..d80ed21e 100644 --- a/src/discord/tool_loop.py +++ b/src/discord/tool_loop.py @@ -493,10 +493,12 @@ async def _run_with_guards( # Fail-closed epilogue: a durability failure (or any escape) # marks the turn FAILED so a half-written checkpoint can never # present itself as resumable. Best-effort — a fence loss here - # just means someone else owns the row now. + # just means someone else owns the row now, and even a + # cancellation delivered inside the settle must not replace the + # original escaping error. try: await st.durability.settle_terminal(cancelled=False, is_error=True) - except Exception: + except BaseException: # noqa: BLE001 — the original error must win log.warning("Durable failure mark failed (non-fatal)") raise diff --git a/src/discord/wiring.py b/src/discord/wiring.py index 57ab5d93..7ca33d80 100644 --- a/src/discord/wiring.py +++ b/src/discord/wiring.py @@ -37,6 +37,8 @@ from ..llm import CodexChatClient, KimiClient, OllamaClient from ..llm.codex_auth import CodexAuthPool from ..llm.cost_tracker import CostTracker +from ..llm.model_breaker import ModelBreakerRegistry +from ..llm.recovery import RecoveryPolicy from ..odin_log import get_logger from ..permissions import PermissionManager from ..permissions.host_access import HostAccessManager @@ -48,6 +50,7 @@ from ..tools.autonomous_loop import LoopManager from ..tools.workspace import DEFAULT_MEMORY_PATH from ..trajectories.saver import TrajectorySaver +from ..turn_state import TurnStateStore from .channel_config import ChannelConfigManager from .channel_logger import ChannelLogger from .channel_state import ChannelStateRegistry @@ -122,9 +125,9 @@ class BotServices: # Turn durability (2026-07-30): the checkpoint/ledger store (None = # disabled) and the model-scoped capacity-breaker registry. Both live # HERE so client rebuilds and live reloads never reset their state. - turn_store: object | None = None - model_breakers: object | None = None - recovery_policy_source: Callable | None = None + turn_store: TurnStateStore | None = None + model_breakers: ModelBreakerRegistry | None = None + recovery_policy_source: Callable[[], RecoveryPolicy] | None = None def build_services(config: Config) -> BotServices: # noqa: PLR0915 — linear composition root @@ -404,9 +407,6 @@ def build_services(config: Config) -> BotServices: # noqa: PLR0915 — linear c # Turn durability (2026-07-30): checkpoint/ledger store + the model- # scoped capacity breakers + the recovery policy source. Services-owned # so provider-client rebuilds and live reloads never reset their state. - from ..llm.model_breaker import ModelBreakerRegistry - from ..llm.recovery import RecoveryPolicy - _lr = config.llm_recovery model_breakers = ModelBreakerRegistry( generation_threshold=_lr.breaker_generation_threshold, @@ -423,8 +423,6 @@ def recovery_policy_source() -> RecoveryPolicy: turn_store = None if config.turn_state.enabled: - from ..turn_state import TurnStateStore - turn_store = TurnStateStore(config.turn_state.db_path) if not turn_store.available: turn_store = None # init failed — feature off, logged loudly diff --git a/src/turn_state/durability.py b/src/turn_state/durability.py index e06deb45..014190d1 100644 --- a/src/turn_state/durability.py +++ b/src/turn_state/durability.py @@ -300,16 +300,17 @@ async def suspend(self, st, reason: str) -> bool: except Exception: pass try: - assert self._store is not None and self._lease is not None + store, lease = self._store, self._lease + assert store is not None and lease is not None payload = await asyncio.to_thread( lambda: snapshot_chat_turn( st, - store_blob=self._store.store_blob_sync, + store_blob=store.store_blob_sync, generation_seq=self.generation_seq, extra=extra, ) ) - await asyncio.to_thread(self._store.suspend_sync, self._lease, payload) + await asyncio.to_thread(store.suspend_sync, lease, payload) except Exception: log.exception("Turn suspension failed — work NOT preserved") return False diff --git a/tests/test_agent_trajectory.py b/tests/test_agent_trajectory.py index 37ea362f..c9fa1898 100644 --- a/tests/test_agent_trajectory.py +++ b/tests/test_agent_trajectory.py @@ -698,7 +698,12 @@ async def test_trajectory_save_error_does_not_crash(self, tmp_path): ) assert agent.state == AgentState.COMPLETED - async def test_trajectory_recovery_attempts(self, tmp_path): + async def test_trajectory_recovery_attempts_field_stays_zero(self, tmp_path): + # Deliberate amendment (2026-07-30): the manager retry ladder is + # gone — transient recovery lives inside the iteration callback + # (src/llm/recovery.py). The trajectory keeps the field for shape + # compatibility; a failure escaping the callback FAILS the agent + # and recovery_attempts stays 0. saver = AgentTrajectorySaver(directory=str(tmp_path)) agent = AgentInfo( id="t12", label="recovery", goal="test recovery tracking", @@ -711,7 +716,7 @@ async def iter_cb(messages, prompt, tools): call_count += 1 if call_count == 1: raise TimeoutError("LLM timeout") - return {"text": "recovered", "tool_calls": []} + return {"text": "would recover", "tool_calls": []} await _run_agent( agent=agent, system_prompt="s", tools=[], @@ -720,9 +725,11 @@ async def iter_cb(messages, prompt, tools): trajectory_saver=saver, ) + assert call_count == 1 # no manager-level second call + assert agent.state == AgentState.FAILED entry = await saver.find_by_agent_id("t12") assert entry is not None - assert entry["recovery_attempts"] >= 1 + assert entry["recovery_attempts"] == 0 async def test_trajectory_tool_error_recorded(self, tmp_path): saver = AgentTrajectorySaver(directory=str(tmp_path)) diff --git a/tests/test_guard_capacity.py b/tests/test_guard_capacity.py index 802da1e7..7a24cf90 100644 --- a/tests/test_guard_capacity.py +++ b/tests/test_guard_capacity.py @@ -118,6 +118,12 @@ def test_transient_never_downgrades_unavailable(self): guard.mark_degraded_transient("llm_codex", "capacity") assert guard.get_state("llm_codex") == SubsystemState.UNAVAILABLE + def test_unregistered_names_are_noops(self): + guard = _guard() + guard.mark_degraded_transient("ghost", "capacity") + guard.mark_degraded("ghost", "manual") + assert guard.get_state("ghost") is None + def test_transient_never_converts_counter_degraded(self, monkeypatch): clock = {"now": 1000.0} monkeypatch.setattr( diff --git a/tests/test_llm_kimi.py b/tests/test_llm_kimi.py index f5809344..4548c0ca 100644 --- a/tests/test_llm_kimi.py +++ b/tests/test_llm_kimi.py @@ -303,3 +303,28 @@ async def test_close(self): c._session = session await c.close() assert session.closed and c._session is None + + +class TestRateLimitExhaustionTyped: + """429-exhausted raises LLMRateLimitError carrying the parsed Retry-After + (typed-error work, 2026-07-30). Quota rotation semantics unchanged.""" + + async def test_numeric_retry_after_header_is_carried(self): + from src.llm.errors import LLMRateLimitError + + c = _client(max_retries=0) # single attempt → immediate exhaustion + _with_session(c, _Resp(429, text="slow down", headers={"Retry-After": "7"})) + with pytest.raises(LLMRateLimitError) as exc_info: + await c._request_with_retry({"messages": []}) + assert exc_info.value.retry_after == 7.0 + assert exc_info.value.provider == "kimi" + assert "rate limited after 1 attempts" in str(exc_info.value) + + async def test_unparseable_retry_after_header_is_none(self): + from src.llm.errors import LLMRateLimitError + + c = _client(max_retries=0) + _with_session(c, _Resp(429, text="slow down", headers={"Retry-After": "soon"})) + with pytest.raises(LLMRateLimitError) as exc_info: + await c._request_with_retry({"messages": []}) + assert exc_info.value.retry_after is None diff --git a/tests/test_write_invariant_integration.py b/tests/test_write_invariant_integration.py index 82ebb771..8a35a7fa 100644 --- a/tests/test_write_invariant_integration.py +++ b/tests/test_write_invariant_integration.py @@ -263,3 +263,201 @@ async def kill_durability(tool_name, tool_input, *, user_id=None): await run_loop(bot, FakeMessage("go")) # The second batch never started: fail closed, not fail quiet. assert executed == ["one"] + + +class TestCancellationBranches: + async def test_stop_before_first_iteration_is_terminal_cancelled(self, tmp_path): + import asyncio + + bot, fake, store = build_with_store([text_response("never")], tmp_path) + msg = FakeMessage("go") + evt = bot.channel_state.cancel_events.setdefault( + str(msg.channel.id), asyncio.Event() + ) + evt.set() + text, *_ = await run_loop(bot, msg) + assert text.startswith("Task stopped by user.") + (status,) = store._conn.execute("SELECT status FROM turns").fetchone() + assert status == TurnStatus.TERMINAL_CANCELLED + + async def test_stop_during_recovery_wait_is_graceful(self, tmp_path): + import asyncio + + bot, fake, store = build_with_store([], tmp_path) + msg = FakeMessage("go") + evt = bot.channel_state.cancel_events.setdefault( + str(msg.channel.id), asyncio.Event() + ) + + def capacity_and_stop(): + evt.set() # /stop lands while the recovery wait begins + raise LLMCapacityError("overloaded", retry_after=5.0) + + fake.responses.append(capacity_and_stop) + text, _, is_error, *_ = await run_loop(bot, msg) + assert text.startswith("Task stopped by user.") + assert is_error is False + (status,) = store._conn.execute("SELECT status FROM turns").fetchone() + assert status == TurnStatus.TERMINAL_CANCELLED + + async def test_task_cancel_with_failing_settle_still_propagates(self, tmp_path): + import asyncio + from unittest.mock import AsyncMock, patch + + from src.turn_state.durability import TurnDurability + + bot, fake, store = build_with_store( + [tool_call_response(("run_command", {"command": "x"}))], tmp_path + ) + started = asyncio.Event() + + async def blocking_tool(tool_name, tool_input, *, user_id=None): + started.set() + await asyncio.sleep(3600) + + bot.tool_executor.execute = blocking_tool + with patch.object( + TurnDurability, + "settle_terminal", + AsyncMock(side_effect=asyncio.CancelledError()), + ): + task = asyncio.get_running_loop().create_task( + run_loop(bot, FakeMessage("go")) + ) + await asyncio.wait_for(started.wait(), timeout=5) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + async def test_escape_settle_failure_keeps_original_error(self, tmp_path): + from unittest.mock import AsyncMock, patch + + from src.tools.result_validator import ToolResult + from src.turn_state.durability import TurnDurability + + bot, fake, store = build_with_store( + [ + tool_call_response(("run_command", {"command": "one"})), + text_response("never"), + ], + tmp_path, + ) + + async def kill_durability(tool_name, tool_input, *, user_id=None): + store._conn.close() + return ToolResult(output="ok", tool_name=tool_name) + + bot.tool_executor.execute = kill_durability + with patch.object( + TurnDurability, + "settle_terminal", + AsyncMock(side_effect=RuntimeError("settle also broken")), + ): + # The ORIGINAL fail-closed error must surface, never the + # settle-bookkeeping failure. + with pytest.raises(TurnStateUnavailableError): + await run_loop(bot, FakeMessage("go")) + + +class TestSuspendFallback: + async def test_failed_suspension_reports_plain_error_never_false_claims( + self, tmp_path + ): + from unittest.mock import AsyncMock, patch + + from src.turn_state.durability import TurnDurability + + bot, fake, store = build_with_store([], tmp_path) + fake.responses.append(capacity_forever(fake)) + with patch.object( + TurnDurability, "suspend", AsyncMock(return_value=False) + ): + text, _, is_error, *_ = await run_loop(bot, FakeMessage("go")) + assert is_error is True + assert text.startswith("LLM API error:") + assert "preserved" not in text # no false preservation claims + (status,) = store._conn.execute("SELECT status FROM turns").fetchone() + assert status == TurnStatus.TERMINAL_FAILED + + +class TestMalformedBatch: + async def test_duplicate_call_ids_bounce_without_execution(self, tmp_path): + from unittest.mock import AsyncMock + + from src.llm.types import LLMResponse, ToolCall + + dup_batch = LLMResponse( + text="", + tool_calls=[ + ToolCall(id="dup", name="run_command", input={"command": "a"}), + ToolCall(id="dup", name="run_command", input={"command": "b"}), + ], + stop_reason="tool_use", + ) + bot, fake, store = build_with_store( + [dup_batch, text_response("recovered with fresh ids")], tmp_path + ) + bot.tool_executor.execute = AsyncMock() + text, *_ = await run_loop(bot, FakeMessage("go")) + assert text == "recovered with fresh ids" + bot.tool_executor.execute.assert_not_awaited() # failed BEFORE execution + bounce = [ + b + for m in fake.calls[-1]["messages"] + if isinstance(m.get("content"), list) + for b in m["content"] + if isinstance(b, dict) and b.get("type") == "tool_result" + ] + assert bounce and all("NOT executed" in b["content"] for b in bounce) + + +class TestHousekeepingSweep: + def test_ttl_sweep_runs_and_swallows_failures(self, tmp_path): + import time as _time + from types import SimpleNamespace + from unittest.mock import MagicMock + + from src.discord.housekeeping import Housekeeping + from src.turn_state import TurnKey + + store = TurnStateStore(tmp_path / "hk" / "turns.sqlite3") + lease = store.admit_turn_sync( + TurnKey("discord", "c1", "m1"), + guild_id=None, user_id="u", content_digest="d", code_version="t", + prompt_policy_hash="p", tool_catalog_hash="t", + session_snapshot=None, + ) + store.suspend_sync(lease, {"p": 1}) + store._conn.execute( + "UPDATE turns SET last_progress_at=?", [_time.time() - 25 * 3600] + ) + store._conn.commit() + + cfg = SimpleNamespace( + turn_state=SimpleNamespace( + resume_ttl_hours=24.0, + payload_retention_days=7.0, + ledger_retention_days=90.0, + ), + attachments=None, + ) + hk = Housekeeping( + get_config=lambda: cfg, + sessions=MagicMock(ids=lambda: []), + channel_state=MagicMock(), + prompt_builder=MagicMock(), + agent_manager=None, + loop_manager=MagicMock(), + loop_agent_bridge=None, + channel_logger=None, + fts_index=None, + turn_store=store, + ) + hk.cleanup_stale() + row = store._conn.execute("SELECT status FROM turns").fetchone() + assert row[0] == TurnStatus.TERMINAL_EXPIRED + + # A raising store must never break housekeeping. + store.ttl_sweep_sync = MagicMock(side_effect=RuntimeError("boom")) + hk.cleanup_stale() + store.close() From ed4f7d2cc8e374bad3ac8cd63522c9c3bd0ca8aa Mon Sep 17 00:00:00 2001 From: ATCharpentier Date: Thu, 30 Jul 2026 15:28:41 -0400 Subject: [PATCH 10/18] =?UTF-8?q?fix:=20review=20round=201=20=E2=80=94=20a?= =?UTF-8?q?ll=20eight=20blockers=20(PR=20#242)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit B1 (ledger fencing): every operation transition is now ONE atomic statement conditional on generation + EXPECTED REVISION + lease token + ACTIVE status + live lease (server-side EXISTS predicate) AND a legal prior op state — APPLIED can never re-enter RUNNING, a stale-revision owner mutates nothing, and intent recording uses plain fenced INSERTs (never OR REPLACE; an id collision raises LedgerIntentError and resets nothing). _verify_lease and its TOCTOU window are gone. B2 (redelivery fail-open): admission returns a disposition; an existing identity REFUSES fresh execution — terminal rows answer 'already processed', live-lease rows 'in progress', suspended (or expired-ACTIVE, swept here) rows point at resume. Existing identity never means 'run without durability'; only store-unavailable does. B3 (cancellability): the in-flight provider attempt is raced against the cancel event — /stop cancels and awaits the attempt task itself, then releases any held probe. The recovery budget still deliberately bounds WAITING, not a healthy in-flight stream (the v3.58.2 wall killed healthy >10-min xhigh generations; client transport limits own the attempt) — documented in the module contract. B4 (stranded ACTIVE): the boot sweep now suspends ALL ACTIVE rows (the store is single-process — nothing can legitimately be in flight at construction), and housekeeping runs a periodic expired-ACTIVE defense sweep. Fast-restart-inside-the-lease reproduced and pinned. B5 (resume stranding): reconstruction (payload restore, tool derivation, transcript repair, cancel check) runs BEFORE the execution lease is acquired; the residual post-acquire window releases the fenced lease back to SUSPENDED (release_acquired_sync). A corrupt payload self-heals to TERMINAL_REJECTED inside load_resumable instead of raising. B6 (auto-resume races): the advance check anchors on the session length captured synchronously AT SUSPENSION (only the suspending turn's own +2 intake bookkeeping is legitimate growth; no sampling window exists), and waiter cleanup is ownership-sensitive (compare-and-pop by task identity). B7 (typed pool exhaustion): CodexAuthPool raises LLMRateLimitError (all rate-limited) / LLMAuthError (all failed / no credentials) at the source — fast-fail at the recovery layer, never an unclassified defect; messages unchanged, RuntimeError-compatible. B8 (storage): turn-state dirs 0700, DB + WAL/SHM + blobs 0600 (secure tmp+rename writes); assistant tool_use arguments are secret-scrubbed at snapshot time (audit-storage parity; results were already scrubbed at source, credential-bearing user messages never reach a turn). CI: executor.py baseline aligned to CI's measurement (env-variant lines); new housekeeping sweep branches covered. Suite 7,886/5, lint/type gates new=0, coverage findings=0 at 88.3%. --- coverage-baseline.json | 8 +- src/discord/housekeeping.py | 13 +- src/discord/tool_loop.py | 27 ++ src/discord/turn_resume.py | 94 ++++--- src/llm/codex_auth.py | 23 +- src/llm/recovery.py | 43 ++- src/turn_state/codec.py | 34 ++- src/turn_state/durability.py | 22 +- src/turn_state/store.py | 324 +++++++++++++++++----- tests/test_codex_auth_rotation.py | 44 +++ tests/test_recovery_policy.py | 59 ++++ tests/test_resume_admission.py | 49 ++++ tests/test_turn_checkpoint_codec.py | 21 ++ tests/test_turn_state_store.py | 181 +++++++++++- tests/test_write_invariant_integration.py | 68 ++++- 15 files changed, 868 insertions(+), 142 deletions(-) diff --git a/coverage-baseline.json b/coverage-baseline.json index 8d43c6bf..7d33c74d 100644 --- a/coverage-baseline.json +++ b/coverage-baseline.json @@ -942,9 +942,9 @@ "statements": 231 }, "src/tools/executor.py": { - "covered": 402, - "missing": 25, - "percent": 94.15, + "covered": 400, + "missing": 27, + "percent": 93.68, "statements": 427 }, "src/tools/git_ops.py": { @@ -1313,4 +1313,4 @@ "percent": 88.98, "statements": 245 } -} +} \ No newline at end of file diff --git a/src/discord/housekeeping.py b/src/discord/housekeeping.py index 6dcf9698..f3368a81 100644 --- a/src/discord/housekeeping.py +++ b/src/discord/housekeeping.py @@ -97,8 +97,19 @@ def cleanup_stale(self) -> None: pass # Turn-state retention: the three clocks (resumable TTL, diagnostic - # payload compaction, ledger expiry — OUTCOME_UNKNOWN never expires). + # payload compaction, ledger expiry — OUTCOME_UNKNOWN never expires) + # plus the expired-ACTIVE defense sweep (a dead owner's turn must + # become visible to the resume path even if the boot sweep missed it). if self._turn_store is not None: + try: + swept_active = self._turn_store.sweep_expired_active_sync() + if any(swept_active.values()): + log.warning( + "Expired-active turn sweep: %s (dead-owner turns " + "suspended)", swept_active, + ) + except Exception: + pass try: config = self._get_config() ts = getattr(config, "turn_state", None) diff --git a/src/discord/tool_loop.py b/src/discord/tool_loop.py index d80ed21e..f08ed14f 100644 --- a/src/discord/tool_loop.py +++ b/src/discord/tool_loop.py @@ -434,6 +434,33 @@ async def run( st = await self._prepare_chat_turn( message, history, system_prompt_override, trace, policy ) + # Admission refusal: this message identity already has durable state + # (terminal / in-flight / suspended). A redelivered or duplicate + # message must never re-run its effects unledgered. + if st.durability.blocked is not None: + self._clear_active(st) + notices = { + "already_processed": ( + "This exact request was already processed — refusing to " + "run it again. Send it as a new message if you want a " + "fresh run." + ), + "in_progress": ( + "This exact request is already being processed elsewhere." + ), + "resumable": ( + "This request has preserved, resumable work — say " + "`resume` to continue it instead of starting over." + ), + } + text = notices.get( + st.durability.blocked, "This request cannot be re-run." + ) + log.warning( + "Turn admission refused (%s) for message %s in channel %s", + st.durability.blocked, st._trajectory.message_id, st._ch_id, + ) + return (text, False, False, [], False) return await self._run_with_guards(st) async def run_resumed(self, st: _ChatTurn) -> tuple[str, bool, bool, list[str], bool]: diff --git a/src/discord/turn_resume.py b/src/discord/turn_resume.py index 5f277e6e..fc230077 100644 --- a/src/discord/turn_resume.py +++ b/src/discord/turn_resume.py @@ -101,23 +101,39 @@ def on_turn_suspended(self, key: TurnKey, generation: str) -> None: """Called by the tool loop when a turn suspends. In-process only.""" if not self._auto_resume_enabled: return + # Session length AT SUSPENSION, captured synchronously inside the + # suspending turn (which still holds the channel lock) — the + # advance-check anchor (review blocker #6a, PR #242). + suspend_len = self._session_len(key.channel_id) existing = self._waiters.pop(key, None) if existing is not None: existing.cancel() task = asyncio.get_running_loop().create_task( - self._auto_resume_waiter(key, generation), + self._auto_resume_waiter(key, generation, suspend_len), name=f"turn-resume:{key.channel_id}:{key.message_id}", ) self._waiters[key] = task - task.add_done_callback(lambda _t: self._waiters.pop(key, None)) - async def _auto_resume_waiter(self, key: TurnKey, generation: str) -> None: + def _cleanup(t: asyncio.Task, *, _key=key) -> None: + # Ownership-sensitive: a cancelled predecessor must never pop + # its successor's registry entry (review blocker #6b, PR #242). + if self._waiters.get(_key) is t: + self._waiters.pop(_key, None) + + task.add_done_callback(_cleanup) + + async def _auto_resume_waiter( + self, key: TurnKey, generation: str, suspend_len: int + ) -> None: """Wait for capacity, then resume IF nothing else happened. - The session baseline is sampled at the FIRST poll (after the - suspending turn's intake bookkeeping settled — the channel lock - serializes us behind it); any later growth means the session - advanced and auto-resume stands down (explicit-only from there). + The advance check anchors on the session length captured AT + SUSPENSION (``suspend_len``): the only legitimate growth afterwards + is the suspending turn's own intake bookkeeping (exactly +2: the + user message and the preserved-marker; a directly-driven turn adds + 0). Anything else — or an unreadable session — stands auto-resume + down fail-safe; explicit resume stays available. No sampling window + exists for a stranger message to sneak into the baseline. Capacity detection is ACTIVE: a quiet breaker is never probed by anyone, so the waiter claims the probe slot itself when the cooldown @@ -129,10 +145,9 @@ async def _auto_resume_waiter(self, key: TurnKey, generation: str) -> None: """ give_up_at = time.monotonic() + self._resume_ttl_hours * 3600.0 breaker = self._llm_gateway.capacity_breaker_for() - await asyncio.sleep(_AUTO_POLL_SECONDS) - async with self._channel_lock(key.channel_id): - baseline = self._session_len(key.channel_id) + allowed = {suspend_len, suspend_len + 2} if suspend_len >= 0 else set() while time.monotonic() < give_up_at: + await asyncio.sleep(_AUTO_POLL_SECONDS) row = await asyncio.to_thread(self._store.load_resumable_sync, key) if row is None or row["generation"] != generation: return # resumed elsewhere / rejected / expired @@ -140,9 +155,7 @@ async def _auto_resume_waiter(self, key: TurnKey, generation: str) -> None: if not isinstance(admission, float): breaker.abandon(admission) # the resume re-acquires for real current = self._session_len(key.channel_id) - if baseline < 0 or current < 0 or current != baseline: - # Advanced OR unreadable — either way auto-resume must - # stand down (fail-safe); explicit resume stays available. + if current not in allowed: log.info( "Auto-resume for %s stands down: session advanced or " "unreadable (explicit resume still available)", key, @@ -150,7 +163,6 @@ async def _auto_resume_waiter(self, key: TurnKey, generation: str) -> None: return await self._run_auto_resume(key, row) return - await asyncio.sleep(min(_AUTO_POLL_SECONDS, float(admission))) log.info("Auto-resume waiter for %s expired", key) def _channel_lock(self, channel_id: str) -> asyncio.Lock: @@ -285,12 +297,11 @@ async def _validate_and_rebuild(self, key: TurnKey, row: dict): ) return None, None, "the original message was edited" - lease = await asyncio.to_thread( - self._store.acquire_resume_lease_sync, key, row["generation"] - ) - if lease is None: - return None, None, "someone else is already resuming it" - + # RECONSTRUCT BEFORE ACQUIRING (review blocker #5, PR #242): every + # fallible step — payload restore, tool derivation, transcript + # repair, cancellation check — runs while the row is still + # SUSPENDED, so a failure rejects/aborts cleanly instead of + # stranding an ACTIVE row invisible to resumable queries. payload = row["payload"] try: fields = restore_field_values( @@ -325,21 +336,36 @@ async def _validate_and_rebuild(self, key: TurnKey, row: dict): if deadline_utc: remaining_budget = max(0.0, float(deadline_utc) - time.time()) - durability = TurnDurability.resumed( - self._store, - lease, - payload.get("generation_seq", 0), - first_generation_budget=remaining_budget, - ) - st = _ChatTurn( - message=original, - policy=CHAT_POLICY, - trace=None, # the old segment was closed into the payload - tools=tools, - _cancel=cancel, - durability=durability, - **fields, + # Acquire LAST — the single-winner transition happens only once + # everything else is ready to run. + lease = await asyncio.to_thread( + self._store.acquire_resume_lease_sync, key, row["generation"] ) + if lease is None: + return None, None, "someone else is already resuming it" + + try: + durability = TurnDurability.resumed( + self._store, + lease, + payload.get("generation_seq", 0), + first_generation_budget=remaining_budget, + ) + st = _ChatTurn( + message=original, + policy=CHAT_POLICY, + trace=None, # the old segment was closed into the payload + tools=tools, + _cancel=cancel, + durability=durability, + **fields, + ) + except Exception: + # Residual post-acquire window: release the fenced lease back to + # SUSPENDED so the turn never strands ACTIVE. + log.exception("Post-acquire turn construction failed — releasing") + await asyncio.to_thread(self._store.release_acquired_sync, lease) + return None, None, "the turn could not be reconstructed" return st, original, None @staticmethod diff --git a/src/llm/codex_auth.py b/src/llm/codex_auth.py index 76aab749..96fdf94d 100644 --- a/src/llm/codex_auth.py +++ b/src/llm/codex_auth.py @@ -11,6 +11,7 @@ import aiohttp from ..odin_log import get_logger +from .errors import LLMAuthError, LLMRateLimitError log = get_logger("codex_auth") @@ -509,7 +510,7 @@ def account_count(self) -> int: @property def current(self) -> CodexAuth: if not self._accounts: - raise RuntimeError("No Codex credentials configured.") + raise LLMAuthError("No Codex credentials configured.", provider="codex") return self._accounts[self._current_index] async def acquire(self) -> tuple[str, str | None, int]: @@ -524,12 +525,12 @@ async def acquire(self) -> tuple[str, str | None, int]: refresh-token reuse. """ if not self._accounts: - raise RuntimeError("No Codex credentials configured.") + raise LLMAuthError("No Codex credentials configured.", provider="codex") errors: list[tuple[int, str]] = [] for _ in range(len(self._accounts)): async with self._pool_lock: if not self._accounts: - raise RuntimeError("No Codex credentials configured.") + raise LLMAuthError("No Codex credentials configured.", provider="codex") self._current_index %= len(self._accounts) idx = self._current_index auth = self._accounts[idx] @@ -551,14 +552,22 @@ async def acquire(self) -> tuple[str, str | None, int]: self._rotate() continue return token, auth.get_account_id(), idx + # Pool exhaustion is part of the typed taxonomy (PR #242 review + # blocker #7): every rotation avenue is spent by the time these + # raise, so the shared recovery must FAST-FAIL them — never treat + # them as unclassified defects, and the subsystem guard must not + # count quota exhaustion as generic subsystem failure. Both types + # subclass RuntimeError, so legacy handlers are unaffected. if errors: - raise RuntimeError( + raise LLMAuthError( f"All {len(self._accounts)} Codex accounts failed: " - + "; ".join(f"#{i}: {err}" for i, err in errors) + + "; ".join(f"#{i}: {err}" for i, err in errors), + provider="codex", ) - raise RuntimeError( + raise LLMRateLimitError( f"All {len(self._accounts)} Codex accounts are rate-limited or " - "backing off; retry shortly." + "backing off; retry shortly.", + provider="codex", ) async def get_access_token(self) -> str: diff --git a/src/llm/recovery.py b/src/llm/recovery.py index 548fbd28..1f7e471c 100644 --- a/src/llm/recovery.py +++ b/src/llm/recovery.py @@ -11,6 +11,13 @@ logical generation. Callers resuming a persisted turn pass the remaining budget (computed from the checkpoint's UTC deadline) via ``deadline_seconds`` — a restart must not grant a fresh five minutes. + The budget bounds WAITING between attempts, deliberately NOT an in-flight + attempt: healthy xhigh generations legitimately stream >10 minutes, and + killing them at a recovery wall is exactly the v3.58.2 disease + (hardcoded 600s total killing healthy generations). In-flight attempts + are bounded by the client's own transport limits (request_timeout / + stream_stall_timeout) and are cancellable via ``cancel_event`` at any + moment (/stop interrupts the await itself, not just the sleeps). - **Retryable**: ``LLMCapacityError``, ``LLMTransportError``, and ``CircuitOpenError`` (the client breaker — recovery waits THROUGH it rather than treating it as terminal). Full-jitter backoff capped at @@ -35,6 +42,7 @@ from __future__ import annotations import asyncio +import contextlib import time from collections.abc import Awaitable, Callable from dataclasses import dataclass @@ -93,6 +101,37 @@ async def _sleep_cancellable(seconds: float, cancel_event: asyncio.Event | None) raise asyncio.CancelledError("cancelled during LLM recovery wait") +async def _attempt_cancellable( + attempt: Callable[[], Awaitable[T]], cancel_event: asyncio.Event | None +) -> T: + """Run one attempt racing the caller's cancel event. + + /stop must interrupt an IN-FLIGHT provider await, not only the waits + between attempts (review blocker #3, PR #242). On cancellation the + attempt task is cancelled and awaited — aiohttp unwinds its transport + cleanly — before CancelledError propagates to the caller (which + releases any held breaker probe). + """ + if cancel_event is None: + return await attempt() + attempt_task = asyncio.ensure_future(attempt()) + cancel_task = asyncio.ensure_future(cancel_event.wait()) + try: + done, _ = await asyncio.wait( + {attempt_task, cancel_task}, return_when=asyncio.FIRST_COMPLETED + ) + if attempt_task in done: + return attempt_task.result() + attempt_task.cancel() + with contextlib.suppress(BaseException): + await attempt_task + raise asyncio.CancelledError("cancelled during LLM attempt") + finally: + cancel_task.cancel() + with contextlib.suppress(BaseException): + await cancel_task + + def _notify_wait( on_wait: Callable[[float, float, BaseException], None] | None, wait: float, @@ -171,9 +210,9 @@ def _exhausted() -> BaseException: admission = breaker.acquire_attempt() token = admission - # -- one attempt --------------------------------------------- + # -- one attempt (raced against /stop) ------------------------ try: - result = await attempt() + result = await _attempt_cancellable(attempt, cancel_event) except asyncio.CancelledError: if breaker is not None and token is not None: breaker.abandon(token) diff --git a/src/turn_state/codec.py b/src/turn_state/codec.py index 7ff83ea9..5d237ab7 100644 --- a/src/turn_state/codec.py +++ b/src/turn_state/codec.py @@ -24,6 +24,7 @@ from dataclasses import asdict from typing import Any +from ..llm.secret_scrubber import scrub_output_secrets from ..odin_log import get_logger log = get_logger("turn_state") @@ -96,6 +97,35 @@ def compute_content_digest(text: str) -> str: return hashlib.sha256((text or "").encode("utf-8", "replace")).hexdigest() +# ── storage redaction ──────────────────────────────────────────────── + + +def _scrub_tool_use_inputs(obj: Any) -> Any: + """Secret-scrub string values inside assistant ``tool_use`` inputs + before they hit durable storage (review blocker #8, PR #242) — the + parity move with audit storage, which deliberately scrubs tool inputs. + + Tool RESULTS are already scrubbed at source (`_run_one_tool` runs + scrub_output_secrets before building the result block), and + credential-bearing USER messages are deleted by the intake secret gate + before a turn ever starts — tool arguments were the remaining + unscrubbed surface. The scrub applies at SNAPSHOT time only, so a + resumed transcript shows the model its own arguments with any embedded + secrets masked; the executed effect already happened and is unaffected. + """ + if isinstance(obj, list): + return [_scrub_tool_use_inputs(x) for x in obj] + if isinstance(obj, dict): + if obj.get("type") == "tool_use" and isinstance(obj.get("input"), dict): + scrubbed = { + k: scrub_output_secrets(v) if isinstance(v, str) else v + for k, v in obj["input"].items() + } + return {**obj, "input": scrubbed} + return {k: _scrub_tool_use_inputs(v) for k, v in obj.items()} + return obj + + # ── image-block externalization ────────────────────────────────────── @@ -250,7 +280,9 @@ def snapshot_chat_turn(st, *, store_blob, generation_seq: int, extra: dict | Non "generation_seq": generation_seq, "fields": { "system_prompt": st.system_prompt, - "messages": _externalize_blocks(st.messages, store_blob), + "messages": _scrub_tool_use_inputs( + _externalize_blocks(st.messages, store_blob) + ), "user_id": st.user_id, "chat_cap": st.chat_cap, "iteration": st.iteration, diff --git a/src/turn_state/durability.py b/src/turn_state/durability.py index 014190d1..74bd4c61 100644 --- a/src/turn_state/durability.py +++ b/src/turn_state/durability.py @@ -70,6 +70,10 @@ def __init__(self, store: TurnStateStore | None, lease: TurnLease | None) -> Non self.cancelled = False # One-shot remaining budget for a resumed generation (see resumed()). self._resume_budget: float | None = None + # Admission refusal disposition (review blocker #2, PR #242): + # "already_processed" / "in_progress" / "resumable" — the loop must + # REFUSE fresh execution, never run unledgered. None = no refusal. + self.blocked: str | None = None # -- construction -------------------------------------------------- @@ -87,7 +91,13 @@ async def admit( tools: list | None, session_snapshot: dict | None, ) -> TurnDurability: - """Admit a fresh Discord chat turn; any refusal → disabled handle.""" + """Admit a fresh Discord chat turn. + + Store-unavailable (feature off / I/O failure) → a disabled handle + (legacy run). An EXISTING identity is different: the handle comes + back with ``blocked`` set and the loop must refuse fresh execution — + a redelivered message must never re-run its effects unledgered. + """ if store is None or not store.available: return cls.disabled() try: @@ -97,7 +107,7 @@ async def admit( message_id=str(message.id), ) tool_names = sorted(t.get("name", "") for t in (tools or [])) - lease = await asyncio.to_thread( + lease, disposition = await asyncio.to_thread( store.admit_turn_sync, key, guild_id=str(getattr(getattr(message, "guild", None), "id", "") or ""), @@ -113,9 +123,13 @@ async def admit( except Exception: log.exception("Turn admission failed — running without durability") return cls.disabled() - if lease is None: + if lease is not None: + return cls(store, lease) + if disposition == "store_unavailable": return cls.disabled() - return cls(store, lease) + handle = cls.disabled() + handle.blocked = disposition + return handle @classmethod def resumed( diff --git a/src/turn_state/store.py b/src/turn_state/store.py index c93ebc74..a42e2d4e 100644 --- a/src/turn_state/store.py +++ b/src/turn_state/store.py @@ -203,13 +203,20 @@ def __init__( self._write_lock = threading.Lock() self._conn: sqlite3.Connection | None = None try: + # Checkpoints carry the model transcript (user content, tool + # arguments) — secret-adjacent material. Everything here is + # owner-only: directories 0700, DB + WAL/SHM + blobs 0600 + # (review blocker #8, PR #242). Path(self.db_path).parent.mkdir(parents=True, exist_ok=True) self._blob_dir.mkdir(parents=True, exist_ok=True) + os.chmod(Path(self.db_path).parent, 0o700) + os.chmod(self._blob_dir, 0o700) conn = sqlite3.connect(self.db_path, check_same_thread=False) conn.execute("PRAGMA journal_mode=WAL") conn.execute("PRAGMA busy_timeout=30000") conn.executescript(_DDL) conn.commit() + self._restrict_db_modes() self._conn = conn swept = self._boot_sweep_sync() if swept["turns"] or swept["ops"]: @@ -225,6 +232,17 @@ def __init__( ) self._conn = None + def _restrict_db_modes(self) -> None: + """0600 on the database and its WAL/SHM sidecars (best-effort — the + sidecars appear lazily; called again from the TTL sweep).""" + for suffix in ("", "-wal", "-shm"): + path = Path(self.db_path + suffix) + try: + if path.exists(): + os.chmod(path, 0o600) + except OSError: + pass + @property def available(self) -> bool: return self._conn is not None @@ -280,20 +298,24 @@ def _op_where(self, lease: TurnLease) -> tuple[str, list]: lease.generation], ) - def _verify_lease(self, lease: TurnLease) -> None: - """Ops-table writes are fenced indirectly: verify the turns row still - carries this lease before touching operations.""" - conn = self._require() - try: - row = conn.execute( - "SELECT lease_token, turn_generation FROM turns " - "WHERE source=? AND channel_id=? AND message_id=?", - [lease.key.source, lease.key.channel_id, lease.key.message_id], - ).fetchone() - except sqlite3.Error as exc: - raise TurnStateUnavailableError(f"turn-state read failed: {exc}") from exc - if row is None or row[0] != lease.token or row[1] != lease.generation: - raise StaleTurnError(f"turn {lease.key}: lease no longer held") + # The full turn fence as a server-side predicate: every ledger write is + # ONE statement conditional on the turn's generation, EXPECTED REVISION, + # current lease token, ACTIVE status, and a live lease — plus a legal + # prior operation state on the row itself. No read-then-write TOCTOU; + # a stale-revision or fenced-out owner mutates nothing (review blocker + # #1, PR #242). + _TURN_FENCE = ( + "EXISTS (SELECT 1 FROM turns WHERE turns.source=? AND turns.channel_id=? " + "AND turns.message_id=? AND turns.turn_generation=? AND turns.revision=? " + "AND turns.lease_token=? AND turns.status=? AND turns.lease_expires_at > ?)" + ) + + def _fence_params(self, lease: TurnLease) -> list: + return [ + lease.key.source, lease.key.channel_id, lease.key.message_id, + lease.generation, lease.revision, lease.token, TurnStatus.ACTIVE, + time.time(), + ] # ── admission / lifecycle (sync bodies) ────────────────────────── @@ -308,12 +330,19 @@ def admit_turn_sync( prompt_policy_hash: str | None, tool_catalog_hash: str | None, session_snapshot: dict | None, - ) -> TurnLease | None: - """Admit a fresh turn. Returns a lease, or None when this message - already has a row (redelivery/crash artifact — caller runs legacy, - loudly) or when the store is unavailable (caller runs legacy).""" + ) -> tuple[TurnLease | None, str]: + """Admit a fresh turn. Returns (lease, disposition). + + Dispositions: ``admitted`` (lease non-None); ``already_processed`` + (a terminal row exists — a redelivered message must REFUSE fresh + execution, never run unledgered — review blocker #2, PR #242); + ``in_progress`` (an ACTIVE row with a live lease — another owner); + ``resumable`` (a SUSPENDED row, or an expired-lease ACTIVE row + swept to SUSPENDED here — the resume path owns it); + ``store_unavailable`` (feature off / I/O failure — legacy run). + """ if self._conn is None: - return None + return None, "store_unavailable" now = time.time() generation = secrets.token_hex(16) token = secrets.token_hex(16) @@ -335,15 +364,51 @@ def admit_turn_sync( ) self._conn.commit() except sqlite3.IntegrityError: - log.warning( - "Turn %s already has a state row — running without " - "durability for this turn", key, - ) - return None + return None, self._classify_existing_row(key, now) except sqlite3.Error: log.exception("Turn admission failed — running without durability") - return None - return TurnLease(key=key, generation=generation, token=token, revision=0) + return None, "store_unavailable" + return ( + TurnLease(key=key, generation=generation, token=token, revision=0), + "admitted", + ) + + def _classify_existing_row(self, key: TurnKey, now: float) -> str: + """Disposition for a message whose identity already has a row. + Lock held by caller.""" + assert self._conn is not None + row = self._conn.execute( + "SELECT status, lease_expires_at FROM turns " + "WHERE source=? AND channel_id=? AND message_id=?", + [key.source, key.channel_id, key.message_id], + ).fetchone() + if row is None: # deleted between INSERT failure and here — treat as busy + return "in_progress" + status, lease_expires_at = row + if status in TurnStatus.TERMINAL: + log.warning("Turn %s was already processed — refusing re-execution", key) + return "already_processed" + if status == TurnStatus.SUSPENDED: + return "resumable" + # ACTIVE: a live lease means another owner; an expired lease is a + # crash artifact — suspend it here so the resume path owns it. + if lease_expires_at is not None and lease_expires_at > now: + return "in_progress" + self._conn.execute( + "UPDATE turns SET status=?, suspended_at=?, lease_token=NULL, " + "lease_expires_at=NULL WHERE source=? AND channel_id=? AND message_id=? " + "AND status=?", + [TurnStatus.SUSPENDED, now, key.source, key.channel_id, + key.message_id, TurnStatus.ACTIVE], + ) + self._conn.execute( + "UPDATE operations SET state=?, updated_at=? " + "WHERE source=? AND channel_id=? AND message_id=? AND state IN (?, ?)", + [OpState.OUTCOME_UNKNOWN, now, key.source, key.channel_id, + key.message_id, OpState.PREPARED, OpState.RUNNING], + ) + self._conn.commit() + return "resumable" def checkpoint_sync( self, @@ -405,8 +470,12 @@ def record_intents_sync( ) -> None: """Persist PREPARED rows for one generation's tool calls. - Validates Odin's rule up front: tool-call ids nonempty and unique - within the generation — malformed duplicates fail BEFORE execution. + Tool-call ids must be nonempty and unique within the generation — + malformed duplicates fail BEFORE execution. Plain INSERTs (never + OR REPLACE — an existing row, e.g. a settled op from a replayed id, + must never be reset to PREPARED); each insert is fenced on the turn + row via INSERT..SELECT..WHERE EXISTS, so a fenced-out owner inserts + nothing. All rows land in one transaction or none do. """ ids = [str(i.get("tool_call_id") or "") for i in intents] if any(not i for i in ids): @@ -414,29 +483,50 @@ def record_intents_sync( if len(set(ids)) != len(ids): raise LedgerIntentError("duplicate tool_call_id in generation") conn = self._require() - self._verify_lease(lease) now = time.time() - rows = [ - (lease.key.source, lease.key.channel_id, lease.key.message_id, - lease.generation, generation_seq, str(i["tool_call_id"]), - OpState.PREPARED, str(i.get("tool_name") or "unknown"), iteration, - effect_fingerprint(str(i.get("tool_name") or ""), i.get("tool_input") or {}), - None, now, now) - for i in intents - ] try: with self._write_lock: - conn.executemany( - "INSERT OR REPLACE INTO operations VALUES " - "(?,?,?,?,?,?,?,?,?,?,?,?,?)", - rows, - ) - conn.commit() + try: + for intent in intents: + cur = conn.execute( + "INSERT INTO operations " + "SELECT ?,?,?,?,?,?,?,?,?,?,?,?,? " + f"WHERE {self._TURN_FENCE}", + [lease.key.source, lease.key.channel_id, + lease.key.message_id, lease.generation, + generation_seq, str(intent["tool_call_id"]), + OpState.PREPARED, + str(intent.get("tool_name") or "unknown"), iteration, + effect_fingerprint( + str(intent.get("tool_name") or ""), + intent.get("tool_input") or {}, + ), + None, now, now, + *self._fence_params(lease)], + ) + if cur.rowcount != 1: + raise StaleTurnError( + f"turn {lease.key}: fence lost recording intents" + ) + conn.commit() + except BaseException: + conn.rollback() + raise + except sqlite3.IntegrityError as exc: + raise LedgerIntentError( + f"intent already recorded for this generation: {exc}" + ) from exc + except StaleTurnError: + raise except sqlite3.Error as exc: raise TurnStateUnavailableError(f"ledger intent write failed: {exc}") from exc def mark_running_sync(self, lease: TurnLease, generation_seq: int, tool_call_id: str) -> None: - self._set_op_state(lease, generation_seq, tool_call_id, OpState.RUNNING, None) + # Legal source: PREPARED only — a settled op can never re-enter RUNNING. + self._set_op_state( + lease, generation_seq, tool_call_id, OpState.RUNNING, None, + legal_sources=(OpState.PREPARED,), + ) def settle_op_sync( self, @@ -453,7 +543,10 @@ def settle_op_sync( # rule: ambiguous outcomes stop, never rerun). if state not in (OpState.APPLIED, OpState.DEFINITELY_FAILED, OpState.OUTCOME_UNKNOWN): raise ValueError(f"settle_op_sync: invalid terminal state {state}") - self._set_op_state(lease, generation_seq, tool_call_id, state, result_text) + self._set_op_state( + lease, generation_seq, tool_call_id, state, result_text, + legal_sources=(OpState.PREPARED, OpState.RUNNING), + ) def settle_interrupted_sync( self, @@ -469,21 +562,12 @@ def settle_interrupted_sync( boundary: it never downgrades an already-settled row, and a missing row (e.g. a parse-error call that had no intent) is tolerated. """ - conn = self._require() - self._verify_lease(lease) - where, params = self._op_where(lease) - try: - with self._write_lock: - conn.execute( - f"UPDATE operations SET state=?, result=?, updated_at=? " - f"WHERE {where} AND generation_seq=? AND tool_call_id=? " - "AND state IN (?, ?)", - [OpState.OUTCOME_UNKNOWN, result_text, time.time(), *params, - generation_seq, tool_call_id, OpState.PREPARED, OpState.RUNNING], - ) - conn.commit() - except sqlite3.Error as exc: - raise TurnStateUnavailableError(f"ledger write failed: {exc}") from exc + self._set_op_state( + lease, generation_seq, tool_call_id, OpState.OUTCOME_UNKNOWN, + result_text, + legal_sources=(OpState.PREPARED, OpState.RUNNING), + tolerate_missing=True, + ) def _set_op_state( self, @@ -492,31 +576,45 @@ def _set_op_state( tool_call_id: str, state: str, result_text: str | None, + *, + legal_sources: tuple[str, ...], + tolerate_missing: bool = False, ) -> None: + """One atomic, fully-fenced operation transition. + + A single UPDATE conditional on the op's legal prior state AND the + turn fence (generation, expected revision, lease token, ACTIVE + status, live lease) — no separate verify step, no TOCTOU window, + and an illegal transition (e.g. APPLIED→RUNNING) matches nothing. + """ conn = self._require() - self._verify_lease(lease) where, params = self._op_where(lease) + placeholders = ", ".join("?" for _ in legal_sources) try: with self._write_lock: if result_text is None: cur = conn.execute( f"UPDATE operations SET state=?, updated_at=? WHERE {where} " - "AND generation_seq=? AND tool_call_id=?", - [state, time.time(), *params, generation_seq, tool_call_id], + f"AND generation_seq=? AND tool_call_id=? " + f"AND state IN ({placeholders}) AND {self._TURN_FENCE}", + [state, time.time(), *params, generation_seq, tool_call_id, + *legal_sources, *self._fence_params(lease)], ) else: cur = conn.execute( f"UPDATE operations SET state=?, result=?, updated_at=? " - f"WHERE {where} AND generation_seq=? AND tool_call_id=?", + f"WHERE {where} AND generation_seq=? AND tool_call_id=? " + f"AND state IN ({placeholders}) AND {self._TURN_FENCE}", [state, result_text, time.time(), *params, generation_seq, - tool_call_id], + tool_call_id, *legal_sources, *self._fence_params(lease)], ) conn.commit() except sqlite3.Error as exc: raise TurnStateUnavailableError(f"ledger write failed: {exc}") from exc - if cur.rowcount != 1: + if cur.rowcount != 1 and not tolerate_missing: raise StaleTurnError( - f"op {tool_call_id} gen_seq {generation_seq}: no PREPARED row to update" + f"op {tool_call_id} gen_seq {generation_seq}: no legally-" + f"transitionable row under the current fence" ) # ── resume (sync bodies) ───────────────────────────────────────── @@ -535,6 +633,15 @@ def load_resumable_sync(self, key: TurnKey) -> dict | None: ).fetchone() if row is None or row[2] is None: return None + try: + payload = json.loads(row[2]) + snapshot = json.loads(row[10] or "{}") + except (json.JSONDecodeError, TypeError): + # An unreadable checkpoint can never be resumed — reject it + # terminally instead of raising into every caller. + log.exception("Corrupt checkpoint payload for %s — rejecting", key) + self.reject_resumable_sync(key, "checkpoint payload unreadable") + return None ops = self._conn.execute( "SELECT generation_seq, tool_call_id, state, tool_name, result " "FROM operations WHERE source=? AND channel_id=? AND message_id=? " @@ -544,7 +651,7 @@ def load_resumable_sync(self, key: TurnKey) -> dict | None: return { "generation": row[0], "revision": row[1], - "payload": json.loads(row[2]), + "payload": payload, "guild_id": row[3], "user_id": row[4], "content_digest": row[5], @@ -552,7 +659,7 @@ def load_resumable_sync(self, key: TurnKey) -> dict | None: "schema_version": row[7], "prompt_policy_hash": row[8], "tool_catalog_hash": row[9], - "session_snapshot": json.loads(row[10] or "{}"), + "session_snapshot": snapshot, "recovery_deadline_utc": row[11], "last_progress_at": row[12], "suspended_at": row[13], @@ -614,6 +721,38 @@ def reject_resumable_sync(self, key: TurnKey, reason: str) -> None: except sqlite3.Error: log.exception("reject_resumable failed (non-fatal)") + def release_acquired_sync( + self, lease: TurnLease, *, terminal_reason: str | None = None + ) -> None: + """Fenced abort for a lease acquired but never run (post-acquire + reconstruction failure — review blocker #5, PR #242). + + With ``terminal_reason`` the turn becomes TERMINAL_REJECTED (payload + compacted); without, it returns to SUSPENDED for a later attempt. + Fenced on the lease so only the acquiring owner can abort.""" + conn = self._require() + if terminal_reason is not None: + set_sql = "status=?, payload=NULL, lease_token=NULL, lease_expires_at=NULL" + first_param: list = [TurnStatus.TERMINAL_REJECTED] + else: + set_sql = "status=?, lease_token=NULL, lease_expires_at=NULL" + first_param = [TurnStatus.SUSPENDED] + try: + with self._write_lock: + cur = conn.execute( + f"UPDATE turns SET {set_sql} " + "WHERE source=? AND channel_id=? AND message_id=? " + "AND turn_generation=? AND lease_token=? AND status=?", + [*first_param, lease.key.source, lease.key.channel_id, + lease.key.message_id, lease.generation, lease.token, + TurnStatus.ACTIVE], + ) + conn.commit() + if cur.rowcount == 1 and terminal_reason: + log.info("Acquired turn %s rejected: %s", lease.key, terminal_reason) + except sqlite3.Error: + log.exception("release_acquired failed (non-fatal)") + def list_suspended_sync(self, source: str | None = None) -> list[dict]: if self._conn is None: return [] @@ -635,16 +774,40 @@ def list_suspended_sync(self, source: str | None = None) -> list[dict]: # ── sweeps ─────────────────────────────────────────────────────── def _boot_sweep_sync(self) -> dict: - """Crash recovery at construction: expired-lease ACTIVE turns become - SUSPENDED; their PREPARED/RUNNING ops become OUTCOME_UNKNOWN (we - cannot know whether the external effect happened — never rerun).""" + """Crash recovery at construction: ALL ACTIVE turns become SUSPENDED + — this store serves a single process, so at construction no turn can + legitimately be in flight, lease expiry notwithstanding (a fast + restart inside the lease TTL used to strand rows ACTIVE forever — + review blocker #4, PR #242). Their PREPARED/RUNNING ops become + OUTCOME_UNKNOWN (we cannot know whether the external effect + happened — never rerun).""" + return self._sweep_active_sync(only_expired=False) + + def sweep_expired_active_sync(self) -> dict: + """Defense-in-depth periodic sweep (housekeeping): any ACTIVE row + whose lease expired belongs to a dead owner — suspend it so the + resume path can see it.""" + if self._conn is None: + return {} + try: + return self._sweep_active_sync(only_expired=True) + except Exception: + log.exception("expired-active sweep failed (non-fatal)") + return {} + + def _sweep_active_sync(self, *, only_expired: bool) -> dict: conn = self._require() now = time.time() + where = "status=?" + params: list = [TurnStatus.ACTIVE] + if only_expired: + where += " AND (lease_expires_at IS NULL OR lease_expires_at < ?)" + params.append(now) with self._write_lock: stale = conn.execute( - "SELECT source, channel_id, message_id, turn_generation FROM turns " - "WHERE status=? AND (lease_expires_at IS NULL OR lease_expires_at < ?)", - [TurnStatus.ACTIVE, now], + f"SELECT source, channel_id, message_id, turn_generation FROM turns " + f"WHERE {where}", + params, ).fetchall() ops = 0 for source, channel_id, message_id, generation in stale: @@ -678,6 +841,7 @@ def ttl_sweep_sync( conn = self._conn now = time.time() expired = payloads = ledger = 0 + self._restrict_db_modes() # WAL/SHM sidecars appear lazily try: with self._write_lock: # Clock 1: resumable window — 24h from last REAL progress. @@ -716,13 +880,19 @@ def ttl_sweep_sync( # ── blobs ──────────────────────────────────────────────────────── def store_blob_sync(self, data: bytes) -> str: - """Content-addressed blob write (tmp+rename). Returns 'blob:'.""" + """Content-addressed blob write (tmp+rename, 0600 — the codex_auth + secure-write discipline). Returns 'blob:'.""" digest = hashlib.sha256(data).hexdigest() path = self._blob_dir / digest if not path.exists(): tmp = path.with_suffix(".tmp") try: - tmp.write_bytes(data) + fd = os.open(str(tmp), os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) + try: + os.write(fd, data) + os.fsync(fd) + finally: + os.close(fd) os.replace(tmp, path) except OSError as exc: raise TurnStateUnavailableError(f"blob write failed: {exc}") from exc diff --git a/tests/test_codex_auth_rotation.py b/tests/test_codex_auth_rotation.py index de353a07..537fe0dd 100644 --- a/tests/test_codex_auth_rotation.py +++ b/tests/test_codex_auth_rotation.py @@ -64,3 +64,47 @@ def test_mark_rate_limited_accepts_custom_window(): auth2 = CodexAuth.__new__(CodexAuth) auth2.mark_rate_limited() assert auth2.is_rate_limited() is True + + +async def test_pool_exhaustion_is_typed_rate_limit(): + """Review blocker #7 (PR #242): all-accounts-rate-limited must raise + LLMRateLimitError (fast-fail at the recovery layer, never an + unclassified defect), with the historical message preserved.""" + import pytest + + from src.llm.errors import LLMRateLimitError + + pool = _pool(2) + for acct in pool._accounts: + acct.is_rate_limited.return_value = True + with pytest.raises(LLMRateLimitError) as exc_info: + await pool.acquire() + assert "rate-limited or backing off" in str(exc_info.value) + assert exc_info.value.provider == "codex" + + +async def test_pool_exhaustion_all_failed_is_typed_auth(): + import pytest + + from src.llm.errors import LLMAuthError + + async def boom(): + raise RuntimeError("refresh dead") + + pool = _pool(2) + for acct in pool._accounts: + acct.is_rate_limited.return_value = False + acct.get_access_token.side_effect = boom + with pytest.raises(LLMAuthError) as exc_info: + await pool.acquire() + assert "accounts failed" in str(exc_info.value) + + +async def test_empty_pool_is_typed_auth(): + import pytest + + from src.llm.errors import LLMAuthError + + pool = _pool(0) + with pytest.raises(LLMAuthError): + await pool.acquire() diff --git a/tests/test_recovery_policy.py b/tests/test_recovery_policy.py index fa7d9fae..399a751f 100644 --- a/tests/test_recovery_policy.py +++ b/tests/test_recovery_policy.py @@ -201,3 +201,62 @@ def hook(wait, remaining, error): result = await generate_with_recovery(attempt, policy=FAST, on_wait=hook) assert result == "ok" assert seen and seen[0][2] == "LLMCapacityError" + + +async def test_cancel_interrupts_an_in_flight_attempt(): + """Review blocker #3 (PR #242): /stop must interrupt the provider await + itself, not just the waits between attempts. The in-flight attempt task + is cancelled and awaited before CancelledError propagates.""" + cancel = asyncio.Event() + attempt_cancelled = asyncio.Event() + + async def slow_attempt(): + try: + await asyncio.sleep(30) # a long healthy stream + except asyncio.CancelledError: + attempt_cancelled.set() + raise + return "never" + + async def fire_cancel(): + await asyncio.sleep(0.05) + cancel.set() + + started = time.monotonic() + canceller = asyncio.create_task(fire_cancel()) + with pytest.raises(asyncio.CancelledError): + await generate_with_recovery( + slow_attempt, + policy=RecoveryPolicy(deadline_seconds=60.0), + cancel_event=cancel, + ) + await canceller + assert time.monotonic() - started < 1.0 # did not wait out the stream + assert attempt_cancelled.is_set() # the attempt itself was unwound + + +async def test_probe_released_when_cancel_interrupts_attempt(): + registry = ModelBreakerRegistry(cooldown_base=0.01) + breaker = registry.for_model("codex", "gpt-5.6-sol") + breaker.record_generation_failure() # open + await asyncio.sleep(0.02) + cancel = asyncio.Event() + + async def hang_forever(): + await asyncio.sleep(30) + + async def fire_cancel(): + await asyncio.sleep(0.05) + cancel.set() + + canceller = asyncio.create_task(fire_cancel()) + with pytest.raises(asyncio.CancelledError): + await generate_with_recovery( + hang_forever, + policy=RecoveryPolicy(deadline_seconds=60.0, backoff_base=0.01), + breaker=breaker, + cancel_event=cancel, + ) + await canceller + # The held probe slot was released — the next caller can probe. + assert breaker.state != "probing" diff --git a/tests/test_resume_admission.py b/tests/test_resume_admission.py index d8aee76a..7e492163 100644 --- a/tests/test_resume_admission.py +++ b/tests/test_resume_admission.py @@ -320,3 +320,52 @@ def test_matched_transcript_is_untouched(self): before = json.loads(json.dumps(messages)) TurnResumeManager._repair_unmatched_tool_use(messages, []) assert messages == before + + +class TestPostAcquireSafety: + async def test_corrupt_checkpoint_rejects_before_acquiring(self, tmp_path): + """Review blocker #5 (PR #242): reconstruction runs BEFORE the + execution lease is acquired, so a corrupt checkpoint becomes + TERMINAL_REJECTED — never a stranded ACTIVE row invisible to + resumable queries.""" + h, original = await suspend_turn(tmp_path) + h.store._conn.execute("UPDATE turns SET payload='{\"broken\": '") + h.store._conn.commit() + result = await h.manager.try_explicit_resume(resume_msg(original)) + # Malformed JSON self-heals inside load_resumable: rejected + # terminally, and the trigger falls through as a normal message. + assert result is None + (status,) = h.store._conn.execute("SELECT status FROM turns").fetchone() + assert status == TurnStatus.TERMINAL_REJECTED + # A structurally-valid-but-unrestorable payload rejects explicitly: + h2, original2 = await suspend_turn(tmp_path / "second") + h2.store._conn.execute( + "UPDATE turns SET payload='{\"fields\": {\"stuck_tracker\": 42}}'" + ) + h2.store._conn.commit() + result2 = await h2.manager.try_explicit_resume(resume_msg(original2)) + assert result2 is not None + assert "could not be restored" in result2[0] + (status,) = h2.store._conn.execute("SELECT status FROM turns").fetchone() + assert status == TurnStatus.TERMINAL_REJECTED # not stranded ACTIVE + + +class TestWaiterRegistry: + async def test_replaced_waiter_cannot_orphan_successor(self, tmp_path, monkeypatch): + """Review blocker #6b (PR #242): a cancelled predecessor's done + callback must not pop its successor's registry entry.""" + monkeypatch.setattr(tr, "_AUTO_POLL_SECONDS", 5.0) # keep waiters parked + h, original = await suspend_turn(tmp_path) + rows = h.store.list_suspended_sync("discord") + from src.turn_state import TurnKey + + key = TurnKey("discord", rows[0]["channel_id"], rows[0]["message_id"]) + h.manager.on_turn_suspended(key, rows[0]["generation"]) + first = h.manager._waiters[key] + h.manager.on_turn_suspended(key, rows[0]["generation"]) # replaces + second = h.manager._waiters[key] + assert second is not first + await asyncio.sleep(0.05) # predecessor's done callback has run + assert h.manager._waiters.get(key) is second # successor survives + second.cancel() + await asyncio.sleep(0) diff --git a/tests/test_turn_checkpoint_codec.py b/tests/test_turn_checkpoint_codec.py index f12aa075..76d7de19 100644 --- a/tests/test_turn_checkpoint_codec.py +++ b/tests/test_turn_checkpoint_codec.py @@ -257,3 +257,24 @@ def test_content_digest_is_full_sha256(): assert len(digest) == 64 assert digest != compute_content_digest("hello ") assert compute_content_digest("") == compute_content_digest(None or "") + + +class TestStorageRedaction: + def test_tool_use_inputs_are_secret_scrubbed_at_snapshot(self): + """Review blocker #8 (PR #242): tool arguments hit durable storage + secret-scrubbed (audit-storage parity). Non-secret args unchanged.""" + blobs, store, load = _blob_dict() + st = _full_turn() + secret = "sk-" + "a" * 24 + st.messages.append({ + "role": "assistant", + "content": [{ + "type": "tool_use", "id": "call_9", "name": "http_post", + "input": {"url": "https://x", "auth": f"api_key={secret}"}, + }], + }) + payload = snapshot_chat_turn(st, store_blob=store, generation_seq=2) + encoded = json.dumps(payload) + assert secret not in encoded + # Innocent arguments are untouched. + assert "https://x" in encoded diff --git a/tests/test_turn_state_store.py b/tests/test_turn_state_store.py index 98a7b351..3b5e2c3f 100644 --- a/tests/test_turn_state_store.py +++ b/tests/test_turn_state_store.py @@ -37,7 +37,7 @@ def store(tmp_path): def _admit(store, key=KEY): - lease = store.admit_turn_sync( + lease, disposition = store.admit_turn_sync( key, guild_id="g1", user_id="u1", @@ -47,6 +47,7 @@ def _admit(store, key=KEY): tool_catalog_hash="tc", session_snapshot={"messages": 3}, ) + assert disposition == "admitted" assert lease is not None return lease @@ -71,14 +72,43 @@ def test_admit_and_checkpoint_roundtrip(self, store): assert revision == 2 assert '"iteration": 2' in payload - def test_double_admission_returns_none(self, store): - _admit(store) - again = store.admit_turn_sync( + def _readmit(self, store): + return store.admit_turn_sync( KEY, guild_id=None, user_id=None, content_digest=None, code_version=None, prompt_policy_hash=None, tool_catalog_hash=None, session_snapshot=None, ) - assert again is None + + def test_existing_identity_dispositions(self, store): + # Review blocker #2 (PR #242): an existing identity must REFUSE + # fresh execution — it must never mean "run without durability". + lease = _admit(store) + assert self._readmit(store) == (None, "in_progress") # live lease + + store.suspend_sync(lease, {"p": 1}) + assert self._readmit(store) == (None, "resumable") + + new_lease = store.acquire_resume_lease_sync(KEY, lease.generation) + store.finish_sync(new_lease) + assert self._readmit(store) == (None, "already_processed") + + def test_expired_active_readmission_sweeps_to_resumable(self, store): + import time as _time + + lease = _admit(store) + store.record_intents_sync( + lease, 0, [{"tool_call_id": "c1", "tool_name": "t", "tool_input": {}}] + ) + store.mark_running_sync(lease, 0, "c1") + store._conn.execute( + "UPDATE turns SET lease_expires_at=?", [_time.time() - 10] + ) + store._conn.commit() + assert self._readmit(store) == (None, "resumable") + row = store._conn.execute("SELECT status FROM turns").fetchone() + assert row[0] == TurnStatus.SUSPENDED + op = store._conn.execute("SELECT state FROM operations").fetchone() + assert op[0] == OpState.OUTCOME_UNKNOWN def test_unavailable_store_admits_none(self, tmp_path): # A file where the DB directory should be → init fails → degraded. @@ -90,7 +120,7 @@ def test_unavailable_store_admits_none(self, tmp_path): KEY, guild_id=None, user_id=None, content_digest=None, code_version=None, prompt_policy_hash=None, tool_catalog_hash=None, session_snapshot=None, - ) is None + ) == (None, "store_unavailable") def test_recovery_deadline_persisted_utc(self, store): lease = _admit(store) @@ -302,18 +332,35 @@ def test_stale_active_suspends_and_ops_go_unknown(self, tmp_path): assert states == {OpState.OUTCOME_UNKNOWN} reopened.close() - def test_live_lease_survives_reopen(self, tmp_path): + def test_fast_restart_inside_lease_ttl_still_suspends(self, tmp_path): + # Deliberate amendment (review blocker #4, PR #242): the store is + # single-process, so at construction NO turn can legitimately be + # in flight — the boot sweep suspends ALL ACTIVE rows, lease expiry + # notwithstanding (a fast restart used to strand them ACTIVE forever). db = tmp_path / "turns.sqlite3" s = TurnStateStore(db, blob_dir=tmp_path / "blobs") _admit(s) s.close() - # Lease still in the future → NOT swept (a concurrent healthy owner - # in another process must not be hijacked). reopened = TurnStateStore(db, blob_dir=tmp_path / "blobs") row = reopened._conn.execute("SELECT status FROM turns").fetchone() - assert row[0] == TurnStatus.ACTIVE + assert row[0] == TurnStatus.SUSPENDED reopened.close() + def test_periodic_expired_active_sweep(self, tmp_path): + import time as _time + + s = TurnStateStore(tmp_path / "t.sqlite3", blob_dir=tmp_path / "blobs") + _admit(s) + # Live lease: the periodic sweep must NOT touch a healthy owner. + assert s.sweep_expired_active_sync() == {"turns": 0, "ops": 0} + s._conn.execute("UPDATE turns SET lease_expires_at=?", [_time.time() - 5]) + s._conn.commit() + out = s.sweep_expired_active_sync() + assert out["turns"] == 1 + row = s._conn.execute("SELECT status FROM turns").fetchone() + assert row[0] == TurnStatus.SUSPENDED + s.close() + class TestTtlSweep: def _age(self, store, key, *, progress_age_s): @@ -393,3 +440,117 @@ def test_effect_fingerprint_is_deterministic_and_sensitive(): c = effect_fingerprint("run_command", {"command": "rm", "host": "web"}) assert a == b # key order irrelevant assert a != c + + +class TestLedgerFencing: + """Review blocker #1 (PR #242): operation transitions are single + atomic statements fenced on generation + expected revision + lease + token + ACTIVE status + live lease + legal prior op state.""" + + def test_settled_op_cannot_reenter_running(self, store): + lease = _admit(store) + store.record_intents_sync( + lease, 0, [{"tool_call_id": "c1", "tool_name": "t", "tool_input": {}}] + ) + store.mark_running_sync(lease, 0, "c1") + store.settle_op_sync(lease, 0, "c1", state=OpState.APPLIED, result_text="ok") + with pytest.raises(StaleTurnError): + store.mark_running_sync(lease, 0, "c1") # APPLIED→RUNNING illegal + row = store._conn.execute( + "SELECT state, result FROM operations WHERE tool_call_id='c1'" + ).fetchone() + assert row == (OpState.APPLIED, "ok") # untouched + + def test_reinsert_never_resets_a_settled_op(self, store): + lease = _admit(store) + store.record_intents_sync( + lease, 0, [{"tool_call_id": "c1", "tool_name": "t", "tool_input": {}}] + ) + store.settle_op_sync(lease, 0, "c1", state=OpState.APPLIED, result_text="ok") + with pytest.raises(LedgerIntentError): + store.record_intents_sync( + lease, 0, + [{"tool_call_id": "c1", "tool_name": "t", "tool_input": {}}], + ) + row = store._conn.execute( + "SELECT state FROM operations WHERE tool_call_id='c1'" + ).fetchone() + assert row[0] == OpState.APPLIED # never reset to PREPARED + + def test_stale_revision_owner_cannot_mutate_ledger(self, store): + import dataclasses + + lease = _admit(store) + store.record_intents_sync( + lease, 0, [{"tool_call_id": "c1", "tool_name": "t", "tool_input": {}}] + ) + stale = dataclasses.replace(lease) + store.checkpoint_sync(lease, {"n": 1}, progressed=True) # revision advances + with pytest.raises(StaleTurnError): + store.mark_running_sync(stale, 0, "c1") # stale revision fenced out + # The current owner still can. + store.mark_running_sync(lease, 0, "c1") + + def test_released_lease_cannot_record_intents(self, store): + lease = _admit(store) + store.suspend_sync(lease, {}) + with pytest.raises(StaleTurnError): + store.record_intents_sync( + lease, 0, [{"tool_call_id": "c1", "tool_name": "t", "tool_input": {}}] + ) + + +class TestReleaseAcquired: + def test_release_returns_to_suspended(self, store): + lease = _admit(store) + store.suspend_sync(lease, {"p": 1}) + acquired = store.acquire_resume_lease_sync(KEY, lease.generation) + store.release_acquired_sync(acquired) + (status,) = store._conn.execute("SELECT status FROM turns").fetchone() + assert status == TurnStatus.SUSPENDED + # Resumable again — a later attempt can win it. + assert store.acquire_resume_lease_sync(KEY, lease.generation) is not None + + def test_release_with_reason_is_terminal(self, store): + lease = _admit(store) + store.suspend_sync(lease, {"p": 1}) + acquired = store.acquire_resume_lease_sync(KEY, lease.generation) + store.release_acquired_sync(acquired, terminal_reason="checkpoint corrupt") + status, payload = store._conn.execute( + "SELECT status, payload FROM turns" + ).fetchone() + assert status == TurnStatus.TERMINAL_REJECTED + assert payload is None + + def test_release_is_fenced(self, store): + import dataclasses + + lease = _admit(store) + store.suspend_sync(lease, {"p": 1}) + acquired = store.acquire_resume_lease_sync(KEY, lease.generation) + thief = dataclasses.replace(acquired) + thief.token = "not-yours" + store.release_acquired_sync(thief) # fenced out — no-op + (status,) = store._conn.execute("SELECT status FROM turns").fetchone() + assert status == TurnStatus.ACTIVE + + +class TestStoragePermissions: + def test_db_dir_and_blobs_are_owner_only(self, tmp_path): + import os + import stat + + s = TurnStateStore(tmp_path / "sec" / "turns.sqlite3", + blob_dir=tmp_path / "sec" / "blobs") + ref = s.store_blob_sync(b"sensitive") + db_mode = stat.S_IMODE(os.stat(s.db_path).st_mode) + dir_mode = stat.S_IMODE(os.stat(tmp_path / "sec").st_mode) + blob_dir_mode = stat.S_IMODE(os.stat(tmp_path / "sec" / "blobs").st_mode) + blob_mode = stat.S_IMODE( + os.stat(tmp_path / "sec" / "blobs" / ref.split(":", 1)[1]).st_mode + ) + assert db_mode == 0o600 + assert dir_mode == 0o700 + assert blob_dir_mode == 0o700 + assert blob_mode == 0o600 + s.close() diff --git a/tests/test_write_invariant_integration.py b/tests/test_write_invariant_integration.py index 8a35a7fa..ce838d9a 100644 --- a/tests/test_write_invariant_integration.py +++ b/tests/test_write_invariant_integration.py @@ -421,12 +421,13 @@ def test_ttl_sweep_runs_and_swallows_failures(self, tmp_path): from src.turn_state import TurnKey store = TurnStateStore(tmp_path / "hk" / "turns.sqlite3") - lease = store.admit_turn_sync( + lease, disposition = store.admit_turn_sync( TurnKey("discord", "c1", "m1"), guild_id=None, user_id="u", content_digest="d", code_version="t", prompt_policy_hash="p", tool_catalog_hash="t", session_snapshot=None, ) + assert disposition == "admitted" store.suspend_sync(lease, {"p": 1}) store._conn.execute( "UPDATE turns SET last_progress_at=?", [_time.time() - 25 * 3600] @@ -457,7 +458,70 @@ def test_ttl_sweep_runs_and_swallows_failures(self, tmp_path): row = store._conn.execute("SELECT status FROM turns").fetchone() assert row[0] == TurnStatus.TERMINAL_EXPIRED - # A raising store must never break housekeeping. + # The expired-active defense sweep runs too: park a dead-owner + # ACTIVE row and let housekeeping suspend it (log-line branch). + lease2, disposition2 = store.admit_turn_sync( + TurnKey("discord", "c2", "m2"), + guild_id=None, user_id="u", content_digest="d", code_version="t", + prompt_policy_hash="p", tool_catalog_hash="t", + session_snapshot=None, + ) + assert disposition2 == "admitted" + store._conn.execute( + "UPDATE turns SET lease_expires_at=? WHERE message_id='m2'", + [_time.time() - 5], + ) + store._conn.commit() + hk.cleanup_stale() + row2 = store._conn.execute( + "SELECT status FROM turns WHERE message_id='m2'" + ).fetchone() + assert row2[0] == TurnStatus.SUSPENDED + + # A raising store must never break housekeeping — both sweep steps. + store.sweep_expired_active_sync = MagicMock(side_effect=RuntimeError("boom")) store.ttl_sweep_sync = MagicMock(side_effect=RuntimeError("boom")) hk.cleanup_stale() store.close() + + +class TestRedeliveryRefusal: + async def test_redelivered_message_never_reruns_effects(self, tmp_path): + """Review blocker #2 (PR #242): after a restart wipes the in-memory + dedup cache, a redelivered message with a terminal ledger row must + REFUSE fresh execution — never run unledgered.""" + from unittest.mock import AsyncMock + + bot, fake, store = build_with_store( + [ + tool_call_response(("run_command", {"command": "deploy"})), + text_response("done"), + ], + tmp_path, + ) + msg = FakeMessage("deploy the thing") + text, *_ = await run_loop(bot, msg) + assert text == "done" + + # Same message identity arrives again (Discord redelivery). + fake.responses.extend([text_response("must never be produced")]) + bot.tool_executor.execute = AsyncMock() + text2, _, is_error2, tools2, _ = await run_loop(bot, msg) + assert "already processed" in text2 + assert is_error2 is False + assert tools2 == [] + bot.tool_executor.execute.assert_not_awaited() + + async def test_suspended_identity_redelivery_points_at_resume(self, tmp_path): + from unittest.mock import AsyncMock + + bot, fake, store = build_with_store([], tmp_path) + fake.responses.append(capacity_forever(fake)) + msg = FakeMessage("long job") + await run_loop(bot, msg) # suspends + bot.tool_executor.execute = AsyncMock() + text2, *_ = await run_loop(bot, msg) # redelivery of the SAME id + assert "resume" in text2 + bot.tool_executor.execute.assert_not_awaited() + (status,) = store._conn.execute("SELECT status FROM turns").fetchone() + assert status == TurnStatus.SUSPENDED # untouched From d62a6857dff3aad07af2ce7afcf1c906662401ba Mon Sep 17 00:00:00 2001 From: ATCharpentier Date: Thu, 30 Jul 2026 16:36:41 -0400 Subject: [PATCH 11/18] =?UTF-8?q?fix:=20review=20round=202=20=E2=80=94=20a?= =?UTF-8?q?ll=20six=20blockers=20(PR=20#242)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit R2-B1 (lease vs long work): the owner now BEATS the lease alive — a per-turn heartbeat task (lease_ttl/3 cadence) started at admission and resume, stopped on suspend/terminal settlement, exiting cleanly on a stolen lease or dead store. _fenced_update carries the COMPLETE live fence (ACTIVE status + unexpired lease joined generation/revision/token), so an expired owner can neither checkpoint (which used to silently renew the lease) nor settle. Pinned: a tool 2.5x the TTL completes; the expired-checkpoint renewal repro is rejected. R2-B2 (runtime admission fail-open): once the store was wired available, an admission I/O failure (raised OR disposition) refuses execution with a bounded notice — identity-unverifiable never means run-unledgered. Legacy runs remain only for durability off/failed at wiring. R2-B3 (auto-resume off-by-one): intake appends the user message BEFORE the turn, so the suspension capture already includes it — legal growth is exactly +1 (the preservation marker). Pinned through the REAL MessagePipeline path end-to-end (suspend via pipeline.run → marker → auto-resume admits → TERMINAL_COMPLETED → reply on the original message). R2-B4 (orphaned attempt): _attempt_cancellable's finally now cancels AND awaits both helper tasks — cancelling the recovery owner can no longer leave a live provider request behind an abandoned probe. Pinned. R2-B5 (scrub coverage): storage scrub = the audit-parity tool-aware redaction (_scrub_tool_input_for_storage, e.g. email bodies) composed with a RECURSIVE secret-pattern scrub, applied to EVERY persisted copy of tool arguments (transcript tool_use blocks AND trajectory iterations). Pinned: nested auth headers and email bodies never persist raw. R2-B6 (OUTCOME_UNKNOWN enforcement): unresolved operations HALT continuation — auto-resume stands down permanently; explicit resume returns the op list, transitions them to MANUAL_RESOLUTION_REQUIRED, and closes the turn TERMINAL_REJECTED. No generation happens (pinned: zero LLM calls). 'Never auto-rerun' is enforcement now, not prompt advice. Also: corrupt checkpoint payloads self-heal to TERMINAL_REJECTED inside load_resumable; two auto-resume tests fixed to wait for TERMINAL status (they raced the in-flight ACTIVE resume under coverage instrumentation). Suite 7,905/5, lint/type gates new=0, coverage findings=0 at 88.4%. --- src/discord/tool_loop.py | 5 ++ src/discord/turn_resume.py | 67 ++++++++++++++-- src/llm/recovery.py | 16 ++-- src/turn_state/codec.py | 77 +++++++++++++++---- src/turn_state/durability.py | 88 +++++++++++++++++++-- src/turn_state/store.py | 38 ++++++++- tests/test_recovery_policy.py | 32 ++++++++ tests/test_resume_admission.py | 94 ++++++++++++++++++++++- tests/test_turn_checkpoint_codec.py | 55 +++++++++++++ tests/test_turn_durability_heartbeat.py | 93 ++++++++++++++++++++++ tests/test_turn_state_store.py | 71 +++++++++++++++++ tests/test_write_invariant_integration.py | 94 +++++++++++++++++++++++ 12 files changed, 687 insertions(+), 43 deletions(-) create mode 100644 tests/test_turn_durability_heartbeat.py diff --git a/src/discord/tool_loop.py b/src/discord/tool_loop.py index f08ed14f..8ca97847 100644 --- a/src/discord/tool_loop.py +++ b/src/discord/tool_loop.py @@ -452,6 +452,11 @@ async def run( "This request has preserved, resumable work — say " "`resume` to continue it instead of starting over." ), + "admission_error": ( + "The durability ledger is unreachable, so I can't verify " + "whether this exact request already ran. Refusing to " + "execute it blind — try again shortly." + ), } text = notices.get( st.durability.blocked, "This request cannot be re-run." diff --git a/src/discord/turn_resume.py b/src/discord/turn_resume.py index fc230077..88069e26 100644 --- a/src/discord/turn_resume.py +++ b/src/discord/turn_resume.py @@ -128,12 +128,14 @@ async def _auto_resume_waiter( """Wait for capacity, then resume IF nothing else happened. The advance check anchors on the session length captured AT - SUSPENSION (``suspend_len``): the only legitimate growth afterwards - is the suspending turn's own intake bookkeeping (exactly +2: the - user message and the preserved-marker; a directly-driven turn adds - 0). Anything else — or an unreadable session — stands auto-resume - down fail-safe; explicit resume stays available. No sampling window - exists for a stranger message to sneak into the baseline. + SUSPENSION (``suspend_len``). Production ordering (round-2 blocker + #3, PR #242): intake appends the USER message BEFORE the turn runs, + so the suspension capture already includes it, and the only + legitimate growth afterwards is the assistant preservation marker — + exactly +1 (a directly-driven turn adds 0). Anything else — or an + unreadable session — stands auto-resume down fail-safe; explicit + resume stays available. No sampling window exists for a stranger + message to sneak into the baseline. Capacity detection is ACTIVE: a quiet breaker is never probed by anyone, so the waiter claims the probe slot itself when the cooldown @@ -145,7 +147,7 @@ async def _auto_resume_waiter( """ give_up_at = time.monotonic() + self._resume_ttl_hours * 3600.0 breaker = self._llm_gateway.capacity_breaker_for() - allowed = {suspend_len, suspend_len + 2} if suspend_len >= 0 else set() + allowed = {suspend_len, suspend_len + 1} if suspend_len >= 0 else set() while time.monotonic() < give_up_at: await asyncio.sleep(_AUTO_POLL_SECONDS) row = await asyncio.to_thread(self._store.load_resumable_sync, key) @@ -180,6 +182,14 @@ def _session_len(self, channel_id: str) -> int: return -1 async def _run_auto_resume(self, key: TurnKey, row: dict) -> None: + if self._unresolved_ops(row): + # Never auto-continue over ambiguous external effects — a human + # must look at this (explicit resume delivers the details). + log.warning( + "Auto-resume for %s stands down permanently: unresolved " + "operations require manual resolution", key, + ) + return async with self._channel_lock(key.channel_id): if self._channel_state.active_requests.get(key.channel_id): log.info("Auto-resume for %s stands down: channel busy", key) @@ -227,6 +237,22 @@ def _append_session( def is_resume_trigger(content: str) -> bool: return (content or "").strip().lower().rstrip("!.") in RESUME_TRIGGERS + @staticmethod + def _unresolved_ops(row: dict) -> list[dict]: + """Operations whose external outcome is not settled. Their presence + HALTS continuation (round-2 blocker #6, PR #242): 'never auto-rerun' + is enforced by not generating, not by asking the model nicely.""" + blocked_states = { + OpState.OUTCOME_UNKNOWN, + OpState.MANUAL_RESOLUTION_REQUIRED, + OpState.PREPARED, + OpState.RUNNING, + } + return [ + op for op in (row.get("operations") or []) + if op.get("state") in blocked_states + ] + async def try_explicit_resume(self, message: Any): """Resume the channel's suspended turn when *message* is a trigger. @@ -260,6 +286,33 @@ async def try_explicit_resume(self, message: Any): waiter = self._waiters.pop(key, None) if waiter is not None: waiter.cancel() + unresolved = self._unresolved_ops(row) + if unresolved: + # Halt: continuation would let a later generation re-issue the + # same effect under a fresh call id. Hand the ambiguity to the + # human and close the turn out. + names = ", ".join( + sorted({str(op.get("tool_name") or "unknown") for op in unresolved}) + ) + moved = await asyncio.to_thread( + self._store.mark_ops_manual_sync, key, row["generation"] + ) + await asyncio.to_thread( + self._store.reject_resumable_sync, key, + f"{len(unresolved)} unresolved operation(s) — manual resolution", + ) + log.warning( + "Resume of %s halted: %d unresolved op(s) (%d moved to manual)", + key, len(unresolved), moved, + ) + return ( + f"I can't safely continue that work: {len(unresolved)} " + f"interrupted operation(s) ({names}) have UNKNOWN outcomes — " + "they may or may not have applied, and I will not re-run " + "them automatically. Verify their current state, then ask " + "fresh for whatever is still needed.", + False, False, [], False, + ) st, _original, reason = await self._validate_and_rebuild(key, row) if st is None: return ( diff --git a/src/llm/recovery.py b/src/llm/recovery.py index 1f7e471c..71abd2d4 100644 --- a/src/llm/recovery.py +++ b/src/llm/recovery.py @@ -122,14 +122,18 @@ async def _attempt_cancellable( ) if attempt_task in done: return attempt_task.result() - attempt_task.cancel() - with contextlib.suppress(BaseException): - await attempt_task raise asyncio.CancelledError("cancelled during LLM attempt") finally: - cancel_task.cancel() - with contextlib.suppress(BaseException): - await cancel_task + # Cancellation of the RECOVERY OWNER task lands here too (round-2 + # blocker #4, PR #242): both helper tasks — including a still-live + # provider attempt — must be cancelled AND awaited before control + # leaves, so shutdown/task-cancellation never orphans an in-flight + # request while the caller releases its breaker probe. + for helper in (attempt_task, cancel_task): + if not helper.done(): + helper.cancel() + with contextlib.suppress(BaseException): + await helper def _notify_wait( diff --git a/src/turn_state/codec.py b/src/turn_state/codec.py index 5d237ab7..e1ac18ac 100644 --- a/src/turn_state/codec.py +++ b/src/turn_state/codec.py @@ -100,28 +100,52 @@ def compute_content_digest(text: str) -> str: # ── storage redaction ──────────────────────────────────────────────── -def _scrub_tool_use_inputs(obj: Any) -> Any: - """Secret-scrub string values inside assistant ``tool_use`` inputs - before they hit durable storage (review blocker #8, PR #242) — the - parity move with audit storage, which deliberately scrubs tool inputs. +def _deep_scrub_strings(value: Any) -> Any: + """Recursive secret-pattern scrub over EVERY nested string value.""" + if isinstance(value, str): + return scrub_output_secrets(value) + if isinstance(value, list): + return [_deep_scrub_strings(v) for v in value] + if isinstance(value, dict): + return {k: _deep_scrub_strings(v) for k, v in value.items()} + return value + + +def scrub_stored_tool_input(tool_name: str, tool_input: Any) -> Any: + """The storage scrub for tool arguments (round-2 blocker #5, PR #242): + the SAME tool-aware privacy redaction audit storage uses + (`_scrub_tool_input_for_storage` — e.g. email bodies, which are not + token-shaped), composed with a recursive secret-pattern scrub so nested + values (auth headers, embedded dicts/lists) are covered too. Applied to + EVERY persisted representation of the arguments. Tool RESULTS are already scrubbed at source (`_run_one_tool` runs scrub_output_secrets before building the result block), and credential-bearing USER messages are deleted by the intake secret gate - before a turn ever starts — tool arguments were the remaining - unscrubbed surface. The scrub applies at SNAPSHOT time only, so a - resumed transcript shows the model its own arguments with any embedded - secrets masked; the executed effect already happened and is unaffected. + before a turn ever starts. The scrub applies at SNAPSHOT time only, so + a resumed transcript shows the model its own arguments with secrets + masked; the executed effect already happened and is unaffected. """ + from ..discord.tool_loop_helpers import _scrub_tool_input_for_storage + + if isinstance(tool_input, dict): + tool_input = _scrub_tool_input_for_storage(tool_name or "", tool_input) + return _deep_scrub_strings(tool_input) + + +def _scrub_tool_use_inputs(obj: Any) -> Any: + """Apply the storage scrub to assistant ``tool_use`` blocks in the + transcript (name-aware — the block carries the tool name).""" if isinstance(obj, list): return [_scrub_tool_use_inputs(x) for x in obj] if isinstance(obj, dict): if obj.get("type") == "tool_use" and isinstance(obj.get("input"), dict): - scrubbed = { - k: scrub_output_secrets(v) if isinstance(v, str) else v - for k, v in obj["input"].items() + return { + **obj, + "input": scrub_stored_tool_input( + str(obj.get("name") or ""), obj["input"] + ), } - return {**obj, "input": scrubbed} return {k: _scrub_tool_use_inputs(v) for k, v in obj.items()} return obj @@ -202,9 +226,30 @@ def import_stuck_tracker(state: dict, tracker_cls): def _trajectory_to_payload(trajectory) -> dict: """FULL trajectory state (unlike TrajectoryTurn.to_dict, which is lossy - by design for the JSONL files). Iterations ride as asdict() rows; the - saved iteration count is the anti-double-append revision (Odin round-2: - "resumption cannot append the same iteration twice").""" + by design for the JSONL files). Iterations ride as asdict() rows with + their tool-call ARGUMENTS passed through the storage scrub — this copy + persists the same raw inputs the transcript does and must get the same + redaction (round-2 blocker #5, PR #242). The saved iteration count is + the anti-double-append revision (Odin round-2: "resumption cannot + append the same iteration twice").""" + + def _scrubbed_iteration(it) -> dict: + row = asdict(it) + calls = row.get("tool_calls") + if isinstance(calls, list): + row["tool_calls"] = [ + { + **tc, + "input": scrub_stored_tool_input( + str(tc.get("name") or ""), tc.get("input") + ), + } + if isinstance(tc, dict) + else tc + for tc in calls + ] + return row + return { "message_id": trajectory.message_id, "channel_id": trajectory.channel_id, @@ -215,7 +260,7 @@ def _trajectory_to_payload(trajectory) -> dict: "user_content": trajectory.user_content, "system_prompt": trajectory.system_prompt, "history": list(trajectory.history or []), - "iterations": [asdict(it) for it in trajectory.iterations], + "iterations": [_scrubbed_iteration(it) for it in trajectory.iterations], "final_response": trajectory.final_response, "tools_used": list(trajectory.tools_used or []), "is_error": trajectory.is_error, diff --git a/src/turn_state/durability.py b/src/turn_state/durability.py index 74bd4c61..08096ceb 100644 --- a/src/turn_state/durability.py +++ b/src/turn_state/durability.py @@ -41,7 +41,15 @@ from ..odin_log import get_logger from .codec import compute_content_digest, snapshot_chat_turn -from .store import OpState, TurnKey, TurnLease, TurnStateStore, TurnStatus +from .store import ( + OpState, + StaleTurnError, + TurnKey, + TurnLease, + TurnStateStore, + TurnStateUnavailableError, + TurnStatus, +) log = get_logger("turn_state") @@ -49,6 +57,10 @@ # transcript; the ledger copy is for replay/reconciliation). _OP_RESULT_CAP = 4000 +# Heartbeat cadence floor — the interval is lease_ttl/3 but never busier +# than this (module-level so tests can drive real beats fast). +_HEARTBEAT_FLOOR_SECONDS = 5.0 + def _hash_text(text: str) -> str: return hashlib.sha256((text or "").encode("utf-8", "replace")).hexdigest() @@ -71,9 +83,14 @@ def __init__(self, store: TurnStateStore | None, lease: TurnLease | None) -> Non # One-shot remaining budget for a resumed generation (see resumed()). self._resume_budget: float | None = None # Admission refusal disposition (review blocker #2, PR #242): - # "already_processed" / "in_progress" / "resumable" — the loop must - # REFUSE fresh execution, never run unledgered. None = no refusal. + # "already_processed" / "in_progress" / "resumable" / + # "admission_error" — the loop must REFUSE fresh execution, never + # run unledgered. None = no refusal. self.blocked: str | None = None + # Lease maintenance (round-2 blocker #1): tools run up to 3660s and + # generations up to 3600s against a 120s lease — the owner beats it + # alive for the turn's whole life. + self._heartbeat_task: asyncio.Task | None = None # -- construction -------------------------------------------------- @@ -121,12 +138,25 @@ async def admit( session_snapshot=session_snapshot, ) except Exception: - log.exception("Turn admission failed — running without durability") - return cls.disabled() + # The store was wired available; an admission failure here means + # the identity could not be checked — refuse (round-2 blocker #2). + log.exception("Turn admission raised — refusing execution (fail closed)") + handle = cls.disabled() + handle.blocked = "admission_error" + return handle if lease is not None: - return cls(store, lease) + handle = cls(store, lease) + handle._start_heartbeats() + return handle if disposition == "store_unavailable": - return cls.disabled() + # The store was wired available but an admission I/O failure + # means this message's identity COULD NOT be checked — refusing + # is the only safe answer (round-2 blocker #2, PR #242). Legacy + # execution is reserved for durability being off/failed at + # wiring, where no ledger can exist to contradict. + handle = cls.disabled() + handle.blocked = "admission_error" + return handle handle = cls.disabled() handle.blocked = disposition return handle @@ -151,6 +181,7 @@ def resumed( handle = cls(store, lease) handle.generation_seq = int(generation_seq) handle._resume_budget = first_generation_budget + handle._start_heartbeats() return handle def pop_resume_budget(self) -> float | None: @@ -174,6 +205,44 @@ def enabled(self) -> bool: def lease(self) -> TurnLease | None: return self._lease + # -- lease maintenance --------------------------------------------- + + def _start_heartbeats(self) -> None: + if self._store is None or self._lease is None: + return + interval = max(_HEARTBEAT_FLOOR_SECONDS, float(self._store.lease_ttl) / 3.0) + store, lease = self._store, self._lease + + async def _beat() -> None: + while True: + await asyncio.sleep(interval) + if not self.enabled: + return + try: + await asyncio.to_thread(store.heartbeat_sync, lease) + except StaleTurnError: + log.warning( + "Turn %s lost its lease during heartbeat — another " + "owner or a sweep took it; the next fenced write " + "fail-closes this turn", lease.key, + ) + return + except TurnStateUnavailableError: + log.warning( + "Turn heartbeat write failed for %s (store trouble); " + "the next fenced write fail-closes the turn", lease.key, + ) + return + + self._heartbeat_task = asyncio.get_running_loop().create_task( + _beat(), name=f"turn-heartbeat:{lease.key.channel_id}:{lease.key.message_id}" + ) + + def _stop_heartbeats(self) -> None: + task, self._heartbeat_task = self._heartbeat_task, None + if task is not None: + task.cancel() + # -- snapshot plumbing (runs in a worker thread: blob writes + sqlite) -- def _checkpoint_sync( @@ -307,6 +376,10 @@ async def suspend(self, st, reason: str) -> bool: preserved checkpoint that does not exist).""" if not self.enabled: return False + # Beats stop first; the suspension write still lands inside the + # last-beat + TTL window, and a failed suspension falls through to + # settle_terminal which needs no live beats either. + self._stop_heartbeats() extra: dict = {"suspend_reason": reason} try: if st.trace is not None: @@ -340,6 +413,7 @@ async def settle_terminal(self, *, cancelled: bool, is_error: bool) -> None: No further external effect follows, so a failure here must not destroy an already-computed reply — log and move on. """ + self._stop_heartbeats() # unconditionally — every turn exit lands here if not self.enabled: return if cancelled or self.cancelled: diff --git a/src/turn_state/store.py b/src/turn_state/store.py index a42e2d4e..15a1fc43 100644 --- a/src/turn_state/store.py +++ b/src/turn_state/store.py @@ -265,13 +265,18 @@ def _require(self) -> sqlite3.Connection: def _fenced_update( self, lease: TurnLease, set_sql: str, params: list, *, bump_revision: bool = True ) -> None: - """Run one UPDATE fenced on (generation, revision, lease token).""" + """Run one UPDATE under the COMPLETE live fence: generation, + expected revision, lease token, ACTIVE status, and unexpired lease + (round-2 blocker #1, PR #242 — an expired-lease owner must never + renew itself through a checkpoint; the fence is identical for turn + and ledger writes).""" conn = self._require() new_revision = lease.revision + 1 if bump_revision else lease.revision sql = ( f"UPDATE turns SET {set_sql}, revision=? " "WHERE source=? AND channel_id=? AND message_id=? " - "AND turn_generation=? AND revision=? AND lease_token=?" + "AND turn_generation=? AND revision=? AND lease_token=? " + "AND status=? AND lease_expires_at > ?" ) try: with self._write_lock: @@ -279,7 +284,7 @@ def _fenced_update( sql, [*params, new_revision, lease.key.source, lease.key.channel_id, lease.key.message_id, lease.generation, lease.revision, - lease.token], + lease.token, TurnStatus.ACTIVE, time.time()], ) conn.commit() except sqlite3.Error as exc: @@ -366,7 +371,10 @@ def admit_turn_sync( except sqlite3.IntegrityError: return None, self._classify_existing_row(key, now) except sqlite3.Error: - log.exception("Turn admission failed — running without durability") + log.exception( + "Turn admission I/O failure — identity unverifiable " + "(caller fail-closes when the store was wired available)" + ) return None, "store_unavailable" return ( TurnLease(key=key, generation=generation, token=token, revision=0), @@ -721,6 +729,28 @@ def reject_resumable_sync(self, key: TurnKey, reason: str) -> None: except sqlite3.Error: log.exception("reject_resumable failed (non-fatal)") + def mark_ops_manual_sync(self, key: TurnKey, generation: str) -> int: + """OUTCOME_UNKNOWN → MANUAL_RESOLUTION_REQUIRED for a turn whose + continuation is being halted (round-2 blocker #6, PR #242). Both + states are in the never-auto-expire set; this transition records + that a human was told and owns the reconciliation now.""" + conn = self._require() + try: + with self._write_lock: + cur = conn.execute( + "UPDATE operations SET state=?, updated_at=? " + "WHERE source=? AND channel_id=? AND message_id=? " + "AND turn_generation=? AND state=?", + [OpState.MANUAL_RESOLUTION_REQUIRED, time.time(), + key.source, key.channel_id, key.message_id, generation, + OpState.OUTCOME_UNKNOWN], + ) + conn.commit() + return cur.rowcount + except sqlite3.Error: + log.exception("mark_ops_manual failed (non-fatal)") + return 0 + def release_acquired_sync( self, lease: TurnLease, *, terminal_reason: str | None = None ) -> None: diff --git a/tests/test_recovery_policy.py b/tests/test_recovery_policy.py index 399a751f..feb9e6ca 100644 --- a/tests/test_recovery_policy.py +++ b/tests/test_recovery_policy.py @@ -260,3 +260,35 @@ async def fire_cancel(): await canceller # The held probe slot was released — the next caller can probe. assert breaker.state != "probing" + + +async def test_parent_cancellation_unwinds_the_attempt(): + """Round-2 blocker #4 (PR #242): cancelling the RECOVERY OWNER task must + cancel and await the in-flight attempt too — an orphaned provider + request must not outlive its abandoned breaker probe.""" + cancel = asyncio.Event() # never fires; parent cancellation is the test + attempt_cancelled = asyncio.Event() + attempt_started = asyncio.Event() + + async def slow_attempt(): + attempt_started.set() + try: + await asyncio.sleep(30) + except asyncio.CancelledError: + attempt_cancelled.set() + raise + return "never" + + owner = asyncio.create_task( + generate_with_recovery( + slow_attempt, + policy=RecoveryPolicy(deadline_seconds=60.0), + cancel_event=cancel, + ) + ) + await asyncio.wait_for(attempt_started.wait(), timeout=5) + owner.cancel() + with pytest.raises(asyncio.CancelledError): + await owner + # The attempt was unwound BEFORE the owner finished — not orphaned. + assert attempt_cancelled.is_set() diff --git a/tests/test_resume_admission.py b/tests/test_resume_admission.py index 7e492163..a6ae19c9 100644 --- a/tests/test_resume_admission.py +++ b/tests/test_resume_admission.py @@ -252,9 +252,9 @@ async def test_auto_resume_fires_when_capacity_returns(self, tmp_path, monkeypat key = TurnKey("discord", rows[0]["channel_id"], rows[0]["message_id"]) h.manager.on_turn_suspended(key, rows[0]["generation"]) - for _ in range(200): + for _ in range(600): await asyncio.sleep(0.02) - if h.row()[0] != TurnStatus.SUSPENDED: + if h.row()[0] in TurnStatus.TERMINAL: break assert h.row()[0] == TurnStatus.TERMINAL_COMPLETED # The reply landed against the ORIGINAL message. @@ -274,8 +274,11 @@ async def test_auto_resume_stands_down_when_session_advances( key = TurnKey("discord", rows[0]["channel_id"], rows[0]["message_id"]) h.manager.on_turn_suspended(key, rows[0]["generation"]) - await asyncio.sleep(0.06) # baseline sampled while the breaker paces + await asyncio.sleep(0.06) # waiter parked while the breaker paces + # A real intervening turn appends user + assistant (+2) — beyond the + # single preservation-marker growth (+1) the waiter tolerates. h.bot.sessions.add_message(str(original.channel.id), "user", "new topic") + h.bot.sessions.add_message(str(original.channel.id), "assistant", "answered") await asyncio.sleep(0.02) make_breaker_probe_ready(h) # capacity "returns" AFTER the advance await asyncio.sleep(0.3) @@ -369,3 +372,88 @@ async def test_replaced_waiter_cannot_orphan_successor(self, tmp_path, monkeypat assert h.manager._waiters.get(key) is second # successor survives second.cancel() await asyncio.sleep(0) + + +class TestUnknownOutcomeHaltsContinuation: + """Round-2 blocker #6 (PR #242): unresolved OUTCOME_UNKNOWN operations + HALT continuation — enforcement, not model-facing advice.""" + + async def _suspend_with_unknown(self, tmp_path): + h, original = await suspend_turn(tmp_path) + h.store._conn.execute( + "UPDATE operations SET state=?", [OpState.OUTCOME_UNKNOWN] + ) + h.store._conn.commit() + return h, original + + async def test_explicit_resume_halts_and_hands_to_human(self, tmp_path): + h, original = await self._suspend_with_unknown(tmp_path) + heal_capacity(h, text_response("must never generate")) + calls_before = len(h.fake.calls) + result = await h.manager.try_explicit_resume(resume_msg(original)) + assert result is not None + assert "UNKNOWN outcomes" in result[0] + assert "parse_time" in result[0] + assert len(h.fake.calls) == calls_before # NO generation happened + status = h.store._conn.execute("SELECT status FROM turns").fetchone()[0] + assert status == TurnStatus.TERMINAL_REJECTED + op_state = h.store._conn.execute( + "SELECT state FROM operations" + ).fetchone()[0] + assert op_state == OpState.MANUAL_RESOLUTION_REQUIRED + + async def test_auto_resume_stands_down_on_unknowns(self, tmp_path, monkeypatch): + monkeypatch.setattr(tr, "_AUTO_POLL_SECONDS", 0.02) + h, original = await self._suspend_with_unknown(tmp_path) + heal_capacity(h, text_response("must never generate")) + rows = h.store.list_suspended_sync("discord") + from src.turn_state import TurnKey + + key = TurnKey("discord", rows[0]["channel_id"], rows[0]["message_id"]) + h.manager.on_turn_suspended(key, rows[0]["generation"]) + await asyncio.sleep(0.3) + status = h.store._conn.execute("SELECT status FROM turns").fetchone()[0] + assert status == TurnStatus.SUSPENDED # untouched, awaiting a human + for task in list(h.manager._waiters.values()): + task.cancel() + + +class TestProductionPipelinePath: + async def test_suspension_bookkeeping_is_plus_one_and_auto_resume_admits( + self, tmp_path, monkeypatch + ): + """Round-2 blocker #3 (PR #242): drive the REAL MessagePipeline + (user message appended BEFORE the turn, preservation marker after + → +1), then prove the waiter's arithmetic admits auto-resume.""" + monkeypatch.setattr(tr, "_AUTO_POLL_SECONDS", 0.05) + h2 = Harness([tool_call_response(("parse_time", {"text": "x"}))], tmp_path) + h2.fake.responses.append(capacity_forever(h2.fake)) + h2.bot.pipeline._turn_resume = h2.manager + original = FakeMessage("please do the long thing") + h2.register(original) + await h2.bot.pipeline.run(original, original.content) + + ch_id = str(original.channel.id) + session = h2.bot.sessions._sessions.get(ch_id) + assert session is not None + # +1 bookkeeping: the suspension callback captured a length that + # already included the user message; only the marker follows. + marker = session.messages[-1].content + assert "PRESERVED" in marker + assert h2.row()[0] == TurnStatus.SUSPENDED + + heal_capacity(h2, text_response("Pipeline auto-finish.")) + # Wait for a TERMINAL status — the row passes through ACTIVE while + # the auto-resume runs (breaking on first non-SUSPENDED raced the + # in-flight resume under coverage instrumentation). + for _ in range(600): + await asyncio.sleep(0.05) + if h2.row()[0] in TurnStatus.TERMINAL: + break + assert h2.row()[0] == TurnStatus.TERMINAL_COMPLETED + assert any( + "Pipeline auto-finish." in (r["content"] or "") + for r in original.replies + ) + for task in list(h2.manager._waiters.values()): + task.cancel() diff --git a/tests/test_turn_checkpoint_codec.py b/tests/test_turn_checkpoint_codec.py index 76d7de19..5b7a78a4 100644 --- a/tests/test_turn_checkpoint_codec.py +++ b/tests/test_turn_checkpoint_codec.py @@ -278,3 +278,58 @@ def test_tool_use_inputs_are_secret_scrubbed_at_snapshot(self): assert secret not in encoded # Innocent arguments are untouched. assert "https://x" in encoded + + +class TestStorageRedactionAllCopies: + """Round-2 blocker #5 (PR #242): the storage scrub covers EVERY + persisted copy of tool arguments, nested values, and tool-aware + privacy fields.""" + + def _payload_with(self, tool_name, tool_input): + blobs, store, load = _blob_dict() + st = _full_turn() + st.messages.append({ + "role": "assistant", + "content": [{ + "type": "tool_use", "id": "call_9", "name": tool_name, + "input": tool_input, + }], + }) + st._trajectory.iterations.append( + ToolIteration( + iteration=1, + tool_calls=[{"id": "call_9", "name": tool_name, + "input": tool_input}], + tool_results=[], llm_text="", input_tokens=0, output_tokens=0, + duration_ms=0, provider="codex", model="m", + reasoning_effort=None, + ) + ) + return json.dumps(snapshot_chat_turn(st, store_blob=store, generation_seq=2)) + + def test_nested_secret_values_are_scrubbed_everywhere(self): + secret = "sk-" + "b" * 24 + encoded = self._payload_with( + "http_post", + {"url": "https://x", "headers": {"Authorization": f"token={secret}"}, + "variants": [{"auth": f"api_key={secret}"}]}, + ) + assert secret not in encoded # transcript AND trajectory copies + assert "https://x" in encoded + + def test_email_body_privacy_redaction_applies(self): + encoded = self._payload_with( + "email_send", + {"to": "a@b.c", "subject": "hi", + "body": "Deeply personal contents that are not token-shaped."}, + ) + assert "Deeply personal contents" not in encoded + assert "redacted email body" in encoded + + def test_non_string_scalars_pass_through_the_scrub(self): + from src.turn_state.codec import scrub_stored_tool_input + + out = scrub_stored_tool_input( + "run_command", {"count": 3, "flag": True, "ratio": 1.5, "none": None} + ) + assert out == {"count": 3, "flag": True, "ratio": 1.5, "none": None} diff --git a/tests/test_turn_durability_heartbeat.py b/tests/test_turn_durability_heartbeat.py new file mode 100644 index 00000000..cbc305d9 --- /dev/null +++ b/tests/test_turn_durability_heartbeat.py @@ -0,0 +1,93 @@ +"""Pins for the durability lease heartbeat (round-2 blocker #1, PR #242). + +The PRODUCTION beat loop is exercised directly: it keeps a long-lived +lease alive, exits cleanly when the lease is stolen (StaleTurnError) or +the store dies (TurnStateUnavailableError), and always stops on terminal +settlement. +""" + +from __future__ import annotations + +import asyncio +import time + +import pytest + +import src.turn_state.durability as dur +from src.turn_state import TurnKey, TurnStateStore +from src.turn_state.durability import TurnDurability + + +@pytest.fixture(autouse=True) +def _fast_beats(monkeypatch): + monkeypatch.setattr(dur, "_HEARTBEAT_FLOOR_SECONDS", 0.05) + + +def make_store(tmp_path, ttl=0.4): + return TurnStateStore( + tmp_path / "hb" / "turns.sqlite3", blob_dir=tmp_path / "hb" / "b", + lease_ttl=ttl, + ) + + +from types import SimpleNamespace + + +def FakeMsg(): # noqa: N802 — message-shaped factory + return SimpleNamespace( + channel=SimpleNamespace(id="c1"), + author=SimpleNamespace(id="u1"), + id="m1", + guild=None, + content="hello", + ) + + +async def admit(store): + handle = await TurnDurability.admit( + store, message=FakeMsg(), system_prompt="s", tools=[], session_snapshot=None + ) + assert handle.enabled + return handle + + +async def test_beats_keep_the_lease_alive_past_the_ttl(tmp_path): + store = make_store(tmp_path, ttl=0.3) + handle = await admit(store) + await asyncio.sleep(0.8) # ~2.5x TTL, several real beats + (expires,) = store._conn.execute( + "SELECT lease_expires_at FROM turns" + ).fetchone() + assert expires > time.time() # still alive + await handle.settle_terminal(cancelled=False, is_error=False) + assert handle._heartbeat_task is None # stopped on settlement + store.close() + + +async def test_beat_exits_when_the_lease_is_stolen(tmp_path): + store = make_store(tmp_path, ttl=0.3) + handle = await admit(store) + task = handle._heartbeat_task + assert task is not None + store._conn.execute("UPDATE turns SET lease_token='stolen'") + store._conn.commit() + await asyncio.wait_for(task, timeout=5) # StaleTurnError branch → clean exit + handle._stop_heartbeats() + store.close() + + +async def test_beat_exits_when_the_store_dies(tmp_path): + store = make_store(tmp_path, ttl=0.3) + handle = await admit(store) + task = handle._heartbeat_task + assert task is not None + store._conn.close() # TurnStateUnavailableError branch + await asyncio.wait_for(task, timeout=5) + handle._stop_heartbeats() + + +async def test_mark_ops_manual_swallows_store_death(tmp_path): + store = make_store(tmp_path) + await admit(store) + store._conn.close() + assert store.mark_ops_manual_sync(TurnKey("discord", "c1", "m1"), "g") == 0 diff --git a/tests/test_turn_state_store.py b/tests/test_turn_state_store.py index 3b5e2c3f..5282f37c 100644 --- a/tests/test_turn_state_store.py +++ b/tests/test_turn_state_store.py @@ -554,3 +554,74 @@ def test_db_dir_and_blobs_are_owner_only(self, tmp_path): assert blob_dir_mode == 0o700 assert blob_mode == 0o600 s.close() + + +class TestFullFenceOnTurnWrites: + """Round-2 blocker #1 (PR #242): every turn write honors the complete + live fence — an expired-lease owner can neither checkpoint (which used + to silently RENEW the lease) nor settle.""" + + def _expire(self, store): + import time as _time + + store._conn.execute("UPDATE turns SET lease_expires_at=?", [_time.time() - 5]) + store._conn.commit() + + def test_expired_lease_checkpoint_is_rejected(self, store): + lease = _admit(store) + self._expire(store) + with pytest.raises(StaleTurnError): + store.checkpoint_sync(lease, {"continued": True}, progressed=True) + # And crucially it did NOT renew the lease. + (expires,) = store._conn.execute( + "SELECT lease_expires_at FROM turns" + ).fetchone() + import time as _time + + assert expires < _time.time() + + def test_expired_lease_settle_is_rejected(self, store): + lease = _admit(store) + store.record_intents_sync( + lease, 0, [{"tool_call_id": "c1", "tool_name": "t", "tool_input": {}}] + ) + store.mark_running_sync(lease, 0, "c1") + self._expire(store) + with pytest.raises(StaleTurnError): + store.settle_op_sync(lease, 0, "c1", state=OpState.APPLIED, + result_text="ok") + + def test_heartbeat_extends_across_the_ttl(self, tmp_path): + import time as _time + + s = TurnStateStore(tmp_path / "hb.sqlite3", blob_dir=tmp_path / "b", + lease_ttl=0.5) + lease = _admit(s) + _time.sleep(0.3) + s.heartbeat_sync(lease) + _time.sleep(0.3) # past the original expiry, inside the beaten one + s.checkpoint_sync(lease, {"alive": True}, progressed=True) # no raise + s.close() + + +class TestMarkOpsManual: + def test_unknowns_move_to_manual(self, store): + lease = _admit(store) + store.record_intents_sync( + lease, 0, + [{"tool_call_id": "a", "tool_name": "t", "tool_input": {}}, + {"tool_call_id": "b", "tool_name": "t", "tool_input": {}}], + ) + store.settle_op_sync(lease, 0, "a", state=OpState.APPLIED, result_text="r") + store._conn.execute( + "UPDATE operations SET state=? WHERE tool_call_id='b'", + [OpState.OUTCOME_UNKNOWN], + ) + store._conn.commit() + moved = store.mark_ops_manual_sync(KEY, lease.generation) + assert moved == 1 + states = dict(store._conn.execute( + "SELECT tool_call_id, state FROM operations" + ).fetchall()) + assert states["a"] == OpState.APPLIED # settled rows untouched + assert states["b"] == OpState.MANUAL_RESOLUTION_REQUIRED diff --git a/tests/test_write_invariant_integration.py b/tests/test_write_invariant_integration.py index ce838d9a..2facb611 100644 --- a/tests/test_write_invariant_integration.py +++ b/tests/test_write_invariant_integration.py @@ -525,3 +525,97 @@ async def test_suspended_identity_redelivery_points_at_resume(self, tmp_path): bot.tool_executor.execute.assert_not_awaited() (status,) = store._conn.execute("SELECT status FROM turns").fetchone() assert status == TurnStatus.SUSPENDED # untouched + + +class TestLeaseHeartbeats: + async def test_long_tool_outlives_the_lease_ttl(self, tmp_path): + """Round-2 blocker #1 (PR #242): the owner beats the lease alive, so + a tool longer than the TTL settles fine and the turn completes.""" + import asyncio + + from src.tools.result_validator import ToolResult + + bot, fake, store = build_with_store( + [ + tool_call_response(("run_command", {"command": "slow"})), + text_response("survived the long tool"), + ], + tmp_path, + ) + store.lease_ttl = 0.4 # heartbeat interval becomes ~5s floor... force lower + # The durability heartbeat floors at 5s; drop the floor via the + # store's ttl AND patch the interval floor for the test. + import src.turn_state.durability as dur_mod + + original_start = dur_mod.TurnDurability._start_heartbeats + + def fast_start(self): + if self._store is None or self._lease is None: + return + store_, lease_ = self._store, self._lease + + async def _beat(): + while True: + await asyncio.sleep(0.1) + if not self.enabled: + return + try: + await asyncio.to_thread(store_.heartbeat_sync, lease_) + except Exception: + return + + self._heartbeat_task = asyncio.get_running_loop().create_task(_beat()) + + dur_mod.TurnDurability._start_heartbeats = fast_start + try: + async def slow_tool(tool_name, tool_input, *, user_id=None): + await asyncio.sleep(1.0) # 2.5x the lease TTL + return ToolResult(output="done slowly", tool_name=tool_name) + + bot.tool_executor.execute = slow_tool + text, _, is_error, *_ = await run_loop(bot, FakeMessage("go")) + assert text == "survived the long tool" + assert is_error is False + (status,) = store._conn.execute("SELECT status FROM turns").fetchone() + assert status == TurnStatus.TERMINAL_COMPLETED + finally: + dur_mod.TurnDurability._start_heartbeats = original_start + + +class TestAdmissionFailClosed: + async def test_runtime_admission_failure_refuses_execution(self, tmp_path): + """Round-2 blocker #2 (PR #242): once the store was wired available, + an admission I/O failure refuses execution — never a fresh-identity + guess.""" + import sqlite3 as _sqlite3 + from unittest.mock import AsyncMock + + bot, fake, store = build_with_store([text_response("never")], tmp_path) + + def raising_admit(*a, **k): + raise _sqlite3.OperationalError("disk I/O error") + + store.admit_turn_sync = raising_admit + bot.tool_executor.execute = AsyncMock() + text, _, is_error, tools, _ = await run_loop(bot, FakeMessage("go")) + assert "can't verify" in text + assert tools == [] + bot.tool_executor.execute.assert_not_awaited() + assert len(fake.calls) == 0 # no generation either + + async def test_store_unavailable_disposition_also_refuses(self, tmp_path): + from unittest.mock import AsyncMock + + bot, fake, store = build_with_store([text_response("never")], tmp_path) + store.admit_turn_sync = lambda *a, **k: (None, "store_unavailable") + bot.tool_executor.execute = AsyncMock() + text, *_ = await run_loop(bot, FakeMessage("go")) + assert "can't verify" in text + assert len(fake.calls) == 0 + + async def test_feature_off_still_runs_legacy(self, tmp_path): + bot, fake, store = build_with_store([text_response("legacy ok")], tmp_path) + bot.tool_loop._turn_store = None # durability off at wiring + text, _, is_error, *_ = await run_loop(bot, FakeMessage("go")) + assert text == "legacy ok" + assert is_error is False From 80fdf854109733d3f11bbd746d29d4fcf0c98e5e Mon Sep 17 00:00:00 2001 From: ATCharpentier Date: Thu, 30 Jul 2026 16:37:02 -0400 Subject: [PATCH 12/18] style: hoist SimpleNamespace import (E402) --- tests/test_turn_durability_heartbeat.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/test_turn_durability_heartbeat.py b/tests/test_turn_durability_heartbeat.py index cbc305d9..d4e715c1 100644 --- a/tests/test_turn_durability_heartbeat.py +++ b/tests/test_turn_durability_heartbeat.py @@ -10,6 +10,7 @@ import asyncio import time +from types import SimpleNamespace import pytest @@ -30,9 +31,6 @@ def make_store(tmp_path, ttl=0.4): ) -from types import SimpleNamespace - - def FakeMsg(): # noqa: N802 — message-shaped factory return SimpleNamespace( channel=SimpleNamespace(id="c1"), From 059e00c58b8682fbabe121d4fd92f9dde560a296 Mon Sep 17 00:00:00 2001 From: ATCharpentier Date: Thu, 30 Jul 2026 17:12:46 -0400 Subject: [PATCH 13/18] =?UTF-8?q?fix:=20review=20round=203=20=E2=80=94=20a?= =?UTF-8?q?ll=20six=20deviations=20(PR=20#242)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 1 (revision race): the in-memory lease revision now advances UNDER the same write lock as the database mutation — a heartbeat interleaved with a checkpoint can no longer read a stale revision, falsely StaleTurnError, and stop beating a lease the turn still owns. Pinned with 40 interleaved concurrent checkpoint+heartbeat cycles. 2 (wired-store death): TurnDurability.admit distinguishes store-None (durability off at wiring — legacy) from a present-but-dead store (was wired available; identity unverifiable — refuses with admission_error). 3 (session TOCTOU): the authoritative session-advance check now runs UNDER the channel lock immediately before reconstruction/acquisition; the pre-lock check only avoids queueing pointlessly. Pinned: a session advancing while auto-resume queues on the lock stands it down. 4 (op-details copy): _op_tool_details persists through the same tool-aware + recursive scrub as the transcript and trajectory copies (input rewritten only when present; result/error strings pattern- scrubbed; non-dict legacy entries pass through). Pinned: nested auth values never persist raw in ANY copy. 5 (structural validation): codec.validate_payload runs inside restore_field_values — codec version, policy, fields envelope, presence of every persisted field, and construction-critical types are checked BEFORE any lease exists; CheckpointInvalidError terminally rejects. The '{}' payload repro now lands TERMINAL_REJECTED, never bounced back to SUSPENDED for an infinite retry loop. 6 (zero-LLM resume path): the explicit-resume check moved ahead of prompt/history assembly in intake — get_task_history (whose compaction can invoke the LLM) is never called for a resume trigger. Pinned at the MessagePipeline level with an instrumented get_task_history proving zero calls end-to-end. Suite 7,914/5, lint/type gates new=0, coverage findings=0 at 88.4%. --- src/discord/intake_pipeline.py | 109 ++++++++++++++---------- src/discord/turn_resume.py | 22 ++++- src/turn_state/codec.py | 82 +++++++++++++++++- src/turn_state/durability.py | 11 ++- src/turn_state/store.py | 9 +- tests/test_resume_admission.py | 109 ++++++++++++++++++++++++ tests/test_turn_checkpoint_codec.py | 16 ++++ tests/test_turn_durability_heartbeat.py | 39 +++++++++ 8 files changed, 341 insertions(+), 56 deletions(-) diff --git a/src/discord/intake_pipeline.py b/src/discord/intake_pipeline.py index 0b6c8d87..4f676b1f 100644 --- a/src/discord/intake_pipeline.py +++ b/src/discord/intake_pipeline.py @@ -559,63 +559,80 @@ async def _run_inner( ) self._sessions.remove_last_message(channel_id, "user") return - _trace = self._turn_recorder._new_context_trace() - if _trace is not None: - with _trace.phase("system_prompt"): + # A bare `resume`/`continue` with preserved work in this + # channel resumes the suspended turn instead of starting a + # fresh one — the trigger is consumed as a command and never + # enters the frozen transcript. Checked BEFORE prompt/history + # assembly (round-3 deviation #6, PR #242): get_task_history + # can trigger compaction's LLM call, and a resume that halts + # on unresolved operations must be a genuine zero-LLM path. + _resumed = None + if self._turn_resume is not None: + try: + _resumed = await self._turn_resume.try_explicit_resume(message) + except Exception: + log.exception( + "Explicit resume attempt failed — falling through " + "to a normal turn" + ) + _resumed = None + _trace = None + _sp = None + task_history: list = [] + if _resumed is None: + _trace = self._turn_recorder._new_context_trace() + if _trace is not None: + with _trace.phase("system_prompt"): + _sp = self._prompt_builder.build_full_prompt( + channel=message.channel, + user_id=user_id, + query=content, + trace=_trace, + ) + else: _sp = self._prompt_builder.build_full_prompt( channel=message.channel, user_id=user_id, query=content, - trace=_trace, ) - else: - _sp = self._prompt_builder.build_full_prompt( - channel=message.channel, - user_id=user_id, - query=content, - ) - log.info("Routing to Codex with tools") - # Use abbreviated history to reduce poisoning from stale responses - # (get_task_history handles compaction internally) - # Pass current message content for relevance scoring — - # older messages unrelated to the current query are dropped - if _trace is not None: - with _trace.phase("history"): + log.info("Routing to Codex with tools") + # Use abbreviated history to reduce poisoning from stale + # responses (get_task_history handles compaction + # internally). Pass current message content for relevance + # scoring — older messages unrelated to the current + # query are dropped + if _trace is not None: + with _trace.phase("history"): + task_history = await self._sessions.get_task_history( + channel_id, + max_messages=160, + current_query=content, + trace=_trace, + ) + else: task_history = await self._sessions.get_task_history( channel_id, max_messages=160, current_query=content, - trace=_trace, ) - else: - task_history = await self._sessions.get_task_history( - channel_id, - max_messages=160, - current_query=content, - ) - if image_blocks and task_history and task_history[-1]["role"] == "user": - last = task_history[-1] - text = ( - last["content"] - if isinstance(last["content"], str) - else str(last["content"]) - ) - # Vision turns legitimately swap str content for - # a block list; the LLM layer handles both shapes. - task_history[-1] = { - "role": "user", - "content": image_blocks + [{"type": "text", "text": text}], # type: ignore[dict-item] - } - log.info("Attached %d image(s) to message for Claude vision", len(image_blocks)) + if image_blocks and task_history and task_history[-1]["role"] == "user": + last = task_history[-1] + text = ( + last["content"] + if isinstance(last["content"], str) + else str(last["content"]) + ) + # Vision turns legitimately swap str content for + # a block list; the LLM layer handles both shapes. + task_history[-1] = { + "role": "user", + "content": image_blocks + [{"type": "text", "text": text}], # type: ignore[dict-item] + } + log.info( + "Attached %d image(s) to message for Claude vision", + len(image_blocks), + ) try: - # A bare `resume`/`continue` with preserved work in this - # channel resumes the suspended turn instead of starting - # a fresh one — the trigger is consumed as a command and - # never enters the frozen transcript. Inherits this - # pipeline's lock/delivery/session plumbing wholesale. - _resumed = None - if self._turn_resume is not None: - _resumed = await self._turn_resume.try_explicit_resume(message) if _resumed is not None: ( response, diff --git a/src/discord/turn_resume.py b/src/discord/turn_resume.py index 88069e26..bf0db859 100644 --- a/src/discord/turn_resume.py +++ b/src/discord/turn_resume.py @@ -156,14 +156,16 @@ async def _auto_resume_waiter( admission = breaker.acquire_attempt() if not isinstance(admission, float): breaker.abandon(admission) # the resume re-acquires for real - current = self._session_len(key.channel_id) - if current not in allowed: + if self._session_len(key.channel_id) not in allowed: log.info( "Auto-resume for %s stands down: session advanced or " "unreadable (explicit resume still available)", key, ) return - await self._run_auto_resume(key, row) + # The authoritative re-check happens UNDER the channel lock + # inside _run_auto_resume — this pre-check just avoids + # queueing on a busy channel for nothing. + await self._run_auto_resume(key, row, allowed) return log.info("Auto-resume waiter for %s expired", key) @@ -181,7 +183,9 @@ def _session_len(self, channel_id: str) -> int: except Exception: return -1 - async def _run_auto_resume(self, key: TurnKey, row: dict) -> None: + async def _run_auto_resume( + self, key: TurnKey, row: dict, allowed: set[int] + ) -> None: if self._unresolved_ops(row): # Never auto-continue over ambiguous external effects — a human # must look at this (explicit resume delivers the details). @@ -194,6 +198,16 @@ async def _run_auto_resume(self, key: TurnKey, row: dict) -> None: if self._channel_state.active_requests.get(key.channel_id): log.info("Auto-resume for %s stands down: channel busy", key) return + # Authoritative session re-check UNDER the lock (round-3 + # deviation #3, PR #242): a message that advanced the session + # while we queued for this lock must stand auto-resume down — + # the pre-lock check alone was a TOCTOU window. + if self._session_len(key.channel_id) not in allowed: + log.info( + "Auto-resume for %s stands down: session advanced while " + "waiting for the channel lock", key, + ) + return st, message, reason = await self._validate_and_rebuild(key, row) if st is None: log.info("Auto-resume for %s rejected: %s", key, reason) diff --git a/src/turn_state/codec.py b/src/turn_state/codec.py index e1ac18ac..0f49f400 100644 --- a/src/turn_state/codec.py +++ b/src/turn_state/codec.py @@ -133,6 +133,28 @@ def scrub_stored_tool_input(tool_name: str, tool_input: Any) -> Any: return _deep_scrub_strings(tool_input) +def _scrub_op_details(details: list) -> list: + """The storage scrub for the op-details copy (round-3 deviation #4, + PR #242): the loop applies the tool-aware redaction when building + `_op_tool_details`, but persistence must ALSO compose the recursive + secret scrub so nested values match the transcript/trajectory copies.""" + out = [] + for detail in details: + if isinstance(detail, dict): + row = dict(detail) + if "input" in row: + row["input"] = scrub_stored_tool_input( + str(row.get("tool") or ""), row.get("input") + ) + for key in ("result", "error"): + if isinstance(row.get(key), str): + row[key] = scrub_output_secrets(row[key]) + out.append(row) + else: + out.append(detail) + return out + + def _scrub_tool_use_inputs(obj: Any) -> Any: """Apply the storage scrub to assistant ``tool_use`` blocks in the transcript (name-aware — the block carries the tool name).""" @@ -343,7 +365,7 @@ def snapshot_chat_turn(st, *, store_blob, generation_seq: int, extra: dict | Non "pending_image_blocks": _externalize_blocks( list(st.pending_image_blocks), store_blob ), - "_op_tool_details": list(st._op_tool_details), + "_op_tool_details": _scrub_op_details(list(st._op_tool_details)), "_pending_validations": list(st._pending_validations), "_validation_required": st._validation_required, "_validation_retries": st._validation_retries, @@ -360,13 +382,65 @@ def snapshot_chat_turn(st, *, store_blob, generation_seq: int, extra: dict | Non return payload +class CheckpointInvalidError(ValueError): + """The payload is structurally unusable — terminally reject, never + retry (round-3 deviation #5, PR #242).""" + + +# Fields whose restored TYPE the codec asserts before any lease exists — +# construction failures after acquisition used to bounce the row back to +# SUSPENDED forever. +_REQUIRED_FIELD_TYPES: dict[str, type | tuple[type, ...]] = { + "system_prompt": str, + "messages": list, + "user_id": str, + "chat_cap": int, + "iteration": int, + "stuck_tracker": dict, + "_trajectory": dict, + "_ch_id": str, + "_req_id": str, +} + + +def validate_payload(payload: Any) -> None: + """Structural schema validation, run BEFORE lease acquisition. + + Raises :class:`CheckpointInvalidError` on any deviation — codec + version, policy name, the fields envelope, presence of every persisted + field, and the basic types construction depends on. + """ + if not isinstance(payload, dict): + raise CheckpointInvalidError("payload is not an object") + version = payload.get("codec_version") + if not isinstance(version, int) or version > CODEC_VERSION or version < 1: + raise CheckpointInvalidError(f"unsupported codec_version: {version!r}") + if payload.get("policy") != "chat": + raise CheckpointInvalidError(f"unsupported policy: {payload.get('policy')!r}") + fields = payload.get("fields") + if not isinstance(fields, dict): + raise CheckpointInvalidError("missing fields envelope") + missing = PERSISTED_FIELDS - fields.keys() + if missing: + raise CheckpointInvalidError(f"missing persisted fields: {sorted(missing)}") + for name, expected in _REQUIRED_FIELD_TYPES.items(): + if not isinstance(fields.get(name), expected): + raise CheckpointInvalidError( + f"field {name!r} has invalid type {type(fields.get(name)).__name__}" + ) + + def restore_field_values(payload: dict, *, load_blob, stuck_tracker_cls) -> dict: """Persisted-field constructor kwargs for `_ChatTurn`. - The resume flow combines these with its RECONSTRUCTED fields (live - message, fresh cancel event, current-policy tools, fresh trace, policy - object) — see the classification at the top of this module. + Validates the payload structurally first (CheckpointInvalidError on + deviation — the resume flow terminally rejects BEFORE acquiring any + lease). The resume flow combines the result with its RECONSTRUCTED + fields (live message, fresh cancel event, current-policy tools, fresh + trace, policy object) — see the classification at the top of this + module. """ + validate_payload(payload) fields = dict(payload.get("fields") or {}) fields["messages"] = _inline_blocks(fields.get("messages") or [], load_blob) fields["pending_image_blocks"] = _inline_blocks( diff --git a/src/turn_state/durability.py b/src/turn_state/durability.py index 08096ceb..f0255e2f 100644 --- a/src/turn_state/durability.py +++ b/src/turn_state/durability.py @@ -115,8 +115,17 @@ async def admit( back with ``blocked`` set and the loop must refuse fresh execution — a redelivered message must never re-run its effects unledgered. """ - if store is None or not store.available: + if store is None: + # Durability was off (or failed) at wiring — legacy run is the + # designed behavior; no ledger can exist to contradict it. return cls.disabled() + if not store.available: + # The store WAS wired available (wiring nulls out failed inits) + # and has since died: this identity cannot be checked — refuse + # (round-3 deviation #2, PR #242). + handle = cls.disabled() + handle.blocked = "admission_error" + return handle try: key = TurnKey( source="discord", diff --git a/src/turn_state/store.py b/src/turn_state/store.py index 15a1fc43..11f9ec2f 100644 --- a/src/turn_state/store.py +++ b/src/turn_state/store.py @@ -287,6 +287,14 @@ def _fenced_update( lease.token, TurnStatus.ACTIVE, time.time()], ) conn.commit() + if cur.rowcount == 1: + # The in-memory revision must advance under the SAME + # lock as the database mutation (round-3 deviation #1, + # PR #242): a concurrent heartbeat sneaking between + # commit and this assignment read the stale revision, + # got StaleTurnError, and stopped beating a lease the + # turn still owned. + lease.revision = new_revision except sqlite3.Error as exc: raise TurnStateUnavailableError(f"turn-state write failed: {exc}") from exc if cur.rowcount != 1: @@ -294,7 +302,6 @@ def _fenced_update( f"turn {lease.key} generation {lease.generation[:8]} " f"rev {lease.revision}: fence lost" ) - lease.revision = new_revision def _op_where(self, lease: TurnLease) -> tuple[str, list]: return ( diff --git a/tests/test_resume_admission.py b/tests/test_resume_admission.py index a6ae19c9..777293c7 100644 --- a/tests/test_resume_admission.py +++ b/tests/test_resume_admission.py @@ -457,3 +457,112 @@ async def test_suspension_bookkeeping_is_plus_one_and_auto_resume_admits( ) for task in list(h2.manager._waiters.values()): task.cancel() + + +class TestSessionRecheckUnderLock: + async def test_advance_while_waiting_for_the_lock_stands_down(self, tmp_path): + """Round-3 deviation #3 (PR #242): the authoritative session check + runs UNDER the channel lock — a message advancing the session while + auto-resume queues for the lock must stand it down.""" + h, original = await suspend_turn(tmp_path) + heal_capacity(h, text_response("stale reply that must never send")) + rows = h.store.list_suspended_sync("discord") + row = h.store.load_resumable_sync( + __import__("src.turn_state", fromlist=["TurnKey"]).TurnKey( + "discord", rows[0]["channel_id"], rows[0]["message_id"] + ) + ) + from src.turn_state import TurnKey + + key = TurnKey("discord", rows[0]["channel_id"], rows[0]["message_id"]) + ch_id = key.channel_id + baseline = h.manager._session_len(ch_id) + allowed = {baseline, baseline + 1} + + lock = h.manager._channel_lock(ch_id) + await lock.acquire() # another turn holds the channel + resume_task = asyncio.get_running_loop().create_task( + h.manager._run_auto_resume(key, row, allowed) + ) + await asyncio.sleep(0.05) # auto-resume is now queued on the lock + # The session advances by a full turn while auto-resume waits. + h.bot.sessions.add_message(ch_id, "user", "new topic") + h.bot.sessions.add_message(ch_id, "assistant", "answered") + lock.release() + await asyncio.wait_for(resume_task, timeout=5) + assert h.row()[0] == TurnStatus.SUSPENDED # stood down under the lock + + +class TestStructuralPayloadValidation: + async def test_empty_object_payload_terminally_rejects(self, tmp_path): + """Round-3 deviation #5 (PR #242): a syntactically-valid but + structurally-invalid payload rejects BEFORE any lease exists — + never bounced back to SUSPENDED for an infinite retry loop.""" + h, original = await suspend_turn(tmp_path) + h.store._conn.execute("UPDATE turns SET payload='{}'") + h.store._conn.commit() + result = await h.manager.try_explicit_resume(resume_msg(original)) + assert result is not None + assert "could not be restored" in result[0] + (status,) = h.store._conn.execute("SELECT status FROM turns").fetchone() + assert status == TurnStatus.TERMINAL_REJECTED + + def test_validate_payload_rejects_each_structural_deviation(self): + import pytest as _pytest + + from src.turn_state.codec import ( + CODEC_VERSION, + CheckpointInvalidError, + validate_payload, + ) + + good_fields = {name: None for name in __import__( + "src.turn_state.codec", fromlist=["PERSISTED_FIELDS"] + ).PERSISTED_FIELDS} + good_fields.update({ + "system_prompt": "s", "messages": [], "user_id": "u", + "chat_cap": 1, "iteration": 0, "stuck_tracker": {}, + "_trajectory": {}, "_ch_id": "c", "_req_id": "r", + }) + base = {"codec_version": CODEC_VERSION, "policy": "chat", + "generation_seq": 0, "fields": good_fields} + validate_payload(base) # sane baseline passes + + for broken in ( + {}, # everything missing + "not-an-object", # payload must be a dict + {**base, "codec_version": CODEC_VERSION + 1}, # future codec + {**base, "policy": "loop"}, # wrong policy + {**base, "fields": "nope"}, # fields envelope must be a dict + {**base, "fields": {}}, # missing persisted fields + {**base, "fields": {**good_fields, "messages": "not-a-list"}}, + ): + with _pytest.raises(CheckpointInvalidError): + validate_payload(broken) + + +class TestExplicitResumeOrdering: + async def test_resume_trigger_skips_history_compaction(self, tmp_path): + """Round-3 deviation #6 (PR #242): the explicit-resume check runs + BEFORE prompt/history assembly — get_task_history (which can invoke + the compaction LLM) must never be called for a resume trigger.""" + h, original = await suspend_turn(tmp_path) + heal_capacity(h, text_response("Resumed through the pipeline.")) + h.bot.pipeline._turn_resume = h.manager + + calls = [] + real_gth = h.bot.sessions.get_task_history + + async def recording_gth(*a, **k): + calls.append(a) + return await real_gth(*a, **k) + + h.bot.sessions.get_task_history = recording_gth + trigger = resume_msg(original) + await h.bot.pipeline.run(trigger, trigger.content) + assert calls == [] # resume never touched history assembly + assert h.row()[0] == TurnStatus.TERMINAL_COMPLETED + assert any( + "Resumed through the pipeline." in (r["content"] or "") + for r in trigger.replies + ) diff --git a/tests/test_turn_checkpoint_codec.py b/tests/test_turn_checkpoint_codec.py index 5b7a78a4..4fc26022 100644 --- a/tests/test_turn_checkpoint_codec.py +++ b/tests/test_turn_checkpoint_codec.py @@ -333,3 +333,19 @@ def test_non_string_scalars_pass_through_the_scrub(self): "run_command", {"count": 3, "flag": True, "ratio": 1.5, "none": None} ) assert out == {"count": 3, "flag": True, "ratio": 1.5, "none": None} + + def test_op_details_copy_is_scrubbed_too(self): + """Round-3 deviation #4 (PR #242): the _op_tool_details copy gets + the same tool-aware + recursive scrub as every other copy.""" + blobs, store, load = _blob_dict() + st = _full_turn() + secret = "sk-" + "c" * 24 + st._op_tool_details.append({ + "tool": "http_post", + "input": {"headers": {"Authorization": f"token={secret}"}}, + "result": f"posted with api_key={secret}", + "error": False, + }) + st._op_tool_details.append("legacy-non-dict-entry") # passthrough arm + payload = snapshot_chat_turn(st, store_blob=store, generation_seq=3) + assert secret not in json.dumps(payload) diff --git a/tests/test_turn_durability_heartbeat.py b/tests/test_turn_durability_heartbeat.py index d4e715c1..ebcaebfa 100644 --- a/tests/test_turn_durability_heartbeat.py +++ b/tests/test_turn_durability_heartbeat.py @@ -89,3 +89,42 @@ async def test_mark_ops_manual_swallows_store_death(tmp_path): await admit(store) store._conn.close() assert store.mark_ops_manual_sync(TurnKey("discord", "c1", "m1"), "g") == 0 + + +async def test_concurrent_checkpoints_and_heartbeats_never_false_stale(tmp_path): + """Round-3 deviation #1 (PR #242): the in-memory revision advances under + the same lock as the DB mutation, so a heartbeat interleaved with + checkpoints can never read a stale revision and falsely lose the lease.""" + import json + + store = make_store(tmp_path, ttl=30.0) + handle = await admit(store) + handle._stop_heartbeats() # drive beats manually, interleaved + lease = handle.lease + + async def checkpoints(): + for i in range(40): + await asyncio.to_thread( + store.checkpoint_sync, lease, {"n": i}, progressed=False + ) + + async def heartbeats(): + for _ in range(40): + await asyncio.to_thread(store.heartbeat_sync, lease) + + await asyncio.gather(checkpoints(), heartbeats()) # no StaleTurnError + (payload,) = store._conn.execute("SELECT payload FROM turns").fetchone() + assert json.loads(payload)["n"] == 39 + store.close() + + +async def test_wired_store_that_died_refuses_admission(tmp_path): + """Round-3 deviation #2 (PR #242): a store that WAS wired available but + has since died must refuse execution, not run legacy.""" + store = make_store(tmp_path) + store.close() # available -> False after successful wiring + handle = await TurnDurability.admit( + store, message=FakeMsg(), system_prompt="s", tools=[], session_snapshot=None + ) + assert handle.enabled is False + assert handle.blocked == "admission_error" From b145ebc13929829d42ea3c5ea4c15f9253b799ca Mon Sep 17 00:00:00 2001 From: ATCharpentier Date: Thu, 30 Jul 2026 17:54:31 -0400 Subject: [PATCH 14/18] =?UTF-8?q?fix:=20review=20round=204=20=E2=80=94=20a?= =?UTF-8?q?ll=20four=20blocking=20classes=20(PR=20#242)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 1 (revision ABA): every revision value — the expected WHERE revision, the new SET revision, and the shared lease publish — now derives from ONE read taken under the write lock. The pre-lock capture let a delayed heartbeat pair a fresh WHERE with a stale SET and REGRESS both the row and the lease (Odin forced DB 1→0). Pinned: a stale lease copy is fenced out entirely and can never regress the row. 2 (validation + fall-through): validate_payload now requires an exact non-negative non-bool generation_seq, exact-int typed fields, and transcript shape (every entry a message object; content a string or a list of block objects — messages=[None] and content=[42] reject). ALL fallible reconstruction (restore, tool derivation, transcript repair) lives inside one pre-acquisition rejection boundary. try_explicit_resume now has a hard contract: once the trigger is recognized against preserved work it NEVER raises and NEVER returns None — an internal failure returns a notice tuple, so a recognized resume can never fall through into a fresh normal turn (Odin reproduced a fresh LLM response; pinned: zero fresh generation, TERMINAL_REJECTED, no SUSPENDED loop). 3 (monotonic session fence): SessionManager gains a per-channel mutation revision bumped on every semantic mutation (append, removal, compaction, fallback compaction, secret scrub, tombstone) — process-local by design, matching its only consumer. Auto-resume anchors on the revision instead of len(messages), closing the add-then-remove ABA (Odin resumed despite an intervening message; pinned: count restored, revision +2, stands down). 4 (key-aware redaction): _deep_scrub_strings adds recursive case-insensitive sensitive-KEY redaction (normalized exact match — auth redacts, author does not) so opaque credentials with no token-shaped signature never persist under Authorization/password/api-key/etc., at any nesting depth, in any persisted copy. Pinned with opaque values. Suite 7,921/5, lint/type gates new=0, coverage findings=0 at 88.4%. --- src/discord/turn_resume.py | 97 ++++++++++++++--------- src/sessions/manager.py | 21 +++++ src/turn_state/codec.py | 77 ++++++++++++++++-- src/turn_state/store.py | 17 ++-- tests/test_resume_admission.py | 117 +++++++++++++++++++++++++++- tests/test_turn_checkpoint_codec.py | 27 +++++++ tests/test_turn_state_store.py | 23 ++++++ 7 files changed, 330 insertions(+), 49 deletions(-) diff --git a/src/discord/turn_resume.py b/src/discord/turn_resume.py index bf0db859..d389ad82 100644 --- a/src/discord/turn_resume.py +++ b/src/discord/turn_resume.py @@ -101,15 +101,15 @@ def on_turn_suspended(self, key: TurnKey, generation: str) -> None: """Called by the tool loop when a turn suspends. In-process only.""" if not self._auto_resume_enabled: return - # Session length AT SUSPENSION, captured synchronously inside the - # suspending turn (which still holds the channel lock) — the - # advance-check anchor (review blocker #6a, PR #242). - suspend_len = self._session_len(key.channel_id) + # Mutation revision AT SUSPENSION, captured synchronously inside + # the suspending turn (which still holds the channel lock) — the + # monotonic advance-check anchor (round-4 blocker #3, PR #242). + suspend_rev = self._session_revision(key.channel_id) existing = self._waiters.pop(key, None) if existing is not None: existing.cancel() task = asyncio.get_running_loop().create_task( - self._auto_resume_waiter(key, generation, suspend_len), + self._auto_resume_waiter(key, generation, suspend_rev), name=f"turn-resume:{key.channel_id}:{key.message_id}", ) self._waiters[key] = task @@ -123,19 +123,19 @@ def _cleanup(t: asyncio.Task, *, _key=key) -> None: task.add_done_callback(_cleanup) async def _auto_resume_waiter( - self, key: TurnKey, generation: str, suspend_len: int + self, key: TurnKey, generation: str, suspend_rev: int ) -> None: """Wait for capacity, then resume IF nothing else happened. - The advance check anchors on the session length captured AT - SUSPENSION (``suspend_len``). Production ordering (round-2 blocker - #3, PR #242): intake appends the USER message BEFORE the turn runs, - so the suspension capture already includes it, and the only - legitimate growth afterwards is the assistant preservation marker — - exactly +1 (a directly-driven turn adds 0). Anything else — or an - unreadable session — stands auto-resume down fail-safe; explicit - resume stays available. No sampling window exists for a stranger - message to sneak into the baseline. + The advance check anchors on the MONOTONIC session mutation + revision captured AT SUSPENSION (``suspend_rev``) — message count + can ABA back to an allowed value via removal or compaction; the + revision only grows (round-4 blocker #3, PR #242). Production + ordering: intake appends the USER message BEFORE the turn runs, so + the capture already includes it and the only legitimate growth is + the assistant preservation marker — exactly +1 (a directly-driven + turn adds 0). Anything else — or an unreadable session — stands + auto-resume down fail-safe; explicit resume stays available. Capacity detection is ACTIVE: a quiet breaker is never probed by anyone, so the waiter claims the probe slot itself when the cooldown @@ -147,7 +147,7 @@ async def _auto_resume_waiter( """ give_up_at = time.monotonic() + self._resume_ttl_hours * 3600.0 breaker = self._llm_gateway.capacity_breaker_for() - allowed = {suspend_len, suspend_len + 1} if suspend_len >= 0 else set() + allowed = {suspend_rev, suspend_rev + 1} if suspend_rev >= 0 else set() while time.monotonic() < give_up_at: await asyncio.sleep(_AUTO_POLL_SECONDS) row = await asyncio.to_thread(self._store.load_resumable_sync, key) @@ -156,7 +156,7 @@ async def _auto_resume_waiter( admission = breaker.acquire_attempt() if not isinstance(admission, float): breaker.abandon(admission) # the resume re-acquires for real - if self._session_len(key.channel_id) not in allowed: + if self._session_revision(key.channel_id) not in allowed: log.info( "Auto-resume for %s stands down: session advanced or " "unreadable (explicit resume still available)", key, @@ -172,14 +172,17 @@ async def _auto_resume_waiter( def _channel_lock(self, channel_id: str) -> asyncio.Lock: return self._channel_state.channel_locks.setdefault(channel_id, asyncio.Lock()) - def _session_len(self, channel_id: str) -> int: - """Peek the channel's session length WITHOUT get_or_create (which - bumps last_active on pure reads). SessionManager keeps its dict at - `_sessions`; a lookup failure returns -1 so a broken peek can never - satisfy the baseline-equality check and wrongly auto-resume.""" + def _session_revision(self, channel_id: str) -> int: + """The channel's MONOTONIC mutation watermark (round-4 blocker #3, + PR #242): message COUNT can ABA back to an allowed value via + removal/compaction; the revision only ever grows. A lookup failure + returns -1 so a broken peek can never satisfy the baseline check + and wrongly auto-resume.""" try: - session = getattr(self._sessions, "_sessions", {}).get(channel_id) - return len(session.messages) if session is not None else 0 + fn = getattr(self._sessions, "mutation_revision", None) + if not callable(fn): + return -1 + return int(fn(channel_id)) except Exception: return -1 @@ -202,7 +205,7 @@ async def _run_auto_resume( # deviation #3, PR #242): a message that advanced the session # while we queued for this lock must stand auto-resume down — # the pre-lock check alone was a TOCTOU window. - if self._session_len(key.channel_id) not in allowed: + if self._session_revision(key.channel_id) not in allowed: log.info( "Auto-resume for %s stands down: session advanced while " "waiting for the channel lock", key, @@ -273,6 +276,12 @@ async def try_explicit_resume(self, message: Any): Returns the run() result tuple, a notice tuple when resume was attempted but rejected, or None when this message is not a resume trigger (normal processing continues). + + Contract (round-4 blocker #2, PR #242): once a trigger IS + recognized against preserved work, this NEVER raises and NEVER + returns None — an internal failure returns a notice tuple, so a + recognized resume command can never fall through into a fresh + normal turn. """ content = getattr(message, "content", "") or "" if not self.is_resume_trigger(content): @@ -281,6 +290,19 @@ async def try_explicit_resume(self, message: Any): row_summary = await self._latest_suspended_for_channel(channel_id) if row_summary is None: return None # nothing to resume — treat as a normal message + try: + return await self._explicit_resume_recognized(message, row_summary) + except Exception: + log.exception("Explicit resume failed internally") + return ( + "I recognized the resume command, but resuming failed " + "internally. The preserved work is untouched — try `resume` " + "again, or ask fresh.", + False, True, [], False, + ) + + async def _explicit_resume_recognized(self, message: Any, row_summary: dict): + channel_id = str(message.channel.id) key = TurnKey( source="discord", channel_id=channel_id, @@ -369,6 +391,11 @@ async def _validate_and_rebuild(self, key: TurnKey, row: dict): # repair, cancellation check — runs while the row is still # SUSPENDED, so a failure rejects/aborts cleanly instead of # stranding an ACTIVE row invisible to resumable queries. + # EVERY fallible reconstruction step lives inside this one + # pre-acquisition rejection boundary (round-4 blocker #2, PR #242): + # schema validation + restore, current-policy tool derivation, and + # transcript repair. A failure here rejects terminally — never a + # SUSPENDED bounce loop, never an escape past the resume flow. payload = row["payload"] try: fields = restore_field_values( @@ -376,22 +403,22 @@ async def _validate_and_rebuild(self, key: TurnKey, row: dict): load_blob=self._store.load_blob_sync, stuck_tracker_cls=self._tool_loop._stuck_loop_tracker_cls, ) + # Current security policy wins: tools re-derived from the live + # catalog + permission filter, never the persisted definitions. + tools = None + if self._get_config().tools.enabled: + tools = self._tool_catalog.merged_definitions() + tools = self._permissions.filter_tools(str(original.author.id), tools) + self._repair_unmatched_tool_use( + fields["messages"], row.get("operations") or [] + ) except Exception: - log.exception("Checkpoint restore failed — rejecting") + log.exception("Checkpoint reconstruction failed — rejecting") await asyncio.to_thread( self._store.reject_resumable_sync, key, "checkpoint unreadable" ) return None, None, "the checkpoint could not be restored" - # Current security policy wins: tools re-derived from the live - # catalog + permission filter, never the persisted definitions. - tools = None - if self._get_config().tools.enabled: - tools = self._tool_catalog.merged_definitions() - tools = self._permissions.filter_tools(str(original.author.id), tools) - - self._repair_unmatched_tool_use(fields["messages"], row.get("operations") or []) - cancel = self._channel_state.cancel_events.setdefault( key.channel_id, asyncio.Event() ) diff --git a/src/sessions/manager.py b/src/sessions/manager.py index 8dc9bba2..fe21d1fa 100644 --- a/src/sessions/manager.py +++ b/src/sessions/manager.py @@ -385,6 +385,12 @@ def __init__( self._continuity_source: dict[str, str] = {} self._sessions: dict[str, Session] = {} self._dirty: set[str] = set() + # Monotonic per-channel mutation watermark (PR #242, round-4 + # blocker #3): bumped on EVERY semantic history mutation (append, + # removal, compaction, secret scrub, tombstone) — unlike message + # count, it can never ABA back to a previous value. Process-local + # by design: its only consumer (auto-resume) is in-process only. + self._mutation_revisions: dict[str, int] = {} self._reflector = reflector self._reflection_tasks: set[asyncio.Task] = set() self._vector_store = vector_store @@ -464,6 +470,15 @@ def _restore_from_archive(self, channel_id: str) -> Session | None: log.error("All %d archive(s) for channel %s were unreadable", len(candidates), channel_id) return None + def _bump_mutation_revision(self, channel_id: str) -> None: + self._mutation_revisions[channel_id] = ( + self._mutation_revisions.get(channel_id, 0) + 1 + ) + + def mutation_revision(self, channel_id: str) -> int: + """Monotonic watermark of semantic history mutations (never ABAs).""" + return self._mutation_revisions.get(channel_id, 0) + def add_message( self, channel_id: str, role: str, content: str, *, user_id: str | None = None, @@ -476,6 +491,7 @@ def add_message( if role == "user" and user_id: session.last_user_id = user_id self._dirty.add(channel_id) + self._bump_mutation_revision(channel_id) def remove_last_message(self, channel_id: str, role: str) -> bool: """Remove the most recent message if it matches *role*. @@ -490,6 +506,7 @@ def remove_last_message(self, channel_id: str, role: str) -> bool: if session.messages[-1].role == role: session.messages.pop() self._dirty.add(channel_id) + self._bump_mutation_revision(channel_id) return True return False @@ -989,6 +1006,7 @@ async def _compact(self, session: Session) -> None: discarded = list(to_summarize) session.messages = to_keep self._dirty.add(session.channel_id) + self._bump_mutation_revision(session.channel_id) log.info( "Compacted %d messages into segment %s for channel %s (%d segments total)", len(discarded), segment["id"], session.channel_id, @@ -1024,6 +1042,7 @@ def _fallback_compact(self, session) -> None: return discarded = session.messages[:-keep] session.messages = session.messages[-keep:] + self._bump_mutation_revision(session.channel_id) # Build a deterministic summary from discarded messages user_ids = set() @@ -1084,6 +1103,7 @@ def _save_reset_epochs(self) -> None: def _tombstone(self, channel_id: str) -> None: """Drop the live session AND block archive restoration up to now.""" self._sessions.pop(channel_id, None) + self._bump_mutation_revision(channel_id) session_file = self.persist_dir / f"{channel_id}.json" try: session_file.unlink(missing_ok=True) @@ -1470,6 +1490,7 @@ def scrub_secrets(self, channel_id: str, content: str) -> bool: if removed: session.messages = filtered self._dirty.add(channel_id) + self._bump_mutation_revision(channel_id) log.warning( "Scrubbed %d message(s) containing secrets from channel %s", removed, diff --git a/src/turn_state/codec.py b/src/turn_state/codec.py index 0f49f400..aedff14a 100644 --- a/src/turn_state/codec.py +++ b/src/turn_state/codec.py @@ -100,14 +100,43 @@ def compute_content_digest(text: str) -> str: # ── storage redaction ──────────────────────────────────────────────── +# Case-insensitive sensitive-KEY redaction (round-4 blocker #4, PR #242): +# opaque credentials carry no `sk-`/`token=` signature for the pattern +# scrub, so any value stored under a credential-shaped key is redacted +# wholesale. Keys are normalized (lowercased, `-`/`_` stripped) and matched +# EXACTLY — "auth" redacts, "author" does not. +_SENSITIVE_KEYS = frozenset({ + "password", "passwd", "pwd", "passphrase", + "secret", "clientsecret", "secretkey", + "token", "apitoken", "accesstoken", "refreshtoken", "sessiontoken", + "idtoken", "bearertoken", "authtoken", + "apikey", "authorization", "auth", "bearer", + "credential", "credentials", "privatekey", "accesskey", "secretaccesskey", + "cookie", "setcookie", "sessionid", "csrftoken", +}) + + +def _is_sensitive_key(key: Any) -> bool: + if not isinstance(key, str): + return False + return key.lower().replace("-", "").replace("_", "") in _SENSITIVE_KEYS + + def _deep_scrub_strings(value: Any) -> Any: - """Recursive secret-pattern scrub over EVERY nested string value.""" + """Recursive storage redaction: secret-PATTERN scrub over every nested + string value, plus KEY-aware wholesale redaction — a value (of any + shape, including nested containers) stored under a sensitive key is + replaced entirely, because opaque credentials defeat pattern matching.""" if isinstance(value, str): return scrub_output_secrets(value) if isinstance(value, list): return [_deep_scrub_strings(v) for v in value] if isinstance(value, dict): - return {k: _deep_scrub_strings(v) for k, v in value.items()} + return { + k: ("[redacted:sensitive-key]" if _is_sensitive_key(k) + else _deep_scrub_strings(v)) + for k, v in value.items() + } return value @@ -412,11 +441,29 @@ def validate_payload(payload: Any) -> None: """ if not isinstance(payload, dict): raise CheckpointInvalidError("payload is not an object") + # bool subclasses int; a checkpoint carrying True where an ordinal + # belongs is corrupt, not truthy — the isinstance pairs are inlined so + # mypy narrows the comparisons. version = payload.get("codec_version") - if not isinstance(version, int) or version > CODEC_VERSION or version < 1: + if ( + not isinstance(version, int) + or isinstance(version, bool) + or version > CODEC_VERSION + or version < 1 + ): raise CheckpointInvalidError(f"unsupported codec_version: {version!r}") if payload.get("policy") != "chat": raise CheckpointInvalidError(f"unsupported policy: {payload.get('policy')!r}") + generation_seq = payload.get("generation_seq") + if ( + not isinstance(generation_seq, int) + or isinstance(generation_seq, bool) + or generation_seq < 0 + ): + raise CheckpointInvalidError(f"invalid generation_seq: {generation_seq!r}") + + def _exact_int(value: Any) -> bool: + return isinstance(value, int) and not isinstance(value, bool) fields = payload.get("fields") if not isinstance(fields, dict): raise CheckpointInvalidError("missing fields envelope") @@ -424,9 +471,29 @@ def validate_payload(payload: Any) -> None: if missing: raise CheckpointInvalidError(f"missing persisted fields: {sorted(missing)}") for name, expected in _REQUIRED_FIELD_TYPES.items(): - if not isinstance(fields.get(name), expected): + value = fields.get(name) + if expected is int: + if not _exact_int(value): + raise CheckpointInvalidError(f"field {name!r} is not an integer") + elif not isinstance(value, expected): + raise CheckpointInvalidError( + f"field {name!r} has invalid type {type(value).__name__}" + ) + # Transcript shape: every entry is a message dict; content is a string + # or a list of block dicts (round-4 blocker #2 — a [None] entry used to + # explode later in transcript repair, outside the rejection boundary). + for i, msg in enumerate(fields["messages"]): + if not isinstance(msg, dict) or not isinstance(msg.get("role"), str): + raise CheckpointInvalidError(f"messages[{i}] is not a message object") + content = msg.get("content") + if isinstance(content, list): + if not all(isinstance(block, dict) for block in content): + raise CheckpointInvalidError( + f"messages[{i}] has a non-object content block" + ) + elif not isinstance(content, str): raise CheckpointInvalidError( - f"field {name!r} has invalid type {type(fields.get(name)).__name__}" + f"messages[{i}] content has invalid type {type(content).__name__}" ) diff --git a/src/turn_state/store.py b/src/turn_state/store.py index 11f9ec2f..d764ff4e 100644 --- a/src/turn_state/store.py +++ b/src/turn_state/store.py @@ -271,7 +271,6 @@ def _fenced_update( renew itself through a checkpoint; the fence is identical for turn and ledger writes).""" conn = self._require() - new_revision = lease.revision + 1 if bump_revision else lease.revision sql = ( f"UPDATE turns SET {set_sql}, revision=? " "WHERE source=? AND channel_id=? AND message_id=? " @@ -280,20 +279,22 @@ def _fenced_update( ) try: with self._write_lock: + # EVERY revision value — the expected WHERE revision, the + # new SET revision, and the shared lease publish — derives + # from ONE read taken under this lock (round-4 blocker #1, + # PR #242): a pre-lock capture let a delayed heartbeat pair + # a fresh WHERE revision with a stale SET revision and + # REGRESS both the row and the lease. + expected_revision = lease.revision + new_revision = expected_revision + 1 if bump_revision else expected_revision cur = conn.execute( sql, [*params, new_revision, lease.key.source, lease.key.channel_id, - lease.key.message_id, lease.generation, lease.revision, + lease.key.message_id, lease.generation, expected_revision, lease.token, TurnStatus.ACTIVE, time.time()], ) conn.commit() if cur.rowcount == 1: - # The in-memory revision must advance under the SAME - # lock as the database mutation (round-3 deviation #1, - # PR #242): a concurrent heartbeat sneaking between - # commit and this assignment read the stale revision, - # got StaleTurnError, and stopped beating a lease the - # turn still owned. lease.revision = new_revision except sqlite3.Error as exc: raise TurnStateUnavailableError(f"turn-state write failed: {exc}") from exc diff --git a/tests/test_resume_admission.py b/tests/test_resume_admission.py index 777293c7..0558d9d5 100644 --- a/tests/test_resume_admission.py +++ b/tests/test_resume_admission.py @@ -476,7 +476,7 @@ async def test_advance_while_waiting_for_the_lock_stands_down(self, tmp_path): key = TurnKey("discord", rows[0]["channel_id"], rows[0]["message_id"]) ch_id = key.channel_id - baseline = h.manager._session_len(ch_id) + baseline = h.manager._session_revision(ch_id) allowed = {baseline, baseline + 1} lock = h.manager._channel_lock(ch_id) @@ -533,9 +533,18 @@ def test_validate_payload_rejects_each_structural_deviation(self): "not-an-object", # payload must be a dict {**base, "codec_version": CODEC_VERSION + 1}, # future codec {**base, "policy": "loop"}, # wrong policy + {**base, "generation_seq": "bad"}, # round-4: exact int required + {**base, "generation_seq": True}, # round-4: bools excluded + {**base, "generation_seq": -1}, {**base, "fields": "nope"}, # fields envelope must be a dict {**base, "fields": {}}, # missing persisted fields {**base, "fields": {**good_fields, "messages": "not-a-list"}}, + {**base, "fields": {**good_fields, "messages": [None]}}, # round-4 + {**base, "fields": {**good_fields, + "messages": [{"role": "user", "content": 42}]}}, + {**base, "fields": {**good_fields, + "messages": [{"role": "user", "content": [42]}]}}, + {**base, "fields": {**good_fields, "chat_cap": True}}, ): with _pytest.raises(CheckpointInvalidError): validate_payload(broken) @@ -566,3 +575,109 @@ async def recording_gth(*a, **k): "Resumed through the pipeline." in (r["content"] or "") for r in trigger.replies ) + + +class TestMonotonicSessionFence: + async def test_add_then_remove_aba_still_stands_down(self, tmp_path): + """Round-4 blocker #3 (PR #242): an add+remove pair returns the + MESSAGE COUNT to an allowed value, but the mutation revision only + grows — the ABA collision must stand auto-resume down.""" + h, original = await suspend_turn(tmp_path) + heal_capacity(h, text_response("stale reply that must never send")) + rows = h.store.list_suspended_sync("discord") + from src.turn_state import TurnKey + + key = TurnKey("discord", rows[0]["channel_id"], rows[0]["message_id"]) + row = h.store.load_resumable_sync(key) + ch_id = key.channel_id + baseline = h.manager._session_revision(ch_id) + allowed = {baseline, baseline + 1} + + lock = h.manager._channel_lock(ch_id) + await lock.acquire() + resume_task = asyncio.get_running_loop().create_task( + h.manager._run_auto_resume(key, row, allowed) + ) + await asyncio.sleep(0.05) + # ABA: append a message, then remove it — count restored, revision +2. + h.bot.sessions.add_message(ch_id, "user", "intervening") + h.bot.sessions.remove_last_message(ch_id, "user") + lock.release() + await asyncio.wait_for(resume_task, timeout=5) + assert h.row()[0] == TurnStatus.SUSPENDED # stood down on the revision + + def test_mutation_revision_is_monotonic_across_mutations(self, tmp_path): + h = Harness([text_response("x")], tmp_path) + sess = h.bot.sessions + assert sess.mutation_revision("chX") == 0 + sess.add_message("chX", "user", "a") + sess.add_message("chX", "assistant", "b") + assert sess.mutation_revision("chX") == 2 + sess.remove_last_message("chX", "assistant") + assert sess.mutation_revision("chX") == 3 # removal GROWS the watermark + + +class TestRecognizedResumeNeverFallsThrough: + async def test_internal_failure_returns_notice_not_fresh_turn(self, tmp_path): + """Round-4 blocker #2 (PR #242): once the resume trigger is + recognized against preserved work, an internal failure yields a + notice — NEVER None (which would run a fresh normal turn).""" + h, original = await suspend_turn(tmp_path) + + async def boom(message, row_summary): + raise RuntimeError("internal resume machinery exploded") + + h.manager._explicit_resume_recognized = boom + result = await h.manager.try_explicit_resume(resume_msg(original)) + assert result is not None + assert "resuming failed internally" in result[0] + assert result[2] is True # surfaced as an error, not silence + # The preserved work is untouched and still resumable. + assert h.row()[0] == TurnStatus.SUSPENDED + + async def test_bad_generation_seq_rejects_terminally_not_loop(self, tmp_path): + """Odin's round-4 repro: generation_seq="bad" used to pass + validation, fail post-acquisition, and bounce back to SUSPENDED + forever. Now it rejects terminally on the FIRST attempt.""" + import json as _json + + h, original = await suspend_turn(tmp_path) + (payload_text,) = h.store._conn.execute( + "SELECT payload FROM turns" + ).fetchone() + payload = _json.loads(payload_text) + payload["generation_seq"] = "bad" + h.store._conn.execute( + "UPDATE turns SET payload=?", [_json.dumps(payload)] + ) + h.store._conn.commit() + result = await h.manager.try_explicit_resume(resume_msg(original)) + assert result is not None + assert "could not be restored" in result[0] + (status,) = h.store._conn.execute("SELECT status FROM turns").fetchone() + assert status == TurnStatus.TERMINAL_REJECTED # no SUSPENDED loop + + async def test_none_message_entry_rejects_inside_the_boundary(self, tmp_path): + """Odin's round-4 repro: messages=[None] used to explode in + transcript repair OUTSIDE the rejection boundary and fall through + to a fresh turn. Now it rejects terminally, zero fresh generation.""" + import json as _json + + h, original = await suspend_turn(tmp_path) + heal_capacity(h, text_response("fresh turn that must never run")) + (payload_text,) = h.store._conn.execute( + "SELECT payload FROM turns" + ).fetchone() + payload = _json.loads(payload_text) + payload["fields"]["messages"] = [None] + h.store._conn.execute( + "UPDATE turns SET payload=?", [_json.dumps(payload)] + ) + h.store._conn.commit() + calls_before = len(h.fake.calls) + result = await h.manager.try_explicit_resume(resume_msg(original)) + assert result is not None + assert "could not be restored" in result[0] + assert len(h.fake.calls) == calls_before # no fresh LLM response + (status,) = h.store._conn.execute("SELECT status FROM turns").fetchone() + assert status == TurnStatus.TERMINAL_REJECTED diff --git a/tests/test_turn_checkpoint_codec.py b/tests/test_turn_checkpoint_codec.py index 4fc26022..a17cacad 100644 --- a/tests/test_turn_checkpoint_codec.py +++ b/tests/test_turn_checkpoint_codec.py @@ -349,3 +349,30 @@ def test_op_details_copy_is_scrubbed_too(self): st._op_tool_details.append("legacy-non-dict-entry") # passthrough arm payload = snapshot_chat_turn(st, store_blob=store, generation_seq=3) assert secret not in json.dumps(payload) + + def test_opaque_credentials_under_sensitive_keys_are_redacted(self): + """Round-4 blocker #4 (PR #242): key-aware redaction — opaque values + with no token-shaped signature must not survive under credential + keys, in ANY persisted copy, at any nesting depth.""" + opaque = "not-pattern-shaped-but-sensitive-7f4d" + encoded = self._payload_with( + "http_post", + { + "url": "https://x", + "headers": {"Authorization": f"Bearer {opaque}"}, + "password": opaque, + "nested": {"config": {"api-key": opaque}}, + }, + ) + assert opaque not in encoded + assert "https://x" in encoded # innocent values untouched + + def test_sensitive_key_matching_is_exact_not_substring(self): + from src.turn_state.codec import _is_sensitive_key + + assert _is_sensitive_key("Authorization") + assert _is_sensitive_key("API-Key") + assert _is_sensitive_key("refresh_token") + assert not _is_sensitive_key("author") # "auth" must not substring-match + assert not _is_sensitive_key("tokenizer") + assert not _is_sensitive_key(42) diff --git a/tests/test_turn_state_store.py b/tests/test_turn_state_store.py index 5282f37c..9bc1c1ab 100644 --- a/tests/test_turn_state_store.py +++ b/tests/test_turn_state_store.py @@ -625,3 +625,26 @@ def test_unknowns_move_to_manual(self, store): ).fetchall()) assert states["a"] == OpState.APPLIED # settled rows untouched assert states["b"] == OpState.MANUAL_RESOLUTION_REQUIRED + + +class TestRevisionAbaClosed: + def test_stale_copy_cannot_regress_the_revision(self, store): + """Round-4 blocker #1 (PR #242): every revision value (expected + WHERE, new SET, shared publish) derives from one read under the + write lock — a stale actor can never pair a fresh WHERE with a + stale SET and regress the row.""" + import dataclasses + + lease = _admit(store) + stale = dataclasses.replace(lease) # captured before the checkpoint + store.checkpoint_sync(lease, {"n": 1}, progressed=True) # rev 0 -> 1 + with pytest.raises(StaleTurnError): + store.heartbeat_sync(stale) # stale revision fenced out entirely + (revision,) = store._conn.execute("SELECT revision FROM turns").fetchone() + assert revision == 1 # never regressed + assert lease.revision == 1 + # The live owner keeps working. + store.heartbeat_sync(lease) + store.checkpoint_sync(lease, {"n": 2}, progressed=True) + (revision,) = store._conn.execute("SELECT revision FROM turns").fetchone() + assert revision == 2 From 3f8fa420f9005fe6a2e7206979707aa7660fab78 Mon Sep 17 00:00:00 2001 From: ATCharpentier Date: Thu, 30 Jul 2026 18:39:45 -0400 Subject: [PATCH 15/18] =?UTF-8?q?fix:=20review=20round=205=20=E2=80=94=20a?= =?UTF-8?q?ll=20three=20blockers=20(PR=20#242)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 1 (fall-through, final hole): _explicit_resume_recognized returning None when the second load comes back empty (corrupt payload self-healed to rejected / another resumer won / expired) now returns a notice tuple — once row_summary establishes preserved work, NO outcome runs the trigger as a fresh tool-capable turn. Odin's corrupt-JSON end-to-end repro pinned: notice delivered, zero fresh generation. 2 (exhaustive checkpoint schema): validate_payload now checks EVERY persisted field with exact types — guard flags via type(v) is bool (his hedging_retried=0 repro re-armed a consumed guard and made two calls; now: TERMINAL_REJECTED, zero generation, pinned end-to-end), non-boolean bounded integers, counter<=cap invariants (continuation, validation, iteration<=chat_cap), typed string/dict lists, the exact stuck-tracker export shape, and the trajectory envelope (iterations list-of-objects, non-negative iteration_revision, history list). The field SET is exact: unknown extras reject too (they used to reach _ChatTurn(**fields) after acquisition and loop back to SUSPENDED). Validation baseline in tests now derives from a REAL snapshot so it can never drift from the schema. 3 (monotonic lease expiry): renewals go through lease_expires_at=MAX(lease_expires_at, ?) in SQL — a delayed same-owner renewal carrying an earlier expiry (same generation/token/revision lineage passes the fence) can no longer move the expiry backward and expose a healthy owner to the expiry sweep. Pinned: the delayed smaller renewal never regresses, checkpoint path included, sweep untouched. Suite 7,923/5, lint/type gates new=0, coverage findings=0 at 88.4%. --- src/discord/turn_resume.py | 13 +++- src/turn_state/codec.py | 118 +++++++++++++++++++++++++++------ src/turn_state/store.py | 11 ++- tests/test_resume_admission.py | 96 +++++++++++++++++++++++---- tests/test_turn_state_store.py | 28 ++++++++ 5 files changed, 227 insertions(+), 39 deletions(-) diff --git a/src/discord/turn_resume.py b/src/discord/turn_resume.py index d389ad82..02cbbf82 100644 --- a/src/discord/turn_resume.py +++ b/src/discord/turn_resume.py @@ -310,7 +310,18 @@ async def _explicit_resume_recognized(self, message: Any, row_summary: dict): ) row = await asyncio.to_thread(self._store.load_resumable_sync, key) if row is None: - return None + # row_summary established preserved work moments ago; the + # detailed load coming back empty means it just became + # unresumable (corrupt payload self-healed to rejected, another + # resumer won, TTL expiry). A recognized resume must NEVER fall + # through into a fresh tool-capable turn (round-5 blocker #1, + # PR #242) — report what happened instead. + return ( + "That preserved work is no longer resumable (it was just " + "rejected as unreadable, claimed by another resume, or " + "expired). Nothing was resumed — ask fresh for what you need.", + False, False, [], False, + ) # Only the original requester may resume their turn. if str(message.author.id) != str(row.get("user_id") or ""): return ( diff --git a/src/turn_state/codec.py b/src/turn_state/codec.py index aedff14a..13aa6411 100644 --- a/src/turn_state/codec.py +++ b/src/turn_state/codec.py @@ -416,19 +416,33 @@ class CheckpointInvalidError(ValueError): retry (round-3 deviation #5, PR #242).""" -# Fields whose restored TYPE the codec asserts before any lease exists — -# construction failures after acquisition used to bounce the row back to -# SUSPENDED forever. -_REQUIRED_FIELD_TYPES: dict[str, type | tuple[type, ...]] = { - "system_prompt": str, - "messages": list, - "user_id": str, - "chat_cap": int, - "iteration": int, - "stuck_tracker": dict, - "_trajectory": dict, - "_ch_id": str, - "_req_id": str, +# The COMPLETE persisted-field schema, validated before any lease exists +# (round-5 blocker #2, PR #242): partial validation let a tampered/corrupt +# `hedging_retried=0` pass, resume, and re-arm a consumed guard budget — +# the hard rule the census exists to protect. Every field gets an exact +# type (bools via `type(v) is bool`, ints excluding bools), bounds, and +# container-element checks; the field SET is exact (extras reject too — +# an unknown field used to reach `_ChatTurn(**fields)` after acquisition). +_GUARD_FLAG_FIELDS = ( + "fabrication_retried", "promise_retried", "unavail_retried", + "hedging_retried", "code_hedging_retried", "premature_failure_retried", + "_validation_required", +) +_NON_NEGATIVE_INT_FIELDS = ( + "iteration", "continuation_count", "max_continuations", + "_validation_retries", "_max_validation_retries", +) +_POSITIVE_INT_FIELDS = ("chat_cap", "_result_store_cap") +_STRING_FIELDS = ("system_prompt", "user_id", "_ch_id", "_req_id") +_STR_LIST_FIELDS = ("tools_used_in_loop", "_pending_validations") +_DICT_LIST_FIELDS = ("pending_image_blocks", "_op_tool_details") +_STUCK_TRACKER_SCHEMA: dict[str, str] = { + "fingerprints": "str_list", + "window_size": "positive_int", + "min_repeats": "positive_int", + "max_cycle_length": "positive_int", + "names_only": "bool", + "warned": "bool", } @@ -470,18 +484,78 @@ def _exact_int(value: Any) -> bool: missing = PERSISTED_FIELDS - fields.keys() if missing: raise CheckpointInvalidError(f"missing persisted fields: {sorted(missing)}") - for name, expected in _REQUIRED_FIELD_TYPES.items(): - value = fields.get(name) - if expected is int: - if not _exact_int(value): - raise CheckpointInvalidError(f"field {name!r} is not an integer") - elif not isinstance(value, expected): - raise CheckpointInvalidError( - f"field {name!r} has invalid type {type(value).__name__}" - ) + extras = fields.keys() - PERSISTED_FIELDS + if extras: + raise CheckpointInvalidError(f"unknown persisted fields: {sorted(extras)}") + + def _fail(name: str, why: str) -> None: + raise CheckpointInvalidError(f"field {name!r} {why}") + + # One-shot guard flags and validation-required: EXACTLY bool. Anything + # else (0, "", None) would restore falsy and RE-ARM a consumed budget. + for name in _GUARD_FLAG_FIELDS: + if type(fields[name]) is not bool: + _fail(name, "must be exactly a bool") + for name in _NON_NEGATIVE_INT_FIELDS: + value = fields[name] + if not _exact_int(value) or value < 0: + _fail(name, "must be a non-negative integer") + for name in _POSITIVE_INT_FIELDS: + value = fields[name] + if not _exact_int(value) or value < 1: + _fail(name, "must be a positive integer") + for name in _STRING_FIELDS: + if not isinstance(fields[name], str): + _fail(name, "must be a string") + for name in _STR_LIST_FIELDS: + value = fields[name] + if not isinstance(value, list) or not all(isinstance(x, str) for x in value): + _fail(name, "must be a list of strings") + for name in _DICT_LIST_FIELDS: + value = fields[name] + if not isinstance(value, list) or not all(isinstance(x, dict) for x in value): + _fail(name, "must be a list of objects") + # Budget invariants — a consumed counter above its cap is corrupt. + if fields["continuation_count"] > fields["max_continuations"]: + _fail("continuation_count", "exceeds max_continuations") + if fields["_validation_retries"] > fields["_max_validation_retries"]: + _fail("_validation_retries", "exceeds _max_validation_retries") + if fields["iteration"] > fields["chat_cap"]: + _fail("iteration", "exceeds chat_cap") + # Stuck-tracker export shape (exact keys, exact types). + tracker = fields["stuck_tracker"] + if not isinstance(tracker, dict) or set(tracker) != set(_STUCK_TRACKER_SCHEMA): + _fail("stuck_tracker", "has an invalid key set") + for key, kind in _STUCK_TRACKER_SCHEMA.items(): + value = tracker[key] + if kind == "bool" and type(value) is not bool: + _fail("stuck_tracker", f"key {key!r} must be exactly a bool") + elif kind == "positive_int" and (not _exact_int(value) or value < 1): + _fail("stuck_tracker", f"key {key!r} must be a positive integer") + elif kind == "str_list" and ( + not isinstance(value, list) + or not all(isinstance(x, str) for x in value) + ): + _fail("stuck_tracker", f"key {key!r} must be a list of strings") + # Trajectory envelope: the pieces its reconstruction touches. + trajectory = fields["_trajectory"] + if not isinstance(trajectory, dict): + _fail("_trajectory", "must be an object") + traj_iters = trajectory.get("iterations") + if not isinstance(traj_iters, list) or not all( + isinstance(x, dict) for x in traj_iters + ): + _fail("_trajectory", "iterations must be a list of objects") + traj_rev = trajectory.get("iteration_revision") + if not _exact_int(traj_rev) or traj_rev < 0: + _fail("_trajectory", "iteration_revision must be a non-negative integer") + if not isinstance(trajectory.get("history", []), list): + _fail("_trajectory", "history must be a list") # Transcript shape: every entry is a message dict; content is a string # or a list of block dicts (round-4 blocker #2 — a [None] entry used to # explode later in transcript repair, outside the rejection boundary). + if not isinstance(fields["messages"], list): + _fail("messages", "must be a list") for i, msg in enumerate(fields["messages"]): if not isinstance(msg, dict) or not isinstance(msg.get("role"), str): raise CheckpointInvalidError(f"messages[{i}] is not a message object") diff --git a/src/turn_state/store.py b/src/turn_state/store.py index d764ff4e..822b50c8 100644 --- a/src/turn_state/store.py +++ b/src/turn_state/store.py @@ -436,7 +436,12 @@ def checkpoint_sync( ) -> None: """Persist the turn payload (fenced). ``progressed`` advances ``last_progress_at``; recovery waits and rewrites must pass False.""" - sets = ["payload=?", "lease_expires_at=?"] + # MAX() makes the renewal MONOTONIC in SQL (round-5 blocker #3, + # PR #242): a delayed same-owner renewal computed earlier can commit + # after a newer one — both pass the fence (same generation/token/ + # revision lineage) — and must never move the expiry BACKWARD, or + # the expiry sweep falsely suspends a healthy owner. + sets = ["payload=?", "lease_expires_at=MAX(lease_expires_at, ?)"] params: list = [json.dumps(payload), time.time() + self.lease_ttl] if progressed: sets.append("last_progress_at=?") @@ -450,7 +455,9 @@ def heartbeat_sync(self, lease: TurnLease) -> None: """Extend the lease. Never advances last_progress_at, never bumps the revision (a heartbeat is not a state change).""" self._fenced_update( - lease, "lease_expires_at=?", [time.time() + self.lease_ttl], + lease, + "lease_expires_at=MAX(lease_expires_at, ?)", # monotonic renewal + [time.time() + self.lease_ttl], bump_revision=False, ) diff --git a/tests/test_resume_admission.py b/tests/test_resume_admission.py index 0558d9d5..ff8d01d8 100644 --- a/tests/test_resume_admission.py +++ b/tests/test_resume_admission.py @@ -335,9 +335,12 @@ async def test_corrupt_checkpoint_rejects_before_acquiring(self, tmp_path): h.store._conn.execute("UPDATE turns SET payload='{\"broken\": '") h.store._conn.commit() result = await h.manager.try_explicit_resume(resume_msg(original)) - # Malformed JSON self-heals inside load_resumable: rejected - # terminally, and the trigger falls through as a normal message. - assert result is None + # Malformed JSON self-heals inside load_resumable (rejected + # terminally). Round-5 blocker #1: once row_summary established + # preserved work, the trigger must get a NOTICE — never fall + # through into a fresh tool-capable turn. + assert result is not None + assert "no longer resumable" in result[0] (status,) = h.store._conn.execute("SELECT status FROM turns").fetchone() assert status == TurnStatus.TERMINAL_REJECTED # A structurally-valid-but-unrestorable payload rejects explicitly: @@ -513,19 +516,18 @@ def test_validate_payload_rejects_each_structural_deviation(self): from src.turn_state.codec import ( CODEC_VERSION, CheckpointInvalidError, + snapshot_chat_turn, validate_payload, ) - - good_fields = {name: None for name in __import__( - "src.turn_state.codec", fromlist=["PERSISTED_FIELDS"] - ).PERSISTED_FIELDS} - good_fields.update({ - "system_prompt": "s", "messages": [], "user_id": "u", - "chat_cap": 1, "iteration": 0, "stuck_tracker": {}, - "_trajectory": {}, "_ch_id": "c", "_req_id": "r", - }) - base = {"codec_version": CODEC_VERSION, "policy": "chat", - "generation_seq": 0, "fields": good_fields} + from tests.test_turn_checkpoint_codec import _blob_dict, _full_turn + + # The baseline is a REAL snapshot — the exhaustive round-5 schema + # validates every persisted field, so a hand-built skeleton cannot + # stay in sync. + _, store_blob, _ = _blob_dict() + base = snapshot_chat_turn(_full_turn(), store_blob=store_blob, + generation_seq=0) + good_fields = base["fields"] validate_payload(base) # sane baseline passes for broken in ( @@ -545,6 +547,42 @@ def test_validate_payload_rejects_each_structural_deviation(self): {**base, "fields": {**good_fields, "messages": [{"role": "user", "content": [42]}]}}, {**base, "fields": {**good_fields, "chat_cap": True}}, + # Round-5: exact bool for guard flags (0 would re-arm a + # consumed budget), no unknown fields, invariants hold. + {**base, "fields": {**good_fields, "hedging_retried": 0}}, + {**base, "fields": {**good_fields, "fabrication_retried": None}}, + {**base, "fields": {**good_fields, "continuation_count": -1}}, + {**base, "fields": {**good_fields, "continuation_count": 99, + "max_continuations": 3}}, + {**base, "fields": {**good_fields, "field_from_nowhere": 1}}, + {**base, "fields": {**good_fields, "tools_used_in_loop": [1]}}, + {**base, "fields": {**good_fields, + "stuck_tracker": {"warned": True}}}, + {**base, "fields": {**good_fields, + "stuck_tracker": {**good_fields["stuck_tracker"], + "warned": 1}}}, + {**base, "fields": {**good_fields, + "_trajectory": {**good_fields["_trajectory"], + "iteration_revision": True}}}, + {**base, "fields": {**good_fields, "system_prompt": 42}}, + {**base, "fields": {**good_fields, "_trajectory": "not-a-dict"}}, + {**base, "fields": {**good_fields, "pending_image_blocks": [1]}}, + {**base, "fields": {**good_fields, "_validation_retries": 5, + "_max_validation_retries": 2}}, + {**base, "fields": {**good_fields, "iteration": 501, + "chat_cap": 500}}, + {**base, "fields": {**good_fields, + "stuck_tracker": {**good_fields["stuck_tracker"], + "window_size": 0}}}, + {**base, "fields": {**good_fields, + "stuck_tracker": {**good_fields["stuck_tracker"], + "fingerprints": [1]}}}, + {**base, "fields": {**good_fields, + "_trajectory": {**good_fields["_trajectory"], + "iterations": [1]}}}, + {**base, "fields": {**good_fields, + "_trajectory": {**good_fields["_trajectory"], + "history": "not-a-list"}}}, ): with _pytest.raises(CheckpointInvalidError): validate_payload(broken) @@ -681,3 +719,33 @@ async def test_none_message_entry_rejects_inside_the_boundary(self, tmp_path): assert len(h.fake.calls) == calls_before # no fresh LLM response (status,) = h.store._conn.execute("SELECT status FROM turns").fetchone() assert status == TurnStatus.TERMINAL_REJECTED + + +class TestGuardRearmImpossible: + async def test_tampered_guard_flag_rejects_instead_of_rearming(self, tmp_path): + """Odin's round-5 repro: hedging_retried changed to integer 0 used + to pass validation, resume, RE-ARM the consumed guard, and make two + LLM calls. Exact-bool validation now rejects it terminally with + zero generation — the hard rule holds.""" + import json as _json + + h, original = await suspend_turn( + tmp_path, script=[text_response("Shall I proceed with the deployment now?")] + ) + heal_capacity(h, text_response("must never generate")) + (payload_text,) = h.store._conn.execute( + "SELECT payload FROM turns" + ).fetchone() + payload = _json.loads(payload_text) + assert payload["fields"]["hedging_retried"] is True # consumed + payload["fields"]["hedging_retried"] = 0 # the tamper + h.store._conn.execute("UPDATE turns SET payload=?", [_json.dumps(payload)]) + h.store._conn.commit() + + calls_before = len(h.fake.calls) + result = await h.manager.try_explicit_resume(resume_msg(original)) + assert result is not None + assert "could not be restored" in result[0] + assert len(h.fake.calls) == calls_before # ZERO fresh generation + (status,) = h.store._conn.execute("SELECT status FROM turns").fetchone() + assert status == TurnStatus.TERMINAL_REJECTED diff --git a/tests/test_turn_state_store.py b/tests/test_turn_state_store.py index 9bc1c1ab..b26b4b91 100644 --- a/tests/test_turn_state_store.py +++ b/tests/test_turn_state_store.py @@ -648,3 +648,31 @@ def test_stale_copy_cannot_regress_the_revision(self, store): store.checkpoint_sync(lease, {"n": 2}, progressed=True) (revision,) = store._conn.execute("SELECT revision FROM turns").fetchone() assert revision == 2 + + +class TestMonotonicLeaseExpiry: + def test_delayed_smaller_renewal_never_regresses_expiry(self, tmp_path): + """Round-5 blocker #3 (PR #242): renewals are monotonic in SQL — + a delayed same-owner renewal carrying an EARLIER expiry (both pass + the fence: same generation/token/revision lineage) must never move + the expiry backward and expose the turn to a false expiry sweep.""" + s = TurnStateStore(tmp_path / "mono.sqlite3", blob_dir=tmp_path / "b", + lease_ttl=100.0) + lease = _admit(s) + s.heartbeat_sync(lease) # newer renewal: now + 100 + (newer,) = s._conn.execute("SELECT lease_expires_at FROM turns").fetchone() + # The delayed renewal computed from a smaller ttl models an earlier + # time.time() capture committing late. + s.lease_ttl = 1.0 + s.heartbeat_sync(lease) + (after,) = s._conn.execute("SELECT lease_expires_at FROM turns").fetchone() + assert after >= newer # never regressed + # The same monotonicity holds through the checkpoint renewal path. + s.checkpoint_sync(lease, {"n": 1}, progressed=True) + (after_ckpt,) = s._conn.execute( + "SELECT lease_expires_at FROM turns" + ).fetchone() + assert after_ckpt >= newer + # And a false expiry sweep cannot touch the still-live owner. + assert s.sweep_expired_active_sync() == {"turns": 0, "ops": 0} + s.close() From fddc6a3fdec3c1eb950564c50bdd68a6c8c03797 Mon Sep 17 00:00:00 2001 From: Odin Date: Thu, 30 Jul 2026 19:19:16 -0400 Subject: [PATCH 16/18] fix: fail closed on explicit-resume lookup errors --- src/discord/turn_resume.py | 20 +++++++++++++------- tests/test_resume_admission.py | 33 +++++++++++++++++++++++++++++++++ 2 files changed, 46 insertions(+), 7 deletions(-) diff --git a/src/discord/turn_resume.py b/src/discord/turn_resume.py index 02cbbf82..637cbcb7 100644 --- a/src/discord/turn_resume.py +++ b/src/discord/turn_resume.py @@ -286,18 +286,24 @@ async def try_explicit_resume(self, message: Any): content = getattr(message, "content", "") or "" if not self.is_resume_trigger(content): return None - channel_id = str(message.channel.id) - row_summary = await self._latest_suspended_for_channel(channel_id) - if row_summary is None: - return None # nothing to resume — treat as a normal message + # From lexical trigger recognition onward, every store read lives + # inside this no-raise boundary. In particular, failure of the + # initial suspended-row lookup must refuse this command rather than + # returning None and letting intake start a fresh, tool-capable turn. + # We cannot prove that no preserved work exists when identity lookup + # failed, so fail closed with a bounded notice (round-6 task 1). try: + channel_id = str(message.channel.id) + row_summary = await self._latest_suspended_for_channel(channel_id) + if row_summary is None: + return None # lookup succeeded: genuinely nothing to resume return await self._explicit_resume_recognized(message, row_summary) except Exception: log.exception("Explicit resume failed internally") return ( - "I recognized the resume command, but resuming failed " - "internally. The preserved work is untouched — try `resume` " - "again, or ask fresh.", + "I recognized the resume command, but resuming failed internally " + "while safely checking the preserved work. Nothing was resumed or " + "started fresh — try `resume` again later.", False, True, [], False, ) diff --git a/tests/test_resume_admission.py b/tests/test_resume_admission.py index ff8d01d8..ea4fa2c2 100644 --- a/tests/test_resume_admission.py +++ b/tests/test_resume_admission.py @@ -655,6 +655,39 @@ def test_mutation_revision_is_monotonic_across_mutations(self, tmp_path): assert sess.mutation_revision("chX") == 3 # removal GROWS the watermark +class TestResumeLookupFailureFailClosed: + async def test_first_store_read_failure_refuses_before_fresh_generation( + self, tmp_path + ): + """Round-6 task 1: the first store read after recognizing `resume` + is inside the no-raise boundary. An unverifiable identity produces a + bounded refusal through the real pipeline, never a fresh LLM turn.""" + h, original = await suspend_turn(tmp_path) + heal_capacity(h, text_response("fresh turn that must never run")) + h.bot.pipeline._turn_resume = h.manager + calls_before = len(h.fake.calls) + + real_list = h.store.list_suspended_sync + failed = False + + def fail_once(source=None): + nonlocal failed + if not failed: + failed = True + raise OSError("forced one-read failure") + return real_list(source) + + h.store.list_suspended_sync = fail_once + trigger = resume_msg(original) + await h.bot.pipeline.run(trigger, trigger.content) + + assert failed is True + assert len(h.fake.calls) == calls_before # ZERO fresh generation + assert h.row()[0] == TurnStatus.SUSPENDED + delivered = "\n".join(trigger.all_delivered_texts()) + assert "Nothing was resumed or started fresh" in delivered + + class TestRecognizedResumeNeverFallsThrough: async def test_internal_failure_returns_notice_not_fresh_turn(self, tmp_path): """Round-4 blocker #2 (PR #242): once the resume trigger is From d663aa2bdb4a49061f54b972e1ef7bad00af56ef Mon Sep 17 00:00:00 2001 From: Odin Date: Thu, 30 Jul 2026 19:24:11 -0400 Subject: [PATCH 17/18] fix: reject tampered checkpoint guard state --- src/turn_state/store.py | 135 ++++++++++++++++++++++++++------- tests/test_resume_admission.py | 99 +++++++++++++++++++----- tests/test_turn_state_store.py | 96 +++++++++++++++++++++++ 3 files changed, 287 insertions(+), 43 deletions(-) diff --git a/src/turn_state/store.py b/src/turn_state/store.py index 822b50c8..068fa950 100644 --- a/src/turn_state/store.py +++ b/src/turn_state/store.py @@ -121,6 +121,22 @@ class TurnLease: revision: int +def _encode_payload(payload: dict) -> str: + """Stable JSON text used for both storage and integrity hashing.""" + return json.dumps(payload, sort_keys=True) + + +def _payload_digest(payload_text: str) -> str: + """Tamper-evident SHA-256 for the exact persisted payload bytes. + + This is deliberately a digest rather than a second mutable safety-state + copy: one atomic fenced write binds the full field census, transcript, + consumed guard flags, and all caps. It detects out-of-band payload edits + without creating another schema that can drift from the checkpoint codec. + """ + return hashlib.sha256(payload_text.encode("utf-8", "surrogatepass")).hexdigest() + + def effect_fingerprint(tool_name: str, tool_input: dict) -> str: """Handler-derived effect fingerprint (secondary reconciliation evidence, NEVER the primary dedup key — deliberate identical invocations are @@ -158,6 +174,7 @@ def effect_fingerprint(tool_name: str, tool_input: dict) -> str: tool_catalog_hash TEXT, session_snapshot TEXT, payload TEXT, + payload_digest TEXT, PRIMARY KEY (source, channel_id, message_id) ); CREATE INDEX IF NOT EXISTS idx_turns_status ON turns(status); @@ -215,6 +232,7 @@ def __init__( conn.execute("PRAGMA journal_mode=WAL") conn.execute("PRAGMA busy_timeout=30000") conn.executescript(_DDL) + self._migrate_schema_sync(conn) conn.commit() self._restrict_db_modes() self._conn = conn @@ -232,6 +250,35 @@ def __init__( ) self._conn = None + @staticmethod + def _migrate_schema_sync(conn: sqlite3.Connection) -> None: + """Add integrity metadata to stores created before round 6. + + Existing payloads are checksummed once at the migration boundary. + Every subsequent payload write updates payload + digest in the same + fenced SQL statement, so out-of-band same-type mutation cannot re-arm + consumed guards or inflate retry/iteration caps undetected. + """ + columns = {row[1] for row in conn.execute("PRAGMA table_info(turns)")} + if "payload_digest" not in columns: + conn.execute("ALTER TABLE turns ADD COLUMN payload_digest TEXT") + rows = conn.execute( + "SELECT source, channel_id, message_id, payload FROM turns " + "WHERE payload IS NOT NULL AND payload_digest IS NULL" + ).fetchall() + for source, channel_id, message_id, payload_text in rows: + conn.execute( + "UPDATE turns SET payload_digest=? WHERE source=? AND channel_id=? " + "AND message_id=? AND payload=? AND payload_digest IS NULL", + [ + _payload_digest(payload_text), + source, + channel_id, + message_id, + payload_text, + ], + ) + def _restrict_db_modes(self) -> None: """0600 on the database and its WAL/SHM sidecars (best-effort — the sidecars appear lazily; called again from the TTL sweep).""" @@ -441,8 +488,17 @@ def checkpoint_sync( # after a newer one — both pass the fence (same generation/token/ # revision lineage) — and must never move the expiry BACKWARD, or # the expiry sweep falsely suspends a healthy owner. - sets = ["payload=?", "lease_expires_at=MAX(lease_expires_at, ?)"] - params: list = [json.dumps(payload), time.time() + self.lease_ttl] + payload_text = _encode_payload(payload) + sets = [ + "payload=?", + "payload_digest=?", + "lease_expires_at=MAX(lease_expires_at, ?)", + ] + params: list = [ + payload_text, + _payload_digest(payload_text), + time.time() + self.lease_ttl, + ] if progressed: sets.append("last_progress_at=?") params.append(time.time()) @@ -463,11 +519,17 @@ def heartbeat_sync(self, lease: TurnLease) -> None: def suspend_sync(self, lease: TurnLease, payload: dict) -> None: now = time.time() + payload_text = _encode_payload(payload) self._fenced_update( lease, - "payload=?, status=?, suspended_at=?, lease_token=NULL, " - "lease_expires_at=NULL", - [json.dumps(payload), TurnStatus.SUSPENDED, now], + "payload=?, payload_digest=?, status=?, suspended_at=?, " + "lease_token=NULL, lease_expires_at=NULL", + [ + payload_text, + _payload_digest(payload_text), + TurnStatus.SUSPENDED, + now, + ], ) def finish_sync(self, lease: TurnLease, status: str = TurnStatus.TERMINAL_COMPLETED) -> None: @@ -477,7 +539,8 @@ def finish_sync(self, lease: TurnLease, status: str = TurnStatus.TERMINAL_COMPLE raise ValueError(f"finish_sync requires a terminal status, got {status}") self._fenced_update( lease, - "status=?, payload=NULL, lease_token=NULL, lease_expires_at=NULL", + "status=?, payload=NULL, payload_digest=NULL, lease_token=NULL, " + "lease_expires_at=NULL", [status], ) @@ -647,18 +710,33 @@ def load_resumable_sync(self, key: TurnKey) -> dict | None: if self._conn is None: return None row = self._conn.execute( - "SELECT turn_generation, revision, payload, guild_id, user_id, " - "content_digest, code_version, schema_version, prompt_policy_hash, " - "tool_catalog_hash, session_snapshot, recovery_deadline_utc, " - "last_progress_at, suspended_at FROM turns " + "SELECT turn_generation, revision, payload, payload_digest, guild_id, " + "user_id, content_digest, code_version, schema_version, " + "prompt_policy_hash, tool_catalog_hash, session_snapshot, " + "recovery_deadline_utc, last_progress_at, suspended_at FROM turns " "WHERE source=? AND channel_id=? AND message_id=? AND status=?", [key.source, key.channel_id, key.message_id, TurnStatus.SUSPENDED], ).fetchone() if row is None or row[2] is None: return None + payload_text, stored_digest = row[2], row[3] + if ( + not isinstance(payload_text, str) + or not isinstance(stored_digest, str) + or not secrets.compare_digest( + stored_digest, _payload_digest(payload_text) + ) + ): + # Exact-type schema validation cannot detect same-type safety + # regressions (True→False guard flips or cap inflation). Every + # fenced payload write atomically replaces this digest, so a + # mismatch is terminal before reconstruction or generation. + log.error("Checkpoint payload integrity mismatch for %s — rejecting", key) + self.reject_resumable_sync(key, "checkpoint payload integrity mismatch") + return None try: - payload = json.loads(row[2]) - snapshot = json.loads(row[10] or "{}") + payload = json.loads(payload_text) + snapshot = json.loads(row[11] or "{}") except (json.JSONDecodeError, TypeError): # An unreadable checkpoint can never be resumed — reject it # terminally instead of raising into every caller. @@ -675,17 +753,17 @@ def load_resumable_sync(self, key: TurnKey) -> dict | None: "generation": row[0], "revision": row[1], "payload": payload, - "guild_id": row[3], - "user_id": row[4], - "content_digest": row[5], - "code_version": row[6], - "schema_version": row[7], - "prompt_policy_hash": row[8], - "tool_catalog_hash": row[9], + "guild_id": row[4], + "user_id": row[5], + "content_digest": row[6], + "code_version": row[7], + "schema_version": row[8], + "prompt_policy_hash": row[9], + "tool_catalog_hash": row[10], "session_snapshot": snapshot, - "recovery_deadline_utc": row[11], - "last_progress_at": row[12], - "suspended_at": row[13], + "recovery_deadline_utc": row[12], + "last_progress_at": row[13], + "suspended_at": row[14], "operations": [ {"generation_seq": o[0], "tool_call_id": o[1], "state": o[2], "tool_name": o[3], "result": o[4]} @@ -733,8 +811,9 @@ def reject_resumable_sync(self, key: TurnKey, reason: str) -> None: try: with self._write_lock: self._conn.execute( - "UPDATE turns SET status=?, payload=NULL, lease_token=NULL, " - "lease_expires_at=NULL WHERE source=? AND channel_id=? " + "UPDATE turns SET status=?, payload=NULL, payload_digest=NULL, " + "lease_token=NULL, lease_expires_at=NULL WHERE source=? " + "AND channel_id=? " "AND message_id=? AND status=?", [TurnStatus.TERMINAL_REJECTED, key.source, key.channel_id, key.message_id, TurnStatus.SUSPENDED], @@ -777,7 +856,10 @@ def release_acquired_sync( Fenced on the lease so only the acquiring owner can abort.""" conn = self._require() if terminal_reason is not None: - set_sql = "status=?, payload=NULL, lease_token=NULL, lease_expires_at=NULL" + set_sql = ( + "status=?, payload=NULL, payload_digest=NULL, lease_token=NULL, " + "lease_expires_at=NULL" + ) first_param: list = [TurnStatus.TERMINAL_REJECTED] else: set_sql = "status=?, lease_token=NULL, lease_expires_at=NULL" @@ -898,7 +980,8 @@ def ttl_sweep_sync( expired = cur.rowcount # Clock 2: diagnostic payload retention — 7d, then tombstone. cur = conn.execute( - "UPDATE turns SET payload=NULL WHERE payload IS NOT NULL " + "UPDATE turns SET payload=NULL, payload_digest=NULL " + "WHERE payload IS NOT NULL " "AND status IN (?, ?, ?, ?, ?) AND last_progress_at < ?", [*sorted(TurnStatus.TERMINAL), now - payload_retention_days * 86400.0], diff --git a/tests/test_resume_admission.py b/tests/test_resume_admission.py index ea4fa2c2..7ab5041a 100644 --- a/tests/test_resume_admission.py +++ b/tests/test_resume_admission.py @@ -106,6 +106,17 @@ def resume_msg(original, content="resume", author=None): return FakeMessage(content, author=author or original.author, channel=original.channel) +def rewrite_payload_with_valid_digest(store, payload_text: str) -> None: + """Bypass payload-integrity only when a test targets codec validation.""" + import hashlib + + digest = hashlib.sha256(payload_text.encode()).hexdigest() + store._conn.execute( + "UPDATE turns SET payload=?, payload_digest=?", [payload_text, digest] + ) + store._conn.commit() + + def make_breaker_probe_ready(h): """Model the production timeline where the breaker cooldown has elapsed by the time a resume happens (suspension→resume is minutes, cooldown is @@ -332,8 +343,7 @@ async def test_corrupt_checkpoint_rejects_before_acquiring(self, tmp_path): TERMINAL_REJECTED — never a stranded ACTIVE row invisible to resumable queries.""" h, original = await suspend_turn(tmp_path) - h.store._conn.execute("UPDATE turns SET payload='{\"broken\": '") - h.store._conn.commit() + rewrite_payload_with_valid_digest(h.store, '{"broken": ') result = await h.manager.try_explicit_resume(resume_msg(original)) # Malformed JSON self-heals inside load_resumable (rejected # terminally). Round-5 blocker #1: once row_summary established @@ -345,10 +355,9 @@ async def test_corrupt_checkpoint_rejects_before_acquiring(self, tmp_path): assert status == TurnStatus.TERMINAL_REJECTED # A structurally-valid-but-unrestorable payload rejects explicitly: h2, original2 = await suspend_turn(tmp_path / "second") - h2.store._conn.execute( - "UPDATE turns SET payload='{\"fields\": {\"stuck_tracker\": 42}}'" + rewrite_payload_with_valid_digest( + h2.store, '{"fields": {"stuck_tracker": 42}}' ) - h2.store._conn.commit() result2 = await h2.manager.try_explicit_resume(resume_msg(original2)) assert result2 is not None assert "could not be restored" in result2[0] @@ -502,8 +511,7 @@ async def test_empty_object_payload_terminally_rejects(self, tmp_path): structurally-invalid payload rejects BEFORE any lease exists — never bounced back to SUSPENDED for an infinite retry loop.""" h, original = await suspend_turn(tmp_path) - h.store._conn.execute("UPDATE turns SET payload='{}'") - h.store._conn.commit() + rewrite_payload_with_valid_digest(h.store, "{}") result = await h.manager.try_explicit_resume(resume_msg(original)) assert result is not None assert "could not be restored" in result[0] @@ -718,10 +726,7 @@ async def test_bad_generation_seq_rejects_terminally_not_loop(self, tmp_path): ).fetchone() payload = _json.loads(payload_text) payload["generation_seq"] = "bad" - h.store._conn.execute( - "UPDATE turns SET payload=?", [_json.dumps(payload)] - ) - h.store._conn.commit() + rewrite_payload_with_valid_digest(h.store, _json.dumps(payload)) result = await h.manager.try_explicit_resume(resume_msg(original)) assert result is not None assert "could not be restored" in result[0] @@ -741,10 +746,7 @@ async def test_none_message_entry_rejects_inside_the_boundary(self, tmp_path): ).fetchone() payload = _json.loads(payload_text) payload["fields"]["messages"] = [None] - h.store._conn.execute( - "UPDATE turns SET payload=?", [_json.dumps(payload)] - ) - h.store._conn.commit() + rewrite_payload_with_valid_digest(h.store, _json.dumps(payload)) calls_before = len(h.fake.calls) result = await h.manager.try_explicit_resume(resume_msg(original)) assert result is not None @@ -772,8 +774,7 @@ async def test_tampered_guard_flag_rejects_instead_of_rearming(self, tmp_path): payload = _json.loads(payload_text) assert payload["fields"]["hedging_retried"] is True # consumed payload["fields"]["hedging_retried"] = 0 # the tamper - h.store._conn.execute("UPDATE turns SET payload=?", [_json.dumps(payload)]) - h.store._conn.commit() + rewrite_payload_with_valid_digest(h.store, _json.dumps(payload)) calls_before = len(h.fake.calls) result = await h.manager.try_explicit_resume(resume_msg(original)) @@ -782,3 +783,67 @@ async def test_tampered_guard_flag_rejects_instead_of_rearming(self, tmp_path): assert len(h.fake.calls) == calls_before # ZERO fresh generation (status,) = h.store._conn.execute("SELECT status FROM turns").fetchone() assert status == TurnStatus.TERMINAL_REJECTED + +class TestCheckpointIntegrityRejectsSameTypeTampering: + @staticmethod + def _tamper_payload(h, mutate): + (payload_text,) = h.store._conn.execute( + "SELECT payload FROM turns" + ).fetchone() + payload = json.loads(payload_text) + mutate(payload["fields"]) + # Deliberately bypass the store write API: the stored digest remains + # bound to the original bytes, as it would under external corruption. + h.store._conn.execute( + "UPDATE turns SET payload=?", [json.dumps(payload, sort_keys=True)] + ) + h.store._conn.commit() + + async def test_consumed_true_guard_flipped_false_halts_zero_generation( + self, tmp_path + ): + """Round-6 task 2: same-type True→False cannot re-arm a guard.""" + h, original = await suspend_turn( + tmp_path, + script=[text_response("Shall I proceed with the deployment now?")], + ) + heal_capacity(h, text_response("must never generate")) + calls_before = len(h.fake.calls) + self._tamper_payload( + h, lambda fields: fields.__setitem__("hedging_retried", False) + ) + + result = await h.manager.try_explicit_resume(resume_msg(original)) + + assert result is not None + assert "no longer resumable" in result[0] + assert len(h.fake.calls) == calls_before + assert h.row()[0] == TurnStatus.TERMINAL_REJECTED + + @pytest.mark.parametrize( + ("cap_name", "expected"), + [ + ("max_continuations", 3), + ("_max_validation_retries", 2), + ("chat_cap", 500), + ], + ) + async def test_same_type_cap_inflation_halts_zero_generation( + self, tmp_path, cap_name, expected + ): + """Continuation, validation, and iteration caps are digest-bound.""" + h, original = await suspend_turn(tmp_path) + heal_capacity(h, text_response("must never generate")) + calls_before = len(h.fake.calls) + + def inflate(fields): + assert fields[cap_name] == expected + fields[cap_name] = expected + 1 + + self._tamper_payload(h, inflate) + result = await h.manager.try_explicit_resume(resume_msg(original)) + + assert result is not None + assert "no longer resumable" in result[0] + assert len(h.fake.calls) == calls_before + assert h.row()[0] == TurnStatus.TERMINAL_REJECTED diff --git a/tests/test_turn_state_store.py b/tests/test_turn_state_store.py index b26b4b91..6b7a2452 100644 --- a/tests/test_turn_state_store.py +++ b/tests/test_turn_state_store.py @@ -676,3 +676,99 @@ def test_delayed_smaller_renewal_never_regresses_expiry(self, tmp_path): # And a false expiry sweep cannot touch the still-live owner. assert s.sweep_expired_active_sync() == {"turns": 0, "ops": 0} s.close() + +class TestCheckpointPayloadIntegrity: + def test_payload_and_digest_advance_atomically_under_the_fence(self, store): + import hashlib + + lease = _admit(store) + store.checkpoint_sync(lease, {"iteration": 1}, progressed=True) + payload, digest = _row(store, cols="payload, payload_digest") + assert digest == hashlib.sha256(payload.encode()).hexdigest() + + store.suspend_sync(lease, {"iteration": 2}) + payload, digest = _row(store, cols="payload, payload_digest") + assert digest == hashlib.sha256(payload.encode()).hexdigest() + + def test_pre_integrity_schema_migrates_existing_payload(self, tmp_path): + import hashlib + import sqlite3 + + db_dir = tmp_path / "legacy" + db_dir.mkdir() + db_path = db_dir / "turns.sqlite3" + conn = sqlite3.connect(db_path) + legacy_ddl = _legacy_turns_ddl_without_payload_digest() + conn.executescript(legacy_ddl) + payload = '{"fields": {"hedging_retried": true}}' + conn.execute( + "INSERT INTO turns (source, channel_id, message_id, turn_generation, " + "revision, status, last_progress_at, created_at, schema_version, payload) " + "VALUES (?,?,?,?,?,?,?,?,?,?)", + ["discord", "c", "m", "g", 0, TurnStatus.SUSPENDED, 1.0, 1.0, 1, + payload], + ) + conn.commit() + conn.close() + + migrated = TurnStateStore(db_path, blob_dir=db_dir / "blobs") + try: + columns = { + row[1] + for row in migrated._conn.execute("PRAGMA table_info(turns)") + } + assert "payload_digest" in columns + stored_payload, digest = migrated._conn.execute( + "SELECT payload, payload_digest FROM turns" + ).fetchone() + assert stored_payload == payload + assert digest == hashlib.sha256(payload.encode()).hexdigest() + finally: + migrated.close() + + +def _legacy_turns_ddl_without_payload_digest() -> str: + """Minimal pre-round-6 schema used to pin the additive migration.""" + return """ + CREATE TABLE turns ( + source TEXT NOT NULL, + channel_id TEXT NOT NULL, + message_id TEXT NOT NULL, + turn_generation TEXT NOT NULL, + revision INTEGER NOT NULL DEFAULT 0, + lease_token TEXT, + lease_expires_at REAL, + status TEXT NOT NULL, + recovery_deadline_utc REAL, + last_progress_at REAL NOT NULL, + created_at REAL NOT NULL, + suspended_at REAL, + guild_id TEXT, + user_id TEXT, + content_digest TEXT, + code_version TEXT, + schema_version INTEGER NOT NULL, + prompt_policy_hash TEXT, + tool_catalog_hash TEXT, + session_snapshot TEXT, + payload TEXT, + PRIMARY KEY (source, channel_id, message_id) + ); + CREATE TABLE operations ( + source TEXT NOT NULL, + channel_id TEXT NOT NULL, + message_id TEXT NOT NULL, + turn_generation TEXT NOT NULL, + generation_seq INTEGER NOT NULL, + tool_call_id TEXT NOT NULL, + state TEXT NOT NULL, + tool_name TEXT NOT NULL, + iteration INTEGER, + effect_fingerprint TEXT, + result TEXT, + created_at REAL NOT NULL, + updated_at REAL NOT NULL, + PRIMARY KEY (source, channel_id, message_id, turn_generation, + generation_seq, tool_call_id) + ); + """ From 928bac2843a0bf9c4d3bde84db94dac689422755 Mon Sep 17 00:00:00 2001 From: Odin Date: Thu, 30 Jul 2026 19:54:13 -0400 Subject: [PATCH 18/18] fix: verify externalized checkpoint blobs --- src/turn_state/store.py | 9 ++++++- tests/test_resume_admission.py | 44 ++++++++++++++++++++++++++++++++++ 2 files changed, 52 insertions(+), 1 deletion(-) diff --git a/src/turn_state/store.py b/src/turn_state/store.py index 068fa950..685a0517 100644 --- a/src/turn_state/store.py +++ b/src/turn_state/store.py @@ -1030,6 +1030,13 @@ def load_blob_sync(self, ref: str) -> bytes: digest = ref.split(":", 1)[1] if ref.startswith("blob:") else ref path = self._blob_dir / digest try: - return path.read_bytes() + data = path.read_bytes() except OSError as exc: raise TurnStateUnavailableError(f"blob read failed: {ref}") from exc + # The checkpoint payload digest binds the blob reference, not the + # externalized bytes. Treat the content-addressed filename as the + # expected digest so an out-of-band blob edit cannot substitute + # transcript content during resume. + if hashlib.sha256(data).hexdigest() != digest: + raise TurnStateUnavailableError(f"blob digest mismatch: {ref}") + return data diff --git a/tests/test_resume_admission.py b/tests/test_resume_admission.py index 7ab5041a..6ac51f7a 100644 --- a/tests/test_resume_admission.py +++ b/tests/test_resume_admission.py @@ -784,6 +784,50 @@ async def test_tampered_guard_flag_rejects_instead_of_rearming(self, tmp_path): (status,) = h.store._conn.execute("SELECT status FROM turns").fetchone() assert status == TurnStatus.TERMINAL_REJECTED + +class TestExternalizedBlobIntegrity: + async def test_tampered_externalized_blob_halts_zero_generation(self, tmp_path): + """Externalized transcript bytes remain bound to their blob ref. + + The inline payload digest covers the content-addressed ref. If the + referenced file is edited out of band, reconstruction must reject + before lease acquisition rather than show substituted content to the + model. + """ + h, original = await suspend_turn(tmp_path) + heal_capacity(h, text_response("must never generate")) + + image_data = b"aGVsbG8=" + ref = h.store.store_blob_sync(image_data) + (payload_text,) = h.store._conn.execute( + "SELECT payload FROM turns" + ).fetchone() + payload = json.loads(payload_text) + payload["fields"]["messages"].append({ + "role": "user", + "content": [{ + "type": "image", + "source": { + "type": "blob_ref", + "media_type": "image/png", + "ref": ref, + }, + }], + }) + rewrite_payload_with_valid_digest(h.store, json.dumps(payload)) + + digest = ref.split(":", 1)[1] + (h.store._blob_dir / digest).write_bytes(b"dGFtcGVyZWQ=") + calls_before = len(h.fake.calls) + + result = await h.manager.try_explicit_resume(resume_msg(original)) + + assert result is not None + assert "could not be restored" in result[0] + assert len(h.fake.calls) == calls_before + assert h.row()[0] == TurnStatus.TERMINAL_REJECTED + + class TestCheckpointIntegrityRejectsSameTypeTampering: @staticmethod def _tamper_payload(h, mutate):