diff --git a/context_intelligence_server/handlers/data_layer_2/session.py b/context_intelligence_server/handlers/data_layer_2/session.py index a39118f1..cf2fb7bc 100644 --- a/context_intelligence_server/handlers/data_layer_2/session.py +++ b/context_intelligence_server/handlers/data_layer_2/session.py @@ -301,32 +301,34 @@ async def _handle_fork( async def _handle_end( self, session_id: str, timestamp: str, data: dict[str, Any] ) -> None: - # Read the session's current labels BEFORE writing the end-event upsert. - # After a flush (the drainer flushes between event batches) the node - # buffer is empty, so get_node falls through to Neo4j and returns the - # real persisted type label (SubSession / ForkedSession). If we upsert - # first, that upsert creates a fresh buffer entry holding only - # ["Session", "SST_EVENT"], which SHADOWS the persisted type on the - # buffer-first get_node read -> _current_type reads None -> stub-recovery - # spuriously adds RootSession (a dual terminal label). Reading first - # mirrors _handle_start and _handle_fork, which both read before writing. + """Terminal handler. Runs TWICE per session and does NOT flush. + + The drainer leaves the ``session:end`` record uncommitted so that + "ended but not finalized" survives a respawn, then re-dispatches it + during finalization (see ``SessionRegistry._process_batch``). Every + write below is therefore a read-then-MERGE and must stay idempotent. + + There is also no ``graph.flush()`` here: the drainer's + ``_flush_barrier`` is the single write boundary, and it is the only + thing holding the Neo4j write semaphore. Flushing from a handler would + write outside that cap. + """ + # Read labels BEFORE the end-event upsert -- upserting first would + # shadow the persisted type label and spuriously trigger stub-recovery. existing = await self.services.graph.get_node(session_id) labels: list[str] = existing.get("labels", []) if existing else [] _warn_if_dual_terminal(labels, session_id) parent_id = _parent_of(data) end_node_data: dict[str, Any] = { - "labels": ["Session", "SST_EVENT"], + # Seed with labels just read so this entry can't shed a + # persisted terminal type and trigger spurious stub-recovery. + "labels": ["Session", "SST_EVENT", *labels], "ended_at": timestamp, "status": "completed", "session_id": session_id, } - # Persist parent_id when the end payload carries one. _handle_start and - # _handle_fork already write parent_id, but a session that reaches - # session:end WITHOUT a captured start/fork previously never had - # parent_id recorded at all — leaving `parent_id IS NULL` ambiguous - # between "genuinely no parent" and "parent never recorded". Only write - # when present in the payload; never fabricate a value when absent. + # Only write parent_id when present; never fabricate a value. if parent_id: end_node_data["parent_id"] = parent_id @@ -358,14 +360,10 @@ async def _handle_end( add_labels=transition.add, ) - # Terminal event: flush directly. There is no hot path after session:end; - # all buffered data must reach the backing store before the process exits. - await self.services.graph.flush() - async def _create_mount_plan( self, session_id: str, data_layer_1_fork_node_id: str ) -> None: - """E04: Session → MountPlan (record existence, no blob dereferencing). + """Session → MountPlan (record existence, no blob dereferencing). SOURCED_FROM: MountPlan → session:fork data_layer_1 event (blob source). """ mount_plan_id = f"{session_id}::mount_plan" diff --git a/context_intelligence_server/main.py b/context_intelligence_server/main.py index 0bf0709f..f172a845 100644 --- a/context_intelligence_server/main.py +++ b/context_intelligence_server/main.py @@ -71,12 +71,12 @@ def build_neo4j_driver(config: Neo4jClientConfig) -> Any: # --------------------------------------------------------------------------- -# Module-level live identity-map stores (T3) +# Module-level live identity-map stores. # -# Set by create_asgi_app() so the future /admin router can mutate the active -# store without needing to carry a reference through the middleware chain. -# Exactly ONE of these is non-None at any time — whichever mode is active. -# The other is always reset to None so accessors return an unambiguous result. +# Set by create_asgi_app() so the /admin router can mutate the active store +# without carrying a reference through the middleware chain. Exactly ONE of +# these is non-None at any time — whichever mode is active. The other is +# always reset to None so accessors return an unambiguous result. # --------------------------------------------------------------------------- _api_key_store: IdentityStore | None = None _entra_identity_store: IdentityStore | None = None @@ -102,6 +102,25 @@ def get_entra_identity_store() -> IdentityStore | None: return _entra_identity_store +def _parse_workspace_and_creator(raw: str | bytes) -> tuple[str, str | None] | None: + """Return ``(workspace, created_by)`` iff ``raw`` parses to a dict with a + non-empty workspace, else ``None``. Total -- never raises.""" + try: + obj = json.loads(raw) + except (ValueError, TypeError): + return None + if not isinstance(obj, dict): + return None + try: + workspace = obj.get("workspace", "") + created_by = obj.get("created_by") + except (AttributeError, TypeError): + return None + if not workspace: + return None + return workspace, created_by + + def _recover_one_session( sid: str, first_line: str | bytes, @@ -113,32 +132,35 @@ def _recover_one_session( real parsing/dispatch logic rather than reimplementing it inline. The queue-read step is handled by the caller (the lifespan loop or the test) - so this function is pure — no I/O, fully synchronous. + so this function is pure -- no I/O, fully synchronous. + + A head line that does not resolve a workspace SKIPS the session rather than + guessing one. ``workspace`` is the graph partition key (every node is + MERGEd on ``{node_id, workspace}``), so dispatching under a substitute + would write the whole session into a partition it does not belong to and + drop its contributor -- worse, and harder to undo, than leaving the data + durable on disk for an operator. The log is untouched and a later boot + reports the session again. Args: sid: Session id being recovered. first_line: The first raw log line (bytes from QueueManager or str from tests). ``json.loads`` accepts both. - get_or_create: The registry callable — ``registry.get_or_create`` in + get_or_create: The registry callable -- ``registry.get_or_create`` in production or a spy in tests. Returns: - True – drainer was (re)spawned via *get_or_create*. - False – session skipped (empty/torn workspace, or malformed JSON line). + True - drainer was (re)spawned via *get_or_create*. + False - session skipped (empty/torn workspace, or malformed JSON line). """ - try: - obj = json.loads(first_line) - workspace: str = obj.get("workspace", "") - created_by: str | None = obj.get("created_by") - except (ValueError, KeyError): - workspace = "" - created_by = None - if not workspace: + parsed = _parse_workspace_and_creator(first_line) + if parsed is None: logger.warning( "recovery_skipped session=%s: torn or empty workspace in first line", sid, ) return False + workspace, created_by = parsed get_or_create(sid, workspace, created_by=created_by) return True @@ -164,8 +186,20 @@ async def _crash_recovery_topup(respawn_limit: int | None) -> int: to_process = recovered if respawn_limit is None else recovered[:respawn_limit] respawned = 0 for sid in to_process: - batch = await registry.queue_manager.read_batch(sid, max_items=1) + # Guarded here (not inside read_batch, which must stay loud for the + # live drainer's hot path) so one bad key can't halt the whole pass. + try: + batch = await registry.queue_manager.read_batch(sid, max_items=1) + except (OSError, ValueError): + logger.exception("crash_recovery_topup_read_failed session=%s", sid) + continue if not batch.lines: + # recover()/read_batch disagreement (e.g. a concurrent finalize + # advanced the offset) -- not a loss, just no longer recoverable. + logger.warning( + "recovery_skipped_empty_batch session=%s reason=empty_batch", + sid, + ) continue # NOTE: _recover_one_session returns True whenever it dispatched to # get_or_create, whether or not a drainer already existed (get_or_create @@ -275,7 +309,7 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: "(un-migrated). Cold start refuses to boot to avoid duplicating " "them on write. Run: context-intelligence-server doctor --fix" ) - # Crash recovery (decisions #5/#6): on startup, respawn one drainer per + # Crash recovery: on startup, respawn one drainer per # session that still has an undrained, complete line. The workspace is # parsed from that session's FIRST log line so the respawned worker is # bound to the same workspace it was originally created with. @@ -307,6 +341,16 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: # call reports them again, and a new event for that session arriving # via POST /events spawns its drainer immediately via get_or_create(), # independent of this startup loop. + # Bound how many drainers this boot respawns: an unbounded backlog can + # respawn every drainer before the server serves a single request, + # driving startup RSS and boot time up with it. None (the default) + # preserves unbounded behaviour. `recovered` is already sorted + # (QueueManager.recover()), so which sessions run this boot vs. defer is + # deterministic across restarts of the same backlog. Deferred sessions + # are untouched -- no read, no write, no drainer -- so they remain fully + # durable and recoverable; a later boot's recover() reports them again, + # and a new event for that session spawns its drainer immediately via + # get_or_create(), independent of this startup loop. respawn_limit = _settings.crash_recovery_respawn_limit if respawn_limit is not None and len(recovered) > respawn_limit: to_process = recovered[:respawn_limit] @@ -322,10 +366,9 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: if _recover_one_session(sid, batch.lines[0], registry.get_or_create): respawned += 1 if deferred_count: - # Loud on purpose (WARNING, not INFO): a deferred backlog must never - # be a silent, un-discoverable fact -- that silence is exactly what - # let the 38 GB spool go unnoticed for two days in the incident this - # guards against. Names the exact counts and the setting to raise. + # WARNING, not INFO, on purpose: a deferred backlog must never be a + # silent, un-discoverable fact. Names the exact counts and the + # setting to raise. logger.warning( "lifespan_startup: crash-recovery respawn cap reached " "(crash_recovery_respawn_limit=%d): %d/%d respawned this boot, " @@ -396,7 +439,7 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: idempotency_cache = EventIdempotencyCache() # Session-less events are keyed by a per-workspace sentinel stem so that events -# from distinct workspaces never collide in one durable log (decision #10). +# from distinct workspaces never collide in one durable log. _NO_SESSION_PREFIX = "_no_session__" @@ -409,7 +452,7 @@ def _workspace_slug(workspace: str) -> str: def _validate_data_timestamp(data: dict[str, Any]) -> None: """Raise HTTPException(400) if data['timestamp'] is missing, empty, or not ISO-8601. - This is the ingest boundary check (Option A). Real Amplifier clients always + This is the ingest boundary check. Real Amplifier clients always supply data.timestamp (verified: 224,530 events on disk, 0 missing). This guard rejects only malformed/hand-rolled payloads with a clear 400, instead of accepting them silently and dead-lettering them later when the graph @@ -431,7 +474,7 @@ def _validate_data_timestamp(data: dict[str, Any]) -> None: def _assert_admin_not_exempt() -> None: - """Startup assertion (TB-07): /admin/* must NEVER be in any exempt set. + """Startup assertion: /admin/* must NEVER be in any exempt set. Called by ``create_asgi_app`` before constructing the middleware. Raises ``RuntimeError`` if any ``/admin`` path or prefix appears in @@ -463,7 +506,7 @@ def _assert_admin_not_exempt() -> None: def _assert_neo4j_clients_explicit(settings: Settings) -> None: - """Startup assertion (doc 11 gap #12): the deployed profile MUST declare the + """Startup assertion: the deployed profile MUST declare the structured neo4j.admin / neo4j.cypher_query clients explicitly. When settings.neo4j_require_explicit_clients is True, refuse to boot if the @@ -477,7 +520,7 @@ def _assert_neo4j_clients_explicit(settings: Settings) -> None: "Neo4j config invariant violated: neo4j_require_explicit_clients=True but " "the structured `neo4j` block (admin + cypher_query) is absent — the server " "would silently fall back to legacy neo4j_* fields. The deployed profile MUST " - "declare both clients explicitly (doc 11 §Backward-compatibility). Set the " + "declare both clients explicitly. Set the " "`neo4j` block in amplifier-online.yaml / server-config.yaml, or unset " "neo4j_require_explicit_clients for a dev/transition deploy." ) @@ -503,8 +546,8 @@ def create_asgi_app( resolver builds a real ``PyJWKClient`` internally. Startup behavior on an EMPTY store: - An empty keystore (static) or empty identity map (entra) NO LONGER - raises — it is a supported bootstrap state. The server BOOTS + An empty keystore (static) or empty identity map (entra) does NOT + raise — it is a supported bootstrap state. The server BOOTS fail-CLOSED and logs a loud startup WARNING; every request 401/403s until the store is populated at runtime via the /admin API. Wide-open pass-through is reachable ONLY via the explicit @@ -512,14 +555,14 @@ def create_asgi_app( credentials configured, which additionally logs a "WIDE OPEN" warning. Raises: - RuntimeError: (TB-07) When any ``/admin`` path or prefix appears in an + RuntimeError: When any ``/admin`` path or prefix appears in an auth-exempt set. The admin API surface must never be unguarded. """ global _api_key_store, _entra_identity_store - # TB-07 structural assertion: /admin must not be in any exempt set. - # This runs before any middleware construction so the failure is loud and - # immediate — no request ever reaches an unauthenticated /admin endpoint. + # Structural assertion: /admin must not be in any exempt set. Runs before + # any middleware construction so the failure is loud and immediate — no + # request ever reaches an unauthenticated /admin endpoint. _assert_admin_not_exempt() s = settings if settings is not None else _settings @@ -534,21 +577,21 @@ def create_asgi_app( app.state.api_key_store = None app.state.entra_identity_store = None - # T5: store auth/admin config on app.state so the require_admin dependency + # Store auth/admin config on app.state so the require_admin dependency # can read it without importing from main (avoids circular import) and so # test-specific settings (passed via create_asgi_app(settings=...)) take # effect without relying on the module-level cached get_settings(). app.state.auth_mode = s.auth_mode app.state.admin_api_key_configured = s.resolve_admin_api_key_digest() is not None app.state.entra_admin_role = s.entra_admin_role - # M2: service capability role names for require_write / require_read deps. + # Service capability role names for require_write / require_read deps. app.state.service_data_role = s.service_data_role app.state.reader_role = s.reader_role # Compute the admin-key digest for the middleware (static mode only). # The middleware checks the bearer token's sha256 against this digest BEFORE # calling the resolver, so the admin key can authenticate even though it is - # not in the data keystore (ROB F1). + # not in the data keystore. # # Storage-at-rest is resolved by Settings: the RECOMMENDED admin_api_key_sha256 # (digest at rest) is used verbatim; the legacy raw admin_api_key (DEPRECATED, @@ -600,36 +643,36 @@ def create_asgi_app( s.entra_identities_store_path, ) - # B4: boot disjointness invariant — each oid must belong to exactly one + # Boot disjointness invariant — each oid must belong to exactly one # identity source. Building the service map here (not inline in the # EntraResolver call) lets us check the overlap BEFORE construction so # the server fails loudly at startup rather than silently misbehaving. - # This is cheap hygiene: B1 already keeps app tokens off the human map - # at request time; this prevents a same-oid-in-both misconfiguration. + # Existing logic already keeps app tokens off the human map at + # request time; this prevents a same-oid-in-both misconfiguration. _service_id_map = s.build_service_identity_map() _entra_oids = set(entra_store.flat_dict.keys()) _service_oids = set(_service_id_map.keys()) _overlap = _entra_oids & _service_oids if _overlap: raise RuntimeError( - f"Boot invariant violated (B4): oid(s) {sorted(_overlap)!r} appear " + f"Boot invariant violated: oid(s) {sorted(_overlap)!r} appear " f"in both entra_identities and service_identities. Each oid must " f"belong to exactly one identity source. Fix the config to remove " f"the overlap before restarting." ) # EntraResolver raises RuntimeError at construction if the JWKS - # prefetch fails (eager fail-closed guard from §8b / crusty gate). + # prefetch fails (eager fail-closed by design). # Pass entra_store.flat_dict (the LIVE dict) so the resolver sees # any put()/delete() made by /admin immediately, no restart required. resolver: StaticKeyResolver | EntraResolver = EntraResolver( s.azure_client_id, # type: ignore[arg-type] — validated non-None by config s.azure_tenant_id, # type: ignore[arg-type] — validated non-None by config entra_store.flat_dict, # live reference — mutations visible immediately - service_identity_map=_service_id_map, # B4: pre-built, disjointness verified - service_data_role=s.service_data_role, # M2: role gate - reader_role=s.reader_role, # M2: role gate - entra_admin_role=s.entra_admin_role, # M2: role gate + service_identity_map=_service_id_map, # pre-built, disjointness verified + service_data_role=s.service_data_role, + reader_role=s.reader_role, + entra_admin_role=s.entra_admin_role, jwks_client=_jwks_client, ) # Entra mode does not use admin_api_key_digest (admin via roles claim). @@ -680,8 +723,8 @@ def create_asgi_app( # Wide-open warning: fires ONLY on the explicit allow_unauthenticated # opt-out combined with no credentials configured. An empty keystore/map - # ALONE no longer triggers this (and no longer refuses to start) — it now - # boots fail-closed instead (see the empty-map/keystore warnings above). + # ALONE does not trigger this and does not refuse to start — it boots + # fail-closed instead (see the empty-map/keystore warnings above). if s.allow_unauthenticated and not resolver.auth_enabled: logger.warning( "allow_unauthenticated=True AND no credentials configured — the " @@ -691,7 +734,7 @@ def create_asgi_app( "(entra) and unset allow_unauthenticated to enforce authentication." ) - # Log admin capability status for operator visibility (E: status surfacing). + # Log admin capability status for operator visibility. if s.auth_mode == "static": _admin_status = ( "enabled" @@ -710,11 +753,11 @@ def create_asgi_app( _admin_status, ) - # T6: store the admin-key digest on app.state so the /admin router handlers + # Store the admin-key digest on app.state so the /admin router handlers # can read it without importing from main (no circular import) and so that # test-specific settings are honoured. In entra mode admin_api_key_digest - # has already been set to None above (line ~385); in static mode it is the - # sha256 of admin_api_key (or None when admin_api_key is not configured). + # has already been set to None above; in static mode it is the sha256 of + # admin_api_key (or None when admin_api_key is not configured). app.state.admin_api_key_digest = admin_api_key_digest return BearerTokenMiddleware( @@ -729,27 +772,21 @@ def create_asgi_app( # Module-level ASGI app used by Gunicorn: context_intelligence_server.main:asgi_app # The raw `app` is kept for internal use and testing against un-authed routes. # -# LAZY construction (PEP 562 module __getattr__), NOT built at import time. -# -# create_asgi_app() enforces the auth guard: it raises RuntimeError when no -# authentication is configured at all (see its docstring / _assert_* helpers). -# That guard is correct and must NOT be weakened. The problem was *timing*: -# this module used to call create_asgi_app() unconditionally at import time, -# which meant the console-script entry point (`context-intelligence-server`) -# imports `main` to reach `main()`, so even `--help`/`--version` constructed -# the whole ASGI app and hit the guard. An operator with a broken/absent -# config couldn't ask the binary what version it was -- exactly when they -# most need to. +# Lazily constructed (PEP 562 module __getattr__), not built at import time: +# create_asgi_app() enforces the auth guard (raises RuntimeError when no +# authentication is configured; see its docstring / _assert_* helpers). The +# console-script entry point imports this module just to reach `main()`, so +# eager construction would make even `--help`/`--version` hit that guard -- +# an operator with a broken/absent config couldn't ask the binary its +# version, exactly when they most need to. # # `_asgi_app` is the cache; `get_asgi_app()` builds-and-caches on first call; -# `__getattr__` makes `context_intelligence_server.main.asgi_app` / -# `from context_intelligence_server.main import asgi_app` keep working for +# `__getattr__` keeps `context_intelligence_server.main.asgi_app` / +# `from context_intelligence_server.main import asgi_app` working for # anything that reads the module attribute directly (gunicorn's `load()`, -# tests) -- construction (and therefore the auth guard) now happens on first -# access instead of at import time. Actually serving (`run()` -> `_App.load()` -# -> `get_asgi_app()`) still triggers it, so an unconfigured server still -# fails loud exactly as before -- only bare import / --help / --version are -# spared. +# tests). Actually serving (`run()` -> `_App.load()` -> `get_asgi_app()`) +# still triggers construction, so an unconfigured server still fails loud -- +# only bare import / --help / --version are spared. _asgi_app: BearerTokenMiddleware | None = None @@ -781,7 +818,7 @@ def __getattr__(name: str) -> Any: # --------------------------------------------------------------------------- -# M2 — service capability dependencies (moved to authz.py to avoid circular import) +# Service capability dependencies (moved to authz.py to avoid circular import) # # require_write, require_read, _is_write_capable are imported from # context_intelligence_server.authz at the top of this file (re-exported here @@ -795,28 +832,25 @@ async def get_status(request: Request) -> dict[str, Any]: response["neo4j_connected"] = await _check_driver_connected( request.app, "neo4j_driver" ) - # Additive (Concern B, council review): surface the query (read-intent) - # driver's connectivity too, so a misconfigured cypher_query client shows - # up here instead of on the first /cypher call. + # Also surface the query (read-intent) driver's connectivity, so a + # misconfigured cypher_query client shows up here instead of on the + # first /cypher call. response["neo4j_query_connected"] = await _check_driver_connected( request.app, "neo4j_query_driver" ) response["neo4j_url"] = _settings.resolve_neo4j_admin().url response["neo4j_browser_url"] = _settings.neo4j_browser_url - # Additive, aggregate-only conservation metrics (D3). /status is - # unauthenticated, so this block must NOT carry the per-key table or the - # dead-letter listing — both are authenticated-only. + # Aggregate-only conservation metrics. /status is unauthenticated, so + # this block must NOT carry the per-key table or the dead-letter + # listing — both are authenticated-only. response["metrics"] = await registry.pipeline_metrics() - # Additive, aggregate-only spool footprint (incident: a 38 GB / 583-file - # durable spool grew completely unnoticed -- the only symptom was a graph - # that had silently stopped updating). Same /status contract as `metrics` - # above: two aggregate integers only, no session ids, no workspace names, - # no per-key table. Cheap by construction (stat-only, short-TTL cached) -- - # see QueueManager.spool_stats() for why this is safe on every poll even - # with a huge spool. + # Aggregate-only spool footprint: two integers only, no session ids, no + # workspace names, no per-key table -- same privacy contract as `metrics` + # above. Cheap by construction (stat-only, short-TTL cached); see + # QueueManager.spool_stats() for why this holds even under a huge spool. response["spool"] = await registry.queue_manager.spool_stats() - # T5 (E): surface auth mode and admin-API capability so operators can - # confirm admin is enabled without tailing startup logs. /status is + # Surface auth mode and admin-API capability so operators can confirm + # admin is enabled without tailing startup logs. /status is # unauthenticated — only config-level boolean flags are exposed here # (no credential values, no key hashes, no token details). _auth_mode = getattr(request.app.state, "auth_mode", _settings.auth_mode) @@ -834,9 +868,7 @@ async def get_status(request: Request) -> dict[str, Any]: _admin_key_set if _auth_mode == "static" else bool(_entra_admin_role) ), # Surface the role names (not secrets) so operators can confirm which - # roles are configured without exposing credential values. Additive: - # existing fields (mode, admin_api_enabled, entra_admin_role) are - # unchanged; reader_role and service_data_role are new in M2. + # roles are configured without exposing credential values. **( { "entra_admin_role": _entra_admin_role, @@ -901,7 +933,7 @@ async def post_events( ) return EventResponse(status="duplicate", session_id=session_id or None) # Empty session_id maps to a per-workspace sentinel stem so session-less - # events from distinct workspaces never collide in one log (decision #10). + # events from distinct workspaces never collide in one log. worker_key = session_id or (_NO_SESSION_PREFIX + _workspace_slug(request.workspace)) # Spawn (or reuse) the sticky drainer keyed by worker_key. registry.get_or_create(worker_key, request.workspace, created_by=contributor_id) @@ -970,8 +1002,7 @@ def main(argv: list[str] | None = None) -> None: ``doctor [--fix]`` diagnoses (and, with ``--fix``, repairs) Neo4j graph health -- the two O(graph-size) migration scans (dedup + :Node backfill) - that used to run unconditionally at cold start now live ONLY here, never - on server boot. See ``context_intelligence_server.doctor``. + live ONLY here, never on server boot. See ``context_intelligence_server.doctor``. """ parser = argparse.ArgumentParser(prog="context-intelligence-server") subparsers = parser.add_subparsers(dest="command") @@ -1039,7 +1070,7 @@ def _validate_single_worker(workers: int | None = None) -> int: f"context-intelligence-server requires exactly one worker, got {effective}. " "The durable drainer assumes one drainer per session per process; unset " "WEB_CONCURRENCY or set WEB_CONCURRENCY=1. Multi-process operation needs a " - "distributed backend (Open Q7)." + "distributed backend." ) return effective diff --git a/context_intelligence_server/queue_manager.py b/context_intelligence_server/queue_manager.py index 6fba186e..a52eca18 100644 --- a/context_intelligence_server/queue_manager.py +++ b/context_intelligence_server/queue_manager.py @@ -1,37 +1,35 @@ -"""On-disk durable queue manager for the event-write pipeline. - -Disk layout (one set of files per session, keyed by ``session_id``): - -- ``.log`` — append-only, newline-terminated, opaque ``bytes``. - Each line is one enqueued record. The log is never rewritten in place. -- ``.offset`` — a single integer: the byte position in the log - that has been durably processed (committed). A missing offset file means 0. -- ``.dead.jsonl`` — append-only dead-letter records for batches - that could not be processed after exhausting retries. - -Durability note: - Appends use a plain durable ``write()``. This gives PROCESS-crash - durability (the bytes are handed to the OS page cache and survive a - process crash). POWER-LOSS durability via ``fsync`` is deliberately - deferred to Phase B3 (fsync group-commit). - -session_id contract: - Every public method validates ``session_id`` and raises ``ValueError`` if - it is empty or contains a path separator (``/`` or ``\\``) or a null byte. - The ``session_id`` is used raw as the filename stem, so it must be a safe, - single path component. +"""On-disk durable queue for the event-write pipeline. + +Per session ````: ``.log`` (append-only, ``\\n``-terminated records), +``.offset`` (committed byte position, missing == 0), ``.dead.jsonl`` (dead +letters). + +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. """ from __future__ import annotations import asyncio import base64 +import contextlib import json +import logging import os +import threading import time +from collections.abc import Coroutine, Iterator from dataclasses import dataclass from pathlib import Path -from typing import Any +from typing import Any, TypeVar + +logger = logging.getLogger(__name__) + +_T = TypeVar("_T") # 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) @@ -41,24 +39,97 @@ _SCAN_CHUNK_BYTES = 1 << 20 +@dataclass(frozen=True) +class Record: + """One log record and the byte range the QUEUE assigned it. + + ``start``/``end`` are opaque cursor values PRODUCED BY THE QUEUE and only + ever handed back to it (``commit``). Callers MUST NOT compute them and + MUST NOT assume ``end - start == len(raw) + 1`` -- that relationship is + the queue's private framing invariant (module docstring), not a public + contract. + """ + + raw: bytes # WITHOUT the terminator, exactly as ``lines`` is today + start: int + end: int + + @dataclass(frozen=True) class Batch: - """A contiguous batch of log lines read from a session's append-only log. + """A contiguous batch of log records read from a session's append-only log. Attributes: - session_id: The session the lines belong to. - lines: Raw, complete log lines WITHOUT their trailing newline. + session_id: The session the records belong to. + records: Queue-produced ``Record``s -- each carries its own opaque + ``start``/``end`` cursor. The queue produces these offsets; a + caller (the registry) only ever hands them back via ``commit``. start_offset: Byte position in the log where this batch begins. - end_offset: Byte position in the log AFTER the last returned line. - This is the value passed to ``commit``. When no complete lines + end_offset: Byte position in the log AFTER the last returned record. + This is the value passed to ``commit``. When no complete records are available, ``end_offset == start_offset``. """ session_id: str - lines: list[bytes] + records: list[Record] start_offset: int end_offset: int + @property + def lines(self) -> list[bytes]: + """Raw record payloads, terminator-stripped -- the pre-Record view. + + Derived from ``records`` so the two can never disagree. Retained + because ~90 call sites across main.py and 12 test files read it. + """ + return [r.raw for r in self.records] + + +@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. + ``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. + ``delete_drained`` refuses to drop the guard while any remain, else two + coroutines could lock the same file under different guards and tear it. + """ + + admission: asyncio.Lock + file_lock: threading.Lock + waiters: int = 0 + + +async def _await_uninterrupted(coro: Coroutine[Any, Any, _T]) -> _T: + """Await ``coro`` to completion even if this coroutine is cancelled. + + ``asyncio.to_thread`` cannot interrupt the OS thread it dispatched, so a + cancellation is absorbed and re-raised only once the write has definitively + succeeded or failed -- otherwise ``append`` would return with bytes still in + flight. This is resource hygiene, not the framing guarantee (that is + ``_KeyGuard.file_lock``). + """ + task = asyncio.ensure_future(coro) + cancelled: asyncio.CancelledError | None = None + while True: + try: + result = await asyncio.shield(task) + break + except asyncio.CancelledError as exc: + if task.done(): + raise # the TASK was cancelled, not us + cancelled = exc # ours: remember it, keep waiting + except BaseException: + if cancelled is not None: + raise cancelled from None # teardown wins over the write's error + raise + if cancelled is not None: + raise cancelled + return result + class QueueManager: """Manages per-session append-only queues on disk.""" @@ -69,8 +140,8 @@ def __init__(self, queues_dir: Path): self._stats_cache: dict[str, Any] | None = None self._stats_cache_at: float = 0.0 self._stats_cache_ttl: float = 1.0 - # Separate cache for spool_stats() (Change 2 / /status spool block). - # A longer TTL than _stats_cache_ttl is fine here: spool_stats() is an + # Separate cache for spool_stats(). A longer TTL than _stats_cache_ttl + # is fine here: spool_stats() is an # operator-facing "is the backlog growing" signal, not a # correctness-sensitive value, so a few extra seconds of staleness is # an acceptable trade for fewer directory scans under frequent @@ -78,6 +149,12 @@ def __init__(self, queues_dir: Path): self._spool_cache: dict[str, int] | None = None self._spool_cache_at: float = 0.0 self._spool_cache_ttl: float = 5.0 + # One _KeyGuard per worker key that has been appended to and + # not yet finalized-and-deleted. Created lazily by _guard(); removed + # ONLY by delete_drained, under the admission lock, gated on identity + # 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] = {} def _log_path(self, session_id: str) -> Path: return self._dir / f"{session_id}.log" @@ -89,6 +166,7 @@ def _dead_path(self, session_id: str) -> Path: return self._dir / f"{session_id}.dead.jsonl" def _read_committed_offset(self, session_id: str) -> int: + """Committed byte offset. A missing or empty ``.offset`` reads 0.""" try: text = self._offset_path(session_id).read_text("utf-8") except FileNotFoundError: @@ -96,20 +174,19 @@ def _read_committed_offset(self, session_id: str) -> int: text = text.strip() return int(text) if text else 0 - def _complete_data_end(self, session_id: str) -> int: - """Byte position after the last complete (newline-terminated) line. + @staticmethod + def _last_complete_end(path: Path) -> int: + """Byte position after the last complete line in ``path`` (0 if none). - A torn trailing line (bytes after the final newline) is ignored: the - returned offset is one past the last ``\\n``, or 0 when the log is - missing or contains no complete line. + A torn trailing fragment (bytes after the final newline) is ignored: + the returned offset is one past the last ``\\n``, or 0 when the file + is missing or contains no complete line. Streams BACKWARD from EOF in fixed chunks to find the last ``\\n`` -- - O(tail) memory and I/O, never O(file). This log can be multi-GB (the - durable spool grew to a 4.9 GB single file in the incident); reading - the whole thing into RAM just to find the final newline is exactly the - boot-time memory blowup this avoids. + O(tail) memory and I/O, never O(file). Path-based (not + session-id-based) so it serves both ``.log`` files (via + ``_complete_data_end``) and ``.dead.jsonl`` files. """ - path = self._log_path(session_id) try: with open(path, "rb") as f: f.seek(0, os.SEEK_END) @@ -126,17 +203,19 @@ def _complete_data_end(self, session_id: str) -> int: except FileNotFoundError: return 0 + def _complete_data_end(self, session_id: str) -> int: + """Byte position after the last complete line of a session's ``.log``.""" + return self._last_complete_end(self._log_path(session_id)) + @staticmethod def _stream_newlines(path: Path, start: int = 0, end: int | None = None) -> int: """Count ``\\n`` bytes in ``path``'s byte range ``[start, end)`` -- streamed. - ``end=None`` counts to EOF. Reads the range in fixed-size chunks - (O(chunk) memory) instead of materialising the whole file (or a slice - copy of it) in RAM, which is what ``read_bytes()`` + - ``data[a:b].count(b"\\n")`` did on multi-GB spool files. Numerically - identical to that slice-count for any range; a missing file counts 0. + ``end=None`` counts to EOF. Reads in fixed-size chunks (O(chunk) + memory, never O(file)); numerically identical to a full slice-count + for any range. A missing file counts 0. - Path-based (not session-id-based) so it serves both ``.log`` scans + Path-based (not session-id-based): serves both ``.log`` scans (``_count_newlines``) and the whole-file ``.dead.jsonl`` count (``_count_dead``). """ @@ -181,16 +260,101 @@ def _validate_session_id(session_id: str) -> None: ): raise ValueError(f"Invalid session_id: {session_id!r}") + @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. + """ + guard = self._guards.get(worker_key) + if guard is None: + guard = _KeyGuard(asyncio.Lock(), threading.Lock()) + self._guards[worker_key] = guard + guard.waiters += 1 + try: + yield guard + finally: + guard.waiters -= 1 + + @staticmethod + def _write_all(fd: int, data: bytes) -> None: + """Write ALL of ``data`` to ``fd``, looping over short writes. + + ``os.write`` may write fewer bytes than requested -- which is + precisely what a network filesystem does with a multi-hundred-KB + buffer -- so one call is an ATTEMPT, not a write. This loop is the + code taking responsibility for what the storage layer does not + promise. + """ + view = memoryview(data) + written = 0 + while written < len(view): + n = os.write(fd, view[written:]) + if n == 0: # never observed, but a 0 would spin forever + raise OSError("os.write returned 0; refusing to spin") + written += n + + @staticmethod + def _discard_partial(fd: int, start: int, path: Path) -> None: + """Newline-terminate a partial write; never truncates -- queue bytes are never removed.""" + try: + QueueManager._write_all(fd, b"\n") + except OSError: + logger.exception( + "append_partial_terminate_failed path=%s start=%d " + "(torn fragment left unterminated; readers skip it, and the " + "next append merges it into one poison line that the drainer " + "dead-letters -- bytes are never silently removed)", + path, + start, + ) + + def _write_record(self, guard: _KeyGuard, path: Path, line: bytes) -> None: + """Append one newline-terminated record as a contiguous byte range. + + Runs in a worker thread and acquires ``guard.file_lock`` itself, so the + lock's lifetime is the thread's write, not the coroutine's await; the + caller must not hold it. ``O_APPEND`` is kept as defence in depth (it + positions every op at server-side EOF), but correctness rests on the + guard, not on its atomicity. Not ``fsync``ed. + """ + with guard.file_lock: + flags = os.O_WRONLY | os.O_CREAT | os.O_APPEND | getattr(os, "O_BINARY", 0) + fd = os.open(path, flags, 0o644) + try: + start = os.fstat(fd).st_size # sole writer: size cannot move under us + try: + self._write_all(fd, line) + except OSError: + self._discard_partial(fd, start, path) + raise + finally: + os.close(fd) + async def append(self, session_id: str, raw: bytes) -> None: + """Durably append one record to ``session_id``'s ``.log``. + + Framing invariant (module docstring) holds under any concurrency, + under cancellation, and regardless of filesystem write atomicity -- + see ``_KeyGuard.file_lock`` and ``_write_record``. + """ self._validate_session_id(session_id) line = raw if raw.endswith(b"\n") else raw + b"\n" path = self._log_path(session_id) - - def _append() -> None: - with open(path, "ab") as f: - f.write(line) - - await asyncio.to_thread(_append) + # ``_guard`` registers this coroutine as a reference-holder + # SYNCHRONOUSLY, before the first await. That ordering is + # load-bearing -- ``delete_drained`` reads ``waiters`` to + # decide whether the guard may be discarded, and a reference taken + # after an await would be invisible to it. + with self._guard(session_id) as guard: + async with guard.admission: + await _await_uninterrupted( + asyncio.to_thread(self._write_record, guard, path, line) + ) async def read_batch(self, session_id: str, max_items: int) -> Batch: self._validate_session_id(session_id) @@ -198,22 +362,24 @@ async def read_batch(self, session_id: str, max_items: int) -> Batch: def _read() -> Batch: start = self._read_committed_offset(session_id) - lines: list[bytes] = [] + records: list[Record] = [] consumed = 0 try: with open(path, "rb") as f: f.seek(start) - while len(lines) < max_items: + 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 - lines.append(raw[:-1]) + 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, lines, start, start + consumed) + return Batch(session_id, records, start, start + consumed) return await asyncio.to_thread(_read) @@ -222,8 +388,8 @@ async def commit(self, session_id: str, new_offset: int) -> None: 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 here: this gives process-crash durability, while - power-loss durability is deferred to Phase B3 (fsync group-commit). + ``fsync`` is issued: the offset survives a process crash but not a + power loss. """ self._validate_session_id(session_id) final = self._offset_path(session_id) @@ -238,14 +404,11 @@ def _commit() -> None: async def dead_letter(self, session_id: str, raw: bytes, error: str) -> None: """Append one dead-letter record for an unprocessable batch line. - The original line is stored under ``payload`` as a UTF-8 string when it - decodes cleanly; otherwise the raw bytes are stored base64-encoded under - ``payload_b64`` (so non-UTF-8 payloads are never silently dropped). Each - record also carries a ``ts`` (epoch seconds) and the ``error`` string. - - This is the dead-letter PRIMITIVE only. The poison-isolation POLICY - (deciding WHEN to dead-letter a line) is Phase B2. The main ``.log`` and - ``.offset`` files are untouched. + Stores the line under ``payload`` (UTF-8) or ``payload_b64`` (raw + bytes), plus ``ts`` and ``error``. Primitive only -- the caller decides + when to dead-letter. Guarded by the same per-key ``_KeyGuard`` as + ``append`` (an unguarded write here is as dangerous as an unguarded + ``.log`` write); the ``.log``/``.offset`` files are untouched. """ self._validate_session_id(session_id) payload = raw[:-1] if raw.endswith(b"\n") else raw @@ -256,35 +419,95 @@ async def dead_letter(self, session_id: str, raw: bytes, error: str) -> None: record["payload_b64"] = base64.b64encode(payload).decode("ascii") line = (json.dumps(record) + "\n").encode("utf-8") path = self._dead_path(session_id) - - def _append() -> None: - with open(path, "ab") as f: - f.write(line) - - await asyncio.to_thread(_append) - - async def delete_drained(self, session_id: str) -> None: - """Remove the drained .log and .offset for a fully-finalized session. - - The .dead.jsonl (if any) is intentionally KEPT — dead-letters are - retained for later inspection/replay (Phase C). Idempotent: missing - files are ignored. + with self._guard(session_id) as guard: + async with guard.admission: + try: + await _await_uninterrupted( + asyncio.to_thread(self._write_record, guard, path, line) + ) + except OSError: + # LOGGING ONLY: re-raise unchanged so propagation is + # unaffected. Without this, a failed dead-letter write + # kills the drainer and surfaces only as a generic + # drain_worker_died, with no hint that the dead-letter + # write itself was the failing operation. + logger.exception("dead_letter_write_failed session=%s", session_id) + raise + + async def delete_drained(self, session_id: str) -> bool: + """Remove the drained ``.log``/``.offset`` for a finalized session. + + Returns True if removed (or already both absent). Takes + ``guard.file_lock`` before unlinking, so it can never race an in-flight + 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. + + The guard-map entry is dropped only when ``waiters == 1`` and identity + matches; otherwise a still-referencing coroutine could later lock a + fresh guard over the same file and tear it. """ self._validate_session_id(session_id) + log = self._log_path(session_id) + offset = self._offset_path(session_id) + + def _delete(guard: _KeyGuard) -> bool: + # Under guard.file_lock: this can never run while a write thread + # for this key owns the fd. Acquired by THIS thread, not + # the coroutine -- same discipline as _write_record. + with guard.file_lock: + try: + size = log.stat().st_size + except FileNotFoundError: + # No log, but a stale .offset must not be left behind + # -- it would make a log recreated later + # start reading past its own end. + try: + offset.unlink() + except FileNotFoundError: + pass + return True + + committed = self._read_committed_offset(session_id) + if size > committed: + logger.warning( + "delete_drained_retained session=%s uncommitted_bytes=%d", + session_id, + size - committed, + ) + return False - def _delete() -> None: - for p in (self._log_path(session_id), self._offset_path(session_id)): try: - p.unlink() + log.unlink() except FileNotFoundError: pass - - await asyncio.to_thread(_delete) + try: + offset.unlink() + except FileNotFoundError: + pass + 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] + return ok 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. + 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. """ self._validate_session_id(session_id) @@ -293,7 +516,26 @@ def _read() -> list[dict]: text = self._dead_path(session_id).read_text(encoding="utf-8") except FileNotFoundError: return [] - return [json.loads(ln) for ln in text.splitlines() if ln.strip()] + records: list[dict] = [] + 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", session_id, skipped + ) + return records return await asyncio.to_thread(_read) @@ -310,8 +552,15 @@ def _scan() -> list[str]: result: list[str] = [] for log in sorted(self._dir.glob("*.log")): session_id = log.stem - if self._read_committed_offset(session_id) < log.stat().st_size: - result.append(session_id) + # Fault-isolate PER KEY -- this is reachable + # from an authenticated route (routers/queues.py), not just + # boot, so the same corrupt-.offset asymmetry the boot paths + # guard against applies here too. + try: + if self._read_committed_offset(session_id) < log.stat().st_size: + result.append(session_id) + except (OSError, ValueError): + logger.error("active_sessions_key_failed session=%s", session_id) return result return await asyncio.to_thread(_scan) @@ -334,9 +583,19 @@ def _scan() -> list[str]: result: list[str] = [] for log in sorted(self._dir.glob("*.log")): session_id = log.stem - committed = self._read_committed_offset(session_id) - if committed < self._complete_data_end(session_id): - result.append(session_id) + # Fault-isolate PER KEY -- a corrupt/unreadable + # `.offset` for one session (NUL-filled, negative, + # non-numeric) must not raise out of a boot-path scan and + # crash-loop the container. Skip just that key, log once. + try: + committed = self._read_committed_offset(session_id) + if committed < self._complete_data_end(session_id): + result.append(session_id) + except (OSError, ValueError): + # Cheap tightening: attach the traceback (was + # message-only) so a repeating corrupt-offset cause is + # visible on a boot-path scan. + logger.exception("recover_key_failed session=%s", session_id) return result return await asyncio.to_thread(_scan) @@ -372,20 +631,10 @@ def _all_worker_keys(self) -> list[str]: async def derive_all_stats(self) -> dict[str, Any]: """Derive live queue stats purely from disk, with a short TTL cache. - Returns an aggregate of per-worker ``in_queue`` (complete, uncommitted - log lines) and ``dead`` (dead-letter records), plus ``in_queue_total`` - and ``dead_total``. No counters are stored: every value is derived from - the files on disk. - - ``in_queue`` is computed with a TAIL READ -- seek to the committed - offset and read only committed->EOF, then count newlines up to the last - ``\\n`` (a torn trailing line has no newline and is not counted). The - whole-file is never read. Results are cached for ``_stats_cache_ttl`` - seconds (monotonic clock) because ``/status`` polls every ~3s; the tail - read plus the cache keep that path cheap under load. - - ``oldest_unflushed_age`` is deferred to C2 and is intentionally NOT - computed or returned here. + Aggregates per-worker ``in_queue`` (complete uncommitted lines) and + ``dead`` (dead-letter records). ``in_queue`` is a tail read from the + committed offset to EOF (the whole file is never read); results are + cached for ``_stats_cache_ttl`` seconds since ``/status`` polls often. """ now = time.monotonic() if ( @@ -402,22 +651,16 @@ def _all() -> dict[str, Any]: try: committed = self._read_committed_offset(worker_key) except (OSError, ValueError): - # /status calls this (via pipeline_metrics); a corrupt or - # transiently-unreadable .offset must NOT 500 the health - # probe. Degrade to 0 for this key's stats -- mirroring the - # existing missing-file->0 convention in - # _read_committed_offset, and tending the conservation - # residual negative (benign, never a false `degraded`). - # Deliberately NO logging here: /status is polled, and a - # per-scan warning on a persistently-corrupt offset would - # flood the hot path. The visibility signal is the aggregate + # Must not 500 the health probe: degrade to 0, mirroring + # the missing-file->0 convention in + # _read_committed_offset. No logging here -- /status is + # polled; a persistently-corrupt offset would flood the + # hot path. Visible instead via the aggregate # `spool.corrupt_offsets` field (see spool_stats()). committed = 0 - # Streamed count of complete lines from committed -> EOF. - # Equivalent to the old f.read() + data[:last_nl+1].count(b"\n") - # (every b"\n" lies at or before the last one), but without - # materialising the undrained tail -- which can be gigabytes - # under a large backlog on this (polled) /status path. + # Streamed count of complete lines from committed -> EOF: + # numerically equivalent to a full-file count, without + # materialising a possibly multi-GB undrained tail. in_queue = self._count_newlines(worker_key, committed) dead = self._count_dead(worker_key) per_key.append( @@ -439,58 +682,16 @@ def _all() -> dict[str, Any]: async def spool_stats(self) -> dict[str, int]: """Cheap, aggregate-only spool footprint for the unauthenticated /status. - Incident context: a durable spool silently grew to 38 GB across 583 - files (largest single file 4.9 GB) with ZERO signal anywhere that it - was happening -- the only symptom was a graph that had stopped - updating. This method exists so that number is always one field away. - - Returns exactly two aggregate integers: - - - ``pending_sessions``: count of worker keys with a ``.log`` file - whose committed offset is strictly less than the file's size, i.e. - there is unconsumed data (mirrors ``active_sessions()``'s - definition, but via ``stat()`` instead of a full scan-and-compare - pass, so it is safe to call on every /status hit). - - ``spool_bytes_total``: total bytes on disk across EVERY file in the - queue directory (``.log`` + ``.offset`` + ``.dead.jsonl``) -- the - same number an operator would get from ``du`` on the spool - directory, without shelling out. - - CHEAP BY CONSTRUCTION: this walks the directory and calls ``stat()`` - on each entry -- O(file count), NEVER O(file bytes). No file content - is read (unlike ``derive_all_stats()``, which tail-reads each log to - count pending lines). This is deliberately how a 38 GB spool can be - sized on every /status poll without walking 38 GB of content. - On top of that, results are cached for ``_spool_cache_ttl`` seconds - (monotonic clock) so a deployment with a very large number of spool - files (thousands of sessions) still does not pay a full directory - scan on every request. - - Per the /status aggregate-only contract (D3): NO session ids, NO - workspace names, and NO per-key table are returned or computable from - this result -- two integers only. - - HEALTH-ENDPOINT SAFE: /status is the unauthenticated health probe (the - ACA liveness surface). This method therefore MUST NOT be able to raise - out to the /status handler -- an uncaught exception there becomes a 500, - a failed health probe, and a container restart loop. Two degradation - rules make that impossible: - - - A directory-level failure (the queue dir missing/unavailable -- e.g. - an Azure Files SMB remount -- or any transient OS error while - scanning) returns the degraded sentinel ``{-1, -1}`` instead of - raising. Unlike every sibling reader, which uses ``glob()`` (empty on - a missing dir), this scan uses ``iterdir()`` (raises on a missing - dir), so the guard is mandatory, not cosmetic. The sentinel is NOT - cached, so the very next poll re-scans and recovers the real numbers - the moment the filesystem is healthy again. - - A per-file failure (a raced delete, or a corrupt/unreadable - ``.offset``) skips just that entry rather than failing the whole - aggregate. - - A ``-1`` in either field is the operator-visible "spool footprint - temporarily unavailable" signal -- distinct from a real ``0`` -- and - never leaks any identifier. + Returns two integers: ``pending_sessions`` (keys whose committed offset + is below the ``.log`` size) and ``spool_bytes_total`` (bytes across all + queue files). Sized via ``stat()`` per file -- O(file count), never + O(bytes) -- and cached for ``_spool_cache_ttl`` seconds. No identifiers + are returned or derivable. + + Must not raise (/status is the unauthenticated health probe): a + directory-level failure returns the uncached sentinel ``{-1, -1}`` and a + per-file failure skips that entry. ``-1`` means "temporarily + unavailable", distinct from a real ``0``. """ now = time.monotonic() if ( @@ -518,19 +719,16 @@ def _scan() -> dict[str, int]: try: committed = self._read_committed_offset(entry.stem) except ValueError: - # The .offset exists but is not a valid integer -- a - # GENUINELY corrupt offset. This is the one visibility - # signal for it (no logging anywhere, to avoid flooding - # the polled health path): surface it as an aggregate - # count on /status so `spool.corrupt_offsets > 0` is the - # operator's alarm. Count this file's bytes; skip its - # pending calc. + # Genuinely corrupt offset (not a valid int). No + # logging (polled health path); surfaced instead via + # the aggregate `spool.corrupt_offsets` count. Count + # this file's bytes; skip its pending calc. corrupt_offsets += 1 continue except OSError: - # A transient/racing FS error reading the offset (NOT - # corruption): count bytes, skip pending calc, and do - # NOT inflate corrupt_offsets with a non-corruption cause. + # Transient/racing FS error, not corruption: count + # bytes, skip pending calc, don't inflate + # corrupt_offsets. continue if committed < size: pending_sessions += 1 @@ -543,12 +741,10 @@ def _scan() -> dict[str, int]: try: stats = await asyncio.to_thread(_scan) except (OSError, ValueError): - # Queue dir missing/unavailable (e.g. Azure Files SMB remount) or a - # transient FS error mid-scan. /status is the health probe and MUST - # return 200 -- degrade to a sentinel and DO NOT cache it, so the - # next poll retries immediately once the filesystem recovers. All - # three fields are -1 = "temporarily unavailable" (distinct from a - # real 0, and from a real corrupt_offsets count). + # Queue dir missing/unavailable, or a transient FS error mid-scan. + # /status must return 200: degrade to an uncached sentinel so the + # next poll retries once the filesystem recovers. -1 means + # "temporarily unavailable", distinct from a real 0. return { "pending_sessions": -1, "spool_bytes_total": -1, @@ -584,64 +780,63 @@ async def purge_dead_letters(self, worker_key: str) -> int: Deletion is routed exclusively through this method: callers must never touch the filesystem directly. + + Guarded by the key's ``_KeyGuard``: an unlink racing a + ``dead_letter`` append from the drainer is the same class of hazard + ``delete_drained`` guards against for the ``.log`` file. """ self._validate_session_id(worker_key) + path = self._dead_path(worker_key) def _purge() -> int: - count = self._count_dead(worker_key) - try: - self._dead_path(worker_key).unlink() - except FileNotFoundError: - pass - return count + with guard.file_lock: + count = self._count_dead(worker_key) + try: + path.unlink() + except FileNotFoundError: + pass + return count - return await asyncio.to_thread(_purge) + with self._guard(worker_key) as guard: + async with guard.admission: + return await _await_uninterrupted(asyncio.to_thread(_purge)) async def recovery_seed_counts(self) -> tuple[int, int]: - """Seed the conservation counters so residual == 0 by construction. - - Returns ``(accepted_seed, written_seed)`` to re-initialise the - accepted/written conservation counters after a crash. Derived purely - from disk so the invariant ``accepted == written + in_queue + dead`` - holds with a zero residual the instant the counters are seeded. - - Per worker key, from disk: - - - ``C`` = complete lines below the committed offset - - ``P`` = complete lines between the committed offset and the end of - complete data (== ``in_queue``) - - ``D`` = dead-letter records - - Formula:: - - written_seed = max(0, C - D) - accepted_seed = written_seed + P + D - - The ``max(0, ...)`` clamp is load-bearing. In a crash/replay window a - dead-but-pending line (dead-lettered, but whose commit has not yet - advanced past it) makes ``C - D`` go negative. The naive formula - ``accepted = C + P`` / ``written = C - D`` yields a negative written - count -- residual ``-1``, a false DEGRADED. Clamping written to zero - and counting the line in BOTH ``P`` and ``D`` absorbs it into - ``accepted_seed`` so the residual stays exactly zero. - - Ordering is load-bearing: this MUST run AFTER ``recovery_reconcile_dead`` - in the lifespan so the dead-letter counts it reads are already settled. + """Seed the conservation counters from disk so residual == 0. + + Returns ``(accepted_seed, written_seed)`` re-derived from disk so + ``accepted == written + in_queue + dead`` holds immediately. Per key, + with C=committed lines, P=pending lines, D=dead records: + ``written_seed = max(0, C - D)``, ``accepted_seed = written_seed + P + + D``. The clamp absorbs a dead-but-not-yet-committed line that would + otherwise drive written negative (a false DEGRADED). Must run after + ``recovery_reconcile_dead`` so the dead counts are settled. """ def _seed() -> tuple[int, int]: accepted = 0 written = 0 for key in self._all_worker_keys(): - committed = self._read_committed_offset(key) - complete_end = self._complete_data_end(key) - dead = self._count_dead(key) - # Streamed newline counts over byte ranges -- numerically - # identical to the old data[:committed].count(b"\n") / - # data[committed:complete_end].count(b"\n"), but without loading - # the whole (possibly multi-GB) log or its slice copies at boot. - before = self._count_newlines(key, 0, committed) - pending = self._count_newlines(key, committed, complete_end) + # Fault-isolate PER KEY -- a corrupt `.offset` must + # not crash the whole seed pass (which runs on every boot, + # BEFORE drainers respawn). A skipped key contributes 0/0. + # The WHOLE per-key body is guarded, not just the offset + # read: a numerically-valid-but-corrupt offset (e.g. + # negative) does not raise when READ, only later when used + # as a seek() position in _count_newlines. + try: + committed = self._read_committed_offset(key) + complete_end = self._complete_data_end(key) + dead = self._count_dead(key) + # Streamed newline counts over byte ranges -- numerically + # identical to the old data[:committed].count(b"\n") / + # data[committed:complete_end].count(b"\n"), but without + # loading the whole (possibly multi-GB) log at boot. + before = self._count_newlines(key, 0, committed) + pending = self._count_newlines(key, committed, complete_end) + except (OSError, ValueError): + logger.error("recovery_seed_counts_key_failed key=%s", key) + continue written_seed = max(0, before - dead) accepted += written_seed + pending + dead written += written_seed @@ -657,74 +852,96 @@ def _dead_payload_set(self, worker_key: str) -> set[bytes]: (base64 of non-UTF-8 bytes). This mirrors ``dead_letter`` and rebuilds the raw line bytes so a reconcile pass can match them against pending log lines. Returns an empty set when no dead-letter file exists. + + A malformed line is skipped (not raised) -- this runs at startup via + ``recovery_reconcile_dead``, and one bad line must not crash-loop the + container. """ try: text = self._dead_path(worker_key).read_text(encoding="utf-8") except FileNotFoundError: return set() payloads: set[bytes] = set() + skipped = 0 for ln in text.splitlines(): if not ln.strip(): continue - record = json.loads(ln) - if "payload" in record: - payloads.add(record["payload"].encode("utf-8")) - elif "payload_b64" in record: - payloads.add(base64.b64decode(record["payload_b64"])) + try: + record = json.loads(ln) + if "payload" in record: + payloads.add(record["payload"].encode("utf-8")) + elif "payload_b64" in record: + payloads.add(base64.b64decode(record["payload_b64"])) + except ( + json.JSONDecodeError, + UnicodeDecodeError, + ValueError, + TypeError, + AttributeError, # a non-string `payload` (e.g. {"payload": 123}) + # raises AttributeError on `.encode()` -- this is a boot-path + # total function; one bad record must not crash-loop the + # container. + ): + skipped += 1 + continue + if skipped: + logger.warning( + "dead_letter_unparseable key=%s skipped=%d", worker_key, skipped + ) return payloads async def recovery_reconcile_dead(self) -> int: """Advance committed offsets past leading already-dead pending lines. - Closes the dead_letter->commit crash window (D2). When the process - crashes after a poison line was dead-lettered but before the commit - advanced past it, the line remains pending in the ``.log``. A naively - respawned drainer would re-read it, re-dead-letter it, and permanently - corrupt the dead count. This pass steps the committed offset over each - LEADING pending line whose raw bytes already appear in the dead-letter - file, stopping at the first non-dead pending line. - - Per worker key with a ``.log`` and a non-empty dead-payload set, walk - from the committed offset toward the end of complete data: for each - leading line whose raw bytes are in the dead-payload set, advance past - it (``skipped += 1``); stop at the first non-dead pending line. If the - offset advanced, persist it atomically (tmp + ``os.replace``, mirroring - ``commit``). Returns the total number of lines skipped across all keys. - - Covers both the crash window (dead_letter then crash before commit) and - the replay window (re-append then crash before purge). - - Ordering is load-bearing: this MUST run ONCE at startup, BEFORE - ``recovery_seed_counts`` and BEFORE drainers respawn. + Closes the dead-letter->commit crash window: a line dead-lettered but + not yet committed past would otherwise be re-read and re-dead-lettered + by a respawned drainer. Per key, steps the committed offset over each + leading pending line whose bytes are already in the dead-letter file, + stopping at the first non-dead line, and persists it atomically. + Returns the total lines skipped. Must run once at startup, before + ``recovery_seed_counts`` and before drainers respawn. """ def _reconcile() -> int: total_skipped = 0 for key in self._all_worker_keys(): - dead_payloads = self._dead_payload_set(key) - if not dead_payloads: - continue + # 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`) has no log to reconcile, so skip the + # expensive read entirely. log_path = self._log_path(key) if not log_path.exists(): 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) + # 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 self._stats_cache = None return total_skipped diff --git a/context_intelligence_server/registry.py b/context_intelligence_server/registry.py index ad078fee..59796802 100644 --- a/context_intelligence_server/registry.py +++ b/context_intelligence_server/registry.py @@ -1,6 +1,7 @@ """Session registry — per-session worker management.""" import asyncio +import functools import json import logging import time @@ -22,11 +23,12 @@ _DRAIN_MAX_BATCH = 100 _DRAIN_POLL_INTERVAL = 0.05 # idle poll cadence; bounded by flush_timeout -# A positive residual must PERSIST this long before it is called degraded. -# Must exceed the worst-case transient-skew window: the derive_all_stats -# cache TTL (1.0s) plus the /status poll cadence (~3s). 15s is >10x the cache -# TTL, so any in-flight two-clock skew clears well before it trips degraded, -# while a genuine (monotonic, non-clearing) silent drop still trips it. +# Bounded retry count for the finalize delete-drained loop; not operator-tunable. +# No backoff between attempts -- sleeping would widen the race window this closes. +_FINALIZE_DELETE_ATTEMPTS = 3 + +# Grace period before a positive residual is flagged degraded -- must exceed +# the stats-cache TTL + poll cadence to avoid false positives from clock skew. _RESIDUAL_DEGRADED_GRACE = 15.0 @@ -41,10 +43,12 @@ class SessionWorker: events_processed: int = 0 started_at: float = field(default_factory=time.time) error_count: int = 0 - # Phase 2 (#278): liveness timestamp — when the flush boundary last - # completed for this worker. Defaults to creation time (NOT 0.0) so a - # brand-new worker reads as fresh, not ancient. Stamped in _flush_barrier. + # Timestamp when the flush boundary last completed; defaults to creation + # time (not 0.0) so a brand-new worker reads as fresh. Set in _flush_barrier. last_successful_flush: float = field(default_factory=time.time) + # Set True by _safe_close, as its FIRST statement. A worker + # whose store has been closed is never revived — see start_drain. + store_closed: bool = False @dataclass @@ -64,26 +68,22 @@ class SessionRegistry: def __init__(self) -> None: self._workers: dict[str, SessionWorker] = {} self._completed: deque[CompletedSession] = deque(maxlen=100) - # Durable-ingest infrastructure, built lazily on first use. The - # module-level registry singleton is constructed at import time, - # before the per-test settings patch applies, so we cannot read - # settings here — see _ensure_infra(). + # Strong refs to fire-and-forget close tasks -- asyncio only holds a + # weak ref, so without this a close can be GC'd mid-execution. Self-discards on done. + self._close_tasks: set[asyncio.Task] = set() + # Durable-ingest infra, built lazily (see _ensure_infra) since the + # module-level singleton is constructed before test settings patches apply. self._queue_manager: QueueManager | None = None self._write_semaphore: asyncio.Semaphore | None = None self._max_delivery_attempts: int = 0 - # Live pipeline-conservation counters (D2): make silently-dropped - # events observable via /status. accepted = events admitted to the - # log; written = events persisted to Neo4j; replayed = events - # re-driven from the log on recovery; write_retries = transient - # write retries attempted by the drainer. + # Live conservation counters surfaced via /status (accepted/written/ + # replayed/write_retries) so silently-dropped events are observable. self._accepted_total: int = 0 self._written_total: int = 0 self._replayed_total: int = 0 self._write_retries_total: int = 0 - # FIX B: monotonic timestamp when the residual first went positive and - # stayed unexplained. None means "clean". Gates the degraded flag so a - # transient two-clock skew never latches; only a sustained positive - # residual (real silent drop) does. + # Monotonic time the residual first went positive (None = clean). + # Gates `degraded` so transient clock skew doesn't latch it. self._residual_positive_since: float | None = None def _ensure_infra(self) -> None: @@ -128,17 +128,10 @@ def record_replayed(self, n: int) -> None: def record_purged(self, n: int) -> None: """Remove n purged dead-letters from the accepted total (conservation). - A bare dead-letter purge unlinks the .dead.jsonl file, dropping `dead` - by n. Those lines were counted in `accepted` at ingest but never - `written`; discarding them from disk must also discard them from - `accepted`, or the residual latches at +n forever. Symmetric to - record_replayed, which moves lines dead -> in_queue and therefore must - NOT touch accepted. - - Clamp: accepted can never fall below written. Under the single-writer - guarantee the clamp can never legitimately engage (a dead line is - accepted-but-not-written, so n <= accepted - written); if it does, log - a warning as an accounting-drift signal rather than silently masking it. + A dead-letter purge drops `dead` by n without ever having been + `written`, so `accepted` must drop too or the residual latches at +n + forever. Clamped so accepted never falls below written; an engaged + clamp logs a warning as an accounting-drift signal. """ if n <= 0: return @@ -178,38 +171,15 @@ def pipeline_counters(self) -> dict[str, int]: } async def pipeline_metrics(self) -> dict[str, Any]: - """Assemble the pipeline-conservation health block for /status (D2/D3). - - Combines the live in-memory counters (pipeline_counters) with the - disk-derived queue/dead aggregate (queue_manager.derive_all_stats) into - a single conservation view. The residual is the count of accepted - events that are neither persisted, nor still queued, nor dead-lettered: - - residual = accepted - written - in_queue - dead - - ``degraded`` is True whenever ``dead > 0`` (an accounted-for loss, no - grace period) OR the residual is POSITIVE and has stayed positive for - at least ``_RESIDUAL_DEGRADED_GRACE`` seconds. A negative residual is - never degraded (it is benign two-clock skew between the live counters - and the cached disk snapshot, clamped to a ``lost`` value of zero), and - a positive residual that clears before the grace window elapses is - treated as the same transient skew rather than real loss. - - IMPORTANT caveats: - - This is a LIVE per-process measure, not an all-time audit. Finalized - session logs are deleted by ``delete_drained``, so their accepted / - written / in_queue contributions leave the disk-derived aggregate. - The in-memory accepted/written counters persist, so the residual - stays conserved for the lifetime of the process (seeded across - restarts via ``seed_counters``). - - It is only valid under the single-worker (single-process) guarantee: - one writer owns the counters and the on-disk queues. - - ``write_retries_total`` is the transient/deadlock proxy — the closest - observable signal for retried (e.g. DeadlockDetected) writes. - - ``deadlock_detected_total`` and ``events_failed_total`` are - intentionally omitted: neither is cleanly trackable at this layer. - - ``oldest_unflushed_age`` is DEFERRED to C2 and is intentionally - absent from this block. + """Assemble the pipeline-conservation health block for /status. + + Combines live counters with the disk-derived queue/dead aggregate. + residual = accepted - written - in_queue - dead. `degraded` is True + when dead > 0, or when a positive residual persists past + `_RESIDUAL_DEGRADED_GRACE` seconds (transient clock skew clears + before then). Live per-process only: finalized session logs are + deleted, but the in-memory counters persist across restarts via + `seed_counters`. """ agg = await self.queue_manager.derive_all_stats() counters = self.pipeline_counters() @@ -270,7 +240,7 @@ async def _process_one( result = "error" error = str(exc) worker.error_count += 1 - raise # Phase B2: propagate so the drainer dead-letters this line + raise # Propagate so the drainer dead-letters this line finally: ring_buffer.add( EventRecord( @@ -284,23 +254,22 @@ async def _process_one( ) async def _flush_barrier(self, worker: SessionWorker) -> None: - """The ONE Neo4j-write boundary: a semaphore-gated, awaited flush. + """The one Neo4j-write boundary: a semaphore-gated, awaited flush. - Acquiring self.write_semaphore caps the number of concurrent Neo4j - write transactions across ALL session drainers (the starvation guard). - The offset must only ever advance AFTER this returns successfully. + The semaphore caps concurrent write transactions across all session + drainers. The offset must only advance after this returns successfully. - Correctness of commit-after-flush depends on neo4j_store._flush_body - snapshotting+clearing the buffer under _flush_lock and RESTORING it on - failure (neo4j_store.py:686-696), plus the empty-buffer early return - (:656-657). We do not modify that file; we rely on it here. + Commit-after-flush is correct only because ``neo4j_store._flush_body`` + snapshots-and-clears the buffer under ``_flush_lock`` and RESTORES it on + failure, plus the empty-buffer early return. That file is not modified + here; this barrier depends on it. This is also the ONLY place a flush + may happen -- a handler that flushes on its own would write outside the + semaphore (see ``SessionHandler._handle_end``). """ async with self.write_semaphore: await worker.services.graph.flush() - # Phase 2 (#278): stamp liveness at the SINGLE flush boundary all - # three success paths funnel through. Marks completion of the flush - # barrier (advances even on an empty-buffer flush = liveness proof - # that the drainer reached and finished the write barrier). + # Stamped here (the one flush boundary) as liveness proof the + # drainer reached and finished the write barrier. worker.last_successful_flush = time.time() async def drain_worker( @@ -308,14 +277,18 @@ async def drain_worker( ) -> None: """Durable drain loop for one session. - Reads the next batch after the committed offset, dispatches each line - through process_event, then runs the single semaphore-gated flush - barrier and commits the offset only on success (the "ack"). A batch - that exhausts its retry budget — or that raises during dispatch — is - isolated ONE LINE AT A TIME and dead-lettered (never silently dropped). - When the log is idle the drainer polls and reaps the session if it has - been idle past the stale timeout. The drainer is the SOLE flush trigger - (process_event no longer self-flushes, Task 6). + Reads batches after the committed offset, dispatches each event, runs + the flush barrier, and commits only on success; an exhausted retry + budget dead-letters the batch line-by-line. + + Any other exception propagates -- `_on_drain_done` is the sole + supervision point (logs, closes, deregisters so a respawn or boot + `recover()` picks it up). A terminal `session:end` record is left + uncommitted so a later drain re-enters `_finalize_session`. + + The queue owns all byte-position math; this registry only chooses + which offset to commit via `qm.commit`/`qm.dead_letter`. When idle, + the drainer polls and reaps the session past the stale timeout. """ handlers = setup_handlers(worker.services) qm = self.queue_manager @@ -328,7 +301,7 @@ async def drain_worker( try: batch = await qm.read_batch(session_id, max_items=_DRAIN_MAX_BATCH) - if not batch.lines: + if not batch.records: await asyncio.sleep(poll_interval) idle_elapsed += poll_interval if idle_elapsed >= flush_timeout: @@ -340,9 +313,10 @@ async def drain_worker( > settings.stale_session_timeout ): logger.info( - "Reaping stale session %s (idle > %s seconds)", + "session_reaped_stale session=%s idle_seconds=%s", session_id, settings.stale_session_timeout, + extra={"session_id": session_id}, ) await self._safe_close(worker) self._deregister(session_id) @@ -353,19 +327,29 @@ async def drain_worker( # --- dispatch + durable write barrier, one error path --- try: - saw_terminal = await self._process_batch(worker, batch, handlers) + safe_count, terminal_at = await self._process_batch( + worker, batch, handlers + ) await self._flush_barrier(worker) except asyncio.CancelledError: + # INFO not ERROR: a cancel here is normally deliberate + # (shutdown, idle reap, test teardown), not a failure. + logger.info( + "drain_worker_cancelled session=%s site=%s", + session_id, + "dispatch", + extra={"session_id": session_id}, + ) await self._safe_close(worker) + # Must deregister so get_or_create builds a fresh worker -- + # else start_drain's store_closed guard refuses it forever. + self._deregister(session_id) return except Exception: attempts += 1 self.record_write_retry() - # Throttle the failure log off the local attempts counter - # (resets to 0 on commit and after exhaustion): the first - # failure gets ONE traceback (WARNING), middle attempts are - # DEBUG, and budget exhaustion gets a single ERROR (no - # per-attempt traceback storm). + # First failure: WARNING w/ traceback. Middle attempts: DEBUG. + # Exhaustion: single ERROR. Avoids a per-attempt traceback storm. if attempts == 1: logger.warning( "drain_batch_failed session=%s attempt=%d", @@ -389,35 +373,51 @@ async def drain_worker( extra={"session_id": session_id}, ) if attempts >= self._max_delivery_attempts: - # Budget spent -> isolate the batch ONE LINE AT A TIME, - # dead-letter the offending line(s), advance past all. - await self._handle_exhausted_batch(worker, batch, handlers) + # Budget spent: isolate the batch line-by-line and dead-letter. + terminal_seen = await self._handle_exhausted_batch( + worker, batch, handlers + ) + if terminal_seen: + # Mirror the normal terminal branch below: the + # session:end record was left uncommitted, so + # finalize instead of resuming the drain loop. + await self._finalize_session(worker, handlers) + return attempts = 0 continue - # Budget NOT yet spent: back off one poll interval before - # re-reading the SAME offset (offset is not committed; the - # idempotent MERGE makes the replay a no-op). The backoff - # avoids a tight Neo4j-hammering retry loop on a transient - # deadlock and keeps retries on the loop's poll cadence. + # Not yet exhausted: back off before re-reading the same + # offset (idempotent MERGE makes the replay a no-op). await asyncio.sleep(poll_interval) continue attempts = 0 - await qm.commit(session_id, batch.end_offset) - self.record_written(len(batch.lines)) + # 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) + counted = len(batch.records) if terminal_at is None else safe_count + self.record_written(counted) logger.debug( "batch_committed events=%d offset=%d", - len(batch.lines), - batch.end_offset, + counted, + commit_to, extra={"session_id": session_id}, ) - if saw_terminal: + if terminal_at is not None: await self._finalize_session(worker, handlers) return except asyncio.CancelledError: + # Cancelled while reading/idle (outer site; never reaches the inner try). + logger.info( + "drain_worker_cancelled session=%s site=%s", + session_id, + "loop", + extra={"session_id": session_id}, + ) await self._safe_close(worker) + self._deregister(session_id) # See the note above. return @staticmethod @@ -428,88 +428,167 @@ def _parse_line(raw: bytes) -> tuple[str, str, dict[str, Any]]: async def _process_batch( self, worker: SessionWorker, batch: Batch, handlers: Any - ) -> bool: - """Dispatch each line in the batch; return True if it contained a - terminal (session:end) event.""" - from context_intelligence_server.pipeline import TERMINAL_EVENTS # noqa: PLC0415 + ) -> tuple[int, int | None]: + """Dispatch every record; report the first terminal boundary. + + Returns ``(safe_count, terminal_at)``: ``terminal_at`` is the queue- + produced start offset of the first ``session:end`` record (or None), + and ``safe_count`` is how many records precede it. Every record is + still dispatched -- a failed terminal dispatch still goes through + the retry/isolation path. + + CONSEQUENCE -- ``session:end`` IS DISPATCHED TWICE PER SESSION. It is + dispatched here, but ``terminal_at`` leaves it UNCOMMITTED so that + "ended but not finalized" survives a respawn; ``_finalize_session`` + then re-reads it via ``_drain_to_eof`` and dispatches it again. Every + handler on the terminal path must therefore be idempotent. The three + that claim ``session:end`` today (the data_layer_1 field lifter, + ``SessionHandler`` in data_layer_2, and ``DelegationHandler`` in + data_layer_3) are read-then-MERGE and tolerate it. Any new terminal + handler with a non-idempotent side effect (counters, appends, + notifications) will double-fire here. + """ + from context_intelligence_server.pipeline import ( + TERMINAL_EVENTS, + ) - saw_terminal = False - for raw in batch.lines: - event, _workspace, data = self._parse_line(raw) + terminal_at: int | None = None + safe_count = 0 + for rec in batch.records: + event, _workspace, data = self._parse_line(rec.raw) await self._process_one(worker, event, data, handlers) - if event in TERMINAL_EVENTS: - saw_terminal = True - return saw_terminal + if terminal_at is None: + if event in TERMINAL_EVENTS: + terminal_at = rec.start + else: + safe_count += 1 + return safe_count, terminal_at async def _handle_exhausted_batch( self, worker: SessionWorker, batch: Batch, handlers: Any - ) -> None: - """Reprocess a poison batch ONE LINE AT A TIME (linear isolation). - - Each line is dispatched + flushed individually under the write - semaphore. A line that still fails (parse error, handler error, or - repeated flush failure) is dead-lettered with its error AND its write - residue is discarded from the store buffer (COE blocker, decision #13); - good lines flush normally. Every line advances the offset past itself - (commit), so the whole batch is accounted for. No silent loss, no - binary shrink, no cross-line contamination. + ) -> bool: + """Reprocess a poison batch one line at a time (linear isolation). + + Each record is dispatched and flushed individually. A record that + fails is dead-lettered and its buffer residue discarded so it can't + contaminate later records. Every non-terminal record advances the + offset to its own queue-produced end, so it is fully accounted for. + + A record that successfully parses as a terminal ``session:end`` + record is NOT dispatched or committed here -- isolation stops + immediately and returns True, leaving that record (and anything + after it) uncommitted, mirroring the normal drain loop's terminal + semantics (see ``drain_worker``/``_process_batch``). The caller must + then call ``_finalize_session`` instead of resuming the drain loop, + exactly like the non-exhausted terminal path: ``_finalize_session``'s + own ``_drain_to_eof`` re-reads and re-dispatches the terminal record. + A record whose bytes fail to parse is NOT terminal -- it is + dead-lettered and committed past like any other poison line. + + Returns False when the whole batch is isolated without ever + reaching a terminal record (unchanged behavior: no finalization). """ qm = self.queue_manager session_id = worker.session_id - # The failed BATCH flush left its writes resident in the store buffer - # (_flush_body restores on failure, neo4j_store.py:686-696). Discard that - # accumulated residue so the FIRST isolated line flushes from a clean - # buffer — otherwise the poison line's residue contaminates line 1. + # 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() - offset = batch.start_offset - for raw in batch.lines: - line_end = offset + len(raw) + 1 # +1 for the newline read_batch strips + for rec in batch.records: + try: + event, _ws, data = self._parse_line(rec.raw) + except Exception as exc: + # Unparseable: can't be a terminal record -- poison as before. + await qm.dead_letter(session_id, rec.raw, str(exc)) # no re-framing + logger.warning( + "dead_letter session=%s error=%s", + session_id, + exc, + exc_info=exc, + extra={"session_id": session_id}, + ) + worker.services.graph.discard_buffer() + await qm.commit(session_id, rec.end) # queue-produced offset + continue + + from context_intelligence_server.pipeline import TERMINAL_EVENTS + + if event in TERMINAL_EVENTS: + return True + + wrote = False try: - event, _ws, data = self._parse_line(raw) await self._process_one(worker, event, data, handlers) await self._flush_barrier(worker) - self.record_written(1) + wrote = True except Exception as exc: - await qm.dead_letter(session_id, raw + b"\n", str(exc)) + await qm.dead_letter(session_id, rec.raw, str(exc)) # no re-framing logger.warning( "dead_letter session=%s error=%s", session_id, exc, + exc_info=exc, extra={"session_id": session_id}, ) - # COE blocker (decision #13): drop the failed line's residue so - # it cannot contaminate the NEXT line's flush. A successful flush - # clears the buffer itself; only the failure path needs this. + # 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, line_end) - offset = line_end + await qm.commit(session_id, rec.end) # queue-produced offset + if wrote: + self.record_written(1) + return False - async def _finalize_session(self, worker: SessionWorker, handlers: Any) -> None: - """session:end seen: drain any tail lines read-to-EOF, then record the - CompletedSession, close the graph, deregister, and DELETE the drained - logs. Panel finding #7: if a tail flush fails, do NOT finalize — return - without recording/closing so the drainer retries (no tail loss).""" + async def _drain_to_eof(self, worker: SessionWorker, handlers: Any) -> bool: + """Drain every remaining record for this session up to EOF. + + Returns True when fully drained. Returns False when a tail flush + failed -- nothing was committed, and the caller must not finalize. + """ qm = self.queue_manager session_id = worker.session_id while True: tail = await qm.read_batch(session_id, max_items=_DRAIN_MAX_BATCH) - if not tail.lines: - break + if not tail.records: + return True try: await self._process_batch(worker, tail, handlers) await self._flush_barrier(worker) except Exception: logger.exception("finalize_tail_flush_failed session=%s", session_id) - return # NOT finalized: keep worker alive, leave tail uncommitted + return False # NOT finalized: keep worker alive, tail uncommitted await qm.commit(session_id, tail.end_offset) - self.record_written(len(tail.lines)) + self.record_written(len(tail.records)) logger.debug( "batch_committed events=%d offset=%d", - len(tail.lines), + len(tail.records), tail.end_offset, extra={"session_id": session_id}, ) + async def _finalize_session(self, worker: SessionWorker, handlers: Any) -> None: + """session:end seen: drain to EOF, record CompletedSession, delete + the drained log, close the graph, then deregister -- in that order. + + If the tail flush fails, finalization is aborted (no record/close) + so a respawn retries. ``delete_drained`` returning False means an + append landed after drain -- retried up to + ``_FINALIZE_DELETE_ATTEMPTS`` times, re-draining each time; a + persistent failure retains the log as a bounded, non-lossy residual. + """ + qm = self.queue_manager + session_id = worker.session_id + if not await self._drain_to_eof(worker, handlers): + # Orphan: still registered, task about to finish. Recoverable -- + # a respawn or boot recover() re-enters _finalize_session. Stays + # registered so orphaned_sessions() surfaces it on /status. + logger.warning( + "finalize_orphan session=%s reason=tail_flush_failed " + "recoverable=respawn", + session_id, + extra={"session_id": session_id}, + ) + return + ended_at = time.time() self._completed.append( CompletedSession( @@ -522,11 +601,40 @@ async def _finalize_session(self, worker: SessionWorker, handlers: Any) -> None: duration_seconds=ended_at - worker.started_at, ) ) + # Reclaim disk (keep .dead.jsonl). delete_drained's return is + # load-bearing: False means an append landed after drain -- retry below. + for attempt in range(1, _FINALIZE_DELETE_ATTEMPTS + 1): + if await qm.delete_drained(session_id): + break + logger.warning( + "finalize_delete_retained session=%s attempt=%d/%d", + session_id, + attempt, + _FINALIZE_DELETE_ATTEMPTS, + extra={"session_id": session_id}, + ) + if attempt == _FINALIZE_DELETE_ATTEMPTS: + # Give up: log retained, picked up by recover() or the sweep. + logger.error( + "finalize_delete_gave_up session=%s retained_log=true " + "pickup=recover_sweep", + session_id, + extra={"session_id": session_id}, + ) + break + if not await self._drain_to_eof(worker, handlers): + # Permanent orphan: this session can never re-enter + # _finalize_session, so nothing will retry it on its own. + # Stays registered so orphaned_sessions() surfaces it on /status. + logger.error( + "finalize_orphan session=%s reason=delete_retry_exhausted " + "permanent=true", + session_id, + extra={"session_id": session_id}, + ) + return # late-tail flush failed: same semantics as the first pass await self._safe_close(worker) - self._deregister(session_id) - # Panel finding #5: reclaim disk — a fully drained, finalized session no - # longer needs its .log/.offset. Keep .dead.jsonl (retained dead-letter). - await qm.delete_drained(session_id) + self._deregister(session_id) # the LAST act -- no await after this logger.info( "session_finalized session=%s events=%d", session_id, @@ -535,15 +643,90 @@ async def _finalize_session(self, worker: SessionWorker, handlers: Any) -> 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 + the await, so there is no suspension point between "we began + closing" and "it is marked".""" + worker.store_closed = True try: await worker.services.graph.close() except Exception: logger.exception("graph.close failed for session %s", worker.session_id) + @staticmethod + def _task_failure(task: asyncio.Task) -> BaseException | None: + """The exception a finished task died with, else None. + + None for a task that is still running, was cancelled, or returned + cleanly. Checking ``cancelled()`` first is mandatory: ``task.exception()`` + RAISES ``CancelledError`` on a cancelled task. + """ + if not task.done() or task.cancelled(): + return None + return task.exception() + + def _on_drain_done(self, worker: SessionWorker, task: asyncio.Task) -> None: + """The ONE supervision point for a finished drain task. + + Synchronous by asyncio contract, invoked via ``call_soon`` exactly + once per task, and only ever AFTER the task is done -- so it can + never race a live drainer. + """ + exc = self._task_failure(task) + if exc is None: + return # cancelled, or a clean return + session_id = worker.session_id + try: + logger.error( + "drain_worker_died session=%s", + session_id, + exc_info=exc, + extra={"session_id": session_id}, + ) + finally: + # Teardown must happen even if logging itself failed. + self._deregister(session_id) # sync; first, so revival unblocks + try: + close_task = asyncio.get_running_loop().create_task( + self._safe_close(worker), name=f"close-{session_id}" + ) + except RuntimeError: # loop already closing at shutdown + logger.warning( + "drain_worker_died_close_skipped session=%s", + session_id, + extra={"session_id": session_id}, + ) + else: + # Hold a strong ref -- asyncio only keeps a weak one, so + # without this the close task can be GC'd mid-execution. + self._close_tasks.add(close_task) + close_task.add_done_callback(self._close_tasks.discard) + def start_drain(self, worker: SessionWorker) -> None: - if worker.task is None or worker.task.done(): - worker.task = asyncio.create_task( - self.drain_worker(worker), name=f"drain-{worker.session_id}" + if worker.store_closed: + # Spent store: draining through it would dead-letter good events. + # The closer MUST also deregister, or this refuses the worker forever. + return + task = worker.task + if task is not None: + if not task.done(): + return # live drainer -- nothing to do + if self._task_failure(task) is not None: + # Crashed; the done-callback owns teardown and will deregister. + return + # A previous, cleanly-finished task means this is a respawn, distinct + # from a brand-new worker (task is None, logged by get_or_create). + respawn = task is not None + new_task = asyncio.create_task( + self.drain_worker(worker), name=f"drain-{worker.session_id}" + ) + new_task.add_done_callback(functools.partial(self._on_drain_done, worker)) + worker.task = new_task + if respawn: + logger.info( + "drainer_respawned session=%s", + worker.session_id, + extra={"session_id": worker.session_id}, ) def get_or_create( @@ -580,10 +763,11 @@ def get_or_create( extra={"session_id": session_id}, ) else: - # Session-ownership invariant: each session_id is owned by exactly one - # contributor; the bound created_by (set once at creation) is load-bearing - # for provenance. Log at ERROR — not WARNING — so monitoring surfaces a - # violation observably; preserve the bound id and don't crash live ingest. + # Respawn on every repeat event so a deregistered-but-not-yet- + # revived worker comes back the moment traffic resumes. + self.start_drain(self._workers[session_id]) + # Each session_id is owned by exactly one contributor. ERROR (not + # WARNING) so monitoring surfaces a violation; ingest still proceeds. if created_by is not None: bound = getattr( self._workers[session_id].services.graph, "created_by", None @@ -606,6 +790,13 @@ def get_or_create( def remove(self, session_id: str) -> None: worker = self._workers.pop(session_id, None) if worker and worker.task and not worker.task.done(): + # Forced removal of a still-live task; the normal path is a graceful finalize. + logger.info( + "drain_worker_remove session=%s had_live_task=%s", + session_id, + True, + extra={"session_id": session_id}, + ) worker.task.cancel() def _deregister(self, session_id: str) -> None: @@ -629,14 +820,11 @@ def workers(self) -> list[SessionWorker]: return list(self._workers.values()) def orphaned_sessions(self) -> list[SessionWorker]: - """Return workers that are still registered but whose drain task has - finished — the silent-stall signal for #278. - - A worker is orphaned iff it is in _workers AND its task has completed - (task.done()). This catches the finalization-path orphan (a tail flush - failure returns early without deregistering, so the task completes but - the worker is never removed) and any unhandled exception that escapes - the drain loop. Deterministic and instant — no timer, no threshold. + """Return workers still registered whose drain task has finished. + + Orphaned iff in ``_workers`` AND ``task.done()`` -- catches a tail- + flush failure that returns early without deregistering, and any + unhandled exception escaping the drain loop. Deterministic, no timer. """ return [ worker diff --git a/pyproject.toml b/pyproject.toml index 2aaa7caf..fd2d754a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "context-intelligence-server" -version = "6.7.0" +version = "6.7.1" description = "Context Intelligence Server for Amplifier" requires-python = ">=3.11" dependencies = [ diff --git a/tests/neo4j/test_end_handler_buffer_shadow.py b/tests/neo4j/test_end_handler_buffer_shadow.py index 9f66fd38..976f9caa 100644 --- a/tests/neo4j/test_end_handler_buffer_shadow.py +++ b/tests/neo4j/test_end_handler_buffer_shadow.py @@ -60,6 +60,22 @@ async def _neo4j_labels(services: Any, node_id: str) -> list[str]: return list(rows[0]["lbls"]) if rows else [] +async def _neo4j_props(services: Any, node_id: str) -> dict[str, Any]: + """Return properties from Neo4j directly (bypasses buffer). + + Used to assert the end-handler's own writes (ended_at/status) genuinely + landed, so the assertion isn't merely re-checking the earlier start/fork + flush's data. + """ + rows = await services.graph.execute_query( + "MATCH (n) WHERE n.node_id = $id AND n.workspace = $workspace " + "RETURN properties(n) AS props", + {"id": node_id, "workspace": services.graph.workspace}, + workspace="*", + ) + return dict(rows[0]["props"]) if rows else {} + + # --------------------------------------------------------------------------- # Store-level tests — prove the shadow mechanism directly # --------------------------------------------------------------------------- @@ -268,7 +284,10 @@ async def test_start_flush_end_yields_single_terminal_label( "timestamp": "2026-01-01T10:05:00Z", }, ) - # _handle_end calls flush() itself at the end. + # _handle_end does not flush itself; the drainer flushes after the + # batch. Flush here so this test's read sees the persisted result, + # as production does via the gated _flush_barrier. + await neo4j_services.graph.flush() # Read final labels directly from Neo4j (bypasses any buffer). final_labels = await _neo4j_labels(neo4j_services, child_id) @@ -287,6 +306,17 @@ async def test_start_flush_end_yields_single_terminal_label( f"Expected exactly one terminal label SubSession; got {terminals} in {final_labels}" ) + # The assertions below must depend on the end-handler's own writes + # having landed -- check ended_at/status so this isn't merely + # re-checking the earlier start+flush. + final_props = await _neo4j_props(neo4j_services, child_id) + assert final_props.get("status") == "completed", ( + f"end-handler's status write missing from Neo4j: {final_props}" + ) + assert final_props.get("ended_at") is not None, ( + f"end-handler's ended_at write missing from Neo4j: {final_props}" + ) + async def test_fork_flush_end_yields_single_terminal_label( self, neo4j_services: Any ) -> None: diff --git a/tests/neo4j/test_handler_flush_concurrency.py b/tests/neo4j/test_handler_flush_concurrency.py new file mode 100644 index 00000000..2695d7e8 --- /dev/null +++ b/tests/neo4j/test_handler_flush_concurrency.py @@ -0,0 +1,251 @@ +"""Neo4j-backed tests for single-terminal-label behaviour on session:end. + +Covers that _handle_end does not flush on its own, that no flush ever runs +outside the drainer's write_semaphore, and that terminal data is durably +flushed before the queue log is deleted. + + uv run pytest tests/neo4j/test_handler_flush_concurrency.py -q -m neo4j +""" + +from __future__ import annotations + +import asyncio +from datetime import datetime +from pathlib import Path +from typing import Any + +import pytest +from neo4j import AsyncGraphDatabase + +from context_intelligence_server.handlers.data_layer_2.session import SessionHandler +from context_intelligence_server.neo4j_store import Neo4jGraphStore, ensure_neo4j_schema +from context_intelligence_server.pipeline import process_event, setup_handlers +from context_intelligence_server.queue_manager import QueueManager +from context_intelligence_server.registry import SessionRegistry, SessionWorker +from context_intelligence_server.services import HookStateService + +pytestmark = pytest.mark.neo4j + +_WS = "handler-flush" +_TS = "2026-01-01T10:00:00+00:00" +_TS2 = "2026-01-01T10:05:00+00:00" +_TS3 = "2026-01-01T10:10:00+00:00" + + +async def _neo4j_labels(services: Any, node_id: str) -> list[str]: + rows = await services.graph.execute_query( + "MATCH (n) WHERE n.node_id = $id AND n.workspace = $workspace " + "RETURN labels(n) AS lbls", + {"id": node_id, "workspace": services.graph.workspace}, + workspace="*", + ) + return list(rows[0]["lbls"]) if rows else [] + + +async def _neo4j_labels_and_props_via_container( + container: dict[str, Any], workspace: str, node_id: str +) -> tuple[list[str], dict[str, Any]]: + # Fresh driver: after finalize, the worker's own store/driver is closed. + driver = AsyncGraphDatabase.driver( + container["bolt_url"], auth=(container["user"], container["password"]) + ) + try: + async with driver.session() as session: + result = await session.run( + "MATCH (n) WHERE n.node_id = $id AND n.workspace = $ws " + "RETURN labels(n) AS lbls, properties(n) AS props", + id=node_id, + ws=workspace, + ) + record = await result.single() + if record is None: + return [], {} + return list(record["lbls"]), dict(record["props"]) + finally: + await driver.close() + + +def _terminals(labels: list[str]) -> list[str]: + return [ + label + for label in labels + if label in ("RootSession", "SubSession", "ForkedSession", "IncompleteSession") + ] + + +def _line(event: str, workspace: str, data: dict[str, Any]) -> bytes: + import json + + return json.dumps({"event": event, "workspace": workspace, "data": data}).encode( + "utf-8" + ) + + +def _build_registry(queues_dir: Path, *, write_concurrency: int = 8) -> SessionRegistry: + reg = SessionRegistry() + reg._queue_manager = QueueManager(queues_dir=queues_dir) + reg._write_semaphore = asyncio.Semaphore(write_concurrency) + reg._max_delivery_attempts = 3 + return reg + + +def _build_worker(container: dict[str, Any], sid: str) -> SessionWorker: + store = Neo4jGraphStore( + uri=container["bolt_url"], + auth=(container["user"], container["password"]), + workspace=_WS, + ) + services = HookStateService(workspace=_WS, graph_store=store) + return SessionWorker(session_id=sid, workspace=_WS, services=services) + + +@pytest.fixture(autouse=True) +async def _schema(neo4j_container: dict[str, Any]) -> None: + driver = AsyncGraphDatabase.driver( + neo4j_container["bolt_url"], + auth=(neo4j_container["user"], neo4j_container["password"]), + ) + try: + await ensure_neo4j_schema(driver) + finally: + await driver.close() + + +class TestSingleTerminalLabel: + async def test_same_batch_end_then_end_via_pipeline( + self, neo4j_services: Any + ) -> None: + # Two session:end for the same session in one batch (no flush between) + # must still leave exactly one terminal label. + services = neo4j_services + handlers = setup_handlers(services) + worker = SessionWorker(session_id="worker", workspace=_WS, services=services) + child_id = "child" + + await process_event( + worker, + "session:start", + {"session_id": child_id, "parent_id": "parent", "timestamp": _TS}, + handlers, + ) + await services.graph.flush() + + lbls_after_start = await _neo4j_labels(services, child_id) + assert "SubSession" in lbls_after_start, lbls_after_start + + await process_event( + worker, "session:end", {"session_id": child_id, "timestamp": _TS2}, handlers + ) + await process_event( + worker, "session:end", {"session_id": child_id, "timestamp": _TS3}, handlers + ) + await services.graph.flush() + + terminals = _terminals(await _neo4j_labels(services, child_id)) + assert terminals == ["SubSession"], terminals + + async def test_same_batch_end_then_end_direct_handler( + self, neo4j_services: Any + ) -> None: + # Same property via a direct handler call (no pipeline touch_session): + # isolates the label-seed guard in _handle_end. + services = neo4j_services + handler = SessionHandler(services) + child_id = "iso-child" + + await handler( + "session:start", + {"session_id": child_id, "parent_id": "iso-parent", "timestamp": _TS}, + ) + await services.graph.flush() + + lbls_after_start = await _neo4j_labels(services, child_id) + assert "SubSession" in lbls_after_start, lbls_after_start + + await handler("session:end", {"session_id": child_id, "timestamp": _TS2}) + await handler("session:end", {"session_id": child_id, "timestamp": _TS3}) + await services.graph.flush() + + terminals = _terminals(await _neo4j_labels(services, child_id)) + assert terminals == ["SubSession"], terminals + + +@pytest.mark.timeout(60) +async def test_no_flush_outside_write_semaphore( + neo4j_container: dict[str, Any], tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + reg = _build_registry(tmp_path / "queues", write_concurrency=1) + qm = reg.queue_manager + sid = "fence" + await qm.append( + sid, _line("session:end", _WS, {"session_id": sid, "timestamp": _TS}) + ) + + worker = _build_worker(neo4j_container, sid) + reg._register_for_test(worker) + + observations: list[bool] = [] + real_flush = worker.services.graph.flush + + async def _spy_flush() -> None: + # The teardown close()->flush() is un-gated by construction; only the + # flushes during live drain must hold the semaphore. + if not worker.store_closed: + observations.append(reg.write_semaphore.locked()) + await real_flush() + + monkeypatch.setattr(worker.services.graph, "flush", _spy_flush) + + reg.start_drain(worker) + assert worker.task is not None + await asyncio.wait_for(worker.task, timeout=30.0) + + assert observations, "no flush was recorded, so the test proves nothing" + assert all(observations), f"a flush ran with the semaphore unlocked: {observations}" + + +@pytest.mark.timeout(60) +async def test_terminal_data_durably_flushed_before_log_deleted( + neo4j_container: dict[str, Any], tmp_path: Path +) -> None: + reg = _build_registry(tmp_path / "queues", write_concurrency=8) + qm = reg.queue_manager + sid = "durability" + await qm.append( + sid, + _line( + "session:start", + _WS, + {"session_id": sid, "parent_id": "parent", "timestamp": _TS}, + ), + ) + await qm.append( + sid, _line("session:end", _WS, {"session_id": sid, "timestamp": _TS2}) + ) + + log_path = tmp_path / "queues" / f"{sid}.log" + offset_path = tmp_path / "queues" / f"{sid}.offset" + assert log_path.stat().st_size > 0 + + worker = _build_worker(neo4j_container, sid) + reg._register_for_test(worker) + reg.start_drain(worker) + assert worker.task is not None + await asyncio.wait_for(worker.task, timeout=30.0) + + # log + offset gone => finalize completed, which only deletes after the + # tail was drained and flushed. + assert not log_path.exists(), "log not deleted: finalize did not complete" + assert not offset_path.exists(), "offset not deleted: finalize did not complete" + + final_labels, props = await _neo4j_labels_and_props_via_container( + neo4j_container, _WS, sid + ) + assert "SubSession" in final_labels, final_labels + assert props.get("status") == "completed", props + + # ended_at returns as neo4j.time.DateTime from a raw session.run. + ended_at = props.get("ended_at") + _to_native = getattr(ended_at, "to_native", None) + ended_at_native: Any = _to_native() if callable(_to_native) else ended_at + assert ended_at_native == datetime.fromisoformat(_TS2), props diff --git a/tests/neo4j/test_oom_regression.py b/tests/neo4j/test_oom_regression.py index b7fd2e1b..fad5e0e8 100644 --- a/tests/neo4j/test_oom_regression.py +++ b/tests/neo4j/test_oom_regression.py @@ -270,6 +270,10 @@ async def test_chunked_flush_drains_same_single_phase_buffer( # --------------------------------------------------------------------------- +# The deterministic three-leg OOM recipe runs ~30s against the capped +# container -- over the global 30s timeout once teardown is added. Give it +# headroom; the global default stays 30s for every other test. +@pytest.mark.timeout(180) async def test_finalization_path_freezes_then_restart_then_drains( neo4j_container_capped: dict[str, Any], caplog: pytest.LogCaptureFixture, diff --git a/tests/neo4j/test_orphan_visibility.py b/tests/neo4j/test_orphan_visibility.py index a2ced1e0..1926e708 100644 --- a/tests/neo4j/test_orphan_visibility.py +++ b/tests/neo4j/test_orphan_visibility.py @@ -266,14 +266,16 @@ async def test_finalization_orphan_surfaces_on_status( "not plain logger.error without exc_info" ) - # 5. Committed offset frozen at the pre-terminal boundary, NOT at tail_end. - # The drain committed the first batch (lines 1-100) but _finalize_session - # returned early without committing the tail (lines 101-200). + # 5. Committed offset is frozen AT session:end's own start (not tail_end): + # the drain commits UP TO session:end, so an unfinalized session stays re-derivable. + terminal_start = first_100.records[-1].start post_drain_batch = await qm.read_batch(sid, 1) committed_offset = post_drain_batch.start_offset - assert committed_offset == boundary, ( - f"Committed offset {committed_offset} must equal boundary {boundary} " - "(drain committed first batch, OOM froze the tail)" + assert committed_offset == terminal_start, ( + f"Committed offset {committed_offset} must equal terminal_start " + f"{terminal_start} (drain committed first batch UP TO session:end, " + "OOM froze the tail -- the offset is parked ON the terminal record " + "so an unfinalized session is durably re-derivable)" ) assert committed_offset != tail_end, ( f"Committed offset {committed_offset} must NOT equal tail_end {tail_end} " diff --git a/tests/test_concurrent_append.py b/tests/test_concurrent_append.py new file mode 100644 index 00000000..9cd6f248 --- /dev/null +++ b/tests/test_concurrent_append.py @@ -0,0 +1,120 @@ +"""Concurrency-correctness proof for QueueManager.append. + +Independent of filesystem append atomicity: correctness rests on +``_KeyGuard.file_lock`` holding across one whole record write. These tests +hammer real on-disk queues with high concurrency (many sessions, many +concurrent writers per session, mixed small/>1 MiB records) and prove every +record survives exactly once, complete, untorn, unmerged. +""" + +from __future__ import annotations + +import asyncio +import hashlib +import json +import os +from pathlib import Path + +from context_intelligence_server.queue_manager import QueueManager + +_SMALL_SIZE = 64 +_LARGE_SIZE = 1_500_000 # > 1 MiB, mixed in with small records +_LARGE_EVERY = 17 # every Nth record (by seq) is oversized + + +def _payload(size_bytes: int, session_id: str, seq: int) -> str: + """Unique-per-record filler; random bytes hex-encoded (no control chars).""" + random_part = os.urandom(max(size_bytes // 2, 8)).hex() + return f"{session_id}:{seq}:{random_part}" + + +def _make_record(session_id: str, seq: int, *, large: bool) -> bytes: + size = _LARGE_SIZE if large else _SMALL_SIZE + payload = _payload(size, session_id, seq) + digest = hashlib.sha256(payload.encode("utf-8")).hexdigest() + obj = {"session_id": session_id, "seq": seq, "payload": payload, "sha256": digest} + return json.dumps(obj, separators=(",", ":")).encode("utf-8") + + +async def _append_range(qm: QueueManager, session_id: str, count: int) -> int: + """Concurrently append `count` uniquely-numbered records; return bytes written.""" + lines = [ + _make_record(session_id, seq, large=(seq % _LARGE_EVERY == 0)) + for seq in range(count) + ] + await asyncio.gather(*(qm.append(session_id, line) for line in lines)) + return sum(len(line) + 1 for line in lines) # +1 per newline terminator + + +def _read_all_lines(path: Path) -> list[bytes]: + """Split a `.log` on newlines, asserting no torn (unterminated) tail.""" + data = path.read_bytes() + assert data.endswith(b"\n"), f"{path}: torn tail -- file does not end on \\n" + lines = data.split(b"\n") + assert lines[-1] == b"" # split() artifact after the trailing terminator + return lines[:-1] + + +def _verify_records(lines: list[bytes]) -> set[tuple[str, int]]: + """Parse every line as exactly one JSON record; return the (session_id, seq) set. + + A merged line (two records concatenated with no newline between them) + fails json.loads with "Extra data"; a torn line fails with a decode + error -- both are zero-tolerance failures here. + """ + seen: set[tuple[str, int]] = set() + for line in lines: + obj = json.loads(line) + digest = hashlib.sha256(obj["payload"].encode("utf-8")).hexdigest() + assert digest == obj["sha256"], "payload hash mismatch -- corrupted record" + key = (obj["session_id"], obj["seq"]) + assert key not in seen, f"duplicate record {key}" + seen.add(key) + return seen + + +async def test_concurrent_appends_many_sessions_no_tear_or_merge_or_loss( + tmp_path: Path, +) -> None: + """>=8 sessions x >=50 records each, all interleaved concurrently, plus + concurrent appends to the SAME session and several >1 MiB payloads + mixed with small ones.""" + qm = QueueManager(queues_dir=tmp_path / "queues") + num_sessions = 10 + records_per_session = 60 + + session_ids = [f"session-{i}" for i in range(num_sessions)] + written = await asyncio.gather( + *(_append_range(qm, sid, records_per_session) for sid in session_ids) + ) + + total_records = 0 + total_bytes = 0 + for sid in session_ids: + log_path = tmp_path / "queues" / f"{sid}.log" + lines = _read_all_lines(log_path) + seen = _verify_records(lines) + assert seen == {(sid, seq) for seq in range(records_per_session)} + assert len(lines) == records_per_session + total_records += len(lines) + total_bytes += log_path.stat().st_size + + assert total_records == num_sessions * records_per_session + assert total_bytes == sum(written) + + +async def test_concurrent_appends_single_session_hammered(tmp_path: Path) -> None: + """Worst-case contention: many concurrent tasks writing ONE session's file.""" + qm = QueueManager(queues_dir=tmp_path / "queues") + session_id = "hot-session" + num_records = 300 + + written = await _append_range(qm, session_id, num_records) + + log_path = tmp_path / "queues" / f"{session_id}.log" + lines = _read_all_lines(log_path) + seen = _verify_records(lines) + + assert seen == {(session_id, seq) for seq in range(num_records)} + assert len(lines) == num_records + assert log_path.stat().st_size == written diff --git a/tests/test_drain_lifecycle_logging.py b/tests/test_drain_lifecycle_logging.py new file mode 100644 index 00000000..92b78319 --- /dev/null +++ b/tests/test_drain_lifecycle_logging.py @@ -0,0 +1,370 @@ +"""Logging-completeness tests: each test asserts that a specific +drain/session-lifecycle event emits exactly one structured log line (right +level, `session=`/`reason=` tokens, session id promoted via `extra`). No +test here asserts on behavior, metrics, or /status -- other files own that. + +Harness: reuses the fakes/helpers from tests/test_drain_supervision.py. +No real Neo4j is used anywhere in this file. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import errno +import logging +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, patch + +import context_intelligence_server.main as main_module +import pytest +from context_intelligence_server.queue_manager import QueueManager +from context_intelligence_server.registry import SessionRegistry + +from tests.test_drain_supervision import ( + _accumulate, + _cancel_and_await, + _FlakyGraph, + _line, + _make_worker, + _pump, + _start_supervised, +) + +pytestmark = pytest.mark.integration + +LOGGER_NAME = "context_intelligence_server" + + +def _has( + caplog: pytest.LogCaptureFixture, + *, + level: int, + contains: list[str], + session_id: str | None = None, +) -> bool: + """True iff some captured record is at ``level`` and its message contains + every string in ``contains`` (and, if given, carries ``session_id`` via + the JsonFormatter-promoted ``extra`` field).""" + for r in caplog.records: + if r.levelno != level: + continue + msg = r.getMessage() + if not all(token in msg for token in contains): + continue + if session_id is not None and getattr(r, "session_id", None) != session_id: + continue + return True + return False + + +# --------------------------------------------------------------------------- +# drain_worker CANCELLED (two distinct sites) +# --------------------------------------------------------------------------- + + +class TestG1DrainWorkerCancelled: + async def test_g1a_cancelled_during_dispatch_logs_info_site_dispatch( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """Cancel while inside the inner try (dispatch/flush) -- the + registry.py inner ``except asyncio.CancelledError`` block.""" + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d3-g1-dispatch" + graph = _FlakyGraph() + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + started = asyncio.Event() + release = asyncio.Event() + + async def _blocking_process( + worker: object, event: str, data: object, handlers: object + ) -> None: + started.set() + await release.wait() # never set -- cancellation always wins here + + await qm.append(sid, _line("e1", "/ws", {"session_id": sid})) + + with ( + patch( + "context_intelligence_server.registry.process_event", + side_effect=_blocking_process, + ), + caplog.at_level(logging.INFO, logger=LOGGER_NAME), + ): + task = _start_supervised(reg, worker) + await started.wait() # deterministic: task is now inside dispatch + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + await _pump() + + assert _has( + caplog, + level=logging.INFO, + contains=["drain_worker_cancelled", "site=dispatch"], + session_id=sid, + ), [r.getMessage() for r in caplog.records] + assert worker.store_closed is True + assert sid not in reg.active_sessions() + + async def test_g1b_cancelled_while_idle_logs_info_site_loop( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """Cancel while idle (no data, polling) -- the registry.py OUTER + ``except asyncio.CancelledError`` block (never the inner one, since + an empty batch never reaches the dispatch/flush try).""" + reg = SessionRegistry() + sid = "d3-g1-loop" + graph = _FlakyGraph() + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + with ( + patch( + "context_intelligence_server.registry.process_event", + side_effect=_accumulate, + ), + caplog.at_level(logging.INFO, logger=LOGGER_NAME), + ): + task = _start_supervised(reg, worker) + await asyncio.sleep(0) + await asyncio.sleep(0) # let it settle into the idle poll + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + await _pump() + + assert _has( + caplog, + level=logging.INFO, + contains=["drain_worker_cancelled", "site=loop"], + session_id=sid, + ), [r.getMessage() for r in caplog.records] + assert worker.store_closed is True + assert sid not in reg.active_sessions() + + +# --------------------------------------------------------------------------- +# remove() cancelling a still-live drain task +# --------------------------------------------------------------------------- + + +class TestG2Remove: + async def test_g2_remove_live_task_logs_info( + self, caplog: pytest.LogCaptureFixture + ) -> None: + reg = SessionRegistry() + sid = "d3-g2-remove" + graph = _FlakyGraph() + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + with ( + patch( + "context_intelligence_server.registry.process_event", + side_effect=_accumulate, + ), + caplog.at_level(logging.INFO, logger=LOGGER_NAME), + ): + task = _start_supervised(reg, worker) + await asyncio.sleep(0) + assert not task.done(), "the drainer must still be live for this test" + reg.remove(sid) + with contextlib.suppress(asyncio.CancelledError): + await task + await _pump() + + assert _has( + caplog, + level=logging.INFO, + contains=["drain_worker_remove", "had_live_task=True"], + session_id=sid, + ), [r.getMessage() for r in caplog.records] + assert sid not in reg.active_sessions() + + +# --------------------------------------------------------------------------- +# start_drain() respawning an already-registered worker +# --------------------------------------------------------------------------- + + +class TestG3Respawn: + async def test_g3_start_drain_respawn_logs_info( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """A worker whose previous task is done (cancelled, not crashed) is + respawned: only this case should log ``drainer_respawned``, not the + brand-new-spawn ``drainer_spawned``.""" + reg = SessionRegistry() + sid = "d3-g3-respawn" + graph = _FlakyGraph() + worker = _make_worker(sid, graph) + + async def _noop() -> None: + return None + + old_task = asyncio.create_task(_noop()) + await old_task # done, cancelled() is False, exception() is None + worker.task = old_task + + with ( + patch( + "context_intelligence_server.registry.process_event", + side_effect=_accumulate, + ), + caplog.at_level(logging.INFO, logger=LOGGER_NAME), + ): + reg.start_drain(worker) + assert worker.task is not old_task, ( + "start_drain must build a fresh task for a done-but-not-crashed worker" + ) + new_task = worker.task + assert new_task is not None + await _cancel_and_await(new_task) + + assert _has( + caplog, + level=logging.INFO, + contains=["drainer_respawned"], + session_id=sid, + ), [r.getMessage() for r in caplog.records] + # No behavior change: drainer_spawned (the pre-existing brand-new-spawn + # log) must NOT have fired for this respawn. + assert not any("drainer_spawned" in r.getMessage() for r in caplog.records) + + +# --------------------------------------------------------------------------- +# _finalize_session leaving an orphaned worker (two sites) +# --------------------------------------------------------------------------- + + +class TestG4FinalizeOrphan: + async def test_g4a_first_pass_tail_flush_failed_logs_warning( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """The FIRST _drain_to_eof call inside _finalize_session fails -> + the recoverable orphan (a respawn/next-drain retries finalize).""" + reg = SessionRegistry() + sid = "d3-g4-tail-flush-failed" + graph = _FlakyGraph() + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + handlers = object() + with ( + patch.object(reg, "_drain_to_eof", AsyncMock(return_value=False)), + caplog.at_level(logging.WARNING, logger=LOGGER_NAME), + ): + await reg._finalize_session(worker, handlers) + + assert _has( + caplog, + level=logging.WARNING, + contains=[ + "finalize_orphan", + "reason=tail_flush_failed", + "recoverable=respawn", + ], + session_id=sid, + ), [r.getMessage() for r in caplog.records] + # No behavior change: still registered, store not closed (respawn will retry). + assert sid in reg.active_sessions() + assert graph.closed is False + + async def test_g4b_delete_retry_exhausted_permanent_orphan_logs_error( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """delete_drained refuses every attempt and the late-tail re-drain + also fails -> the permanent-retention orphan.""" + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d3-g4-permanent" + graph = _FlakyGraph() + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + handlers = object() + with ( + patch.object(reg, "_drain_to_eof", AsyncMock(side_effect=[True, False])), + patch.object(qm, "delete_drained", AsyncMock(return_value=False)), + caplog.at_level(logging.WARNING, logger=LOGGER_NAME), + ): + await reg._finalize_session(worker, handlers) + + assert _has( + caplog, + level=logging.ERROR, + contains=[ + "finalize_orphan", + "reason=delete_retry_exhausted", + "permanent=true", + ], + session_id=sid, + ), [r.getMessage() for r in caplog.records] + assert sid in reg.active_sessions(), "permanent orphan stays registered" + assert graph.closed is False + + +# --------------------------------------------------------------------------- +# QueueManager.dead_letter()'s own write failing +# --------------------------------------------------------------------------- + + +class TestG5DeadLetterWriteFailure: + async def test_g5_dead_letter_write_oserror_logs_error_and_reraises( + self, + caplog: pytest.LogCaptureFixture, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + qm = QueueManager(queues_dir=tmp_path / "queues") + sid = "d3-g5-dead-letter" + injected = OSError(errno.EIO, "Input/output error") + monkeypatch.setattr(qm, "_write_record", MagicMock(side_effect=injected)) + + with ( + caplog.at_level(logging.ERROR, logger=LOGGER_NAME), + pytest.raises(OSError) as ei, + ): + await qm.dead_letter(sid, b"bad-line", "boom") + + assert ei.value is injected, "propagation must be the SAME exception object" + matches = [ + r + for r in caplog.records + if r.levelno == logging.ERROR + and "dead_letter_write_failed" in r.getMessage() + ] + assert matches, [r.getMessage() for r in caplog.records] + rec = matches[0] + assert sid in rec.getMessage() + assert rec.exc_info is not None, "dead_letter_write_failed must carry exc_info" + + +# --------------------------------------------------------------------------- +# Crash-recovery topup: recover()/read_batch disagreement +# --------------------------------------------------------------------------- + + +class TestG8RecoverySkippedEmptyBatch: + async def test_g8_recover_reports_session_but_read_batch_empty_logs_warning( + self, caplog: pytest.LogCaptureFixture + ) -> None: + sid = "d3-g8-empty-batch" + qm = main_module.registry.queue_manager + with ( + patch.object(qm, "recover", AsyncMock(return_value=[sid])), + caplog.at_level(logging.WARNING, logger=LOGGER_NAME), + ): + result = await main_module._crash_recovery_topup(None) + + assert result == 0 + assert any( + r.levelno == logging.WARNING + and "recovery_skipped_empty_batch" in r.getMessage() + and sid in r.getMessage() + for r in caplog.records + ), [r.getMessage() for r in caplog.records] diff --git a/tests/test_drain_supervision.py b/tests/test_drain_supervision.py new file mode 100644 index 00000000..3c439972 --- /dev/null +++ b/tests/test_drain_supervision.py @@ -0,0 +1,1098 @@ +"""Drain supervision + offset-ownership tests. + +Verifies that an unexpected exception raised inside ``drain_worker`` is +never silent: ``add_done_callback``/``_on_drain_done`` ensures it is +logged, the store is closed, and the worker is deregistered so a fresh +one can be created. No real Neo4j is used anywhere in this file. + +Determinism rule: no wall-clock sleeps as synchronisation -- every wait is +either a bounded poll on an observable condition, or an ``asyncio.Event`` +the injected fake sets right before blocking. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import errno +import functools +import json +import logging +from collections.abc import Awaitable, Callable +from typing import Any +from unittest.mock import patch + +import neo4j.exceptions as neo4j_exc +import pytest + +from context_intelligence_server.queue_manager import QueueManager +from context_intelligence_server.registry import SessionRegistry, SessionWorker +from context_intelligence_server.services import HookStateService + +pytestmark = pytest.mark.integration + + +# --------------------------------------------------------------------------- +# Wire format + fakes +# --------------------------------------------------------------------------- + + +def _line(event: str, workspace: str, data: dict) -> bytes: + """Encode an appended event line exactly as POST /events stores it + (mirrors ``tests/test_registry.py::_line``).""" + return json.dumps({"event": event, "workspace": workspace, "data": data}).encode( + "utf-8" + ) + + +class _FlakyGraph: + """A faithful accumulating-buffer graph fake (NOT a hollow mock). + + Mirrors ``tests/test_registry.py::_AccumBufferGraph`` / + ``tests/test_large_event_tail_drop.py::_FaultInjectableGraph``: + writes accumulate in ``buffer`` until ``flush()`` succeeds, at which + point they move into ``flushed`` -- a SET, modeling a real store's + idempotent id-keyed MERGE (replaying the same event after a respawn + must never show up twice). ``fail_when`` decides whether ``flush()`` + raises for the CURRENT buffer contents; the default rejects only + multi-event batches (a single isolated line always succeeds), which is + what drives a batch through retries -> exhaustion -> per-line isolation + without every individual line being unprocessable. + """ + + def __init__(self, *, fail_when: Callable[[set[str]], bool] | None = None) -> None: + self.workspace = "/ws" + self.created_by: str | None = None + self.buffer: set[str] = set() + self.flushed: set[str] = set() + self.discards = 0 + self.closed = False + self._fail_when = fail_when or (lambda buf: len(buf) > 1) + + async def flush(self) -> None: + if not self.buffer: + return # empty-buffer early return (mirrors neo4j_store.py:1501-1502) + if self._fail_when(self.buffer): + raise RuntimeError(f"flush rejected for buffer={sorted(self.buffer)}") + self.flushed |= self.buffer + self.buffer.clear() # success clears + + def discard_buffer(self) -> None: + self.buffer.clear() + self.discards += 1 + + async def close(self) -> None: + self.closed = True + + +class _SequencedFlushGraph: + """flush() raises each exception in ``sequence`` in order, then succeeds + forever after. Used only for the real-neo4j-exception-type test.""" + + def __init__(self, sequence: list[BaseException]) -> None: + self.workspace = "/ws" + self.created_by: str | None = None + self.buffer: set[str] = set() + self.flushed: set[str] = set() + self.discards = 0 + self.closed = False + self._sequence = list(sequence) + self._calls = 0 + + async def flush(self) -> None: + if self._calls < len(self._sequence): + exc = self._sequence[self._calls] + self._calls += 1 + raise exc + self._calls += 1 + self.flushed |= self.buffer + self.buffer.clear() + + def discard_buffer(self) -> None: + self.buffer.clear() + self.discards += 1 + + async def close(self) -> None: + self.closed = True + + +async def _accumulate( + worker: SessionWorker, event: str, data: object, handlers: object +) -> None: + """Stand-in for ``process_event``: buffers the event name on the fake + graph, exactly like ``_FaultInjectableGraph``'s harness in the sibling + large-event tail-drop test file.""" + worker.services.graph.buffer.add(event) + + +def _make_worker(sid: str, graph: Any, workspace: str = "/ws") -> SessionWorker: + worker = SessionWorker( + session_id=sid, + workspace=workspace, + services=HookStateService(workspace=workspace), + ) + worker.services.graph = graph # type: ignore[assignment] + return worker + + +def _flaky( + original: Callable[..., Awaitable[Any]], + exc: BaseException, + *, + on_call: int = 1, +) -> Callable[..., Awaitable[Any]]: + """Return an async wrapper around ``original`` that raises ``exc`` on + its ``on_call``-th invocation (1-based) and delegates to ``original`` + for every other call. This is the fault-injection primitive used by + every site test below -- always a REAL exception type, never a bare + ``Exception()``, so a test can never pass by accident on an + over-broad except clause.""" + state = {"n": 0} + + async def _wrapper(*args: Any, **kwargs: Any) -> Any: + state["n"] += 1 + if state["n"] == on_call: + raise exc + return await original(*args, **kwargs) + + return _wrapper + + +def _start_supervised( + reg: SessionRegistry, worker: SessionWorker, *, flush_timeout: float = 10.0 +) -> asyncio.Task: + """Mirror production ``start_drain`` (registry.py) EXACTLY: create the + task, attach the done-callback, bind ``worker.task``. Needed because + this file drives ``drain_worker`` directly (for injection control) the + same way ``tests/test_large_event_tail_drop.py::_drive_drain_to_quiescence`` + does, and production's supervision is only real if the binding matches + production's own ``start_drain``.""" + task = asyncio.create_task( + reg.drain_worker(worker, flush_timeout=flush_timeout), + name=f"drain-{worker.session_id}", + ) + task.add_done_callback(functools.partial(reg._on_drain_done, worker)) + worker.task = task + return task + + +async def _pump(n: int = 5) -> None: + """Let ``n`` event-loop iterations pass -- long enough for a + ``call_soon``-scheduled done-callback to actually run.""" + for _ in range(n): + await asyncio.sleep(0) + + +async def _await_death(task: asyncio.Task) -> None: + """Wait for ``task`` to finish (absorbing whatever it raised), then pump + the loop so its done-callback has actually executed before we assert + anything about its effects.""" + with contextlib.suppress(BaseException): + await task + await _pump() + + +async def _drain_until_idle( + reg: SessionRegistry, + qm: QueueManager, + worker: SessionWorker, + sid: str, + *, + max_polls: int = 400, + poll_sleep: float = 0.01, +) -> asyncio.Task: + """Poll (bounded, never a bare sleep as the only wait) until EITHER the + task finishes on its own OR the queue is fully drained -- mirrors + ``tests/test_large_event_tail_drop.py::_drive_drain_to_quiescence``.""" + task = worker.task + assert task is not None + for _ in range(max_polls): + await asyncio.sleep(poll_sleep) + if task.done(): + break + if (await qm.read_batch(sid, 10)).lines == []: + break + return task + + +async def _cancel_and_await(task: asyncio.Task | None) -> None: + assert task is not None + if not task.done(): + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + + +def _died_loudly( + caplog: pytest.LogCaptureFixture, sid: str, exc_type: type[BaseException] +) -> logging.LogRecord: + """Assert (a): a ``drain_worker_died`` ERROR was logged, carrying the + session id (in both the message and ``extra``) and the injected + exception's real type via ``exc_info``. Returns the matching record.""" + matches = [ + r + for r in caplog.records + if r.levelno == logging.ERROR + and "drain_worker_died" in r.getMessage() + and sid in r.getMessage() + and getattr(r, "session_id", None) == sid + ] + assert matches, ( + f"expected a drain_worker_died ERROR with session_id={sid!r}; " + f"caplog had: {[r.getMessage() for r in caplog.records]}" + ) + rec = matches[0] + assert rec.exc_info is not None, "drain_worker_died must carry exc_info" + actual_type = rec.exc_info[0] + assert actual_type is not None and issubclass(actual_type, exc_type), ( + f"expected exc_info type {exc_type}, got {actual_type}" + ) + return rec + + +# --------------------------------------------------------------------------- +# Injection matrix: one test per unguarded failure site. +# +# Every test asserts: +# (a) not silently dead -- drain_worker_died ERROR w/ session id + exc_info +# (b) store closed -- worker.store_closed is True, fake.closed is True +# (c) deregistered+respawn drains the exact remaining suffix, no gap/dup +# (d) finalize re-runs to completion (terminal-path tests only) +# (e) no second live drainer +# --------------------------------------------------------------------------- + + +class TestInjectionMatrix: + async def test_read_batch_failure_is_supervised_and_respawns( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """qm.read_batch raises OSError(EIO) once.""" + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d1-s1-read-batch" + graph = _FlakyGraph() + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + qm.read_batch = _flaky( # type: ignore[method-assign] + qm.read_batch, OSError(errno.EIO, "Input/output error") + ) + + written_before = reg.pipeline_counters()["written_total"] + + with ( + patch( + "context_intelligence_server.registry.process_event", + side_effect=_accumulate, + ), + caplog.at_level(logging.ERROR, logger="context_intelligence_server"), + ): + await qm.append(sid, _line("e1", "/ws", {"session_id": sid})) + task = _start_supervised(reg, worker) + await _await_death(task) + + _died_loudly(caplog, sid, OSError) + assert worker.store_closed is True + assert graph.closed is True + assert sid not in reg.active_sessions() + + # (c) respawn: a fresh worker over the SAME queue + SAME fake + # (the fake models the accumulating write buffer; flushed is a + # SET, so a replayed line can never show up twice). + worker2 = _make_worker(sid, graph) + reg._register_for_test(worker2) + reg.start_drain(worker2) + await _drain_until_idle(reg, qm, worker2, sid) + await _cancel_and_await(worker2.task) + + assert graph.flushed == {"e1"} + assert reg.pipeline_counters()["written_total"] == written_before + 1 + assert (await qm.read_batch(sid, 10)).lines == [] + assert worker2.task is not task + + async def test_commit_failure_is_supervised_and_respawns( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """qm.commit raises OSError(ESTALE) once. No duplicate node in + fake.flushed after the replay (flushed is a set), and written_total + counts the line exactly once.""" + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d1-s2-commit" + graph = _FlakyGraph() + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + qm.commit = _flaky( # type: ignore[method-assign] + qm.commit, OSError(errno.ESTALE, "Stale file handle") + ) + + written_before = reg.pipeline_counters()["written_total"] + + with ( + patch( + "context_intelligence_server.registry.process_event", + side_effect=_accumulate, + ), + caplog.at_level(logging.ERROR, logger="context_intelligence_server"), + ): + await qm.append(sid, _line("e1", "/ws", {"session_id": sid})) + task = _start_supervised(reg, worker) + await _await_death(task) + + _died_loudly(caplog, sid, OSError) + assert worker.store_closed is True + assert sid not in reg.active_sessions() + + worker2 = _make_worker(sid, graph) # SAME fake -- dedup proof + reg._register_for_test(worker2) + reg.start_drain(worker2) + await _drain_until_idle(reg, qm, worker2, sid) + await _cancel_and_await(worker2.task) + + assert graph.flushed == {"e1"}, "no duplicate: flushed is a set" + assert reg.pipeline_counters()["written_total"] == written_before + 1 + assert (await qm.read_batch(sid, 10)).lines == [] + + async def test_dead_letter_failure_no_longer_kills_the_drainer( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """qm.dead_letter raises OSError(EIO) once, in the dead-letter + except-clause. The poison line ends up in read_dead_letters and + tail-1/tail-2 are persisted after respawn.""" + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d1-s3-dead-letter" + graph = _FlakyGraph(fail_when=lambda buf: "oversized" in buf) + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + qm.dead_letter = _flaky( # type: ignore[method-assign] + qm.dead_letter, OSError(errno.EIO, "disk unavailable") + ) + + with ( + patch( + "context_intelligence_server.registry.process_event", + side_effect=_accumulate, + ), + caplog.at_level(logging.ERROR, logger="context_intelligence_server"), + ): + await qm.append(sid, _line("small-1", "/ws", {"session_id": sid})) + await qm.append(sid, _line("small-2", "/ws", {"session_id": sid})) + await qm.append(sid, _line("oversized", "/ws", {"session_id": sid})) + await qm.append(sid, _line("tail-1", "/ws", {"session_id": sid})) + await qm.append(sid, _line("tail-2", "/ws", {"session_id": sid})) + + task = _start_supervised(reg, worker) + await _await_death(task) + + _died_loudly(caplog, sid, OSError) + assert worker.store_closed is True + assert sid not in reg.active_sessions() + + worker2 = _make_worker(sid, graph) + reg._register_for_test(worker2) + reg.start_drain(worker2) + await _drain_until_idle(reg, qm, worker2, sid) + await _cancel_and_await(worker2.task) + + dead = await qm.read_dead_letters(sid) + assert len(dead) == 1 + assert json.loads(dead[0]["payload"])["event"] == "oversized" + assert graph.flushed == {"small-1", "small-2", "tail-1", "tail-2"} + assert (await qm.read_batch(sid, 10)).lines == [] + + async def test_isolation_commit_failure_is_supervised_and_respawns( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """qm.commit raises TimeoutError once on the first isolation-path + commit call. No line is dead-lettered twice, none is lost.""" + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d1-s4-isolation-commit" + graph = _FlakyGraph() # batch of 2 forces exhaustion -> isolation + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + qm.commit = _flaky( # type: ignore[method-assign] + qm.commit, TimeoutError("SMB operation timed out") + ) + + written_before = reg.pipeline_counters()["written_total"] + + with ( + patch( + "context_intelligence_server.registry.process_event", + side_effect=_accumulate, + ), + caplog.at_level(logging.ERROR, logger="context_intelligence_server"), + ): + await qm.append(sid, _line("g1", "/ws", {"session_id": sid})) + await qm.append(sid, _line("g2", "/ws", {"session_id": sid})) + + task = _start_supervised(reg, worker) + await _await_death(task) + + _died_loudly(caplog, sid, TimeoutError) + assert worker.store_closed is True + assert sid not in reg.active_sessions() + + worker2 = _make_worker(sid, graph) + reg._register_for_test(worker2) + reg.start_drain(worker2) + await _drain_until_idle(reg, qm, worker2, sid) + await _cancel_and_await(worker2.task) + + assert graph.flushed == {"g1", "g2"} + assert reg.pipeline_counters()["written_total"] == written_before + 2 + dead = await qm.read_dead_letters(sid) + assert dead == [], "neither line is poison -- nothing should be dead-lettered" + assert (await qm.read_batch(sid, 10)).lines == [] + + async def test_finalize_tail_read_failure_refinalizes_after_respawn( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """qm.read_batch raises OSError(ESTALE) once, on _finalize_session's + own tail read (its second call). The terminal batch was already + committed up to session:end, so finalize re-runs after respawn to + full completion (assert (d)).""" + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d1-s5-finalize-read" + # A clean (always-succeeds) graph: this test targets read_batch, not + # flush -- a batch-size-sensitive fake would force exhaustion on the + # very first (2-record) batch and never reach the terminal batch's + # commit at all. + graph = _FlakyGraph(fail_when=lambda buf: False) + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + qm.read_batch = _flaky( # type: ignore[method-assign] + qm.read_batch, OSError(errno.ESTALE, "Stale file handle"), on_call=2 + ) + + with ( + patch( + "context_intelligence_server.registry.process_event", + side_effect=_accumulate, + ), + caplog.at_level(logging.ERROR, logger="context_intelligence_server"), + ): + await qm.append(sid, _line("tool:pre", "/ws", {"session_id": sid})) + await qm.append(sid, _line("session:end", "/ws", {"session_id": sid})) + + task = _start_supervised(reg, worker) + await _await_death(task) + + _died_loudly(caplog, sid, OSError) + assert worker.store_closed is True + assert sid not in reg.active_sessions() + + worker2 = _make_worker(sid, graph) + reg._register_for_test(worker2) + reg.start_drain(worker2) + assert worker2.task is not None + await asyncio.wait_for(worker2.task, timeout=5.0) + + # (d) finalize re-ran to completion. + assert len(reg.completed_sessions()) == 1 + assert graph.flushed == {"tool:pre", "session:end"} + assert sid not in reg.active_sessions() + + async def test_finalize_tail_commit_failure_refinalizes_after_respawn( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """qm.commit raises OSError(ENOSPC) once, on _finalize_session's own + tail commit (its second call, after the terminal batch's own + up-to-session:end commit succeeds).""" + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d1-s6-finalize-commit" + # Clean graph: this test targets commit, not flush -- same reasoning + # as the previous test. + graph = _FlakyGraph(fail_when=lambda buf: False) + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + qm.commit = _flaky( # type: ignore[method-assign] + qm.commit, OSError(errno.ENOSPC, "No space left on device"), on_call=2 + ) + + with ( + patch( + "context_intelligence_server.registry.process_event", + side_effect=_accumulate, + ), + caplog.at_level(logging.ERROR, logger="context_intelligence_server"), + ): + await qm.append(sid, _line("tool:pre", "/ws", {"session_id": sid})) + await qm.append(sid, _line("session:end", "/ws", {"session_id": sid})) + + task = _start_supervised(reg, worker) + await _await_death(task) + + _died_loudly(caplog, sid, OSError) + assert worker.store_closed is True + assert sid not in reg.active_sessions() + + worker2 = _make_worker(sid, graph) + reg._register_for_test(worker2) + reg.start_drain(worker2) + assert worker2.task is not None + await asyncio.wait_for(worker2.task, timeout=5.0) + + assert len(reg.completed_sessions()) == 1 + assert graph.flushed == {"tool:pre", "session:end"} + assert sid not in reg.active_sessions() + + async def test_delete_drained_failure_is_supervised_without_second_drainer( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """qm.delete_drained raises PermissionError once. CompletedSession + was already recorded (appended before delete_drained runs); the + callback deregisters + closes; no second drainer.""" + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d1-s7-delete-drained" + graph = _FlakyGraph() + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + qm.delete_drained = _flaky( # type: ignore[method-assign] + qm.delete_drained, PermissionError(errno.EACCES, "Permission denied") + ) + + with ( + patch( + "context_intelligence_server.registry.process_event", + side_effect=_accumulate, + ), + caplog.at_level(logging.ERROR, logger="context_intelligence_server"), + ): + await qm.append(sid, _line("session:end", "/ws", {"session_id": sid})) + task = _start_supervised(reg, worker) + await _await_death(task) + + _died_loudly(caplog, sid, PermissionError) + assert worker.store_closed is True + assert graph.closed is True + assert sid not in reg.active_sessions() + assert len(reg.completed_sessions()) == 1, ( + "CompletedSession is appended BEFORE delete_drained; it must " + "survive delete_drained raising" + ) + # (e): no second live drainer exists for this session. + assert all( + w.session_id != sid or (w.task is None or w.task.done()) + for w in reg.workers() + ) + + async def test_prologue_exception_is_supervised( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """setup_handlers raises before the while loop even starts. A later + get_or_create-equivalent still builds a working worker.""" + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d1-s8-prologue" + graph = _FlakyGraph() + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + with ( + patch( + "context_intelligence_server.registry.setup_handlers", + side_effect=RuntimeError("handler wiring failed"), + ), + caplog.at_level(logging.ERROR, logger="context_intelligence_server"), + ): + task = _start_supervised(reg, worker) + await _await_death(task) + + _died_loudly(caplog, sid, RuntimeError) + assert worker.store_closed is True + assert graph.closed is True + assert sid not in reg.active_sessions() + + # A later worker (setup_handlers no longer patched) works normally. + worker2 = _make_worker(sid, graph) + reg._register_for_test(worker2) + with patch( + "context_intelligence_server.registry.process_event", + side_effect=_accumulate, + ): + await qm.append(sid, _line("e1", "/ws", {"session_id": sid})) + reg.start_drain(worker2) + await _drain_until_idle(reg, qm, worker2, sid) + await _cancel_and_await(worker2.task) + + assert graph.flushed == {"e1"} + + async def test_flush_failure_uses_real_neo4j_exception_types(self) -> None: + """Not a crash scenario -- proves the inner retry path tolerates + real neo4j driver exception types, not just a bare Exception.""" + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d1-flush-real-exceptions" + graph = _SequencedFlushGraph( + [ + neo4j_exc.ServiceUnavailable("db unreachable"), + neo4j_exc.TransientError("deadlock, retry"), + ] + ) + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + with patch( + "context_intelligence_server.registry.process_event", + side_effect=_accumulate, + ): + await qm.append(sid, _line("e1", "/ws", {"session_id": sid})) + reg.start_drain(worker) + await _drain_until_idle(reg, qm, worker, sid) + assert not worker.task.done(), ( # type: ignore[union-attr] + "within-budget retries must not kill the task" + ) + await _cancel_and_await(worker.task) # type: ignore[arg-type] + + assert graph.flushed == {"e1"} + assert reg.pipeline_counters()["write_retries_total"] >= 2 + assert (await qm.read_batch(sid, 10)).lines == [] + + +# --------------------------------------------------------------------------- +# Mechanism-specific tests +# --------------------------------------------------------------------------- + + +class TestMechanismSpecific: + async def test_committed_offset_freezes_AT_the_terminal_line(self) -> None: + """After the terminal batch commits, the first pending + record parses to session:end -- the offset is frozen AT the + boundary, not past it. _finalize_session is stubbed out so we can + inspect queue state before delete_drained would remove the log.""" + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d1-committed-at-terminal" + # Clean graph: this test is about the commit boundary, not flush + # failure -- a batch-size-sensitive fake would force the 2-record + # batch through poison isolation instead of a normal commit. + graph = _FlakyGraph(fail_when=lambda buf: False) + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + with ( + patch( + "context_intelligence_server.registry.process_event", + side_effect=_accumulate, + ), + patch.object(reg, "_finalize_session", autospec=True) as mock_finalize, + ): + mock_finalize.return_value = None + await qm.append(sid, _line("tool:pre", "/ws", {"session_id": sid})) + await qm.append(sid, _line("session:end", "/ws", {"session_id": sid})) + task = _start_supervised(reg, worker) + await asyncio.wait_for(task, timeout=5.0) + + mock_finalize.assert_awaited_once() + pending = await qm.read_batch(sid, 10) + assert len(pending.records) == 1 + event, _ws, _data = reg._parse_line(pending.records[0].raw) + assert event == "session:end" + + async def test_recover_reports_a_terminal_but_unfinalized_session(self) -> None: + """recover() reports a session frozen at its terminal line as + recoverable (committed < complete_data_end).""" + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d1-recover-terminal-unfinalized" + graph = _FlakyGraph(fail_when=lambda buf: False) + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + with ( + patch( + "context_intelligence_server.registry.process_event", + side_effect=_accumulate, + ), + patch.object(reg, "_finalize_session", autospec=True) as mock_finalize, + ): + mock_finalize.return_value = None + await qm.append(sid, _line("tool:pre", "/ws", {"session_id": sid})) + await qm.append(sid, _line("session:end", "/ws", {"session_id": sid})) + task = _start_supervised(reg, worker) + await asyncio.wait_for(task, timeout=5.0) + + recoverable = await qm.recover() + assert sid in recoverable + + async def test_finalize_reruns_to_completion_after_a_transient_finalize_failure( + self, + ) -> None: + """After a transient finalize-tail read failure and a respawn, the + session is fully finalized -- CompletedSession recorded, + delete_drained ran, every line persisted exactly once.""" + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d1-finalize-full-completion" + graph = _FlakyGraph(fail_when=lambda buf: False) + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + qm.read_batch = _flaky( # type: ignore[method-assign] + qm.read_batch, OSError(errno.ESTALE, "Stale file handle"), on_call=2 + ) + + with patch( + "context_intelligence_server.registry.process_event", + side_effect=_accumulate, + ): + await qm.append(sid, _line("tool:pre", "/ws", {"session_id": sid})) + await qm.append(sid, _line("session:end", "/ws", {"session_id": sid})) + task = _start_supervised(reg, worker) + await _await_death(task) + + worker2 = _make_worker(sid, graph) + reg._register_for_test(worker2) + reg.start_drain(worker2) + assert worker2.task is not None + await asyncio.wait_for(worker2.task, timeout=5.0) + + assert len(reg.completed_sessions()) == 1 + assert graph.flushed == {"tool:pre", "session:end"} + assert not qm._log_path(sid).exists(), "delete_drained must have run" + assert not qm._offset_path(sid).exists() + + async def test_no_second_drainer_during_the_finalize_window(self) -> None: + """While qm.delete_drained is parked, get_or_create must be a no-op + (no new task, worker.task unchanged); finalization then completes.""" + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d1-no-second-drainer" + graph = _FlakyGraph() + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + entered_delete = asyncio.Event() + release = asyncio.Event() + original_delete = qm.delete_drained + + async def _parked_delete(session_id: str) -> bool: + entered_delete.set() + await release.wait() + return await original_delete(session_id) + + qm.delete_drained = _parked_delete # type: ignore[method-assign] + + with patch( + "context_intelligence_server.registry.process_event", + side_effect=_accumulate, + ): + await qm.append(sid, _line("session:end", "/ws", {"session_id": sid})) + task = _start_supervised(reg, worker) + + await asyncio.wait_for(entered_delete.wait(), timeout=5.0) + assert sid in reg.active_sessions(), ( + "worker must still be registered while delete_drained is parked " + "(_deregister is the LAST act of finalization)" + ) + + pre_task = worker.task + reg.get_or_create(sid, "/ws") + assert worker.task is pre_task, ( + "a concurrent get_or_create during the finalize window must be " + "a no-op: the live task is not done() yet" + ) + + release.set() + await asyncio.wait_for(task, timeout=5.0) + + assert len(reg.completed_sessions()) == 1 + assert sid not in reg.active_sessions() + + +# --------------------------------------------------------------------------- +# Spent-worker guard: a cancelled worker must never be revived +# --------------------------------------------------------------------------- + + +class TestSpentWorkerGuard: + async def test_cancelled_worker_is_not_revived_through_a_closed_store( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """task.cancel() -- drain_worker swallows CancelledError and returns + cleanly, worker.store_closed is True, no drain_worker_died ERROR, + and start_drain(worker) creates no new task: the same worker object + must never be revived once its store is closed.""" + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d1-cancelled-not-revived" + graph = _FlakyGraph() + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + with ( + patch( + "context_intelligence_server.registry.process_event", + side_effect=_accumulate, + ), + caplog.at_level(logging.ERROR, logger="context_intelligence_server"), + ): + await qm.append(sid, _line("e1", "/ws", {"session_id": sid})) + task = _start_supervised(reg, worker) + await asyncio.sleep(0) # let it actually start running + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + await _pump() + + assert task.cancelled() is False, ( + "drain_worker catches CancelledError and returns -- the task " + "carries a clean result, not a cancellation" + ) + assert graph.closed is True + assert worker.store_closed is True + assert sid not in reg.active_sessions(), ( + "the cancellation handler must deregister, or this worker " + "is wedged (found by get_or_create, refused by start_drain, " + "forever)" + ) + assert not any("drain_worker_died" in r.getMessage() for r in caplog.records), ( + "a clean cancellation is not a crash -- no ERROR expected" + ) + + pre_task = worker.task + reg.start_drain(worker) # attempt to revive the SAME spent worker + assert worker.task is pre_task, ( + "start_drain must refuse to revive a store_closed worker" + ) + + +class TestPoisonLineIsolation: + async def test_poison_line_is_dead_lettered_and_advances_without_dying( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """A merged/unparseable line is dead-lettered via the isolation + path, the offset advances past it, the task stays alive, and no + drain_worker_died fires.""" + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d1-poison-isolated" + graph = _FlakyGraph() + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + with ( + patch( + "context_intelligence_server.registry.process_event", + side_effect=_accumulate, + ), + caplog.at_level(logging.ERROR, logger="context_intelligence_server"), + ): + await qm.append(sid, _line("good-1", "/ws", {"session_id": sid})) + await qm.append(sid, b"{ this is not valid json") + await qm.append(sid, _line("good-2", "/ws", {"session_id": sid})) + + task = _start_supervised(reg, worker) + await _drain_until_idle(reg, qm, worker, sid) + assert not task.done(), "the drainer must stay alive after isolation" + await _cancel_and_await(task) + + dead = await qm.read_dead_letters(sid) + assert len(dead) == 1 + assert graph.flushed == {"good-1", "good-2"} + assert not any("drain_worker_died" in r.getMessage() for r in caplog.records) + + +class TestNoDoubleCountOnReplay: + async def test_no_double_count_on_replay(self) -> None: + """Force the isolation-path commit to fail once. written_total after + the replay equals the number of distinct committed lines.""" + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d1-no-double-count" + graph = _FlakyGraph() + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + qm.commit = _flaky( # type: ignore[method-assign] + qm.commit, TimeoutError("SMB operation timed out") + ) + + written_before = reg.pipeline_counters()["written_total"] + + with patch( + "context_intelligence_server.registry.process_event", + side_effect=_accumulate, + ): + await qm.append(sid, _line("g1", "/ws", {"session_id": sid})) + await qm.append(sid, _line("g2", "/ws", {"session_id": sid})) + + task = _start_supervised(reg, worker) + await _await_death(task) + + worker2 = _make_worker(sid, graph) + reg._register_for_test(worker2) + reg.start_drain(worker2) + await _drain_until_idle(reg, qm, worker2, sid) + await _cancel_and_await(worker2.task) + + assert reg.pipeline_counters()["written_total"] == written_before + 2, ( + "each of the 2 distinct lines must be counted exactly once" + ) + + +class TestCloseTaskReferenced: + async def test_close_task_is_referenced_until_it_completes(self) -> None: + """registry._close_tasks holds the fire-and-forget close task while + it is pending, and is empty once it finishes -- without this, + asyncio's weak reference could let it be garbage-collected mid-close.""" + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d1-close-task-referenced" + graph = _FlakyGraph() + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + close_started = asyncio.Event() + close_release = asyncio.Event() + + async def _slow_close() -> None: + close_started.set() + await close_release.wait() + graph.closed = True + + graph.close = _slow_close # type: ignore[method-assign] + + qm.read_batch = _flaky( # type: ignore[method-assign] + qm.read_batch, OSError(errno.EIO, "boom") + ) + + with patch( + "context_intelligence_server.registry.process_event", + side_effect=_accumulate, + ): + await qm.append(sid, _line("e1", "/ws", {"session_id": sid})) + task = _start_supervised(reg, worker) + with contextlib.suppress(BaseException): + await task + + await asyncio.wait_for(close_started.wait(), timeout=5.0) + assert len(reg._close_tasks) == 1, ( + "the close task must be referenced while pending" + ) + + close_release.set() + for _ in range(200): + await asyncio.sleep(0.005) + if not reg._close_tasks: + break + + assert reg._close_tasks == set() + assert graph.closed is True + + +class TestTerminalBatchFlushExhaustion: + async def test_terminal_batch_flush_exhaustion_still_finalizes(self) -> None: + """A batch that exhausts the retry budget AND contains session:end + must still finalize the session -- a CompletedSession is recorded and + delete_drained runs -- instead of committing past session:end with no + finalization, which would leak the fully-drained log forever + (recover()'s strict `<` excludes it once the offset reaches EOF). + + The 2-record batch (tool:pre, session:end) dispatches into one + buffer of size 2, which the default ``_FlakyGraph.fail_when`` + rejects -- forcing every batch-level attempt to fail until the + retry budget is spent and ``_handle_exhausted_batch`` isolates it + line by line. Isolation flushes ``tool:pre`` alone (buffer size 1, + succeeds) and, on reaching ``session:end``, must leave it + uncommitted rather than dispatching/committing past it -- + ``_finalize_session``'s own ``_drain_to_eof`` re-dispatch (the same + re-dispatch the NORMAL non-exhausted terminal path already relies + on) then completes it. + """ + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d1-terminal-exhaustion-finalizes" + graph = _FlakyGraph() # default fail_when: len(buf) > 1 + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + with patch( + "context_intelligence_server.registry.process_event", + side_effect=_accumulate, + ): + await qm.append(sid, _line("tool:pre", "/ws", {"session_id": sid})) + await qm.append(sid, _line("session:end", "/ws", {"session_id": sid})) + task = _start_supervised(reg, worker) + await asyncio.wait_for(task, timeout=5.0) + + assert len(reg.completed_sessions()) == 1 + assert graph.flushed == {"tool:pre", "session:end"} + assert not qm._log_path(sid).exists(), "delete_drained must have run" + assert not qm._offset_path(sid).exists() + recoverable = await qm.recover() + assert sid not in recoverable + dead = await qm.read_dead_letters(sid) + assert dead == [], ( + "session:end must never be dead-lettered by isolation -- it is " + "left uncommitted for _finalize_session to re-dispatch" + ) + + async def test_exhausted_batch_without_terminal_isolates_without_finalizing( + self, + ) -> None: + """Unchanged behavior: a poison batch that exhausts the retry + budget but contains NO session:end record must dead-letter every + record and NOT finalize the session -- guards against the fix + over-triggering finalization for a batch that never reaches a + terminal record.""" + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d1-exhausted-no-terminal" + graph = _FlakyGraph(fail_when=lambda buf: True) # flush ALWAYS fails + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + with patch( + "context_intelligence_server.registry.process_event", + side_effect=_accumulate, + ): + await qm.append(sid, _line("tool:pre", "/ws", {"session_id": sid})) + await qm.append(sid, _line("tool:post", "/ws", {"session_id": sid})) + task = _start_supervised(reg, worker) + await _drain_until_idle(reg, qm, worker, sid) + assert not task.done(), "no terminal record: the drainer stays alive" + await _cancel_and_await(task) + + dead = await qm.read_dead_letters(sid) + assert len(dead) == 2 + assert reg.completed_sessions() == [] + + async def test_poison_line_before_terminal_is_dead_lettered_and_session_finalizes( + self, + ) -> None: + """An unparseable (poison) line sitting before session:end in the + same exhausted batch is still dead-lettered by isolation, and the + session still finalizes once the terminal record is reached.""" + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d1-poison-before-terminal-finalizes" + graph = _FlakyGraph() # default fail_when: len(buf) > 1 + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + with patch( + "context_intelligence_server.registry.process_event", + side_effect=_accumulate, + ): + await qm.append(sid, b"{ this is not valid json") + await qm.append(sid, _line("session:end", "/ws", {"session_id": sid})) + task = _start_supervised(reg, worker) + await asyncio.wait_for(task, timeout=5.0) + + assert len(reg.completed_sessions()) == 1 + dead = await qm.read_dead_letters(sid) + assert len(dead) == 1 + assert not qm._log_path(sid).exists(), "delete_drained must have run" diff --git a/tests/test_durable_append_framing.py b/tests/test_durable_append_framing.py new file mode 100644 index 00000000..18040c2a --- /dev/null +++ b/tests/test_durable_append_framing.py @@ -0,0 +1,946 @@ +"""Durable append-log framing: QueueManager serializes every write to a +session's files through a per-key ``_KeyGuard``, so concurrent or +split (SMB-style short) writes can never merge or tear a record. +``_SplitWriteOS`` models a non-atomic filesystem by splitting each +``os.write`` into multiple short writes against the real fd. +""" + +from __future__ import annotations + +import asyncio +import json +import logging +import os +import threading +import time +from pathlib import Path +from typing import Any + +import pytest + +from context_intelligence_server import queue_manager as qm_module +from context_intelligence_server.queue_manager import QueueManager, _KeyGuard +from context_intelligence_server.registry import SessionRegistry + +pytestmark = pytest.mark.integration + +SESSION = "71afde0c-f061-4f4c-b124-9323f9d5b110" +WORKSPACE = "-Users-samule-repo-team-pulse-structure" + +# Bounded wait for cross-thread handshakes (never a bare sleep); exists so a +# broken handshake fails loud instead of hanging. +_HANDSHAKE_TIMEOUT_S = 10.0 + +# Models one SMB write op; a small value keeps tests fast while still +# forcing multiple storage ops per logical write. +_SMB_OP_BYTES = 64 * 1024 + + +def _event_bytes(event: str, *, filler: int = 0) -> bytes: + """One event line, JSON-encoded; `filler` pads the payload to make a large record.""" + obj: dict[str, Any] = { + "event": event, + "workspace": WORKSPACE, + "data": {"session_id": SESSION, "payload": "x" * filler}, + "created_by": "samueljklee", + } + return json.dumps(obj, separators=(",", ":")).encode("utf-8") + + +def _parses(raw: bytes) -> bool: + """True iff the real drain-side parser accepts this line.""" + try: + SessionRegistry._parse_line(raw) + except Exception: # noqa: BLE001 - mirrors the real parser's own broad catch + return False + return True + + +async def _poll_until( + pred: Any, timeout: float = _HANDSHAKE_TIMEOUT_S, interval: float = 0.005 +) -> None: + """Poll a predicate until true, or fail loud on timeout.""" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if pred(): + return + await asyncio.sleep(interval) + raise AssertionError(f"condition not met within {timeout}s") + + +class _SplitWriteOS: + """Splits each os.write call into multiple short writes against the + real fd, modelling a non-atomic filesystem (e.g. SMB). ``on_chunk`` + fires after each chunk lands, for deterministic test handshakes. + """ + + def __init__( + self, real_write: Any, *, chunk: int = _SMB_OP_BYTES, on_chunk: Any = None + ) -> None: + self._real_write = real_write + self._chunk = chunk + self._on_chunk = on_chunk + self._counts: dict[int, int] = {} + self._lock = threading.Lock() + + def __call__(self, fd: int, data: Any) -> int: + buf = bytes(data) + to_write = buf[: self._chunk] if len(buf) > self._chunk else buf + n = self._real_write(fd, to_write) + with self._lock: + idx = self._counts.get(fd, 0) + self._counts[fd] = idx + 1 + if self._on_chunk is not None: + self._on_chunk(fd, idx, to_write[:n], len(buf)) + return n + + +# --------------------------------------------------------------------------- +# CONTROL -- baseline, real filesystem, no shim +# --------------------------------------------------------------------------- + + +async def test_control_local_o_append_is_atomic(tmp_path: Path) -> None: + """BASELINE: on a local filesystem O_APPEND happens to be atomic, so + every line parses even without the guard doing any work. Passing alone + is not proof the guard works -- see test_smb_split_write_no_longer_merges_records. + """ + qm = QueueManager(tmp_path) + + records: list[bytes] = [] + for i in range(12): + records.append(_event_bytes(f"llm:request:{i}", filler=300 * 1024)) + records.append(_event_bytes(f"tool:call:{i}a")) + records.append(_event_bytes(f"tool:call:{i}b")) + + await asyncio.gather(*(qm.append(SESSION, r) for r in records)) + + batch = await qm.read_batch(SESSION, max_items=1000) + + assert len(batch.lines) == len(records), ( + f"expected {len(records)} complete lines, got {len(batch.lines)}" + ) + bad = [i for i, ln in enumerate(batch.lines) if not _parses(ln)] + assert not bad, f"lines failed _parse_line on a LOCAL filesystem: {bad}" + assert sorted(batch.lines) == sorted(records) + + +# --------------------------------------------------------------------------- +# CONTROL: the shim tears without the gate -- isolates the guard from luck +# --------------------------------------------------------------------------- + + +def _raw_write_no_lock(path: Path, line: bytes) -> None: + """Unguarded append: open O_APPEND, write, close -- no serialization.""" + flags = os.O_WRONLY | os.O_CREAT | os.O_APPEND | getattr(os, "O_BINARY", 0) + fd = os.open(path, flags, 0o644) + try: + view = memoryview(line) + written = 0 + while written < len(view): + n = os.write(fd, view[written:]) + if n == 0: + raise OSError("os.write returned 0; refusing to spin") + written += n + finally: + os.close(fd) + + +def test_smb_shim_tears_WITHOUT_the_gate_control( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """CONTROL: without the guard's serialization, two threads writing to + the same file via the split-write shim interleave and produce a line + that fails to parse -- proving the shim actually models non-atomic + append (and that the guard, not luck, is what fixes it next door). + """ + real_write = os.write + large_first_op = threading.Event() + small_landed = threading.Event() + + def _on_chunk(fd: int, idx: int, chunk: bytes, total_len: int) -> None: + if idx == 0 and total_len > _SMB_OP_BYTES: + large_first_op.set() + assert small_landed.wait(_HANDSHAKE_TIMEOUT_S), ( + "handshake failed: the small write never landed" + ) + + monkeypatch.setattr( + qm_module.os, "write", _SplitWriteOS(real_write, on_chunk=_on_chunk) + ) + + log_path = tmp_path / "control-no-gate-key.log" + record_large = _event_bytes("llm:request", filler=300 * 1024) + record_small = _event_bytes("llm:stream_block_start") + line_large = record_large if record_large.endswith(b"\n") else record_large + b"\n" + line_small = record_small if record_small.endswith(b"\n") else record_small + b"\n" + assert len(line_large) > _SMB_OP_BYTES + + def _write_large() -> None: + _raw_write_no_lock(log_path, line_large) + + def _write_small() -> None: + assert large_first_op.wait(_HANDSHAKE_TIMEOUT_S), ( + "the large writer's first sub-op never landed" + ) + _raw_write_no_lock(log_path, line_small) + small_landed.set() + + t_large = threading.Thread(target=_write_large) + t_small = threading.Thread(target=_write_small) + t_large.start() + t_small.start() + t_large.join(_HANDSHAKE_TIMEOUT_S) + t_small.join(_HANDSHAKE_TIMEOUT_S) + assert not t_large.is_alive() and not t_small.is_alive(), ( + "a writer thread failed to finish -- the handshake deadlocked" + ) + + raw_log = log_path.read_bytes() + assert raw_log != line_large + line_small, ( + "control is VACUOUS: bytes landed byte-exact/sequential -- the " + "forced handshake did not actually interleave the two writes" + ) + + physical_lines = raw_log.split(b"\n") + if physical_lines and physical_lines[-1] == b"": + physical_lines = physical_lines[:-1] + bad = [ln for ln in physical_lines if not _parses(ln)] + assert bad, ( + "expected at least one merged/torn physical line to FAIL " + "_parse_line when nothing serializes the two writers -- if this " + "assertion fails, the shim is not modelling non-atomic append, " + "which would make T2's green PASS vacuous" + ) + + +# --------------------------------------------------------------------------- +# the inverted repro: concurrent write under contention +# --------------------------------------------------------------------------- + + +async def test_smb_split_write_no_longer_merges_records( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Framing holds under genuine contention: parks a large write mid-record + via the shim while a second appender for the same key waits on the + guard; both lines land whole, in order, and parse cleanly. + """ + qm = QueueManager(tmp_path) + real_write = os.write + + large_first_op = threading.Event() + small_landed = threading.Event() + + def _on_chunk(fd: int, idx: int, chunk: bytes, total_len: int) -> None: + if idx == 0 and total_len > _SMB_OP_BYTES: + large_first_op.set() + assert small_landed.wait(_HANDSHAKE_TIMEOUT_S), ( + "handshake failed: the small append never landed" + ) + + monkeypatch.setattr( + qm_module.os, "write", _SplitWriteOS(real_write, on_chunk=_on_chunk) + ) + + large = _event_bytes("llm:request", filler=300 * 1024) + small = _event_bytes("llm:stream_block_start") + assert len(large) > _SMB_OP_BYTES + + async def _append_large() -> None: + await qm.append(SESSION, large) + + async def _append_small() -> None: + await asyncio.to_thread( + lambda: large_first_op.wait(_HANDSHAKE_TIMEOUT_S) or None + ) + guard = qm._guards[SESSION] + assert guard.admission.locked() + task = asyncio.ensure_future(qm.append(SESSION, small)) + await _poll_until(lambda: guard.waiters == 2) + small_landed.set() + await task + + await asyncio.gather(_append_large(), _append_small()) + + raw_log = (tmp_path / f"{SESSION}.log").read_bytes() + assert raw_log == large + b"\n" + small + b"\n", ( + "expected byte-exact [large]\\n[small]\\n -- no merged/torn line" + ) + + batch = await qm.read_batch(SESSION, max_items=100) + assert len(batch.lines) == 2, f"expected 2 lines, got {len(batch.lines)}" + assert _parses(batch.lines[0]) and _parses(batch.lines[1]) + + +# --------------------------------------------------------------------------- +# concurrent appends over both key shapes, all parse +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("key", [SESSION, "_nosession_-workspace-abc"]) +async def test_concurrent_appends_under_smb_shim_all_parse( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, key: str +) -> None: + qm = QueueManager(tmp_path) + monkeypatch.setattr(qm_module.os, "write", _SplitWriteOS(os.write)) + + records: list[bytes] = [] + for i in range(12): + records.append(_event_bytes(f"llm:request:{i}", filler=300 * 1024)) + for i in range(24): + records.append(_event_bytes(f"tool:call:{i}")) + + # precondition: no raw newline anywhere except the trailing terminator + for r in records: + stored = r if r.endswith(b"\n") else r + b"\n" + assert b"\n" not in stored[:-1], f"P1 violated by record: {r[:80]!r}" + + await asyncio.gather(*(qm.append(key, r) for r in records)) + + batch = await qm.read_batch(key, max_items=1000) + assert len(batch.lines) == len(records) + bad = [i for i, ln in enumerate(batch.lines) if not _parses(ln)] + assert not bad, f"lines failed _parse_line: {bad}" + assert sorted(batch.lines) == sorted(records) + + +# --------------------------------------------------------------------------- +# cancellation cannot reintroduce the tear +# --------------------------------------------------------------------------- + + +async def test_cancel_mid_write_never_releases_the_file_lock( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + qm = QueueManager(tmp_path) + key = SESSION + + entered_write = threading.Event() + proceed = threading.Event() + ops: list[tuple[str, int]] = [] + ops_lock = threading.Lock() + + record_a = _event_bytes("llm:request:A", filler=300 * 1024) + record_b = _event_bytes("tool:call:B") + # exact on-disk bytes (append() adds the trailing newline), used below to + # identify which record a given short chunk belongs to + line_a = record_a if record_a.endswith(b"\n") else record_a + b"\n" + line_b = record_b if record_b.endswith(b"\n") else record_b + b"\n" + real_write = os.write + + def _on_chunk(fd: int, idx: int, chunk: bytes, total_len: int) -> None: + label = ( + "A" + if chunk and line_a.startswith(chunk) + else ("B" if chunk and line_b.startswith(chunk) else "?") + ) + with ops_lock: + ops.append((label, len(chunk))) + if idx == 0 and total_len > _SMB_OP_BYTES: + entered_write.set() + assert proceed.wait(_HANDSHAKE_TIMEOUT_S), "proceed handshake timed out" + + monkeypatch.setattr( + qm_module.os, "write", _SplitWriteOS(real_write, on_chunk=_on_chunk) + ) + + task_a = asyncio.ensure_future(qm.append(key, record_a)) + await asyncio.to_thread(lambda: entered_write.wait(_HANDSHAKE_TIMEOUT_S) or None) + + guard = qm._guards[key] + assert guard.file_lock.locked(), ( + "file_lock must be held while A's write is mid-flight" + ) + + task_a.cancel() + + # dispatch B now -- it must not reach os.write while A's thread holds file_lock + task_b = asyncio.ensure_future(qm.append(key, record_b)) + await asyncio.sleep(0.05) + assert guard.file_lock.locked(), ( + "file_lock unexpectedly released before A's write finished" + ) + assert all(label != "B" for label, _ in ops), ( + "B's bytes landed before A released file_lock" + ) + + proceed.set() # let A's write finish + + a_landed_before_cancel_observed = False + with pytest.raises(asyncio.CancelledError): + try: + await task_a + except asyncio.CancelledError: + # A's bytes must already be on disk at the moment cancellation is observed here + current = (tmp_path / f"{key}.log").read_bytes() + a_landed_before_cancel_observed = current.startswith(record_a + b"\n") + raise + + assert a_landed_before_cancel_observed, ( + "CancelledError observed before A's bytes landed" + ) + + await task_b + + # (c) both records whole, in order, and _parse_line-clean. + raw_log = (tmp_path / f"{key}.log").read_bytes() + assert raw_log == record_a + b"\n" + record_b + b"\n" + batch = await qm.read_batch(key, max_items=10) + assert len(batch.lines) == 2 + assert _parses(batch.lines[0]) and _parses(batch.lines[1]) + + # No interleaving: every A-chunk fully precedes every B-chunk. + labels = [label for label, _ in ops] + last_a = max(i for i, lbl in enumerate(labels) if lbl == "A") + first_b = min(i for i, lbl in enumerate(labels) if lbl == "B") + assert last_a < first_b, f"ops interleaved: {labels}" + + +# --------------------------------------------------------------------------- +# distinct keys append concurrently -- per-key parallelism preserved +# --------------------------------------------------------------------------- + + +async def test_distinct_keys_append_concurrently( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + qm = QueueManager(tmp_path) + key_a = "key-a" + key_b = "key-b" + + entered = threading.Event() + proceed = threading.Event() + + def _on_chunk(fd: int, idx: int, chunk: bytes, total_len: int) -> None: + if idx == 0 and total_len > _SMB_OP_BYTES: + entered.set() + assert proceed.wait(_HANDSHAKE_TIMEOUT_S) + + monkeypatch.setattr( + qm_module.os, "write", _SplitWriteOS(os.write, on_chunk=_on_chunk) + ) + + large_a = _event_bytes("large-a", filler=300 * 1024) + small_b = _event_bytes("small-b") + + task_a = asyncio.ensure_future(qm.append(key_a, large_a)) + await asyncio.to_thread(lambda: entered.wait(_HANDSHAKE_TIMEOUT_S) or None) + + # key B is a DIFFERENT key -- must complete while A is still parked. + await qm.append(key_b, small_b) + batch_b = await qm.read_batch(key_b, max_items=10) + assert batch_b.lines == [small_b] + + proceed.set() + await task_a + batch_a = await qm.read_batch(key_a, max_items=10) + assert batch_a.lines == [large_a] + + +# --------------------------------------------------------------------------- +# partial write failure discards the record; failure is loud +# --------------------------------------------------------------------------- + + +async def test_partial_write_failure_discards_the_record( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + qm = QueueManager(tmp_path) + key = "fail-key" + real_write = os.write + calls = {"n": 0} + + def _flaky_write(fd: int, data: Any) -> int: + calls["n"] += 1 + if calls["n"] == 1: + return real_write(fd, bytes(data)[:8]) + if calls["n"] == 2: + raise OSError("simulated append failure") + return real_write(fd, bytes(data)) # rollback's newline write succeeds + + monkeypatch.setattr(qm_module.os, "write", _flaky_write) + + record = _event_bytes("will-fail") + with pytest.raises(OSError): + await qm.append(key, record) + + # Fragment is newline-terminated, never truncated -- the queue never + # removes bytes it already wrote. + assert qm._log_path(key).read_bytes() == record[:8] + b"\n" + + monkeypatch.setattr(qm_module.os, "write", real_write) + good = _event_bytes("will-succeed") + await qm.append(key, good) + batch = await qm.read_batch(key, max_items=10) + assert batch.lines == [record[:8], good] + assert not _parses(batch.lines[0]), "malformed fragment must not parse" + assert _parses(batch.lines[1]) + + +async def test_partial_write_failure_logs_when_newline_terminate_fails( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + qm = QueueManager(tmp_path) + key = "fail-key-2" + real_write = os.write + calls = {"n": 0} + + def _flaky_write(fd: int, data: Any) -> int: + calls["n"] += 1 + if calls["n"] == 1: + return real_write(fd, bytes(data)[:8]) + raise OSError("simulated write failure") + + monkeypatch.setattr(qm_module.os, "write", _flaky_write) + + record = _event_bytes("will-fail-hard") + with ( + caplog.at_level( + logging.ERROR, logger="context_intelligence_server.queue_manager" + ), + pytest.raises(OSError), + ): + await qm.append(key, record) + + errors = [r.message for r in caplog.records if r.levelno >= logging.ERROR] + assert any("append_partial_terminate_failed" in m for m in errors), ( + "newline-terminate failure must be logged at ERROR" + ) + + # Torn tail left untouched -- queue bytes are never removed. Readers skip an + # unterminated trailing fragment; the next append merges it into a single + # poison line that the drainer dead-letters. + assert qm._log_path(key).read_bytes() == record[:8] + + +# _discard_partial never truncates: a peer writer's committed line and this +# writer's own prior records survive a rollback. + + +async def test_discard_partial_never_destroys_a_peer_process_committed_line( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Two independent processes append to the same file (writer-lease is + detect-only -- see _write_record's own docstring, no cross-process lock). + While this writer's partial write is rolled back, a peer process + completes and closes its own fully-formed, already-acknowledged record. + _discard_partial must never remove those bytes. + """ + qm = QueueManager(tmp_path) + key = "race-key" + path = qm._log_path(key) + path.write_bytes(b'{"payload":"PRIOR-COMMITTED"}\n') + + peer_can_go = threading.Event() + peer_done = threading.Event() + real_write = os.write + + def _flaky_write(fd: int, data: Any) -> int: + buf = bytes(data) + real_write(fd, buf[: len(buf) // 2]) + peer_can_go.set() + assert peer_done.wait(_HANDSHAKE_TIMEOUT_S) + raise OSError("simulated mid-record failure") + + def _peer_append() -> None: + assert peer_can_go.wait(_HANDSHAKE_TIMEOUT_S) + flags = os.O_WRONLY | os.O_CREAT | os.O_APPEND + fd = os.open(path, flags, 0o644) + try: + real_write(fd, b'{"payload":"PEER-COMMITTED"}\n') + finally: + os.close(fd) + peer_done.set() + + peer = threading.Thread(target=_peer_append) + peer.start() + + monkeypatch.setattr(qm_module.os, "write", _flaky_write) + record = _event_bytes("mine", filler=5000) + with pytest.raises(OSError): + await qm.append(key, record) + monkeypatch.setattr(qm_module.os, "write", real_write) + + peer.join(_HANDSHAKE_TIMEOUT_S) + assert not peer.is_alive() + + final = path.read_bytes() + assert b"PEER-COMMITTED" in final, ( + "a peer's already-acknowledged line must never be destroyed" + ) + assert b"PRIOR-COMMITTED" in final, ( + "pre-existing committed data must never be destroyed" + ) + + +async def test_discard_partial_single_writer_preserves_prior_records( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Single-writer case: a partial write rolled back leaves prior COMPLETE + records intact and the fragment newline-terminated (not merged into the + next record); a subsequent drain dead-letters the fragment rather than + crashing. + """ + qm = QueueManager(tmp_path) + key = "single-writer-key" + prior = _event_bytes("prior-committed") + await qm.append(key, prior) + + real_write = os.write + calls = {"n": 0} + + def _flaky_write(fd: int, data: Any) -> int: + calls["n"] += 1 + if calls["n"] == 1: + return real_write(fd, bytes(data)[:8]) + if calls["n"] == 2: + raise OSError("simulated append failure") + return real_write(fd, bytes(data)) # rollback's newline write succeeds + + monkeypatch.setattr(qm_module.os, "write", _flaky_write) + record = _event_bytes("torn-fragment") + with pytest.raises(OSError): + await qm.append(key, record) + monkeypatch.setattr(qm_module.os, "write", real_write) + + good = _event_bytes("after-recovery") + await qm.append(key, good) + + batch = await qm.read_batch(key, max_items=10) + assert batch.lines == [prior, record[:8], good], ( + "no committed record lost; fragment isolated on its own line" + ) + assert _parses(batch.lines[0]) + assert not _parses(batch.lines[1]), ( + "malformed fragment is dead-lettered, not crashed on" + ) + assert _parses(batch.lines[2]) + + +# --------------------------------------------------------------------------- +# delete_drained cannot race append; retains on uncommitted bytes +# --------------------------------------------------------------------------- + + +async def test_delete_drained_cannot_race_an_in_flight_append( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + qm = QueueManager(tmp_path) + key = "race-delete-key" + + entered = threading.Event() + proceed = threading.Event() + + def _on_chunk(fd: int, idx: int, chunk: bytes, total_len: int) -> None: + if idx == 0 and total_len > _SMB_OP_BYTES: + entered.set() + assert proceed.wait(_HANDSHAKE_TIMEOUT_S) + + monkeypatch.setattr( + qm_module.os, "write", _SplitWriteOS(os.write, on_chunk=_on_chunk) + ) + + record = _event_bytes("in-flight", filler=300 * 1024) + task_append = asyncio.ensure_future(qm.append(key, record)) + await asyncio.to_thread(lambda: entered.wait(_HANDSHAKE_TIMEOUT_S) or None) + + guard = qm._guards[key] + assert guard.admission.locked() + + delete_task = asyncio.ensure_future(qm.delete_drained(key)) + await _poll_until(lambda: guard.waiters == 2) # append(1) + delete parked(1) + + proceed.set() + await task_append + ok = await delete_task + + # nothing committed yet, so delete_drained must retain the record fully + # -- admission serialization blocks delete until append releases file_lock + assert ok is False + assert qm._log_path(key).exists() + batch = await qm.read_batch(key, max_items=10) + assert batch.lines == [record] + assert _parses(batch.lines[0]) + + +async def test_delete_drained_retains_a_log_with_uncommitted_bytes( + tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: + qm = QueueManager(tmp_path) + key = "uncommitted-key" + first = _event_bytes("first") + 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.append(key, second) # uncommitted tail + + with caplog.at_level( + logging.WARNING, logger="context_intelligence_server.queue_manager" + ): + ok = await qm.delete_drained(key) + + assert ok is False + assert qm._log_path(key).exists() + assert qm._offset_path(key).exists() + assert any("delete_drained_retained" in r.message for r in caplog.records) + + recoverable = await qm.recover() + assert key in recoverable + + +# --------------------------------------------------------------------------- +# guard map is bounded and ABA-proof +# --------------------------------------------------------------------------- + + +async def test_guard_map_is_released_on_delete_drained_and_identity_checked( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + qm = QueueManager(tmp_path) + keys = [f"session-{i}" for i in range(5)] + 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) + assert set(qm._guards.keys()) == set(keys) + + for k in keys: + assert await qm.delete_drained(k) is True + assert qm._guards == {} + + # ABA probe: swap the guard map entry for a foreign object while + # delete_drained is mid-flight; its identity check must refuse to remove it + 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) + g_orig = qm._guards[key] + + real_stat = Path.stat + entered = threading.Event() + proceed = threading.Event() + + def _paused_stat(self: Path, *a: Any, **kw: Any) -> Any: + if self == qm._log_path(key): + entered.set() + assert proceed.wait(_HANDSHAKE_TIMEOUT_S) + return real_stat(self, *a, **kw) + + monkeypatch.setattr(Path, "stat", _paused_stat) + + task = asyncio.ensure_future(qm.delete_drained(key)) + await asyncio.to_thread(lambda: entered.wait(_HANDSHAKE_TIMEOUT_S) or None) + + foreign = _KeyGuard(asyncio.Lock(), threading.Lock()) + qm._guards[key] = foreign # simulate a swap while delete_drained is mid-flight + + proceed.set() + ok = await task + + assert ok is True # the unlink itself still completed against g_orig + assert qm._guards.get(key) is foreign, ( + "foreign entry must survive the identity check" + ) + assert g_orig is not foreign + + +# --------------------------------------------------------------------------- +# dead-letter path is covered and un-crashable +# --------------------------------------------------------------------------- + + +async def test_dead_letter_record_is_framed_under_smb_shim( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + qm = QueueManager(tmp_path) + key = "dl-key" + monkeypatch.setattr(qm_module.os, "write", _SplitWriteOS(os.write)) + large_raw = _event_bytes("bad:record", filler=300 * 1024) + await qm.dead_letter(key, large_raw, "simulated parse error") + records = await qm.read_dead_letters(key) + assert len(records) == 1 + assert records[0]["error"] == "simulated parse error" + assert records[0]["payload"] == large_raw.decode("utf-8") + + +async def test_dead_letter_parsing_survives_a_malformed_line(tmp_path: Path) -> None: + qm = QueueManager(tmp_path) + key = "dl-malformed-key" + dead_path = qm._dead_path(key) + good = json.dumps({"ts": 1.0, "error": "e", "payload": "ok"}) + dead_path.write_text(good + "\n" + "{not json" + "\n", encoding="utf-8") + + records = await qm.read_dead_letters(key) # must not raise + assert len(records) == 1 + assert records[0]["payload"] == "ok" + + payload_set = qm._dead_payload_set(key) # must not raise + assert payload_set == {b"ok"} + + +# --------------------------------------------------------------------------- +# guard survives a delete racing a parked appender +# --------------------------------------------------------------------------- + + +async def test_guard_survives_a_delete_that_races_a_parked_appender( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """delete_drained racing a parked appender: the guard map entry must + survive while appender A still holds a reference, so a subsequent + appender B is served by the SAME guard rather than a disjoint one. + """ + qm = QueueManager(tmp_path) + key = "race-key" + + 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 + guard = qm._guards[key] + + real_stat = Path.stat + entered_delete = threading.Event() + proceed_delete = threading.Event() + + def _paused_stat(self: Path, *a: Any, **kw: Any) -> Any: + if self == qm._log_path(key): + entered_delete.set() + assert proceed_delete.wait(_HANDSHAKE_TIMEOUT_S), ( + "delete handshake timed out" + ) + return real_stat(self, *a, **kw) + + monkeypatch.setattr(Path, "stat", _paused_stat) + + delete_task = asyncio.ensure_future(qm.delete_drained(key)) + await asyncio.to_thread(lambda: entered_delete.wait(_HANDSHAKE_TIMEOUT_S) or None) + assert qm._guards.get(key) is guard # not yet removed -- delete hasn't returned + + append_a = _event_bytes("A") + task_a = asyncio.ensure_future(qm.append(key, append_a)) + await _poll_until(lambda: guard.waiters == 2) # delete(1, itself) + A parked(1) + assert guard.admission.locked() + + monkeypatch.undo() # restore Path.stat before it resumes for real + proceed_delete.set() + + ok = await delete_task + assert ok is True # size == committed -> genuinely drained -> unlink succeeds + await task_a + + append_b = _event_bytes("B") + await qm.append(key, append_b) + + assert qm._guards.get(key) is guard, ( + "B must be served by the SAME _KeyGuard object A used -- if the " + "guard was discarded while A held a reference (v2's identity-only " + "removal condition), B gets a fresh guard with a disjoint " + "file_lock, which reproduces the same torn/merged-line append corruption" + ) + + raw = qm._log_path(key).read_bytes() + assert raw == append_a + b"\n" + append_b + b"\n" + batch = await qm.read_batch(key, max_items=10) + assert len(batch.lines) == 2 + assert _parses(batch.lines[0]) and _parses(batch.lines[1]) + + +# --------------------------------------------------------------------------- +# captured production artifacts +# --------------------------------------------------------------------------- + + +def _seed_dir() -> Path | None: + """Locate captured dead-letter seed files (kept outside this repo); + walk upward so the lookup survives any checkout depth. + """ + for parent in Path(__file__).resolve().parents: + candidate = parent / "docs" / "04-deadletter-artifacts" / "seeds" + if candidate.is_dir(): + return candidate + return None + + +_SEEDS = _seed_dir() +_requires_seeds = pytest.mark.skipif( + _SEEDS is None, + reason="captured dead-letter seeds not present", +) + + +@_requires_seeds +async def test_read_batch_over_a_pre_existing_merged_middle_line( + tmp_path: Path, +) -> None: + """A pre-existing merged middle line still fails _parse_line -- consuming + it (dead-letter, commit past it) is the drainer's job, not this fix's. + """ + qm = QueueManager(tmp_path) + key = "merged-middle-key" + assert _SEEDS is not None + merged = (_SEEDS / "seed_corrupt_merged_line_1.0MiB.raw").read_bytes() + good_before = _event_bytes("before") + good_after = _event_bytes("after") + log_path = qm._log_path(key) + log_path.write_bytes(good_before + b"\n" + merged + b"\n" + good_after + b"\n") + + batch = await qm.read_batch(key, max_items=10) + assert len(batch.lines) == 3 + assert batch.lines[0] == good_before + assert batch.lines[1] == merged + assert batch.lines[2] == good_after + assert not _parses(merged), ( + "the framing fix does not claim to fix a pre-existing merged middle line -- the drainer consumes it" + ) + assert _parses(good_before) and _parses(good_after) + + +@_requires_seeds +def test_captured_corrupt_seed_is_rejected_by_the_real_parser() -> None: + """A real captured corrupt production line still fails _parse_line, at + the same offset and with the same error text as the original dead-letter. + """ + assert _SEEDS is not None + raw = (_SEEDS / "seed_corrupt_merged_line_1.0MiB.raw").read_bytes() + + # one physical line containing two merged records, no separator between them + assert raw.count(b"\n") == 0, "seed is a single physical line by construction" + starts = [i for i in range(len(raw)) if raw.startswith(b'{"event":', i)] + assert len(starts) == 2, f"expected two merged records, found {len(starts)}" + assert starts[0] == 0 + boundary = starts[1] + assert raw[boundary - 1 : boundary] != b"\n", ( + "record B is preceded by a newline -- that would be normal framing, not a tear" + ) + + with pytest.raises(json.JSONDecodeError) as exc: + SessionRegistry._parse_line(raw) + + # parse fails near the merge boundary; here A tore inside a quoted string + # so B's leading bytes are swallowed, breaking a couple bytes later + assert boundary <= exc.value.pos <= boundary + 8, ( + f"parse failed at {exc.value.pos}, merge boundary is {boundary}" + ) + assert exc.value.pos == 1050641, "matches the recorded dead-letter error offset" + assert "Expecting ',' delimiter" in str(exc.value), ( + "matches the recorded dead-letter error text" + ) + + +@_requires_seeds +def test_captured_valid_large_event_parses_cleanly() -> None: + """A real large event of similar size parses fine -- isolating the + defect to framing, not payload size. + """ + assert _SEEDS is not None + raw = (_SEEDS / "seed_valid_large_event_1.04MiB.json").read_bytes() + assert len(raw) > 1024 * 1024 + assert raw.count(b"\n") == 0 + + event, workspace, data = SessionRegistry._parse_line(raw) + assert event + assert workspace + assert isinstance(data, dict) diff --git a/tests/test_finalize_delete_ordering.py b/tests/test_finalize_delete_ordering.py new file mode 100644 index 00000000..3815a3ad --- /dev/null +++ b/tests/test_finalize_delete_ordering.py @@ -0,0 +1,515 @@ +"""`_finalize_session` delete-ordering race: a late append landing in the +finalize window must be drained then deleted, or, if every attempt sees a +late append, the bounded retry gives up and the log is retained. + +Uses a deterministic window-injection technique: wrap `qm.delete_drained` +with a spy that appends a late line before delegating to the real method. +No real Neo4j is used anywhere in this file. +""" + +from __future__ import annotations + +import asyncio +import json +import logging +from collections.abc import Awaitable, Callable +from typing import Any +from unittest.mock import patch + +import pytest +from context_intelligence_server.queue_manager import QueueManager +from context_intelligence_server.registry import SessionRegistry, SessionWorker +from context_intelligence_server.services import HookStateService + +pytestmark = pytest.mark.integration + + +# --------------------------------------------------------------------------- +# Wire format + fakes (mirrors tests/test_drain_supervision.py's style) +# --------------------------------------------------------------------------- + + +def _line(event: str, workspace: str, data: dict) -> bytes: + """Encode an appended event line exactly as POST /events stores it.""" + return json.dumps({"event": event, "workspace": workspace, "data": data}).encode( + "utf-8" + ) + + +class _AccumGraph: + """Accumulating-buffer graph fake. Writes accumulate in ``buffer`` until + ``flush()`` moves them into ``flushed`` (a SET, so a replayed event never + shows up twice). ``fail_on_call``, if given, makes the Nth non-empty + ``flush()`` call raise (1-based).""" + + def __init__(self, *, fail_on_call: int | None = None) -> None: + self.workspace = "/ws" + self.created_by: str | None = None + self.buffer: set[str] = set() + self.flushed: set[str] = set() + self.discards = 0 + self.closed = False + self._fail_on_call = fail_on_call + self._calls = 0 + + async def flush(self) -> None: + if not self.buffer: + return # empty-buffer early return (GraphStore Protocol guarantee #5) + self._calls += 1 + if self._fail_on_call is not None and self._calls == self._fail_on_call: + raise RuntimeError(f"simulated flush failure on call {self._calls}") + self.flushed |= self.buffer + self.buffer.clear() + + def discard_buffer(self) -> None: + self.buffer.clear() + self.discards += 1 + + async def close(self) -> None: + self.closed = True + + +async def _accumulate( + worker: SessionWorker, event: str, data: object, handlers: object +) -> None: + """Stand-in for ``process_event``: buffers the event name on the fake graph.""" + worker.services.graph.buffer.add(event) + + +def _make_worker(sid: str, graph: Any, workspace: str = "/ws") -> SessionWorker: + worker = SessionWorker( + session_id=sid, + workspace=workspace, + services=HookStateService(workspace=workspace), + ) + worker.services.graph = graph # type: ignore[assignment] + return worker + + +def _delete_drained_injector( + qm: QueueManager, + late_lines: list[bytes], + inject_on: set[int], +) -> tuple[Callable[[str], Awaitable[bool]], dict[str, int]]: + """Wrap ``qm.delete_drained`` so, on the given 1-based call numbers, it + appends the next late line before delegating to the real method. Returns + ``(wrapper, calls)`` where ``calls["count"]`` tracks invocations.""" + original = qm.delete_drained + calls = {"count": 0} + injected = {"count": 0} + + async def _wrapper(session_id: str) -> bool: + calls["count"] += 1 + attempt = calls["count"] + if attempt in inject_on and injected["count"] < len(late_lines): + await qm.append(session_id, late_lines[injected["count"]]) + injected["count"] += 1 + return await original(session_id) + + return _wrapper, calls + + +def _start_supervised( + reg: SessionRegistry, worker: SessionWorker, *, flush_timeout: float = 10.0 +) -> asyncio.Task: + """Mirror production ``start_drain`` (registry.py) exactly: create the + task, attach the done-callback, bind ``worker.task``.""" + import functools + + task = asyncio.create_task( + reg.drain_worker(worker, flush_timeout=flush_timeout), + name=f"drain-{worker.session_id}", + ) + task.add_done_callback(functools.partial(reg._on_drain_done, worker)) + worker.task = task + return task + + +# --------------------------------------------------------------------------- +# A late append is drained, then deleted +# --------------------------------------------------------------------------- + + +async def test_late_append_in_finalize_window_is_drained_then_deleted( + caplog: pytest.LogCaptureFixture, +) -> None: + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d5-t1-late-append" + graph = _AccumGraph() + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + late_line = _line("late:event", "/ws", {"session_id": sid}) + wrapper, calls = _delete_drained_injector(qm, [late_line], inject_on={1}) + qm.delete_drained = wrapper # type: ignore[method-assign] + + with ( + patch( + "context_intelligence_server.registry.process_event", + side_effect=_accumulate, + ), + caplog.at_level(logging.WARNING, logger="context_intelligence_server"), + ): + await qm.append(sid, _line("tool:pre", "/ws", {"session_id": sid})) + await qm.append(sid, _line("session:end", "/ws", {"session_id": sid})) + await reg._finalize_session(worker, handlers=object()) + + assert calls["count"] == 2, ( + "attempt 1 must retain (late append lands inside the window); " + "attempt 2 must succeed after the re-drain persists it" + ) + assert "late:event" in graph.flushed, ( + "the late event must be dispatched BEFORE the log is deleted" + ) + assert not qm._log_path(sid).exists() + assert not qm._offset_path(sid).exists() + assert any( + r.levelno == logging.WARNING + and "finalize_delete_retained" in r.getMessage() + and getattr(r, "session_id", None) == sid + for r in caplog.records + ), "the retained-on-attempt-1 WARNING must be logged exactly once" + + +# --------------------------------------------------------------------------- +# A late append on every attempt exhausts the retry +# --------------------------------------------------------------------------- + + +async def test_late_append_on_every_attempt_retains_and_is_recoverable( + caplog: pytest.LogCaptureFixture, +) -> None: + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d5-t2-give-up" + graph = _AccumGraph() + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + late_lines = [ + _line(f"late:event:{i}", "/ws", {"session_id": sid}) for i in range(3) + ] + wrapper, calls = _delete_drained_injector(qm, late_lines, inject_on={1, 2, 3}) + qm.delete_drained = wrapper # type: ignore[method-assign] + + with ( + patch( + "context_intelligence_server.registry.process_event", + side_effect=_accumulate, + ), + caplog.at_level(logging.WARNING, logger="context_intelligence_server"), + ): + await qm.append(sid, _line("session:end", "/ws", {"session_id": sid})) + await reg._finalize_session(worker, handlers=object()) + + assert calls["count"] == 3, ( + "must give up after exactly _FINALIZE_DELETE_ATTEMPTS(=3) delete calls" + ) + assert qm._log_path(sid).exists(), "log must be RETAINED on give-up (never lost)" + assert any( + r.levelno == logging.ERROR + and "finalize_delete_gave_up" in r.getMessage() + and getattr(r, "session_id", None) == sid + for r in caplog.records + ), "the give-up must be logged loudly at ERROR" + + recoverable = await qm.recover() + assert sid in recoverable, ( + "a retained log with a complete uncommitted line must be recover()-reportable" + ) + assert sid not in reg.active_sessions(), ( + "give-up still deregisters + closes (unchanged teardown path, spec 3.5)" + ) + assert worker.store_closed is True + + +# --------------------------------------------------------------------------- +# No double-delete + delete/close/deregister ordering preserved on the clean path +# --------------------------------------------------------------------------- + + +async def test_no_double_delete_and_call_b_ordering_preserved() -> None: + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d5-t4-call-b-ordering" + graph = _AccumGraph() + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + sequence: list[str] = [] + original_delete = qm.delete_drained + original_close = graph.close + original_deregister = reg._deregister + + async def _spy_delete(session_id: str) -> bool: + sequence.append("delete_drained") + return await original_delete(session_id) + + async def _spy_close() -> None: + sequence.append("graph.close") + await original_close() + + def _spy_deregister(session_id: str) -> None: + sequence.append("_deregister") + original_deregister(session_id) + + qm.delete_drained = _spy_delete # type: ignore[method-assign] + graph.close = _spy_close # type: ignore[method-assign] + reg._deregister = _spy_deregister # type: ignore[method-assign] + + with patch( + "context_intelligence_server.registry.process_event", + side_effect=_accumulate, + ): + await qm.append(sid, _line("session:end", "/ws", {"session_id": sid})) + await reg._finalize_session(worker, handlers=object()) + + assert sequence.count("delete_drained") == 1, ( + "exactly one delete_drained call on the clean path -- never a double-delete" + ) + assert sequence == ["delete_drained", "graph.close", "_deregister"], ( + "ordering: delete -> close -> deregister, deregister LAST" + ) + + +# --------------------------------------------------------------------------- +# Regression: the common no-late-append finalize still deletes +# --------------------------------------------------------------------------- + + +async def test_clean_finalize_still_deletes_and_tears_down( + caplog: pytest.LogCaptureFixture, +) -> None: + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d5-t6-clean-finalize" + graph = _AccumGraph() + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + delete_calls = {"count": 0} + original_delete = qm.delete_drained + + async def _counting_delete(session_id: str) -> bool: + delete_calls["count"] += 1 + return await original_delete(session_id) + + qm.delete_drained = _counting_delete # type: ignore[method-assign] + + with ( + patch( + "context_intelligence_server.registry.process_event", + side_effect=_accumulate, + ), + caplog.at_level(logging.INFO, logger="context_intelligence_server"), + ): + await qm.append(sid, _line("tool:pre", "/ws", {"session_id": sid})) + await qm.append(sid, _line("session:end", "/ws", {"session_id": sid})) + await reg._finalize_session(worker, handlers=object()) + + assert not qm._log_path(sid).exists() + assert not qm._offset_path(sid).exists() + assert delete_calls["count"] == 1 + assert len(reg.completed_sessions()) == 1 + assert any( + r.levelno == logging.INFO and "session_finalized" in r.getMessage() + for r in caplog.records + ) + assert sid not in reg.active_sessions() + assert worker.store_closed is True + + +# --------------------------------------------------------------------------- +# A first-pass tail flush failure returns before CompletedSession is recorded +# --------------------------------------------------------------------------- + + +async def test_first_pass_tail_flush_failure_returns_before_completed_session( + caplog: pytest.LogCaptureFixture, +) -> None: + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d5-t7-tail-flush-failure" + graph = _AccumGraph(fail_on_call=1) + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + with ( + patch( + "context_intelligence_server.registry.process_event", + side_effect=_accumulate, + ), + caplog.at_level(logging.ERROR, logger="context_intelligence_server"), + ): + await qm.append(sid, _line("tool:pre", "/ws", {"session_id": sid})) + await qm.append(sid, _line("session:end", "/ws", {"session_id": sid})) + await reg._finalize_session(worker, handlers=object()) + + assert any( + r.levelno == logging.ERROR and "finalize_tail_flush_failed" in r.getMessage() + for r in caplog.records + ) + assert len(reg.completed_sessions()) == 0, ( + "CompletedSession must NOT be recorded when the FIRST pass's tail flush fails" + ) + assert sid in reg.active_sessions(), ( + "worker must remain registered so a respawn retries" + ) + assert worker.store_closed is False + assert qm._log_path(sid).exists(), "the tail must remain uncommitted on disk" + + +# --------------------------------------------------------------------------- +# The retry loop terminates in at most _FINALIZE_DELETE_ATTEMPTS DELETE +# attempts, regardless of a continuously-appending client +# --------------------------------------------------------------------------- + + +@pytest.mark.timeout(30) +async def test_retry_loop_terminates_under_continuous_append() -> None: + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d5-t8-bounded-termination" + graph = _AccumGraph() + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + # More late lines than attempts are possible, to prove the bound is real + # even if the injector *could* keep going. + late_lines = [ + _line(f"late:event:{i}", "/ws", {"session_id": sid}) for i in range(10) + ] + wrapper, calls = _delete_drained_injector(qm, late_lines, inject_on={1, 2, 3, 4, 5}) + qm.delete_drained = wrapper # type: ignore[method-assign] + + with patch( + "context_intelligence_server.registry.process_event", + side_effect=_accumulate, + ): + await qm.append(sid, _line("session:end", "/ws", {"session_id": sid})) + await reg._finalize_session(worker, handlers=object()) + + assert calls["count"] == 3, ( + "delete attempts must be bounded to _FINALIZE_DELETE_ATTEMPTS(=3) " + "regardless of how many late lines a continuous appender could supply" + ) + + +# --------------------------------------------------------------------------- +# Permanent retention: the retry's own re-drain can itself suffer a tail +# flush failure after CompletedSession was already recorded, returning early +# and never reaching _safe_close/_deregister. orphaned_sessions() is the +# honest signal. +# --------------------------------------------------------------------------- + + +async def test_permanent_retention_when_retrys_own_redrain_flush_fails( + caplog: pytest.LogCaptureFixture, +) -> None: + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d5-t9-permanent-retention" + # 1st flush (tool:pre + session:end, the initial _drain_to_eof) succeeds; + # 2nd flush (the retry's re-drain of the late event) fails. + graph = _AccumGraph(fail_on_call=2) + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + late_line = _line("late:event", "/ws", {"session_id": sid}) + wrapper, calls = _delete_drained_injector(qm, [late_line], inject_on={1}) + qm.delete_drained = wrapper # type: ignore[method-assign] + + with ( + patch( + "context_intelligence_server.registry.process_event", + side_effect=_accumulate, + ), + caplog.at_level(logging.ERROR, logger="context_intelligence_server"), + ): + await qm.append(sid, _line("tool:pre", "/ws", {"session_id": sid})) + await qm.append(sid, _line("session:end", "/ws", {"session_id": sid})) + # A real, completed Task is required for orphaned_sessions() (it + # checks worker.task is not None and worker.task.done()). + task = asyncio.create_task(reg._finalize_session(worker, handlers=object())) + worker.task = task + await task + + assert calls["count"] == 1, ( + "delete_drained called exactly once; the re-drain's OWN flush " + "failure returns early before a second delete attempt" + ) + assert len(reg.completed_sessions()) == 1, ( + "CompletedSession was already recorded BEFORE the retry loop began" + ) + assert sid in reg.active_sessions(), ( + "the early return never reaches _deregister -- permanently registered" + ) + assert worker.store_closed is False, ( + "the early return never reaches _safe_close either" + ) + assert any( + r.levelno == logging.ERROR and "finalize_tail_flush_failed" in r.getMessage() + for r in caplog.records + ) + orphans = reg.orphaned_sessions() + assert any(w.session_id == sid for w in orphans), ( + "orphaned_sessions() is the honest signal for this permanent-" + "retention residual -- registered, task done, never re-entered" + ) + assert qm._log_path(sid).exists(), ( + "the late event's log is RETAINED -- never lost, never re-attempted" + ) + + +# --------------------------------------------------------------------------- +# A fresh drainer over the same on-disk retained log dispatches the late +# event and drains fully. +# --------------------------------------------------------------------------- + + +async def test_retained_log_is_picked_up_by_a_fresh_drainer() -> None: + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d5-t3-retained-log-pickup" + graph = _AccumGraph() + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + late_lines = [ + _line(f"late:event:{i}", "/ws", {"session_id": sid}) for i in range(3) + ] + original_delete = qm.delete_drained + wrapper, calls = _delete_drained_injector(qm, late_lines, inject_on={1, 2, 3}) + qm.delete_drained = wrapper # type: ignore[method-assign] + + with patch( + "context_intelligence_server.registry.process_event", + side_effect=_accumulate, + ): + await qm.append(sid, _line("session:end", "/ws", {"session_id": sid})) + await reg._finalize_session(worker, handlers=object()) + + # Give-up end-state: log retained with late:event:2 undrained. + assert calls["count"] == 3 + assert qm._log_path(sid).exists() + assert "late:event:2" not in graph.flushed + + # Restore the REAL delete_drained for the fresh drainer -- the + # injector's job (landing a late append inside the *original* + # finalize's window) is done; a second drainer must not re-inject. + qm.delete_drained = original_delete # type: ignore[method-assign] + + worker2 = _make_worker(sid, graph) + reg._register_for_test(worker2) + reg.start_drain(worker2) + assert worker2.task is not None + await asyncio.wait_for(worker2.task, timeout=5.0) + + assert "late:event:2" in graph.flushed, ( + "the fresh drainer must dispatch the previously-retained late event" + ) + assert qm._read_committed_offset(sid) == qm._complete_data_end(sid), ( + "log ends fully drained -- committed advances to complete_data_end, " + "proving outcome (B) is real pick-up, not merely asserted" + ) diff --git a/tests/test_graph_store.py b/tests/test_graph_store.py index 1669fe0b..0b088125 100644 --- a/tests/test_graph_store.py +++ b/tests/test_graph_store.py @@ -17,12 +17,30 @@ class MinimalGraphStore: - """Conforming implementation of GraphStore with all required members.""" + """Conforming implementation of GraphStore with all required members. + + The Protocol declares a settable ``workspace`` and a ``created_by`` + getter/setter -- both required for a real isinstance() conformance + check to pass, so this fixture must carry them too (a + runtime_checkable Protocol only checks attribute PRESENCE). + """ @property def workspace(self) -> str: return "test-workspace" + @workspace.setter + def workspace(self, value: str) -> None: + pass + + @property + def created_by(self) -> str | None: + return None + + @created_by.setter + def created_by(self, value: str | None) -> None: + pass + async def upsert_node(self, node_id: str, data: dict[str, Any]) -> None: pass @@ -76,12 +94,27 @@ async def close(self) -> None: class MinimalQueryableStore: - """Conforming implementation of QueryableStore with all required members.""" + """Conforming implementation of QueryableStore with all required members. + + See MinimalGraphStore's docstring. + """ @property def workspace(self) -> str: return "test-workspace" + @workspace.setter + def workspace(self, value: str) -> None: + pass + + @property + def created_by(self) -> str | None: + return None + + @created_by.setter + def created_by(self, value: str | None) -> None: + pass + @property def supported_dialects(self) -> frozenset[str]: return frozenset({"cypher", "sparql"}) diff --git a/tests/test_large_event_tail_drop.py b/tests/test_large_event_tail_drop.py new file mode 100644 index 00000000..3c89e2c1 --- /dev/null +++ b/tests/test_large_event_tail_drop.py @@ -0,0 +1,169 @@ +"""Tail-drop under load: POST /events durably appends and returns before any +Neo4j write; a per-session `drain_worker` task later flushes. An unguarded +`dead_letter` call inside `_handle_exhausted_batch`'s except-clause can +raise past `drain_worker`'s only guard (`asyncio.CancelledError` only), +killing the task and stranding the uncommitted tail on disk. `start_drain`'s +done-callback makes that death loud (`drain_worker_died`) and self-healing: +it deregisters the worker so the next event or a boot recovery respawns a +fresh drainer that resumes at the stranded tail. Large (>1MB) events reach +this more often since their own slow solo transaction is more likely to +exhaust the retry budget. No real Neo4j is required. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import functools +import json +from unittest.mock import patch + +import pytest + +from context_intelligence_server.queue_manager import QueueManager +from context_intelligence_server.registry import SessionRegistry, SessionWorker +from context_intelligence_server.services import HookStateService + +pytestmark = pytest.mark.integration + + +def _line(event: str, workspace: str, data: dict) -> bytes: + """Encode an appended event line exactly as POST /events stores it.""" + return json.dumps({"event": event, "workspace": workspace, "data": data}).encode( + "utf-8" + ) + + +class _FaultInjectableGraph: + """Models a real store's accumulating write buffer: writes accumulate; + `flush()` fails while the designated poison event is resident (modeling + a Neo4j write rejection on an oversized event's own solo transaction); a + successful flush clears the buffer; `discard_buffer()` clears it without + flushing, so poison residue can't contaminate the next line.""" + + def __init__(self, poison_event: str) -> None: + self.workspace = "/ws" + self.poison_event = poison_event + self.buffer: set[str] = set() + self.flushed: list[str] = [] + self.discards = 0 + self.closed = False + + async def flush(self) -> None: + if not self.buffer: + return # empty-buffer early return, mirroring the real store + if self.poison_event in self.buffer: + raise RuntimeError( + f"neo4j write rejected for {self.poison_event!r} " + "(oversized solo transaction exhausted retries)" + ) + self.flushed.extend(sorted(self.buffer)) + self.buffer.clear() # success clears + + def discard_buffer(self) -> None: + self.buffer.clear() + self.discards += 1 + + async def close(self) -> None: + self.closed = True + + +async def _drive_drain_to_quiescence( + reg: SessionRegistry, + qm: QueueManager, + worker: SessionWorker, + sid: str, + *, + flush_timeout: float = 10.0, + max_polls: int = 400, + poll_sleep: float = 0.01, +) -> asyncio.Task: + """Start the real drain_worker as a background task and poll (never a + bare sleep) until either it finishes on its own (defect: an unguarded + exception kills it) or the queue drains (control: idle-polls forever, + must be cancelled). Attaches the same done-callback `start_drain` does, + so Case B can observe the real loud-death + self-heal contract.""" + task = asyncio.create_task( + reg.drain_worker(worker, flush_timeout=flush_timeout), name=f"drain-{sid}" + ) + # mirror start_drain's own bindings so a later supervision check is meaningful + task.add_done_callback(functools.partial(reg._on_drain_done, worker)) + worker.task = task + for _ in range(max_polls): + await asyncio.sleep(poll_sleep) + if task.done(): + break + if (await qm.read_batch(sid, 10)).lines == []: + break + return task + + +async def _cancel_and_await(task: asyncio.Task) -> None: + if not task.done(): + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + + +# --------------------------------------------------------------------------- +# Case A -- CONTROL: the designed resilience path. Proves isolation + +# dead-letter isolates an oversized line while the prefix and tail persist. +# --------------------------------------------------------------------------- + + +class TestOversizedEventControlPathIsolatesAndContinues: + async def test_prefix_and_tail_persist_oversized_dead_lettered(self) -> None: + """Batch [small-1, small-2, OVERSIZED, tail-1, tail-2] against a store + whose flush() only rejects the oversized event: retries exhaust, + driving isolation with dead_letter and commit both healthy. All four + small events persist, the oversized event is dead-lettered (not + dropped), the offset advances past all 5 lines, and the drain task + stays healthy -- contrasts directly with Case B below. + """ + reg = SessionRegistry() + qm = reg.queue_manager + sid = "large-event-control" + + fake = _FaultInjectableGraph(poison_event="oversized") + worker = SessionWorker( + session_id=sid, + workspace="/ws", + services=HookStateService(workspace="/ws"), + ) + worker.services.graph = fake # type: ignore[assignment] + reg._register_for_test(worker) + + async def _process(w: object, event: str, data: object, h: object) -> None: + fake.buffer.add(event) + + with patch( + "context_intelligence_server.registry.process_event", side_effect=_process + ): + await qm.append(sid, _line("small-1", "/ws", {"session_id": sid})) + await qm.append(sid, _line("small-2", "/ws", {"session_id": sid})) + await qm.append(sid, _line("oversized", "/ws", {"session_id": sid})) + await qm.append(sid, _line("tail-1", "/ws", {"session_id": sid})) + await qm.append(sid, _line("tail-2", "/ws", {"session_id": sid})) + + task = await _drive_drain_to_quiescence(reg, qm, worker, sid) + await _cancel_and_await(task) + + # --- pin: full prefix AND tail persisted (no drop, no truncation) --- + assert fake.flushed == ["small-1", "small-2", "tail-1", "tail-2"] + + # --- pin: oversized event isolated + dead-lettered, not lost --- + dead = await qm.read_dead_letters(sid) + assert len(dead) == 1 + assert json.loads(dead[0]["payload"])["event"] == "oversized" + + # --- pin: offset advanced past the whole batch --- + assert (await qm.read_batch(sid, 10)).lines == [] + + # --- pin: the drain task is healthy -- no unhandled exception --- + + +# The failure-mode counterpart (dead_letter raising mid-drain -> loud task +# death, respawn drains the stranded tail, oversized event dead-lettered) is +# proven by test_drain_supervision.py:: +# test_dead_letter_failure_no_longer_kills_the_drainer over this same event +# shape; not duplicated here. diff --git a/tests/test_main.py b/tests/test_main.py index 302fa1bf..f8d71322 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -797,9 +797,11 @@ async def test_lifespan_recovers_and_respawns_drainers( async def test_lifespan_skips_recovery_for_empty_workspace( monkeypatch: pytest.MonkeyPatch, ) -> None: - """A session whose first line has an empty workspace is NOT respawned - (spawning a workspace='' worker would violate the non-empty-workspace - invariant).""" + """A session whose only line has an empty workspace is SKIPPED -- never + dispatched under workspace='' and never under a substitute. workspace is + the graph partition key, so a guessed value would write the session into a + partition it does not belong to; the log is left durable on disk and a + later boot reports the session again.""" sid = "sess-empty-ws" qm = registry.queue_manager body = json.dumps( diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 32ef699a..dfe3bffb 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -27,7 +27,6 @@ import pytest - # --------------------------------------------------------------------------- # NOTE: ToolCallHandler stub injection is performed in conftest.py so it # fires before any test module loads, regardless of pytest collection order. @@ -153,8 +152,8 @@ def test_setup_handlers_returns_pipeline_handlers() -> None: def test_setup_handlers_has_default_handler_with_services() -> None: - from context_intelligence_server.pipeline import setup_handlers from context_intelligence_server.handlers.data_layer_1.default import DefaultHandler + from context_intelligence_server.pipeline import setup_handlers from context_intelligence_server.services import HookStateService services = HookStateService(workspace="test") @@ -185,17 +184,17 @@ def test_setup_handlers_enricher_count() -> None: def test_setup_handlers_enricher_order() -> None: """Enrichers must be [SessionHandler, OrchestratorRunHandler, IterationHandler, ContentBlockHandler, ToolCallHandler] in that dispatch order.""" - from context_intelligence_server.pipeline import setup_handlers - from context_intelligence_server.handlers.data_layer_2.session import SessionHandler - from context_intelligence_server.handlers.data_layer_2.orchestrator_run import ( - OrchestratorRunHandler, + from context_intelligence_server.handlers.data_layer_2.content_block import ( + ContentBlockHandler, ) from context_intelligence_server.handlers.data_layer_2.iteration import ( IterationHandler, ) - from context_intelligence_server.handlers.data_layer_2.content_block import ( - ContentBlockHandler, + from context_intelligence_server.handlers.data_layer_2.orchestrator_run import ( + OrchestratorRunHandler, ) + from context_intelligence_server.handlers.data_layer_2.session import SessionHandler + from context_intelligence_server.pipeline import setup_handlers from context_intelligence_server.services import HookStateService services = HookStateService(workspace="test") @@ -210,19 +209,19 @@ def test_setup_handlers_enricher_order() -> None: def test_setup_handlers_l3_enricher_order() -> None: """Layer 3 enrichers must be appended after all Layer 2 enrichers in correct order: [DelegationHandler, SkillLoadHandler, RecipeRunHandler, RecipeStepHandler].""" - from context_intelligence_server.pipeline import setup_handlers from context_intelligence_server.handlers.data_layer_3.delegation import ( DelegationHandler, ) - from context_intelligence_server.handlers.data_layer_3.skill_load import ( - SkillLoadHandler, - ) from context_intelligence_server.handlers.data_layer_3.recipe_run import ( RecipeRunHandler, ) from context_intelligence_server.handlers.data_layer_3.recipe_step import ( RecipeStepHandler, ) + from context_intelligence_server.handlers.data_layer_3.skill_load import ( + SkillLoadHandler, + ) + from context_intelligence_server.pipeline import setup_handlers from context_intelligence_server.services import HookStateService services = HookStateService(workspace="test") @@ -388,6 +387,46 @@ async def test_process_event_terminal_does_not_self_flush( mock_worker.services.graph.flush.assert_not_called() +async def test_process_event_terminal_does_not_self_flush_real_handlers() -> None: + """Wires the real ``setup_handlers(services)`` enrichers (not + ``_StubEnricher``) to verify ``SessionHandler`` never self-flushes; + also asserts the session node was actually written.""" + from context_intelligence_server.pipeline import process_event, setup_handlers + from context_intelligence_server.registry import SessionWorker + from context_intelligence_server.services import HookStateService + + services = HookStateService(workspace="test-real-handlers") + handlers = setup_handlers(services) + worker = SessionWorker( + session_id="sess-real-1", workspace="test-real-handlers", services=services + ) + + real_flush = services.graph.flush + flush_calls: list[None] = [] + + async def _counting_flush() -> None: + flush_calls.append(None) + await real_flush() + + services.graph.flush = _counting_flush # type: ignore[method-assign] + + data = {"session_id": "sess-real-1", "timestamp": "2026-01-01T00:00:00Z"} + await process_event(worker, "session:end", data, handlers) + + assert flush_calls == [], ( + f"SessionHandler._handle_end (via the REAL setup_handlers enrichers) " + f"called graph.flush directly {len(flush_calls)} time(s) -- " + f"process_event must not self-flush; the drainer's gated " + f"_flush_barrier is the sole trigger" + ) + + # Confirm the real SessionHandler actually ran and wrote the node. + node = await services.graph.get_node("sess-real-1") + assert node is not None and node.get("status") == "completed", ( + f"real SessionHandler did not run (mis-wired fixture?): {node}" + ) + + async def test_process_event_non_terminal_does_not_self_flush( mock_worker: MagicMock, pipeline_handlers: Any, @@ -412,7 +451,7 @@ async def test_process_event_default_handler_exception_propagates( mock_worker: MagicMock, default_handler: _StubDefaultHandler, ) -> None: - """Phase B2: a default-handler (step 4) error must PROPAGATE so the drainer + """A default-handler (step 4) error must PROPAGATE so the drainer can dead-letter the line instead of committing the offset past a never-persisted event (no silent loss).""" from context_intelligence_server.pipeline import PipelineHandlers, process_event @@ -426,7 +465,6 @@ async def test_process_event_default_handler_exception_propagates( ) -# NOTE (Task 6): test_process_event_flush_exception_propagates was removed. # process_event no longer flushes at all — the drainer's gated _flush_barrier is # the sole write trigger, so flush-failure-propagation is now a drainer contract # covered by tests/test_registry.py::TestDurableDrainLoop @@ -438,7 +476,7 @@ async def test_process_event_propagates_handler_error( mock_worker: MagicMock, pipeline_handlers: Any, ) -> None: - """Phase B2 (USER DECISION option a): a handler error in steps 2-6 must + """A handler error in steps 2-6 must PROPAGATE, not be swallowed — here ensure_session_node (step 2) raises and process_event must re-raise so the drainer routes the line to dead-letter rather than committing the offset past a never-persisted event.""" diff --git a/tests/test_queue_manager.py b/tests/test_queue_manager.py index bd270601..e84c1646 100644 --- a/tests/test_queue_manager.py +++ b/tests/test_queue_manager.py @@ -1,12 +1,15 @@ -"""Tests for the on-disk durable queue manager (Phase B1).""" +"""Tests for the on-disk durable queue manager.""" from __future__ import annotations import time import pytest - -from context_intelligence_server.queue_manager import Batch, QueueManager +from context_intelligence_server.queue_manager import ( + Batch, + QueueManager, + Record, +) @pytest.fixture @@ -22,13 +25,81 @@ def test_constructor_creates_queues_dir(tmp_path): def test_batch_holds_its_fields(): - batch = Batch(session_id="s1", lines=[b"a", b"b"], start_offset=0, end_offset=4) + """``batch.lines`` is derived from ``batch.records``.""" + batch = Batch( + session_id="s1", + records=[Record(b"a", 0, 2), Record(b"b", 2, 4)], + start_offset=0, + end_offset=4, + ) assert batch.session_id == "s1" assert batch.lines == [b"a", b"b"] assert batch.start_offset == 0 assert batch.end_offset == 4 +# --------------------------------------------------------------------------- +# Record / Batch.records: offsets are queue-produced and read-only for callers. +# --------------------------------------------------------------------------- + + +async def test_read_batch_records_carry_queue_produced_offsets(qm, tmp_path): + """Each record's start equals the previous record's end, the first/last + records bound the batch's start/end_offset, and no record's raw payload + has a trailing newline.""" + await qm.append("s1", b"one") + await qm.append("s1", b"two") + await qm.append("s1", b"three") + + batch = await qm.read_batch("s1", max_items=10) + + assert len(batch.records) == 3 + assert batch.records[0].start == batch.start_offset + assert batch.records[-1].end == batch.end_offset + for i in range(1, len(batch.records)): + assert batch.records[i].start == batch.records[i - 1].end + for rec in batch.records: + assert not rec.raw.endswith(b"\n") + assert [r.raw for r in batch.records] == [b"one", b"two", b"three"] + + +async def test_batch_lines_is_derived_from_records(qm, tmp_path): + """``batch.lines`` always matches ``[r.raw for r in batch.records]``.""" + await qm.append("s1", b"alpha") + await qm.append("s1", b"beta") + + batch = await qm.read_batch("s1", max_items=10) + + assert batch.lines == [r.raw for r in batch.records] + + +async def test_read_batch_records_survive_a_torn_trailing_line(qm, tmp_path): + """A log ending in a partial (torn) line yields records only for the + complete lines that precede it; end_offset stops on the line boundary.""" + log = tmp_path / "queues" / "s1.log" + log.parent.mkdir(parents=True, exist_ok=True) + log.write_bytes(b"complete-one\ncomplete-two\ntorn-no-newline-yet") + + batch = await qm.read_batch("s1", max_items=10) + + assert [r.raw for r in batch.records] == [b"complete-one", b"complete-two"] + assert batch.end_offset == len(b"complete-one\ncomplete-two\n") + + +async def test_committing_rec_end_advances_exactly_one_record(qm, tmp_path): + """``commit(sid, records[0].end)`` then a fresh ``read_batch`` returns + records ``[1:]``.""" + await qm.append("s1", b"first") + await qm.append("s1", b"second") + await qm.append("s1", b"third") + + batch = await qm.read_batch("s1", max_items=10) + await qm.commit("s1", batch.records[0].end) + + remaining = await qm.read_batch("s1", max_items=10) + assert [r.raw for r in remaining.records] == [b"second", b"third"] + + async def test_append_writes_line_with_trailing_newline(qm, tmp_path): await qm.append("s1", b'{"e":1}') log = tmp_path / "queues" / "s1.log" @@ -163,6 +234,42 @@ async def test_commit_is_atomic_no_temp_leftover(qm, tmp_path): assert list(qdir.glob("*.tmp")) == [] +# _read_committed_offset parses the bare-int form written by commit(). A +# present-but-unusable offset must raise, never silently return 0 (0 would +# force a full re-drain). + + +async def test_read_committed_offset_accepts_bare_int_unchanged(qm): + """Bare-int offsets (the current write format) still parse exactly.""" + qm._offset_path("s1").write_text("980582046", encoding="utf-8") + assert qm._read_committed_offset("s1") == 980582046 + + +async def test_read_committed_offset_missing_file_is_zero(qm): + assert qm._read_committed_offset("never-written") == 0 + + +async def test_read_committed_offset_empty_file_is_zero(qm): + qm._offset_path("s1").write_text("", encoding="utf-8") + assert qm._read_committed_offset("s1") == 0 + + +async def test_read_committed_offset_legacy_json_without_usable_offset_raises(qm): + """A JSON object present but with no usable integer "offset" must raise + ValueError -- the same as any other unparseable offset -- rather than + silently returning 0 (which would trigger a full re-drain).""" + qm._offset_path("s1").write_text('{"v":1,"cursor":{}}', encoding="utf-8") + with pytest.raises(ValueError): + qm._read_committed_offset("s1") + + +async def test_read_committed_offset_garbage_still_raises(qm): + """Genuinely unparseable text (not JSON, not an int) still raises.""" + qm._offset_path("s1").write_text("not-a-number", encoding="utf-8") + with pytest.raises(ValueError): + qm._read_committed_offset("s1") + + 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") @@ -317,7 +424,7 @@ def counting(): assert calls["n"] == 2 -# --- recovery_seed_counts (D2): residual-0-by-construction crash-recovery seed --- +# --- recovery_seed_counts: residual-0-by-construction crash-recovery seed --- async def test_recovery_seed_counts_pending_and_committed(qm): @@ -393,7 +500,7 @@ async def test_recovery_seed_counts_crash_window_residual_zero(qm): assert residual == 0 -# --- recovery_reconcile_dead (D2): close the dead_letter->commit crash window --- +# --- recovery_reconcile_dead: close the dead_letter->commit crash window --- async def test_recovery_reconcile_dead_advances_past_already_dead_pending(qm): diff --git a/tests/test_registry.py b/tests/test_registry.py index c31efeb4..901e00e7 100644 --- a/tests/test_registry.py +++ b/tests/test_registry.py @@ -17,7 +17,7 @@ import context_intelligence_server.registry as registry_module from context_intelligence_server.blob_store import AsyncDiskBlobStore from context_intelligence_server.config import get_settings -from context_intelligence_server.queue_manager import QueueManager +from context_intelligence_server.queue_manager import QueueManager, Record from context_intelligence_server.registry import ( CompletedSession, SessionRegistry, @@ -25,7 +25,6 @@ ) from context_intelligence_server.services import HookStateService - # --------------------------------------------------------------------------- # Factory helper # --------------------------------------------------------------------------- @@ -503,8 +502,7 @@ async def test_dead_letter_logs_warning_with_session_id( # A malformed line makes _parse_line raise inside _handle_exhausted_batch, # triggering the dead-letter path. poison = MagicMock() - poison.lines = [b"{ this is not valid json"] - poison.start_offset = 0 + 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()) @@ -961,7 +959,9 @@ async def test_session_end_finalizes_and_deregisters( cs = reg._completed[0] assert cs.session_id == sid assert cs.workspace == "/ws" - assert cs.events_processed == 2 + # session:end is dispatched twice: once live, once via the + # finalize re-read. + assert cs.events_processed == 3 assert cs.error_count == 0 assert cs.ended_at > 0.0 assert cs.duration_seconds >= 0.0 @@ -1268,7 +1268,7 @@ class _S: queues_path = str(tmp_path / "queues") neo4j_url = "bolt://unused:7687" neo4j_user = "neo4j" - neo4j_password = "unused" # noqa: S105 - test stub, not a real secret + neo4j_password = "unused" stale_session_timeout = 3600.0 write_concurrency = 2 max_delivery_attempts = 3 @@ -1403,7 +1403,7 @@ async def test_handler_error_is_dead_lettered_not_silently_committed( # --------------------------------------------------------------------------- -# Task 5 (D2): live conservation counters on SessionRegistry. These feed the +# Live conservation counters on SessionRegistry. These feed the # pipeline-conservation snapshot in /status so silently-dropped events become # observable (accepted vs written vs replayed, plus write retries). # --------------------------------------------------------------------------- @@ -1448,7 +1448,7 @@ def test_seed_counters_adds(self) -> None: # --------------------------------------------------------------------------- -# Task 6 (D2/D3): SessionRegistry.pipeline_metrics() assembles the live +# Task 6: SessionRegistry.pipeline_metrics() assembles the live # conservation counters with disk-derived queue/dead aggregates into a single # health block (residual + degraded). This is the /status aggregate that makes # silent loss observable. LIVE per-process measure (not an all-time audit): @@ -1620,7 +1620,7 @@ async def test_sustained_drop_reports_degraded( # --------------------------------------------------------------------------- -# Task 7 (D2): written/retry counter increments wired into the drainer at the +# Task 7: written/retry counter increments wired into the drainer at the # four real commit/retry sites: (1) normal-path commit, (2) retry on flush # failure, (3) per-line success during exhausted-batch isolation, and (4) the # finalize tail commit. These prove the live conservation counters actually @@ -1952,8 +1952,10 @@ async def test_session_finalized_logs_info_with_session_id( worker.services.graph.close = AsyncMock() # type: ignore[method-assign] reg._register_for_test(worker) # Isolate from the real queue: no tail to drain, no real disk I/O. + # (records=[] -- _finalize_session's tail loop now iterates Batch.records; + # lines is a derived property, so an empty records is an empty tail.) reg.queue_manager.read_batch = AsyncMock( # type: ignore[method-assign] - return_value=MagicMock(lines=[]) + return_value=MagicMock(records=[], lines=[]) ) reg.queue_manager.commit = AsyncMock() # type: ignore[method-assign] reg.queue_manager.delete_drained = AsyncMock() # type: ignore[method-assign] @@ -2063,31 +2065,29 @@ def test_get_or_create_reuse_keeps_bound_created_by( reg = SessionRegistry() - with caplog.at_level(logging.ERROR, logger="context_intelligence_server"): - with ( - patch( - "context_intelligence_server.registry.Neo4jGraphStore" - ) as MockStore, - patch( - "context_intelligence_server.registry.AsyncDiskBlobStore" - ) as MockBlob, - patch( - "context_intelligence_server.registry.HookStateService" - ) as MockService, - ): - MockStore.return_value = MagicMock() - MockBlob.return_value = MagicMock() - mock_svc = MagicMock() - MockService.return_value = mock_svc - reg.start_drain = MagicMock() - - # First call — creates the worker, binds "alice" - reg.get_or_create("sess-same", "/ws", created_by="alice") - # Simulate what the real HookStateService sets on graph_store - mock_svc.graph.created_by = "alice" - - # Second call — same contributor, must be silent - worker = reg.get_or_create("sess-same", "/ws", created_by="alice") + with ( + caplog.at_level(logging.ERROR, logger="context_intelligence_server"), + patch("context_intelligence_server.registry.Neo4jGraphStore") as MockStore, + patch( + "context_intelligence_server.registry.AsyncDiskBlobStore" + ) as MockBlob, + patch( + "context_intelligence_server.registry.HookStateService" + ) as MockService, + ): + MockStore.return_value = MagicMock() + MockBlob.return_value = MagicMock() + mock_svc = MagicMock() + MockService.return_value = mock_svc + reg.start_drain = MagicMock() + + # First call — creates the worker, binds "alice" + reg.get_or_create("sess-same", "/ws", created_by="alice") + # Simulate what the real HookStateService sets on graph_store + mock_svc.graph.created_by = "alice" + + # Second call — same contributor, must be silent + worker = reg.get_or_create("sess-same", "/ws", created_by="alice") error_records = [r for r in caplog.records if r.levelno == logging.ERROR] assert error_records == [], ( @@ -2105,30 +2105,28 @@ def test_get_or_create_reuse_ignores_new_created_by_no_error( reg = SessionRegistry() - with caplog.at_level(logging.ERROR, logger="context_intelligence_server"): - with ( - patch( - "context_intelligence_server.registry.Neo4jGraphStore" - ) as MockStore, - patch( - "context_intelligence_server.registry.AsyncDiskBlobStore" - ) as MockBlob, - patch( - "context_intelligence_server.registry.HookStateService" - ) as MockService, - ): - MockStore.return_value = MagicMock() - MockBlob.return_value = MagicMock() - mock_svc = MagicMock() - MockService.return_value = mock_svc - reg.start_drain = MagicMock() - - # First call — creates the worker bound to "alice" - reg.get_or_create("sess-none", "/ws", created_by="alice") - mock_svc.graph.created_by = "alice" - - # Second call — created_by=None must never trigger the guard - worker = reg.get_or_create("sess-none", "/ws", created_by=None) + with ( + caplog.at_level(logging.ERROR, logger="context_intelligence_server"), + patch("context_intelligence_server.registry.Neo4jGraphStore") as MockStore, + patch( + "context_intelligence_server.registry.AsyncDiskBlobStore" + ) as MockBlob, + patch( + "context_intelligence_server.registry.HookStateService" + ) as MockService, + ): + MockStore.return_value = MagicMock() + MockBlob.return_value = MagicMock() + mock_svc = MagicMock() + MockService.return_value = mock_svc + reg.start_drain = MagicMock() + + # First call — creates the worker bound to "alice" + reg.get_or_create("sess-none", "/ws", created_by="alice") + mock_svc.graph.created_by = "alice" + + # Second call — created_by=None must never trigger the guard + worker = reg.get_or_create("sess-none", "/ws", created_by=None) error_records = [r for r in caplog.records if r.levelno == logging.ERROR] assert error_records == [], ( @@ -2146,30 +2144,28 @@ def test_invariant_violation_is_observed_and_not_overwritten( reg = SessionRegistry() - with caplog.at_level(logging.ERROR, logger="context_intelligence_server"): - with ( - patch( - "context_intelligence_server.registry.Neo4jGraphStore" - ) as MockStore, - patch( - "context_intelligence_server.registry.AsyncDiskBlobStore" - ) as MockBlob, - patch( - "context_intelligence_server.registry.HookStateService" - ) as MockService, - ): - MockStore.return_value = MagicMock() - MockBlob.return_value = MagicMock() - mock_svc = MagicMock() - MockService.return_value = mock_svc - reg.start_drain = MagicMock() - - # First call — creates the worker bound to "alice" - reg.get_or_create("sess-conflict", "/ws", created_by="alice") - mock_svc.graph.created_by = "alice" - - # Second call — conflicting contributor "bob" arrives - worker = reg.get_or_create("sess-conflict", "/ws", created_by="bob") + with ( + caplog.at_level(logging.ERROR, logger="context_intelligence_server"), + patch("context_intelligence_server.registry.Neo4jGraphStore") as MockStore, + patch( + "context_intelligence_server.registry.AsyncDiskBlobStore" + ) as MockBlob, + patch( + "context_intelligence_server.registry.HookStateService" + ) as MockService, + ): + MockStore.return_value = MagicMock() + MockBlob.return_value = MagicMock() + mock_svc = MagicMock() + MockService.return_value = mock_svc + reg.start_drain = MagicMock() + + # First call — creates the worker bound to "alice" + reg.get_or_create("sess-conflict", "/ws", created_by="alice") + mock_svc.graph.created_by = "alice" + + # Second call — conflicting contributor "bob" arrives + worker = reg.get_or_create("sess-conflict", "/ws", created_by="bob") # 1. An ERROR record must have been emitted error_records = [r for r in caplog.records if r.levelno == logging.ERROR] diff --git a/tests/test_touch_session_no_root_contention.py b/tests/test_touch_session_no_root_contention.py index ac559153..5e5e99ae 100644 --- a/tests/test_touch_session_no_root_contention.py +++ b/tests/test_touch_session_no_root_contention.py @@ -14,12 +14,20 @@ class FakeGraph: - """Minimal async graph store that records which nodes get upserted.""" + """Minimal async graph store that records which nodes get upserted. + + Conforms fully to the ``GraphStore`` Protocol: the ``graph_store`` + constructor parameter is typed as ``GraphStore | None``, so a fake + passed to it must structurally satisfy the Protocol even though this + test only exercises get_node/upsert_node. The extra members are + no-ops -- this test's behavior is unchanged. + """ def __init__(self, nodes: dict[str, dict[str, Any]]) -> None: self.nodes = nodes self.touched: list[str] = [] self.workspace = "test" + self.created_by: str | None = None async def get_node(self, node_id: str) -> dict[str, Any] | None: return self.nodes.get(node_id) @@ -28,6 +36,26 @@ async def upsert_node(self, node_id: str, data: dict[str, Any]) -> None: self.touched.append(node_id) self.nodes.setdefault(node_id, {}).update(data) + async def upsert_edge(self, src_id: str, dst_id: str, data: dict[str, Any]) -> None: + pass + + async def get_edge(self, src_id: str, dst_id: str) -> dict[str, Any] | None: + return None + + async def find_delegation_by_sub_session( + self, sub_session_id: str, workspace: str + ) -> dict[str, Any] | None: + return None + + def discard_buffer(self) -> None: + pass + + async def flush(self) -> None: + pass + + async def close(self) -> None: + pass + async def test_touch_session_updates_only_direct_node() -> None: """Touching a child must update only the child, never the shared root.""" diff --git a/uv.lock b/uv.lock index 06f8fbaf..9dd31e09 100644 --- a/uv.lock +++ b/uv.lock @@ -233,7 +233,7 @@ wheels = [ [[package]] name = "context-intelligence-server" -version = "6.7.0" +version = "6.7.1" source = { editable = "." } dependencies = [ { name = "aiofiles" },