From d67166c6f164b4fd85ae5b4fcb802f19dc719314 Mon Sep 17 00:00:00 2001 From: colombod Date: Mon, 24 Aug 2026 15:51:11 +0000 Subject: [PATCH 1/2] feat(ingest): durable per-record cursor, retry-dedup, and run-id tiebreaker MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Persist a cross-handler cursor atomically with the committed offset so a rebuilt drain worker resumes its data-layer counters exactly where it left off, instead of resetting them and manufacturing duplicate nodes on replay. Retry-dedup and the run-id tiebreaker on parallel handlers keep a re-delivered batch idempotent. The cursor is only durable if EVERY writer of the offset record preserves it, not just commit(). The idle-compaction rebase/restore and the dead-letter reconcile route through the same atomic offset-record writer, and the boot RESET_OFFSET reclaim now preserves the committed cursor instead of wiping it on a re-drainable log. The dead-letter/exhausted-batch redrive rolls the counters back to their pre-batch snapshot before isolating records, and that snapshot is deep-copied so a later attempt's in-place mutation cannot corrupt the baseline the next rollback restores. commit() writes under the same per-key file lock as compaction, closing a race where a concurrent commit during compaction could be silently erased. Cross-handler run-id resolution is shared by the orchestrator-run, iteration, and content-block handlers, so a partial cursor after a rebuild re-derives the run's real id (with its sequence) instead of dropping the HAS_PART edge or orphaning ContentBlock nodes. A genuinely corrupt/unparseable offset quarantines the one affected drain worker (a transient I/O error still supervises and respawns as before). A crash-then-respawn mid-isolation does not re-dead-letter an already-dead leading record: the filesystem backend marks a session dead-unreconciled in-memory on each dead_letter, and the next read_batch reconciles past the leading dead lines exactly once. This is an internal backend detail, kept OFF the QueueManager Protocol -- a broker backend has no leading-dead-line window to reconcile. Process restart is still covered by the boot-time global recovery_reconcile_dead(). The exhausted-batch buffer discard also invalidates the seen-session cache, so the isolated re-dispatch re-issues the Session node instead of early-returning. Cursor restore is all-or-nothing: both replacement dataclasses are built before either live field is reassigned. Breaking Protocol change: QueueManager.commit() takes a REQUIRED cursor argument (offset and cursor are written in one atomic record), and read_cursor joins the backend-neutral surface. These are the genuine general primitives; log-structured-queue details (dead-line reconciliation, is_fully_drained) stay off the Protocol until a real cross-backend caller needs them. Every in-repo caller is updated; the tolerant reader stays forward- and backward-compatible, so a legacy bare-integer .offset still reads correctly. Version 6.8.0. 🤖 Generated with [Amplifier](https://github.com/microsoft/amplifier) Co-Authored-By: Amplifier <240397093+microsoft-amplifier@users.noreply.github.com> --- CHANGELOG.md | 57 +++ .../handlers/data_layer_2/content_block.py | 17 +- .../handlers/data_layer_2/iteration.py | 67 ++- .../handlers/data_layer_2/orchestrator_run.py | 31 +- .../handlers/data_layer_2/state.py | 59 ++- .../queue_manager/filesystem.py | 275 +++++++++--- .../queue_manager/protocol.py | 6 +- context_intelligence_server/registry.py | 105 ++++- context_intelligence_server/services.py | 84 ++++ pyproject.toml | 2 +- .../data_layer_2/test_content_block.py | 26 +- tests/handlers/data_layer_2/test_iteration.py | 316 ++++++++++++- .../data_layer_2/test_orchestrator_run.py | 81 +++- .../test_data_layer_3_delegation_skill.py | 5 +- tests/integration/test_turn_chain.py | 4 +- tests/test_boot_safety.py | 12 +- tests/test_drain_supervision.py | 196 ++++++++ tests/test_durable_append_framing.py | 8 +- tests/test_ingest_cursor_hardening.py | 422 ++++++++++++++++++ tests/test_main.py | 4 +- tests/test_queue_manager.py | 226 +++++++++- tests/test_registry.py | 7 +- tests/test_services.py | 85 +++- tests/test_steady_state_reclaim.py | 36 +- uv.lock | 2 +- 25 files changed, 1938 insertions(+), 195 deletions(-) create mode 100644 CHANGELOG.md create mode 100644 tests/test_ingest_cursor_hardening.py diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 00000000..e922bc94 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,57 @@ +# 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). + +## [6.8.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. The minor version bump + reflects the breaking Protocol surface change. +- **`QueueManager` Protocol gains `read_cursor`, `is_fully_drained`, and + `reconcile_dead`** as members every backend must provide. `read_cursor` + returns the persisted cursor for a rebuilt worker; `is_fully_drained` reports + whether a session has undrained log data; `reconcile_dead` advances one + session's committed offset past its leading already-dead lines. + +### Fixed + +- The durable cursor now survives every offset-mutation path, not just + `commit()`: idle compaction, the boot `RESET_OFFSET` reclaim, and finalize's + `delete_drained` preserve a non-empty cursor instead of wiping it with a bare + offset write. +- `commit()` writes under the same per-key file lock as compaction, closing a + race where a concurrent commit during compaction could be silently erased. +- A corrupt/unparseable `.offset` file now quarantines the one affected drain + worker (logged, closed, deregistered) instead of crash-looping it; other + sessions keep draining. +- 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. +- 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..bf105c18 100644 --- a/context_intelligence_server/queue_manager/filesystem.py +++ b/context_intelligence_server/queue_manager/filesystem.py @@ -170,6 +170,14 @@ 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] = {} + # 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 +207,69 @@ 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 _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. + + Writes ``{"v": 1, "offset": offset, "cursor": cursor}`` to a 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._dir / f"{session_id}.offset.tmp" + record = {"v": 1, "offset": offset, "cursor": cursor} + 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] @staticmethod def _last_complete_end(path: Path) -> int: @@ -506,8 +559,18 @@ async def append(self, session_id: str, raw: bytes) -> None: async def read_batch(self, session_id: str, max_items: int) -> Batch: self._validate_session_id(session_id) path = self._log_path(session_id) + # 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. Clear the flag + # unconditionally: _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. + reconcile = session_id in self._dead_unreconciled def _read() -> Batch: + if reconcile: + self._reconcile_dead_key(session_id) + self._dead_unreconciled.discard(session_id) start = self._read_committed_offset(session_id) records: list[Record] = [] consumed = 0 @@ -530,23 +593,45 @@ def _read() -> Batch: 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). - - 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. + 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 +666,10 @@ 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) async def delete_drained(self, session_id: str) -> bool: """Remove the drained ``.log``/``.offset`` for a finalized session. @@ -627,6 +716,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: @@ -667,14 +762,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 +825,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 +842,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", @@ -1126,12 +1219,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: @@ -1381,6 +1487,23 @@ 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: + return self._read_committed_offset(session_id) >= 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 +1929,46 @@ 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. + """ + # 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 + 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 +1984,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..509c8e18 100644 --- a/context_intelligence_server/registry.py +++ b/context_intelligence_server/registry.py @@ -353,10 +353,41 @@ 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) + try: + batch = await qm.read_batch(session_id, max_items=_DRAIN_MAX_BATCH) + except ValueError: + # A ValueError here means the committed offset read parsed a + # genuinely corrupt/unparseable .offset -- a PERMANENT fault, + # raised before the guarded cursor read below is reached. + # Quarantine this one session (close + deregister, exit + # cleanly) so it does not crash-loop (respawn -> same corrupt + # read -> die, forever). A transient OSError is NOT caught + # here: it propagates to _on_drain_done for supervised + # respawn, exactly as before. Other sessions keep draining; + # the boot RESET_OFFSET pass heals the offset (cursor-preserving). + logger.exception( + "drain_worker_quarantined session=%s reason=corrupt_offset", + session_id, + extra={"session_id": session_id}, + ) + await self._safe_close(worker) + self._deregister(session_id) + return if not batch.records: # Idle compaction runs before the dry-exit check below, so a @@ -402,6 +433,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 +512,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 +522,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 +534,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( @@ -541,7 +606,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 +635,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 +657,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 +685,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 @@ -632,7 +711,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", 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..ed7bb43b 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 = "6.8.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..ceac2908 --- /dev/null +++ b/tests/test_ingest_cursor_hardening.py @@ -0,0 +1,422 @@ +"""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 +from typing import Any, cast + +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 +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 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..28a00ee0 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,7 +405,7 @@ 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") @@ -473,7 +499,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 +510,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 +534,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 +627,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 +667,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 +692,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 +851,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 +953,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..705457a3 100644 --- a/uv.lock +++ b/uv.lock @@ -233,7 +233,7 @@ wheels = [ [[package]] name = "context-intelligence-server" -version = "6.7.2" +version = "6.8.0" source = { editable = "." } dependencies = [ { name = "aiofiles" }, From 56a959b793e04cbe6077f2d606e0e18fef2091d2 Mon Sep 17 00:00:00 2001 From: colombod Date: Wed, 26 Aug 2026 08:56:12 +0000 Subject: [PATCH 2/2] fix(ingest): serialise every .offset write behind one locked writer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A session's committed offset was written from several paths, but the lock that protects it was the CALLER's job. commit() and compact_committed_prefix() took the per-key file_lock; the dead-letter reconcile reached through read_batch() took none, and read_batch itself ran wholly unguarded. A commit racing that reconcile was silently rolled back to the offset the reconcile had read before it -- 248 of 300 concurrent runs, with 2 more leaving a half-written record. The staging file compounded it: every writer staged through the same .offset.tmp, so one writer could rename a file another was still filling. Move the invariant into the writer. _write_offset_record now takes the key's file_lock itself and stages through a per-write temp name, so no caller can be the one that forgets, and no two writers collide on the staging path. The lock is reentrant (_KeyLock) because compaction and the reconcile legitimately hold it across a wider atomic sequence; it keeps locked() so the framing guarantees stay assertable. Guard creation moves behind a mutex: guards are now minted from worker threads as well as the event loop, and two guards for one key would mean two locks over the same bytes. read_batch joins the guarded paths -- it can move the offset, and even its pure read must not observe the .log mid-swap. This also makes the backgrounded boot reconcile safe against live drainers without blocking first request on it. Corrupt-offset quarantine now covers every read in the drain loop, not the main-loop read alone: the idle dry-exit recheck and the session:end tail drain escaped to the supervisor and crash-looped the session. All three route through one SessionRegistry._read_batch. Finalizing a session retires its dead letters out of the session's own name. A session id can be reused, and the previous session's payloads made the new session's reconcile skip log lines whose bytes merely matched an old dead payload, committing past events it never processed. Nothing is discarded -- the retired file still expires on its own schedule and read_dead_letters still reports it under the session. CHANGELOG corrected to match the code: is_fully_drained and reconcile_dead are NOT on the Protocol (the commit that introduced them says so explicitly), and delete_drained drops the cursor by design rather than preserving it. Version 6.8.0 -> 7.0.0. commit() gained a required argument on a Protocol this repo documents as a backend extension point, and the CHANGELOG asserts SemVer adherence; a minor bump for a breaking contract change makes the version number lie about compatibility. New guard tests, each red on 17b1756 and green here: the commit/reconcile race, staging-path uniqueness and its reclaim attribution, the two previously unguarded quarantine sites, and dead-letter retirement under id reuse. 🤖 Generated with [Amplifier](https://github.com/microsoft/amplifier) Co-Authored-By: Amplifier <240397093+microsoft-amplifier@users.noreply.github.com> --- CHANGELOG.md | 43 +- .../queue_manager/filesystem.py | 538 +++++++++++++----- context_intelligence_server/registry.py | 62 +- pyproject.toml | 2 +- tests/test_ingest_cursor_hardening.py | 384 ++++++++++++- tests/test_queue_manager.py | 5 +- uv.lock | 2 +- 7 files changed, 845 insertions(+), 191 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e922bc94..d2e6da36 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,7 +5,7 @@ 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). -## [6.8.0] +## [7.0.0] ### Changed (breaking) @@ -15,25 +15,33 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 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. The minor version bump - reflects the breaking Protocol surface change. -- **`QueueManager` Protocol gains `read_cursor`, `is_fully_drained`, and - `reconcile_dead`** as members every backend must provide. `read_cursor` - returns the persisted cursor for a rebuilt worker; `is_fully_drained` reports - whether a session has undrained log data; `reconcile_dead` advances one - session's committed offset past its leading already-dead lines. + 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, not just - `commit()`: idle compaction, the boot `RESET_OFFSET` reclaim, and finalize's - `delete_drained` preserve a non-empty cursor instead of wiping it with a bare - offset write. -- `commit()` writes under the same per-key file lock as compaction, closing a - race where a concurrent commit during compaction could be silently erased. +- 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. + 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. @@ -43,6 +51,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - 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. diff --git a/context_intelligence_server/queue_manager/filesystem.py b/context_intelligence_server/queue_manager/filesystem.py index bf105c18..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,11 @@ 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 @@ -207,33 +257,67 @@ def _compact_tmp_path(self, session_id: str) -> Path: """ return self._dir / f"{session_id}.log.compact.tmp" + 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. - Writes ``{"v": 1, "offset": offset, "cursor": cursor}`` to a 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. + 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._dir / f"{session_id}.offset.tmp" + tmp = self._offset_tmp_path(session_id) record = {"v": 1, "offset": offset, "cursor": cursor} - 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 + 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)``. @@ -271,6 +355,23 @@ 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: """Byte position after the last complete line in ``path`` (0 if none). @@ -366,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: @@ -557,41 +681,54 @@ 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) - # 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. Clear the flag - # unconditionally: _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. - reconcile = session_id in self._dead_unreconciled - - def _read() -> Batch: - if reconcile: - self._reconcile_dead_key(session_id) - self._dead_unreconciled.discard(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) - return await asyncio.to_thread(_read) + 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) + + 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 @@ -671,6 +808,43 @@ async def dead_letter(self, session_id: str, raw: bytes, error: str) -> None: # 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. @@ -679,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 @@ -705,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", @@ -730,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( @@ -896,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) @@ -1325,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 @@ -1498,9 +1719,14 @@ async def is_fully_drained(self, session_id: str) -> bool: self._validate_session_id(session_id) def _check() -> bool: - return self._read_committed_offset(session_id) >= self._complete_data_end( - session_id - ) + 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) @@ -1937,6 +2163,17 @@ def _reconcile_dead_key(self, key: str) -> int: 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 @@ -1948,22 +2185,23 @@ def _reconcile_dead_key(self, key: str) -> int: dead_payloads = self._dead_payload_set(key) if not dead_payloads: return 0 - 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) + 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) diff --git a/context_intelligence_server/registry.py b/context_intelligence_server/registry.py index 509c8e18..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 @@ -368,26 +374,7 @@ async def drain_worker( while True: try: - try: - batch = await qm.read_batch(session_id, max_items=_DRAIN_MAX_BATCH) - except ValueError: - # A ValueError here means the committed offset read parsed a - # genuinely corrupt/unparseable .offset -- a PERMANENT fault, - # raised before the guarded cursor read below is reached. - # Quarantine this one session (close + deregister, exit - # cleanly) so it does not crash-loop (respawn -> same corrupt - # read -> die, forever). A transient OSError is NOT caught - # here: it propagates to _on_drain_done for supervised - # respawn, exactly as before. Other sessions keep draining; - # the boot RESET_OFFSET pass heals the offset (cursor-preserving). - logger.exception( - "drain_worker_quarantined session=%s reason=corrupt_offset", - session_id, - extra={"session_id": session_id}, - ) - await self._safe_close(worker) - self._deregister(session_id) - return + batch = await self._read_batch(worker, _DRAIN_MAX_BATCH) if not batch.records: # Idle compaction runs before the dry-exit check below, so a @@ -400,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) @@ -560,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( @@ -702,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: @@ -799,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/pyproject.toml b/pyproject.toml index ed7bb43b..7e771dfc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "context-intelligence-server" -version = "6.8.0" +version = "7.0.0" description = "Context Intelligence Server for Amplifier" requires-python = ">=3.11" dependencies = [ diff --git a/tests/test_ingest_cursor_hardening.py b/tests/test_ingest_cursor_hardening.py index ceac2908..7456ad82 100644 --- a/tests/test_ingest_cursor_hardening.py +++ b/tests/test_ingest_cursor_hardening.py @@ -11,8 +11,13 @@ 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, ) @@ -20,7 +25,11 @@ OrchestratorRunHandler, ) from context_intelligence_server.queue_manager import FileSystemQueueManager -from context_intelligence_server.registry import SessionRegistry, SessionWorker +from context_intelligence_server.registry import ( + SessionRegistry, + SessionWorker, + _SessionQuarantined, +) from context_intelligence_server.services import HookStateService @@ -420,3 +429,376 @@ async def _handler(w, event, data, handlers): 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_queue_manager.py b/tests/test_queue_manager.py index 28a00ee0..a488c4e7 100644 --- a/tests/test_queue_manager.py +++ b/tests/test_queue_manager.py @@ -412,7 +412,10 @@ async def test_delete_drained_removes_log_and_offset_keeps_dead(tmp_path) -> Non 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 diff --git a/uv.lock b/uv.lock index 705457a3..be93615c 100644 --- a/uv.lock +++ b/uv.lock @@ -233,7 +233,7 @@ wheels = [ [[package]] name = "context-intelligence-server" -version = "6.8.0" +version = "7.0.0" source = { editable = "." } dependencies = [ { name = "aiofiles" },