diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 00000000..d2e6da36 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,70 @@ +# Changelog + +All notable changes to the Context Intelligence Server are recorded here. + +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), +and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). + +## [7.0.0] + +### Changed (breaking) + +- **`QueueManager.commit()` now takes a required `cursor` argument** — + `commit(session_id, new_offset, cursor)`. The committed byte offset and the + cross-handler cursor are written in one atomic offset record so they can + never skew. This is a breaking change to the `QueueManager` Protocol that any + alternate backend must implement; all in-repo callers are updated. A legacy + bare-integer `.offset` written by an older build still reads its committed + position correctly, so an in-place upgrade is safe. +- **`QueueManager` Protocol gains `read_cursor`** as a member every backend + must provide: it returns the persisted cursor a rebuilt worker resumes from. + Log-structured-queue details stay off the Protocol -- dead-line + reconciliation is a filesystem-backend concern (a broker backend has no + leading-dead-line window), so it is reached through the concrete backend, not + the Protocol. + +### Fixed + +- The durable cursor now survives every offset-mutation path that keeps the + session alive, not just `commit()`: idle compaction and the boot + `RESET_OFFSET` reclaim preserve a non-empty cursor instead of wiping it with + a bare offset write. Finalize's `delete_drained` still drops it, which is + terminal cleanup, not loss: the session is over. +- Every write to a session's `.offset` now happens through one writer that + takes the session's file lock itself and stages through a per-write temp + file. Previously the lock was the caller's job, and the dead-letter + reconcile reached by `read_batch` took none -- a commit racing it was + silently rolled back to the offset the reconcile had read (measured: 248 of + 300 concurrent runs), and the shared temp name let one writer publish + another's half-written record. +- A corrupt/unparseable `.offset` file now quarantines the one affected drain + worker (logged, closed, deregistered) instead of crash-looping it; other + sessions keep draining. Every read in the drain loop is covered, including + the idle dry-exit recheck and the `session:end` tail drain, which previously + escaped to the supervisor. +- Cross-handler run-id resolution is consistent across the orchestrator-run, + iteration, and content-block handlers, so a partial cursor after a worker + rebuild no longer drops the `HAS_PART` edge or orphans ContentBlock nodes. +- `restore_cursor` deep-copies mutable cursor fields, so an in-place mutation on + a later retry attempt can no longer corrupt the pre-batch baseline the next + rollback restores. +- A crash-then-respawn mid-isolation no longer re-dead-letters an + already-dead-lettered record: the drain worker reconciles the session's + leading dead lines on every (re)spawn, not only at boot. +- Finalizing a session retires its dead letters out of the session's own name + instead of leaving them in place. A later session reusing the id no longer + reconciles against the previous session's dead payloads and commits past + events it never processed. The payloads are retained, still reported by + `GET /queues/dead-letter/{key}`, and still expire on their own schedule. +- The Session node survives an exhausted-batch isolation: discarding the failed + batch's buffer now also invalidates the seen-session cache, so the isolated + re-dispatch re-issues the node instead of early-returning. + +## [6.7.3] + +### Added + +- Durable per-record cursor persisted atomically with the committed offset, so + a rebuilt drain worker resumes its cross-handler counters instead of + re-minting node ids. Retry-dedup and the run-id tiebreaker keep a re-delivered + batch idempotent. diff --git a/context_intelligence_server/handlers/data_layer_2/content_block.py b/context_intelligence_server/handlers/data_layer_2/content_block.py index c0238f79..c37ee00f 100644 --- a/context_intelligence_server/handlers/data_layer_2/content_block.py +++ b/context_intelligence_server/handlers/data_layer_2/content_block.py @@ -49,11 +49,18 @@ async def __call__(self, event: str, data: dict[str, Any]) -> HookResult: return HookResult(action="continue") block_index = data.get("block_index") - iteration_id = self.services.data_layer_2.active_iteration_id - # ID format is "{session_id}::iteration::{n}"; [-1] extracts the iteration number. - # If the cursor format ever changes, this extraction must be updated to match. - iteration_n = iteration_id.split("::")[-1] if iteration_id else "0" - block_node_id = f"{session_id}::block::{iteration_n}::{block_index}" + # Key the block off the FULL active_iteration_id so a block inherits the + # iteration's run scope (including the run tiebreaker). Using only the + # trailing iteration number would let two runs that share an iteration + # number collide on the same block_node_id and MERGE-overwrite. + # Resolve the iteration id the same seq-aware way the Iteration handler + # does, so after a worker rebuild a block inherits the run's real + # iteration scope rather than the run-less fallback. + iteration_id = self.services.data_layer_2.resolve_active_iteration_id( + session_id + ) + iteration_key = iteration_id if iteration_id else f"{session_id}::iteration::0" + block_node_id = f"{iteration_key}::block::{block_index}" if event == "content_block:start": await self._handle_start(session_id, block_node_id, block_index, data) diff --git a/context_intelligence_server/handlers/data_layer_2/iteration.py b/context_intelligence_server/handlers/data_layer_2/iteration.py index 4c10d1e8..14898c24 100644 --- a/context_intelligence_server/handlers/data_layer_2/iteration.py +++ b/context_intelligence_server/handlers/data_layer_2/iteration.py @@ -36,6 +36,18 @@ class IterationHandler: def __init__(self, services: HookStateService) -> None: self.services = services + def _current_iteration_scope(self) -> str: + """The 'run' | 'unscoped' discriminator, sourced from the SAME cursor + field (``execution_start_ts``) used to decide the node_id shape, so all + three upsert_node call sites (provider:request, llm:request, + llm:response) always agree. + """ + return ( + "run" + if self.services.data_layer_2.execution_start_ts is not None + else "unscoped" + ) + async def __call__(self, event: str, data: dict[str, Any]) -> HookResult: """Dispatch to the appropriate sub-handler. @@ -64,22 +76,55 @@ async def _handle_provider_request( ) -> None: """Create Iteration node and set active_iteration_id cursor. - - Computes iteration_id as '{session_id}::iteration::{iteration_number}' + - Computes iteration_id as run-scoped, reusing the active run's full + orch_run_id (which carries the run tiebreaker): + '{orch_run_id}::iteration::{iteration_number}' when a run is active, + falling back to the bare '{session_id}::iteration::{iteration_number}' + shape when no run is active. - Sets active_iteration_id cursor on DataLayer2State - Creates Iteration:SST_EVENT node with session_id, iteration_number, started_at - - Conditionally creates E06: OrchestratorRun -[:HAS_PART {sst_semantic: 'CONTAINS'}]-> - Iteration when execution_start_ts cursor is set + - Conditionally creates the OrchestratorRun -[:HAS_PART]-> Iteration edge + when a run is active. + + Without run-scoping, iteration_number alone (a counter scoped to the whole + session, not the run) can repeat across orchestrator runs -- e.g. after a + drainer restart resets the in-memory counter -- causing distinct runs' + Iteration nodes to MERGE onto the same node_id and their usage figures to + clobber each other. """ # Increment counter to get the next iteration number self.services.data_layer_2.iteration_count += 1 iteration_number = self.services.data_layer_2.iteration_count timestamp: str = data.get("timestamp", "") - iteration_id = f"{session_id}::iteration::{iteration_number}" + + # Resolve the run id the same seq-aware way every other handler does: + # after a worker rebuild active_orch_run_id may be absent while the run + # is still active, in which case this re-derives the run's real id (with + # its seq) rather than falling through to the seq-less, run-less shape. + orch_run_id = self.services.data_layer_2.resolve_active_orch_run_id(session_id) + if orch_run_id is not None: + iteration_id = f"{orch_run_id}::iteration::{iteration_number}" + else: + iteration_id = f"{session_id}::iteration::{iteration_number}" # Set cursor so llm:request and llm:response can find this iteration self.services.data_layer_2.active_iteration_id = iteration_id + # A queryable discriminator between a run-scoped iteration and a + # legitimate loop-basic session with no active orchestrator run. + iteration_scope = self._current_iteration_scope() + if iteration_scope == "unscoped": + # INFO not WARNING: a loop-basic session with no execution:start is + # a normal case, not an alert-worthy anomaly. + logger.info( + "unscoped_iteration_emitted session=%s iteration_number=%d iteration_id=%s", + session_id, + iteration_number, + iteration_id, + extra={"session_id": session_id}, + ) + # Create the Iteration node await self.services.graph.upsert_node( iteration_id, @@ -88,13 +133,12 @@ async def _handle_provider_request( "session_id": session_id, "iteration_number": iteration_number, "started_at": timestamp, + "iteration_scope": iteration_scope, }, ) # E06 (conditional): OrchestratorRun -[:HAS_PART {sst_semantic: 'CONTAINS'}]-> Iteration - execution_start_ts = self.services.data_layer_2.execution_start_ts - if execution_start_ts is not None: - orch_run_id = f"{session_id}::orch_run::{execution_start_ts}" + if orch_run_id is not None: await self.services.graph.upsert_edge( orch_run_id, iteration_id, @@ -127,6 +171,12 @@ async def _handle_llm_request(self, data: dict[str, Any]) -> None: "model": data.get("model"), "message_count": data.get("message_count"), "has_system": data.get("has_system"), + # Stamp independently of provider:request's own write -- an + # Iteration node must never be created/updated without a scope + # value, even if this write is the first one + # to ever reach the node (e.g. a dead-lettered provider:request + # whose cursor mutation nonetheless survived). + "iteration_scope": self._current_iteration_scope(), }, ) @@ -157,6 +207,9 @@ async def _handle_llm_response(self, data: dict[str, Any]) -> None: "usage_input": usage.get("input_tokens"), "usage_output": usage.get("output_tokens"), "usage_cache_write": usage.get("cache_creation_input_tokens"), + # See _handle_llm_request -- same completeness rationale applies + # to this, the third of the three sites. + "iteration_scope": self._current_iteration_scope(), }, ) diff --git a/context_intelligence_server/handlers/data_layer_2/orchestrator_run.py b/context_intelligence_server/handlers/data_layer_2/orchestrator_run.py index 45cb7bb3..24f233ff 100644 --- a/context_intelligence_server/handlers/data_layer_2/orchestrator_run.py +++ b/context_intelligence_server/handlers/data_layer_2/orchestrator_run.py @@ -62,12 +62,24 @@ async def __call__(self, event: str, data: dict[str, Any]) -> HookResult: # Sub-handlers # ------------------------------------------------------------------ + def _resolve_run_id(self, session_id: str) -> str: + """Resolve the OrchestratorRun id an enrichment event must target. + + Delegates to the shared seq-aware resolver so every handler agrees on + the id shape. The enrichment callers only reach here after guarding on + ``execution_start_ts is not None`` (a run is active), so the resolver + never returns None here. + """ + run_id = self.services.data_layer_2.resolve_active_orch_run_id(session_id) + assert run_id is not None # guarded by execution_start_ts on every caller + return run_id + async def _handle_execution_start( self, session_id: str, data: dict[str, Any] ) -> None: """Create OrchestratorRun node and wire E01 (and optionally E14). - - Computes orch_run_id as '{session_id}::orch_run::{timestamp}' + - Computes orch_run_id as '{session_id}::orch_run::{timestamp}::{seq}' - Sets execution_start_ts cursor on DataLayer2State - Creates OrchestratorRun:SST_EVENT node with session_id + started_at - Creates E01: Session -[:HAS_EXECUTION {sst_semantic: 'CONTAINS'}]-> OrchestratorRun @@ -75,10 +87,16 @@ async def _handle_execution_start( OrchestratorRun when last_prompt_id cursor is set """ timestamp: str = data.get("timestamp", "") - orch_run_id = f"{session_id}::orch_run::{timestamp}" - - # Store cursor so execution:end and orchestrator:complete can find this run + # Tiebreaker: two runs sharing an identical timestamp (coarse clock or a + # replayed execution:start) must not collide on the same orch_run_id. + self.services.data_layer_2.orch_run_seq += 1 + seq = self.services.data_layer_2.orch_run_seq + orch_run_id = f"{session_id}::orch_run::{timestamp}::{seq}" + + # Store cursors so execution:end, orchestrator:complete, and the + # Iteration/ContentBlock handlers all reuse this exact id. self.services.data_layer_2.execution_start_ts = timestamp + self.services.data_layer_2.active_orch_run_id = orch_run_id # Create the OrchestratorRun node await self.services.graph.upsert_node( @@ -130,7 +148,7 @@ async def _handle_execution_end( if ts is None: return - orch_run_id = f"{session_id}::orch_run::{ts}" + orch_run_id = self._resolve_run_id(session_id) timestamp: str = data.get("timestamp", "") node_data: dict[str, Any] = { @@ -167,7 +185,7 @@ async def _handle_orchestrator_complete( if ts is None: return - orch_run_id = f"{session_id}::orch_run::{ts}" + orch_run_id = self._resolve_run_id(session_id) timestamp: str = data.get("timestamp", "") orchestrator: str = data.get("orchestrator", "") turn_count: Any = data.get("turn_count") @@ -213,3 +231,4 @@ async def _handle_orchestrator_complete( # Update cursors self.services.data_layer_2.last_completed_orch_run_id = orch_run_id self.services.data_layer_2.execution_start_ts = None + self.services.data_layer_2.active_orch_run_id = None diff --git a/context_intelligence_server/handlers/data_layer_2/state.py b/context_intelligence_server/handlers/data_layer_2/state.py index 47b363c2..bb12d9a1 100644 --- a/context_intelligence_server/handlers/data_layer_2/state.py +++ b/context_intelligence_server/handlers/data_layer_2/state.py @@ -16,6 +16,19 @@ class DataLayer2State: # OrchestratorRun identity execution_start_ts: str | None = None + # Monotonic per-session run counter, incremented on each execution:start and + # folded into orch_run_id. Two runs that share an identical execution_start_ts + # (coarse clock, or a replayed execution:start) would otherwise collide on the + # same orch_run_id and MERGE-overwrite each other's nodes. Lives here so the + # durable cursor persists it across a worker rebuild. + orch_run_seq: int = 0 + + # The full orch_run_id of the active run, including the tiebreaker. Set at + # execution:start and read back by execution:end/orchestrator:complete and + # the Iteration/ContentBlock handlers, so every node in one run agrees on the + # same id without any site recomputing it from execution_start_ts alone. + active_orch_run_id: str | None = None + # Iteration cursor read by ContentBlockHandler + ToolCallHandler active_iteration_id: str | None = None @@ -28,6 +41,48 @@ class DataLayer2State: # E15 OrchestratorRun→Prompt turn-flow cursor last_completed_orch_run_id: str | None = None - # Iteration counter — incremented on each provider:request; used to compute - # iteration_id as '{session_id}::iteration::{iteration_count}' + # Iteration counter — incremented on each provider:request; combined with + # execution_start_ts (via IterationHandler) to compute the run-scoped + # iteration_id '{session_id}::orch_run::{execution_start_ts}::{seq}::iteration::{iteration_count}' + # (falls back to the bare '{session_id}::iteration::{iteration_count}' shape when no + # orchestrator run is active). Scoped per-session, not reset per run: uniqueness across + # runs comes from the orch_run_id prefix, not from this counter — + # resetting it would collide with ContentBlockHandler's block_node_id derivation, which + # keys solely off this counter's value, not the run. iteration_count: int = 0 + + def resolve_active_orch_run_id(self, session_id: str) -> str | None: + """Return the active run's orch_run_id, or None when no run is active. + + Prefers the stored ``active_orch_run_id`` cursor. When that is absent + but a run IS active (``execution_start_ts`` set) — a partial cursor + after a worker rebuild — re-derives the same + ``{session_id}::orch_run::{execution_start_ts}::{seq}`` shape + ``execution:start`` produces, so every handler in one run agrees on the + id and the HAS_PART edge is never dropped onto a non-existent node. + Returns None only when no run is active (a loop-basic session). + """ + if self.active_orch_run_id is not None: + return self.active_orch_run_id + if self.execution_start_ts is not None: + return ( + f"{session_id}::orch_run::{self.execution_start_ts}" + f"::{self.orch_run_seq}" + ) + return None + + def resolve_active_iteration_id(self, session_id: str) -> str | None: + """Return the active iteration_id, or None when no iteration is active. + + Prefers the stored ``active_iteration_id`` cursor. When that is absent + but a run+iteration is reconstructible (a partial cursor after a + rebuild), re-derives ``{orch_run_id}::iteration::{iteration_count}`` + from the seq-aware run id, so a ContentBlock inherits the same run + scope every other handler agrees on rather than an orphan key. + """ + if self.active_iteration_id is not None: + return self.active_iteration_id + run_id = self.resolve_active_orch_run_id(session_id) + if run_id is not None and self.iteration_count > 0: + return f"{run_id}::iteration::{self.iteration_count}" + return None diff --git a/context_intelligence_server/queue_manager/filesystem.py b/context_intelligence_server/queue_manager/filesystem.py index c77f7135..75972abc 100644 --- a/context_intelligence_server/queue_manager/filesystem.py +++ b/context_intelligence_server/queue_manager/filesystem.py @@ -5,11 +5,12 @@ letters), ``.log.compact.tmp`` (transient compaction copy). Framing (one event == one ``\\n``-terminated byte range) holds only while a -single process writes the directory; each key's ``file_lock`` (a -``threading.Lock`` held on the writing thread) serialises its writes. Records -must contain no raw ``0x0A`` except the terminator. ``session_id`` is the raw -filename stem and is rejected if empty or containing a separator or null byte. -Appends are not ``fsync``ed: crash-durable, not power-loss-durable. +single process writes the directory; each key's ``file_lock`` (held on the +writing thread) serialises its writes, including every write of the ``.offset`` +record. Records must contain no raw ``0x0A`` except the terminator. +``session_id`` is the raw filename stem and is rejected if empty or containing a +separator or null byte. Appends are not ``fsync``ed: crash-durable, not +power-loss-durable. """ from __future__ import annotations @@ -22,11 +23,12 @@ import os import threading import time +import uuid from collections.abc import Callable, Coroutine, Iterator from dataclasses import dataclass from enum import Enum from pathlib import Path -from typing import Any, TypeVar +from typing import Any, Self, TypeVar from context_intelligence_server.config import get_settings from context_intelligence_server.queue_manager.protocol import Batch, Record @@ -35,6 +37,10 @@ _T = TypeVar("_T") +# Ending shared by every ``.offset`` staging file, so one glob still reaps +# strays left by a crash mid-write (``reclaim_orphans``). +_OFFSET_TMP_SUFFIX = ".offset.tmp" + # Fixed buffer size for streaming scans over a session ``.log`` (last-newline # search and newline counting). Bounds boot-time and /status memory to O(chunk) # instead of O(file): a durable log can be multi-GB (4.9 GB in the incident), @@ -97,12 +103,51 @@ def _reclaim_redrain_max_bytes() -> int: return get_settings().reclaim_redrain_max_bytes +class _KeyLock: + """A reentrant lock whose held state can be observed. + + Reentrancy is required because ``_write_offset_record`` takes this lock + itself while callers that need a wider atomic sequence (compaction, the + dead-letter reconcile) already hold it. ``threading.RLock`` alone would do, + but exposes no way to ask whether it is held, and the framing guarantees + here are asserted against exactly that. + + ``depth`` is only ever changed while the lock is held, so a concurrent + reader sees either 0 (free) or a positive count (held). + """ + + def __init__(self) -> None: + self._lock = threading.RLock() + self._depth = 0 + + def acquire(self, blocking: bool = True, timeout: float = -1) -> bool: + if not self._lock.acquire(blocking, timeout): + return False + self._depth += 1 + return True + + def release(self) -> None: + self._depth -= 1 + self._lock.release() + + def locked(self) -> bool: + """Whether any thread currently holds this lock.""" + return self._depth > 0 + + def __enter__(self) -> Self: + self.acquire() + return self + + def __exit__(self, *exc_info: object) -> None: + self.release() + + @dataclass class _KeyGuard: """Serializes access to one worker key's files. - ``file_lock`` (``threading.Lock``): correctness lock for the bytes, held on - the writing thread so no coroutine cancellation can release it mid-write. + ``file_lock`` (``_KeyLock``): correctness lock for the bytes, held on the + writing thread so no coroutine cancellation can release it mid-write. ``admission`` (``Semaphore(1)``): caps dispatched threads per key so one key cannot occupy the shared executor; not a correctness lock. ``waiters``: exact count of coroutines referencing this guard. @@ -111,7 +156,7 @@ class _KeyGuard: """ admission: asyncio.Lock - file_lock: threading.Lock + file_lock: _KeyLock waiters: int = 0 @@ -170,6 +215,19 @@ def __init__(self, queues_dir: Path): # AND waiters == 1 (see _guard / delete_drained). No sweeper, no # timer, no refcount map, no eviction on the hot path. self._guards: dict[str, _KeyGuard] = {} + # Serialises get-or-create of _guards. Guards are minted from the event + # loop (_guard) AND from worker threads (_key_guard, via the offset + # writer), and two guards for one key would mean two locks over the same + # bytes -- exactly the tear the guard exists to prevent. + self._guards_mutex = threading.Lock() + # Sessions with a dead-letter written since their last read_batch. A + # crash in the dead_letter->commit gap leaves a leading already-dead + # line uncommitted; the NEXT read_batch reconciles past it (once, then + # clears the flag) so a respawn cannot re-read and re-dead-letter it. + # In-memory only: process restart is covered by the boot-time global + # recovery_reconcile_dead(). Keeping this a dirty flag (not a per-poll + # scan) is what keeps the full .dead.jsonl read off the hot path. + self._dead_unreconciled: set[str] = set() @property def queues_dir(self) -> Path: @@ -199,24 +257,120 @@ def _compact_tmp_path(self, session_id: str) -> Path: """ return self._dir / f"{session_id}.log.compact.tmp" - def _read_committed_offset(self, session_id: str) -> int: - """Committed byte offset; reads bare-int and legacy JSON offset files.""" + def _offset_tmp_path(self, session_id: str) -> Path: + """A staging path used by exactly one ``.offset`` write. + + The uniqueness component is what makes concurrent writers independent: + a shared staging name lets one writer ``os.replace`` a file the other is + still filling, publishing a half-written record. Keeps the + ``.offset.tmp`` ending so ``reclaim_orphans`` still reaps strays. + """ + unique = f"{os.getpid()}-{uuid.uuid4().hex}" + return self._dir / f"{session_id}.{unique}{_OFFSET_TMP_SUFFIX}" + + def _offset_tmp_owner(self, path: Path) -> str: + """Session id a ``.offset.tmp`` staging file belongs to. + + Current names carry exactly one uniqueness component + (``..offset.tmp``); one left by an older build carries + none. Resolve to whichever form has a live ``.log`` so a staging file + beside a live session is never mistaken for an orphan. + """ + stem = path.name[: -len(_OFFSET_TMP_SUFFIX)] + stripped = stem.rsplit(".", 1)[0] + if stripped != stem and self._log_path(stripped).exists(): + return stripped + return stem + + def _write_offset_record( + self, session_id: str, offset: int, cursor: dict[str, Any] | None + ) -> None: + """Sole writer of the ``.offset`` file: one atomic JSON record. + + Takes the key's ``file_lock`` ITSELF rather than trusting callers to, + which is the invariant that keeps the record consistent: every path that + moves a session's offset -- commit, compaction's rebase, the + dead-letter reconcile -- is serialised here, at the one place it cannot + be forgotten. Reentrant, so a caller holding the lock across a wider + sequence still nests safely. + + Writes ``{"v": 1, "offset": offset, "cursor": cursor}`` to a per-write + temp file and ``os.replace``s it in, so a reader never sees a torn + record. Folding the cursor into the SAME record as the offset (not a + sidecar file) is what keeps the two from ever skewing: a crash loses + both together, never one without the other. No ``fsync`` -- + process-crash-durable, not power-durable, matching the rest of this + module. + """ + if cursor is not None and not isinstance(cursor, dict): + raise TypeError( + f"cursor must be a dict or None, got {type(cursor).__name__}" + ) + final = self._offset_path(session_id) + tmp = self._offset_tmp_path(session_id) + record = {"v": 1, "offset": offset, "cursor": cursor} + with self._key_guard(session_id).file_lock: + tmp.write_text(json.dumps(record, separators=(",", ":")), encoding="utf-8") + try: + os.replace(tmp, final) + except OSError: + # A failed rename leaves the prior committed record intact; drop + # the staged temp so a crashed write leaves nothing behind. + tmp.unlink(missing_ok=True) + raise + + def _read_offset_record(self, session_id: str) -> tuple[int, dict[str, Any] | None]: + """Read the ``.offset`` file, returning ``(offset, cursor)``. + + Accepts the current JSON record and the legacy bare-integer shape; a + missing or empty file yields ``(0, None)``. An unreadable cursor + degrades to ``None`` (an unknown ``v`` or non-dict ``cursor`` keeps the + offset but drops the cursor -- a corrupt cursor must not crash boot). An + unreadable OFFSET is NOT degraded: a malformed record raises, because + silently resetting a corrupt offset to 0 would replay the whole log and + manufacture duplicate nodes -- a worse, quieter failure than a loud one. + """ try: text = self._offset_path(session_id).read_text("utf-8") except FileNotFoundError: - return 0 + return 0, None text = text.strip() if not text: - return 0 + return 0, None if text[0] == "{": try: - cursor = json.loads(text) - return int(cursor["offset"]) + rec = json.loads(text) + offset = int(rec["offset"]) except (json.JSONDecodeError, KeyError, TypeError, ValueError): raise ValueError( f"unparseable legacy offset document for session {session_id!r}" ) from None - return int(text) + cursor = rec.get("cursor") + if rec.get("v") == 1 and isinstance(cursor, dict): + return offset, cursor + return offset, None + return int(text), None + + def _read_committed_offset(self, session_id: str) -> int: + """Committed byte offset; reads envelope, bare-int, and legacy JSON.""" + return self._read_offset_record(session_id)[0] + + def _read_committed_offset_degrading(self, session_id: str) -> int | None: + """Committed offset, or ``None`` when the ``.offset`` is unreadable. + + Factors the ``(OSError, ValueError)`` degrade the drain-loop reads + already apply (``read_batch`` -> the registry's ``_read_batch`` + quarantine) so the FINALIZE path -- ``delete_drained`` and + ``is_fully_drained``, reached via ``_finalize_session`` on every normal + ``session:end`` -- degrades to a safe no-op instead of letting a corrupt + ``.offset`` raise out to the drain supervisor. The boot ``RESET_OFFSET`` + reclaim (cursor-preserving) is what actually heals the file. + """ + try: + return self._read_committed_offset(session_id) + except (OSError, ValueError): + logger.exception("committed_offset_unreadable session=%s", session_id) + return None @staticmethod def _last_complete_end(path: Path) -> int: @@ -313,23 +467,46 @@ def _validate_session_id(session_id: str) -> None: @contextlib.contextmanager def _guard(self, worker_key: str) -> Iterator[_KeyGuard]: - """Get-or-create this key's guard and register this coroutine as a holder. - - The lookup and the ``waiters`` increment are one synchronous step with - no ``await`` between them, so an uncounted reference is impossible; the - ``finally`` decrements. Keep both statements synchronous -- a yield - point between them reintroduces the race. Every guarded operation uses - this. + """Get-or-create this key's guard and register this holder in ``waiters``. + + The get-or-create AND the ``waiters`` increment happen together under + ``_guards_mutex``, and ``delete_drained``'s eviction gate takes the same + mutex to read ``waiters`` and drop the entry. That mutual exclusion is + load-bearing now that ``_guard`` is called from WORKER THREADS as well as + the event loop (the boot reconcile and ``reclaim_orphans``): without it, + an eviction could slip into the window between the map lookup and the + increment, drop the guard this caller is about to hold, and let the next + writer mint a SECOND ``_KeyLock`` over the same file (the guard-eviction + clobber). The ``finally`` decrements under the mutex too. """ - guard = self._guards.get(worker_key) - if guard is None: - guard = _KeyGuard(asyncio.Lock(), threading.Lock()) - self._guards[worker_key] = guard - guard.waiters += 1 + with self._guards_mutex: + guard = self._guards.get(worker_key) + if guard is None: + guard = _KeyGuard(asyncio.Lock(), _KeyLock()) + self._guards[worker_key] = guard + guard.waiters += 1 try: yield guard finally: - guard.waiters -= 1 + with self._guards_mutex: + guard.waiters -= 1 + + def _key_guard(self, worker_key: str) -> _KeyGuard: + """This key's guard, minted on first use. Safe from any thread. + + Callers on the event loop should use ``_guard`` instead, so their + reference is counted in ``waiters``; this is the raw accessor for code + already running on a worker thread. + """ + guard = self._guards.get(worker_key) + if guard is not None: + return guard + with self._guards_mutex: + guard = self._guards.get(worker_key) + if guard is None: + guard = _KeyGuard(asyncio.Lock(), _KeyLock()) + self._guards[worker_key] = guard + return guard @staticmethod def _write_all(fd: int, data: bytes) -> None: @@ -504,49 +681,94 @@ async def append(self, session_id: str, raw: bytes) -> None: ) async def read_batch(self, session_id: str, max_items: int) -> Batch: + """Read up to ``max_items`` uncommitted records, oldest first. + + Runs under the key's guard, like every other path that touches these + bytes: it can advance the offset (the dead-letter reconcile below), and + even its pure read must not observe the ``.log`` mid-swap while + compaction rebuilds it. + """ self._validate_session_id(session_id) path = self._log_path(session_id) - def _read() -> Batch: - start = self._read_committed_offset(session_id) - records: list[Record] = [] - consumed = 0 - try: - with open(path, "rb") as f: - f.seek(start) - while len(records) < max_items: - raw = f.readline() - if not raw or not raw.endswith(b"\n"): - # EOF, or a torn trailing line with no newline yet: - # ignore the partial line and stop on a line boundary. - break - rec_start = start + consumed - consumed += len(raw) - rec_end = start + consumed - records.append(Record(raw[:-1], rec_start, rec_end)) - except FileNotFoundError: - pass - return Batch(session_id, records, start, start + consumed) - - return await asyncio.to_thread(_read) - - async def commit(self, session_id: str, new_offset: int) -> None: - """Atomically and durably persist ``new_offset`` (the ack). + def _read(guard: _KeyGuard) -> Batch: + with guard.file_lock: + # Reconcile past any leading already-dead line ONLY when a + # dead-letter was written since the last read (the dirty flag) + # -- never per poll, so the full .dead.jsonl read stays off the + # hot path. Take-and-clear in one step so a dead_letter racing + # this read leaves the flag set for the NEXT read rather than + # having its request consumed here. Clearing is unconditional: + # _reconcile_dead_key is fault-isolated (swallows + # OSError/ValueError -> 0), so a retry would not help, and + # leaving the flag set would re-scan every poll. + if session_id in self._dead_unreconciled: + self._dead_unreconciled.discard(session_id) + self._reconcile_dead_key(session_id) + start = self._read_committed_offset(session_id) + records: list[Record] = [] + consumed = 0 + try: + with open(path, "rb") as f: + f.seek(start) + while len(records) < max_items: + raw = f.readline() + if not raw or not raw.endswith(b"\n"): + # EOF, or a torn trailing line with no newline + # yet: ignore the partial line and stop on a + # line boundary. + break + rec_start = start + consumed + consumed += len(raw) + rec_end = start + consumed + records.append(Record(raw[:-1], rec_start, rec_end)) + except FileNotFoundError: + pass + return Batch(session_id, records, start, start + consumed) - Writes the offset to a temp file and uses ``os.replace`` for an atomic - rename, so a reader never observes a torn or partial offset file. No - ``fsync`` is issued: the offset survives a process crash but not a - power loss. + with self._guard(session_id) as guard: + async with guard.admission: + return await _await_uninterrupted(asyncio.to_thread(_read, guard)) + + async def commit( + self, session_id: str, new_offset: int, cursor: dict[str, Any] | None + ) -> None: + """Atomically and durably persist ``new_offset`` (the ack) + ``cursor``. + + ``cursor`` has NO default: every call site must pass it explicitly so it + is never silently omitted -- a missed site (or a plain rolling deploy + against an older signature) would null the cursor and reset the + cross-handler counters, manufacturing duplicate nodes on the next run. + Offset and cursor are written in the SAME atomic record, so they can + never skew. No ``fsync``: process-crash-durable, not power-durable. """ self._validate_session_id(session_id) - final = self._offset_path(session_id) - tmp = self._dir / f"{session_id}.offset.tmp" + # cursor type is enforced by _write_offset_record (raises TypeError on a + # non-dict, non-None cursor) -- the sole writer, so the check lives in + # one place and survives -O rather than as a bypassable assert here. + + def _commit(guard: _KeyGuard) -> None: + # Same file_lock compact_committed_prefix takes: a commit that ran + # unlocked could interleave with compaction's read-committed / + # rebase-to-0 / replace-log sequence and be silently erased (or + # erase the rebase). Serializing both under the per-key lock closes + # that race. + with guard.file_lock: + self._write_offset_record(session_id, new_offset, cursor) - def _commit() -> None: - tmp.write_text(str(new_offset), encoding="utf-8") - os.replace(tmp, final) + with self._guard(session_id) as guard: + async with guard.admission: + await _await_uninterrupted(asyncio.to_thread(_commit, guard)) + + async def read_cursor(self, session_id: str) -> dict[str, Any] | None: + """Return the persisted cursor for ``session_id``, or ``None``. - await asyncio.to_thread(_commit) + ``None`` covers: no offset file, a legacy bare-integer offset file, or a + JSON record whose cursor is absent/unreadable. This is the read side a + rebuilt worker uses to restore its cross-handler counters. + """ + self._validate_session_id(session_id) + return await asyncio.to_thread(lambda: self._read_offset_record(session_id)[1]) async def dead_letter(self, session_id: str, raw: bytes, error: str) -> None: """Append one dead-letter record for an unprocessable batch line. @@ -581,6 +803,47 @@ async def dead_letter(self, session_id: str, raw: bytes, error: str) -> None: # (traceback carries the exception; exc is not repeated.) logger.exception("dead_letter_write_failed session=%s", session_id) raise + # The dead-letter is durable but the offset that accounts for it may not + # be committed yet (a crash in the gap leaves a leading already-dead + # line). Mark the session so the NEXT read_batch reconciles past it once. + self._dead_unreconciled.add(session_id) + + def _retired_dead_paths(self, session_id: str) -> list[Path]: + """This session's retired dead-letter files, oldest first. + + The retirement timestamp sorts lexically, so plain sorting is append + order. + """ + return sorted(self._dir.glob(f"{session_id}.finalized-*.dead.jsonl")) + + def _retire_dead_letters(self, session_id: str) -> None: + """Move a finalized session's dead letters out of its own name. + + A session id can be reused. Left in place, the previous session's + ``.dead.jsonl`` would make the new session's reconcile pass skip log + lines whose bytes merely match an OLD dead payload -- committing past + events that were never processed. Renaming ends that: the reconcile + only ever reads the live name. Nothing is discarded -- the retired name + still ends in ``.dead.jsonl``, so ``expire_dead_letters`` ages it out on + its own schedule and ``read_dead_letters`` still reports it under this + session. Also clears the in-memory reconcile flag, which would + otherwise make the reused id pay for a scan it cannot benefit from. + + Caller must hold the key's ``file_lock`` (the rename must not race a + ``dead_letter`` append). Failure is logged, not raised: retiring is + hygiene, and failing the delete over it would retain a drained log. + """ + self._dead_unreconciled.discard(session_id) + dead = self._dead_path(session_id) + retired = self._dir / f"{session_id}.finalized-{time.time_ns()}.dead.jsonl" + try: + os.replace(dead, retired) + except FileNotFoundError: + return + except OSError: + logger.exception("dead_letter_retire_failed session=%s", session_id) + return + logger.info("dead_letters_retired session=%s path=%s", session_id, retired.name) async def delete_drained(self, session_id: str) -> bool: """Remove the drained ``.log``/``.offset`` for a finalized session. @@ -590,8 +853,10 @@ async def delete_drained(self, session_id: str) -> bool: append. Refuses (returns False) if the log still has uncommitted bytes; the caller re-drains and retries a bounded number of times, and ``recover()`` picks up any give-up. A missing ``.log`` still unlinks a - stale ``.offset`` (else a recreated log reads past its own end). Keeps - ``.dead.jsonl``. Idempotent. + stale ``.offset`` (else a recreated log reads past its own end). + Dead letters are retained, but retired out of the session's own name so + a later session reusing the id starts clean (``_retire_dead_letters``). + Idempotent. The guard-map entry is dropped only when ``waiters == 1`` and identity matches; otherwise a still-referencing coroutine could later lock a @@ -616,9 +881,21 @@ def _delete(guard: _KeyGuard) -> bool: offset.unlink() except FileNotFoundError: pass + self._retire_dead_letters(session_id) return True - committed = self._read_committed_offset(session_id) + committed = self._read_committed_offset_degrading(session_id) + if committed is None: + # Corrupt .offset at finalize: we cannot prove the log is + # fully committed, so do NOT delete and do NOT raise out to + # the drain supervisor. Retained + recovered by the boot + # RESET_OFFSET pass; returning False routes into finalize's + # bounded retain/give-up path, identical to uncommitted-bytes. + logger.warning( + "delete_drained_retained session=%s reason=corrupt_offset", + session_id, + ) + return False if size > committed: logger.warning( "delete_drained_retained session=%s uncommitted_bytes=%d", @@ -627,6 +904,12 @@ def _delete(guard: _KeyGuard) -> bool: ) return False + # Terminal cleanup: remove BOTH files. The cursor is intentionally + # dropped here -- a finalized session is done, and an orch_run_id + # is scoped by execution_start_ts, so a genuinely-new post-finalize + # run gets a distinct id regardless of the seq counter (no + # collision). Keeping the .offset would only leak a file per + # finalized session, which is why finalize removes it. try: log.unlink() except FileNotFoundError: @@ -635,17 +918,29 @@ def _delete(guard: _KeyGuard) -> bool: offset.unlink() except FileNotFoundError: pass + self._retire_dead_letters(session_id) return True with self._guard(session_id) as guard: async with guard.admission: ok = await _await_uninterrupted(asyncio.to_thread(_delete, guard)) # Still holding admission: apply the three-part removal - # condition. waiters == 1 is THIS call itself; - # anything higher means another coroutine holds the guard - # and removal must be skipped. - if ok and guard.waiters == 1 and self._guards.get(session_id) is guard: - del self._guards[session_id] + # condition UNDER ``_guards_mutex`` -- the same lock ``_guard`` + # increments ``waiters`` under. That mutual exclusion is what + # makes ``eviction can't happen underneath a live holder`` true + # even against a holder running on a worker thread (the boot + # reconcile / reclaim): the check-and-delete can no longer slip + # into the window between another caller's map lookup and its + # ``waiters`` increment. waiters == 1 is THIS call itself; + # anything higher means another holder exists and removal is + # skipped. + with self._guards_mutex: + if ( + ok + and guard.waiters == 1 + and self._guards.get(session_id) is guard + ): + del self._guards[session_id] return ok async def compact_committed_prefix( @@ -667,14 +962,12 @@ async def compact_committed_prefix( """ self._validate_session_id(session_id) log = self._log_path(session_id) - offset = self._offset_path(session_id) - offset_tmp = self._dir / f"{session_id}.offset.tmp" tmp = self._compact_tmp_path(session_id) def _compact(_guard: _KeyGuard) -> int: with _guard.file_lock: try: - c = self._read_committed_offset(session_id) + c, cursor = self._read_offset_record(session_id) except (OSError, ValueError): return 0 try: @@ -732,10 +1025,11 @@ def _compact(_guard: _KeyGuard) -> int: return 0 # Step 5: rebase the offset to 0 FIRST -- the point of no - # return. + # return. The cursor rides along unchanged: rebasing the log to + # its undrained tail does not roll back cross-handler state, so + # a rebuild after compaction must still see the committed cursor. try: - offset_tmp.write_text("0", encoding="utf-8") - os.replace(offset_tmp, offset) + self._write_offset_record(session_id, 0, cursor) except OSError: logger.exception( "compact_offset_rebase_failed session=%s", session_id @@ -748,13 +1042,12 @@ def _compact(_guard: _KeyGuard) -> int: try: os.replace(tmp, log) except OSError: - # R3: restore the offset to C so this becomes a PURE - # NO-OP -- never an in-process re-drive (which would - # double-count `written` and drive the residual - # negative). + # Restore the offset to the committed value so a failed + # compaction is a pure no-op, never an in-process re-drive + # (which would double-count `written` and drive the + # residual negative). The cursor is restored with it. try: - offset_tmp.write_text(str(c), encoding="utf-8") - os.replace(offset_tmp, offset) + self._write_offset_record(session_id, c, cursor) logger.error( "compact_replace_failed session=%s committed=%d " "action=offset_restored", @@ -803,38 +1096,46 @@ def _compact(_guard: _KeyGuard) -> int: async def read_dead_letters(self, session_id: str) -> list[dict]: """Return all dead-letter records for ``session_id`` in append order. - Returns an empty list when no dead-letter file exists. A malformed - line is skipped (logged once, not per line) rather than raising -- - reached by ``GET /queues/dead-letter/{key}`` and the replay path, and - a malformed record must not 500 an operator endpoint or abort a - replay. + Spans the retired files an earlier finalize left behind + (``_retire_dead_letters``), oldest first, so an operator still sees + everything ever dead-lettered under this id. Returns an empty list when + the session has none. A malformed line is skipped (logged once per + file, not per line) rather than raising -- reached by ``GET + /queues/dead-letter/{key}`` and the replay path, and a malformed record + must not 500 an operator endpoint or abort a replay. """ self._validate_session_id(session_id) def _read() -> list[dict]: - try: - text = self._dead_path(session_id).read_text(encoding="utf-8") - except FileNotFoundError: - return [] records: list[dict] = [] - skipped = 0 - for ln in text.splitlines(): - if not ln.strip(): - continue + for path in [ + *self._retired_dead_paths(session_id), + self._dead_path(session_id), + ]: try: - records.append(json.loads(ln)) - except ( - json.JSONDecodeError, - UnicodeDecodeError, - ValueError, - TypeError, - ): - skipped += 1 + text = path.read_text(encoding="utf-8") + except FileNotFoundError: continue - if skipped: - logger.warning( - "dead_letter_unparseable key=%s skipped=%d", session_id, skipped - ) + skipped = 0 + for ln in text.splitlines(): + if not ln.strip(): + continue + try: + records.append(json.loads(ln)) + except ( + json.JSONDecodeError, + UnicodeDecodeError, + ValueError, + TypeError, + ): + skipped += 1 + continue + if skipped: + logger.warning( + "dead_letter_unparseable key=%s skipped=%d", + path.name, + skipped, + ) return records return await asyncio.to_thread(_read) @@ -1126,12 +1427,25 @@ def _apply(guard: _KeyGuard) -> bool: ) try: if c.verdict is Verdict.RESET_OFFSET: - # Unlink ONLY the offset (+ any stray .offset.tmp) -- - # the .log stays; the next drain re-reads from 0. + # Reset the committed offset to 0 (the .log stays and + # re-drains from 0), PRESERVING the cursor. After + # compaction the .log is the undrained tail and the + # committed cursor is the correct starting state for it; + # wiping it would re-mint ids already persisted for the + # reclaimed prefix. A missing/unreadable cursor (a truly + # unparseable .offset) degrades to a bare removal, which + # reads back as committed 0 all the same. try: - offset.unlink() - except FileNotFoundError: - pass + _, _cursor = self._read_offset_record(c.key) + except (OSError, ValueError): + _cursor = None + if _cursor is not None: + self._write_offset_record(c.key, 0, _cursor) + else: + try: + offset.unlink() + except FileNotFoundError: + pass try: offset_tmp.unlink() except FileNotFoundError: @@ -1219,38 +1533,51 @@ def _scan() -> dict[str, int]: *((p, "orphan_offset") for p in offset_paths), *((p, "orphan_offset_tmp") for p in tmp_paths), ]: - stem = path.name[ - : -len(".offset.tmp" if reason.endswith("tmp") else ".offset") - ] - if self._log_path(stem).exists(): - continue - try: - size = path.stat().st_size - except OSError: - failed += 1 - logger.exception( - "boot_reclaim_failed reason=%s path=%s", reason, path - ) - continue - logger.warning( - "boot_reclaimed reason=%s path=%s session=%s bytes=%d action=%s", - reason, - path, - stem, - size, - action, + stem = ( + self._offset_tmp_owner(path) + if reason.endswith("tmp") + else path.name[: -len(".offset")] ) - if not enabled: + if self._log_path(stem).exists(): continue - try: - path.unlink() - reclaimed += 1 - reclaimed_bytes += size - except OSError: - failed += 1 - logger.exception( - "boot_reclaim_failed reason=%s path=%s", reason, path + # Cheap pre-filter above avoids minting a guard for obviously + # live keys; the authoritative decision is re-made UNDER the + # same counted per-key lock every .offset writer takes, so this + # unlink can never race a concurrent offset write/os.replace (or + # delete) that would otherwise bypass the lock. `continue` + # inside the `with` still runs the guard's finally (decrements + # `waiters`) and releases file_lock. + with self._guard(stem) as guard, guard.file_lock: + if self._log_path(stem).exists(): + continue + try: + size = path.stat().st_size + except OSError: + failed += 1 + logger.exception( + "boot_reclaim_failed reason=%s path=%s", reason, path + ) + continue + logger.warning( + "boot_reclaimed reason=%s path=%s session=%s bytes=%d " + "action=%s", + reason, + path, + stem, + size, + action, ) + if not enabled: + continue + try: + path.unlink() + reclaimed += 1 + reclaimed_bytes += size + except OSError: + failed += 1 + logger.exception( + "boot_reclaim_failed reason=%s path=%s", reason, path + ) for path in torn_paths: try: mtime = path.stat().st_mtime @@ -1381,6 +1708,28 @@ def _scan() -> list[str]: return await asyncio.to_thread(_scan) + async def is_fully_drained(self, session_id: str) -> bool: + """True iff the session has no undrained log data left. + + Compares the committed offset against the end of complete (newline- + terminated) data, read straight from disk -- so it is independent of + in-memory worker liveness and stays correct across a crash + restart. A + session with no ``.log`` reads as drained (0 >= 0): it never had data. + """ + self._validate_session_id(session_id) + + def _check() -> bool: + committed = self._read_committed_offset_degrading(session_id) + if committed is None: + # Corrupt .offset: we cannot prove the session is drained, so + # report NOT drained (conservative) instead of raising out to + # the finalize/drain supervisor. Same degrade delete_drained + # applies; the boot RESET_OFFSET pass heals the file. + return False + return committed >= self._complete_data_end(session_id) + + return await asyncio.to_thread(_check) + async def recover(self) -> list[str]: """Return sorted session_ids that have a complete unprocessed line. @@ -1806,6 +2155,58 @@ def _dead_payload_set(self, worker_key: str) -> set[bytes]: ) return payloads + def _reconcile_dead_key(self, key: str) -> int: + """Advance ONE key's committed offset past its leading already-dead lines. + + Fault-isolated: a corrupt/unreadable dead-payload set or offset for this + key returns 0 (logged) rather than raising, so neither the boot pass nor + the per-session respawn path it feeds can crash-loop on one bad key. The + committed cursor rides along unchanged -- skipping already-dead lines + does not change cross-handler state. + + Holds the key's ``file_lock`` across the read AND the write via the + COUNTED ``_guard()`` -- not the raw ``_key_guard()`` accessor. This is a + read-modify-write of the offset, so a commit landing between the two + would be silently rolled back to the value read here. Routing through + ``_guard()`` also registers this caller in ``waiters``, so + ``delete_drained``'s eviction gate (``waiters == 1``) cannot drop the + guard-map entry out from under an in-flight reconcile -- which would let + the next writer mint a SECOND ``_KeyLock`` over the same file and + clobber it (the guard-eviction race). That counting is what lets the + boot sweep run concurrently with live drainers. + """ + # Q-6: check `.log` existence BEFORE reading the whole `.dead.jsonl` + # into RAM. A key with only a `.dead.jsonl` (the common shape left by + # `delete_drained`) pays nothing here. + log_path = self._log_path(key) + if not log_path.exists(): + return 0 + try: + dead_payloads = self._dead_payload_set(key) + if not dead_payloads: + return 0 + with self._guard(key) as guard, guard.file_lock: + committed, cursor = self._read_offset_record(key) + complete_end = self._complete_data_end(key) + pos = committed + skipped = 0 + with open(log_path, "rb") as f: + f.seek(committed) + while pos < complete_end: + raw = f.readline() + if not raw or not raw.endswith(b"\n"): + break + if raw[:-1] not in dead_payloads: + break + pos += len(raw) + skipped += 1 + if pos > committed: + self._write_offset_record(key, pos, cursor) + return skipped + except (OSError, ValueError): + logger.exception("recovery_reconcile_dead_key_failed key=%s", key) + return 0 + async def recovery_reconcile_dead(self) -> int: """Advance committed offsets past leading already-dead pending lines. @@ -1821,43 +2222,7 @@ async def recovery_reconcile_dead(self) -> int: def _reconcile() -> int: total_skipped = 0 for key in self._all_worker_keys(): - # Q-6: check `.log` existence BEFORE reading the whole - # `.dead.jsonl` into RAM (+ a payload set at ~3.6x its size). - # A key with only a `.dead.jsonl` (the common shape left by - # `delete_drained`) previously paid that read for nothing, - # every boot, forever. Free fix. - log_path = self._log_path(key) - if not log_path.exists(): - continue - # Fault-isolate this key -- a corrupt/unreadable - # dead-payload set or offset for ONE key must not abort the - # reconcile pass for every other key. The boot - # hook this feeds must never crash-loop the share it reads. - try: - dead_payloads = self._dead_payload_set(key) - if not dead_payloads: - continue - committed = self._read_committed_offset(key) - complete_end = self._complete_data_end(key) - pos = committed - with open(log_path, "rb") as f: - f.seek(committed) - while pos < complete_end: - raw = f.readline() - if not raw or not raw.endswith(b"\n"): - break - if raw[:-1] not in dead_payloads: - break - pos += len(raw) - total_skipped += 1 - if pos > committed: - final = self._offset_path(key) - tmp = self._dir / f"{key}.offset.tmp" - tmp.write_text(str(pos), encoding="utf-8") - os.replace(tmp, final) - except (OSError, ValueError): - logger.exception("recovery_reconcile_dead_key_failed key=%s", key) - continue + total_skipped += self._reconcile_dead_key(key) self._stats_cache = None return total_skipped diff --git a/context_intelligence_server/queue_manager/protocol.py b/context_intelligence_server/queue_manager/protocol.py index 769993a8..a439267e 100644 --- a/context_intelligence_server/queue_manager/protocol.py +++ b/context_intelligence_server/queue_manager/protocol.py @@ -80,7 +80,11 @@ async def append(self, session_id: str, raw: bytes) -> None: ... async def read_batch(self, session_id: str, max_items: int) -> Batch: ... - async def commit(self, session_id: str, new_offset: int) -> None: ... + async def commit( + self, session_id: str, new_offset: int, cursor: dict[str, Any] | None + ) -> None: ... + + async def read_cursor(self, session_id: str) -> dict[str, Any] | None: ... async def dead_letter(self, session_id: str, raw: bytes, error: str) -> None: ... diff --git a/context_intelligence_server/registry.py b/context_intelligence_server/registry.py index cbef0304..fe0ef21b 100644 --- a/context_intelligence_server/registry.py +++ b/context_intelligence_server/registry.py @@ -39,6 +39,12 @@ _RESIDUAL_DEGRADED_GRACE = 15.0 +class _SessionQuarantined(Exception): + """Raised after a session has been closed and deregistered for a permanent + fault. The drain task must unwind to its entry point and return without + touching the worker again.""" + + @dataclass class SessionWorker: session_id: str @@ -353,10 +359,22 @@ async def drain_worker( poll_interval = min(flush_timeout, _DRAIN_POLL_INTERVAL) idle_elapsed = 0.0 attempts = 0 + # A (re)built worker starts with empty cross-handler counters. Restore the + # cursor the last commit persisted -- lazily, just before the FIRST real + # batch is processed -- so a rebuild resumes those counters instead of + # restarting from zero (which would remint node ids and duplicate them). + # Kept off the idle path so an idle worker never does this read. + cursor_restored = False + # Cursor state as it was BEFORE the first attempt on the current batch. + # A failed attempt leaves the cross-handler counters advanced; replaying + # the same batch from that advanced state mints fresh node ids and + # duplicates the Iteration. Restoring this snapshot before each retry + # makes the replay reproduce the SAME ids (idempotent MERGE). + pre_batch_cursor: dict[str, Any] | None = None while True: try: - batch = await qm.read_batch(session_id, max_items=_DRAIN_MAX_BATCH) + batch = await self._read_batch(worker, _DRAIN_MAX_BATCH) if not batch.records: # Idle compaction runs before the dry-exit check below, so a @@ -369,7 +387,7 @@ async def drain_worker( # Dry-exit for a recovered drainer with no terminal record: re-read # after the await closes the race with a live POST arriving mid-check. if not worker.live_event_seen: - recheck = await qm.read_batch(session_id, max_items=1) + recheck = await self._read_batch(worker, 1) if not recheck.records and not worker.live_event_seen: await self._safe_close(worker) self._deregister(session_id) @@ -402,6 +420,35 @@ async def drain_worker( idle_elapsed = 0.0 + # Restore the durable cursor once, before this worker's FIRST + # real batch is processed (and before the pre-attempt snapshot + # below, so that snapshot captures the restored state). A + # brand-new session has no committed cursor (read_cursor -> None + # -> no-op). + if not cursor_restored: + # A corrupt or legacy-unparseable offset record must degrade + # the CURSOR to None (restore nothing, keep defaults), never + # kill the worker: the committed offset is read separately by + # read_batch, so a lost cursor costs at most a bounded replay, + # while an unhandled raise here would crash-loop the drainer. + try: + persisted_cursor = await qm.read_cursor(session_id) + except (OSError, ValueError): + logger.warning( + "cursor_read_failed session=%s action=degrade_to_none", + session_id, + exc_info=True, + extra={"session_id": session_id}, + ) + persisted_cursor = None + worker.services.restore_cursor(persisted_cursor) + cursor_restored = True + + # Snapshot once, before the first attempt on this batch, so a + # retry can roll the counters back (see pre_batch_cursor above). + if attempts == 0: + pre_batch_cursor = worker.services.snapshot_cursor() + # --- dispatch + durable write barrier, one error path --- try: safe_count, terminal_at = await self._process_batch( @@ -452,7 +499,7 @@ async def drain_worker( if attempts >= self._max_delivery_attempts: # Budget spent: isolate the batch line-by-line and dead-letter. terminal_seen = await self._handle_exhausted_batch( - worker, batch, handlers + worker, batch, handlers, pre_batch_cursor ) if terminal_seen: # Mirror the normal terminal branch below: the @@ -462,8 +509,11 @@ async def drain_worker( return attempts = 0 continue - # Not yet exhausted: back off before re-reading the same - # offset (idempotent MERGE makes the replay a no-op). + # Not yet exhausted: roll the cross-handler counters back to + # their pre-attempt state so the replay reproduces the same + # node ids (idempotent MERGE) instead of duplicating them, + # then back off before re-reading the same offset. + worker.services.restore_cursor(pre_batch_cursor) await asyncio.sleep(poll_interval) continue @@ -471,7 +521,9 @@ async def drain_worker( # Commit only up to session:end -- leaving it uncommitted makes # "ended but not finalized" durable across a respawn/recover(). commit_to = batch.end_offset if terminal_at is None else terminal_at - await qm.commit(session_id, commit_to) + await qm.commit( + session_id, commit_to, worker.services.snapshot_cursor() + ) counted = len(batch.records) if terminal_at is None else safe_count self.record_written(counted) logger.debug( @@ -495,6 +547,11 @@ async def drain_worker( await self._finalize_session(worker, handlers) return + except _SessionQuarantined: + # _read_batch already closed, deregistered, and logged. Reached + # from any read in this loop, including the finalize tail drain. + return + except asyncio.CancelledError: # Cancelled while reading/idle (outer site; never reaches the inner try). logger.info( @@ -541,7 +598,11 @@ async def _process_batch( return safe_count, terminal_at async def _handle_exhausted_batch( - self, worker: SessionWorker, batch: Batch, handlers: Any + self, + worker: SessionWorker, + batch: Batch, + handlers: Any, + pre_batch_cursor: dict[str, Any] | None, ) -> bool: """Reprocess a poison batch one line at a time (linear isolation). @@ -566,9 +627,15 @@ async def _handle_exhausted_batch( """ qm = self.queue_manager session_id = worker.session_id + # Roll the cross-handler counters back to their pre-batch state, exactly + # as the in-place-retry branch does. The failed batch attempts advanced + # those counters in memory; without this rollback, each isolated + # record's per-record commit would snapshot the already-advanced state + # and mint fresh ids, double-counting on the dead-letter redrive. + worker.services.restore_cursor(pre_batch_cursor) # The failed batch flush left writes resident in the store buffer -- # discard so the first isolated record flushes from a clean buffer. - worker.services.graph.discard_buffer() + worker.services.discard_buffered_writes() for rec in batch.records: try: event, _ws, data = self._parse_line(rec.raw) @@ -582,8 +649,10 @@ async def _handle_exhausted_batch( exc_info=exc, extra={"session_id": session_id}, ) - worker.services.graph.discard_buffer() - await qm.commit(session_id, rec.end) # queue-produced offset + worker.services.discard_buffered_writes() + await qm.commit( + session_id, rec.end, worker.services.snapshot_cursor() + ) # queue-produced offset continue from context_intelligence_server.pipeline import TERMINAL_EVENTS @@ -608,8 +677,10 @@ async def _handle_exhausted_batch( # Drop the failed record's residue so it cannot contaminate # the NEXT record's flush. A successful flush clears the # buffer itself; only the failure path needs this. - worker.services.graph.discard_buffer() - await qm.commit(session_id, rec.end) # queue-produced offset + worker.services.discard_buffered_writes() + await qm.commit( + session_id, rec.end, worker.services.snapshot_cursor() + ) # queue-produced offset if wrote: self.record_written(1) return False @@ -623,7 +694,7 @@ async def _drain_to_eof(self, worker: SessionWorker, handlers: Any) -> bool: qm = self.queue_manager session_id = worker.session_id while True: - tail = await qm.read_batch(session_id, max_items=_DRAIN_MAX_BATCH) + tail = await self._read_batch(worker, _DRAIN_MAX_BATCH) if not tail.records: return True try: @@ -632,7 +703,9 @@ async def _drain_to_eof(self, worker: SessionWorker, handlers: Any) -> bool: except Exception: logger.exception("finalize_tail_flush_failed session=%s", session_id) return False # NOT finalized: keep worker alive, tail uncommitted - await qm.commit(session_id, tail.end_offset) + await qm.commit( + session_id, tail.end_offset, worker.services.snapshot_cursor() + ) self.record_written(len(tail.records)) logger.debug( "batch_committed events=%d offset=%d", @@ -718,6 +791,32 @@ async def _finalize_session(self, worker: SessionWorker, handlers: Any) -> None: extra={"session_id": session_id}, ) + async def _read_batch(self, worker: SessionWorker, max_items: int) -> Batch: + """Read a batch for ``worker``, quarantining it on a corrupt offset. + + A ``ValueError`` means the committed offset is genuinely unparseable -- + a PERMANENT fault, so respawning would only reproduce it forever. The + session is closed and deregistered and ``_SessionQuarantined`` is + raised, which unwinds the drain task without touching the worker again; + other sessions keep draining and the boot RESET_OFFSET pass heals the + offset (cursor-preserving). A transient ``OSError`` is deliberately NOT + caught: it propagates to ``_on_drain_done`` for supervised respawn. + + Every read in the drain loop goes through here, so no call site can be + the one that lets a corrupt offset escape to the supervisor. + """ + try: + return await self.queue_manager.read_batch(worker.session_id, max_items) + except ValueError: + logger.exception( + "drain_worker_quarantined session=%s reason=corrupt_offset", + worker.session_id, + extra={"session_id": worker.session_id}, + ) + await self._safe_close(worker) + self._deregister(worker.session_id) + raise _SessionQuarantined(worker.session_id) from None + async def _safe_close(self, worker: SessionWorker) -> None: """Close the graph store. A worker whose store has been closed is never revived (see ``start_drain``'s guard) -- mark it FIRST, before diff --git a/context_intelligence_server/services.py b/context_intelligence_server/services.py index d861ca6a..7182af82 100644 --- a/context_intelligence_server/services.py +++ b/context_intelligence_server/services.py @@ -7,8 +7,11 @@ from __future__ import annotations +import copy +import dataclasses import fnmatch import logging +from dataclasses import asdict from datetime import datetime from typing import Any @@ -240,10 +243,91 @@ def __init__( self.data_layer_2 = DataLayer2State() self.data_layer_3 = DataLayer3State() + # ------------------------------------------------------------------ + # Durable cursor (survives a worker rebuild) + # ------------------------------------------------------------------ + + def snapshot_cursor(self) -> dict[str, Any]: + """Return a JSON-safe snapshot of cross-handler cursor state. + + Snapshots the whole ``data_layer_2``/``data_layer_3`` dataclasses via + ``asdict`` (every field is JSON-native) rather than a hand-picked + allowlist that would silently miss a newly added field. + """ + return { + "dl2": asdict(self.data_layer_2), + "dl3": asdict(self.data_layer_3), + } + + def restore_cursor(self, record: dict[str, Any] | None) -> None: + """Restore cross-handler cursor state from a persisted snapshot. + + No-op on ``record is None`` (a brand-new session, or a legacy + ``.offset`` with no cursor). Otherwise only field NAMES present on the + current dataclass are assigned -- unknown/renamed keys are dropped and a + field absent from the record keeps its default, so the persisted format + tolerates dataclass evolution in both directions without a version bump. + + All-or-nothing: both replacement dataclasses are built and validated + BEFORE either live field is reassigned, so a failure partway through can + never leave a half-restored hybrid. Any failure is caught and logged, + leaving the dataclasses untouched: a corrupt or unexpected cursor must + never crash boot. + """ + if record is None: + return + try: + rebuilt: list[tuple[str, Any]] = [] + for key, target in ( + ("dl2", self.data_layer_2), + ("dl3", self.data_layer_3), + ): + value = record.get(key) + if not isinstance(value, dict): + continue + valid_fields = {f.name for f in dataclasses.fields(type(target))} + # Deep-copy every override value: pre_batch_cursor is snapshotted + # ONCE and replayed on up to _max_delivery_attempts retries, so a + # mutable field (e.g. pending_tool_block_ids dict, + # active_recipe_run_stack list) shared by reference would let an + # attempt's in-place mutation corrupt the baseline the NEXT + # rollback restores. Copying severs the live state from the + # snapshot so each restore reproduces the identical baseline. + overrides = { + name: copy.deepcopy(v) + for name, v in value.items() + if name in valid_fields + } + # replace() builds a fresh instance; nothing is mutated in place + # until every target below has been built successfully. + rebuilt.append((key, dataclasses.replace(target, **overrides))) + except Exception: + logger.warning("cursor_restore_failed", exc_info=True) + return + for key, new_state in rebuilt: + if key == "dl2": + self.data_layer_2 = new_state + else: + self.data_layer_3 = new_state + # ------------------------------------------------------------------ # Session node management # ------------------------------------------------------------------ + def discard_buffered_writes(self) -> None: + """Drop the graph store's buffered writes AND the seen-session cache. + + ``ensure_session_node`` marks a session id in ``_seen_sessions`` after a + merely BUFFERED upsert. When the isolation path discards that buffer the + node was never flushed, so the cache is now a lie: the re-dispatch would + early-return and never re-issue the Session node (losing its + status/started_at/StubSession). Clearing the cache with the buffer keeps + the two in lockstep -- a re-issued node is an idempotent MERGE, never a + duplicate. + """ + self.graph.discard_buffer() + self._seen_sessions.clear() + async def ensure_session_node(self, session_id: str, data: dict[str, Any]) -> None: """Idempotently create a Session node in the graph for *session_id*. diff --git a/pyproject.toml b/pyproject.toml index e1fe9974..7e771dfc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "context-intelligence-server" -version = "6.7.2" +version = "7.0.0" description = "Context Intelligence Server for Amplifier" requires-python = ">=3.11" dependencies = [ diff --git a/tests/handlers/data_layer_2/test_content_block.py b/tests/handlers/data_layer_2/test_content_block.py index 53746570..32baefef 100644 --- a/tests/handlers/data_layer_2/test_content_block.py +++ b/tests/handlers/data_layer_2/test_content_block.py @@ -3,7 +3,7 @@ Covers: - handled_events == frozenset({'content_block:start', 'content_block:end'}) - content_block:start creates ContentBlock:SST_EVENT node keyed as - '{session_id}::block::{iteration_n}::{block_index}' with session_id, block_index, + '{iteration_id}::block::{block_index}' with session_id, block_index, started_at; iteration_n extracted from active_iteration_id cursor (split('::')[-1]) - E07: Iteration -[:HAS_PART {sst_semantic: 'CONTAINS'}]-> ContentBlock created when active_iteration_id is set; NOT created when no active iteration (zero edges) @@ -53,7 +53,7 @@ class TestContentBlockStartCreatesNode: async def test_node_created_with_correct_compound_key( self, services: HookStateService ) -> None: - """content_block:start must create node at '{session_id}::block::{iteration_n}::{block_index}'.""" + """content_block:start must create node at '{iteration_id}::block::{block_index}'.""" services.data_layer_2.active_iteration_id = "s1::iteration::1" handler = ContentBlockHandler(services) await handler( @@ -64,7 +64,7 @@ async def test_node_created_with_correct_compound_key( "block_index": 0, }, ) - node_id = "s1::block::1::0" + node_id = "s1::iteration::1::block::0" node = await services.graph.get_node(node_id) assert node is not None, f"content_block:start must create node at '{node_id}'" @@ -82,7 +82,7 @@ async def test_node_has_content_block_and_sst_event_labels( "block_index": 0, }, ) - node = await services.graph.get_node("s1::block::1::0") + node = await services.graph.get_node("s1::iteration::1::block::0") assert node is not None assert "ContentBlock" in node["labels"], ( f"ContentBlock label missing. Got: {node['labels']}" @@ -105,7 +105,7 @@ async def test_node_has_session_id_property( "block_index": 0, }, ) - node = await services.graph.get_node("s1::block::1::0") + node = await services.graph.get_node("s1::iteration::1::block::0") assert node is not None assert node.get("session_id") == "s1", ( f"session_id property missing or wrong. Got: {node!r}" @@ -125,7 +125,7 @@ async def test_node_has_block_index_and_started_at_properties( "block_index": 2, }, ) - node = await services.graph.get_node("s1::block::1::2") + node = await services.graph.get_node("s1::iteration::1::block::2") assert node is not None assert node.get("block_index") == 2, ( f"block_index property missing or wrong. Got: {node!r}" @@ -158,7 +158,7 @@ async def test_e07_edge_created_when_active_iteration_id_is_set( }, ) iteration_id = "s1::iteration::1" - block_id = "s1::block::1::0" + block_id = "s1::iteration::1::block::0" edge = await services.graph.get_edge(iteration_id, block_id) assert edge is not None, ( f"E07 HAS_PART edge from '{iteration_id}' to '{block_id}' must exist " @@ -228,7 +228,7 @@ async def test_content_block_end_sets_block_type( "block": {"type": "text"}, }, ) - node = await services.graph.get_node("s1::block::1::0") + node = await services.graph.get_node("s1::iteration::1::block::0") assert node is not None assert node.get("block_type") == "text", ( f"content_block:end must set block_type from block.type. Got: {node!r}" @@ -258,7 +258,7 @@ async def test_content_block_end_sets_ended_at( "block": {"type": "text"}, }, ) - node = await services.graph.get_node("s1::block::1::0") + node = await services.graph.get_node("s1::iteration::1::block::0") assert node is not None assert node.get("ended_at") == "2026-01-01T00:01:00Z", ( f"content_block:end must set ended_at. Got: {node!r}" @@ -302,9 +302,9 @@ async def test_tool_call_block_with_id_is_cached( ) assert ( services.data_layer_2.pending_tool_block_ids["tool-block-abc"] - == "s1::block::1::0" + == "s1::iteration::1::block::0" ), ( - "pending_tool_block_ids['tool-block-abc'] must map to the block node id 's1::block::1::0'" + "pending_tool_block_ids['tool-block-abc'] must map to the block node id 's1::iteration::1::block::0'" ) async def test_text_block_not_cached(self, services: HookStateService) -> None: @@ -448,7 +448,7 @@ async def test_content_block_start_creates_sourced_from_edge( "block_index": 0, }, ) - block_node_id = "s1::block::1::0" + block_node_id = "s1::iteration::1::block::0" data_layer_1_node_id = make_node_id("s1", "content_block:start", timestamp) edge = await services.graph.get_edge(block_node_id, data_layer_1_node_id) assert edge is not None, ( @@ -484,7 +484,7 @@ async def test_content_block_end_creates_sourced_from_edge( "block": {"type": "text"}, }, ) - block_node_id = "s1::block::1::0" + block_node_id = "s1::iteration::1::block::0" data_layer_1_node_id = make_node_id("s1", "content_block:end", end_timestamp) edge = await services.graph.get_edge(block_node_id, data_layer_1_node_id) assert edge is not None, ( diff --git a/tests/handlers/data_layer_2/test_iteration.py b/tests/handlers/data_layer_2/test_iteration.py index 77512c4f..1e511025 100644 --- a/tests/handlers/data_layer_2/test_iteration.py +++ b/tests/handlers/data_layer_2/test_iteration.py @@ -3,10 +3,15 @@ Covers: - handled_events == frozenset({'provider:request', 'llm:request', 'llm:response'}) - provider:request creates Iteration:SST_EVENT node keyed as - '{session_id}::iteration::{iteration_number}' with session_id, iteration_number, - and started_at; sets active_iteration_id cursor + '{session_id}::iteration::{iteration_number}' (no active orchestrator run) or + '{session_id}::orch_run::{execution_start_ts}::iteration::{iteration_number}' + (run-scoped when a run is active) with session_id, + iteration_number, and started_at; sets active_iteration_id cursor - E06: OrchestratorRun -[:HAS_PART {sst_semantic: 'CONTAINS'}]-> Iteration - created when execution_start_ts cursor is set; NOT created when None + created when execution_start_ts cursor is set (target is the run-scoped + iteration_id); NOT created when None +- two iterations sharing the same iteration_number under DIFFERENT + orchestrator runs get DISTINCT node_ids (no cross-run collision) - llm:request enriches active Iteration with provider, model, message_count, has_system; noop when active_iteration_id is None - llm:response enriches active Iteration with usage_input, usage_output, usage_cache_write; @@ -16,11 +21,15 @@ from __future__ import annotations +import logging + from context_intelligence_server.handlers.data_layer_2.iteration import IterationHandler +from context_intelligence_server.handlers.data_layer_2.orchestrator_run import ( + OrchestratorRunHandler, +) from context_intelligence_server.services import HookStateService from context_intelligence_server.utils import make_node_id - # --------------------------------------------------------------------------- # 1. TestIterationHandlerHandledEvents # --------------------------------------------------------------------------- @@ -172,10 +181,14 @@ class TestE06HasPartEdge: async def test_e06_has_part_edge_created_when_execution_start_ts_is_set( self, services: HookStateService ) -> None: - """E06 edge must be created when execution_start_ts cursor is set before provider:request.""" + """E06 edge must be created when execution_start_ts cursor is set before provider:request. + + The edge target is the run-scoped iteration_id, not the bare shape. + """ handler = IterationHandler(services) - # Simulate that execution:start previously fired and set the cursor + # Simulate that execution:start previously fired and set both cursors. services.data_layer_2.execution_start_ts = "2026-01-01T00:00:00Z" + services.data_layer_2.active_orch_run_id = "s1::orch_run::2026-01-01T00:00:00Z::1" await handler( "provider:request", @@ -184,8 +197,8 @@ async def test_e06_has_part_edge_created_when_execution_start_ts_is_set( "timestamp": "2026-01-01T00:00:01Z", }, ) - orch_run_id = "s1::orch_run::2026-01-01T00:00:00Z" - iteration_id = "s1::iteration::1" + orch_run_id = "s1::orch_run::2026-01-01T00:00:00Z::1" + iteration_id = "s1::orch_run::2026-01-01T00:00:00Z::1::iteration::1" edge = await services.graph.get_edge(orch_run_id, iteration_id) assert edge is not None, ( f"E06 HAS_PART edge from '{orch_run_id}' to '{iteration_id}' must exist " @@ -197,6 +210,9 @@ async def test_e06_has_part_edge_created_when_execution_start_ts_is_set( assert edge.get("sst_semantic") == "CONTAINS", ( f"E06 edge must have sst_semantic='CONTAINS'. Got: {edge.get('sst_semantic')}" ) + assert services.data_layer_2.active_iteration_id == iteration_id, ( + "active_iteration_id cursor must be set to the run-scoped iteration_id" + ) async def test_e06_not_created_when_execution_start_ts_is_none( self, services: HookStateService @@ -218,6 +234,118 @@ async def test_e06_not_created_when_execution_start_ts_is_none( f"Only SOURCED_FROM edge should exist when execution_start_ts is None. " f"Got {len(services.graph._edges)} edges: {list(services.graph._edges.keys())}" ) + # Falls back to the bare (pre-fix-shaped) id when no orchestrator run is active + assert services.data_layer_2.active_iteration_id == "s1::iteration::1" + + +# --------------------------------------------------------------------------- +# 3b. TestIterationRunScopingP21 +# --------------------------------------------------------------------------- + + +class TestIterationRunScopingP21: + """Iteration node_id is run-scoped and does not collide across runs.""" + + async def test_same_iteration_number_under_different_runs_gets_distinct_node_ids( + self, services: HookStateService + ) -> None: + """Two iterations sharing iteration_number=1 under DIFFERENT orchestrator runs + must produce DISTINCT node_ids and both nodes must independently exist. + + This reproduces the real-world collision: iteration_count is a per-session + (not per-run) counter, so after a drainer restart/replay recreates + DataLayer2State the counter can restart from zero and reproduce a prior + run's iteration_number under a NEW orchestrator run. Before run-scoping, + both runs' first iteration would MERGE onto the bare + 's1::iteration::1' node_id. After the fix, each run's iteration_number=1 is + prefixed with its own orch_run_id and the two nodes stay distinct. + """ + orch = OrchestratorRunHandler(services) + handler = IterationHandler(services) + + # --- Run 1: a real execution:start mints the run + its tiebreaker (seq 1). + await orch( + "execution:start", + {"session_id": "s1", "timestamp": "2026-01-01T00:00:00Z"}, + ) + await handler( + "provider:request", + {"session_id": "s1", "timestamp": "2026-01-01T00:00:01Z"}, + ) + run1_iteration_id = services.data_layer_2.active_iteration_id + await orch( + "orchestrator:complete", + {"session_id": "s1", "timestamp": "2026-01-01T00:00:02Z"}, + ) + + # --- Simulate a drainer restart: a fresh DataLayer2State restarts + # iteration_count from 0. The durable orch_run_seq is what a restored + # cursor carries; leave it as-is so a genuinely NEW run keeps advancing it. + services.data_layer_2.iteration_count = 0 + # A SECOND run that shares run 1's identical timestamp (coarse clock / + # replayed execution:start) -- the exact collision the tiebreaker exists for. + await orch( + "execution:start", + {"session_id": "s1", "timestamp": "2026-01-01T00:00:00Z"}, + ) + await handler( + "provider:request", + {"session_id": "s1", "timestamp": "2026-01-01T00:00:03Z"}, + ) + run2_iteration_id = services.data_layer_2.active_iteration_id + + # Same timestamp, different seq -> distinct ids (tiebreaker load-bearing). + assert run1_iteration_id == "s1::orch_run::2026-01-01T00:00:00Z::1::iteration::1" + assert run2_iteration_id == "s1::orch_run::2026-01-01T00:00:00Z::2::iteration::1" + assert run1_iteration_id != run2_iteration_id, ( + "Two runs sharing an identical execution_start_ts must still get " + "distinct Iteration ids (the run tiebreaker)." + ) + + node1 = await services.graph.get_node(run1_iteration_id) + node2 = await services.graph.get_node(run2_iteration_id) + assert node1 is not None, ( + f"Run 1's Iteration node '{run1_iteration_id}' must exist" + ) + assert node2 is not None, ( + f"Run 2's Iteration node '{run2_iteration_id}' must exist" + ) + assert node1.get("iteration_number") == 1 + assert node2.get("iteration_number") == 1 + assert node1.get("started_at") == "2026-01-01T00:00:01Z", ( + "Run 1's Iteration node must retain its own started_at, not run 2's " + "(proves the two nodes were never merged together)" + ) + assert node2.get("started_at") == "2026-01-01T00:00:03Z" + + # --- Regression guard: exactly ONE distinct HAS_PART parent per Iteration. + run1_orch_run_id = "s1::orch_run::2026-01-01T00:00:00Z::1" + run2_orch_run_id = "s1::orch_run::2026-01-01T00:00:00Z::2" + + def has_part_parents(iteration_id: str) -> list[str]: + """Distinct HAS_PART parent ids pointing at *iteration_id* in the fake graph.""" + return [ + src + for (src, dst), data in services.graph._edges.items() + if dst == iteration_id and data.get("type") == "HAS_PART" + ] + + run1_parents = has_part_parents(run1_iteration_id) + run2_parents = has_part_parents(run2_iteration_id) + + assert run1_parents == [run1_orch_run_id], ( + f"Run 1's Iteration node '{run1_iteration_id}' must have exactly ONE " + f"HAS_PART parent (its own OrchestratorRun). Got: {run1_parents!r}" + ) + assert run2_parents == [run2_orch_run_id], ( + f"Run 2's Iteration node '{run2_iteration_id}' must have exactly ONE " + f"HAS_PART parent (its own OrchestratorRun). Got: {run2_parents!r}" + ) + assert set(run1_parents).isdisjoint(run2_parents), ( + "No Iteration node may be shared (MERGEd) across the two " + "OrchestratorRuns -- each run's iterations must be distinct nodes " + "with distinct, non-overlapping HAS_PART parents." + ) # --------------------------------------------------------------------------- @@ -597,3 +725,175 @@ async def test_llm_response_creates_sourced_from_edge( assert edge.get("type") == "SOURCED_FROM", ( f"Edge type must be 'SOURCED_FROM'. Got: {edge.get('type')!r}" ) + + +# --------------------------------------------------------------------------- +# iteration_scope completeness +# +# The Iteration node is upsert_node'd from THREE sites: provider:request, +# llm:request, llm:response. Each site stamps the additive 'iteration_scope' +# ('run' | 'unscoped') property independently, sourced from the SAME cursor +# field (execution_start_ts) -- so a node can never be created/updated +# without a scope value, regardless of which of the three sites happens to +# be the one that actually writes it (e.g. a dead-lettered provider:request +# whose active_iteration_id mutation nonetheless survives to a later +# llm:request/llm:response). +# --------------------------------------------------------------------------- + + +class TestIterationScopeCompleteness: + """iteration_scope stamped at ALL THREE upsert_node call sites.""" + + async def test_provider_request_unscoped_with_no_execution_start( + self, services: HookStateService + ) -> None: + """provider:request with NO preceding execution:start -> + iteration_scope == 'unscoped'.""" + assert services.data_layer_2.execution_start_ts is None + handler = IterationHandler(services) + await handler( + "provider:request", + {"session_id": "s1", "timestamp": "2026-01-01T00:00:00Z"}, + ) + node = await services.graph.get_node("s1::iteration::1") + assert node is not None + assert node.get("iteration_scope") == "unscoped", ( + f"Expected iteration_scope='unscoped'. Got: {node!r}" + ) + + async def test_provider_request_run_scoped_with_execution_start( + self, services: HookStateService + ) -> None: + """A normal (execution:start already seen) run stamps 'run'.""" + services.data_layer_2.execution_start_ts = "2026-01-01T00:00:00Z" + services.data_layer_2.active_orch_run_id = "s1::orch_run::2026-01-01T00:00:00Z::1" + handler = IterationHandler(services) + await handler( + "provider:request", + {"session_id": "s1", "timestamp": "2026-01-01T00:00:01Z"}, + ) + node_id = "s1::orch_run::2026-01-01T00:00:00Z::1::iteration::1" + node = await services.graph.get_node(node_id) + assert node is not None + assert node.get("iteration_scope") == "run", ( + f"Expected iteration_scope='run'. Got: {node!r}" + ) + + async def test_unscoped_emitted_log_is_info_not_warning( + self, services: HookStateService, caplog + ) -> None: + """The log line for an unscoped iteration must be + INFO (or rate-limited), NEVER WARNING -- a loop-basic unscoped + session is a normal case, not an alert-worthy anomaly.""" + handler = IterationHandler(services) + with caplog.at_level( + logging.INFO, + logger="context_intelligence_server.handlers.data_layer_2.iteration", + ): + await handler( + "provider:request", + {"session_id": "s1", "timestamp": "2026-01-01T00:00:00Z"}, + ) + unscoped_records = [ + r for r in caplog.records if "unscoped_iteration_emitted" in r.message + ] + assert unscoped_records, ( + "expected an 'unscoped_iteration_emitted' log record to be emitted" + ) + assert all(r.levelno == logging.INFO for r in unscoped_records), ( + "unscoped_iteration_emitted must log at INFO, got levels: " + f"{[r.levelname for r in unscoped_records]}" + ) + assert not any(r.levelno >= logging.WARNING for r in caplog.records), ( + "no WARNING (or higher) should be emitted for a normal unscoped iteration" + ) + + async def test_llm_request_stamps_unscoped_independent_of_provider_request( + self, services: HookStateService + ) -> None: + """llm:request must stamp iteration_scope + on its OWN upsert_node call, even when the node was never created by + provider:request (e.g. a dead-lettered provider:request whose + active_iteration_id cursor mutation nonetheless survives). Simulated + here by setting the cursor directly, bypassing provider:request.""" + services.data_layer_2.active_iteration_id = "s1::iteration::99" + handler = IterationHandler(services) + await handler( + "llm:request", + { + "session_id": "s1", + "timestamp": "2026-01-01T00:00:00Z", + "provider": "anthropic", + "model": "claude", + }, + ) + node = await services.graph.get_node("s1::iteration::99") + assert node is not None + assert node.get("iteration_scope") == "unscoped", ( + f"llm:request must stamp iteration_scope='unscoped'. Got: {node!r}" + ) + + async def test_llm_request_stamps_run_scope_when_execution_start_ts_set( + self, services: HookStateService + ) -> None: + services.data_layer_2.execution_start_ts = "2026-01-01T00:00:00Z" + services.data_layer_2.active_orch_run_id = "s1::orch_run::2026-01-01T00:00:00Z::1" + node_id = "s1::orch_run::2026-01-01T00:00:00Z::1::iteration::1" + services.data_layer_2.active_iteration_id = node_id + handler = IterationHandler(services) + await handler( + "llm:request", + { + "session_id": "s1", + "timestamp": "2026-01-01T00:00:01Z", + "provider": "anthropic", + "model": "claude", + }, + ) + node = await services.graph.get_node(node_id) + assert node is not None + assert node.get("iteration_scope") == "run", ( + f"llm:request must stamp iteration_scope='run'. Got: {node!r}" + ) + + async def test_llm_response_stamps_unscoped_independent_of_provider_request( + self, services: HookStateService + ) -> None: + """Same completeness gap as llm:request, for the third site.""" + services.data_layer_2.active_iteration_id = "s1::iteration::99" + handler = IterationHandler(services) + await handler( + "llm:response", + { + "session_id": "s1", + "timestamp": "2026-01-01T00:00:00Z", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + node = await services.graph.get_node("s1::iteration::99") + assert node is not None + assert node.get("iteration_scope") == "unscoped", ( + f"llm:response must stamp iteration_scope='unscoped'. Got: {node!r}" + ) + + async def test_llm_response_stamps_run_scope_when_execution_start_ts_set( + self, services: HookStateService + ) -> None: + services.data_layer_2.execution_start_ts = "2026-01-01T00:00:00Z" + services.data_layer_2.active_orch_run_id = "s1::orch_run::2026-01-01T00:00:00Z::1" + node_id = "s1::orch_run::2026-01-01T00:00:00Z::1::iteration::1" + services.data_layer_2.active_iteration_id = node_id + handler = IterationHandler(services) + await handler( + "llm:response", + { + "session_id": "s1", + "timestamp": "2026-01-01T00:00:01Z", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + node = await services.graph.get_node(node_id) + assert node is not None + assert node.get("iteration_scope") == "run", ( + f"llm:response must stamp iteration_scope='run'. Got: {node!r}" + ) diff --git a/tests/handlers/data_layer_2/test_orchestrator_run.py b/tests/handlers/data_layer_2/test_orchestrator_run.py index 44bd54f1..34c3ecd6 100644 --- a/tests/handlers/data_layer_2/test_orchestrator_run.py +++ b/tests/handlers/data_layer_2/test_orchestrator_run.py @@ -66,7 +66,7 @@ async def test_node_created_with_correct_compound_key( "timestamp": "2026-01-01T00:00:00Z", }, ) - node_id = "s1::orch_run::2026-01-01T00:00:00Z" + node_id = "s1::orch_run::2026-01-01T00:00:00Z::1" node = await services.graph.get_node(node_id) assert node is not None, f"execution:start must create node at '{node_id}'" @@ -82,7 +82,7 @@ async def test_node_has_orchestrator_run_and_sst_event_labels( "timestamp": "2026-01-01T00:00:00Z", }, ) - node = await services.graph.get_node("s1::orch_run::2026-01-01T00:00:00Z") + node = await services.graph.get_node("s1::orch_run::2026-01-01T00:00:00Z::1") assert node is not None assert "OrchestratorRun" in node["labels"], ( f"OrchestratorRun label missing. Got: {node['labels']}" @@ -103,7 +103,7 @@ async def test_node_has_session_id_and_started_at( "timestamp": "2026-01-01T00:00:00Z", }, ) - node = await services.graph.get_node("s1::orch_run::2026-01-01T00:00:00Z") + node = await services.graph.get_node("s1::orch_run::2026-01-01T00:00:00Z::1") assert node is not None assert node.get("session_id") == "s1", ( f"session_id property missing or wrong. Got: {node!r}" @@ -124,7 +124,7 @@ async def test_e01_has_execution_edge_created( "timestamp": "2026-01-01T00:00:00Z", }, ) - orch_run_id = "s1::orch_run::2026-01-01T00:00:00Z" + orch_run_id = "s1::orch_run::2026-01-01T00:00:00Z::1" edge = await services.graph.get_edge("s1", orch_run_id) assert edge is not None, ( f"E01 HAS_EXECUTION edge from 's1' to '{orch_run_id}' must exist" @@ -191,7 +191,7 @@ async def test_execution_end_sets_ended_at( "status": "completed", }, ) - node = await services.graph.get_node("s1::orch_run::2026-01-01T00:00:00Z") + node = await services.graph.get_node("s1::orch_run::2026-01-01T00:00:00Z::1") assert node is not None assert node.get("ended_at") == "2026-01-01T00:01:00Z", ( f"ended_at must be set by execution:end. Got: {node!r}" @@ -212,7 +212,7 @@ async def test_execution_end_sets_status(self, services: HookStateService) -> No "status": "completed", }, ) - node = await services.graph.get_node("s1::orch_run::2026-01-01T00:00:00Z") + node = await services.graph.get_node("s1::orch_run::2026-01-01T00:00:00Z::1") assert node is not None assert node.get("status") == "completed", ( f"status must be set by execution:end. Got: {node!r}" @@ -236,7 +236,7 @@ async def test_execution_end_sets_response_properties( "response": "final answer text", }, ) - node = await services.graph.get_node("s1::orch_run::2026-01-01T00:00:00Z") + node = await services.graph.get_node("s1::orch_run::2026-01-01T00:00:00Z::1") assert node is not None assert node.get("response") == "final answer text", ( f"response must be set by execution:end. Got: {node!r}" @@ -269,7 +269,7 @@ async def test_orchestrator_complete_enriches_name_turn_count_completed_at( "turn_count": 3, }, ) - node = await services.graph.get_node("s1::orch_run::2026-01-01T00:00:00Z") + node = await services.graph.get_node("s1::orch_run::2026-01-01T00:00:00Z::1") assert node is not None assert node.get("orchestrator_name") == "my-orchestrator", ( f"orchestrator_name must be set by orchestrator:complete. Got: {node!r}" @@ -366,7 +366,7 @@ async def test_orchestrator_complete_sets_last_completed_orch_run_id( "turn_count": 1, }, ) - expected_run_id = "s1::orch_run::2026-01-01T00:00:00Z" + expected_run_id = "s1::orch_run::2026-01-01T00:00:00Z::1" assert services.data_layer_2.last_completed_orch_run_id == expected_run_id, ( f"last_completed_orch_run_id must be set to '{expected_run_id}'. " f"Got: {services.data_layer_2.last_completed_orch_run_id!r}" @@ -420,7 +420,7 @@ async def test_e14_prompt_triggers_orchestrator_run_edge_created( "timestamp": "2026-01-01T00:01:00Z", }, ) - orch_run_id = "s1::orch_run::2026-01-01T00:01:00Z" + orch_run_id = "s1::orch_run::2026-01-01T00:01:00Z::1" prompt_id = "s1::prompt::2026-01-01T00:00:00Z" edge = await services.graph.get_edge(prompt_id, orch_run_id) assert edge is not None, ( @@ -509,7 +509,7 @@ async def test_execution_start_creates_sourced_from_edge( "timestamp": "2026-01-01T00:00:00Z", }, ) - orch_run_id = "s1::orch_run::2026-01-01T00:00:00Z" + orch_run_id = "s1::orch_run::2026-01-01T00:00:00Z::1" data_layer_1_node_id = make_node_id( "s1", "execution:start", "2026-01-01T00:00:00Z" ) @@ -539,7 +539,7 @@ async def test_execution_end_creates_sourced_from_edge( "status": "completed", }, ) - orch_run_id = "s1::orch_run::2026-01-01T00:00:00Z" + orch_run_id = "s1::orch_run::2026-01-01T00:00:00Z::1" data_layer_1_node_id = make_node_id( "s1", "execution:end", "2026-01-01T00:01:00Z" ) @@ -570,7 +570,7 @@ async def test_orchestrator_complete_creates_sourced_from_edge( "turn_count": 3, }, ) - orch_run_id = "s1::orch_run::2026-01-01T00:00:00Z" + orch_run_id = "s1::orch_run::2026-01-01T00:00:00Z::1" data_layer_1_node_id = make_node_id( "s1", "orchestrator:complete", "2026-01-01T00:02:00Z" ) @@ -581,3 +581,58 @@ async def test_orchestrator_complete_creates_sourced_from_edge( assert edge.get("type") == "SOURCED_FROM", ( f"Edge must have type='SOURCED_FROM'. Got: {edge.get('type')}" ) + + +# --------------------------------------------------------------------------- +# Run-id tiebreaker survives a worker rebuild (no seq-less collision) +# --------------------------------------------------------------------------- + + +class TestRunIdSurvivesRebuild: + """Two same-timestamp runs must get distinct ids, and an enrichment event + after a cursor rebuild must re-derive the run id WITH the seq -- never the + seq-less form that both collides same-timestamp runs and targets a node + execution:start never created.""" + + async def test_same_timestamp_runs_get_distinct_ids( + self, services: HookStateService + ) -> None: + handler = OrchestratorRunHandler(services) + ts = "2026-01-01T00:00:00Z" + await handler("execution:start", {"session_id": "s1", "timestamp": ts}) + await handler("orchestrator:complete", {"session_id": "s1", "timestamp": ts}) + await handler("execution:start", {"session_id": "s1", "timestamp": ts}) + + first = await services.graph.get_node(f"s1::orch_run::{ts}::1") + second = await services.graph.get_node(f"s1::orch_run::{ts}::2") + assert first is not None + assert second is not None # distinct id despite identical timestamp + + async def test_enrichment_after_rebuild_redereives_id_with_seq( + self, services: HookStateService + ) -> None: + handler = OrchestratorRunHandler(services) + ts = "2026-01-01T00:00:00Z" + await handler("execution:start", {"session_id": "s1", "timestamp": ts}) + + # Simulate a rebuilt worker restored from a partial/legacy cursor: the + # active_orch_run_id is missing, but execution_start_ts and orch_run_seq + # survive (they are what the cursor carries). + services.data_layer_2.active_orch_run_id = None + assert services.data_layer_2.execution_start_ts == ts + assert services.data_layer_2.orch_run_seq == 1 + + await handler( + "execution:end", + {"session_id": "s1", "timestamp": ts, "status": "ok"}, + ) + + # The enrichment must land on the seq-qualified node execution:start + # created, not a seq-less orphan. + enriched = await services.graph.get_node(f"s1::orch_run::{ts}::1") + assert enriched is not None + assert enriched.get("ended_at") == ts # RED before fix: enrichment + assert enriched.get("status") == "ok" # went to the seq-less node + # And the seq-less collision id must NOT have been created. + orphan = await services.graph.get_node(f"s1::orch_run::{ts}") + assert orphan is None diff --git a/tests/integration/test_data_layer_3_delegation_skill.py b/tests/integration/test_data_layer_3_delegation_skill.py index ad56d110..14883482 100644 --- a/tests/integration/test_data_layer_3_delegation_skill.py +++ b/tests/integration/test_data_layer_3_delegation_skill.py @@ -236,8 +236,9 @@ async def test_skill_loaded_during_iteration_creates_e05(self) -> None: handlers, ) - # iteration_id = SESSION_ID::iteration::1 (first iteration) - iteration_id = f"{SESSION_ID}::iteration::1" + # A run is active (execution:start fired), so the iteration id is + # run-scoped and carries the run tiebreaker (seq 1 for the first run). + iteration_id = f"{SESSION_ID}::orch_run::{T1}::1::iteration::1" skill_load_id = f"{SESSION_ID}::skill::{SKILL_NAME}::{T2}" # Verify active_iteration_id was set by IterationHandler diff --git a/tests/integration/test_turn_chain.py b/tests/integration/test_turn_chain.py index 06ad16f8..04b2d4fd 100644 --- a/tests/integration/test_turn_chain.py +++ b/tests/integration/test_turn_chain.py @@ -48,9 +48,9 @@ # Computed node IDs matching the handlers' key conventions PROMPT_1_ID = f"{SESSION_ID}::prompt::{T1}" -RUN_1_ID = f"{SESSION_ID}::orch_run::{T2}" +RUN_1_ID = f"{SESSION_ID}::orch_run::{T2}::1" PROMPT_2_ID = f"{SESSION_ID}::prompt::{T8}" -RUN_2_ID = f"{SESSION_ID}::orch_run::{T9}" +RUN_2_ID = f"{SESSION_ID}::orch_run::{T9}::2" PROMPT_3_ID = f"{SESSION_ID}::prompt::{T11}" diff --git a/tests/test_boot_safety.py b/tests/test_boot_safety.py index 488bfd64..fd22e057 100644 --- a/tests/test_boot_safety.py +++ b/tests/test_boot_safety.py @@ -740,7 +740,7 @@ async def _drain_to_dry(worker: SessionWorker) -> None: drain loop would otherwise never finish).""" batch = await qm.read_batch(worker.session_id, max_items=10) if batch.records: - await qm.commit(worker.session_id, batch.end_offset) + await qm.commit(worker.session_id, batch.end_offset, None) reg._deregister(worker.session_id) def _get_or_create( @@ -908,7 +908,9 @@ async def test_dry_exit_fires_for_recovered_drainer_over_drained_log( reg._queue_manager = qm line = _line() await qm.append("sess-x", line) - await qm.commit("sess-x", len(line) + 1) # fully drained, no terminal record + await qm.commit( + "sess-x", len(line) + 1, None + ) # fully drained, no terminal record worker = _make_worker("sess-x", live_event_seen=False) # recovered=True shape reg._register_for_test(worker) @@ -927,7 +929,7 @@ async def test_dry_exit_negative_control_live_created_worker_never_exits( reg._queue_manager = qm line = _line() await qm.append("sess-live", line) - await qm.commit("sess-live", len(line) + 1) + await qm.commit("sess-live", len(line) + 1, None) worker = _make_worker("sess-live", live_event_seen=True) # live path (default) reg._register_for_test(worker) @@ -1162,7 +1164,7 @@ async def test_boot_reclaim_auto_reclaims_drained_log_under_shipped_defaults( qm = FileSystemQueueManager(queues_dir=tmp_path) line = _line() await qm.append("drained-key", line) - await qm.commit("drained-key", len(line)) + await qm.commit("drained-key", len(line), None) monkeypatch.setattr(main_module.registry, "_queue_manager", qm) monkeypatch.setattr(main_module._settings, "reclaim_enabled", False) @@ -1258,7 +1260,7 @@ async def test_boot_reclaim_drained_log_with_live_worker_is_skipped( qm = FileSystemQueueManager(queues_dir=tmp_path) line = _line() await qm.append("live-drained-key", line) - await qm.commit("live-drained-key", len(line)) + await qm.commit("live-drained-key", len(line), None) reg = SessionRegistry() reg._queue_manager = qm diff --git a/tests/test_drain_supervision.py b/tests/test_drain_supervision.py index 3c439972..f75bf56d 100644 --- a/tests/test_drain_supervision.py +++ b/tests/test_drain_supervision.py @@ -1096,3 +1096,199 @@ async def test_poison_line_before_terminal_is_dead_lettered_and_session_finalize dead = await qm.read_dead_letters(sid) assert len(dead) == 1 assert not qm._log_path(sid).exists(), "delete_drained must have run" + + +class TestRetryDoesNotDuplicateIteration: + """In-place retry rolls the cross-handler counter back to its pre-attempt + state, so a replayed batch reproduces the SAME node ids instead of + duplicating them.""" + + async def test_in_place_retry_yields_exactly_one_iteration(self) -> None: + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d1-retry-dedup" + # Fail the first multi-record flush exactly once, then succeed (in-place + # retry, no respawn). flushed is a SET: a duplicate id shows as an extra + # member, and a non-rolled-back counter yields iter::3/iter::4 too. + flush_calls = {"n": 0} + + def _fail_first_multi(buf: set[str]) -> bool: + if len(buf) > 1: + flush_calls["n"] += 1 + return flush_calls["n"] == 1 + return False + + graph = _FlakyGraph(fail_when=_fail_first_multi) + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + async def _advance_and_buffer( + w: SessionWorker, event: str, data: object, handlers: object + ) -> None: + # Mimic a real enricher: bump the iteration counter, then buffer the + # id it produces. Without the pre-retry rollback the counter keeps + # climbing across attempts and the replay emits a DIFFERENT id. + w.services.data_layer_2.iteration_count += 1 + w.services.graph.buffer.add( + f"iter::{w.services.data_layer_2.iteration_count}" + ) + + with patch( + "context_intelligence_server.registry.process_event", + side_effect=_advance_and_buffer, + ): + await qm.append(sid, _line("e1", "/ws", {"session_id": sid})) + await qm.append(sid, _line("e2", "/ws", {"session_id": sid})) + reg.start_drain(worker) + await _drain_until_idle(reg, qm, worker, sid) + await _cancel_and_await(worker.task) + + # Two records => ids iter::1, iter::2 written exactly once each. A missing + # rollback would additionally leave iter::3/iter::4 from the failed attempt. + assert graph.flushed == {"iter::1", "iter::2"} + assert (await qm.read_batch(sid, 10)).lines == [] + + +class TestCursorRestoredOnWorkerRebuild: + """A rebuilt worker restores the durable cursor the last commit persisted, + so cross-handler counters resume instead of restarting from zero (which + would remint node ids and duplicate them).""" + + async def test_orch_run_seq_is_restored_before_a_rebuilt_worker_processes( + self, + ) -> None: + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d1-cursor-rebuild" + + # A prior worker committed a cursor carrying orch_run_seq=5, leaving one + # event still undrained (offset 0) for the rebuilt worker to pick up. + await qm.append(sid, _line("e1", "/ws", {"session_id": sid})) + await qm.commit(sid, 0, {"dl2": {"orch_run_seq": 5}, "dl3": {}}) + assert (await qm.read_cursor(sid))["dl2"]["orch_run_seq"] == 5 + + # A REBUILT worker: a brand-new SessionWorker whose DataLayer2State starts + # at orch_run_seq=0. drain_worker must restore the persisted cursor before + # processing, so the counter is 5 (not 0) by the time an event is handled + # -- without the restore the next run would remint the already-used seq. + graph = _FlakyGraph() + worker = _make_worker(sid, graph) + assert worker.services.data_layer_2.orch_run_seq == 0 + reg._register_for_test(worker) + + captured: dict[str, int] = {} + + async def _capture(w: SessionWorker, event, data, handlers) -> None: + captured["seq_at_process"] = w.services.data_layer_2.orch_run_seq + w.services.graph.buffer.add("e1") + + with patch( + "context_intelligence_server.registry.process_event", + side_effect=_capture, + ): + reg.start_drain(worker) + await _drain_until_idle(reg, qm, worker, sid) + await _cancel_and_await(worker.task) + + assert captured.get("seq_at_process") == 5, ( + "drain_worker must restore the durable cursor (orch_run_seq=5) before " + "processing any event; a rebuilt worker that starts from 0 would " + "remint already-used run ids" + ) + + +class TestExhaustedBatchCursorRollback: + """The dead-letter/isolation redrive must restore the pre-batch cursor, + exactly like the in-place-retry branch -- otherwise per-record commits + snapshot already-advanced counters and double-count on the redrive.""" + + async def test_exhausted_batch_restores_pre_batch_cursor(self) -> None: + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d1-exhausted-dedup" + + # Default _FlakyGraph rejects every multi-record flush, so the 2-record + # batch exhausts its retry budget and falls into _handle_exhausted_batch, + # which then isolates each single line (single-line flushes succeed). + graph = _FlakyGraph() + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + async def _advance_and_buffer( + w: SessionWorker, event: str, data: object, handlers: object + ) -> None: + w.services.data_layer_2.iteration_count += 1 + w.services.graph.buffer.add( + f"iter::{w.services.data_layer_2.iteration_count}" + ) + + with patch( + "context_intelligence_server.registry.process_event", + side_effect=_advance_and_buffer, + ): + await qm.append(sid, _line("e1", "/ws", {"session_id": sid})) + await qm.append(sid, _line("e2", "/ws", {"session_id": sid})) + reg.start_drain(worker) + await _drain_until_idle(reg, qm, worker, sid) + await _cancel_and_await(worker.task) + + # The isolation redrive must reproduce iter::1/iter::2. Without the + # pre-batch restore the failed batch's advanced counters would leak + # through and the isolated lines would flush iter::3/iter::4. + assert graph.flushed == {"iter::1", "iter::2"} + assert (await qm.read_batch(sid, 10)).lines == [] + + +class TestReadCursorFailureQuarantines: + """A corrupt/unparseable cursor read degrades to None instead of killing + the drain worker -- the committed offset is read separately, so the loss + costs at most a bounded replay, never a crash-loop.""" + + async def test_read_cursor_valueerror_is_quarantined_not_fatal( + self, caplog + ) -> None: + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d1-cursor-corrupt" + + await qm.append(sid, _line("e1", "/ws", {"session_id": sid})) + + graph = _FlakyGraph() + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + # Simulate a corrupt/legacy-unparseable cursor record: read_cursor raises + # ValueError. Before the fix this propagates and kills the drainer. + original_read_cursor = qm.read_cursor + calls = {"n": 0} + + async def _raising_read_cursor(session_id: str): + calls["n"] += 1 + raise ValueError("unparseable offset document") + + async def _capture(w, event, data, handlers) -> None: + w.services.graph.buffer.add("e1") + + with ( + patch.object(qm, "read_cursor", side_effect=_raising_read_cursor), + patch( + "context_intelligence_server.registry.process_event", + side_effect=_capture, + ), + caplog.at_level(logging.WARNING), + ): + reg.start_drain(worker) + await _drain_until_idle(reg, qm, worker, sid) + await _cancel_and_await(worker.task) + + # The worker degraded (logged a warning) and still drained the event, + # rather than dying with a drain_worker_died ERROR. + assert calls["n"] >= 1 + assert graph.flushed == {"e1"} + assert any( + "cursor_read_failed" in r.getMessage() for r in caplog.records + ), "read_cursor failure must log a degrade warning, not crash the worker" + assert not any( + "drain_worker_died" in r.getMessage() for r in caplog.records + ), "a corrupt cursor read must never kill the drainer" + _ = original_read_cursor diff --git a/tests/test_durable_append_framing.py b/tests/test_durable_append_framing.py index 2725bdd8..2e41bb88 100644 --- a/tests/test_durable_append_framing.py +++ b/tests/test_durable_append_framing.py @@ -731,7 +731,7 @@ async def test_delete_drained_retains_a_log_with_uncommitted_bytes( second = _event_bytes("second") await qm.append(key, first) batch = await qm.read_batch(key, max_items=10) - await qm.commit(key, batch.end_offset) # commits only `first` + await qm.commit(key, batch.end_offset, None) # commits only `first` await qm.append(key, second) # uncommitted tail with caplog.at_level( @@ -761,7 +761,7 @@ async def test_guard_map_is_released_on_delete_drained_and_identity_checked( for k in keys: await qm.append(k, _event_bytes("ev")) batch = await qm.read_batch(k, max_items=10) - await qm.commit(k, batch.end_offset) + await qm.commit(k, batch.end_offset, None) assert set(qm._guards.keys()) == set(keys) for k in keys: @@ -773,7 +773,7 @@ async def test_guard_map_is_released_on_delete_drained_and_identity_checked( key = "aba-key" await qm.append(key, _event_bytes("ev")) batch = await qm.read_batch(key, max_items=10) - await qm.commit(key, batch.end_offset) + await qm.commit(key, batch.end_offset, None) g_orig = qm._guards[key] real_stat = Path.stat @@ -856,7 +856,7 @@ async def test_guard_survives_a_delete_that_races_a_parked_appender( seed = _event_bytes("seed") await qm.append(key, seed) seed_batch = await qm.read_batch(key, max_items=10) - await qm.commit(key, seed_batch.end_offset) # fully drained: size == committed + await qm.commit(key, seed_batch.end_offset, None) # fully drained: size == committed guard = qm._guards[key] real_stat = Path.stat diff --git a/tests/test_ingest_cursor_hardening.py b/tests/test_ingest_cursor_hardening.py new file mode 100644 index 00000000..7456ad82 --- /dev/null +++ b/tests/test_ingest_cursor_hardening.py @@ -0,0 +1,804 @@ +"""Regression guards for the durable-cursor hardening round. + +Each test drives the real code path of a distinct, execution-proven defect: +cross-handler run-id desync, cursor loss on offset-deletion paths, the +commit/compaction lock race, a corrupt-offset crash-loop, snapshot aliasing, +respawn double-dead-letter, and the session-node lost on isolation. +""" + +from __future__ import annotations + +import asyncio +import json +import logging +import sys +import threading +import time +from pathlib import Path +from typing import Any, cast + +import pytest +from context_intelligence_server.handlers.data_layer_2.iteration import ( + IterationHandler, +) +from context_intelligence_server.handlers.data_layer_2.orchestrator_run import ( + OrchestratorRunHandler, +) +from context_intelligence_server.queue_manager import FileSystemQueueManager +from context_intelligence_server.registry import ( + SessionRegistry, + SessionWorker, + _SessionQuarantined, +) +from context_intelligence_server.services import HookStateService + + +def _line(event: str, data: dict[str, Any]) -> bytes: + return json.dumps({"event": event, "workspace": "/ws", "data": data}).encode( + "utf-8" + ) + + +# --------------------------------------------------------------------------- +# Fix 1: cross-handler run-id consistency (E06 HAS_PART survives a rebuild) +# --------------------------------------------------------------------------- + + +async def test_e06_has_part_edge_survives_worker_rebuild( + services: HookStateService, +) -> None: + """After a rebuild leaves a partial cursor (active_orch_run_id lost, but + execution_start_ts + orch_run_seq preserved), the Iteration handler must + re-derive the run's real id and still wire the E06 HAS_PART edge, instead + of falling through to the run-less shape and dropping the edge.""" + orch = OrchestratorRunHandler(services) + it = IterationHandler(services) + ts = "2026-01-01T00:00:00Z" + + await orch("execution:start", {"session_id": "s1", "timestamp": ts}) + run_id = f"s1::orch_run::{ts}::1" + + # Simulate the rebuild: the run is still active (execution_start_ts + seq + # survived in the cursor) but the full id cursor was not restored. + services.data_layer_2.active_orch_run_id = None + + await it("provider:request", {"session_id": "s1", "timestamp": ts}) + + iteration_id = f"{run_id}::iteration::1" + edge = await services.graph.get_edge(run_id, iteration_id) + assert edge is not None, ( + "E06 HAS_PART edge dropped: the Iteration handler did not re-derive the " + "run id after a partial-cursor rebuild" + ) + assert edge["type"] == "HAS_PART" + + +async def test_content_block_inherits_run_scope_after_rebuild( + services: HookStateService, +) -> None: + """A ContentBlock whose active_iteration_id cursor was lost must reconstruct + the run-scoped iteration key, not the run-less ``::iteration::0`` fallback.""" + orch = OrchestratorRunHandler(services) + it = IterationHandler(services) + ts = "2026-01-01T00:00:00Z" + + await orch("execution:start", {"session_id": "s1", "timestamp": ts}) + await it("provider:request", {"session_id": "s1", "timestamp": ts}) + run_id = f"s1::orch_run::{ts}::1" + expected_iteration_id = f"{run_id}::iteration::1" + + resolved = services.data_layer_2.resolve_active_iteration_id("s1") + assert resolved == expected_iteration_id + + # Drop the iteration cursor (partial rebuild) -- resolution must reconstruct + # the same run-scoped key from run id + iteration_count. + services.data_layer_2.active_iteration_id = None + reconstructed = services.data_layer_2.resolve_active_iteration_id("s1") + assert reconstructed == expected_iteration_id + + +# --------------------------------------------------------------------------- +# Fix 2: cursor survives the offset-deletion paths +# --------------------------------------------------------------------------- + + +async def test_reset_offset_preserves_cursor(tmp_path) -> None: + """A RESET_OFFSET boot-reclaim (bad offset, re-drainable log) must reset the + committed offset to 0 while PRESERVING the committed cursor.""" + qm = FileSystemQueueManager(queues_dir=tmp_path) + sid = "s-reset" + cursor = {"dl2": {"orch_run_seq": 4, "active_orch_run_id": "s-reset::r"}, "dl3": {}} + + await qm.append(sid, b"one-real-line") + # A valid record whose offset points PAST end-of-log -> classify RESET_OFFSET + # while the cursor stays perfectly readable. + log_size = (tmp_path / f"{sid}.log").stat().st_size + await qm.commit(sid, log_size + 10_000, cursor) + + classification = await qm.classify_session(sid, head_is_resumable=lambda _b: True) + assert classification.verdict.value == "reset_offset" + reclaimed = await qm.reclaim(classification, is_owned=lambda: False) + assert reclaimed is True + + assert qm._read_committed_offset(sid) == 0 # re-drains from the start + assert await qm.read_cursor(sid) == cursor # cursor preserved across the reset + + +async def test_delete_drained_is_terminal_and_clears_offset(tmp_path) -> None: + """delete_drained is terminal cleanup: it removes BOTH the log and the + offset even when the cursor is non-empty. A finalized session's cursor is + intentionally dropped -- an orch_run_id is scoped by execution_start_ts, so + a genuinely-new post-finalize run gets a distinct id regardless of the seq + counter, and keeping the offset would leak a file per finalized session.""" + qm = FileSystemQueueManager(queues_dir=tmp_path) + sid = "s-final" + line = b"a\n" + await qm.append(sid, line) + await qm.commit(sid, len(line), {"dl2": {"orch_run_seq": 9}, "dl3": {}}) + + assert await qm.delete_drained(sid) is True + + assert not (tmp_path / f"{sid}.log").exists() + assert not (tmp_path / f"{sid}.offset").exists() + assert await qm.read_cursor(sid) is None + + +# --------------------------------------------------------------------------- +# Fix 3: commit serializes under the same per-key lock as compaction +# --------------------------------------------------------------------------- + + +async def test_commit_holds_file_lock_during_write(tmp_path, monkeypatch) -> None: + """commit() must write the offset record under the per-key file_lock -- the + same lock compaction takes -- so a commit concurrent with a compaction can + never be silently erased. Verified by observing the lock is held at the + moment commit writes the record.""" + qm = FileSystemQueueManager(queues_dir=tmp_path) + sid = "s-lock" + await qm.append(sid, b"a") + + observed: dict[str, bool] = {} + real_write = qm._write_offset_record + + def _spy_write(session_id, offset, cursor): + guard = qm._guards.get(session_id) + observed["guard_exists"] = guard is not None + observed["locked"] = bool(guard is not None and guard.file_lock.locked()) + return real_write(session_id, offset, cursor) + + monkeypatch.setattr(qm, "_write_offset_record", _spy_write) + await qm.commit(sid, 2, {"dl2": {}, "dl3": {}}) + + assert observed.get("guard_exists") is True + assert observed.get("locked") is True, ( + "commit wrote the offset record without holding the per-key file_lock -- " + "a concurrent compaction could erase it" + ) + + +# --------------------------------------------------------------------------- +# Fix 4: a corrupt .offset quarantines the worker, never crash-loops it +# --------------------------------------------------------------------------- + + +class _InertGraph: + workspace = "/ws" + created_by: str | None = None + + async def flush(self) -> None: # pragma: no cover - never reached + return None + + def discard_buffer(self) -> None: # pragma: no cover + return None + + async def close(self) -> None: + return None + + +def _make_worker(sid: str, graph: Any) -> SessionWorker: + worker = SessionWorker( + session_id=sid, workspace="/ws", services=HookStateService(workspace="/ws") + ) + worker.services.graph = graph # type: ignore[assignment] + return worker + + +async def test_corrupt_offset_quarantines_worker_without_dying(caplog) -> None: + reg = SessionRegistry() + qm = cast(FileSystemQueueManager, reg.queue_manager) + sid = "s-corrupt" + await qm.append(sid, _line("e1", {"session_id": sid})) + # A genuinely corrupt offset document on disk: read_batch reads the offset + # first and raises ValueError before the guarded cursor read is reached. + qm._offset_path(sid).write_text('{"v":1,"offset":"NOT-AN-INT"}', encoding="utf-8") + + worker = _make_worker(sid, _InertGraph()) + reg._register_for_test(worker) + + with caplog.at_level(logging.ERROR, logger="context_intelligence_server"): + # Must RETURN (quarantine), never raise out of the drain loop. + await reg.drain_worker(worker, flush_timeout=0.05) + + assert any("drain_worker_quarantined" in r.getMessage() for r in caplog.records), ( + "a corrupt .offset must quarantine the worker with a loud log" + ) + assert not any("drain_worker_died" in r.getMessage() for r in caplog.records), ( + "a corrupt .offset must never crash-loop the worker" + ) + assert sid not in reg._workers # deregistered cleanly + + +# --------------------------------------------------------------------------- +# Fix 5: restore_cursor deep-copies mutable fields (no snapshot aliasing) +# --------------------------------------------------------------------------- + + +async def test_restore_cursor_does_not_alias_snapshot_mutables( + services: HookStateService, +) -> None: + """restore_cursor must deep-copy mutable fields so a later in-place mutation + of the LIVE state cannot corrupt the snapshot replayed on the next retry.""" + snapshot = { + "dl2": {"pending_tool_block_ids": {"b1": "n1"}}, + "dl3": {"active_recipe_run_stack": ["r1"]}, + } + services.restore_cursor(snapshot) + + # A handler mutates the live containers in place on a later attempt. + services.data_layer_2.pending_tool_block_ids["b2"] = "n2" + services.data_layer_3.active_recipe_run_stack.append("r2") + + # The snapshot (the pre_batch_cursor baseline) must be untouched. + assert snapshot["dl2"]["pending_tool_block_ids"] == {"b1": "n1"} + assert snapshot["dl3"]["active_recipe_run_stack"] == ["r1"] + + # And a second restore reproduces the identical baseline, not the mutation. + services.restore_cursor(snapshot) + assert services.data_layer_2.pending_tool_block_ids == {"b1": "n1"} + assert services.data_layer_3.active_recipe_run_stack == ["r1"] + + +# --------------------------------------------------------------------------- +# Fix 6: a respawn does not re-dead-letter an already-dead leading record +# --------------------------------------------------------------------------- + + +class _PoisonGraph: + """Flush always succeeds; the poison is injected in the handler.""" + + workspace = "/ws" + created_by: str | None = None + + def __init__(self) -> None: + self.flushed = 0 + + async def flush(self) -> None: + self.flushed += 1 + + def discard_buffer(self) -> None: + return None + + async def close(self) -> None: + return None + + +async def test_respawn_does_not_double_dead_letter(monkeypatch) -> None: + reg = SessionRegistry() + qm = cast(FileSystemQueueManager, reg.queue_manager) + reg._max_delivery_attempts = 1 + sid = "s-respawn" + + e1 = _line("e1", {"session_id": sid}) # poison + e2 = _line("e2", {"session_id": sid}) # good + await qm.append(sid, e1) + await qm.append(sid, e2) + + # Reconstruct the exact post-crash state: e1 was dead-lettered but the + # offset was never advanced past it (the crash landed in the + # dead_letter -> commit gap). dead_letter() marks the session + # dead-unreconciled in-memory, so the NEXT read_batch reconciles past e1 + # once (the internalized dirty-flag path, no registry-level call). + await qm.dead_letter(sid, e1, "poison") + assert qm._read_committed_offset(sid) == 0 + + async def _handler(w, event, data, handlers): + if event == "e1": + raise RuntimeError("poison record") + + worker = _make_worker(sid, _PoisonGraph()) + reg._register_for_test(worker) + + with monkeypatch.context() as m: + m.setattr( + "context_intelligence_server.registry.process_event", + _handler, + ) + reg.start_drain(worker) + for _ in range(200): + await asyncio.sleep(0.01) + if (await qm.read_batch(sid, 10)).lines == []: + break + task = worker.task + if task is not None and not task.done(): + task.cancel() + try: + if task is not None: + await task + except asyncio.CancelledError: + pass + + dead = await qm.read_dead_letters(sid) + payloads = [d.get("payload") for d in dead] + assert payloads.count(e1.decode("utf-8")) == 1, ( + f"e1 was dead-lettered {payloads.count(e1.decode('utf-8'))} times; a " + "respawn must reconcile the already-dead leading record, not re-dead-letter it" + ) + + +# --------------------------------------------------------------------------- +# Fix 7: the Session node survives an exhausted-batch isolation +# --------------------------------------------------------------------------- + + +class _BufferedGraph: + """Models Neo4jGraphStore's buffer/flush split: upsert_node BUFFERS, flush + promotes the buffer to the store, discard_buffer DROPS the buffer, and + get_node reads buffer-first. The in-memory GraphState has no such split, so + only a faithful buffered fake can exercise the isolation buffer-discard.""" + + workspace = "/ws" + created_by: str | None = None + + def __init__(self) -> None: + self._store: dict[str, dict[str, Any]] = {} + self._buffer: dict[str, dict[str, Any]] = {} + + async def upsert_node(self, node_id: str, data: dict[str, Any]) -> None: + node = self._buffer.setdefault(node_id, {}) + if "labels" in data: + node["labels"] = sorted(set(node.get("labels", [])) | set(data["labels"])) + for k, v in data.items(): + if k != "labels": + node[k] = v + + async def get_node(self, node_id: str) -> dict[str, Any] | None: + if node_id in self._buffer: + return dict(self._buffer[node_id]) + if node_id in self._store: + return dict(self._store[node_id]) + return None + + async def upsert_edge(self, *args: Any) -> None: + return None + + async def flush(self) -> None: + for nid, data in self._buffer.items(): + self._store.setdefault(nid, {}).update(data) + self._buffer.clear() + + def discard_buffer(self) -> None: + self._buffer.clear() + + async def close(self) -> None: + return None + + +async def test_session_node_survives_isolation(monkeypatch) -> None: + """When a failed batch is isolated line-by-line, the buffer discard must + also invalidate the seen-session cache, so the re-dispatched session:start + re-issues the Session node (status/started_at) instead of early-returning.""" + reg = SessionRegistry() + sid = "s-node" + ts = "2026-02-02T00:00:00Z" + + worker = _make_worker(sid, _BufferedGraph()) + reg._register_for_test(worker) + + async def _handler(w, event, data, handlers): + if event == "session:start": + await w.services.ensure_session_node(data["session_id"], data) + elif event == "poison": + raise RuntimeError("poison") + + # Prime the exact pre-isolation condition: the failed batch already + # dispatched session:start, buffering the node and caching the id (the node + # is NOT yet flushed to the store). + await worker.services.ensure_session_node(sid, {"session_id": sid, "timestamp": ts}) + assert sid in worker.services._seen_sessions + pre_batch_cursor = worker.services.snapshot_cursor() + + from context_intelligence_server.queue_manager.protocol import Batch, Record + + batch = Batch( + session_id=sid, + records=[ + Record(_line("session:start", {"session_id": sid, "timestamp": ts}), 0, 40), + Record(_line("poison", {"session_id": sid}), 40, 80), + ], + start_offset=0, + end_offset=80, + ) + + with monkeypatch.context() as m: + m.setattr("context_intelligence_server.registry.process_event", _handler) + await reg._handle_exhausted_batch( + worker, batch, handlers=None, pre_batch_cursor=pre_batch_cursor + ) + + node = await worker.services.graph.get_node(sid) + assert node is not None, "Session node was lost on isolation" + assert node.get("status") == "running" + assert node.get("started_at") == ts + + +# --------------------------------------------------------------------------- +# Fix 8: every path that moves a session's .offset serialises on its file_lock +# --------------------------------------------------------------------------- + + +async def _commit_racing_reconcile(queues_dir: Path) -> tuple[int, dict | None]: + """Run one commit concurrent with the read_batch dead-letter reconcile. + + Both move the same session's ``.offset``: the commit writes it, and the + reconcile read-modify-writes it. Returns the record left on disk. + """ + qm = FileSystemQueueManager(queues_dir=queues_dir) + sid = "s-race" + poison = b"poison-line" + + await qm.append(sid, poison) + for i in range(20): + await qm.append(sid, b'{"n":%d}' % i) + # Arms the reconcile: the next read_batch walks the dead payloads and + # advances the offset past the already-dead leading line. + await qm.dead_letter(sid, poison, "boom") + + await asyncio.gather( + qm.commit(sid, 999_999, {"dl2": {"orch_run_seq": 3}, "dl3": {}}), + qm.read_batch(sid, max_items=10), + ) + return qm._read_offset_record(sid) + + +async def test_concurrent_offset_writers_never_clobber_or_tear(tmp_path) -> None: + """A commit racing the reconcile must survive intact. + + The reconcile reaches the offset through read_batch, on a different thread + than the commit. With both writes funnelled through the one locked writer, + the commit is never rolled back to the value the reconcile read before it + and the record is never half-written. Repeated because a race that survives + one interleaving proves nothing. + """ + for attempt in range(25): + offset, cursor = await _commit_racing_reconcile(tmp_path / f"run{attempt}") + assert offset == 999_999, ( + f"attempt {attempt}: the reconcile rolled the committed offset back " + f"to {offset} -- the session re-drains and duplicates those records" + ) + assert cursor == {"dl2": {"orch_run_seq": 3}, "dl3": {}}, ( + f"attempt {attempt}: committed cursor lost to a concurrent write" + ) + + +async def test_offset_writers_do_not_share_a_staging_path(tmp_path) -> None: + """Concurrent writers must stage to different temp files. + + A shared staging name lets one writer rename a file the other is still + filling, publishing a half-written record -- corruption that no amount of + caller-side locking would prevent. The name still ends in ``.offset.tmp`` + so ``reclaim_orphans`` keeps reaping strays. + """ + qm = FileSystemQueueManager(queues_dir=tmp_path) + await qm.append("s1", b"a") + paths = {qm._offset_tmp_path("s1") for _ in range(50)} + + assert len(paths) == 50, "staging paths collided between writes" + assert all(p.name.endswith(".offset.tmp") for p in paths) + assert all(qm._offset_tmp_owner(p) == "s1" for p in paths), ( + "reclaim_orphans could no longer tell which session a staging file " + "belongs to, and would reap one belonging to a live session" + ) + + # A staging file beside a live log is kept; a log-less stray is reaped. + live = qm._offset_tmp_path("s1") + live.write_text("{}", encoding="utf-8") + stray = qm._offset_tmp_path("s-gone") + stray.write_text("{}", encoding="utf-8") + + await qm.reclaim_orphans(before_ts=time.time() + 60) + + assert live.exists(), "reclaim reaped a staging file belonging to a live session" + assert not stray.exists(), "reclaim left a log-less staging file behind" + + +# --------------------------------------------------------------------------- +# Fix 9: every drain-loop read quarantines a corrupt offset, not just the first +# --------------------------------------------------------------------------- + + +async def test_corrupt_offset_at_idle_recheck_quarantines_worker( + caplog, monkeypatch +) -> None: + """The dry-exit recheck read must quarantine like the main-loop read. + + A recovered worker with no backlog rechecks before exiting; an offset that + goes unparseable between the two reads used to escape to the supervisor and + crash-loop the session. + """ + reg = SessionRegistry() + qm = cast(FileSystemQueueManager, reg.queue_manager) + sid = "s-corrupt-recheck" + worker = _make_worker(sid, _InertGraph()) + worker.live_event_seen = False + reg._register_for_test(worker) + + reads = {"n": 0} + real_read = qm._read_offset_record + + def _corrupt_after_first(session_id: str): + reads["n"] += 1 + if reads["n"] > 1: + raise ValueError(f"unparseable offset document for {session_id!r}") + return real_read(session_id) + + monkeypatch.setattr(qm, "_read_offset_record", _corrupt_after_first) + + with caplog.at_level(logging.ERROR, logger="context_intelligence_server"): + await reg.drain_worker(worker, flush_timeout=0.05) + + assert any("drain_worker_quarantined" in r.getMessage() for r in caplog.records) + assert sid not in reg._workers + + +async def test_corrupt_offset_during_finalize_tail_quarantines_worker( + caplog, +) -> None: + """The finalize tail-drain read must quarantine like the main-loop read. + + This is the common ``session:end`` path; an unguarded ValueError here + escaped to the supervisor instead of quarantining the one bad session. + """ + reg = SessionRegistry() + qm = cast(FileSystemQueueManager, reg.queue_manager) + sid = "s-corrupt-finalize" + await qm.append(sid, _line("session:end", {"session_id": sid})) + qm._offset_path(sid).write_text('{"v":1,"offset":"NOT-AN-INT"}', encoding="utf-8") + + worker = _make_worker(sid, _InertGraph()) + reg._register_for_test(worker) + + with ( + caplog.at_level(logging.ERROR, logger="context_intelligence_server"), + pytest.raises(_SessionQuarantined), + ): + await reg._drain_to_eof(worker, handlers=None) + + assert any("drain_worker_quarantined" in r.getMessage() for r in caplog.records) + assert sid not in reg._workers + assert worker.store_closed is True + + +# --------------------------------------------------------------------------- +# Fix 10: a finalized session leaves no dead state behind for a reused id +# --------------------------------------------------------------------------- + + +async def test_delete_drained_retires_dead_letters_for_id_reuse(tmp_path) -> None: + """Finalizing must retire the dead letters out of the session's own name. + + A session id can be reused. Left in place, the previous session's dead + payloads make the new session's first read skip its own leading lines -- + committing past events that were never processed. + """ + qm = FileSystemQueueManager(queues_dir=tmp_path) + sid = "s-reuse" + poison = b"poison-line" + + await qm.append(sid, poison) + await qm.dead_letter(sid, poison, "boom") + await qm.commit(sid, len(poison) + 1, None) + assert await qm.delete_drained(sid) is True + + assert not (tmp_path / f"{sid}.dead.jsonl").exists() + retired = list(tmp_path.glob(f"{sid}.finalized-*.dead.jsonl")) + assert len(retired) == 1, "dead-letter payloads must be retained, not deleted" + assert sid not in qm._dead_unreconciled + + # The reused id must read its own first line, not skip it as already-dead. + await qm.append(sid, poison) + batch = await qm.read_batch(sid, max_items=10) + assert [r.raw for r in batch.records] == [poison], ( + "the reused session skipped its own line against the previous " + "session's dead payloads" + ) + + +# --------------------------------------------------------------------------- +# Fix 11: the boot-sweep reconcile is COUNTED, so delete_drained's eviction +# cannot drop the guard out from under it (the guard-eviction clobber, one +# layer down from Fix 8). Also: reclaim_orphans unlinks under the key lock, +# and the finalize offset reads degrade on a corrupt .offset. +# --------------------------------------------------------------------------- + + +def _seed_reconcilable(qm: FileSystemQueueManager, key: str, dead_lines: int) -> int: + """Seed a fully-committed log of ``dead_lines`` already-dead lines + a dead + file naming that payload, so ``_reconcile_dead_key`` does a real RMW and + ``delete_drained`` accepts (size == committed) and evicts. Returns the log + byte size.""" + body = b"D\n" * dead_lines + qm._log_path(key).write_bytes(body) + qm._offset_path(key).write_text( + json.dumps({"v": 1, "offset": len(body), "cursor": None}), encoding="utf-8" + ) + qm._dead_path(key).write_text( + json.dumps({"ts": time.time(), "error": "poison", "payload": "D"}) + "\n", + encoding="utf-8", + ) + return len(body) + + +async def test_boot_reconcile_counted_guard_survives_eviction_race( + tmp_path, monkeypatch +) -> None: + """The boot-sweep reconcile must hold the COUNTED guard. + + ``recovery_reconcile_dead`` reaches ``_reconcile_dead_key`` on a worker + thread. Before the fix it took ``file_lock`` through the RAW ``_key_guard`` + accessor, so ``delete_drained``'s ``waiters == 1`` eviction gate could not + see it, evicted the guard-map entry mid-reconcile, and the next writer + minted a SECOND ``_KeyLock`` over the same file -- two locks, no mutual + exclusion, clobber. Drive the reconcile concurrently with an + eviction-then-remint on the same key, under maximally adverse scheduling, + and assert NO key ever had two distinct ``_KeyLock`` objects held at once. + """ + # Passive detector: two distinct _KeyLock objects held for one key at the + # same instant is exactly the bypass. Patched on the class, auto-restored. + from context_intelligence_server.queue_manager import filesystem as fsmod + + held: dict[str, set[int]] = {} + held_lock = threading.Lock() + violations: list[str] = [] + real_acquire = fsmod._KeyLock.acquire + real_release = fsmod._KeyLock.release + + def spy_acquire(self, blocking: bool = True, timeout: float = -1) -> bool: + ok = real_acquire(self, blocking, timeout) + key = getattr(self, "_probe_key", None) + if ok and key is not None: + with held_lock: + s = held.setdefault(key, set()) + s.add(id(self)) + if len(s) > 1: + violations.append(f"{key}:{sorted(s)}") + return ok + + def spy_release(self) -> None: + key = getattr(self, "_probe_key", None) + if key is not None: + with held_lock: + held.get(key, set()).discard(id(self)) + real_release(self) + + real_key_guard = FileSystemQueueManager._key_guard + + def tagging_key_guard(self, worker_key: str): + g = real_key_guard(self, worker_key) + # Tag the lock so the spy can attribute it to a key. + g.file_lock._probe_key = worker_key # type: ignore[attr-defined] + return g + + monkeypatch.setattr(fsmod._KeyLock, "acquire", spy_acquire) + monkeypatch.setattr(fsmod._KeyLock, "release", spy_release) + monkeypatch.setattr(FileSystemQueueManager, "_key_guard", tagging_key_guard) + + old_interval = sys.getswitchinterval() + sys.setswitchinterval(1e-6) # maximally adverse: switch every few bytecodes + torn = 0 + rounds = 300 + try: + qm = FileSystemQueueManager(queues_dir=tmp_path) + for r in range(rounds): + key = f"k{r}" + _seed_reconcilable(qm, key, dead_lines=64) + barrier = threading.Barrier(2) + + def _reconcile_thread(k: str = key, b: threading.Barrier = barrier) -> None: + b.wait() + # The boot-sweep entry point (iterates keys off worker thread). + qm._reconcile_dead_key(k) + + t = threading.Thread(target=_reconcile_thread, name="boot-reconcile") + t.start() + + async def _evict_then_remint( + k: str = key, b: threading.Barrier = barrier + ) -> None: + await asyncio.to_thread(b.wait) + await qm.delete_drained(k) # the eviction path + await qm.append(k, b"D\n") # the next writer that could re-mint + + await _evict_then_remint() + t.join(10) + + # The final .offset must be a parseable record, never torn. + try: + qm._read_offset_record(key) + except (ValueError, OSError): + torn += 1 + for p in tmp_path.glob(f"{key}*"): + p.unlink(missing_ok=True) + finally: + sys.setswitchinterval(old_interval) + + assert violations == [], ( + f"two distinct _KeyLock objects were held for one key at once over " + f"{rounds} rounds -- the guard was evicted underneath a live reconcile: " + f"{violations[:3]}" + ) + assert torn == 0, f"{torn}/{rounds} rounds left a torn .offset record" + + +async def test_reclaim_orphans_unlinks_offset_under_key_lock( + tmp_path, monkeypatch +) -> None: + """``reclaim_orphans`` must unlink an orphan ``.offset`` while holding the + key's ``file_lock`` -- the same lock every offset writer takes -- so the + unlink can never race a concurrent offset write. Observed by checking the + lock state at the moment of the unlink.""" + qm = FileSystemQueueManager(queues_dir=tmp_path) + stem = "orphan-key" + # An orphan .offset (NO .log beside it) -> a reclaim candidate. + qm._offset_path(stem).write_text( + json.dumps({"v": 1, "offset": 5, "cursor": None}), encoding="utf-8" + ) + # Same guard object reclaim's counted _guard(stem) will resolve to. + guard = qm._key_guard(stem) + target = qm._offset_path(stem) + + observed: dict[str, bool] = {} + real_unlink = Path.unlink + + def spy_unlink(self: Path, *args: Any, **kwargs: Any) -> None: + if self == target: + observed["locked_at_unlink"] = guard.file_lock.locked() + return real_unlink(self, *args, **kwargs) + + monkeypatch.setattr(Path, "unlink", spy_unlink) + + result = await qm.reclaim_orphans(before_ts=time.time() + 60) + + assert observed.get("locked_at_unlink") is True, ( + "reclaim_orphans unlinked the orphan .offset WITHOUT holding the key's " + "file_lock -- a concurrent offset write could race the unlink" + ) + assert not target.exists() + assert result["reclaimed"] >= 1 + + +async def test_finalize_offset_reads_degrade_on_corrupt_offset(tmp_path) -> None: + """``delete_drained`` and ``is_fully_drained`` read the committed offset on + the finalize path (``session:end``). A corrupt ``.offset`` must DEGRADE + (retain / report-not-drained) instead of raising out to the drain + supervisor and crashing the worker.""" + qm = FileSystemQueueManager(queues_dir=tmp_path) + sid = "s-finalize-corrupt" + await qm.append(sid, b"one-real-line\n") + # A genuinely corrupt committed offset (non-int) beside a present log. + qm._offset_path(sid).write_text('{"v":1,"offset":"NOT-AN-INT"}', encoding="utf-8") + + # Neither call may raise; both degrade to the safe, conservative answer. + drained = await qm.is_fully_drained(sid) + assert drained is False, ( + "is_fully_drained must report NOT drained (conservative) on a corrupt " + ".offset, never raise" + ) + + deleted = await qm.delete_drained(sid) + assert deleted is False, ( + "delete_drained must retain (return False) on a corrupt .offset so " + "finalize takes its bounded retain/give-up path -- never crash" + ) + # Files retained for the boot RESET_OFFSET pass to heal, not destroyed. + assert qm._log_path(sid).exists() + assert qm._offset_path(sid).exists() diff --git a/tests/test_main.py b/tests/test_main.py index 1bfa6129..4db529fc 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -1051,7 +1051,7 @@ async def test_crash_recovery_topup_drains_deferred_tail_across_passes( # stops reporting them (exactly what a real drainer does on completion). for sid in sids[:2]: batch = await qm.read_batch(sid, max_items=10) - await qm.commit(sid, batch.end_offset) + await qm.commit(sid, batch.end_offset, None) # Pass 2: the previously-DEFERRED tail is now dispatched -- not stranded. spawned.clear() @@ -1324,7 +1324,7 @@ async def test_lifespan_seeds_counters_from_disk(tmp_path: Path) -> None: await seed_qm.append(sid, line1) await seed_qm.append(sid, line2) committed = len(line1) + 1 # +1 for the appended trailing newline - await seed_qm.commit(sid, committed) + await seed_qm.commit(sid, committed, None) # Fresh registry reusing the same on-disk queue dir. reg = SessionRegistry() diff --git a/tests/test_queue_manager.py b/tests/test_queue_manager.py index 33f3790a..a488c4e7 100644 --- a/tests/test_queue_manager.py +++ b/tests/test_queue_manager.py @@ -95,7 +95,7 @@ async def test_committing_rec_end_advances_exactly_one_record(qm, tmp_path): await qm.append("s1", b"third") batch = await qm.read_batch("s1", max_items=10) - await qm.commit("s1", batch.records[0].end) + await qm.commit("s1", batch.records[0].end, None) remaining = await qm.read_batch("s1", max_items=10) assert [r.raw for r in remaining.records] == [b"second", b"third"] @@ -208,7 +208,7 @@ async def test_commit_advances_offset(qm): await qm.append("s1", b"a") await qm.append("s1", b"b") first = await qm.read_batch("s1", max_items=1) - await qm.commit("s1", first.end_offset) + await qm.commit("s1", first.end_offset, None) await qm.append("s1", b"c") second = await qm.read_batch("s1", max_items=10) assert second.lines == [b"b", b"c"] @@ -221,7 +221,7 @@ async def test_commit_persists_across_a_new_instance(tmp_path): await qm1.append("s1", b"a") await qm1.append("s1", b"b") batch = await qm1.read_batch("s1", max_items=1) - await qm1.commit("s1", batch.end_offset) + await qm1.commit("s1", batch.end_offset, None) qm2 = FileSystemQueueManager(queues_dir=qdir) # simulate restart resumed = await qm2.read_batch("s1", max_items=10) assert resumed.lines == [b"b"] @@ -229,9 +229,9 @@ async def test_commit_persists_across_a_new_instance(tmp_path): async def test_commit_is_atomic_no_temp_leftover(qm, tmp_path): await qm.append("s1", b"a") - await qm.commit("s1", 2) + await qm.commit("s1", 2, None) qdir = tmp_path / "queues" - assert (qdir / "s1.offset").read_text("utf-8") == "2" + assert (qdir / "s1.offset").read_text("utf-8") == '{"v":1,"offset":2,"cursor":null}' assert list(qdir.glob("*.tmp")) == [] @@ -315,11 +315,37 @@ async def test_active_sessions_excludes_fully_committed(qm): await qm.append("s_active", b"x") # appended, never committed -> undrained await qm.append("s_done", b"y") done = await qm.read_batch("s_done", max_items=10) - await qm.commit("s_done", done.end_offset) # drained + await qm.commit("s_done", done.end_offset, None) # drained active = await qm.active_sessions() assert active == ["s_active"] +async def test_is_fully_drained_true_for_unknown_session(qm): + assert await qm.is_fully_drained("never_seen") is True + + +async def test_is_fully_drained_false_while_uncommitted(qm): + await qm.append("s1", b"x") + assert await qm.is_fully_drained("s1") is False + + +async def test_is_fully_drained_true_after_commit(qm): + await qm.append("s1", b"x") + batch = await qm.read_batch("s1", max_items=10) + await qm.commit("s1", batch.end_offset, None) + assert await qm.is_fully_drained("s1") is True + + +async def test_is_fully_drained_ignores_torn_trailing_fragment(qm, tmp_path): + # A torn tail (bytes after the final newline) is not complete data, so a + # session whose complete lines are all committed reads as drained. + log = tmp_path / "queues" / "s1.log" + log.write_bytes(b"a\nb\nTORN_PARTIAL") + batch = await qm.read_batch("s1", max_items=10) + await qm.commit("s1", batch.end_offset, None) + assert await qm.is_fully_drained("s1") is True + + async def test_recover_empty_dir_is_safe(qm): assert await qm.recover() == [] @@ -328,7 +354,7 @@ async def test_recover_reports_session_with_uncommitted_complete_line(qm, tmp_pa log = tmp_path / "queues" / "s1.log" log.write_bytes(b"a\nb\nTORN") # two complete lines + torn tail assert await qm.recover() == ["s1"] - await qm.commit("s1", 4) # past 'a\nb\n' == 4 bytes + await qm.commit("s1", 4, None) # past 'a\nb\n' == 4 bytes assert await qm.recover() == [] # only torn tail remains -> not recoverable @@ -359,7 +385,7 @@ async def test_read_batch_rejects_unsafe_session_id(qm, bad_id): @pytest.mark.parametrize("bad_id", ["", "a/b", "a\\b", "a\x00b"]) async def test_commit_rejects_unsafe_session_id(qm, bad_id): with pytest.raises(ValueError): - await qm.commit(bad_id, 0) + await qm.commit(bad_id, 0, None) @pytest.mark.parametrize("bad_id", ["", "a/b", "a\\b", "a\x00b"]) @@ -379,14 +405,17 @@ async def test_delete_drained_removes_log_and_offset_keeps_dead(tmp_path) -> Non qm = FileSystemQueueManager(queues_dir=tmp_path) await qm.append("s", b"line") - await qm.commit("s", 5) + await qm.commit("s", 5, None) await qm.dead_letter("s", b"bad\n", "boom") await qm.delete_drained("s") assert not (tmp_path / "s.log").exists() assert not (tmp_path / "s.offset").exists() - assert (tmp_path / "s.dead.jsonl").exists() # retained + # Retained, but retired out of the live name so a session reusing this id + # cannot reconcile against these payloads. Still reported for the session. + assert not (tmp_path / "s.dead.jsonl").exists() + assert len(list(tmp_path.glob("s.finalized-*.dead.jsonl"))) == 1 assert len(await qm.read_dead_letters("s")) == 1 @@ -473,7 +502,7 @@ async def test_recovery_seed_counts_pending_and_committed(qm): await qm.append("s1", b"a") await qm.append("s1", b"b") await qm.append("s1", b"c") - await qm.commit("s1", 4) # commit the first two complete lines + await qm.commit("s1", 4, None) # commit the first two complete lines accepted, written = await qm.recovery_seed_counts() @@ -484,7 +513,7 @@ async def test_recovery_seed_counts_pending_and_committed(qm): async def test_recovery_seed_counts_committed_includes_dead(qm): # C=1 committed, P=0 pending, D=1 dead. before-dead == 0. await qm.append("s2", b"a") - await qm.commit("s2", 2) + await qm.commit("s2", 2, None) await qm.dead_letter("s2", b"a", error="boom") accepted, written = await qm.recovery_seed_counts() @@ -508,10 +537,10 @@ async def test_recovery_seed_counts_residual_is_zero_mixed_shape(qm): await qm.append("a", b"1") await qm.append("a", b"2") await qm.append("a", b"3") - await qm.commit("a", 4) + await qm.commit("a", 4, None) # Key B: 1 committed + 1 dead. await qm.append("b", b"x") - await qm.commit("b", 2) + await qm.commit("b", 2, None) await qm.dead_letter("b", b"x", error="boom") # Key C: dead-only (log reclaimed). await qm.dead_letter("c", b"poison", error="boom") @@ -601,7 +630,7 @@ async def test_recovery_seed_counts_replay_window_residual_zero(qm): # log = [line0 committed][line0 re-appended pending]. C=1, P=1, D=1. # The re-appended line is absorbed into accepted_seed (counted in P and D). await qm.append("s6", b"a") - await qm.commit("s6", 2) + await qm.commit("s6", 2, None) await qm.dead_letter("s6", b"a", error="boom") await qm.append("s6", b"a") # re-append the dead line for replay @@ -641,7 +670,7 @@ async def test_spool_stats_fully_committed_session_not_pending(qm): (spool_bytes_total still reflects them).""" await qm.append("s1", b"a") line = b"a\n" - await qm.commit("s1", len(line)) + await qm.commit("s1", len(line), None) stats = await qm.spool_stats() @@ -666,7 +695,7 @@ async def test_spool_stats_multiple_sessions_aggregate(qm): await qm.append("s1", b"a") # pending await qm.append("s2", b"b") line = b"b\n" - await qm.commit("s2", len(line)) # fully committed, not pending + await qm.commit("s2", len(line), None) # fully committed, not pending await qm.append("s3", b"c") # pending stats = await qm.spool_stats() @@ -825,7 +854,7 @@ async def test_recovery_seed_counts_unchanged_under_streaming(qm): await qm.append("s1", b"a") await qm.append("s1", b"bb") line1 = b"a\n" - await qm.commit("s1", len(line1)) # 1 written, 1 still pending + await qm.commit("s1", len(line1), None) # 1 written, 1 still pending accepted, written = await qm.recovery_seed_counts() @@ -927,9 +956,173 @@ async def test_spool_stats_healthy_offsets_report_zero_corrupt(qm): fire on the normal committed-offset path).""" await qm.append("s1", b"a") line = b"a\n" - await qm.commit("s1", len(line)) # writes a valid numeric .offset + await qm.commit("s1", len(line), None) # writes a valid numeric .offset qm._spool_cache = None stats = await qm.spool_stats() assert stats["corrupt_offsets"] == 0 + + +# --------------------------------------------------------------------------- +# Durable cursor folded into the atomic offset write +# --------------------------------------------------------------------------- + + +async def test_commit_requires_cursor_argument(qm): + """cursor has no default: omitting it fails loudly at call time. + + A default would silently null the cursor on a missed migration site or a + rolling deploy against an older signature -- the exact silent-loss class. + """ + with pytest.raises(TypeError): + await qm.commit("s1", 0) # type: ignore[call-arg] + + +async def test_commit_rejects_wrong_typed_cursor(qm): + """A non-dict, non-None cursor fails loud at write time. + + Without this, a wrong-typed cursor is written verbatim and silently reads + back as None -- the same silent cross-handler-counter reset the required + arg exists to prevent. + """ + await qm.append("s1", b"a") + with pytest.raises(TypeError): + await qm.commit("s1", 2, "not-a-dict") # type: ignore[arg-type] + + +async def test_commit_persists_cursor_in_same_record_as_offset(qm): + cursor = {"dl2": {"iteration_count": 4}, "dl3": {}} + await qm.append("s1", b"a") + await qm.commit("s1", 2, cursor) + assert await qm.read_cursor("s1") == cursor + assert qm._read_committed_offset("s1") == 2 # offset + cursor never skew + + +async def test_read_cursor_is_none_for_bare_int_and_missing(qm): + assert await qm.read_cursor("never_seen") is None + qm._offset_path("s1").write_text("42", encoding="utf-8") # legacy bare int + assert await qm.read_cursor("s1") is None + assert qm._read_committed_offset("s1") == 42 # bare-int still parses + + +async def test_rolling_upgrade_bare_int_then_envelope(qm): + # An old worker wrote a bare-int offset; the new worker commits an envelope + # on top of the same session -- both are readable, the cursor now persists. + qm._offset_path("s1").write_text("10", encoding="utf-8") + assert await qm.read_cursor("s1") is None + await qm.commit("s1", 20, {"dl2": {"iteration_count": 9}, "dl3": {}}) + assert qm._read_committed_offset("s1") == 20 + assert await qm.read_cursor("s1") == {"dl2": {"iteration_count": 9}, "dl3": {}} + + +async def test_corrupt_offset_record_raises_not_silently_zero(qm): + # A malformed record must raise, never degrade to 0 (0 replays the whole + # log and manufactures duplicate nodes -- a worse, quieter failure). + qm._offset_path("s1").write_text('{"v":1,"offset":"NaN"}', encoding="utf-8") + with pytest.raises(ValueError): + qm._read_committed_offset("s1") + + +async def test_commit_is_atomic_across_a_mid_write_crash(qm, tmp_path, monkeypatch): + # Simulate os.replace failing mid-commit: the previously committed record + # must survive intact (no torn/partial offset file), and no .tmp leaks. + await qm.append("s1", b"a") + await qm.commit("s1", 2, {"dl2": {"iteration_count": 1}, "dl3": {}}) + + import context_intelligence_server.queue_manager.filesystem as qmmod + + def _boom(src, dst): + raise OSError("simulated crash during os.replace") + + monkeypatch.setattr(qmmod.os, "replace", _boom) + with pytest.raises(OSError): + await qm.commit("s1", 99, {"dl2": {"iteration_count": 2}, "dl3": {}}) + monkeypatch.undo() + + # The pre-crash record is intact; the crashed write left nothing behind. + assert qm._read_committed_offset("s1") == 2 + assert await qm.read_cursor("s1") == {"dl2": {"iteration_count": 1}, "dl3": {}} + assert list((tmp_path / "queues").glob("*.tmp")) == [] + + +# --------------------------------------------------------------------------- +# Every .offset writer preserves the cursor (not just commit()) +# --------------------------------------------------------------------------- + + +async def test_compaction_preserves_committed_cursor(qm): + """Idle compaction must NOT wipe the cursor the last commit persisted. + + compact_committed_prefix rebases the .offset to 0; if it writes a bare + "0" it destroys active_orch_run_id/orch_run_seq moments after commit wrote + them. Since queue_compact_enabled defaults True, this is the common path. + """ + cursor = { + "dl2": {"active_orch_run_id": "s1::orch_run::T::1", "orch_run_seq": 1}, + "dl3": {}, + } + await qm.append("s1", b"a") # bytes [0, 2) + await qm.append("s1", b"b") # bytes [2, 4) -- an undrained tail to keep + await qm.commit("s1", 2, cursor) # commit past the first line only + + reclaimed = await qm.compact_committed_prefix("s1", 0) + + assert reclaimed == 2 # the committed prefix was reclaimed + assert qm._read_committed_offset("s1") == 0 # offset rebased + assert await qm.read_cursor("s1") == cursor # RED before fix: cursor wiped + + +async def test_compaction_restore_on_replace_failure_preserves_cursor(qm, monkeypatch): + """A failed log-replace restores offset := C AND the cursor with it.""" + cursor = {"dl2": {"orch_run_seq": 7}, "dl3": {}} + await qm.append("s1", b"a") + await qm.append("s1", b"b") + await qm.commit("s1", 2, cursor) + + import context_intelligence_server.queue_manager.filesystem as qmmod + + real_replace = qmmod.os.replace + + def _fail_log_replace(src, dst): + # Fail only the .log replace (Step 6); let the .offset writes through. + if str(dst).endswith(".log"): + raise OSError("simulated log replace failure") + return real_replace(src, dst) + + monkeypatch.setattr(qmmod.os, "replace", _fail_log_replace) + reclaimed = await qm.compact_committed_prefix("s1", 0) + monkeypatch.undo() + + assert reclaimed == 0 # compaction was a no-op + assert qm._read_committed_offset("s1") == 2 # offset restored to C + assert await qm.read_cursor("s1") == cursor # cursor restored with it + + +async def test_recovery_reconcile_dead_preserves_cursor(qm): + """Advancing the offset past already-dead lines must keep the cursor.""" + cursor = {"dl2": {"orch_run_seq": 3}, "dl3": {}} + line = b'{"event":"x","workspace":"/ws","data":{}}' + await qm.append("s1", line) # one pending line at offset 0 + await qm.commit("s1", 0, cursor) # committed=0, cursor persisted + # Dead-letter that exact payload so reconcile steps the offset past it. + await qm.dead_letter("s1", line, "poison") + + skipped = await qm.recovery_reconcile_dead() + + assert skipped == 1 # the leading dead line was skipped + assert qm._read_committed_offset("s1") == len(line) + 1 # advanced past it + assert await qm.read_cursor("s1") == cursor # RED before fix: cursor wiped + + +async def test_read_cursor_unknown_version_degrades_to_none(qm): + """An unknown envelope version keeps the offset but drops the cursor. + + Previously-untested: a forward-version record must read back as (offset, + None) rather than handing a foreign-shaped cursor to restore_cursor. + """ + qm._offset_path("s1").write_text( + '{"v":2,"offset":5,"cursor":{"dl2":{"orch_run_seq":9}}}', encoding="utf-8" + ) + assert await qm.read_cursor("s1") is None + assert qm._read_committed_offset("s1") == 5 # offset preserved diff --git a/tests/test_registry.py b/tests/test_registry.py index f707d973..84ebdf0d 100644 --- a/tests/test_registry.py +++ b/tests/test_registry.py @@ -505,7 +505,12 @@ async def test_dead_letter_logs_warning_with_session_id( poison.records = [Record(b"{ this is not valid json", 0, 25)] with caplog.at_level(logging.WARNING, logger="context_intelligence_server"): - await reg._handle_exhausted_batch(worker, poison, handlers=MagicMock()) + await reg._handle_exhausted_batch( + worker, + poison, + handlers=MagicMock(), + pre_batch_cursor=worker.services.snapshot_cursor(), + ) reg.queue_manager.dead_letter.assert_awaited() records = [ diff --git a/tests/test_services.py b/tests/test_services.py index 7b37014b..1a78d01a 100644 --- a/tests/test_services.py +++ b/tests/test_services.py @@ -2,16 +2,16 @@ from __future__ import annotations -import pytest from unittest.mock import AsyncMock, patch +import pytest +from context_intelligence_server.handlers.data_layer_3.state import DataLayer3State from context_intelligence_server.services import ( GraphState, HookConfig, HookStateService, ) - # --------------------------------------------------------------------------- # HookConfig tests # --------------------------------------------------------------------------- @@ -710,3 +710,84 @@ def test_created_by_propagated_to_custom_graph_store(self) -> None: store = GraphState() svc = HookStateService(workspace="/ws", graph_store=store, created_by="carol") assert svc.graph.created_by == "carol" + + +class TestDurableCursor: + def test_snapshot_restore_round_trips_cursor_state(self) -> None: + src = HookStateService(workspace="/ws") + src.data_layer_2.iteration_count = 7 + src.data_layer_2.execution_start_ts = "2026-01-01T00:00:02+00:00" + src.data_layer_2.active_iteration_id = "sid::iteration::7" + src.data_layer_3.active_recipe_run_stack = ["run-1", "run-2"] + + snapshot = src.snapshot_cursor() + + # A fresh worker (empty in-process state) restores from the snapshot. + dst = HookStateService(workspace="/ws") + dst.restore_cursor(snapshot) + assert dst.data_layer_2.iteration_count == 7 + assert dst.data_layer_2.execution_start_ts == "2026-01-01T00:00:02+00:00" + assert dst.data_layer_2.active_iteration_id == "sid::iteration::7" + assert dst.data_layer_3.active_recipe_run_stack == ["run-1", "run-2"] + + def test_restore_none_is_a_noop(self) -> None: + svc = HookStateService(workspace="/ws") + svc.data_layer_2.iteration_count = 3 + svc.restore_cursor(None) # legacy .offset with no cursor + assert svc.data_layer_2.iteration_count == 3 + + def test_restore_drops_unknown_keys_and_keeps_defaults(self) -> None: + svc = HookStateService(workspace="/ws") + svc.restore_cursor({"dl2": {"iteration_count": 5, "gone_field": "x"}}) + assert svc.data_layer_2.iteration_count == 5 + assert not hasattr(svc.data_layer_2, "gone_field") + # A field absent from the record keeps its dataclass default. + assert svc.data_layer_2.execution_start_ts is None + + def test_corrupt_record_never_raises(self) -> None: + svc = HookStateService(workspace="/ws") + svc.restore_cursor({"dl2": "not-a-dict", "dl3": 123}) # type: ignore[dict-item] + assert svc.data_layer_2.iteration_count == 0 + + def test_restore_cursor_all_or_nothing_under_partial_failure( + self, monkeypatch + ) -> None: + """A failure partway through the restore (the second target's rebuild + raises) must leave BOTH live dataclasses at their pre-restore state -- + never a half-restored hybrid where dl2 was replaced but dl3 was not.""" + import dataclasses as _dc + + svc = HookStateService(workspace="/ws") + svc.data_layer_2.iteration_count = 3 + svc.data_layer_2.orch_run_seq = 5 + svc.data_layer_2.active_orch_run_id = "OLD" + before_dl2 = svc.data_layer_2 + before_dl3 = svc.data_layer_3 + + real_replace = _dc.replace + + def _fail_on_dl3(obj, **kw): + if isinstance(obj, DataLayer3State): + raise RuntimeError("injected mid-rebuild failure on dl3") # noqa: TRY004 + return real_replace(obj, **kw) + + monkeypatch.setattr( + "context_intelligence_server.services.dataclasses.replace", _fail_on_dl3 + ) + svc.restore_cursor( + { + "dl2": { + "iteration_count": 99, + "orch_run_seq": 77, + "active_orch_run_id": "NEW", + }, + "dl3": {}, + } + ) + + # Neither live object was reassigned, and dl2's values are untouched. + assert svc.data_layer_2 is before_dl2 + assert svc.data_layer_3 is before_dl3 + assert svc.data_layer_2.iteration_count == 3 + assert svc.data_layer_2.orch_run_seq == 5 + assert svc.data_layer_2.active_orch_run_id == "OLD" diff --git a/tests/test_steady_state_reclaim.py b/tests/test_steady_state_reclaim.py index d11f9c96..e50d2424 100644 --- a/tests/test_steady_state_reclaim.py +++ b/tests/test_steady_state_reclaim.py @@ -89,7 +89,7 @@ async def test_b_undrained_tail_never_reclaimed_past_committed_c_less_than_tail( first_batch = await qm.read_batch(sid, max_items=40) assert len(first_batch.records) == 40 - await qm.commit(sid, first_batch.end_offset) + await qm.commit(sid, first_batch.end_offset, None) c = first_batch.end_offset log_path = tmp_path / f"{sid}.log" e = log_path.stat().st_size @@ -117,7 +117,7 @@ async def test_b_undrained_tail_never_reclaimed_past_committed_c_greater_than_ta first_batch = await qm.read_batch(sid, max_items=70) assert len(first_batch.records) == 70 - await qm.commit(sid, first_batch.end_offset) + await qm.commit(sid, first_batch.end_offset, None) c = first_batch.end_offset log_path = tmp_path / f"{sid}.log" e = log_path.stat().st_size @@ -139,7 +139,7 @@ async def test_b_below_min_prefix_bytes_is_a_noop(tmp_path: Path) -> None: for i in range(10): await qm.append(sid, _fixed(i)) batch = await qm.read_batch(sid, max_items=5) - await qm.commit(sid, batch.end_offset) + await qm.commit(sid, batch.end_offset, None) log_path = tmp_path / f"{sid}.log" before = log_path.read_bytes() @@ -164,13 +164,13 @@ async def test_c_mid_copy_oserror_is_a_pure_noop( for i in range(9): await qm.append(sid, _fixed(i)) batch = await qm.read_batch(sid, max_items=3) - await qm.commit(sid, batch.end_offset) + await qm.commit(sid, batch.end_offset, None) log_path = tmp_path / f"{sid}.log" offset_path = tmp_path / f"{sid}.offset" log_before = log_path.read_bytes() offset_before = offset_path.read_text(encoding="utf-8") - assert offset_before.strip() == str(batch.end_offset) + assert offset_before.strip() == f'{{"v":1,"offset":{batch.end_offset},"cursor":null}}' def _raise(fd: int, data: bytes) -> None: raise OSError("simulated mid-copy failure") @@ -198,7 +198,7 @@ async def test_c_window2_offset_rebased_before_log_replaced_bounded_redrive( for ev in events: await qm.append(sid, ev) batch = await qm.read_batch(sid, max_items=3) - await qm.commit(sid, batch.end_offset) # committed = 3 events (30 bytes) + await qm.commit(sid, batch.end_offset, None) # committed = 3 events (30 bytes) # Simulate the crash: offset already rebased to 0 (step 5 completed), # log NOT yet replaced (step 6 never ran). @@ -226,7 +226,7 @@ async def test_c_control_rejected_log_then_offset_order_loses_data( for ev in events: await qm.append(sid, ev) batch = await qm.read_batch(sid, max_items=3) - await qm.commit(sid, batch.end_offset) # committed = 30 bytes (C == 30) + await qm.commit(sid, batch.end_offset, None) # committed = 30 bytes (C == 30) log_path = tmp_path / f"{sid}.log" c = qm._read_committed_offset(sid) @@ -270,11 +270,10 @@ async def test_i_replace_failure_restores_offset_zero_accounting_drift( for ev in events: await qm.append(sid, ev) batch = await qm.read_batch(sid, max_items=3) - await qm.commit(sid, batch.end_offset) + await qm.commit(sid, batch.end_offset, None) c = batch.end_offset log_path = tmp_path / f"{sid}.log" - offset_path = tmp_path / f"{sid}.offset" log_before = log_path.read_bytes() real_replace = os.replace @@ -293,8 +292,11 @@ def _flaky_replace(src: Any, dst: Any) -> None: reclaimed = await qm.compact_committed_prefix(sid, 0) # must not raise assert reclaimed == 0 - # Offset restored to C -- a pure no-op, not a re-drive. - assert offset_path.read_text(encoding="utf-8").strip() == str(c) + # Offset restored to C -- a pure no-op, not a re-drive. Read via the + # committed-offset accessor: the restore now writes the same atomic + # ``{"v":1,"offset":C,"cursor":...}`` record every other .offset writer + # uses, so the value (not the raw byte shape) is what must equal C. + assert qm._read_committed_offset(sid) == c # Log completely untouched. assert log_path.read_bytes() == log_before assert any( @@ -321,19 +323,21 @@ async def test_i_double_replace_failure_logs_restore_failed_honestly( for i in range(9): await qm.append(sid, _fixed(i)) batch = await qm.read_batch(sid, max_items=3) - await qm.commit(sid, batch.end_offset) + await qm.commit(sid, batch.end_offset, None) log_path = tmp_path / f"{sid}.log" def _always_raise(src: Any, dst: Any) -> None: # let the first offset rebase-to-0 write through, then fail both the - # log replace and the subsequent restore-to-C write + # log replace and the subsequent restore-to-C write. Both offset writes + # are now atomic ``{"v":1,"offset":N,...}`` records, so the rebase-to-0 + # is the one carrying offset 0 -- everything else is the restore. src_content = ( Path(src).read_text(encoding="utf-8") if Path(src).exists() else "" ) if str(dst) == str(log_path): raise OSError("simulated persistent log-replace failure") - if str(dst).endswith(".offset") and src_content != "0": + if str(dst).endswith(".offset") and '"offset":0' not in src_content: raise OSError("simulated persistent offset-restore failure") monkeypatch.setattr(queue_manager_module.os, "replace", _always_raise, raising=True) @@ -363,7 +367,7 @@ async def test_j_large_tail_does_not_block_prefix_reclaim( for i in range(20): await qm.append(sid, _fixed(i)) batch = await qm.read_batch(sid, max_items=2) # C = 20 bytes - await qm.commit(sid, batch.end_offset) + await qm.commit(sid, batch.end_offset, None) log_path = tmp_path / f"{sid}.log" assert batch.end_offset == 20 @@ -637,7 +641,7 @@ async def test_f_status_not_blocked_by_an_in_progress_compaction( qm = main_module.registry.queue_manager sid = "s-status-lock" await qm.append(sid, _fixed(0)) - await qm.commit(sid, 10) + await qm.commit(sid, 10, None) with qm._guard(sid) as guard: loop = asyncio.get_event_loop() diff --git a/uv.lock b/uv.lock index a21b650a..be93615c 100644 --- a/uv.lock +++ b/uv.lock @@ -233,7 +233,7 @@ wheels = [ [[package]] name = "context-intelligence-server" -version = "6.7.2" +version = "7.0.0" source = { editable = "." } dependencies = [ { name = "aiofiles" },