diff --git a/AGENTS.md b/AGENTS.md index a6edbc3c..b80acb4c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -82,19 +82,22 @@ local runs. Scope is intentionally narrow — do this and nothing else. context_intelligence_server/ # FastAPI ingestion server ├── main.py # App factory, routes, lifespan ├── config.py # Settings (YAML + env vars via Pydantic) -├── queue_manager.py # Durable per-session append-log (persist-then-202) -├── registry.py # Per-session drainers (drain_worker, write semaphore, retry/dead-letter) +├── queue_manager.py # Durable per-session append-log (persist-then-202); per-key _KeyGuard file lock (atomic single-writer append), compaction, dead-letter expiry, GC candidate scan +├── registry.py # Per-session drainers (drain_worker, write semaphore, retry/dead-letter, drain-task done-callback supervision) ├── pipeline.py # Per-event dispatch spine (invoked by the drainer) ├── neo4j_store.py # Managed-transaction Neo4j writes ├── blob_store.py # Async disk blob storage +├── idempotency.py # Request dedupe (seen/store split — a key is stored only AFTER a durable append) +├── writer_lease.py # Single-writer lease guard for the queue directory (heartbeat + staleness; enforce/detect/off, default enforce; surfaces on /status.writer_lease) ├── handlers/ # Event handlers (data_layer_1/2/3) │ ├── data_layer_1/ # Session/tool-call handlers │ ├── data_layer_2/ # Graph enrichment handlers │ └── data_layer_3/ # High-level insight handlers -├── routers/ # API routers (queues.py = dead-letter inspect/replay/purge; admin.py = /admin/* identity-map CRUD) +├── routers/ # API routers (queues.py = dead-letter inspect/replay/purge only \u2014 no operator GC endpoint; admin.py = /admin/* identity-map CRUD; version.py) ├── auth.py # Bearer-token auth middleware (StaticKeyResolver / EntraResolver via PrincipalResolver; BearerTokenMiddleware; admin-key recognition) +├── authz.py # Per-route capability gates (require_read / require_write) ├── identity_store.py # Durable JSON identity map (write-file-then-swap, fail-closed load, live flat_dict) -├── status.py # Status/version plumbing (EventRingBuffer, build_status_response, SERVER_VERSION, ring_buffer) +├── status.py # Status/version plumbing (EventRingBuffer, build_status_response, SERVER_VERSION, ring_buffer, BootState/boot_state — the boot-phase singleton behind /status.boot) └── models.py # Pydantic request/response models docs/ # ⚠️ PRODUCT DOCUMENTATION ONLY @@ -432,6 +435,11 @@ with admin"). Runtime runbook: `docs/identity-management.md`. ## Key Concepts - **Event pipeline** — `POST /events` persists the raw event to a durable per-session append-log (`queue_manager.py`) and returns `202` immediately (persist-then-202). A single drainer per session (`registry.py`) processes batches and flushes them to Neo4j under a global write semaphore, with transient/deadlock retry, dead-letter isolation of poison events, and crash recovery (replay + counter re-seed) on startup. Each handler invoked by the per-event dispatch spine is a Python class in `handlers/data_layer_*/`. +- **Append framing** — appends to one worker key are serialized by a per-key file lock (`_KeyGuard.file_lock`, held by the writing *thread*, not the coroutine), so a record lands as one contiguous newline-terminated line or not at all; a partial write is discarded and the error raised. The idempotency key is stored **only after** the durable append succeeds (`idempotency.py`: `seen()` / `store()` are separate) so a failed write plus a client retry is honoured, not falsely refused. +- **Drain supervision** — a drain task's done-callback (`registry.py`) is the single supervision point: a task that dies logs `drain_worker_died` at ERROR with the session id + traceback, deregisters the session, and closes its store. A poison line is dead-lettered and the offset advanced past it, so draining continues. +- **Boot phases** — boot recovery runs as a supervised background task (`main._boot_reconcile`), so `/status` and `/version` answer from the first phase. `status.BootState` (module singleton `boot_state`) tracks `recovering → heal → reclaim → expire → reconcile → seed → topup → sweep → ready`, or terminates at `failed` — and `failed` keeps serving. The phase is surfaced additively on `/status.boot`; while booting, `/status` does zero disk reads and reports `metrics`/`spool` as `null` with `status_detail.reason == "booting"`. +- **Self-shrinking queue storage** — a live session's committed prefix is reclaimed continuously (compaction); a fully-drained log is reclaimed automatically at boot regardless of `reclaim_enabled`; log-less dead-letters expire on an mtime window, opt-in (`dead_letter_expiry_enabled`, default off — a dead-letter may be the last surviving copy of an un-recovered event). There is no operator GC endpoint — reclamation is automatic. `reclaim_enabled` (default `false`) gates only the more aggressive boot reclaim of unresumable/reset-offset logs. +- **Writer lease** — `writer_lease.py` is a single-writer lease guard on the queue directory: a backstop against an accidental second writer (not a rolling-deploy coordinator; deployment is single-replica). Default mode `enforce` refuses to boot against a *live* foreign lease, releases the lease on clean shutdown, and takes over a *stale* foreign lease automatically. `detect` only surfaces a conflict on `/status.writer_lease` within a heartbeat, never refusing to boot. - **Graph model** — session sub-labels: `RootSession`, `SubSession`, `ForkedSession`, `IncompleteSession` (health marker; not a terminal). Full schema with all node and edge types: see `docs/architecture/03-graph-model.dot` and `docs/architecture/README.md`. - **Blob storage** — Large event payloads are written to disk and referenced by URI to avoid graph bloat. - **Configuration** — Pydantic Settings reads from `server-config.yaml` first, then environment variables. See `config.py`. diff --git a/README.md b/README.md index 01b95ffb..2d64e43d 100644 --- a/README.md +++ b/README.md @@ -18,6 +18,11 @@ Amplifier CLI sessions | - Retry + dead-letter + crash recovery | | 5 node / 8 edge types| | - Blob storage (large payloads to disk) | +----------------------+ | - API docs (/docs) + Cypher proxy | +| - Atomic single-writer append (per-key | +| file lock) + supervised drain workers | +| - Self-shrinking queue storage | +| (continuous compaction + dead-letter | +| retention) | +------------------------------------------+ ``` @@ -27,6 +32,20 @@ under a global write semaphore, retrying transient/deadlock failures and dead-le events. See [docs/architecture/05-durable-ingest-queue.png](docs/architecture/05-durable-ingest-queue.png) for the full ingest/drain flow. +Each append is serialised by a per-key file lock, so one worker key is written by exactly one +writer at a time and a record lands as one contiguous, newline-terminated line or not at all — +a write that fails part-way is discarded (and the failure surfaced to the caller) rather than +left as a torn or merged line. Drain workers are supervised: a worker that dies is logged at +`ERROR` with its session id and traceback, its session is deregistered and its store closed, +and an unparseable/poison line is written to that key's dead-letter file so draining continues +past it instead of halting. + +Queue storage shrinks itself. A live session's already-committed prefix is reclaimed +continuously (compaction) rather than only when the session ends, so a `.log` tracks the +undrained tail rather than the whole session history; dead-letter files that no longer have a +`.log` beside them expire on an mtime-based retention window. Both run automatically as part of +normal operation and need no operator action. + --- ## Upgrading: Cold Start No Longer Auto-Migrates @@ -311,7 +330,7 @@ Full runtime onboarding/offboarding runbook and the `/admin/*` API: | Method | Path | Description | |--------|------|-------------| | `POST` | `/events` | Ingest a session event (returns 202 immediately) | -| `GET` | `/status` | Server health, active sessions, completed history, error counts, `neo4j_connected`, `neo4j_query_connected` (reflects the read/cypher_query driver's connection health), `neo4j_url`, `neo4j_browser_url` | +| `GET` | `/status` | Server health, active sessions, completed history, error counts, `neo4j_connected`, `neo4j_query_connected` (reflects the read/cypher_query driver's connection health), `neo4j_url`, `neo4j_browser_url`, plus `boot` (boot-reconciliation phase + counters), `writer_lease` (queue-directory single-writer lease guard), `spool` (aggregate spool footprint) and `metrics` (aggregate conservation counters). `boot` and `writer_lease` are always present; **`metrics` and `spool` are `null` while the server is still booting**, and `status_detail` is then `{"reason": "booting"}`. HTTP `200` and `status: "ok"` at every boot phase — the boot phase is informational, never a liveness signal | | `GET` | `/version` | Server version (`{"version": "..."}`) — always unauthenticated | | `GET` | `/docs` | Swagger UI — always on (auth-exempt) | | `GET` | `/openapi.json` | OpenAPI spec — always on (auth-exempt) | @@ -399,6 +418,20 @@ Values are resolved with this priority (highest first): | `AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_QUEUES_PATH` | `queues_path` | `/data/queues` | Directory for the durable per-session append-logs (persist-then-202 ingest); mirrors `blob_path`. | | `AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_WRITE_CONCURRENCY` | `write_concurrency` | `8` | Max concurrent Neo4j write flushes across all session drainers (starvation guard). | | `AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_MAX_DELIVERY_ATTEMPTS` | `max_delivery_attempts` | `5` | Flush retries for one batch before its offending line is dead-lettered. | +| `AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_CRASH_RECOVERY_RESPAWN_LIMIT` | `crash_recovery_respawn_limit` | `8` | **Behaviour change for existing installs — this default was previously unbounded.** Ceiling on how many recovered sessions get a drainer respawned per boot/sweep pass. The remainder is *deferred*: left untouched on disk, still durable, and drained by a later sweep pass or the moment a new event for that session arrives. Deferral is never silent (a `WARNING` names the respawned/deferred counts). `null` restores unbounded; `0` disables automatic respawn at boot. Negative values are a hard startup error. | +| `AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_CRASH_RECOVERY_SWEEP_INTERVAL_SECONDS` | `crash_recovery_sweep_interval_seconds` | `60` | **Behaviour change for existing installs — this default was previously 300.** Interval at which the deferred backlog is re-scanned and the recovered-drainer pool topped back up to the ceiling. Only runs when `crash_recovery_respawn_limit` is finite — change the two together or neither. `0` disables the sweep (the deferred tail then drains only on restart or on a new event). | +| `AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_RECLAIM_ENABLED` | `reclaim_enabled` | `false` | Whether the boot-time reclaim pass may actually delete / reset anything. **Ships disabled**: boot still classifies every pre-existing queue key and logs the same audit line with `action=dry_run`, but nothing is unlinked. Boot once with this off, read the reclaim summary in the logs, reconcile it against expectations, then opt in. Surfaced on `/status` as `boot.reclaim_enabled`. | +| `AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_RECLAIM_REDRAIN_MAX_BYTES` | `reclaim_redrain_max_bytes` | `67108864` (64 MiB) | Size ceiling for a `.log` with a **negative or past-EOF** `.offset`: below the ceiling it is **re-drained from byte 0** instead of deleted — bounded, and idempotent by construction, so re-driving costs nothing. At or above the ceiling the `.log` is deleted instead. `0` means always delete. An **unparseable** `.offset` is a separate case and always re-drains from byte 0 at any size — never deleted. | +| `AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_QUEUE_COMPACT_ENABLED` | `queue_compact_enabled` | `true` | Continuously reclaim a live session's already-committed prefix, so a `.log` holds the undrained tail rather than the whole session history. Turning it off is a config change + restart (a kill switch, not a live toggle). | +| `AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_QUEUE_COMPACT_MIN_PREFIX_BYTES` | `queue_compact_min_prefix_bytes` | `8388608` (8 MiB) | Bounds compaction **frequency** on a continuously busy session: below this many committed bytes the rewrite is skipped. The idle path closes the remaining gap for free the moment the session goes quiet. `0` is a valid explicit opt-out; negative is a hard startup error. | +| `AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_DEAD_LETTER_EXPIRY_ENABLED` | `dead_letter_expiry_enabled` | `false` | Expire dead-letter files that have **no `.log` beside them** once they age past the retention window below. **Opt-in (default off)**: a dead-letter may be the only surviving copy of an un-recovered event, so it is never auto-deleted out of the box. Deliberately independent of `reclaim_enabled`: the predicate here is two plain filesystem facts, not a heuristic classification. | +| `AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_DEAD_LETTER_RETENTION_SECONDS` | `dead_letter_retention_seconds` | `2592000.0` (30 days) | How long a log-less `.dead.jsonl` survives (by mtime) before expiry — long enough to notice it via `GET /queues/dead-letter` and purge or replay it, short enough that it cannot accumulate forever. `<= 0` disables expiry outright; negative is a hard startup error. | +| `AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_WRITER_LEASE_MODE` | `writer_lease_mode` | `enforce` | Single-writer lease guard for the queue directory — a backstop against an accidental second writer, not a rolling-deploy coordinator (deployment is single-replica). `enforce` (default) **refuses to boot** against a *live* foreign lease, releases the lease on clean shutdown so a restart reacquires immediately, and takes over a *stale* foreign lease automatically once it ages past the staleness window. `detect` acquires the lease best-effort, heartbeats it, and only surfaces a conflict on `/status.writer_lease` — it never refuses to boot. `off` disables the guard. See [docs/operational-hardening.md](docs/operational-hardening.md). | +| `AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_WRITER_LEASE_HEARTBEAT_SECONDS` | `writer_lease_heartbeat_seconds` | `5.0` | Lease renew + re-read interval. Under `detect`, sets how fast a conflict becomes visible on `/status` (within one heartbeat); under `enforce`, sets how fast a foreign takeover is noticed. Must be `> 0`. | +| `AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_WRITER_LEASE_STALENESS_MULTIPLIER` | `writer_lease_staleness_multiplier` | `3.0` | Staleness window = heartbeat x this (15.0s at defaults). Must survive two consecutive missed ticks without a false "stale" verdict, so values **below `2.0` are a hard startup error**. | +| `AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_WRITER_LEASE_CONFIRM_DELAY_SECONDS` | `writer_lease_confirm_delay_seconds` | `1.0` | Settle delay between writing the lease and the confirming re-read, to exceed write→other-reader visibility latency on a shared SMB mount. Must be `>= 0`. | +| `AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_WRITER_LEASE_ACQUIRE_TIMEOUT_SECONDS` | `writer_lease_acquire_timeout_seconds` | `5.0` | Hard bound on an entire acquire (and, per tick, an entire renew), so a hung mount can never block startup or wedge the heartbeat. **Must exceed `writer_lease_confirm_delay_seconds`** — otherwise every acquire would time out and silently disarm the guard, so that combination is a hard startup error. | +| `AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_WRITER_LEASE_FORCE_ACQUIRE` | `writer_lease_force_acquire` | `false` | One-boot operator escape hatch: take over a **fresh** (non-stale) foreign lease instead of refusing (`enforce`) or merely latching (`detect`). Logs a `WARNING` on every boot while set and is surfaced on `/status.writer_lease`, so leaving it on by accident is never silent. | | `AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_LOG_PATH` | `log_path` | `/data/logs/server.jsonl` | Structured log file path | | `AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_LOG_LEVEL` | `log_level` | `INFO` | Log level (`DEBUG`/`INFO`/`WARNING`/`ERROR`) | | `AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_SERVER_HOST` | `server_host` | `0.0.0.0` | Bind host | @@ -433,6 +466,17 @@ recovered rather than lost across a restart. The durable per-session logs — no just Neo4j — are the record for events that have been accepted but not yet written to the graph. +**A `.log` is a transient buffer, not a session archive.** Once a batch is +committed, its bytes have reached the graph and the queue no longer needs them: the +already-committed prefix of a *live, still-open* session's `.log` is reclaimed +continuously (compaction), so the file tracks the undrained tail rather than the +whole session history. A fully-drained log is removed at session finalize. +Dead-letter files are likewise bounded: a +`.dead.jsonl` with no `.log` beside it expires once it ages past +`dead_letter_retention_seconds`. Read the queue directory as "what has not reached +the graph yet", never as "everything this server ever ingested" — the graph and the +blob store are the durable record. + For persistence on Azure Container Apps (persistent storage + Neo4j on AuraDB), see [docs/azure-deployment.md](docs/azure-deployment.md). diff --git a/amplifier-online.yaml b/amplifier-online.yaml index f32ae44d..b2ec3cb6 100644 --- a/amplifier-online.yaml +++ b/amplifier-online.yaml @@ -47,7 +47,10 @@ services: # blobs, and logs. Survives restarts/revision swaps. Single-instance enforced. volume: mount_path: /data - size_gib: 100 # Premium minimum is 100 GiB (billed on provisioned size) + size_gib: 1024 # Standing size 1 TiB — data grows; live quota was raised to + # 1 TiB out-of-band, so declare it here to keep manifest == actual + # (a re-apply then grows/holds, never shrinks below live usage). + # (Premium minimum is 100 GiB; billed on provisioned size.) tier: premium # Premium_LRS SSD share. Standard Azure Files (SMB) per-op # latency caps the write-heavy queue/blob throughput and # builds a standing write backlog under pressure. diff --git a/context_intelligence_server/auth.py b/context_intelligence_server/auth.py index a395a1d3..90492a31 100644 --- a/context_intelligence_server/auth.py +++ b/context_intelligence_server/auth.py @@ -6,20 +6,15 @@ from collections.abc import Callable, MutableMapping from typing import Any -import jwt # pyjwt[crypto] — added in T1; used by EntraResolver +import jwt # pyjwt[crypto] — used by EntraResolver from jwt import PyJWKClient from typing_extensions import Protocol _log = logging.getLogger(__name__) -# JWKS signing-key cache TTL passed to PyJWKClient. -# -# PyJWKClient handles per-kid caching and lifespan-bounded refresh natively; -# making the value explicit here keeps the contract visible in code even -# though 300 s matches the library default. No custom per-kid dedup lock -# or global refresh cap is built for the pilot (council/cranky: pilot-scale -# does not justify that complexity — revisit at scale if stampede behaviour -# is observed in production metrics). +# JWKS signing-key cache TTL passed to PyJWKClient (matches its own default; +# kept explicit so the contract is visible in code). PyJWKClient handles +# per-kid caching and refresh natively -- no custom dedup lock needed. JWKS_CACHE_LIFESPAN_SECONDS: int = 300 # Paths that are exempt from authentication: health checks, version info, and @@ -39,12 +34,9 @@ # longer serves any static assets (the dashboard's /static/ mount was removed). _EXEMPT_PREFIXES: tuple[str, ...] = () -# Route prefix of the admin router (mirrors routers/admin.py: -# ``APIRouter(prefix="/admin", ...)``). The static-mode admin-key fast-path is -# scoped to these paths: the admin key is an administration credential, NOT a -# data-ingestion identity, so it must only short-circuit auth on /admin/* — never -# on data routes like POST /events (where it would otherwise stamp a synthetic -# ``created_by="admin"`` and let a bare admin key post events). +# The admin-key fast-path is scoped to these paths only: the admin key is an +# administration credential, not a data-ingestion identity, so it must never +# short-circuit auth on data routes like POST /events. _ADMIN_ROUTE_PREFIX: str = "/admin" @@ -61,21 +53,10 @@ def _is_admin_route(path: str) -> bool: class AuthError(Exception): """Authentication/authorisation failure with a specific HTTP status code. - Raised by :class:`EntraResolver` (and may be raised by future resolvers) - to communicate *why* a request was rejected, not just *that* it was. - - ``status_code``: - 401 — token is missing, malformed, expired, has wrong audience/issuer/ - tenant, fails signature verification, uses a disallowed algorithm, - or is missing required claims (``oid``, ``scp``). - 403 — token is cryptographically valid but the ``oid`` is not in the - identity map (``bearer_identity_unbound``), or a service token - with no qualifying App Role. - - ``reason``: - Short human-readable message for logging / response bodies. The 403 - reason MUST name the unbound ``oid`` so operators can diagnose and add - the missing entry. + 401: token missing/malformed/expired/wrong-aud-or-iss/wrong-tenant/bad-sig + or missing required claims. 403: token is cryptographically valid but the + ``oid`` is unmapped, or a service token has no qualifying App Role -- the + ``reason`` must name the unbound ``oid`` so operators can diagnose it. """ def __init__(self, status_code: int, reason: str) -> None: @@ -110,31 +91,14 @@ def _resolve_token(token: str, keystore: dict[str, str]) -> str | None: class PrincipalResolver(Protocol): """Resolves a raw bearer token string to a contributor id. - Returns ``(contributor_id, roles, is_service)`` on success, or ``None`` when - the token is not recognised (caller should respond 401). Implementations may - also raise :class:`AuthError` to signal specific auth failures (401 or 403); - the middleware dispatches the exact ``status_code`` from the exception. - - The ``roles`` element (a list of strings) carries the token's App-Role - claim for the entra resolver, or an empty list for the static resolver. - The middleware stores these on ``scope["state"]["roles"]`` so downstream - dependencies (e.g. ``require_admin``) can read them without re-parsing. - - ``is_service`` is ``True`` for app/service tokens resolved by the service - branch, ``False`` for delegated user tokens and static-key tokens. The - middleware writes this onto ``scope["state"]["is_service"]`` so route - capability deps (``require_write`` / ``require_read``) can gate service - principals without re-parsing the token. - - Only one concrete implementation exists today: :class:`StaticKeyResolver`. - The ``EntraResolver`` (JWT via JWKS) is added in T4. Do NOT add a third - resolver without a separate design review. - - M2 protocol change: ``resolve()`` previously returned - ``tuple[str, list[str]] | None``. Changed to - ``tuple[str, list[str], bool] | None`` to carry the ``is_service`` flag - so the middleware can set capability state without re-parsing the token. - ``StaticKeyResolver`` always returns ``is_service=False``. + Returns ``(contributor_id, roles, is_service)`` on success, or ``None`` + when the token is not recognised (caller should respond 401). + Implementations may also raise :class:`AuthError` for a specific status. + + ``roles`` carries the token's App-Role claim (entra) or ``[]`` (static), + stored on ``scope["state"]["roles"]``. ``is_service`` is ``True`` for + app/service tokens, stored on ``scope["state"]["is_service"]`` so route + capability deps can gate service principals without re-parsing. """ @property @@ -162,13 +126,11 @@ def resolve( routed through the service branch); ``False`` for delegated user tokens and static-key tokens. - ``admin_path`` (keyword-only, default ``False``) signals that the - request targets an ``/admin/*`` route. When ``True``, ``EntraResolver`` - relaxes ONLY the identity-map membership check (an unbound-but-valid - oid is admitted so an IdentityAdmin role-holder can bootstrap the map); - NO token-authenticity check (signature/issuer/audience/expiry/tenant/ - scope/oid-presence) is ever relaxed. ``StaticKeyResolver`` ignores this - parameter entirely — its admin authorization is via a separate + ``admin_path`` (keyword-only, default ``False``) signals a request + targeting ``/admin/*``. When ``True``, ``EntraResolver`` relaxes only + the identity-map membership check (never token authenticity), so an + IdentityAdmin role-holder can bootstrap the map. ``StaticKeyResolver`` + ignores this parameter -- its admin authorization is a separate admin-key fast-path, not map membership. """ ... @@ -177,52 +139,36 @@ def resolve( class EntraResolver: """Resolves Entra RS256 bearer tokens to contributor ids. - Mirrors ``validate_entra_token()`` from Team Pulse - (``amplifier-app-team-pulse`` / ``team_pulse/identity/extractors.py``): PyJWKClient → ``jwt.decode`` with ``algorithms=["RS256"]``, dual audience, explicit ``tid`` + ``scp`` + ``oid`` checks, then ``oid → contributor_id`` - lookup via *identity_map*. - - M2: adds a second branch for app/service tokens (no ``scp``) selected by - a ``scp``-presence discriminator, authorized by App-Role alone, with a - fail-loud ``created_by`` derived from stable claims. + lookup via *identity_map*. A second branch handles app/service tokens (no + ``scp``), authorized by App-Role alone, with `created_by` derived from + stable claims. Raises: - :class:`AuthError` (401): Token is missing/malformed/expired/wrong-aud/ - wrong-iss/wrong-tid/fails-sig-verification/wrong-alg, [B1] anomaly, - or missing/invalid identity claim. - :class:`AuthError` (403): Token is cryptographically valid but the - lowercased ``oid`` is not in *identity_map* (user branch), or no - qualifying App Role is present (service branch). - RuntimeError: At construction if eager JWKS prefetch fails — the server - must refuse to start rather than lazily fail at first request (§8b). + :class:`AuthError` (401): token invalid/expired/wrong-aud-or-iss/ + wrong-tenant, an anomalous scp+idtyp=app combo, or missing/invalid + identity claim. + :class:`AuthError` (403): token valid but ``oid`` unmapped (user + branch), or no qualifying App Role (service branch). + RuntimeError: at construction if eager JWKS prefetch fails -- the + server must refuse to start rather than fail lazily. Args: client_id: Azure App Registration client ID (GUID). tenant_id: Azure AD tenant ID (GUID). - identity_map: ``{oid_lower -> contributor_id}`` — built by - :meth:`~context_intelligence_server.config.Settings.build_identity_map`. - MAY be empty at construction: an empty map is a - supported bootstrap state (the server boots - fail-closed and is populated at runtime via the - IdentityAdmin-gated /admin/identities API). A live - reference is passed so runtime PUT/DELETE are - visible immediately. On a data route an unmapped - oid still 403s; the map-miss is exempted ONLY for - /admin/* paths (``admin_path=True``) so a role- - holder can bootstrap the first identity. + identity_map: ``{oid_lower -> contributor_id}``. May be empty + at construction (supported bootstrap state, + populated at runtime via /admin/identities). A + live reference so runtime mutation is visible + immediately; the map-miss exemption is scoped + to ``admin_path=True`` only. service_identity_map: ``{oid_lower -> contributor_id}`` for service - principals. Optional; ``{}`` = no service map. - service_data_role: App Role name granting write access. ``""`` disables. - reader_role: App Role name granting read-only access. ``""`` disables. - entra_admin_role: App Role name granting admin access. ``""`` disables. + principals. Optional; ``{}`` = no service map. + service_data_role: App Role name granting write access. ``""`` disables. + reader_role: App Role name granting read-only access. ``""`` disables. + entra_admin_role: App Role name granting admin access. ``""`` disables. jwks_client: Injectable JWKS client for tests. - - Note — JWKS caching (T5): - Per-``kid`` caching and lifespan-bounded refresh are handled by - ``PyJWKClient`` (``lifespan=JWKS_CACHE_LIFESPAN_SECONDS``). No custom - per-``kid`` dedup lock or global refresh cap is built for the pilot; - revisit at scale if stampede behaviour appears in production metrics. """ def __init__( @@ -231,23 +177,22 @@ def __init__( tenant_id: str, identity_map: dict[str, str], *, - service_identity_map: dict[str, str] | None = None, # NEW (M2) - service_data_role: str = "", # NEW (M2) - reader_role: str = "", # NEW (M2) - entra_admin_role: str = "", # NEW (M2) + service_identity_map: dict[str, str] | None = None, + service_data_role: str = "", + reader_role: str = "", + entra_admin_role: str = "", jwks_client: Any = None, ) -> None: self._client_id = client_id self._tenant_id = tenant_id self._identity_map = identity_map - # M2 service-path config — fail-closed defaults (empty disables each role). + # Fail-closed defaults: empty string disables each role. self._service_identity_map: dict[str, str] = service_identity_map or {} self._service_data_role = service_data_role self._reader_role = reader_role self._entra_admin_role = entra_admin_role # Accept both the bare client GUID (ID-token aud) and the api:// form - # (access-token aud when access_as_user scope is exposed). Matches - # the Team Pulse mirror and the Q-AUD confirmation from §2b. + # (access-token aud when access_as_user scope is exposed). self._expected_aud = [client_id, f"api://{client_id}"] self._expected_issuer = f"https://login.microsoftonline.com/{tenant_id}/v2.0" @@ -257,32 +202,23 @@ def __init__( ) jwks_client = PyJWKClient(jwks_uri, lifespan=JWKS_CACHE_LIFESPAN_SECONDS) - # Eager prefetch — fail-closed at startup (§8b / crusty gate). - # Called regardless of whether the client was injected or built by - # default so that tests can inject a _FailingJWKSClient and verify - # the fail-closed guarantee. - # Per-kid caching and lifespan-bounded refresh are handled by - # PyJWKClient (lifespan=JWKS_CACHE_LIFESPAN_SECONDS). No custom - # per-kid dedup lock or global refresh cap is built for the pilot. + # Eager prefetch, fail-closed at startup -- runs regardless of whether + # the client was injected, so tests can verify the fail-closed guarantee. try: jwks_client.fetch_data() - except Exception as exc: # noqa: BLE001 — any failure is fatal here + except Exception as exc: raise RuntimeError( f"EntraResolver: JWKS prefetch failed for tenant " f"{tenant_id!r} — server cannot start without a reachable " f"JWKS endpoint. Cause: {exc}" ) from exc - # Guard: a reachable-but-empty JWKS ({"keys": []}) would let - # construction succeed but then 401 every request lazily. Detect it - # here so the server refuses to start rather than silently degrading. - # Uses get_jwk_set() if available; stubs that pre-date this check - # (AttributeError) are tolerated — all production PyJWKClient - # instances expose the method. + # Guard: a reachable-but-empty JWKS would let construction succeed + # but then 401 every request lazily -- detect and refuse to start. try: jwk_set = jwks_client.get_jwk_set() except AttributeError: - pass # Pre-existing stub without get_jwk_set() — skip the check + pass # stub without get_jwk_set() — skip the check else: if not jwk_set.keys: raise RuntimeError( @@ -303,72 +239,46 @@ def resolve( ) -> tuple[str, list[str], bool]: """Validate Entra JWT; return (contributor_id, roles, is_service). - Implements the dual-path discriminator (M2): - - Tokens with ``scp`` present → USER / delegated branch (unchanged from V1). - - Tokens without ``scp`` → SERVICE / app branch (new in M2). - The ``[B1]`` anomaly check fires first when both ``scp`` and - ``idtyp=app`` are present — neither branch can claim such a token - (fail-closed, 401). - - Args: - token: Raw bearer token string (``Authorization: Bearer`` prefix - already stripped by the middleware). - - Returns: - A ``(contributor_id, roles, is_service)`` 3-tuple. - *contributor_id* is mapped from ``oid`` (user branch) or derived - via ``service_identity_map``/``appid``/``azp``/``oid`` (service - branch). *roles* is the token's App Role assignments as a list of - strings. *is_service* is ``True`` for app/service tokens, - ``False`` for delegated user tokens. + Tokens with ``scp`` present → user/delegated branch. Tokens without + ``scp`` → service/app branch. A token carrying both ``scp`` and + ``idtyp=app`` is anomalous and rejected (401) before either branch + can claim it. Raises: - AuthError(401): JWT validation failure, wrong tenant, [B1] anomaly, - missing/invalid ``oid`` (user branch), or no resolvable identity - (service branch). - AuthError(403): Valid user token with unmapped ``oid``, or valid - service token with no qualifying App Role [D7]. + AuthError(401): JWT validation failure, wrong tenant, the + scp+idtyp anomaly, missing/invalid ``oid``, or no resolvable + service identity. + AuthError(403): valid user token with unmapped ``oid``, or valid + service token with no qualifying App Role. """ - # ---- SHARED VALIDATION (both paths) — UNCHANGED from V1 ---- + # ---- Shared validation (both paths) ---- try: key = self._jwks_client.get_signing_key_from_jwt(token).key claims = jwt.decode( token, key, - algorithms=["RS256"], # H1: pin RS256, reject alg=none / HS256 - audience=self._expected_aud, # B7: aud enforced here - issuer=self._expected_issuer, # B7: iss enforced here + algorithms=["RS256"], # pin RS256, reject alg=none / HS256 + audience=self._expected_aud, + issuer=self._expected_issuer, options={"require": ["exp", "iss", "aud"]}, ) except jwt.PyJWTError as exc: - # Covers: InvalidSignatureError, ExpiredSignatureError, - # InvalidAudienceError, InvalidIssuerError, InvalidAlgorithmError, - # MissingRequiredClaimError, PyJWKClientError, ImmatureSignatureError - # (nbf), and all other PyJWT validation failures. raise AuthError(401, f"Invalid bearer token: {exc}") from exc - # Explicit tid check — defense-in-depth alongside the issuer pin. - # A v2 Entra token's iss already encodes the tenant, but the explicit - # check makes the tenant binding self-documenting and mirrors TP. + # Explicit tid check, defense-in-depth alongside the issuer pin. if claims.get("tid") != self._tenant_id: raise AuthError(401, "Token from wrong tenant") - # ---- DISCRIMINATOR (M2) — scp PRIMARY, idtyp CONFIRMATION ---- - # scp normalization is IDENTICAL to V1 (non-string -> ""). _scp_raw = claims.get("scp") scp: str = _scp_raw if isinstance(_scp_raw, str) else "" - has_scp: bool = bool( - scp.split() - ) # any whitespace-delimited scope token present + has_scp: bool = bool(scp.split()) - # idtyp normalization [B2]: non-string -> "", then strip().lower(). _idtyp_raw = claims.get("idtyp") idtyp: str = _idtyp_raw.strip().lower() if isinstance(_idtyp_raw, str) else "" - # [B1] Branches MUST be mutually exclusive. A token bearing BOTH a - # delegated scope AND idtyp=="app" is anomalous (no legitimate Entra - # token does this) -> fail closed. Checked FIRST so neither branch - # can claim it. + # Branches must be mutually exclusive: a token with both a delegated + # scope and idtyp=="app" is anomalous and rejected before either + # branch can claim it. if has_scp and idtyp == "app": raise AuthError( 401, @@ -377,38 +287,22 @@ def resolve( ) if has_scp: - # ========================================================= - # USER / DELEGATED BRANCH — BYTE-FOR-BYTE auth.py V1 logic - # (only the return arity changes: append is_service=False) - # ========================================================= + # User / delegated branch. if "access_as_user" not in scp.split(): raise AuthError( 401, f"Token missing required scope 'access_as_user' (got scp={scp!r})", ) - # oid is required. Missing/non-string/whitespace oid -> 401, NOT 403 - # (AC12). isinstance guard prevents AttributeError on non-string oid - # (e.g. int 42, list) — FAIL-1/FAIL-3 fix. oid = claims.get("oid") if not isinstance(oid, str) or not oid.strip(): raise AuthError(401, "Token missing or invalid 'oid' claim") - # Map oid -> contributor — unrecognised oid is a 403 (identity_unbound). # Both sides lowercased: config validator lowercases keys at build time. oid_lower = oid.lower() contributor_id = self._identity_map.get(oid_lower) if contributor_id is None: - # BOOTSTRAP EXEMPTION — /admin/* paths ONLY (admin_path=True): - # a cryptographically-valid delegated token whose oid is not yet - # bound is admitted to routing so an IdentityAdmin role-holder can - # populate the map on a fresh deployment. Authorization is still - # enforced downstream by require_admin on the `roles` claim; a - # non-admin unbound token reaches /admin and is 403'd there. - # - # SECURITY: this relaxes ONLY the oid->id map-membership lookup. - # All JWT authenticity checks (signature, issuer, audience, - # expiry, tenant, access_as_user scope, oid presence) have already - # passed above and are NOT affected. On every non-admin path - # admin_path is False, so an unmapped oid is still a hard 403. + # Bootstrap exemption, /admin/* only: an unbound-but-valid oid + # is admitted so an IdentityAdmin role-holder can populate the + # map; require_admin still enforces the roles claim downstream. if not admin_path: raise AuthError( 403, @@ -416,27 +310,18 @@ def resolve( f"identity map; contact the server administrator to add this " f"identity (tenant {self._tenant_id!r})", ) - # Provisional contributor id = the oid itself, so the admin audit - # log records who performed the bootstrap mutation even though - # they are not yet a mapped contributor. + # Provisional contributor id = the oid itself, for audit trail. contributor_id = oid_lower - # Roles: list[str] normalization — only `roles`, never `groups` (TB-09). + # Only `roles`, never `groups`. _roles_raw = claims.get("roles") roles: list[str] = ( [r for r in _roles_raw if isinstance(r, str)] if isinstance(_roles_raw, list) else [] ) - return ( - contributor_id, - roles, - False, - ) # <-- only delta from V1: third element - - # ========================================================= - # SERVICE / APP BRANCH (NEW, M2) — scp ABSENT - # ========================================================= - # Roles normalization is identical to the user branch. + return (contributor_id, roles, False) + + # Service / app branch (scp absent). _roles_raw = claims.get("roles") roles = ( [r for r in _roles_raw if isinstance(r, str)] @@ -444,17 +329,14 @@ def resolve( else [] ) - # --- Authorization = ROLE ALONE [D7]. Admit iff a qualifying configured - # role is present. Empty name disables that role (config.py:504-513). + # Authorization = role alone; empty name disables that role. authorized = ( (self._service_data_role and self._service_data_role in roles) or (self._reader_role and self._reader_role in roles) or (self._entra_admin_role and self._entra_admin_role in roles) ) if not authorized: - # 403 message: name the rejected principal (appid preferred over oid) - # and the required App Roles. Do NOT echo roles=[...] in the response - # body — that is an internal claim value, not operator guidance (R1). + # Name the rejected principal, not the raw roles claim. _appid_raw = claims.get("appid") _oid_raw_msg = claims.get("oid") _principal = ( @@ -478,9 +360,8 @@ def resolve( f"in Azure Entra, then re-request a token.", ) - # --- created_by derivation [B6/B8]: stable claims, truthiness chaining, - # NEVER app_displayname (spoofable in Entra — B8), fail-loud. - # Order: service_map[oid] > appid > azp > oid. + # created_by derivation: stable claims only, never app_displayname + # (spoofable in Entra). Order: service_map[oid] > appid > azp > oid. _oid_raw = claims.get("oid") oid_str = _oid_raw if isinstance(_oid_raw, str) and _oid_raw.strip() else "" oid_lower = oid_str.lower() @@ -506,13 +387,8 @@ def resolve( class StaticKeyResolver: """Resolves tokens via a pre-built ``{sha256_hex(token) -> contributor_id}`` keystore. - This is a pure extraction of the inline logic that previously lived in - :class:`BearerTokenMiddleware.__call__`. Behaviour is byte-for-byte - identical to the previous implementation. - - The keystore is built by :meth:`~context_intelligence_server.config.Settings.build_keystore` - and maps the SHA-256 hex digest of each raw bearer token to the owner's - contributor id string. Raw tokens are never stored here. + Built by :meth:`~context_intelligence_server.config.Settings.build_keystore`. + Raw tokens are never stored here. """ def __init__(self, keystore: dict[str, str]) -> None: @@ -522,24 +398,16 @@ def __init__(self, keystore: dict[str, str]) -> None: def auth_enabled(self) -> bool: """True when at least one key is configured (authentication is active). - ``False`` means the keystore is empty. This alone NO LONGER makes the - server pass requests through: an empty keystore now boots fail-CLOSED - (a supported bootstrap state) and every request 401s until keys are - onboarded via the /admin/keys API. The ONLY way requests pass through - unauthenticated is the explicit ``allow_unauthenticated=True`` opt-out - combined with this returning ``False`` — see - :class:`BearerTokenMiddleware` and - :func:`~context_intelligence_server.main.create_asgi_app`. + An empty keystore boots fail-closed (401 until onboarded via + /admin/keys); the only unauthenticated path is the explicit + ``allow_unauthenticated=True`` opt-out combined with this returning + ``False``. """ return bool(self._keystore) @property def is_empty(self) -> bool: - """True when no keys are configured. - - Kept for backward compatibility with existing tests. Prefer - ``auth_enabled`` (its logical inverse) for new code. - """ + """True when no keys are configured. Prefer ``auth_enabled`` (inverse) for new code.""" return not self._keystore def resolve( @@ -547,17 +415,9 @@ def resolve( ) -> tuple[str, list[str], bool] | None: """Return ``(contributor_id, [], False)`` for *token*, or ``None`` on a miss. - The roles list is always empty for static-key auth — admin authority is - signalled via ``scope["state"]["is_admin"]`` by the middleware (which - recognises the admin key before calling this resolver). - - ``is_service`` is always ``False`` for static-key tokens — they behave - like humans (always write-capable), preserving static-mode behavior. - - ``admin_path`` is accepted for Protocol compatibility with - :class:`EntraResolver` but is unused here: static-mode admin - authorization goes through the admin-key fast-path (matched before - this resolver is ever called), not identity-map membership. + Roles is always empty and is_service always False for static-key + auth. ``admin_path`` is accepted for Protocol compatibility but + unused: static-mode admin goes through a separate key fast-path. """ _ = admin_path # unused: static-mode admin uses the admin-key fast-path contributor_id = _resolve_token(token, self._keystore) @@ -569,50 +429,23 @@ def resolve( class BearerTokenMiddleware: """ASGI middleware that validates ``Authorization: Bearer `` headers. - Accepts a :class:`PrincipalResolver` via the *resolver* keyword argument - (preferred — used by :func:`~context_intelligence_server.main.create_asgi_app`), - or a raw *keystore* dict for backward compatibility with tests that - construct the middleware directly. - - Fail-open pass-through happens ONLY when BOTH conditions hold: the - middleware was constructed with ``allow_unauthenticated=True`` (the explicit - test/dev opt-out) AND the resolver's ``auth_enabled`` property is ``False`` - (i.e. a :class:`StaticKeyResolver` built with an empty keystore). An empty - keystore ALONE no longer passes requests through — with the production - default (``allow_unauthenticated=False``) an empty keystore fail-CLOSES: - the request falls through to token extraction and the resolver returns - ``None`` → 401. Entra mode is unaffected: ``EntraResolver.auth_enabled`` is - always ``True``, so this fail-open branch can never fire there. - - On a successful match the following keys are injected into - ``scope["state"]`` so downstream handlers can read authenticated identity - without re-resolving: - - * ``contributor_id`` (str): the resolved contributor. - * ``is_admin`` (bool): ``True`` only when the static-mode admin key was - used. Always ``False`` for regular data keys and for entra tokens (where - admin authority is carried in ``roles`` instead). - * ``roles`` (list[str]): App Role assignments from the Entra ``roles`` - claim, or ``[]`` for static-mode tokens. The ``require_admin`` - dependency checks this list for the ``IdentityAdmin`` role. - * ``is_service`` (bool): ``True`` for app/service tokens resolved by the - service branch; ``False`` for human/delegated and static-key tokens. - Used by ``require_write`` / ``require_read`` in ``main.py`` to gate - service principals without re-parsing the token. - - T5 — static-mode admin key (ROB F1): - - The admin key (``admin_api_key_digest``) is not in the data keystore, so - it would normally fail the resolver and produce a 401. Instead, the - middleware checks the bearer token's sha256 against the admin-key digest - BEFORE delegating to the resolver. A match authenticates the request - directly with ``contributor_id="admin"`` and ``is_admin=True``. The - token still reaches data-API endpoints (it is a valid principal) but - ``require_admin`` passes only for admin-key bearers. - - Several paths are always exempt (see ``_EXEMPT_PATHS``) so health checks, - monitoring tools, and public-facing pages continue working without - credentials. + Accepts a :class:`PrincipalResolver` via *resolver* (preferred), or a raw + *keystore* dict for tests that construct the middleware directly. + + Fail-open pass-through happens only when the middleware was constructed + with ``allow_unauthenticated=True`` AND the resolver's ``auth_enabled`` is + ``False`` (an empty static keystore). Entra mode can never fire this + branch (``EntraResolver.auth_enabled`` is always ``True``). + + On a match, injects into ``scope["state"]``: ``contributor_id``, + ``is_admin`` (True only for the static-mode admin key), ``roles`` (Entra + App Role assignments, ``[]`` for static mode), and ``is_service``. + + The admin key digest is checked against the bearer token's sha256 before + delegating to the resolver, authenticating as ``contributor_id="admin"``, + ``is_admin=True`` directly. + + ``_EXEMPT_PATHS`` bypasses auth for health checks and public-facing pages. """ def __init__( @@ -626,31 +459,19 @@ def __init__( allow_unauthenticated: bool = False, ) -> None: self.app = app - # Explicit opt-out (test/dev ONLY): when True AND the resolver has no - # credentials configured (auth_enabled is False), ALL requests pass - # through unauthenticated. An empty keystore ALONE no longer fails - # open — see the fail-open check in __call__ for the full rationale. + # Test/dev opt-out only: see the fail-open check in __call__. self._allow_unauthenticated: bool = allow_unauthenticated if resolver is not None: - # Preferred path: caller explicitly constructed and wired the resolver. self.resolver: PrincipalResolver = resolver else: - # Backward-compat path: construct a StaticKeyResolver from the - # provided (or defaulted-to-empty) keystore dict. + # Backward-compat: construct a StaticKeyResolver from the keystore. ks: dict[str, str] = keystore if keystore is not None else {} self.resolver = StaticKeyResolver(ks) - # Exempt paths: which exact paths bypass auth entirely. Defaults to - # the single API-only set (_EXEMPT_PATHS: /status, /version, /docs, - # /openapi.json). Path prefixes (_EXEMPT_PREFIXES) are always applied - # in addition to this set. self._exempt_paths: frozenset[str] = ( exempt_paths if exempt_paths is not None else _EXEMPT_PATHS ) - # Admin key digest (T5 / ROB F1): sha256 hex of the raw admin_api_key. - # When set, a bearer token whose sha256 matches this digest is - # authenticated as the "admin" principal with is_admin=True and bypasses - # the data keystore. None means admin key is not configured (static - # mode) or is irrelevant (entra mode — admin is via roles claim). + # sha256 hex of admin_api_key; None when not configured or irrelevant + # (entra mode -- admin comes via the roles claim instead). self._admin_api_key_digest: str | None = admin_api_key_digest async def __call__( @@ -660,26 +481,16 @@ async def __call__( await self.app(scope, receive, send) return - # Fail-open pass-through ONLY when the operator EXPLICITLY opted out via - # allow_unauthenticated=True (test/dev) AND the resolver has no - # credentials configured. An empty keystore ALONE no longer fails open: - # with allow_unauthenticated=False (production default) an empty static - # keystore fail-CLOSES — the request falls through to token extraction - # and the resolver returns None -> 401. This is the change that makes an - # empty-keystore boot a SAFE bootstrap state instead of a wide-open one. - # - # SECURITY: entra mode is unaffected — EntraResolver.auth_enabled is - # always True, so `not self.resolver.auth_enabled` is always False and - # this branch can never fire in entra mode regardless of the flag. + # Fail-open only when the operator explicitly opted out AND the + # resolver has no credentials configured. Entra mode can never fire + # this (EntraResolver.auth_enabled is always True). if self._allow_unauthenticated and not self.resolver.auth_enabled: await self.app(scope, receive, send) return path: str = scope.get("path", "") - # Compute once: is this an /admin/* route? Used by BOTH the static-mode - # admin-key fast-path below AND the entra bootstrap exemption passed into - # resolver.resolve(admin_path=...). Scoping the map-membership exemption - # to admin paths is what keeps every data route hard-gated. + # Used by both the admin-key fast-path and the entra bootstrap + # exemption; scoping to admin paths keeps data routes hard-gated. is_admin_path: bool = _is_admin_route(path) if path in self._exempt_paths or any( path.startswith(p) for p in _EXEMPT_PREFIXES @@ -693,27 +504,10 @@ async def __call__( await _send_401(send) return - # T5 / ROB F1 — static-mode admin key check, SCOPED to /admin/* routes. - # - # The admin key gates the /admin/* endpoints ONLY — it is an - # administration credential, not a data-ingestion identity. It is NOT in - # the data keystore, so on an admin route it would fail the resolver and - # produce a 401; the fast-path below checks the bearer token's sha256 - # against the admin-key digest BEFORE the resolver so an admin-key bearer - # authenticates with is_admin=True (which require_admin needs). - # - # Restricting this to admin routes is the fix for the identity-conflation - # bug: on a data route (e.g. POST /events) the admin key MUST NOT - # short-circuit auth. Instead it falls through to the resolver, so: - # - a token that is ALSO a registered data key resolves to its real - # contributor id (created_by reflects that id, never "admin"); and - # - a bare admin key that is not a data key is correctly rejected (401) - # rather than posting events attributed to a synthetic "admin". - # is_admin is only ever read on /admin/* (require_admin), so scoping the - # fast-path there loses no authorization capability. - # - # Note: entra mode does not use admin_api_key_digest (it is always - # None in entra mode); admin authority comes from the roles claim. + # Admin key gates /admin/* only -- it is an administration credential, + # not a data-ingestion identity. On a data route it must fall through + # to the resolver instead, so a registered data key still resolves to + # its real contributor id rather than a synthetic "admin". if self._admin_api_key_digest is not None and is_admin_path: token_digest = hashlib.sha256(token.encode()).hexdigest() if token_digest == self._admin_api_key_digest: @@ -721,21 +515,13 @@ async def __call__( state["contributor_id"] = "admin" state["is_admin"] = True state["roles"] = [] - state["is_service"] = False # admin key behaves like a human (M2) + state["is_service"] = False await self.app(scope, receive, send) return try: - # admin_path=True relaxes ONLY the oid->id map-membership lookup for - # /admin/* (bootstrap): an unbound-but-valid delegated token reaches - # require_admin, which then authorizes on the role claim. On every - # non-admin path admin_path=False, so an unmapped oid is still 403. result = self.resolver.resolve(token, admin_path=is_admin_path) except AuthError as exc: - # EntraResolver (and future resolvers) raise AuthError to communicate - # 401 vs 403. Dispatch the status code directly. - # Log at INFO — distinguishable from unexpected errors (ERROR). - # auth_event=auth_denied is a greppable marker for "bad token rejected". _log.info( "auth_event=auth_denied: %s (status=%d)", exc.reason, @@ -743,13 +529,9 @@ async def __call__( ) await _send_error(send, exc.status_code, exc.reason) return - except Exception: - # Defense-in-depth catch-all: any unexpected exception from the - # resolver (e.g. a transient library bug) must not propagate as a - # 500 — respond fail-closed (401) and log loudly for operators. - # auth_event=resolver_unexpected_exception distinguishes this from - # a normal auth denial so operators can grep specifically for it. - # The raw token is intentionally NOT logged (credential hygiene). + except Exception: # noqa: BLE001 -- defense-in-depth catch-all + # Any unexpected resolver exception must not propagate as a 500 -- + # respond fail-closed and log loudly. Raw token is never logged. _log.error( "auth_event=resolver_unexpected_exception: unexpected error in " "resolver.resolve() — denying request fail-closed " @@ -760,15 +542,8 @@ async def __call__( return if result is None: - # Backward-compat path: StaticKeyResolver returns None on a miss. - # Log it with the SAME greppable auth_event=auth_denied marker as the - # AuthError branch above so a static-key rejection is not invisible in - # server.jsonl. Without this line a genuine 401 leaves zero trace, - # which made a real "Bearer [REDACTED]" rejection look impossible to - # diagnose. The raw token is intentionally NOT logged (credential - # hygiene); a short sha256 fingerprint is emitted so operators can - # correlate the rejected credential (e.g. the redaction sentinel - # "[REDACTED]" has a recognisable digest) without exposing a secret. + # Raw token is never logged; a short sha256 fingerprint lets + # operators correlate the rejected credential without exposing it. _log.info( "auth_event=auth_denied: static key not recognized (status=401) " "token_sha256=%s", @@ -777,16 +552,15 @@ async def __call__( await _send_401(send) return - contributor_id, roles, is_service = result # M2: unpack 3-tuple + contributor_id, roles, is_service = result - # Inject authenticated identity and auth metadata into scope state. - # is_admin is False for regular data keys and for entra tokens (admin - # authority for entra is signalled via the roles list, not this flag). + # is_admin is False here for entra tokens too -- admin authority for + # entra is signalled via the roles list instead. state = scope.setdefault("state", {}) state["contributor_id"] = contributor_id state["is_admin"] = False state["roles"] = roles - state["is_service"] = is_service # M2: capability signal for route deps + state["is_service"] = is_service await self.app(scope, receive, send) diff --git a/context_intelligence_server/config.py b/context_intelligence_server/config.py index 48b20e97..6c776c9b 100644 --- a/context_intelligence_server/config.py +++ b/context_intelligence_server/config.py @@ -11,6 +11,7 @@ import hashlib import logging +import math import os import re from functools import lru_cache @@ -26,21 +27,15 @@ SettingsConfigDict, ) -# Environment variable used to locate the YAML configuration file. -# This variable is intentionally NOT covered by the AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_ -# prefix — it is read directly from the environment before the Settings class is -# instantiated, so the prefix-based machinery cannot apply. +# Read directly from the environment (not the AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_ +# prefix) before Settings is constructed. _CONFIG_FILE_ENV = "AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_CONFIG_FILE" _CONFIG_FILE_DEFAULT = "server-config.yaml" logger = logging.getLogger(__name__) -# Config keys that USED to exist and were removed or renamed in the headless -# refactor. Pydantic's settings sources silently drop unknown keys, so an -# operator upgrading a live deployment whose YAML still carries one of these -# would get a SILENT behaviour change (e.g. a customized timeout reverting to -# default, or a dashboard toggle becoming a no-op). We warn loudly instead of -# failing -- the server must still boot -- so the drop is never invisible. +# Removed/renamed config keys. Pydantic silently drops unknown keys, so warn +# loudly instead of failing -- the server must still boot. _REMOVED_CONFIG_KEYS: dict[str, str] = { "web_ui_enabled": ( "removed -- the server is headless-only and has no web UI toggle; " @@ -55,13 +50,10 @@ # --------------------------------------------------------------------------- # GUID validation helpers (Entra identities) # --------------------------------------------------------------------------- -# Anchored pattern for lowercase hex groups of 8-4-4-4-12. -# re.fullmatch() anchors the match to the full string, so braces, urn:uuid: -# prefixes, and trailing junk are all rejected without explicit anchors in the -# pattern. +# 8-4-4-4-12 lowercase hex, fullmatch()'d so braces/urn:uuid: prefixes/trailing +# junk are rejected without explicit anchors. _GUID_RE = re.compile(r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}") -# The all-zeros sentinel is explicitly rejected — a placeholder accidentally -# left in config should never authorize anyone. +# Placeholder sentinel; never a valid identity. _ALL_ZEROS_GUID = "00000000-0000-0000-0000-000000000000" @@ -73,30 +65,16 @@ def _validate_identity_map( ) -> dict[str, dict[str, str]] | None: """Shared validator for GUID-keyed identity maps (entra_identities, service_identities). - Enforces the same rules for both fields so they stay in sync: - - - ``None`` passes through (field is optional). - - An empty dict is rejected (fail-closed; omit or null-out to disable), - UNLESS ``allow_empty=True`` (entra_identities only), in which case an - empty dict is accepted and returned as ``{}`` so the server can boot on a - fresh /data volume and be populated at runtime via the /admin API. - - Every key must be a valid lowercase GUID in 8-4-4-4-12 form after - normalization (rejects braces, urn:uuid: prefixes, trailing junk). - - The all-zeros GUID is rejected (placeholder sentinel). - - Every value must carry a non-empty, non-whitespace ``id`` string. - - Keys are normalized to lowercase and returned as such. - - ``field_name`` is included verbatim in error messages so operators can tell - which field failed at startup. + - ``None`` passes through (optional field). + - Empty dict is rejected unless ``allow_empty=True`` (entra_identities only). + - Keys must be valid lowercase GUIDs (8-4-4-4-12); all-zeros is rejected. + - Every value must carry a non-empty, non-whitespace ``id``. + - Keys are normalized to lowercase. """ if v is None: return None if len(v) == 0: if allow_empty: - # entra_identities ONLY: an explicit empty map is permitted so the - # server boots on a fresh /data volume and is populated at runtime - # via PUT /admin/identities (bootstrap). service_identities does NOT - # pass allow_empty, so {} there remains a fail-closed startup error. return {} raise ValueError( f"{field_name} must contain at least one entry if specified; " @@ -143,21 +121,8 @@ def _build_identity_map_from( def _default_identity_store_path(filename: str) -> str: """Return a host-install-writable default path for an identity-map store file. - These paths used to default to "/data/identity/" -- an Azure - Files volume path baked in for the container deployment. On a plain - (non-container) host install nothing mounts /data, so the - seed-on-first-boot write in IdentityStore.seed() silently failed with - PermissionError: the server kept running fail-closed on the in-memory - map, but nothing was ever persisted to disk, and any key added later via - the /admin API would vanish on restart. - - Default to the invoking user's own writable data dir instead, using the - same ~/.local/share/ci-server/... layout already illustrated as the - host-install convention in YamlConfigSettingsSource's docstring above - (its blob_path / log_path example values). Container deployments are - unaffected: they set these paths explicitly via env/YAML (e.g. - amplifier-online.yaml sets entra_identities_store_path to the mounted - /data volume) -- this default only matters when nothing overrides it. + Defaults under the invoking user's own data dir rather than a container-only + /data mount. Container deployments override these paths explicitly. """ return str(Path.home() / ".local" / "share" / "ci-server" / "identity" / filename) @@ -165,23 +130,12 @@ def _default_identity_store_path(filename: str) -> str: class YamlConfigSettingsSource(PydanticBaseSettingsSource): """Load settings from a YAML configuration file. - The file path is resolved in this order: + Path resolution order: constructor ``yaml_file`` arg, then + ``AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_CONFIG_FILE``, then + ``server-config.yaml`` in cwd (skipped if absent). - 1. The ``yaml_file`` argument passed to the constructor (for tests / explicit use). - 2. The ``AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_CONFIG_FILE`` environment variable. - 3. ``server-config.yaml`` in the current working directory (silently skipped if - it does not exist). - - Keys in the YAML file correspond to the field names in :class:`Settings` without - the ``AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_`` prefix. Unknown keys are ignored. - - Example ``server-config.yaml``:: - - neo4j_url: neo4j://localhost:7687 - neo4j_password: "" - blob_path: /home/user/.local/share/ci-server/blobs - log_path: /home/user/.local/share/ci-server/logs/server.jsonl - Environment variables always take precedence over values in the YAML file. + Keys match :class:`Settings` field names (no prefix); unknown keys are + ignored. Environment variables take precedence over YAML values. """ def __init__( @@ -229,18 +183,13 @@ def __call__(self) -> dict[str, Any]: class Neo4jClientConfig(BaseModel): - """One Neo4j logical client (admin OR cypher_query). Same shape for both. - - Future RBAC fields (rbac_role, database, etc.) land in THIS object -- no - new top-level knobs later (doc 11 Structured config). - """ + """One Neo4j logical client (admin OR cypher_query). Same shape for both.""" url: str username: str = "neo4j" password: str = "" - # access_mode steers session routing/intent. "WRITE" for admin, "READ" for - # cypher_query. On a Community single instance over bolt:// this is a - # routing HINT, not server-side enforcement (doc 11 Honest caveat). + # "WRITE" for admin, "READ" for cypher_query. On a Community single + # instance over bolt:// this is a routing hint, not server-side enforcement. access_mode: Literal["READ", "WRITE"] = "WRITE" @property @@ -256,9 +205,8 @@ def auth(self) -> tuple[str, str] | None: class Neo4jConfig(BaseModel): """The structured `neo4j` block: two same-shaped clients. - When present, BOTH sub-clients are required (pydantic enforces this), so - the only fallback case the startup guard must detect is - `Settings.neo4j is None`. + Both sub-clients are required when present, so the startup guard's only + fallback case to detect is `Settings.neo4j is None`. """ admin: Neo4jClientConfig @@ -266,18 +214,10 @@ class Neo4jConfig(BaseModel): @model_validator(mode="after") def _validate_access_modes(self) -> "Neo4jConfig": - """Enforce role/access_mode correctness -- fail loud, not silent. - - `Neo4jClientConfig.access_mode` defaults to "WRITE", so a `cypher_query` - block that is a copy-paste of `admin` (or simply omits `access_mode`) - would silently behave as a WRITE-capable "read" client -- defeating the - entire point of the two-client split. Reject that at construction time: + """Enforce `admin.access_mode == "WRITE"` and `cypher_query.access_mode == "READ"`. - - `admin.access_mode` MUST be "WRITE" (the read/write client). - - `cypher_query.access_mode` MUST be "READ" (the read-intent client). - - Both violations are reported together so the operator sees one clear - message naming exactly which client has the wrong access_mode. + Fails loud (both violations reported together) rather than letting a + copy-pasted `cypher_query` block silently behave as WRITE-capable. """ errors: list[str] = [] if self.admin.access_mode != "WRITE": @@ -311,23 +251,12 @@ class Settings(BaseSettings): server_host: str = "0.0.0.0" server_port: int = 8000 - # Gunicorn worker timeouts (run() in main.py). Both were hardcoded until - # the incident below made that a problem: a durable-spool boot whose - # crash-recovery work is legitimately O(backlog size) (see - # crash_recovery_respawn_limit above) can take minutes on a large - # backlog, and gunicorn's own worker-timeout watchdog cannot distinguish - # "still doing legitimate startup work" from "hung" -- it just SIGKILLs - # the worker either way, which then gets restarted by systemd and repeats - # the same slow boot forever. Defaults (30s / 10s) are UNCHANGED from the - # previous hardcoded values, so this PR is a no-op unless an operator - # opts in to raise them for a deployment that expects a slow/large-backlog - # boot. + # Gunicorn worker timeouts (run() in main.py). Crash-recovery boot work is + # O(backlog size) and can take minutes; raise these for a deployment that + # expects a slow/large-backlog boot. # - # gunicorn_worker_timeout: seconds gunicorn allows a worker to go silent - # (no heartbeat) before killing it. See gunicorn's `timeout` setting. - # gunicorn_graceful_timeout: seconds gunicorn waits for a worker to finish - # handling in-flight work after SIGTERM before force-killing it. See - # gunicorn's `graceful_timeout` setting. + # gunicorn_worker_timeout: seconds a silent worker is allowed before kill. + # gunicorn_graceful_timeout: seconds to finish in-flight work after SIGTERM. gunicorn_worker_timeout: int = 30 gunicorn_graceful_timeout: int = 10 @@ -342,7 +271,7 @@ def _normalize_api_key(cls, v: str | None) -> str | None: """Normalize empty string to None so that api_key: '' in config disables auth.""" return None if v == "" else v - # Per-contributor API keys (NESTED form, design D4): the keystore is keyed by + # Per-contributor API keys (NESTED form): the keystore is keyed by # the SHA-256 hex digest of the raw token (64 lowercase hex chars), and each # value is a metadata dict carrying at least ``id`` (the contributor id). The # nested shape leaves room to add ``role`` / ``label`` later without a breaking @@ -362,23 +291,10 @@ def _validate_api_keys( ) -> dict[str, dict[str, str]] | None: """Fail-closed: raise unless every entry is ``<64-hex> -> {"id": }``. - Rejects (by raising ``ValueError``): - - an explicitly empty dict (omit or null-out to disable authentication); - - a key that is not exactly 64 lowercase-hex characters after normalization - (whitespace characters are rejected because they are not valid hex digits); - - a value whose ``id`` is missing, empty, or whitespace-only. - - Non-dict values are already rejected by pydantic's ``dict[str, dict[str, str]]`` - coercion before this validator runs (``mode="after"``), so no extra - ``isinstance`` check is needed here. - - Digest keys are normalized to lowercase before validation and returned as - lowercase so an UPPERCASE digest in a config file maps correctly to the - lowercase hexdigest produced by ``hashlib.sha256(...).hexdigest()``. - - NOTE: Duplicate digest keys in YAML/dict collapse to last-wins at the YAML - parse level, before this validator sees the data. Detection is not possible - here. + Rejects an empty dict, a key that isn't 64 lowercase-hex chars after + normalization, or a value with a missing/empty ``id``. Digest keys are + lowercased so an uppercase digest still maps to + ``hashlib.sha256(...).hexdigest()``. """ if v is None: return None @@ -432,35 +348,20 @@ def build_keystore(self) -> dict[str, str]: # ------------------------------------------------------------------------- # Entra authentication (auth_mode=entra) # ------------------------------------------------------------------------- - # auth_mode selects which resolver is active: "static" = today's sha256 - # keystore; "entra" = JWT validation via Entra / JWKS. Exactly one mode - # is active at a time — no "both". Choosing "entra" without the required - # supporting fields is a hard startup error (AC7 / §8b). + # "static" = sha256 keystore; "entra" = JWT validation via Entra/JWKS. + # Exactly one mode is active at a time. auth_mode: Literal["static", "entra"] = "static" - # allow_unauthenticated: the SOLE fail-open trigger in the whole server. - # - # An empty keystore / identity map alone NO LONGER fails open: the server - # BOOTS fail-CLOSED with zero credentials (a supported bootstrap state, - # announced by a loud startup WARNING) and every request 401/403s until the - # store is populated at runtime via the /admin API. - # - # The ONLY way to make the server pass EVERY request through unauthenticated - # is to set this flag to True AND leave credentials unconfigured - # (auth_enabled=False). In that case the middleware fails open and - # create_asgi_app() emits a loud "WIDE OPEN" warning at startup. - # - # This flag exists ONLY for the test harness and local dev environments - # where auth is intentionally disabled. Never set it in production. - # (In auth_mode=entra it has no effect: EntraResolver.auth_enabled is always - # True, so the fail-open branch can never fire regardless of this flag.) + # The sole fail-open trigger in the server. An empty keystore/identity map + # boots fail-closed (401/403 until populated via /admin). Setting this True + # with no credentials configured makes every request pass unauthenticated + # ("WIDE OPEN" warning at startup) -- test/dev only, never production. No + # effect in auth_mode=entra (EntraResolver.auth_enabled is always True). allow_unauthenticated: bool = False - # azure_client_id / azure_tenant_id: the App Registration coordinates. - # Both are required when auth_mode="entra". Empty / whitespace-only - # strings are normalized to None so that a template placeholder in a YAML - # file (e.g. azure_client_id: "") behaves identically to omitting the field - # and triggers a clear startup error rather than a silent wrong-value lookup. + # App Registration coordinates; both required when auth_mode="entra". + # Empty/whitespace strings normalize to None so a blank YAML placeholder + # triggers a clear startup error instead of a silent wrong-value lookup. azure_client_id: str | None = None azure_tenant_id: str | None = None @@ -474,21 +375,11 @@ def _normalize_azure_field(cls, v: Any) -> str | None: return None return v - # entra_identities: the oid→contributor map — exact parity with api_keys. - # - # Shape: { "": {"id": ""} } (value = {id} only) - # - # Key = the user's Azure AD object ID (oid), stored verbatim (public id); - # not hashed — hashing buys nothing for a public id and hurts auditability. - # Value = {"id": ""} matching the api_keys payload — same - # contributor string space, same write-once provenance semantics. - # - # Many oids → one contributor works automatically: each oid is its own key - # with the same "id" value (e.g. two AD identities for the same person). - # - # NOTE: oid is a persistent personal identifier. Do NOT commit real oid - # values to product repos — use env/secret injection or a git-ignored map - # (see §3 PII note in the auth plan). + # oid -> contributor map: { "": {"id": ""} }. + # oid is stored verbatim (not hashed -- it's already a public id). Many + # oids may map to one contributor. + # NOTE: oid is a persistent personal identifier -- do not commit real + # values to product repos; use env/secret injection or a git-ignored map. entra_identities: dict[str, dict[str, str]] | None = None @field_validator("entra_identities", mode="after") @@ -496,34 +387,15 @@ def _normalize_azure_field(cls, v: Any) -> str | None: def _validate_entra_identities( cls, v: dict[str, dict[str, str]] | None ) -> dict[str, dict[str, str]] | None: - """Fail-closed: raise unless every entry is `` -> {"id": }``. - - Delegates to the shared ``_validate_identity_map()`` helper which enforces - GUID key validation, the all-zeros sentinel rejection, non-empty ``id`` - requirement, and key lowercasing. See that function's docstring for the - full rule set. - - This validator runs in ``mode="after"``, so pydantic has already coerced - the field as ``dict[str, dict[str, str]]`` before this function is called. - Non-dict values and non-string ``id`` values are caught by pydantic before - reaching this code. - """ + """Fail-closed: raise unless every entry is `` -> {"id": }``.""" return _validate_identity_map(v, "entra_identities", allow_empty=True) @model_validator(mode="after") def _validate_entra_config(self) -> "Settings": - """Cross-field startup validator for auth_mode='entra' (AC7). - - When auth_mode is 'entra' ALL of the following must be present and - non-None after normalization: - - azure_client_id - - azure_tenant_id - - entra_identities is NOT required: an empty/omitted identity map is a - supported bootstrap state (populate at runtime via /admin/identities). + """When auth_mode='entra', require azure_client_id and azure_tenant_id. - A single ValueError names every missing field so the operator sees one - clear startup message rather than cryptic downstream failures. + entra_identities is not required (empty/omitted is a supported + bootstrap state). Names every missing field in one ValueError. """ if self.auth_mode == "entra": errors: list[str] = [] @@ -539,11 +411,6 @@ def _validate_entra_config(self) -> "Settings": "set AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_AZURE_TENANT_ID " "or azure_tenant_id in the config file" ) - # NOTE: entra_identities is intentionally NOT required here. - # An empty/omitted map is a SUPPORTED bootstrap state: the server - # boots and operators onboard the first oid at runtime via the - # IdentityAdmin-gated /admin/identities API. main.create_asgi_app() - # logs a loud warning when the effective map is empty at startup. if errors: raise ValueError( "Entra auth misconfiguration (startup refused): " @@ -554,32 +421,17 @@ def _validate_entra_config(self) -> "Settings": def build_identity_map(self) -> dict[str, str]: """Return ``{oid_lower -> contributor_id}`` for all configured Entra identities. - Mirrors ``build_keystore()`` — returns a plain ``{key: contributor_id}`` - dict that the EntraResolver can use for O(1) lookup after extracting the + Mirrors ``build_keystore()`` for O(1) lookup after extracting the ``oid`` claim from a validated JWT. - - Keys are lowercased as a belt-and-suspenders guarantee: the field - validator already normalizes them, but the resolver also lowercases the - JWT ``oid`` claim before lookup, so both sides use the same casing. """ return _build_identity_map_from(self.entra_identities) # ------------------------------------------------------------------------- # M2 non-interactive auth: service / app-token identity path # ------------------------------------------------------------------------- - # service_identities: the OID → contributor map for service principals / - # managed identities. Same shape as entra_identities; lives in config - # only (no durable store — service identities don't need runtime mutation). - # - # Shape: { "": {"id": ""} } - # - # Validation rules are identical to entra_identities (both delegate to the - # shared _validate_identity_map() helper) — GUID keys, non-empty id, no - # all-zeros sentinel. - # - # This field is OPTIONAL. The service identity path never participates in - # the _validate_entra_config cross-field check, so auth_mode=entra boots - # with only client_id / tenant_id / entra_identities. + # OID -> contributor map for service principals / managed identities. + # Same shape and validation as entra_identities; config-only (no durable + # store). Optional; doesn't participate in _validate_entra_config. service_identities: dict[str, dict[str, str]] | None = None @field_validator("service_identities", mode="after") @@ -587,31 +439,18 @@ def build_identity_map(self) -> dict[str, str]: def _validate_service_identities( cls, v: dict[str, dict[str, str]] | None ) -> dict[str, dict[str, str]] | None: - """Fail-closed: same GUID-map rules as entra_identities (shared helper). - - Delegates to ``_validate_identity_map()``. See that function's docstring - for the full rule set. - """ + """Fail-closed: same GUID-map rules as entra_identities (shared helper).""" return _validate_identity_map(v, "service_identities") def build_service_identity_map(self) -> dict[str, str]: - """Return ``{oid_lower -> contributor_id}`` for all configured service identities. - - Mirrors ``build_identity_map()`` — returns a plain ``{key: contributor_id}`` - dict for O(1) lookup after extracting the ``oid`` claim from an app token. - - Returns ``{}`` when ``service_identities`` is ``None`` or empty. - """ + """Return ``{oid_lower -> contributor_id}`` for all configured service identities.""" return _build_identity_map_from(self.service_identities) # ------------------------------------------------------------------------- # Admin API key (static mode only — gates /admin/* map-mutation endpoints) # ------------------------------------------------------------------------- - # admin_api_key is a separate credential from the data-auth api_keys. - # It is set via the YAML config file (same CONFIG_FILE that carries api_keys) - # and/or the env var below (env overrides YAML — standard pydantic-settings - # priority). Empty string is normalised to None so that admin_api_key: "" - # in a YAML template behaves identically to omitting the field. + # Separate credential from the data-auth api_keys. Empty string normalizes + # to None so admin_api_key: "" behaves like omitting the field. admin_api_key: str | None = None @field_validator("admin_api_key", mode="before") @@ -620,21 +459,10 @@ def _normalize_admin_api_key(cls, v: object) -> str | None: """Normalize empty string to None (mirrors _normalize_api_key).""" return None if v == "" else v # type: ignore[return-value] - # admin_api_key_sha256 is the RECOMMENDED way to configure the admin key: - # store the SHA-256 hex digest of the admin token at rest, never the raw - # token — mirroring how the data-auth ``api_keys`` map stores digests, not - # tokens (see docs/managing-api-keys.md). A leak of the config file then - # yields only a one-way digest, not a usable admin credential. - # - # The legacy raw ``admin_api_key`` above still works for back-compat (it is - # hashed at load time, exactly like the legacy singular ``api_key``), but is - # DEPRECATED because it stores the secret in plaintext at rest. When both - # are set, ``admin_api_key_sha256`` wins and the raw field is ignored - # (surfaced as a startup warning in create_asgi_app). - # - # Set via YAML or the env var - # ``AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_ADMIN_API_KEY_SHA256``. Empty - # string normalises to None (mirrors admin_api_key). + # Recommended way to configure the admin key: store the SHA-256 hex digest + # at rest, never the raw token. The legacy raw ``admin_api_key`` still + # works (hashed at load time) but is deprecated; when both are set, this + # wins. Empty string normalizes to None. admin_api_key_sha256: str | None = None @field_validator("admin_api_key_sha256", mode="before") @@ -691,16 +519,9 @@ def resolve_admin_api_key_digest(self) -> str | None: # ------------------------------------------------------------------------- # Entra admin role (entra mode only — gates /admin/* map-mutation endpoints) # ------------------------------------------------------------------------- - # entra_admin_role is the Entra App Role name whose presence in a token's - # `roles` claim grants access to /admin/* endpoints. The role is created - # in the App Registration (approles-patch.json). - # - # Empty string ("") means the admin API is DISABLED in entra mode (callers - # receive 503). The default "IdentityAdmin" matches the App Registration - # role defined for the pilot. Override via YAML or env var to rename. - # - # NOTE: the check ONLY reads the `roles` claim — NEVER `groups`. A value - # in the `groups` claim must NOT grant admin access (TB-09 / design §6). + # Entra App Role name whose presence in a token's `roles` claim grants + # access to /admin/* endpoints. Empty string disables the admin API in + # entra mode (503). Checks ONLY the `roles` claim — never `groups`. entra_admin_role: str = "IdentityAdmin" @field_validator("entra_admin_role", mode="before") @@ -711,18 +532,9 @@ def _normalize_entra_admin_role(cls, v: object) -> str: return "" return str(v) - # M2 service role names - # - # service_data_role: the Entra App Role name whose presence in an app token's - # ``roles`` claim grants the standard Contributor-level data access. This - # mirrors what a delegated user gets via entra_identities, but for service - # principals. Empty string ('') disables the service data path entirely. - # - # reader_role: the Entra App Role name granting read-only access. Empty string - # disables read-only app-token gating. Default 'Reader' matches the App - # Registration role defined for the M2 service path. - # - # Both fields normalize None → '' (same pattern as entra_admin_role). + # Entra App Roles gating service/app-token access: service_data_role for + # standard Contributor-level data access, reader_role for read-only. + # Empty string disables the respective path. service_data_role: str = "Contributor" reader_role: str = "Reader" @@ -737,15 +549,7 @@ def _normalize_service_role_fields(cls, v: object) -> str: # ------------------------------------------------------------------------- # Durable identity-map store paths # ------------------------------------------------------------------------- - # These paths control where the two JSON identity-map files live. Both are - # env/YAML overridable to allow non-default layouts in development, - # custom deployments, or containers -- e.g. amplifier-online.yaml sets - # entra_identities_store_path explicitly to the mounted Azure Files - # volume (/data/identity/entra-identities.json). - # - # The DEFAULT (see _default_identity_store_path()) is a host-writable - # per-user path, not /data/... -- see that helper's docstring for why. - # + # Where the two JSON identity-map files live; both env/YAML overridable. # api_keys_store_path: SHA-256 digest → contributor map (static mode) # entra_identities_store_path: OID → contributor map (entra mode) api_keys_store_path: str = _default_identity_store_path("api-keys.json") @@ -761,18 +565,12 @@ def _normalize_service_role_fields(cls, v: object) -> str: neo4j_password: str = "password" neo4j_browser_url: str = "http://localhost:7474" - # Structured two-client config (doc 11). OPTIONAL for backward-compat: when - # absent, BOTH clients fall back to the legacy flat neo4j_* fields above. - # The real amplifier-online.yaml MUST set this explicitly (see - # neo4j_require_explicit_clients + the startup guard). + # Structured two-client config. Optional for backward-compat: when absent, + # both clients fall back to the legacy flat neo4j_* fields above. neo4j: Neo4jConfig | None = None - # Deployed-profile signal (gap #12). When True, the startup guard REFUSES to - # boot on the legacy fallback (i.e. neo4j is None) -- the deployed system must - # declare admin + cypher_query explicitly, even pointing at the same instance. - # Default False so existing deployments / server-config.yaml keep booting on - # the legacy fallback during the transition. - # Env: AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_NEO4J_REQUIRE_EXPLICIT_CLIENTS=true + # When True, the startup guard refuses to boot on the legacy fallback + # (neo4j is None) -- admin + cypher_query must be declared explicitly. neo4j_require_explicit_clients: bool = False def resolve_neo4j_admin(self) -> Neo4jClientConfig: @@ -797,6 +595,32 @@ def resolve_neo4j_query(self) -> Neo4jClientConfig: access_mode="READ", ) + # Upper bound on concurrent bolt connections for a driver shared across many + # logical callers (the lifespan admin driver, the registry's per-session + # driver). Well under the server's default bolt thread-pool size so a + # driver leak can no longer starve it. + neo4j_max_connection_pool_size: int = 50 + + # Recycles a pooled connection after this many seconds, so a long-idle + # connection cannot accumulate indefinitely on the server side. + neo4j_max_connection_lifetime: float = 3600.0 + + @field_validator("neo4j_max_connection_pool_size") + @classmethod + def _validate_neo4j_max_connection_pool_size(cls, v: int) -> int: + """Fail loud on a non-positive pool size.""" + if v <= 0: + raise ValueError(f"neo4j_max_connection_pool_size must be > 0, got {v}") + return v + + @field_validator("neo4j_max_connection_lifetime") + @classmethod + def _validate_neo4j_max_connection_lifetime(cls, v: float) -> float: + """Fail loud on a non-positive lifetime (must be finite so idle connections recycle).""" + if v <= 0: + raise ValueError(f"neo4j_max_connection_lifetime must be > 0, got {v}") + return v + # ------------------------------------------------------------------------- # Storage paths # ------------------------------------------------------------------------- @@ -806,47 +630,40 @@ def resolve_neo4j_query(self) -> Neo4jClientConfig: # ------------------------------------------------------------------------- # Durable ingest queue # ------------------------------------------------------------------------- - # Conservative working defaults pending tuning (design Open Question 4). write_concurrency: int = 8 # global cap on concurrent Neo4j-write flushes max_delivery_attempts: int = 5 # flush retries for one batch before dead-letter - # Sub-transaction chunk bounds for _flush_body (issue #278). - # A chunk closes when EITHER bound trips first: cardinality or payload size. + # Sub-transaction chunk bounds for _flush_body. A chunk closes when EITHER + # bound trips first: cardinality or payload size. neo4j_flush_chunk_rows: int = ( 100 # max rows per sub-transaction (cardinality bound) ) neo4j_flush_chunk_bytes: int = ( 4_194_304 # max serialized bytes per sub-tx (4 MiB payload bound) ) - neo4j_lock_timeout: float = ( - 30.0 # per-transaction server-side timeout in seconds (Layer B) - ) - # A conservative default matching max_transaction_retry_time=30s. Prevents - # a blocked flush from parking indefinitely when db.lock.acquisition.timeout=0 - # (Neo4j default) holds all write_semaphore permits and stalls the pipeline. - # Set to 0 to disable (no per-transaction timeout). - - # Crash-recovery respawn ceiling (incident: a 38 GB / 583-file durable - # spool made cold start respawn 94/94 drainers before the server could - # accept a single request; startup took ~4 minutes and RSS peaked at - # 43.9 GB, tripping the kernel OOM killer -- which systemd then restarted, - # repeating the same unbounded respawn and never letting the backlog - # shrink). This mirrors write_concurrency's role as a hard ceiling on a - # startup-time resource cost, but bounds the RESPAWN LOOP in - # lifespan() (main.py) rather than write-flush concurrency: at most this - # many sessions from the recovered backlog get a drainer respawned on - # THIS boot; the remainder are DEFERRED -- left completely untouched on - # disk (still durable, still recoverable on a later boot, or instantly - # via get_or_create() the moment a new event for that session arrives - # through POST /events). A deferred backlog is never silent: lifespan() - # logs a WARNING naming the exact respawned/deferred counts and this - # setting, and /status's spool block (pending_sessions, spool_bytes_total) - # makes the backlog observable continuously, not just at boot. + neo4j_lock_timeout: float = 30.0 # per-transaction server-side timeout in seconds + # Prevents a blocked flush from parking indefinitely when + # db.lock.acquisition.timeout=0 holds all write_semaphore permits and + # stalls the pipeline. 0 disables the per-transaction timeout. + + # Hard ceiling on drainers respawned from the recovered backlog on THIS + # boot (mirrors write_concurrency's role, but bounds the respawn loop in + # lifespan() rather than write-flush concurrency). The remainder are + # deferred, left untouched on disk (still durable; recoverable on a later + # boot, or instantly via get_or_create() on a new event). Never silent: + # lifespan() logs the respawned/deferred counts, and /status's spool + # block makes the backlog observable continuously. # - # None (the default) preserves TODAY'S BEHAVIOUR EXACTLY: unbounded, - # every recovered session is respawned on this boot, matching every - # existing deployment -- this PR is a no-op unless an operator opts in - # by setting a finite ceiling. - crash_recovery_respawn_limit: int | None = None + # None means unbounded; with an unbounded ceiling no sweep task starts + # (see crash_recovery_sweep_interval_seconds below), so a recovered drainer + # that runs dry with no terminal record + # (the common shape of a legacy/crashed backlog) never freed its slot + # and was never replaced: there was no real protection against the OOM + # this field exists to prevent. 8 is a pessimistic, single-line- + # overridable default: an operator who can safely respawn more may raise + # it; one who cannot is now protected OUT OF THE BOX. 0 remains a valid + # explicit opt-out (never respawn automatically at boot; the deferred + # tail then drains only via a new event or the periodic sweep). + crash_recovery_respawn_limit: int | None = 8 @field_validator("crash_recovery_respawn_limit") @classmethod @@ -859,25 +676,12 @@ def _validate_crash_recovery_respawn_limit(cls, v: int | None) -> int | None: ) return v - # Crash-recovery deferred-backlog SWEEP interval (seconds). Only relevant - # when crash_recovery_respawn_limit is FINITE. Without this, a finite cap - # would drain the head of the backlog on boot and leave the deferred tail - # untouched until either a restart or a NEW event for that exact session - # arrives -- so a backlog of already-COMPLETED sessions (the incident's - # shape) would never drain at all, and a cap of 0 would strand EVERYTHING - # permanently. This sweep periodically re-runs recover() and tops the - # drainer pool back up to the ceiling: because respawn is idempotent - # (get_or_create) and recover() drops sessions the moment they finish, the - # number of live recovered drainers stays <= the ceiling while the deferred - # tail advances in deterministic sorted order as head sessions drain. - # - # Default 300s applies ONLY when a finite ceiling is set; with the default - # crash_recovery_respawn_limit=None (unbounded) there is no deferred tail - # and NO sweep task is ever started -- so every existing deployment is - # completely unaffected. Set to 0 to DISABLE the sweep even under a finite - # ceiling (the deferred tail then drains only on restart or a new event -- - # an explicit, documented choice, not a silent surprise). - crash_recovery_sweep_interval_seconds: int = 300 + # Deferred-backlog sweep interval (seconds); only relevant when + # crash_recovery_respawn_limit is finite. Periodically re-runs recover() + # and tops the drainer pool back up to the ceiling so the deferred tail + # keeps advancing instead of stalling until a restart or new event. + # 0 disables the sweep (deferred tail then drains only on restart/new event). + crash_recovery_sweep_interval_seconds: int = 60 @field_validator("crash_recovery_sweep_interval_seconds") @classmethod @@ -890,6 +694,150 @@ def _validate_crash_recovery_sweep_interval(cls, v: int) -> int: ) return v + # A bad `.offset` (unparseable, negative, or past-EOF) below this many + # bytes is reset (re-drained from byte 0, bounded and idempotent) rather + # than deleted outright; at/above the threshold the `.log` is deleted too. + # 0 means "always delete". + reclaim_redrain_max_bytes: int = 64 * 1024 * 1024 + + # Reclaim pass's DELETE/RESET_OFFSET actions ship disabled by default. + # With this False, boot still classifies every key and logs the same + # audit line (action=dry_run), but nothing is unlinked. First-deploy + # sequence: boot dry-run, review boot_reclaim_histogram, then opt in. + reclaim_enabled: bool = False + + # The queue is a transient buffer, not an archive -- an open, + # actively-draining session's already-committed prefix is reclaimed + # continuously (not just at session:end). Ships True: the decision input + # is the committed offset, the single value the durability design already + # trusts, and delete_drained -- already shipped, always on -- deletes the + # entire file on exactly this same evidence. False is a + # config-change-plus-restart kill switch, not a live toggle. + queue_compact_enabled: bool = True + + # Bounds compaction frequency on a continuously-hot session: below this + # many committed bytes, the rewrite is skipped (the idle path closes the + # gap for free once the session goes idle). + queue_compact_min_prefix_bytes: int = 8 * 1024 * 1024 + + # Separate flag from reclaim_enabled: the log-less + stale-mtime predicate + # is a structural proof (not a heuristic), so gating this on + # reclaim_enabled would let dead-letters accumulate forever by default. + # Ships False: a dead-letter may be the only surviving copy of an + # un-recovered event -- never auto-delete it without an explicit opt-in. + dead_letter_expiry_enabled: bool = False + + # How long a log-less `.dead.jsonl` survives before being expired. + # Long enough that an operator who notices a dead-letter via + # `GET /queues/dead-letter` or `/status`'s dead count has time to purge or + # replay it; short enough that the file cannot accumulate indefinitely. + # <=0 disables expiry outright (an explicit opt-out, not a silent one). + dead_letter_retention_seconds: float = 30 * 86400.0 + + @field_validator("queue_compact_min_prefix_bytes") + @classmethod + def _validate_queue_compact_bytes(cls, v: int) -> int: + """Fail loud on a negative value; 0 is a valid explicit opt-out.""" + if v < 0: + raise ValueError( + f"queue_compact_min_prefix_bytes must be a non-negative integer, got {v}" + ) + return v + + @field_validator("dead_letter_retention_seconds") + @classmethod + def _validate_dead_letter_retention_seconds(cls, v: float) -> float: + """Fail loud on a negative retention; 0 (disabled) is valid.""" + if v < 0: + raise ValueError( + "dead_letter_retention_seconds must be a non-negative number " + f"(0 disables expiry), got {v}" + ) + return v + + # Hard ceiling per _boot_reconcile phase (heal/reclaim/expire/reconcile/ + # seed/topup). A phase that hangs (e.g. a blocking stat/read on a + # degraded mount) would otherwise leave boot phase stuck pre-ready + # forever, latching /status's spool/metrics at null. <=0 disables the + # per-phase timeout (unbounded wait, pre-existing behavior). + boot_phase_timeout_seconds: float = 300.0 + + @field_validator("boot_phase_timeout_seconds") + @classmethod + def _validate_boot_phase_timeout_seconds(cls, v: float) -> float: + """Fail loud on non-finite input; <=0 is the documented opt-out.""" + if not math.isfinite(v): + raise ValueError(f"boot_phase_timeout_seconds must be finite, got {v}") + return v + + # ------------------------------------------------------------------------- + # Writer lease + # ------------------------------------------------------------------------- + # Refuses boot against a live foreign lease (takes over a stale one); + # `detect` only observes+heartbeats without ever refusing; `off` disables it. + writer_lease_mode: Literal["off", "detect", "enforce"] = "enforce" + # Renew + re-read interval; fixed at 5s for "conflict visible within one + # heartbeat" well inside a typical revision-overlap window. + writer_lease_heartbeat_seconds: float = 5.0 + # Staleness window = heartbeat_seconds * this multiplier. Must survive two + # consecutive missed ticks without a false "stale" verdict. + writer_lease_staleness_multiplier: float = 3.0 + # Post-write settle delay before the acquire's confirming re-read, to + # exceed write -> other-reader visibility latency on the shared mount. + writer_lease_confirm_delay_seconds: float = 1.0 + # Hard bound on the entire acquire/renew (reads, write, confirm sleep) so + # a hung mount can never block `lifespan` before its `yield` forever. + writer_lease_acquire_timeout_seconds: float = 5.0 + # One-boot operator escape hatch: force-acquire over a fresh foreign + # lease. Logs a warning every boot while set and is surfaced on /status. + writer_lease_force_acquire: bool = False + + @field_validator("writer_lease_heartbeat_seconds") + @classmethod + def _validate_writer_lease_heartbeat_seconds(cls, v: float) -> float: + """Fail loud on a non-positive heartbeat interval.""" + if v <= 0: + raise ValueError(f"writer_lease_heartbeat_seconds must be > 0, got {v}") + return v + + @field_validator("writer_lease_staleness_multiplier") + @classmethod + def _validate_writer_lease_staleness_multiplier(cls, v: float) -> float: + """Fail loud below 2.0x (one missed tick would yield a false stale verdict).""" + if v < 2.0: + raise ValueError( + "writer_lease_staleness_multiplier must be >= 2.0 (below 2, " + f"one missed heartbeat tick yields a false stale verdict), got {v}" + ) + return v + + @field_validator("writer_lease_confirm_delay_seconds") + @classmethod + def _validate_writer_lease_confirm_delay_seconds(cls, v: float) -> float: + """Fail loud on a negative confirm delay.""" + if v < 0: + raise ValueError( + f"writer_lease_confirm_delay_seconds must be >= 0, got {v}" + ) + return v + + @model_validator(mode="after") + def _validate_writer_lease_timeout_exceeds_confirm_delay(self) -> "Settings": + """Cross-field guard: timeout <= confirm delay makes every acquire + time out, silently disarming the detector -- must fail at config load.""" + if ( + self.writer_lease_acquire_timeout_seconds + <= self.writer_lease_confirm_delay_seconds + ): + raise ValueError( + "writer_lease_acquire_timeout_seconds must exceed " + "writer_lease_confirm_delay_seconds (otherwise every acquire " + f"times out): got acquire_timeout=" + f"{self.writer_lease_acquire_timeout_seconds}, confirm_delay=" + f"{self.writer_lease_confirm_delay_seconds}" + ) + return self + # ------------------------------------------------------------------------- # Logging # ------------------------------------------------------------------------- diff --git a/context_intelligence_server/graph_store.py b/context_intelligence_server/graph_store.py index aa491398..70e6008f 100644 --- a/context_intelligence_server/graph_store.py +++ b/context_intelligence_server/graph_store.py @@ -58,7 +58,31 @@ class GraphStore(Protocol): @property def workspace(self) -> str: - """Workspace this store is bound to (set at construction, read-only).""" + """Workspace this store is bound to. + + Settable: HookStateService binds it immediately after construction + (services.py). Reads never return None -- an unset workspace resolves + to "default". + """ + ... + + @workspace.setter + def workspace(self, value: str) -> None: ... + + @property + def created_by(self) -> str | None: + """Authenticated contributor id for write-once provenance (None when unset).""" + ... + + @created_by.setter + def created_by(self, value: str | None) -> None: ... + + def discard_buffer(self) -> None: + """Drop all buffered writes without persisting them (guarantee #13). + + MUST NOT perform I/O and MUST NOT raise. In-memory implementations with + no backing store may treat this as a no-op. + """ ... async def upsert_node(self, node_id: str, data: dict[str, Any]) -> None: diff --git a/context_intelligence_server/handlers/data_layer_2/session.py b/context_intelligence_server/handlers/data_layer_2/session.py index a39118f1..e2f5771d 100644 --- a/context_intelligence_server/handlers/data_layer_2/session.py +++ b/context_intelligence_server/handlers/data_layer_2/session.py @@ -71,15 +71,9 @@ class SessionLabelStateMachine: def classify( self, current_type: str | None, event: str, has_parent: bool ) -> LabelTransition: - # NOTE on "StubSession" removal below: StubSession is a plain observability - # marker (added by services.ensure_session_node when a node is created from - # a reference — delegation, fork/start parent — before its own lifecycle - # events arrive). It is not part of the RootSession/SubSession/ForkedSession - # terminal lattice. Every branch below that assigns a REAL terminal label - # (including IncompleteSession, a confirmed-if-incomplete terminal) also - # clears StubSession, so genuine late enrichment removes the marker. Removing - # a label that isn't present is a silent no-op (see GraphState.set_labels / - # Neo4jGraphStore.set_labels), so it is always safe to include in `remove`. + # StubSession is a plain observability marker, not part of the + # terminal lattice; every branch assigning a real terminal label + # also clears it (removing an absent label is a no-op). if event == "start": if current_type in ("ForkedSession", "SubSession"): return LabelTransition() @@ -118,17 +112,8 @@ def classify( if event == "end": if current_type is not None: return LabelTransition() - # Bare session: session:start/fork was permanently lost. Rather than - # fabricating a real terminal (Sub/Root), mark it explicitly so it - # stays outside the clean terminal space and surfaces as a health signal. - # IncompleteSession IS a confirmed (if incomplete) terminal outcome — the - # node is no longer an orphaned stub, it is a diagnosed data-loss case — - # so StubSession is cleared here too. - # - # NOTE: if a real start/fork ever arrives AFTER this end (out-of-order, - # vanishingly rare), _handle_start/_handle_fork will classify normally - # and add the real terminal. IncompleteSession may then coexist as an - # audit trail — that is acceptable; no special stripping is needed. + # Bare session: start/fork was permanently lost. Mark + # IncompleteSession rather than fabricating a real terminal. return LabelTransition( add=["IncompleteSession", "SST_EVENT"], remove=["StubSession"] ) @@ -184,9 +169,8 @@ async def _handle_start( _warn_if_dual_terminal(labels, session_id) current_type = _current_type(labels) - # Always enrich started_at and session identity. "Session" MUST be in - # labels so neo4j_store routes this to MERGE (n:Session {...}), the same - # bucket as ensure_session_node. + # "Session" MUST be in labels so neo4j_store routes this to MERGE, + # the same bucket as ensure_session_node. await self.services.graph.upsert_node( session_id, { @@ -203,7 +187,6 @@ async def _handle_start( session_id, data_layer_1_node_id, {"type": "SOURCED_FROM"} ) - # Label decision is owned by the state machine. transition = self._label_machine.classify( current_type, "start", bool(parent_id) ) @@ -214,9 +197,8 @@ async def _handle_start( add_labels=transition.add, ) - # Edge rule: a session becoming a SubSession under a parent gets a - # HAS_SUBSESSION edge. Covers both Root->Sub and bare->Sub. Root (no - # parent) and the terminal no-ops create no edge. + # A session becoming a SubSession under a parent gets a + # HAS_SUBSESSION edge; Root (no parent) creates no edge. if "SubSession" in transition.add and parent_id: await self.services.ensure_session_node(parent_id, {}) await self.services.graph.upsert_edge( @@ -249,8 +231,8 @@ async def _handle_fork( _warn_if_dual_terminal(labels, session_id) current_type = _current_type(labels) - # Always enrich. "Session" MUST be in labels for the same MERGE-bucket - # reason as _handle_start. + # "Session" MUST be in labels for the same MERGE-bucket reason as + # _handle_start. await self.services.graph.upsert_node( session_id, { @@ -276,11 +258,8 @@ async def _handle_fork( add_labels=transition.add, ) - # Edge rule: a session becoming a ForkedSession under a parent gets a - # FORKED edge. If this was a reclassification of an already-typed - # Root/Sub node, drop the stale parent edge FIRST. Keyed on current_type - # (not transition.remove) to mirror the legacy code exactly. The terminal - # ForkedSession no-op has empty transition.add, so creates no edge. + # A session becoming a ForkedSession gets a FORKED edge; a + # reclassified Root/Sub node drops its stale parent edge first. if "ForkedSession" in transition.add and parent_id: if current_type in ("RootSession", "SubSession"): self.services.graph.remove_edge(parent_id, session_id) @@ -301,32 +280,22 @@ 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. + # 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 @@ -338,10 +307,8 @@ async def _handle_end( session_id, data_layer_1_node_id, {"type": "SOURCED_FROM"} ) - # Stub recovery: if session:start was permanently missed (bare Session), - # mark the session as IncompleteSession instead of fabricating a real - # terminal label (Sub/Root). This keeps the guess out of the clean - # Root/Sub/Forked terminal space and surfaces a health signal. + # Stub recovery: if start was permanently missed, mark + # IncompleteSession rather than fabricating a real terminal label. transition = self._label_machine.classify( _current_type(labels), "end", bool(parent_id) ) @@ -358,10 +325,6 @@ 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: @@ -378,7 +341,6 @@ async def _create_mount_plan( {"type": "HAS_PART", "sst_semantic": "CONTAINS"}, ) # SOURCED_FROM bridge: MountPlan -> data_layer_1 session:fork event - # The session:fork event contains data.raw (blob) with the full mount plan config await self.services.graph.upsert_edge( mount_plan_id, data_layer_1_fork_node_id, diff --git a/context_intelligence_server/idempotency.py b/context_intelligence_server/idempotency.py index 7efacaae..cca666ed 100644 --- a/context_intelligence_server/idempotency.py +++ b/context_intelligence_server/idempotency.py @@ -1,9 +1,16 @@ -"""In-memory request deduplication for event ingestion.""" +"""In-memory request deduplication for event ingestion. + +Keys are recorded only after a successful durable append; a key in the +cache means the event is on disk. +""" from __future__ import annotations +import asyncio import time from collections import OrderedDict +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager class EventIdempotencyCache: @@ -16,17 +23,39 @@ def __init__( self._max_entries = max_entries self._seen: OrderedDict[str, float] = OrderedDict() - def check_and_store(self, key: str, now: float | None = None) -> bool: - """Return True if *key* is new and store it, False if it is a duplicate.""" + def seen(self, key: str, now: float | None = None) -> bool: + """Return True if *key* was already recorded as durably accepted. + + This is the READ half of the former ``check_and_store``, byte for + byte: it purges expired entries first and refreshes LRU recency on a + hit, so eviction behaviour is unchanged. It NEVER records a key -- + recording is ``store``'s job, and happens only after a successful + durable append. + + NOTE THE POLARITY: this returns True for a DUPLICATE, whereas + ``check_and_store`` returned True for a NEW key. The name states the + meaning; the caller's guard is ``if seen(...): return duplicate``. + """ current_time = time.time() if now is None else now self._purge(current_time) if key in self._seen: self._seen.move_to_end(key) - return False + return True + return False + + def store(self, key: str, now: float | None = None) -> None: + """Record *key* as durably accepted. + + MUST be called ONLY after the event's durable append has returned + successfully. Burning a key before durability is a silent-loss + bug: the client's retry is answered "duplicate" for an event that is + nowhere on disk, and no recovery path can resurrect it -- the bytes + never reached the log. + """ + current_time = time.time() if now is None else now self._seen[key] = current_time self._seen.move_to_end(key) self._trim() - return True def clear(self) -> None: """Remove all remembered keys.""" @@ -42,3 +71,27 @@ def _purge(self, now: float) -> None: def _trim(self) -> None: while len(self._seen) > self._max_entries: self._seen.popitem(last=False) + + +class KeyedAsyncLocks: + """Per-key asyncio.Lock registry; a key's lock is dropped once idle.""" + + def __init__(self) -> None: + self._locks: dict[str, asyncio.Lock] = {} + self._waiters: dict[str, int] = {} + + @asynccontextmanager + async def acquire(self, key: str) -> AsyncIterator[None]: + lock = self._locks.setdefault(key, asyncio.Lock()) + self._waiters[key] = self._waiters.get(key, 0) + 1 + try: + async with lock: + yield + finally: + self._waiters[key] -= 1 + if self._waiters[key] <= 0: + self._waiters.pop(key, None) + # Only drop the lock if nothing else grabbed a reference to + # it in the meantime (it is not currently locked/awaited). + if not lock.locked(): + self._locks.pop(key, None) diff --git a/context_intelligence_server/main.py b/context_intelligence_server/main.py index 0bf0709f..912c3bb4 100644 --- a/context_intelligence_server/main.py +++ b/context_intelligence_server/main.py @@ -9,14 +9,16 @@ import sys import time from collections.abc import AsyncGenerator -from contextlib import asynccontextmanager, suppress +from contextlib import asynccontextmanager, nullcontext, suppress +from dataclasses import dataclass from datetime import datetime +from functools import partial from pathlib import Path from typing import Any from fastapi import Depends, FastAPI, HTTPException, Request from fastapi.responses import JSONResponse, Response -from neo4j import READ_ACCESS, WRITE_ACCESS, AsyncGraphDatabase +from neo4j import READ_ACCESS, WRITE_ACCESS from context_intelligence_server import __version__ from context_intelligence_server.auth import ( @@ -32,7 +34,10 @@ ) from context_intelligence_server.blob_store import AsyncDiskBlobStore from context_intelligence_server.config import Neo4jClientConfig, Settings, get_settings -from context_intelligence_server.idempotency import EventIdempotencyCache +from context_intelligence_server.idempotency import ( + EventIdempotencyCache, + KeyedAsyncLocks, +) from context_intelligence_server.identity_store import IdentityStore from context_intelligence_server.logging_config import setup_logging from context_intelligence_server.models import ( @@ -41,6 +46,7 @@ EventResponse, ) from context_intelligence_server.neo4j_store import ( + build_bounded_neo4j_driver, count_untagged_nodes, ensure_neo4j_schema, ) @@ -48,7 +54,12 @@ from context_intelligence_server.routers.admin import router as admin_router from context_intelligence_server.routers.queues import router as queues_router from context_intelligence_server.routers.version import router as version_router -from context_intelligence_server.status import build_status_response +from context_intelligence_server.status import boot_state, build_status_response +from context_intelligence_server.writer_lease import ( + WriterLeaseConflict, + shutdown_lease_io, + writer_lease, +) _settings = get_settings() @@ -61,23 +72,23 @@ def _neo4j_access_const(mode: str) -> str: def build_neo4j_driver(config: Neo4jClientConfig) -> Any: - """Construct an AsyncGraphDatabase driver from a resolved Neo4j client config. + """Construct the pool-bounded admin AsyncGraphDatabase driver. Shared by ``lifespan()`` (the admin driver, on every server boot) and ``doctor.run_doctor()`` (the CLI), so the two entry points can never - construct the connection differently. + construct the connection differently. Delegates the actual driver + construction to ``build_bounded_neo4j_driver`` so the pool-bounding kwargs + have one source of truth, shared with ``SessionRegistry``'s driver. """ - return AsyncGraphDatabase.driver(config.url, auth=config.auth) + return build_bounded_neo4j_driver( + config, + max_connection_pool_size=_settings.neo4j_max_connection_pool_size, + max_connection_lifetime=_settings.neo4j_max_connection_lifetime, + ) -# --------------------------------------------------------------------------- -# Module-level live identity-map stores (T3) -# -# 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. -# --------------------------------------------------------------------------- +# Module-level live identity-map stores. Exactly one is non-None at a time -- +# the other is reset to None so accessors return an unambiguous result. _api_key_store: IdentityStore | None = None _entra_identity_store: IdentityStore | None = None @@ -102,106 +113,498 @@ def get_entra_identity_store() -> IdentityStore | None: return _entra_identity_store +# Last-resort workspace sentinel when no head line resolves a workspace. +# Dispatching under it still isolates the bad line and drains the rest. +_RECOVERY_FALLBACK_WORKSPACE = "unknown-recovered" + + +def _head_is_resumable(raw: bytes) -> bool: + """Total predicate: does ``raw`` parse to a dict with a workspace? + + Shared by ``_recover_one_session`` and ``QueueManager.classify_session`` + (injected as a pure callable so the queue never learns the event schema). + Never raises -- valid-but-non-dict JSON must not escape as an AttributeError. + """ + try: + obj = json.loads(raw) + except (ValueError, TypeError): + return False + if not isinstance(obj, dict): + return False + try: + return bool(obj.get("workspace", "")) + except (AttributeError, TypeError): + return False + + +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, get_or_create: Any, + first_log_line: bytes | None = None, + *, + recovered: bool = True, ) -> bool: """Parse the first queued line for *sid* and respawn a drainer when valid. - Extracted from the lifespan startup recovery loop so tests can exercise the - real parsing/dispatch logic rather than reimplementing it inline. + Falls back to ``first_log_line`` (byte-0) when ``first_line`` doesn't + resolve a workspace, then to the ``_RECOVERY_FALLBACK_WORKSPACE`` + sentinel -- so an unparseable head never blocks recovery of the data + behind it. Returns True if a drainer was (re)spawned, False if skipped. + """ + parsed = _parse_workspace_and_creator(first_line) + if parsed is not None: + workspace, created_by = parsed + get_or_create(sid, workspace, created_by=created_by, recovered=recovered) + return True + + if first_log_line is not None: + byte0_parsed = _parse_workspace_and_creator(first_log_line) + if byte0_parsed is not None: + workspace, created_by = byte0_parsed + logger.warning("recovery_fallback_workspace session=%s source=byte0", sid) + get_or_create(sid, workspace, created_by=created_by, recovered=recovered) + return True + # Last resort: dispatch under the sentinel anyway -- the drainer + # dead-letters the unparseable head and drains everything behind it. + logger.warning("recovery_fallback_workspace session=%s source=sentinel", sid) + get_or_create( + sid, + _RECOVERY_FALLBACK_WORKSPACE, + created_by=None, + recovered=recovered, + ) + return True + + logger.warning( + "recovery_skipped session=%s: torn or empty workspace in first line", + sid, + ) + return False - 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. - 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 - production or a spy in tests. +@dataclass +class TopupResult: + """Result of one ``_crash_recovery_topup`` pass. - Returns: - True – drainer was (re)spawned via *get_or_create*. - False – session skipped (empty/torn workspace, or malformed JSON line). + ``dispatched``: sessions dispatched this pass (idempotent, so an upper + bound on newly-spawned drainers). ``recovered``: total size of this + pass's ``recover()`` report, before ceiling slicing. ``deferred``: + ``recovered`` minus how many were processed (0 when unbounded). """ - 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: - logger.warning( - "recovery_skipped session=%s: torn or empty workspace in first line", - sid, - ) - return False - get_or_create(sid, workspace, created_by=created_by) - return True + dispatched: int + recovered: int + deferred: int -async def _crash_recovery_topup(respawn_limit: int | None) -> int: + +async def _crash_recovery_topup(respawn_limit: int | None) -> TopupResult: """One bounded crash-recovery pass: respawn drainers for up to ``respawn_limit`` recovered sessions (all of them when ``None``). - This is the shared body of the boot-time recovery and the periodic sweep. - It is SAFE to call repeatedly on a live server because respawn is - idempotent -- ``registry.get_or_create`` returns the existing worker for a - session that already has a live drainer (no duplicate drainer, no reset). - And because ``recover()`` reports only sessions that still have undrained - data, a session drops out the moment it finishes, so the number of live - RECOVERED drainers stays <= ``respawn_limit`` while the deferred tail - advances in deterministic sorted order as head sessions drain. - - Returns the number of sessions DISPATCHED to get_or_create on this pass -- - an upper bound on newly-spawned drainers, since get_or_create is a no-op - for a session that already has a live drainer (see NOTE in the loop). + Shared by the boot-time recovery and the periodic sweep. Safe to call + repeatedly -- ``get_or_create`` is idempotent, and ``recover()`` only + reports sessions with undrained data, so live recovered drainers stay + bounded as the deferred tail advances. Falls back to the session's + byte-0 line when the head doesn't resolve a workspace. """ recovered = await registry.queue_manager.recover() to_process = recovered if respawn_limit is None else recovered[:respawn_limit] - respawned = 0 + deferred_count = ( + 0 if respawn_limit is None else max(0, len(recovered) - respawn_limit) + ) + dispatched = 0 for sid in to_process: - batch = await registry.queue_manager.read_batch(sid, max_items=1) + try: + # 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 pass. + 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 compaction + # 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 - # is idempotent). So this count is "sessions dispatched this pass", an - # upper bound on newly-spawned drainers -- fine for an INFO log. - if _recover_one_session(sid, batch.lines[0], registry.get_or_create): - respawned += 1 - return respawned + # An upper bound on newly-spawned drainers: get_or_create is + # idempotent, so "dispatched" may include already-live workers. + dispatched_ok = _recover_one_session( + sid, batch.lines[0], registry.get_or_create + ) + if not dispatched_ok: + first_log_line = await registry.queue_manager.read_first_line(sid) + dispatched_ok = _recover_one_session( + sid, + batch.lines[0], + registry.get_or_create, + first_log_line=first_log_line, + ) + if dispatched_ok: + dispatched += 1 + if deferred_count: + # WARNING (not INFO): a deferred backlog must never be silently + # undiscoverable. 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 pass, " + "%d session(s) deferred to a later pass (untouched on disk, " + "still fully recoverable). Raise crash_recovery_respawn_limit " + "to respawn more per pass.", + respawn_limit, + dispatched, + len(to_process), + deferred_count, + ) + return TopupResult( + dispatched=dispatched, recovered=len(recovered), deferred=deferred_count + ) + + +async def _ensure_schema_ready() -> None: + """Attempt Neo4j schema init once; a no-op once already ready. + + Sets ``app.state.schema_ready`` on success. A connectivity failure + (Neo4j unreachable) is logged and swallowed here -- schema stays + not-ready, retried later (boot's sweep phase) instead of crash-looping + the server. Raises ``RuntimeError`` only for a genuine data conflict + (graph reachable but un-migrated) -- the one refusal this still + preserves, now recorded via ``boot_state.fail()`` by the caller instead + of aborting ASGI startup. + """ + if getattr(app.state, "schema_ready", False): + return + try: + await ensure_neo4j_schema(app.state.neo4j_driver, fail_on_data_conflict=True) + except RuntimeError: + raise # genuine data conflict: fatal, let the caller record it + except Exception as exc: # noqa: BLE001 - Neo4j unreachable, not fatal + logger.warning( + "schema_init_unreachable: Neo4j not reachable, will retry: %s", exc + ) + return + # Catches nodes lacking the :Node label (the other un-migrated shape the + # constraint above can't see). A probe failure is logged at DEBUG, not + # treated as confirmed-bad -- the flush path's self-heal still covers it. + try: + untagged = await count_untagged_nodes(app.state.neo4j_driver) + except Exception as exc: # noqa: BLE001 - connectivity probe, not a confirmed bad state + logger.debug( + "schema_init: untagged-node probe skipped (graph unreachable?): %s", exc + ) + untagged = 0 + if untagged: + raise RuntimeError( + f"Neo4j graph has {untagged} node(s) lacking the :Node label " + "(un-migrated). Cold start refuses to boot to avoid duplicating " + "them on write. Run: context-intelligence-server doctor --fix" + ) + app.state.schema_ready = True + logger.info("lifespan_startup: Neo4j schema initialized") async def _crash_recovery_sweep_loop(interval: int, respawn_limit: int) -> None: - """Periodically top the recovered-drainer pool back up to the ceiling so a - finite ``crash_recovery_respawn_limit`` cannot permanently strand the - deferred backlog (the tail only advances as head sessions finish draining). - - Started by ``lifespan`` ONLY when a finite ceiling is configured and the - interval is > 0; with the default unbounded ceiling there is no deferred - tail and this loop never runs. A single failed tick must never kill the - loop, so the body is guarded (CancelledError propagates for clean - shutdown; everything else is logged and the loop continues). + """Periodically top the recovered-drainer pool back up to the ceiling so + a finite ``crash_recovery_respawn_limit`` cannot permanently strand the + deferred backlog. Also the retry mechanism for a schema that wasn't + ready at boot: each tick retries schema init first, and only tops up + (and marks boot ready) once it succeeds. A single failed tick is logged + and retried; ``CancelledError`` propagates for clean shutdown. """ while True: try: await asyncio.sleep(interval) - respawned = await _crash_recovery_topup(respawn_limit) - if respawned: - logger.info( - "crash_recovery_sweep: dispatched %d recovered session(s) " - "(ceiling=%d) -- draining deferred backlog", - respawned, - respawn_limit, - ) + if not app.state.schema_ready: + try: + await _ensure_schema_ready() + except Exception as exc: # noqa: BLE001 - retried next tick + logger.warning( + "crash_recovery_sweep: schema still not ready, will retry: %s", + exc, + ) + else: + if app.state.schema_ready: + logger.info( + "crash_recovery_sweep: schema now ready -- " + "draining deferred backlog" + ) + # Drainer start stays gated on schema; disk-only work below + # (expire) does not and must run every tick regardless. + if app.state.schema_ready: + result = await _crash_recovery_topup(respawn_limit) + if result.dispatched: + logger.info( + "crash_recovery_sweep: dispatched %d recovered session(s) " + "(ceiling=%d) -- draining deferred backlog", + result.dispatched, + respawn_limit, + ) + if boot_state.phase == "awaiting_schema": + boot_state.finish() + # Live counters here (unlike boot) must record_purged expired + # records, or the accepted/written residual latches at +n. + expire_result = await registry.queue_manager.expire_dead_letters( + time.time(), + _settings.dead_letter_retention_seconds, + _settings.dead_letter_expiry_enabled, + ) + if expire_result["expired_records"]: + registry.record_purged(expire_result["expired_records"]) except asyncio.CancelledError: raise - except Exception as exc: # noqa: BLE001 - a sweep tick must never kill the loop - logger.warning("crash_recovery_sweep: tick failed, will retry: %s", exc) + except Exception as exc: + logger.warning( + "crash_recovery_sweep: tick failed, will retry: %s", + exc, + exc_info=True, + ) + + +async def _writer_lease_boot() -> None: + """The only writer-lease call on the synchronous boot path. + + Only ``WriterLeaseConflict`` (the intended enforce-mode refusal) is + allowed to escape. The heartbeat task starts whenever + ``writer_lease_mode != "off"``, even after a failed acquire -- that is + the re-arm mechanism a transient fault depends on. + """ + try: + await writer_lease.acquire(_settings, lambda: registry.queues_dir_path) + except WriterLeaseConflict: + raise # the ONE intended abort (enforce mode only) + except Exception as exc: # noqa: BLE001 - a detector must never crash-loop the server + logger.error( + "writer_lease: boot acquire failed -- the writer-lease detector " + "is NOT ARMED for this process: %r", + exc, + ) + writer_lease.mark_unarmed(repr(exc), _settings.writer_lease_mode) + if _settings.writer_lease_mode != "off": + app.state.lease_task = asyncio.create_task(writer_lease.heartbeat_loop()) + + +async def _boot_reclaim() -> None: + """Log-then-delete every un-resumable/already-drained key, resume with + fallback for a recoverable-but-unparseable head, and reset a bounded + bad-offset key. Skips any key with a live registry worker. Classify + always runs; the actual unlink/reset only runs when ``reclaim_enabled``. + """ + qm = registry.queue_manager + # Module-level `_settings`, not a fresh get_settings() -- keeps this in + # sync with test monkeypatches bound to the same object. + settings = _settings + boot_state.reclaim_enabled = settings.reclaim_enabled + # Iterate the QueueManager's own directory, not settings.queues_path -- + # the two can differ (tests do this routinely). + keys = sorted(p.stem for p in qm.queues_dir.glob("*.log")) + reclaimed = 0 + reclaimed_bytes = 0 + kept = 0 + failed = 0 + for key in keys: + if registry.has_worker(key): + kept += 1 + continue + try: + c = await qm.classify_session(key, _head_is_resumable) + except (OSError, ValueError) as exc: # pragma: no cover -- defence in depth + logger.error("boot_reclaim_classify_failed session=%s error=%s", key, exc) + failed += 1 + continue + if c.verdict.value == "resumable": + kept += 1 + if c.reason == "fallback_workspace": + if c.fallback_source == "byte0": + boot_state.fallback_workspace_byte0 += 1 + elif c.fallback_source == "sentinel": + boot_state.fallback_workspace_sentinel += 1 + continue + if c.verdict.value == "unreadable": + failed += 1 + logger.warning("boot_reclaim_kept reason=%s session=%s", c.reason, key) + continue + if c.verdict.value == "keep": + kept += 1 + logger.warning("boot_reclaim_kept reason=%s session=%s", c.reason, key) + continue + # verdict in (unresumable, drained, reset_offset): actionable. + # drained is the same evidence delete_drained already acts on + # unconditionally at session finalize -- safe to auto-reclaim + # regardless of reclaim_enabled. unresumable/reset_offset stay + # gated: they can act on a log whose offset was merely unreadable. + if c.verdict.value != "drained" and not settings.reclaim_enabled: + logger.warning( + "boot_reclaimed reason=%s path=%s session=%s bytes=%d action=dry_run", + c.reason, + Path(settings.queues_path) / f"{key}.log", + key, + c.size, + ) + kept += 1 + continue + ok = await qm.reclaim(c, partial(registry.has_worker, key)) + if ok: + reclaimed += 1 + reclaimed_bytes += c.size + else: + kept += 1 + # reclaim_orphans itself gates on reclaim_enabled and reports only real + # unlinks (0 when disabled) -- no further gating needed here. + orphan_result = await qm.reclaim_orphans(_start_time, settings.reclaim_enabled) + reclaimed += orphan_result["reclaimed"] + reclaimed_bytes += orphan_result["reclaimed_bytes"] + failed += orphan_result["failed"] + boot_state.reclaimed += reclaimed + boot_state.reclaimed_bytes += reclaimed_bytes + boot_state.kept += kept + boot_state.failed += failed + logger.info( + "boot_reclaim_summary reclaimed=%d bytes=%d kept=%d failed=%d mode=%s", + reclaimed, + reclaimed_bytes, + kept, + failed, + "live" if settings.reclaim_enabled else "dry_run", + ) + + +async def _phase_run(coro: Any) -> Any: + """Run one boot-phase awaited call under ``boot_phase_timeout_seconds``. + + A hung mount-touching call (a blocking stat/read on a degraded mount) + would otherwise leave ``boot_state.phase`` stuck pre-ready forever, + latching /status's spool/metrics at null. On timeout this raises + ``TimeoutError`` -- left to propagate to ``_boot_reconcile``'s own + except-Exception handler, which records it via ``boot_state.fail()`` + exactly like any other phase failure. ``<= 0`` disables the timeout + (unbounded wait, pre-existing behavior). + """ + timeout = _settings.boot_phase_timeout_seconds + if timeout is not None and timeout > 0: + return await asyncio.wait_for(coro, timeout=timeout) + return await coro + + +async def _boot_reconcile() -> None: + """The backgrounded, exception-safe boot-recovery body. + + Runs schema -> heal -> reclaim -> expire -> reconcile -> seed -> topup -> + sweep, then phase=ready. Spawned from ``lifespan``, not awaited, so the + server serves its first request while this still runs. Any exception is + recorded via ``boot_state.fail()``; the server keeps serving. + + ``schema`` is the one phase gating something real: drainer start + (``topup``) requires ``app.state.schema_ready``, since the Session/:Node + uniqueness constraints must be active before any flush() MERGE. The + disk-only phases (heal/reclaim/expire/reconcile/seed) need no schema and + run regardless. A schema left not-ready (Neo4j unreachable) is retried + by the periodic sweep, not by blocking this pass. + """ + boot_state.begin() + # Defensive: a direct call (bypassing lifespan's own init) must not + # AttributeError on the topup-phase read below. + app.state.schema_ready = getattr(app.state, "schema_ready", False) + try: + boot_state.phase = "schema" + await _ensure_schema_ready() + + boot_state.phase = "heal" + _heal_result = await _phase_run(registry.queue_manager.heal_torn_tails()) + logger.info("lifespan_startup: heal_torn_tails result=%s", _heal_result) + + boot_state.phase = "reclaim" + await _phase_run(_boot_reclaim()) + + boot_state.phase = "expire" + # Runs before recovery_seed_counts, so expired lines are simply never + # counted into accepted_seed -- record_purged must not be called here. + await _phase_run( + registry.queue_manager.expire_dead_letters( + time.time(), + _settings.dead_letter_retention_seconds, + _settings.dead_letter_expiry_enabled, + ) + ) + + boot_state.phase = "reconcile" + await _phase_run(registry.queue_manager.recovery_reconcile_dead()) + + boot_state.phase = "seed" + ( + _accepted_seed, + _written_seed, + ) = await _phase_run(registry.queue_manager.recovery_seed_counts()) + registry.seed_counters(_accepted_seed, _written_seed) + + boot_state.phase = "topup" + respawn_limit = _settings.crash_recovery_respawn_limit + if app.state.schema_ready: + result = await _phase_run(_crash_recovery_topup(respawn_limit)) + boot_state.resumed += result.dispatched + boot_state.deferred += result.deferred + logger.info( + "lifespan_startup: crash recovery respawned %d/%d drainers", + result.dispatched, + result.recovered, + ) + else: + logger.warning( + "crash_recovery_topup_skipped phase=topup reason=schema_not_ready " + "-- drainers deferred until Neo4j schema init succeeds " + "(retried by the periodic sweep)" + ) + + boot_state.phase = "sweep" + _sweep_interval = _settings.crash_recovery_sweep_interval_seconds + if respawn_limit is not None and _sweep_interval > 0: + app.state.sweep_task = asyncio.create_task( + _crash_recovery_sweep_loop(_sweep_interval, respawn_limit) + ) + logger.info( + "crash_recovery_sweep: enabled (interval=%ds, ceiling=%d) -- " + "deferred backlog will drain progressively, not just on restart", + _sweep_interval, + respawn_limit, + ) + if app.state.schema_ready: + # Finish unconditionally here (after starting the loop), so a + # forever-running sweep never leaves phase stuck at "sweep". + boot_state.finish() + else: + # Schema never came up this pass -- stay visibly NOT ready + # (never silently reported as "ready") until the sweep loop + # above retries schema + topup and marks it ready itself. + boot_state.phase = "awaiting_schema" + except asyncio.CancelledError: + raise + except Exception as exc: # a boot hook must never crash-loop the server + failed_step = boot_state.phase + boot_state.fail(failed_step, exc) + logger.exception( + "boot_reconcile_failed phase=failed failed_step=%s", failed_step + ) @asynccontextmanager @@ -216,170 +619,75 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: _query.url, _query.access_mode, ) - # Admin (read/write): schema init + all mutation paths. Keep the existing - # app.state.neo4j_driver NAME so nothing that reads it silently breaks. - # build_neo4j_driver() is the SAME helper doctor.run_doctor() uses, so the - # server and the doctor CLI can never construct this connection differently. + # Admin (read/write): schema init + all mutation paths. Shares + # build_neo4j_driver() with doctor.run_doctor() so the two never diverge. app.state.neo4j_driver = build_neo4j_driver(_admin) - # Cypher-query (read-intent): /cypher + dashboard reads. - app.state.neo4j_query_driver = AsyncGraphDatabase.driver( - _query.url, auth=_query.auth + # Cypher-query (read-intent): /cypher + dashboard reads. Bounded through the + # same helper as the admin driver so every process-wide pool shares one cap. + app.state.neo4j_query_driver = build_bounded_neo4j_driver( + _query, + max_connection_pool_size=_settings.neo4j_max_connection_pool_size, + max_connection_lifetime=_settings.neo4j_max_connection_lifetime, ) # Stash the resolved query access_mode so /cypher opens READ sessions without # re-resolving settings on every request. app.state.neo4j_query_access_mode = _query.access_mode - # Initialize schema (indexes + uniqueness constraints) BEFORE the server starts - # accepting requests. This ensures the Session uniqueness constraint is active - # before any concurrent flush() transactions execute MERGE, which prevents the - # duplicate-Session-node race condition observed under concurrent upload load. - logger.info( - "lifespan_startup: initializing Neo4j schema (indexes + uniqueness constraints)" - ) - # Cold start FAILS LOUD on schema/data corruption that requires - # `doctor --fix` -- an un-migrated graph (duplicate legacy nodes OR - # nodes lacking the universal :Node label). Nothing has been written yet - # at cold start, so refusing to boot loses no data: this is the safest - # possible moment to surface an impossible state as an un-missable - # signal rather than a log line someone greps for later. Contrast with - # the flush path (Neo4jGraphStore._ensure_schema), which must keep - # self-healing and never raise (Salil's blocker -- raising there would - # dead-letter real in-flight activity records). fail_on_data_conflict=True - # here mirrors run_repair's contract: a :Node constraint data conflict - # raises a RuntimeError naming `doctor --fix` instead of being logged - # and swallowed. - await ensure_neo4j_schema(app.state.neo4j_driver, fail_on_data_conflict=True) - logger.info("lifespan_startup: Neo4j schema initialized") - # Fail-loud migration-health guard: duplicate nodes are already caught - # above by the :Node constraint (fail_on_data_conflict=True); this catches - # the OTHER un-migrated shape the constraint can't see on its own -- - # nodes that simply lack the :Node label altogether, which violate no - # constraint and so raise nothing by themselves. O(1) via the counts - # store (see count_untagged_nodes) -- this must never regress into the - # AllNodesScan stall PR #67 removed from the write path. - # - # A connectivity/probe failure here is NOT the same as "confirmed - # un-migrated" -- it means graph state could not be determined, not that - # it was determined to be bad -- so it is logged at DEBUG and swallowed - # rather than treated as a corruption finding; the flush path's - # self-heal still covers a genuinely dirty graph once it becomes - # reachable. - try: - untagged = await count_untagged_nodes(app.state.neo4j_driver) - except Exception as exc: # noqa: BLE001 - connectivity probe, not a confirmed bad state - _LOG_MSG = "startup migration-health probe skipped (graph unreachable?): %s" - logger.debug(_LOG_MSG, exc) - untagged = 0 - if untagged: - raise RuntimeError( - f"Neo4j graph has {untagged} node(s) lacking the :Node label " - "(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 - # 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. - # - # Conservation-counter recovery runs FIRST, and its two steps are - # order-load-bearing: reconcile MUST precede seed. recovery_reconcile_dead - # advances committed offsets past already-dead pending lines so the - # dead-letter counts are settled; only then does recovery_seed_counts read - # disk to reconstruct the accepted/written baseline. Seeding before - # reconciling would leave a residual==1 false DEGRADED. Both run before the - # respawn loop so the respawned drainers start from a conserved baseline. - await registry.queue_manager.recovery_reconcile_dead() - _accepted_seed, _written_seed = await registry.queue_manager.recovery_seed_counts() - registry.seed_counters(_accepted_seed, _written_seed) - recovered = await registry.queue_manager.recover() - # Bound how many drainers this boot respawns (incident: an unbounded - # backlog respawned 94/94 drainers before the server could serve a - # single request, driving a ~4 minute boot and 43.9 GB RSS that tripped - # the OOM killer -- which then never let the backlog shrink because - # every restart repeated the same unbounded respawn). None (the default) - # preserves today's behaviour exactly: every recovered session is - # processed on this boot, unbounded. `recovered` is already sorted - # (QueueManager.recover()), so which sessions are processed this boot - # vs. deferred is deterministic across restarts of the same backlog. - # - # Deferred sessions are NOT touched in any way here -- no read, no - # write, no drainer -- so they remain exactly as durable and - # recoverable as they were before this boot: a later boot's recover() - # 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. - respawn_limit = _settings.crash_recovery_respawn_limit - if respawn_limit is not None and len(recovered) > respawn_limit: - to_process = recovered[:respawn_limit] - deferred_count = len(recovered) - respawn_limit - else: - to_process = recovered - deferred_count = 0 - respawned = 0 - for sid in to_process: - batch = await registry.queue_manager.read_batch(sid, max_items=1) - if not batch.lines: - continue - 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. - logger.warning( - "lifespan_startup: crash-recovery respawn cap reached " - "(crash_recovery_respawn_limit=%d): %d/%d respawned this boot, " - "%d session(s) deferred to a later boot (untouched on disk, " - "still fully recoverable). Raise crash_recovery_respawn_limit " - "to respawn more per boot.", - respawn_limit, - respawned, - len(to_process), - deferred_count, - ) - logger.info( - "lifespan_startup: crash recovery respawned %d/%d drainers", - respawned, - len(recovered), - ) - # Periodic deferred-backlog sweep: only meaningful under a FINITE ceiling - # (a deferred tail can exist). With the default unbounded ceiling - # (respawn_limit is None) there is no deferred tail, so NO background task - # is started -- existing deployments are completely unaffected. When a - # finite ceiling IS set, this drains the deferred tail over time instead of - # stranding it until a restart or a new event (see _crash_recovery_sweep_loop - # and config.crash_recovery_sweep_interval_seconds). - _sweep_task: asyncio.Task[None] | None = None - _sweep_interval = _settings.crash_recovery_sweep_interval_seconds - if respawn_limit is not None and _sweep_interval > 0: - _sweep_task = asyncio.create_task( - _crash_recovery_sweep_loop(_sweep_interval, respawn_limit) - ) - logger.info( - "crash_recovery_sweep: enabled (interval=%ds, ceiling=%d) -- " - "deferred backlog will drain progressively, not just on restart", - _sweep_interval, - respawn_limit, - ) + # Schema init (indexes + the Session/:Node uniqueness constraints) no + # longer runs synchronously here -- a Neo4j connectivity failure must + # never raise out of lifespan (ASGI startup abort -> crash-loop). It now + # runs as _boot_reconcile's first phase ("schema"), backgrounded like + # the rest of boot recovery; app.state.schema_ready gates drainer start + # (_crash_recovery_topup) until it succeeds. + app.state.schema_ready = False + # Acquires the lease before any boot-recovery pass mutates the shared + # directory; awaited synchronously so an enforce-mode refusal can't be + # silently downgraded by _boot_reconcile's own exception-safety. + await _writer_lease_boot() + # Every share-reading recovery pass moves off the critical path to first + # request: spawned as a background task, not awaited, so /status and + # /version answer while it still runs. + boot_state.begin() + app.state.boot_task = asyncio.create_task(_boot_reconcile()) try: yield finally: + # Ordering is load-bearing: sweep stops before reconcile (it can + # re-enter its work), and every task stops before the drivers close. + _sweep_task = getattr(app.state, "sweep_task", None) if _sweep_task is not None: _sweep_task.cancel() with suppress(asyncio.CancelledError): await _sweep_task + _boot_task = getattr(app.state, "boot_task", None) + if _boot_task is not None: + _boot_task.cancel() + with suppress(asyncio.CancelledError): + await _boot_task + # Released last; its heartbeat is cancelled first so no in-flight + # tick can regain the gate mid-shutdown. release() never raises. + _lease_task = getattr(app.state, "lease_task", None) + if _lease_task is not None: + _lease_task.cancel() + with suppress(asyncio.CancelledError): + await _lease_task + await writer_lease.release() + shutdown_lease_io() logger.info("lifespan_shutdown: closing Neo4j drivers") await app.state.neo4j_driver.close() await app.state.neo4j_query_driver.close() + # The registry's shared per-session driver is independent of the two + # above (its own pool, built from settings.resolve_neo4j_admin() the + # first time a session is created) -- close it here too so no bolt + # connection outlives the process. + await registry.close_neo4j_driver() app = FastAPI( title="Context Intelligence Server", version=__version__, lifespan=lifespan, - # Headless server: no browser-facing UI. The OpenAPI contract + Swagger UI - # are the developer surface and are always registered; ReDoc is a redundant - # second doc UI and is intentionally left off (docs_url=None equivalent). + # Headless server: Swagger UI is the dev surface; ReDoc is a redundant + # second doc UI, intentionally left off. docs_url="/docs", redoc_url=None, openapi_url="/openapi.json", @@ -389,14 +697,16 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: app.include_router(queues_router) _start_time = time.time() registry = SessionRegistry() -# Expose the registry singleton on app.state so routers can read it via -# request.app.state.registry instead of importing the module-level name -# (avoids a circular import between main and the routers package). +# Expose the registry singleton on app.state so routers can read it without +# importing the module-level name (avoids a circular import). app.state.registry = registry idempotency_cache = EventIdempotencyCache() +# Serializes the seen()->append->store() sequence per idempotency_key so +# concurrent same-key requests cannot both durably append (see post_events). +_idempotency_locks = KeyedAsyncLocks() # 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__" @@ -407,13 +717,9 @@ 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 - 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 - drainer calls make_node_id() on an empty string. + """Raise HTTPException(400) if data['timestamp'] is missing, empty, or + not ISO-8601. This is the ingest boundary check -- reject malformed + payloads with a clear 400 instead of dead-lettering them later. """ value = data.get("timestamp") if value is None or not isinstance(value, str) or not value.strip(): @@ -431,15 +737,11 @@ 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. + """/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 - ``_EXEMPT_PATHS`` or ``_EXEMPT_PREFIXES``, because that would make the - admin API accessible without authentication. - - This is a defence-in-depth structural check: it is impossible to - accidentally ship an unauthenticated admin surface. + ``_EXEMPT_PATHS`` or ``_EXEMPT_PREFIXES`` -- a defence-in-depth check + against accidentally shipping an unauthenticated admin surface. """ import context_intelligence_server.auth as _auth_module @@ -463,14 +765,10 @@ 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 - structured neo4j.admin / neo4j.cypher_query clients explicitly. - - When settings.neo4j_require_explicit_clients is True, refuse to boot if the - server silently fell back to the legacy flat neo4j_* fields (settings.neo4j is - None). Back-compat fallback is allowed ONLY when the flag is False (dev / test / - transition). This makes a silent partial-config fallback impossible in the - deployed profile. + """The deployed profile must declare structured neo4j.admin / + neo4j.cypher_query clients explicitly. When + ``neo4j_require_explicit_clients`` is True, refuse to boot on a silent + fallback to legacy flat neo4j_* fields. """ if settings.neo4j_require_explicit_clients and settings.neo4j is None: raise RuntimeError( @@ -490,70 +788,43 @@ def create_asgi_app( ) -> BearerTokenMiddleware: """Return the ASGI app wrapped with auth middleware. - This is the single strategy-selection point. *settings* defaults to - the module-level ``_settings`` (the cached production config). Pass an - explicit :class:`~context_intelligence_server.config.Settings` instance - from tests to exercise specific configurations without touching the live - cached settings. - - *_jwks_client* is an injectable JWKS client used **only** when - ``auth_mode="entra"`` — intended for tests that need to construct an - :class:`~context_intelligence_server.auth.EntraResolver` without making - real network calls. Production deployments leave it as ``None``; the - 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 - 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 - ``settings.allow_unauthenticated=True`` opt-out combined with no - credentials configured, which additionally logs a "WIDE OPEN" warning. - - Raises: - RuntimeError: (TB-07) When any ``/admin`` path or prefix appears in an - auth-exempt set. The admin API surface must never be unguarded. + *settings* defaults to the module-level ``_settings``; tests pass an + explicit instance to exercise a config without touching the live cache. + *_jwks_client* injects a JWKS client for ``auth_mode="entra"`` tests only. + + An empty keystore/identity map is a supported bootstrap state: the + server boots fail-closed and every request 401/403s until populated via + the /admin API, unless ``allow_unauthenticated=True`` with no + credentials configured (wide-open, logged loudly). """ 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: runs before middleware construction so the + # failure is loud and immediate. _assert_admin_not_exempt() s = settings if settings is not None else _settings _assert_neo4j_clients_explicit(s) - # Reset both stores; the active mode sets exactly one of them below. - # app.state.* mirrors the module-level globals so the /admin router can - # access the live stores via request.app.state without importing from main - # (which would create a circular import). + # Reset both stores; the active mode sets exactly one below. app.state.* + # mirrors the globals so /admin can read them without importing main. _api_key_store = None _entra_identity_store = None 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 - # 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(). + # Store auth/admin config on app.state so dependencies can read it + # without importing main, and test-specific settings take effect. 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. 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). - # - # 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, - # plaintext at rest) is hashed by the resolver. Surface the deprecation and - # precedence as one-time startup warnings so operators can migrate. + # Admin-key digest for the middleware (static mode only): checked against + # the bearer token's sha256 before the resolver, so the admin key + # authenticates even though it isn't in the data keystore. admin_api_key_digest: str | None = s.resolve_admin_api_key_digest() if s.admin_api_key is not None and s.admin_api_key_sha256 is not None: logger.warning( @@ -575,9 +846,8 @@ def create_asgi_app( entra_store = IdentityStore(Path(s.entra_identities_store_path)) entra_store.load() if not entra_store.path.exists(): - # First boot: seed in-process map from config. Converts the flat - # {oid -> contributor_id} from build_identity_map() to the rich - # {oid -> {"id": contributor_id}} format that IdentityStore expects. + # First boot: seed from config, converting flat {oid: contributor_id} + # to the rich {oid: {"id": contributor_id}} format IdentityStore expects. config_map = s.build_identity_map() if config_map: rich_seed = {oid: {"id": cid} for oid, cid in config_map.items()} @@ -585,11 +855,8 @@ def create_asgi_app( _entra_identity_store = entra_store app.state.entra_identity_store = entra_store - # Bootstrap visibility: announce an EMPTY identity map loudly at startup. - # This is a SUPPORTED state, not an error — the server is up and serving. - # Delegated (human) tokens will 403 until an IdentityAdmin role-holder - # binds the first oid via PUT /admin/identities/{oid}. Without this line - # an empty map would be silent and look like a misconfiguration. + # Supported bootstrap state, not an error -- without this the empty + # map would be silent and look like a misconfiguration. if not entra_store.flat_dict: logger.warning( "entra identity map is EMPTY at startup (0 bound oids) — server " @@ -600,12 +867,9 @@ def create_asgi_app( s.entra_identities_store_path, ) - # B4: 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. + # Disjointness invariant: each oid belongs to exactly one identity + # source. Built here (not inline in EntraResolver) so the overlap can + # be checked before construction, failing loud at startup. _service_id_map = s.build_service_identity_map() _entra_oids = set(entra_store.flat_dict.keys()) _service_oids = set(_service_id_map.keys()) @@ -618,30 +882,28 @@ def create_asgi_app( f"the overlap before restarting." ) - # EntraResolver raises RuntimeError at construction if the JWKS - # prefetch fails (eager fail-closed guard from §8b / crusty gate). - # Pass entra_store.flat_dict (the LIVE dict) so the resolver sees - # any put()/delete() made by /admin immediately, no restart required. + # EntraResolver raises at construction if the JWKS prefetch fails + # (fail-closed). Pass the live flat_dict so /admin mutations are + # visible 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 + 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, # pre-built, disjointness verified + service_data_role=s.service_data_role, # role gate + reader_role=s.reader_role, # role gate + entra_admin_role=s.entra_admin_role, # role gate jwks_client=_jwks_client, ) - # Entra mode does not use admin_api_key_digest (admin via roles claim). + # Entra mode: admin is via roles claim, not admin_api_key_digest. admin_api_key_digest = None else: # Build and load the API-key store. key_store = IdentityStore(Path(s.api_keys_store_path)) key_store.load() if not key_store.path.exists(): - # First boot: seed from config. Converts the flat - # {sha256_hex -> contributor_id} from build_keystore() to the - # rich {sha256_hex -> {"id": contributor_id}} format. + # First boot: seed from config, converting flat {sha256: contributor_id} + # to the rich {sha256: {"id": contributor_id}} format. config_ks = s.build_keystore() if config_ks: rich_seed = {digest: {"id": cid} for digest, cid in config_ks.items()} @@ -649,9 +911,8 @@ def create_asgi_app( _api_key_store = key_store app.state.api_key_store = key_store - # Bootstrap visibility: announce an EMPTY keystore loudly at startup. - # This is a SUPPORTED state (fail-CLOSED, not fail-open) — the server - # is up and serving, but every request 401s until keys are onboarded. + # Supported bootstrap state (fail-closed, not fail-open): server is + # up but every request 401s until keys are onboarded. if not key_store.flat_dict: if s.resolve_admin_api_key_digest() is not None: logger.warning( @@ -678,10 +939,8 @@ def create_asgi_app( # put()/delete() made by /admin immediately, no restart required. resolver = StaticKeyResolver(key_store.flat_dict) - # 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). + # Fires only on the explicit allow_unauthenticated opt-out combined with + # no credentials configured; an empty store alone boots fail-closed instead. if s.allow_unauthenticated and not resolver.auth_enabled: logger.warning( "allow_unauthenticated=True AND no credentials configured — the " @@ -691,7 +950,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 +969,8 @@ def create_asgi_app( _admin_status, ) - # T6: 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). + # Store the admin-key digest on app.state so the /admin router can read + # it without importing main; None in entra mode, sha256 in static mode. app.state.admin_api_key_digest = admin_api_key_digest return BearerTokenMiddleware( @@ -727,40 +983,16 @@ 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. -# -# `_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 -# 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. +# Lazily constructed (PEP 562 __getattr__) so bare import / --help / --version +# don't trigger create_asgi_app()'s auth guard; get_asgi_app() builds-and-caches. _asgi_app: BearerTokenMiddleware | None = None def get_asgi_app() -> BearerTokenMiddleware: """Return the module-level ASGI app, constructing it on first call. - This is the single lazy-construction point. Internal code (``_App.load()`` - below) MUST call this function rather than referencing a bare ``asgi_app`` - global -- a bare name reference is a normal global-variable lookup and - would NOT go through ``__getattr__``, so it would raise ``NameError`` - once the unconditional module-level assignment is removed. + Internal code must call this rather than referencing a bare ``asgi_app`` + global -- that lookup would not go through ``__getattr__``. """ global _asgi_app if _asgi_app is None: @@ -780,13 +1012,8 @@ def __getattr__(name: str) -> Any: raise AttributeError(f"module {__name__!r} has no attribute {name!r}") -# --------------------------------------------------------------------------- -# M2 — 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 -# so tests and existing imports from main still work). -# --------------------------------------------------------------------------- +# require_write, require_read, _is_write_capable live in authz.py (avoids a +# circular import) and are re-exported here for existing imports from main. @app.get("/status") @@ -795,30 +1022,34 @@ 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. + # Surface the query (read-intent) driver's connectivity too, so a + # misconfigured cypher_query client shows up here, not on first /cypher. 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. - 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. - 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 - # unauthenticated — only config-level boolean flags are exposed here - # (no credential values, no key hashes, no token details). + # Gated on boot being OVER (ready or failed), not SUCCEEDED -- gating on + # `ready` alone would permanently null the spool alarm after any reconcile failure. + response["boot"] = boot_state.snapshot() + # Pure in-memory (never touches disk), so this is safe at every boot + # phase. Kept out of `spool`, which is a live cache dict returned by reference. + response["writer_lease"] = writer_lease.snapshot() + if boot_state.phase in ("ready", "failed"): + # /status is unauthenticated: only aggregate-only conservation + # metrics, no per-key table or dead-letter listing. + response["metrics"] = await registry.pipeline_metrics() + # Same contract: aggregate integers only, cheap (stat-only, + # short-TTL cached) even with a huge spool. + response["spool"] = await registry.queue_manager.spool_stats() + else: + # While booting, /status performs zero disk reads. metrics/spool stay + # present but null, so an absent key is never confused with a version skew. + response["metrics"] = None + response["spool"] = None + response["status_detail"] = {"reason": "booting"} + # Surface auth mode/admin capability so operators can confirm admin is + # enabled without tailing logs -- boolean flags only, no credentials. _auth_mode = getattr(request.app.state, "auth_mode", _settings.auth_mode) _admin_key_set = getattr( request.app.state, @@ -833,10 +1064,8 @@ async def get_status(request: Request) -> dict[str, Any]: "admin_api_enabled": ( _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. + # Surface role names (not secrets) so operators can confirm what's + # configured without exposing credential values. **( { "entra_admin_role": _entra_admin_role, @@ -868,7 +1097,7 @@ async def _check_driver_connected(app_instance: FastAPI, attr_name: str) -> bool try: await driver.verify_connectivity() return True - except Exception: + except Exception: # noqa: BLE001 -- status must never 500 return False @@ -889,34 +1118,39 @@ async def post_events( # Validate data.timestamp at the ingest boundary (fail loud, not silent dead-letter). # Real Amplifier clients always supply this field; 400 only hits malformed payloads. _validate_data_timestamp(request.data) - # Idempotency-cache check + replay stay BEFORE the durable append so a - # duplicate is rejected without persisting a second log line. - if request.idempotency_key and not replay: - is_new = idempotency_cache.check_and_store(request.idempotency_key) - if not is_new: + # Serialize seen->append->store per key so concurrent same-key requests + # cannot both append; store only after a successful append. + dedup_key = request.idempotency_key if not replay else None + lock_ctx = _idempotency_locks.acquire(dedup_key) if dedup_key else nullcontext() + async with lock_ctx: + if dedup_key and idempotency_cache.seen(dedup_key): logger.info( "event_duplicate_skipped: event=%s session_id=%s", request.event, session_id, ) 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). - 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) - # Re-parse the raw validated body bytes, stamp created_by (server-assigned, - # unconditional overwrite — kills any client-supplied spoofed value), then - # re-serialize compact JSON before persisting to the durable queue. - # IMPORTANT: re-parse raw bytes (not the pydantic model) so client extra - # fields are preserved. body() is cached by Starlette after the first read. - body = await http_request.body() - body_obj = json.loads(body) - body_obj["created_by"] = contributor_id # overwrite, never setdefault - body = json.dumps(body_obj, separators=(",", ":")).encode() - await registry.queue_manager.append(worker_key, body) - registry.record_accepted() # count the durably-accepted event - return EventResponse(status="queued", 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. + 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) + # Re-parse raw bytes (not the pydantic model) so client extra fields + # survive; stamp created_by server-side, overwriting any spoofed value. + body = await http_request.body() + body_obj = json.loads(body) + body_obj["created_by"] = contributor_id # overwrite, never setdefault + body = json.dumps(body_obj, separators=(",", ":")).encode() + await registry.queue_manager.append(worker_key, body) + # Bytes are on disk: the key may be burned now (a failed append + # simply never reaches this line -- the lock is still released, + # via the `async with`, WITHOUT storing). + if dedup_key: + idempotency_cache.store(dedup_key) + registry.record_accepted() # count the durably-accepted event + return EventResponse(status="queued", session_id=session_id or None) @app.get("/blobs/{session_id}", dependencies=[Depends(require_read)]) @@ -955,23 +1189,19 @@ async def post_cypher(body: CypherRequest, request: Request) -> Response: rows.append(dict(record)) serialized = json.dumps({"results": rows}, default=str) return Response(content=serialized, media_type="application/json") - except Exception as exc: # catch all Neo4j and serialization errors + except Exception as exc: # noqa: BLE001 -- catch all Neo4j and serialization errors raise HTTPException(status_code=500, detail=str(exc)) def main(argv: list[str] | None = None) -> None: """CLI entrypoint. - INVARIANT: no subcommand (or the explicit ``serve`` subcommand) starts the - ingestion server. This MUST hold because the systemd unit (and the - macOS launchd agent) invoke the bare console script - ``context-intelligence-server`` with NO arguments -- that call dispatches - to ``serve`` unchanged. + No subcommand (or the explicit ``serve``) starts the ingestion server -- + the systemd unit and macOS launchd agent invoke the bare console script + with no arguments, dispatching to ``serve``. - ``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``. + ``doctor [--fix]`` diagnoses (and repairs) Neo4j graph health; see + ``context_intelligence_server.doctor``. """ parser = argparse.ArgumentParser(prog="context-intelligence-server") subparsers = parser.add_subparsers(dest="command") @@ -995,9 +1225,7 @@ def main(argv: list[str] | None = None) -> None: return # Deferred import: doctor.py imports build_neo4j_driver back from this - # module, so importing it at module load time (rather than here, inside - # main()) would be a circular import at import time. By the time main() - # runs, this module has already finished executing top-to-bottom. + # module, so a top-level import here would be circular. from context_intelligence_server import doctor as _doctor sys.exit(asyncio.run(_doctor.run_doctor(fix=args.fix))) @@ -1048,9 +1276,8 @@ def run() -> None: """Start the server using gunicorn + uvicorn worker for graceful SIGTERM shutdown.""" from gunicorn.app.base import BaseApplication - # Read WEB_CONCURRENCY and fail loud if it would run != 1 worker. The same - # value is fed into gunicorn below so the guard and the live config are one - # source of truth (they can never diverge). + # Fail loud if WEB_CONCURRENCY would run != 1 worker; the same value + # feeds gunicorn below so the guard and config can never diverge. workers = _validate_single_worker() class _App(BaseApplication): @@ -1066,9 +1293,8 @@ def load_config(self) -> None: self.cfg.set(key, value) def load(self) -> Any: - # get_asgi_app() (not the bare `asgi_app` global) -- this is - # where lazy construction actually happens for a real serve, - # and where the auth guard still fires if unconfigured. + # get_asgi_app(), not the bare `asgi_app` global -- this is where + # lazy construction happens and the auth guard still fires. return get_asgi_app() _App().run() diff --git a/context_intelligence_server/neo4j_store.py b/context_intelligence_server/neo4j_store.py index 52ddc7d9..08b7cc3e 100644 --- a/context_intelligence_server/neo4j_store.py +++ b/context_intelligence_server/neo4j_store.py @@ -22,8 +22,32 @@ from neo4j import unit_of_work as _unit_of_work from neo4j.exceptions import DriverError, Neo4jError +from context_intelligence_server.config import Neo4jClientConfig + _LOG = logging.getLogger(__name__) + +def build_bounded_neo4j_driver( + config: Neo4jClientConfig, + *, + max_connection_pool_size: int, + max_connection_lifetime: float, +) -> Any: + """Construct an AsyncGraphDatabase driver with a bounded connection pool. + + Single source of truth for the pool-bounding kwargs applied to any driver + meant to be shared across many logical callers (the lifespan admin driver, + the registry's per-session driver). Both ``main.build_neo4j_driver`` and + ``SessionRegistry`` call this so the two never diverge. + """ + return AsyncGraphDatabase.driver( + config.url, + auth=config.auth, + max_connection_pool_size=max_connection_pool_size, + max_connection_lifetime=max_connection_lifetime, + ) + + # --------------------------------------------------------------------------- # Cypher identifier validation # --------------------------------------------------------------------------- @@ -1179,12 +1203,15 @@ def __init__( flush_chunk_rows: int = 100, flush_chunk_bytes: int = 4_194_304, neo4j_lock_timeout: float | None = None, + driver: Any | None = None, ) -> None: - """Initialise the store and create the async Neo4j driver. + """Initialise the store, reusing or creating the async Neo4j driver. Args: uri: Bolt/neo4j URI, e.g. ``bolt://localhost:7687``. + Ignored when ``driver`` is provided. auth: ``(username, password)`` tuple, or ``None`` for no-auth. + Ignored when ``driver`` is provided. database: Target Neo4j database name (default: ``"neo4j"``). workspace: Workspace to scope writes to. ``None`` resolves to ``"default"`` via the ``workspace`` property. @@ -1195,18 +1222,29 @@ def __init__( so a blocked flush raises ``Neo4jError`` instead of parking forever. ``None`` disables the timeout (default: no per-transaction limit). - Also sets ``connection_acquisition_timeout`` on - the driver to the same value so pool-exhaustion - failures also surface quickly. + When the store builds its own driver (``driver`` + not provided), this also sets + ``connection_acquisition_timeout`` on it to the + same value so pool-exhaustion failures surface + quickly. + driver: A pre-built async driver to reuse instead of + constructing a new one. When provided, this store + does not own the driver's lifecycle: ``close()`` + flushes and no-ops on the driver itself, leaving + it open for other stores sharing it. """ - # Explicit auto-retry budget for transient errors (e.g. deadlocks) so the - # managed-transaction retry window is deliberate and reviewable rather than - # relying on the driver default implicitly. 30.0s is a working default; - # design Open Question #3 — verify driver 6.1.0 backoff constants before tuning. - driver_kwargs: dict[str, Any] = {"max_transaction_retry_time": 30.0} - if neo4j_lock_timeout is not None and neo4j_lock_timeout > 0: - driver_kwargs["connection_acquisition_timeout"] = neo4j_lock_timeout - self._driver = AsyncGraphDatabase.driver(uri, auth=auth, **driver_kwargs) + if driver is not None: + self._driver = driver + self._owns_driver = False + else: + # Explicit auto-retry budget for transient errors (e.g. deadlocks) so + # the managed-transaction retry window is deliberate and reviewable + # rather than relying on the driver default implicitly. + driver_kwargs: dict[str, Any] = {"max_transaction_retry_time": 30.0} + if neo4j_lock_timeout is not None and neo4j_lock_timeout > 0: + driver_kwargs["connection_acquisition_timeout"] = neo4j_lock_timeout + self._driver = AsyncGraphDatabase.driver(uri, auth=auth, **driver_kwargs) + self._owns_driver = True self._database = database self._workspace = workspace self._created_by: str | None = None @@ -1224,6 +1262,19 @@ def __init__( else None ) + # ------------------------------------------------------------------ + # owns_driver property + # ------------------------------------------------------------------ + + @property + def owns_driver(self) -> bool: + """True when this store built its own driver; False when injected. + + Governs ``close()``: a store that does not own its driver must never + close it, since other stores may still be using it. + """ + return self._owns_driver + # ------------------------------------------------------------------ # workspace property # ------------------------------------------------------------------ @@ -1633,10 +1684,15 @@ async def _ensure_schema(self) -> None: # once Neo4j is reachable / duplicates are cleared by the dedup pass). async def close(self) -> None: - """Flush pending writes, await any background task, and close the driver. + """Flush pending writes and close the driver, if this store owns it. Handles event-loop mismatch gracefully when closing the driver from a different loop context. Sets ``_closed`` on completion. + + When the driver was injected (``owns_driver`` is False), the driver is + left open: it is shared with other stores/callers and closing it here + would break them out from under their own in-flight work. The shared + driver's owner is responsible for closing it exactly once. """ # Final flush to persist remaining buffer contents try: @@ -1646,11 +1702,12 @@ async def close(self) -> None: "Final flush failed during close; buffered writes may be lost" ) - # Close the driver, ignoring event-loop mismatch errors - try: - await self._driver.close() - except RuntimeError: - pass + if self._owns_driver: + # Close the driver, ignoring event-loop mismatch errors + try: + await self._driver.close() + except RuntimeError: + pass self._closed = True diff --git a/context_intelligence_server/queue_manager.py b/context_intelligence_server/queue_manager.py index 6fba186e..301bc4d6 100644 --- a/context_intelligence_server/queue_manager.py +++ b/context_intelligence_server/queue_manager.py @@ -1,37 +1,38 @@ -"""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), ``.log.compact.tmp`` (transient compaction copy). + +Framing (one event == one ``\\n``-terminated byte range) holds only while a +single process writes the directory; each key's ``file_lock`` (a +``threading.Lock`` held on the writing thread) serialises its writes. Records +must contain no raw ``0x0A`` except the terminator. ``session_id`` is the raw +filename stem and is rejected if empty or containing a separator or null byte. +Appends are not ``fsync``ed: crash-durable, not power-loss-durable. """ from __future__ import annotations import asyncio import base64 +import contextlib import json +import logging import os +import threading import time +from collections.abc import Callable, Coroutine, Iterator from dataclasses import dataclass +from enum import Enum from pathlib import Path -from typing import Any +from typing import Any, TypeVar + +from context_intelligence_server.config import get_settings + +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 +42,151 @@ _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] + + +class Verdict(str, Enum): + """Boot-safety classifier verdict. + + RESUMABLE -- keep; a drainer can/should be dispatched for this key. + UNRESUMABLE -- delete: the `.log` cannot reach a drainer at all, or has + nothing left to persist that a reset wouldn't re-derive. + DRAINED -- delete: fully committed, nothing left to persist + (`fully_drained`). + RESET_OFFSET -- delete the `.offset` ONLY; the `.log` re-drains from + byte 0 (bounded re-drain, gated on size + an empty + `.dead.jsonl`). + KEEP -- keep, counted (not resumed, not deleted): either an + inert-but-harmless bucket (`bad_offset_with_dead`, + `unclassifiable`) or a genuinely unowned decision left + to a later pass. + UNREADABLE -- keep, counted as `failed`: a transient FS error, NOT a + corruption finding, and must never be laundered into a + deletion. + """ + + RESUMABLE = "resumable" + UNRESUMABLE = "unresumable" + DRAINED = "drained" + RESET_OFFSET = "reset_offset" + KEEP = "keep" + UNREADABLE = "unreadable" + + +@dataclass(frozen=True) +class Classification: + """One side-effect-free ``classify_session`` verdict. + + ``size`` is the ``.log`` ``st_size`` at classify time; ``reclaim`` re-stats + inside its guarded body and refuses to apply if the size has drifted. + ``dead_empty`` records whether ``.dead.jsonl`` was empty. ``fallback_source`` + is set only when ``reason == "fallback_workspace"``. + """ + + key: str + verdict: Verdict + reason: str # one token from a closed vocabulary; "" for plain RESUMABLE + size: int + dead_empty: bool + fallback_source: str | None = None + + +# A bad-offset log at/below this many +# bytes is RESET (re-drained from byte 0) rather than deleted outright. +# Read from Settings so it stays a single, operator-overridable config knob +# rather than a second hardcoded constant. +def _reclaim_redrain_max_bytes() -> int: + return get_settings().reclaim_redrain_max_bytes + + +@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.""" @@ -78,6 +206,21 @@ 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] = {} + + @property + def queues_dir(self) -> Path: + """The queue directory this manager owns (for main._boot_reclaim's + `*.log` glob, so it never needs to recompute the path from settings -- + it reads it from the SAME QueueManager instance registry.queue_manager + already resolved, avoiding drift from a test/instance that points the + registry's queue manager at a different directory).""" + return self._dir def _log_path(self, session_id: str) -> Path: return self._dir / f"{session_id}.log" @@ -88,28 +231,49 @@ def _offset_path(self, session_id: str) -> Path: def _dead_path(self, session_id: str) -> Path: return self._dir / f"{session_id}.dead.jsonl" + def _compact_tmp_path(self, session_id: str) -> Path: + """The tmp used by ``compact_committed_prefix``'s tail copy. + + Matches no glob any existing pass uses (``*.log``, ``*.offset``, + ``*.offset.tmp``, ``*.dead.jsonl``, ``*.log.torn-*.bin``) -- a stray + left by a crash between steps 3-5 is inert until ``reclaim_orphans`` + (``orphan_compact_tmp``) reaps it. + """ + return self._dir / f"{session_id}.log.compact.tmp" + def _read_committed_offset(self, session_id: str) -> int: + """Committed byte offset; reads bare-int and legacy JSON offset files.""" try: text = self._offset_path(session_id).read_text("utf-8") except FileNotFoundError: return 0 text = text.strip() - return int(text) if text else 0 + if not text: + return 0 + if text[0] == "{": + try: + cursor = json.loads(text) + return int(cursor["offset"]) + except (json.JSONDecodeError, KeyError, TypeError, ValueError): + raise ValueError( + f"unparseable legacy offset document for session {session_id!r}" + ) from None + return int(text) - 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``, a pure delegating refactor) and + ``.dead.jsonl`` files (``heal_torn_tails``). """ - path = self._log_path(session_id) try: with open(path, "rb") as f: f.seek(0, os.SEEK_END) @@ -126,6 +290,14 @@ 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``. + + Pure delegation to ``_last_complete_end`` -- no behaviour change from + the pre-refactor inline version. + """ + 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. @@ -181,16 +353,197 @@ 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 tail left; heal_torn_tails will remove it at next boot)", + 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) + + @staticmethod + def _heal_one(path: Path) -> tuple[int, bool]: + """Heal one torn tail. Returns ``(bytes_discarded, healed)``. + + Ordering is load-bearing: copy the torn bytes to a quarantine sidecar, + verify it is byte-complete, then truncate -- never the reverse. A + partial or failed quarantine leaves the file untouched (readers already + skip the tail; next boot retries). Raises ``OSError`` on any failure; + the caller catches it per file. + """ + end = QueueManager._last_complete_end(path) + size = path.stat().st_size + if size <= end: + return 0, False + torn_bytes = size - end + quarantine = path.with_name(f"{path.name}.torn-{time.time_ns()}.bin") + + with open(path, "rb") as src: + src.seek(end) + data = src.read() + if len(data) != torn_bytes: + raise OSError( + f"short read quarantining {path}: expected {torn_bytes} bytes, " + f"got {len(data)}" + ) + + flags = os.O_WRONLY | os.O_CREAT | os.O_TRUNC | getattr(os, "O_BINARY", 0) + fd = os.open(quarantine, flags, 0o644) + try: + QueueManager._write_all(fd, data) + finally: + os.close(fd) + + written = quarantine.stat().st_size + if written != torn_bytes: + raise OSError( + f"quarantine incomplete for {path}: wrote {written} of {torn_bytes} bytes" + ) + + # ONLY now, with a verified-complete quarantine on disk, is it safe + # to shorten the original file. + os.truncate(path, end) + return torn_bytes, True + + async def heal_torn_tails(self) -> dict[str, int]: + """One-time startup pass: truncate every queue file back to its last + complete line, quarantining the removed bytes. + + Runs once per boot before any reader/writer is live, and is the only + place a queue file is shortened. Each ``*.log``/``*.dead.jsonl`` is + healed independently; a per-file failure is logged and skipped. Must + not raise (the caller is a lifespan hook; a raise would restart-loop on + the share being healed). Returns + ``{"files_healed", "bytes_discarded", "files_failed"}``. + """ + + def _heal_all() -> dict[str, int]: + files_healed = 0 + bytes_discarded = 0 + files_failed = 0 + try: + paths = sorted(self._dir.glob("*.log")) + sorted( + self._dir.glob("*.dead.jsonl") + ) + except OSError: + logger.exception("heal_torn_tails_scan_failed dir=%s", self._dir) + return { + "files_healed": 0, + "bytes_discarded": 0, + "files_failed": 0, + } + for path in paths: + try: + discarded, healed = self._heal_one(path) + except OSError: + files_failed += 1 + logger.exception("torn_tail_heal_failed path=%s", path) + continue + if healed: + files_healed += 1 + bytes_discarded += discarded + logger.warning( + "torn_tail_healed path=%s discarded=%d", + path, + discarded, + ) + result = { + "files_healed": files_healed, + "bytes_discarded": bytes_discarded, + "files_failed": files_failed, + } + logger.info("heal_torn_tails result=%s", result) + if files_failed: + logger.error("heal_torn_tails_incomplete files_failed=%d", files_failed) + return result + + return await asyncio.to_thread(_heal_all) + 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" + line = raw if raw.endswith(b"\n") else raw + b"\n" # unchanged (was :186) 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 (v2.1 G1) -- ``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 +551,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 +577,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 +593,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 +608,248 @@ 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: + # Previously ZERO logging here -- an OSError + # killed the drainer and showed only as a generic + # drain_worker_died with no hint the FAILING write was a + # dead-letter. LOGGING ONLY: re-raise unchanged so + # propagation behavior is identical to before this line. + # (traceback carries the exception; exc is not repeated.) + 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 + 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 compact_committed_prefix( + self, session_id: str, min_prefix_bytes: int = 0 + ) -> int: + """Rewrite ``.log`` to keep only its undrained tail. + + Reclaims the committed prefix ``[0, C)`` while the session stays live + (unlike ``delete_drained``, which removes the whole file at finalize). + Returns the reclaimed prefix byte count ``C``; ``0`` means nothing was + done (no prefix, or failure). Reclaims regardless of tail size -- a + large tail only costs more lock-hold time, never a skipped reclaim. + Never raises. + + Crash ordering: rebase ``.offset`` to 0 first (atomic tmp + replace), + then replace the ``.log`` with the verified tail; if that fails, restore + ``.offset := C``. Every window degrades to a bounded re-drive, not a + loss. The guard is kept (the log still exists and the session is live). + """ + self._validate_session_id(session_id) + log = self._log_path(session_id) + offset = self._offset_path(session_id) + offset_tmp = self._dir / f"{session_id}.offset.tmp" + tmp = self._compact_tmp_path(session_id) + + def _compact(_guard: _KeyGuard) -> int: + with _guard.file_lock: + try: + c = self._read_committed_offset(session_id) + except (OSError, ValueError): + return 0 + try: + e = log.stat().st_size + except OSError: + return 0 + + # Step 2: bail (return 0) unless C >= min_prefix_bytes and + # 0 < C <= E. Reclaimed regardless of tail size. + if not (c >= min_prefix_bytes and 0 < c <= e): + return 0 + tail = e - c + + # Step 3: copy [C, E) into a fresh tmp. O_TRUNC so a stray + # tmp from a previous attempt is never appended to. + flags = ( + os.O_WRONLY | os.O_CREAT | os.O_TRUNC | getattr(os, "O_BINARY", 0) + ) + try: + with open(log, "rb") as src: + src.seek(c) + fd = os.open(tmp, flags, 0o644) + try: + remaining = tail + while remaining > 0: + chunk = src.read(min(_SCAN_CHUNK_BYTES, remaining)) + if not chunk: + break + self._write_all(fd, chunk) + remaining -= len(chunk) + finally: + os.close(fd) + except OSError: + logger.exception("compact_copy_failed session=%s", session_id) + with contextlib.suppress(OSError): + tmp.unlink() + return 0 + + # Step 4: verify byte-complete BEFORE destroying anything -- + # _heal_one's ordering (quarantine, verify, only then act). + try: + written = tmp.stat().st_size + except OSError: + written = -1 + if written != tail: + with contextlib.suppress(OSError): + tmp.unlink() + logger.error( + "compact_aborted session=%s cause=short_copy " + "expected=%d got=%d", + session_id, + tail, + written, + ) + return 0 + + # Step 5: rebase the offset to 0 FIRST -- the point of no + # return. + try: + offset_tmp.write_text("0", encoding="utf-8") + os.replace(offset_tmp, offset) + except OSError: + logger.exception( + "compact_offset_rebase_failed session=%s", session_id + ) + with contextlib.suppress(OSError): + tmp.unlink() + return 0 - await asyncio.to_thread(_delete) + # Step 6: replace the log with the verified tail copy. + try: + os.replace(tmp, log) + except OSError: + # R3: restore the offset to C so this becomes a PURE + # NO-OP -- never an in-process re-drive (which would + # double-count `written` and drive the residual + # negative). + try: + offset_tmp.write_text(str(c), encoding="utf-8") + os.replace(offset_tmp, offset) + logger.error( + "compact_replace_failed session=%s committed=%d " + "action=offset_restored", + session_id, + c, + ) + except OSError: + logger.error( + "compact_restore_failed session=%s committed=%d " + "redrive_expected=true", + session_id, + c, + ) + with contextlib.suppress(OSError): + tmp.unlink() + return 0 + + logger.info( + "queue_compacted session=%s reclaimed=%d tail=%d", + session_id, + c, + tail, + ) + return c + + try: + with self._guard(session_id) as guard: + async with guard.admission: + reclaimed = await _await_uninterrupted( + asyncio.to_thread(_compact, guard) + ) + except Exception: + # Precision 1: a blanket catch so an OSError mid-copy can never + # escape into drain_worker and kill a healthy drainer. + # asyncio.CancelledError derives from BaseException, so it still + # propagates -- required, or the drainer's cancellation paths + # would break. + logger.exception("compact_failed session=%s", session_id) + return 0 + + if reclaimed: + self._stats_cache = None + self._spool_cache = None + return reclaimed 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,10 +858,545 @@ 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) + @staticmethod + def _is_parseable_line(raw: bytes) -> bool: + """Total: True iff ``raw`` is valid JSON. Never raises.""" + try: + json.loads(raw) + except (ValueError, TypeError): + return False + return True + + async def read_first_line(self, key: str) -> bytes | None: + """Return the FIRST (byte-0) line of ``key``'s `.log`, or None. + + One `open` + one `readline` at offset 0 -- bounded, cheap. Returns + None when the file is missing, empty, or its first line is not yet + newline-terminated (torn/incomplete). Never raises (Q-13): + this feeds the boot-path workspace-fallback resolution + (``main._recover_one_session``), which must be a total function. + """ + self._validate_session_id(key) + path = self._log_path(key) + + def _read() -> bytes | None: + try: + with open(path, "rb") as f: + raw = f.readline() + except OSError: + return None + if not raw or not raw.endswith(b"\n"): + return None + return raw[:-1] + + return await asyncio.to_thread(_read) + + async def classify_session( + self, + key: str, + head_is_resumable: Callable[[bytes], bool], + ) -> Classification: + """Boot-safety classifier. Side-effect-free (reads only). + + Bounded I/O: never a whole-log read. Must not raise -- an + unattributable OSError/ValueError becomes ``Verdict.UNREADABLE`` (the + caller is a boot hook; a raise would restart-loop on the share). ``key`` + is a ``.log`` stem; log-less keys are ``reclaim_orphans``'s to own. + """ + self._validate_session_id(key) + log_path = self._log_path(key) + dead_path = self._dead_path(key) + threshold = _reclaim_redrain_max_bytes() + + def _dead_empty() -> bool: + try: + return dead_path.stat().st_size == 0 + except FileNotFoundError: + return True # absent .dead.jsonl counts as empty + except OSError: + # Cannot prove empty -> the conservative, non-destructive + # answer is "not empty" (refuse RESET_OFFSET, fall to KEEP). + return False + + def _bad_offset(reason: str, size: int, dead_empty: bool) -> Classification: + # unparseable_offset: the .log bytes are unexamined here, so an + # unreadable sidecar alone must never delete them -- size-gate + # only the reasons where the parsed value itself is bad. + if size <= threshold or reason == "unparseable_offset": + if dead_empty: + return Classification( + key, Verdict.RESET_OFFSET, reason, size, dead_empty + ) + return Classification( + key, Verdict.KEEP, "bad_offset_with_dead", size, dead_empty + ) + return Classification(key, Verdict.UNRESUMABLE, reason, size, dead_empty) + + def _classify() -> Classification: + try: + size = log_path.stat().st_size + except FileNotFoundError: + # A classify-time race (unlinked between the + # `*.log` glob and this call) -- benign, not corruption. + return Classification(key, Verdict.UNREADABLE, "log_vanished", 0, True) + except OSError: + return Classification( + key, Verdict.UNREADABLE, "unreadable_offset", 0, _dead_empty() + ) + + dead_empty = _dead_empty() + + try: + committed = self._read_committed_offset(key) + except OSError: + return Classification( + key, Verdict.UNREADABLE, "unreadable_offset", size, dead_empty + ) + except ValueError: + return _bad_offset("unparseable_offset", size, dead_empty) + + if committed < 0: + return _bad_offset("negative_offset", size, dead_empty) + if committed > size: + return _bad_offset("offset_past_eof", size, dead_empty) + if size == 0: + return Classification( + key, Verdict.UNRESUMABLE, "empty_log", size, dead_empty + ) + + complete_end = self._complete_data_end(key) + if committed >= complete_end == size: + return Classification( + key, Verdict.DRAINED, "fully_drained", size, dead_empty + ) + if committed >= complete_end < size: + # Un-newline-terminated remainder heal didn't reach this + # boot (files_failed > 0) -- not our data to delete; heal + # retries next boot. Not resumable YET by recover()'s own + # predicate, but classify must not judge it un-resumable. + return Classification(key, Verdict.RESUMABLE, "", size, dead_empty) + + # Head-record check (the common, 99% case): the first + # UNCOMMITTED line, exactly what `_recover_one_session` parses. + try: + with open(log_path, "rb") as f: + f.seek(committed) + head_raw = f.readline() + except OSError: + return Classification( + key, Verdict.UNREADABLE, "unreadable_offset", size, dead_empty + ) + if head_raw.endswith(b"\n") and head_is_resumable(head_raw[:-1]): + return Classification(key, Verdict.RESUMABLE, "", size, dead_empty) + + # The head is unparseable/torn/lacks a workspace -- resume + # with a FALLBACK workspace instead of deleting. + # Step 10: byte 0 of the SAME file. + try: + with open(log_path, "rb") as f: + byte0_raw = f.readline() + except OSError: + byte0_raw = b"" + if byte0_raw.endswith(b"\n") and head_is_resumable(byte0_raw[:-1]): + return Classification( + key, + Verdict.RESUMABLE, + "fallback_workspace", + size, + dead_empty, + fallback_source="byte0", + ) + + # Step 11: any parseable line within the first _SCAN_CHUNK_BYTES + # from byte 0 (bounded -- never a whole-log read). + window = min(size, _SCAN_CHUNK_BYTES) + try: + with open(log_path, "rb") as f: + buf = f.read(window) + except OSError: + buf = b"" + found_parseable = any( + self._is_parseable_line(line) for line in buf.split(b"\n")[:-1] + ) + if found_parseable: + return Classification( + key, + Verdict.RESUMABLE, + "fallback_workspace", + size, + dead_empty, + fallback_source="sentinel", + ) + + # Step 12/13: bounding decides DELETE vs KEEP -- DELETE only when + # the probe window provably covered the WHOLE file. + if size <= _SCAN_CHUNK_BYTES: + # `main._recover_one_session` is MORE LENIENT than + # this classifier was -- it unconditionally sentinel- + # dispatches whenever byte 0 is a COMPLETE (newline- + # terminated) line, regardless of JSON parseability, because + # the drainer dead-letters the unparseable head and drains + # everything behind it. Deleting here would destroy data the + # recovery path would have kept. DELETE is reserved for the + # genuinely-unrecoverable case: no complete line at byte 0 + # at all (a torn-from-the-start log -- heal's territory, not + # ours) -- which cannot co-occur with `complete_end > + # committed >= 0` (already established above) but is kept as + # a defensive, provably-safe fallback rather than assumed. + if byte0_raw.endswith(b"\n"): + return Classification( + key, + Verdict.RESUMABLE, + "fallback_workspace", + size, + dead_empty, + fallback_source="sentinel", + ) + return Classification( + key, Verdict.UNRESUMABLE, "no_parseable_line", size, dead_empty + ) + return Classification(key, Verdict.KEEP, "unclassifiable", size, dead_empty) + + return await asyncio.to_thread(_classify) + + async def reclaim( + self, + c: Classification, + is_owned: Callable[[], bool], + ) -> bool: + """Apply a ``classify_session`` verdict: log-then-delete, or reset. + + Emits the ``boot_reclaimed`` audit line BEFORE any unlink, so a crash + mid-unlink still records the intent. Re-verifies inside the guarded body + (the server runs concurrently with this pass): ownership still False, + live ``.log`` size still matches ``c.size``, and for RESET_OFFSET the + ``.dead.jsonl`` still empty; any drift applies nothing and returns + False. Actionable only for UNRESUMABLE/DRAINED (delete) and + RESET_OFFSET (reset). Must not raise. + """ + self._validate_session_id(c.key) + if c.verdict not in ( + Verdict.UNRESUMABLE, + Verdict.DRAINED, + Verdict.RESET_OFFSET, + ): + return False + + log = self._log_path(c.key) + offset = self._offset_path(c.key) + offset_tmp = self._dir / f"{c.key}.offset.tmp" + dead = self._dead_path(c.key) + action = "reset_offset" if c.verdict is Verdict.RESET_OFFSET else "delete" + + def _apply(guard: _KeyGuard) -> bool: + with guard.file_lock: + # Ownership, re-checked FRESH inside the guarded + # body -- the registry owns live sessions; a session that acquired a + # live worker after classify-time must never be reclaimed. + if is_owned(): + logger.warning( + "boot_reclaim_skipped_changed session=%s reason=%s cause=owned", + c.key, + c.reason, + ) + return False + try: + live_size = log.stat().st_size + except FileNotFoundError: + live_size = 0 + except OSError: + logger.error( + "boot_reclaim_failed reason=%s path=%s session=%s error=stat_failed", + c.reason, + log, + c.key, + ) + return False + if live_size != c.size: + logger.warning( + "boot_reclaim_skipped_changed session=%s reason=%s cause=size_drift", + c.key, + c.reason, + ) + return False + if c.verdict is Verdict.RESET_OFFSET: + try: + dead_empty_now = dead.stat().st_size == 0 + except FileNotFoundError: + dead_empty_now = True + except OSError: + dead_empty_now = False + if not dead_empty_now: + logger.warning( + "boot_reclaim_skipped_changed session=%s reason=%s " + "cause=dead_nonempty", + c.key, + c.reason, + ) + return False + + logger.warning( + "boot_reclaimed reason=%s path=%s session=%s bytes=%d action=%s", + c.reason, + log, + c.key, + c.size, + action, + ) + try: + if c.verdict is Verdict.RESET_OFFSET: + # Unlink ONLY the offset (+ any stray .offset.tmp) -- + # the .log stays; the next drain re-reads from 0. + try: + offset.unlink() + except FileNotFoundError: + pass + try: + offset_tmp.unlink() + except FileNotFoundError: + pass + else: + # Unlink order is load-bearing: .log -> .offset -> + # .offset.tmp. A crash after the .log unlink leaves + # an orphan .offset, self-healed by the NEXT boot's + # reclaim_orphans; the reverse order would leave a + # .log with no .offset (committed==0) -- a full + # replay of data just judged un-resumable. Never. + try: + log.unlink() + except FileNotFoundError: + pass + try: + offset.unlink() + except FileNotFoundError: + pass + try: + offset_tmp.unlink() + except FileNotFoundError: + pass + except OSError: + logger.exception( + "boot_reclaim_failed reason=%s path=%s session=%s", + c.reason, + log, + c.key, + ) + return False + return True + + with self._guard(c.key) as guard: + async with guard.admission: + ok = await _await_uninterrupted(asyncio.to_thread(_apply, guard)) + if ( + ok + and c.verdict is not Verdict.RESET_OFFSET + and guard.waiters == 1 + and self._guards.get(c.key) is guard + ): + # The guard-removal condition, applied identically + # here: only when this call is the SOLE holder of the + # guard AND the identity check passes. A RESET_OFFSET + # leaves the .log in place, so its guard stays live. + del self._guards[c.key] + if ok: + self._stats_cache = None + self._spool_cache = None + return ok + + async def reclaim_orphans( + self, before_ts: float, enabled: bool = True + ) -> dict[str, int]: + r"""One ``stat()``-only directory pass over log-less artifacts. + + Classifies and (when ``enabled``) unlinks: ``.offset``/``.offset.tmp`` + with no ``.log``; ``*.torn-*.bin`` sidecars older than ``before_ts`` + (this boot's own are kept); and ``*.log.compact.tmp`` older than + ``before_ts`` (age-gated, not log-less-gated, since a compaction tmp + sits beside a live log). With ``enabled`` False every candidate is only + logged as a dry-run. ``reclaimed``/``reclaimed_bytes`` count real + unlinks only. Must not raise. + """ + + def _scan() -> dict[str, int]: + reclaimed = 0 + reclaimed_bytes = 0 + failed = 0 + action = "delete" if enabled else "dry_run" + try: + offset_paths = sorted(self._dir.glob("*.offset")) + tmp_paths = sorted(self._dir.glob("*.offset.tmp")) + torn_paths = sorted(self._dir.glob("*.log.torn-*.bin")) + compact_tmp_paths = sorted(self._dir.glob("*.log.compact.tmp")) + except OSError: + logger.exception("reclaim_orphans_scan_failed dir=%s", self._dir) + return { + "reclaimed": 0, + "reclaimed_bytes": 0, + "failed": 0, + } + for path, reason in [ + *((p, "orphan_offset") for p in offset_paths), + *((p, "orphan_offset_tmp") for p in tmp_paths), + ]: + stem = path.name[ + : -len(".offset.tmp" if reason.endswith("tmp") else ".offset") + ] + if self._log_path(stem).exists(): + continue + try: + size = path.stat().st_size + except OSError: + failed += 1 + logger.exception( + "boot_reclaim_failed reason=%s path=%s", reason, path + ) + continue + logger.warning( + "boot_reclaimed reason=%s path=%s session=%s bytes=%d action=%s", + reason, + path, + stem, + size, + action, + ) + if not enabled: + continue + try: + path.unlink() + reclaimed += 1 + reclaimed_bytes += size + except OSError: + failed += 1 + logger.exception( + "boot_reclaim_failed reason=%s path=%s", reason, path + ) + for path in torn_paths: + try: + mtime = path.stat().st_mtime + except OSError: + # Previously silent -- matches the + # neighbouring boot_reclaim_failed shape (the size-stat + # and unlink failures just below already log this way). + failed += 1 + logger.exception( + "boot_reclaim_failed reason=torn_sidecar path=%s (mtime stat)", + path, + ) + continue + if mtime >= before_ts: + continue # created by THIS boot's heal -- keep it + try: + size = path.stat().st_size + except OSError: + failed += 1 + logger.exception( + "boot_reclaim_failed reason=torn_sidecar path=%s", path + ) + continue + logger.warning( + "boot_reclaimed reason=torn_sidecar path=%s session=%s " + "bytes=%d action=%s", + path, + path.name.split(".log.torn-")[0], + size, + action, + ) + if not enabled: + continue + try: + path.unlink() + reclaimed += 1 + reclaimed_bytes += size + except OSError: + failed += 1 + logger.exception( + "boot_reclaim_failed reason=torn_sidecar path=%s", path + ) + for path in compact_tmp_paths: + # NOT gated on + # log-absence -- a stray tmp coexists with a very much + # still-live `.log`. Its OWN suffix strip + # (".log.compact.tmp", NOT the shorter ".offset.tmp" slice + # length) so the audit line names the right session. + try: + mtime = path.stat().st_mtime + except OSError: + # Previously silent -- matches the + # neighbouring boot_reclaim_failed shape. + failed += 1 + logger.exception( + "boot_reclaim_failed reason=orphan_compact_tmp " + "path=%s (mtime stat)", + path, + ) + continue + if mtime >= before_ts: + continue # would-be THIS boot's own compaction -- keep it + stem = path.name[: -len(".log.compact.tmp")] + try: + size = path.stat().st_size + except OSError: + failed += 1 + logger.exception( + "boot_reclaim_failed reason=orphan_compact_tmp path=%s", + path, + ) + continue + logger.warning( + "boot_reclaimed reason=orphan_compact_tmp path=%s session=%s " + "bytes=%d action=%s", + path, + stem, + size, + action, + ) + if not enabled: + continue + try: + path.unlink() + reclaimed += 1 + reclaimed_bytes += size + except OSError: + failed += 1 + logger.exception( + "boot_reclaim_failed reason=orphan_compact_tmp path=%s", + path, + ) + return { + "reclaimed": reclaimed, + "reclaimed_bytes": reclaimed_bytes, + "failed": failed, + } + + result = await asyncio.to_thread(_scan) + if result["reclaimed"]: + self._stats_cache = None + self._spool_cache = None + return result + async def active_sessions(self) -> list[str]: """Return sorted session_ids with undrained data. @@ -310,8 +1410,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 +1441,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 +1489,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 ( @@ -439,58 +1546,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 ( @@ -584,64 +1649,152 @@ 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 - - return await 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: + with guard.file_lock: + count = self._count_dead(worker_key) + try: + path.unlink() + except FileNotFoundError: + pass + return count - - ``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 + with self._guard(worker_key) as guard: + async with guard.admission: + return await _await_uninterrupted(asyncio.to_thread(_purge)) + + async def expire_dead_letters( + self, now: float, retention_seconds: float, enabled: bool + ) -> dict[str, int]: + r"""Expire log-less dead-letter files older than ``retention_seconds``. + + Expired iff ``.log`` is absent AND + ``now - mtime(.dead.jsonl) > retention_seconds``. Whole-file mtime + is safe because ``dead_letter`` only appends, so an old mtime means + every line is old (an actively-failing session stays visible by + design). ``retention_seconds <= 0`` disables expiry. With ``enabled`` + False every candidate is only logged as a dry-run; the counts reflect + real unlinks only. Deletes route through ``purge_dead_letters``. Must + not raise. + """ + zeros = { + "expired_keys": 0, + "expired_records": 0, + "expired_bytes": 0, + "failed": 0, + } + if retention_seconds <= 0: + return zeros - Formula:: + def _scan() -> list[str]: + candidates: list[str] = [] + for dead_path in sorted(self._dir.glob("*.dead.jsonl")): + key = dead_path.name[: -len(".dead.jsonl")] + if self._log_path(key).exists(): + continue # Never touch a key with a live .log + try: + mtime = dead_path.stat().st_mtime + except OSError: + continue + if now - mtime > retention_seconds: + candidates.append(key) + return candidates - written_seed = max(0, C - D) - accepted_seed = written_seed + P + D + try: + candidates = await asyncio.to_thread(_scan) + except OSError: + logger.exception("dead_letter_expire_scan_failed dir=%s", self._dir) + return dict(zeros, failed=1) + + expired_keys = 0 + expired_records = 0 + expired_bytes = 0 + failed = 0 + action = "delete" if enabled else "dry_run" + for key in candidates: + dead_path = self._dead_path(key) + try: + size = dead_path.stat().st_size + records = self._count_dead(key) + except OSError: + failed += 1 + logger.exception("dead_letter_expire_stat_failed key=%s", key) + continue + age_seconds = now - dead_path.stat().st_mtime + logger.warning( + "dead_letter_expired key=%s records=%d bytes=%d " + "age_seconds=%.0f action=%s", + key, + records, + size, + age_seconds, + action, + ) + if not enabled: + continue + try: + await self.purge_dead_letters(key) + except (OSError, ValueError): + failed += 1 + logger.exception("dead_letter_expire_purge_failed key=%s", key) + continue + expired_keys += 1 + expired_records += records + expired_bytes += size - 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. + if expired_keys: + self._stats_cache = None + self._spool_cache = None + return { + "expired_keys": expired_keys, + "expired_records": expired_records, + "expired_bytes": expired_bytes, + "failed": failed, + } - Ordering is load-bearing: this MUST run AFTER ``recovery_reconcile_dead`` - in the lifespan so the dead-letter counts it reads are already settled. + async def recovery_seed_counts(self) -> tuple[int, int]: + """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 +1810,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, # Q-5: 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 + # Q-6: check `.log` existence BEFORE reading the whole + # `.dead.jsonl` into RAM (+ a payload set at ~3.6x its size). + # A key with only a `.dead.jsonl` (the common shape left by + # `delete_drained`) previously paid that read for nothing, + # every boot, forever. Free fix. log_path = self._log_path(key) if not log_path.exists(): continue - 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..3f215584 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 @@ -11,22 +12,26 @@ from context_intelligence_server.blob_store import AsyncDiskBlobStore from context_intelligence_server.config import get_settings -from context_intelligence_server.status import EventRecord, ring_buffer -from context_intelligence_server.neo4j_store import Neo4jGraphStore +from context_intelligence_server.neo4j_store import ( + Neo4jGraphStore, + build_bounded_neo4j_driver, +) from context_intelligence_server.pipeline import process_event, setup_handlers from context_intelligence_server.queue_manager import Batch, QueueManager from context_intelligence_server.services import HookStateService +from context_intelligence_server.status import EventRecord, ring_buffer logger = logging.getLogger("context_intelligence_server") _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 +46,18 @@ 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 + # False only for a crash-recovery-dispatched worker that hasn't seen a live + # POST yet; gates the dry-exit in drain_worker once its backlog drains. + live_event_seen: bool = True + # Scheduling flag: a commit landed since the last compaction attempt. + # Lets an idle drainer skip a no-op compaction each poll tick. + compact_pending: bool = False @dataclass @@ -64,26 +77,27 @@ 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. + # Shared, pool-bounded Neo4j driver for every per-session Neo4jGraphStore + # (see _ensure_neo4j_driver). Built lazily for the same reason as + # _queue_manager; kept separate from _ensure_infra so the two concerns + # can evolve independently. + self._neo4j_driver: Any | None = None + # 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: @@ -106,6 +120,51 @@ def queue_manager(self) -> QueueManager: assert self._queue_manager is not None return self._queue_manager + @property + def queues_dir_path(self) -> Path: + """Queue directory path, resolved without constructing a QueueManager. + + Unlike ``queue_manager``, never calls ``_ensure_infra`` -- avoids a + race where an observer builds a second QueueManager for the same + directory. Falls back to the same expression ``_ensure_infra`` uses, + so the two can never disagree. + """ + if self._queue_manager is not None: + return self._queue_manager.queues_dir + return Path(get_settings().queues_path) + + def _ensure_neo4j_driver(self) -> Any: + """Build the shared, pool-bounded Neo4j driver on first use. + + Lazy for the same reason as ``_ensure_infra``. Kept as its own method + (not folded into ``_ensure_infra``) so the two constructions stay + independent edits. + """ + if self._neo4j_driver is None: + settings = get_settings() + admin = settings.resolve_neo4j_admin() + self._neo4j_driver = build_bounded_neo4j_driver( + admin, + max_connection_pool_size=settings.neo4j_max_connection_pool_size, + max_connection_lifetime=settings.neo4j_max_connection_lifetime, + ) + return self._neo4j_driver + + @property + def neo4j_driver(self) -> Any: + """The single shared, pool-bounded driver used by every per-session + Neo4jGraphStore -- never closed by a per-session finalize.""" + return self._ensure_neo4j_driver() + + async def close_neo4j_driver(self) -> None: + """Close the shared driver exactly once, at process shutdown. + + No-op if the driver was never built (no session has run yet). + """ + if self._neo4j_driver is not None: + await self._neo4j_driver.close() + self._neo4j_driver = None + @property def write_semaphore(self) -> asyncio.Semaphore: """The single shared global cap on concurrent Neo4j-write flushes.""" @@ -128,17 +187,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 +230,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 +299,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 +313,16 @@ 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. - - 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. + The semaphore caps concurrent write transactions across all session + drainers. The offset must only advance after this returns successfully; + correctness relies on the GraphStore protocol's flush-failure isolation. """ 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 +330,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 +354,27 @@ 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: + # Idle compaction runs before the dry-exit check below, so a + # recovered drainer compacts before exiting instead of leaking its log. + if worker.compact_pending: + worker.compact_pending = False + settings = get_settings() + if settings.queue_compact_enabled: + await qm.compact_committed_prefix(session_id, 0) + # Dry-exit for a recovered drainer with no terminal record: re-read + # after the await closes the race with a live POST arriving mid-check. + if not worker.live_event_seen: + recheck = await qm.read_batch(session_id, max_items=1) + if not recheck.records and not worker.live_event_seen: + await self._safe_close(worker) + self._deregister(session_id) + logger.info( + "recovered_drainer_exited session=%s reason=drained", + session_id, + extra={"session_id": session_id}, + ) + return await asyncio.sleep(poll_interval) idle_elapsed += poll_interval if idle_elapsed >= flush_timeout: @@ -340,9 +386,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 +400,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 +446,61 @@ 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 None: + # Skipped for a terminal batch: it goes straight to + # _finalize_session -> delete_drained, so compacting first is waste. + worker.compact_pending = True + settings = get_settings() + if settings.queue_compact_enabled: + await qm.compact_committed_prefix( + session_id, settings.queue_compact_min_prefix_bytes + ) + + 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 +511,156 @@ 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. + """ + 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 +673,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 +715,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( @@ -551,7 +806,15 @@ def get_or_create( session_id: str, workspace: str, created_by: str | None = None, + *, + recovered: bool = False, ) -> SessionWorker: + """Get-or-create the sticky drainer for ``session_id``. + + ``recovered=True`` (from the crash-recovery/sweep path) is the one + place ``live_event_seen`` is set False, letting a never-live worker + dry-exit once its backlog drains (see ``drain_worker``). + """ if session_id not in self._workers: settings = get_settings() blob_store = AsyncDiskBlobStore(root=settings.blob_path) @@ -559,6 +822,7 @@ def get_or_create( neo4j_store = Neo4jGraphStore( uri=_admin.url, auth=_admin.auth, + driver=self.neo4j_driver, flush_chunk_rows=settings.neo4j_flush_chunk_rows, flush_chunk_bytes=settings.neo4j_flush_chunk_bytes, neo4j_lock_timeout=settings.neo4j_lock_timeout, @@ -572,6 +836,7 @@ def get_or_create( blob_store=blob_store, graph_store=neo4j_store, ), + live_event_seen=not recovered, ) self.start_drain(self._workers[session_id]) logger.info( @@ -580,10 +845,15 @@ 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]) + # Guarded by the parameter: only a call that omits `recovered` + # (a real live POST) flips this True. + if not recovered: + self._workers[session_id].live_event_seen = True + # 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 +876,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: @@ -628,15 +905,20 @@ def workers(self) -> list[SessionWorker]: """Return the list of all active SessionWorker objects.""" return list(self._workers.values()) + def has_worker(self, session_id: str) -> bool: + """The public read for boot-reclaim's ownership gate. + + Lets ``main._boot_reclaim`` check live ownership WITHOUT reaching + into ``_workers`` and without paying O(n) per key over ``workers()``. + """ + return session_id in self._workers + 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/context_intelligence_server/routers/queues.py b/context_intelligence_server/routers/queues.py index 184319fb..6201cfe5 100644 --- a/context_intelligence_server/routers/queues.py +++ b/context_intelligence_server/routers/queues.py @@ -12,7 +12,7 @@ import logging from typing import Any -from fastapi import APIRouter, Depends, HTTPException # noqa: F401 (HTTPException per spec) +from fastapi import APIRouter, Depends, HTTPException from fastapi.requests import Request from context_intelligence_server.authz import require_read, require_write diff --git a/context_intelligence_server/services.py b/context_intelligence_server/services.py index b5f0230b..d861ca6a 100644 --- a/context_intelligence_server/services.py +++ b/context_intelligence_server/services.py @@ -12,6 +12,7 @@ from datetime import datetime from typing import Any +from context_intelligence_server.graph_store import GraphStore from context_intelligence_server.handlers.data_layer_2.state import DataLayer2State from context_intelligence_server.handlers.data_layer_3.state import DataLayer3State @@ -154,15 +155,8 @@ async def find_delegation_by_sub_session( ) -> dict[str, Any] | None: """Return a copy of the Delegation node whose sub_session_id matches, or None. - Scans the in-memory node store for a node carrying the ``Delegation`` - label with a matching ``sub_session_id`` property -- the parent - Delegation that spawned *sub_session_id*. - - ``GraphState`` is a single-workspace store (workspace is fixed at - construction), so the *workspace* argument is accepted for parity with - other ``GraphStore`` implementations (e.g. ``Neo4jGraphStore``) but is - not used to filter here -- every node in this store already belongs to - the same workspace. + *workspace* is accepted for parity with other ``GraphStore`` + implementations but unused here -- this store is single-workspace. """ for data in self._nodes.values(): if ( @@ -225,13 +219,16 @@ class HookStateService: def __init__( self, workspace: str = "default", - graph_store: Any | None = None, + graph_store: GraphStore | None = None, *, created_by: str | None = None, raw_config: dict[str, Any] | None = None, blob_store: Any | None = None, ) -> None: self.config = HookConfig(raw_config or {}) + # Explicitly `Any`: tests reach into GraphState's private buffers + # directly, so narrowing this type ripples into many test-file errors. + self.graph: Any if graph_store is not None: self.graph = graph_store else: @@ -250,41 +247,18 @@ def __init__( async def ensure_session_node(self, session_id: str, data: dict[str, Any]) -> None: """Idempotently create a Session node in the graph for *session_id*. - Uses a two-tier lookup for replay resilience: - - 1. Fast path — if *session_id* is already in the in-memory - ``_seen_sessions`` cache, return immediately. - 2. Graph query — call ``graph.get_node(session_id)``. If the node - already exists (e.g. from a previous run), repopulate the cache and - return without overwriting any data. If the node is absent, create - it with labels ``["Session"]`` and ``status = 'running'``. - - This method is a safety net that creates a minimal session node if it - doesn't exist. ``SessionHandler`` is the sole authority on session - type labels (``RootSession``, ``SubSession``, ``ForkedSession``). - ``ensure_session_node`` always creates a bare ``Session`` node; - ``SessionHandler`` enriches it with the correct type label via a - subsequent upsert. - - Only caches session_id after a successful write to ensure retry - resilience on write failure. + Safety net only: ``SessionHandler`` is the sole authority on session + type labels (``RootSession``, ``SubSession``, ``ForkedSession``); this + always creates a bare ``Session`` node for later enrichment. Caches + ``session_id`` only after a successful write, for retry resilience. """ - # Tier 1: fast path — warm cache hit if session_id in self._seen_sessions: return - # Tier 2: graph query — check durable state existing = await self.graph.get_node(session_id) if existing is not None: - # Node already in graph. Also upsert a bare stub to this worker's buffer - # so that the current worker's flush uses MERGE (idempotent) rather than - # creating a second node. This prevents the asyncio race condition where: - # 1. Worker A flushes a stub node — tx is in-flight. - # 2. Worker B calls get_node — falls through to Neo4j, finds the node. - # 3. Without this upsert, Worker B's _node_buffer stays empty. - # 4. Worker B's flush later issues a fresh MERGE → duplicate node. - # upsert_node uses union-merge for labels, so existing type labels - # (e.g. "RootSession") are preserved — this call never strips labels. + # Upsert a stub so this worker's own flush uses MERGE (idempotent) + # instead of racing a second worker into creating a duplicate node. await self.graph.upsert_node( session_id, {"labels": ["Session"], "status": "running", "session_id": session_id}, @@ -292,18 +266,10 @@ async def ensure_session_node(self, session_id: str, data: dict[str, Any]) -> No self._seen_sessions.add(session_id) return - # Node absent from both cache and graph — create it as a bare Session node. - # ensure_session_node is a safety net; SessionHandler is the sole authority - # on session type labels (RootSession, SubSession, ForkedSession). - # - # "StubSession" marks this node as created by a reference (delegation, - # fork/start parent) BEFORE its own lifecycle events arrived. If those - # events never arrive, the node stays bare forever — indistinguishable - # by label from a node about to be enriched. StubSession makes that - # permanently-orphaned state observable. It is a plain marker, not a - # terminal label: SessionLabelStateMachine.classify() removes it the - # moment a real terminal label (RootSession/SubSession/ForkedSession/ - # IncompleteSession) is assigned via genuine lifecycle enrichment. + # "StubSession" marks a node created by reference (delegation, + # fork/start parent) before its own lifecycle events arrived, so a + # permanently-orphaned node stays observable. Removed by + # SessionLabelStateMachine.classify() once a real terminal label lands. node_data: dict[str, Any] = { "labels": ["Session", "StubSession"], "status": "running", @@ -323,28 +289,20 @@ async def ensure_session_node(self, session_id: str, data: dict[str, Any]) -> No async def touch_session(self, session_id: str, timestamp: str) -> None: """Update last_updated on the direct Session node only. - Updates exactly one node — the session named by *session_id*. There is - deliberately NO ancestor/parent_id propagation: the previous parent-chain - walk SET last_updated on the shared root :Session node for every child - event, so many independent writers contended on that one node's exclusive - lock — the source of the Neo4j deadlock that silently dropped events. - Root/session attributes (started_at/status/parent_id) are written once at - session:start by SessionHandler, and staleness reaping uses - worker.last_event_time — neither depends on ancestor last_updated — so - dropping propagation costs nothing while removing the contention hot spot. - - Skips the write when the stored last_updated is already at or ahead of - *timestamp*. Never raises — errors are logged at WARNING level. + No ancestor/parent_id propagation -- staleness reaping uses + worker.last_event_time, which doesn't depend on it, so propagating + would only add lock contention on the shared root node. + + Skips the write when stored last_updated is already at or ahead of + *timestamp*. Never raises -- errors are logged at WARNING level. """ try: node = await self.graph.get_node(session_id) if node is None: return current = node.get("last_updated") - # Compare using stdlib datetime only; the store's read path normalises - # driver DateTime objects to Python datetime, but the in-memory store returns - # whatever was written (often a str), so coerce both sides defensively. - # No driver-specific datetime types here. + # Coerce both sides defensively: the in-memory store may return + # a str; only stdlib datetime is compared here. ts = ( datetime.fromisoformat(timestamp) if isinstance(timestamp, str) diff --git a/context_intelligence_server/status.py b/context_intelligence_server/status.py index 0ae0049f..19f8fbc7 100644 --- a/context_intelligence_server/status.py +++ b/context_intelligence_server/status.py @@ -61,6 +61,87 @@ def recent(self) -> list[EventRecord]: ring_buffer: EventRingBuffer = EventRingBuffer() +# --------------------------------------------------------------------------- +# BootState +# --------------------------------------------------------------------------- + +# `sweep`/`topup` are momentary step labels only -- the sweep loop runs +# forever once started, so `_boot_reconcile` sets phase="ready" right after. +# `schema` is the new first phase (Neo4j schema init); `awaiting_schema` is +# the terminal-but-not-ready state when Neo4j stayed unreachable this pass -- +# the periodic sweep retries schema + topup and marks "ready" once it lands. +_BOOT_PHASES = ( + "recovering", + "schema", + "heal", + "reclaim", + "expire", # dead-letter expiry, before reconcile + "reconcile", + "seed", + "topup", + "sweep", + "awaiting_schema", + "ready", + "failed", +) + + +@dataclasses.dataclass +class BootState: + """Boot-safety progress, surfaced (additively) on /status. + + Module-level singleton; all mutation happens on the event loop inside + ``_boot_reconcile`` between awaits, so plain ints need no lock. + + ``phase`` defaults to ``"recovering"``, never ``"ready"``, so a bare ASGI + test client that never runs the real lifespan still gets a true value. + ``status`` stays ``"ok"``/200 at every phase including ``"failed"`` -- + the boot phase is informational, never a liveness signal. + """ + + phase: str = "recovering" + started_at: float = 0.0 + completed_at: float | None = None + reclaimed: int = 0 + reclaimed_bytes: int = 0 + kept: int = 0 + failed: int = 0 + resumed: int = 0 + deferred: int = 0 + error: str | None = None + failed_step: str | None = None + fallback_workspace_byte0: int = 0 + fallback_workspace_sentinel: int = 0 + reclaim_enabled: bool = False + + def begin(self) -> None: + """Mark the start of boot reconciliation (called once, at boot).""" + self.phase = "recovering" + self.started_at = time.time() + self.completed_at = None + self.error = None + self.failed_step = None + + def finish(self) -> None: + """Mark boot reconciliation as complete: phase -> "ready".""" + self.phase = "ready" + self.completed_at = time.time() + + def fail(self, step: str, exc: BaseException) -> None: + """Mark boot reconciliation as FAILED. The server keeps serving.""" + self.phase = "failed" + self.completed_at = time.time() + self.failed_step = step + self.error = f"{type(exc).__name__}: {exc}" + + def snapshot(self) -> dict[str, Any]: + """Read-only view for /status. Plain dict, no I/O.""" + return dataclasses.asdict(self) + + +boot_state: BootState = BootState() + + # --------------------------------------------------------------------------- # error_count_last_hour # --------------------------------------------------------------------------- @@ -85,26 +166,10 @@ def build_status_response( ) -> dict[str, Any]: """Build a status response dict from registry state and recent events. - Args: - registry: The active SessionRegistry. - start_time: Server start time as a Unix timestamp (from time.time()). - - Returns: - A dict with keys: status, uptime_seconds, active_sessions, sessions, - recent_events, completed_sessions, error_count_last_hour, server_version, - orphaned_sessions. - - Each entry in ``sessions`` includes the keys: session_id, workspace, - last_event, last_event_time, events_processed, orphaned, - last_successful_flush. - - Note: ``orphaned_sessions`` is the count of ALL registered workers whose - drain task has completed (``task.done()``). A worker filtered *out* of - the visible ``sessions`` list by ``status_inactive_timeout`` still - contributes to this count but will not appear with ``orphaned: True`` in - any per-session dict. For a fresh OOM orphan this asymmetry is - irrelevant (OOM orphans are recent by definition); it can surface for - long-running orphans whose ``last_event_time`` ages past the timeout. + ``orphaned_sessions`` counts ALL workers whose drain task is done, even + ones filtered out of the visible ``sessions`` list by + ``status_inactive_timeout`` -- so the count and the per-session + ``orphaned`` flags can disagree for long-idle orphans. """ settings = get_settings() now = time.time() diff --git a/context_intelligence_server/writer_lease.py b/context_intelligence_server/writer_lease.py new file mode 100644 index 00000000..f6110897 --- /dev/null +++ b/context_intelligence_server/writer_lease.py @@ -0,0 +1,582 @@ +"""Writer-lease DETECTOR -- not a mutex. + +Durable append-log framing is correct only while exactly one process writes +the queue directory (per-key serialization is in-process, ``queue_manager.py``). +A rolling/blue-green overlap on a shared mount can silently violate that. + +``enforce`` (default): refuses boot against a LIVE foreign lease, takes over a +STALE one. ``detect`` (opt-in): best-effort acquire + heartbeat, latches a +conflict on ``/status``, never refuses to boot. ``off``: disabled. + +Honest limits: staleness tolerance is heartbeat * multiplier, not a mutex; a +share fault degrades to "not armed" rather than crash-looping; never +constructs the queue directory (pure path read); all I/O runs on a private +single-thread executor so a hung mount leaks at most one thread; clean +shutdown releases (bounded, best-effort) -- the staleness window is the +backstop only when release itself fails or is skipped (e.g. a crash). +""" + +from __future__ import annotations + +import asyncio +import concurrent.futures +import dataclasses +import json +import logging +import os +import socket +import time +import uuid +from collections.abc import Callable +from pathlib import Path +from typing import Any, Literal, Protocol + +from context_intelligence_server.status import SERVER_VERSION + +logger = logging.getLogger("context_intelligence_server") + + +class WriterLeaseSettings(Protocol): + """The six fields `WriterLease.acquire` reads. + + Structural (not nominal) on purpose: `acquire()` needs nothing else from + a settings object, and the real `Settings` model (config.py) satisfies + this Protocol by construction. Tests may pass any duck-typed stub + carrying just these six fields without needing the full pydantic + model.""" + + writer_lease_mode: Literal["off", "detect", "enforce"] + writer_lease_heartbeat_seconds: float + writer_lease_staleness_multiplier: float + writer_lease_confirm_delay_seconds: float + writer_lease_acquire_timeout_seconds: float + writer_lease_force_acquire: bool + + +LEASE_FILENAME = ".writer.lease" +LEASE_TMP_FILENAME = ".writer.lease.tmp" +_LEASE_VERSION = 1 + +# Private, single-thread executor: all lease I/O runs here, never on the +# shared default pool the append/commit path and `spool_stats` also use. +_LEASE_IO = concurrent.futures.ThreadPoolExecutor( + max_workers=1, thread_name_prefix="writer-lease-io" +) + + +def shutdown_lease_io() -> None: + """Shut the private lease-I/O executor down (lifespan's `finally`). + + ``wait=False`` is deliberate: waiting would hang shutdown on exactly the + hung mount this design bounds. ``cancel_futures=True`` + drops anything still queued (there is at most one slot, so this is at + most one item). + """ + _LEASE_IO.shutdown(wait=False, cancel_futures=True) + + +class WriterLeaseConflict(RuntimeError): + """Raised ONLY when an `enforce`-mode boot must refuse. The one intended + abort -- see `_acquire_once`'s `refuse` parameter, which is the ONLY + place either raise site is reachable from.""" + + +class WriterLeaseBusy(RuntimeError): + """Raised by `_io()` when a previous lease op has not yet completed (the + one-slot in-flight gate is closed), or when submitting a new op to the + private executor itself failed. Handled identically to a filesystem + fault by every caller -- never a conflict, never wedges anything.""" + + +@dataclasses.dataclass +class LeaseRecord: + """Parsed view of one on-disk `.writer.lease` line. + + `unreadable=True` marks a synthetic record standing in for a torn or + hand-mangled lease (JSONDecodeError / missing key / wrong type / unknown + `lease_version`) -- treated at fresh-foreign strength, never at face + value.""" + + owner: str + host: str + pid: int + started_at: float + heartbeat: float + revision: str | None + server_version: str + lease_version: int + unreadable: bool = False + + +def _now() -> float: + return time.time() + + +class WriterLease: + """One process's writer-lease detector. Module-level singleton below. + + `__init__` performs NO I/O and reads NO settings -- only + `uuid.uuid4()` / `os.getpid()` / `socket.gethostname()` / `time.time()`, + so importing this module (and therefore `main`) stays exactly as cheap + and side-effect-free as it is today. Every settings-derived field is + `None` until `acquire()` assigns it, which it does as the FIRST thing it + does, before any fallible step -- so a prelude + death can never leave the object half-built. + """ + + def __init__(self) -> None: + # Identity -- literals only, no I/O. + self.owner: str = uuid.uuid4().hex + self.host: str = socket.gethostname() + self.pid: int = os.getpid() + self.started_at: float = time.time() + + # Settings-derived state -- None until acquire()'s I/O-free prelude. + self.mode: str | None = None + self.heartbeat_seconds: float | None = None + self.staleness_seconds: float | None = None + self.force_acquire: bool = False + self._confirm_delay: float | None = None + self._acquire_timeout: float | None = None + self._dir_source: Callable[[], Path] | None = None + + self._dir: Path | None = None + self._path: Path | None = None + + # Observable state. + self.acquired: bool = False + self.ever_acquired: bool = False + self.conflict: bool = False + self.conflict_source: str | None = None # "boot" | "reacquire" | "runtime" + self.observed_owner: str | None = None + self.observed_at: float | None = None + self.took_over_stale: bool = False + self.superseded_owner: str | None = None + self.superseded_age_seconds: float | None = None + self.error: str | None = None + self.last_renewed: float | None = None + + # The one-slot in-flight gate. + self._io_inflight: bool = False + + @property + def path(self) -> Path: + assert self._path is not None + return self._path + + # ----------------------------------------------------------------- + # Sync I/O primitives -- run ONLY via `_io()`, on the private executor. + # ----------------------------------------------------------------- + + def _read(self) -> LeaseRecord | None: + assert self._path is not None + try: + text = self._path.read_text(encoding="utf-8") + except FileNotFoundError: + # A missing lease means "free directory", not a share fault. + return None + try: + data = json.loads(text.strip()) + return LeaseRecord( + owner=str(data["owner"]), + host=str(data.get("host", "")), + pid=int(data.get("pid", 0)), + started_at=float(data.get("started_at", 0.0)), + heartbeat=float(data["heartbeat"]), + revision=data.get("revision"), + server_version=str(data.get("server_version", "")), + lease_version=int(data.get("lease_version", -1)), + ) + except (json.JSONDecodeError, KeyError, TypeError, ValueError): + # Torn/malformed lease is treated as fresh-and-foreign, same + # strength as a genuine live peer. + return LeaseRecord( + owner="", + host="", + pid=0, + started_at=0.0, + heartbeat=0.0, + revision=None, + server_version="", + lease_version=-1, + unreadable=True, + ) + + def _write(self, heartbeat: float) -> None: + assert self._dir is not None + assert self._path is not None + record = { + "lease_version": _LEASE_VERSION, + "owner": self.owner, + "host": self.host, + "pid": self.pid, + "started_at": self.started_at, + "heartbeat": heartbeat, + "revision": os.environ.get("CONTAINER_APP_REVISION"), + "server_version": SERVER_VERSION, + } + tmp = self._dir / LEASE_TMP_FILENAME + tmp.write_text( + json.dumps(record, separators=(",", ":")) + "\n", encoding="utf-8" + ) + os.replace(tmp, self._path) + + def _unlink_if_owned(self) -> None: + """Best-effort, owner-gated unlink -- release()'s sync body. + + Never unlinks a foreign lease: if a peer stole it, deleting theirs + would actively hand the directory to a third process.""" + if self._path is None: + return + rec = self._read() + if rec is not None and not rec.unreadable and rec.owner == self.owner: + try: + self._path.unlink() + except FileNotFoundError: + pass + + # ----------------------------------------------------------------- + # The dedicated single-thread executor + one-slot in-flight gate. + # ----------------------------------------------------------------- + + async def _io(self, fn: Callable[[], Any]) -> Any: + """Run ONE blocking lease op on the private single-thread executor. + + The gate flips False->True synchronously on the event-loop thread + with no await between the check and the set, so two concurrent + submissions are unrepresentable. It is cleared by the FUTURE'S + done-callback -- NOT by the awaiting coroutine -- so a `wait_for` + cancellation leaves the gate CLOSED until the syscall actually + returns; that is what makes the <=1-leaked-thread bound exact. + + If `submit` itself raises (a dead/shutdown executor), no future is + ever created and the done-callback would never fire -- so THIS is + the one path that clears the gate directly, converting the failure + into `WriterLeaseBusy` rather than a permanent silent disarm. + """ + if self._io_inflight: + raise WriterLeaseBusy("lease I/O still in flight (mount not responding)") + self._io_inflight = True + try: + fut = _LEASE_IO.submit(fn) + except Exception as exc: + self._io_inflight = False + raise WriterLeaseBusy(f"failed to submit lease I/O: {exc!r}") from exc + fut.add_done_callback(lambda _f: setattr(self, "_io_inflight", False)) + return await asyncio.wrap_future(fut) + + # ----------------------------------------------------------------- + # Boot acquisition + # ----------------------------------------------------------------- + + async def acquire( + self, settings: WriterLeaseSettings, dir_source: Callable[[], Path] + ) -> None: + """Acquire the lease at boot. Raises `WriterLeaseConflict` ONLY in + `enforce` mode against a fresh foreign (or unreadable) lease, or on + losing the confirm-handshake race in `enforce` mode. Every other + fault -- OSError, a hung mount, a busy gate -- is absorbed and + surfaced via `error`/`conflict`, never raised. + """ + # I/O-free prelude: attribute reads only. `_dir_source` assigned + # first so a later prelude failure still leaves a real re-arm source. + self._dir_source = dir_source + self.mode = settings.writer_lease_mode + self.heartbeat_seconds = settings.writer_lease_heartbeat_seconds + self.staleness_seconds = ( + self.heartbeat_seconds * settings.writer_lease_staleness_multiplier + ) + self.force_acquire = settings.writer_lease_force_acquire + self._confirm_delay = settings.writer_lease_confirm_delay_seconds + self._acquire_timeout = settings.writer_lease_acquire_timeout_seconds + + if self.force_acquire: + # Log on every boot while set so it's never silently forgotten. + logger.warning( + "writer_lease: FORCE_ACQUIRE IS ENABLED -- the boot refusal is " + "disabled for this process. Unset " + "AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_WRITER_LEASE_FORCE_ACQUIRE." + ) + + if self.mode == "off": + logger.info("writer_lease: mode=off -- detector disabled") + return + + refuse = self.mode == "enforce" and not self.force_acquire + await self._acquire_once(refuse=refuse, source="boot") + + async def _acquire_once(self, *, refuse: bool, source: str) -> None: + """One bounded acquire attempt. Never raises except + `WriterLeaseConflict` when `refuse=True` (boot's `enforce` path only; + `tick()`'s re-arm always passes `refuse=False`).""" + assert self._acquire_timeout is not None + try: + await asyncio.wait_for( + self._acquire_once_inner(refuse=refuse, source=source), + timeout=self._acquire_timeout, + ) + except WriterLeaseConflict: + raise + except TimeoutError: + self._apply_fault_policy( + f"acquire timed out after {self._acquire_timeout}s" + ) + except (OSError, WriterLeaseBusy) as exc: + self._apply_fault_policy(repr(exc)) + + def _apply_fault_policy(self, error_repr: str) -> None: + """A share fault or busy gate is EVIDENCE OF NOTHING -- never a + conflict. Continue boot; the detector is simply unarmed for this + process until the next tick's retry.""" + logger.error( + "writer_lease: acquire failed on a filesystem error -- the " + "writer-lease detector is NOT ARMED for this process: %s", + error_repr, + ) + self.acquired = False + self.error = error_repr + # self.conflict is deliberately left UNTOUCHED here. + + async def _acquire_once_inner(self, *, refuse: bool, source: str) -> None: + assert self._dir_source is not None + assert self.staleness_seconds is not None + assert self._confirm_delay is not None + + # Pure path read, zero syscalls -- this detector constructs nothing. + self._dir = self._dir_source() + self._path = self._dir / LEASE_FILENAME + + rec = await self._io(self._read) + if rec is not None and rec.owner != self.owner: + age = 0.0 if rec.unreadable else (_now() - rec.heartbeat) + if age < self.staleness_seconds: + # Fresh (or unreadable, or future-dated) foreign lease. + if refuse: + msg = self._refusal_message(rec, age) + logger.error("writer_lease_refused_boot %s", msg) + raise WriterLeaseConflict(msg) + self._latch_conflict(source, rec.owner) + logger.error( + "writer_lease_conflict at boot: taking over a FRESH foreign " + "lease owner=%s host=%s pid=%s revision=%s age=%.1fs -- TWO " + "WRITERS ARE SHARING THIS DIRECTORY -- concurrent writers " + "corrupt the append log", + rec.owner, + rec.host, + rec.pid, + rec.revision, + age, + ) + else: + self.took_over_stale = True + self.superseded_owner = rec.owner + self.superseded_age_seconds = age + logger.warning( + "writer_lease: took over a STALE lease owner=%s age=%.1fs", + rec.owner, + age, + ) + + heartbeat = _now() + await self._io(lambda: self._write(heartbeat)) + await asyncio.sleep(self._confirm_delay) + rec2 = await self._io(self._read) + if rec2 is None or rec2.owner != self.owner: + if refuse: + msg = f"lost the acquire race to owner={rec2.owner if rec2 else None}" + logger.error("writer_lease_refused_boot %s", msg) + raise WriterLeaseConflict(msg) + self._latch_conflict(source, rec2.owner if rec2 else None) + self.acquired = False + logger.error( + "writer_lease_conflict: lost the acquire race to owner=%s", + rec2.owner if rec2 else None, + ) + return + + self.acquired = True + self.ever_acquired = True + self.last_renewed = heartbeat + logger.info("writer lease acquired owner=%s", self.owner) + + def _latch_conflict(self, source: str, observed_owner: str | None) -> None: + self.conflict = True + self.observed_owner = observed_owner + self.observed_at = _now() + # Upgrade ladder: boot/reacquire -> runtime, never back. + if self.conflict_source != "runtime": + self.conflict_source = source + + def _refusal_message(self, rec: LeaseRecord, age: float) -> str: + assert self._dir is not None + assert self.staleness_seconds is not None + return ( + "Refusing to boot: another writer holds the queue-directory lease.\n" + f" dir = {self._dir}\n" + f" foreign owner = {rec.owner} (host={rec.host} pid={rec.pid} " + f"revision={rec.revision} version={rec.server_version})\n" + f" lease age = {age:.1f}s (stale after " + f"{self.staleness_seconds:.1f}s)\n" + "This server serializes durable appends IN-PROCESS; two processes " + "writing the same queue directory corrupts it. Wait " + "for the previous revision to drain and exit, or set " + "AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_WRITER_LEASE_FORCE_ACQUIRE=" + "true for ONE boot if you are certain the previous writer is gone." + ) + + # ----------------------------------------------------------------- + # Heartbeat / runtime conflict detection + # ----------------------------------------------------------------- + + async def tick(self) -> None: + """ONE observation + at most one renewal. Never raises (this is what + lets `heartbeat_loop` stay a thin, supervised wrapper).""" + if self.mode is None or self.mode == "off": + return + + if not self.acquired: + if self.ever_acquired: + # Held-then-lost: keep reading, but never write again -- + # renewing would ping-pong the lease with the peer. + await self._observe_only() + return + # Never-acquired: nothing to protect, so try again this tick. + await self._acquire_once(refuse=False, source="reacquire") + return + + await self._renew_once() + + async def _renew_once(self) -> None: + """The 'we hold it' branch -- also bounded by the acquire timeout + so a hung mount gives the heartbeat loop ONE + bounded WARNING per tick and re-arms, rather than wedging forever + inside `tick()`.""" + assert self._acquire_timeout is not None + try: + await asyncio.wait_for( + self._renew_once_inner(), timeout=self._acquire_timeout + ) + except TimeoutError: + logger.warning( + "writer_lease: tick timed out after %.1fs -- will retry", + self._acquire_timeout, + ) + except (OSError, WriterLeaseBusy) as exc: + logger.warning("writer_lease: tick failed (%s), will retry", exc) + + async def _renew_once_inner(self) -> None: + rec = await self._io(self._read) + if rec is None or rec.owner != self.owner: + self.conflict = True + self.conflict_source = "runtime" # unconditional upgrade + self.observed_owner = rec.owner if rec else None + self.observed_at = _now() + self.acquired = False # we LOST it; ever_acquired stays True + logger.error( + "writer_lease_conflict: lease taken by owner=%s", + self.observed_owner, + ) + return + heartbeat = _now() + await self._io(lambda: self._write(heartbeat)) + self.last_renewed = heartbeat + + async def _observe_only(self) -> None: + """Held-then-lost: read-only, best-effort, bounded. Never writes.""" + assert self._acquire_timeout is not None + try: + rec = await asyncio.wait_for( + self._io(self._read), timeout=self._acquire_timeout + ) + except (OSError, WriterLeaseBusy, TimeoutError): + return + if rec is not None: + self.observed_owner = rec.owner + self.observed_at = _now() + + async def heartbeat_loop(self) -> None: + """Sleep -> tick -> forever, supervised. + + Interval is read once before the loop; an unset value (prelude never + ran) is a loud single-shot return rather than an uncapped busy-loop.""" + interval = self.heartbeat_seconds + if not interval or interval <= 0: + logger.error( + "writer_lease: heartbeat loop NOT STARTED -- heartbeat_seconds " + "is unset (acquire() never completed its prelude). The " + "writer-lease detector is NOT ARMED for this process." + ) + self.error = "heartbeat loop not started: heartbeat_seconds unset" + return + while True: + try: + await asyncio.sleep(interval) + await self.tick() + except asyncio.CancelledError: + raise + except Exception as exc: + logger.warning( + "writer_lease_heartbeat: tick failed, will retry: %s", + exc, + exc_info=True, + ) + + # ----------------------------------------------------------------- + # Shutdown + # ----------------------------------------------------------------- + + async def release(self) -> None: + """Owner-gated, best-effort release. A failed release is not a + failed shutdown -- the next boot just waits out the staleness window. + Bounded by the acquire timeout so a hung mount can never block + shutdown.""" + if self.mode is None or self.mode == "off" or self._path is None: + return + timeout = self._acquire_timeout if self._acquire_timeout is not None else 5.0 + try: + await asyncio.wait_for(self._io(self._unlink_if_owned), timeout=timeout) + except TimeoutError: + logger.warning("writer_lease: release timed out after %.1fs", timeout) + except (OSError, WriterLeaseBusy) as exc: + logger.warning("writer_lease: release failed (best-effort): %s", exc) + + def mark_unarmed(self, error_repr: str, mode: str | None = None) -> None: + """Called when a fault escapes `acquire()`'s own fault policy. + `conflict` is left untouched.""" + self.acquired = False + self.error = error_repr + if self.mode is None and mode is not None: + self.mode = mode + + # ----------------------------------------------------------------- + # /status + # ----------------------------------------------------------------- + + def snapshot(self) -> dict[str, Any]: + """Pure in-memory dict build -- no I/O, so it cannot raise on the + unauthenticated health path.""" + last_renewed = self.last_renewed + lease_age = None if last_renewed is None else max(0.0, _now() - last_renewed) + return { + "mode": self.mode, + "acquired": self.acquired, + "owner": self.owner, + "conflict": self.conflict, + "conflict_source": self.conflict_source, + "observed_owner": self.observed_owner, + "observed_at": self.observed_at, + "took_over_stale": self.took_over_stale, + "superseded_owner": self.superseded_owner, + "superseded_age_seconds": self.superseded_age_seconds, + "force_acquire": self.force_acquire, + "error": self.error, + "last_renewed": last_renewed, + "lease_age_seconds": lease_age, + "heartbeat_seconds": self.heartbeat_seconds, + "staleness_seconds": self.staleness_seconds, + } + + +# Module singleton. No I/O at import, so a bare ASGI test client that never +# runs the real lifespan still gets a coherent, non-lying /status. +writer_lease: WriterLease = WriterLease() diff --git a/docs/architecture/05-durable-ingest-queue.dot b/docs/architecture/05-durable-ingest-queue.dot index 96be9338..e322dc87 100644 --- a/docs/architecture/05-durable-ingest-queue.dot +++ b/docs/architecture/05-durable-ingest-queue.dot @@ -24,7 +24,7 @@ digraph durable_ingest_queue { ] // ============================================================ - // PHASE 1 \u2014 HTTP Layer (synchronous, returns fast) + // PHASE 1 -- HTTP Layer (synchronous, returns fast) // ============================================================ subgraph cluster_http { graph [ @@ -43,12 +43,12 @@ digraph durable_ingest_queue { ] BearerMiddleware [ - label = "BearerTokenMiddleware\n(ASGI middleware \u2014 see diagram 06 for full auth flow)\nverifies bearer token; stamps contributor_id into scope[\"state\"]\nexempt paths bypass auth (/status, /version, etc.)" + label = "BearerTokenMiddleware\n(ASGI middleware -- see diagram 06 for full auth flow)\nverifies bearer token; stamps contributor_id into scope[\"state\"]\nexempt paths bypass auth (/status, /version, etc.)" fillcolor = "#B2EBF2" ] Auth401Exit [ - label = "HTTP 401 / 403\nauth_event=auth_denied\n(or resolver_unexpected_exception \u2192 401)" + label = "HTTP 401 / 403\nauth_event=auth_denied\n(or resolver_unexpected_exception -> 401)" shape = oval fillcolor = "#FFCCBC" ] @@ -60,7 +60,7 @@ digraph durable_ingest_queue { ] ValidateTimestamp [ - label = "_validate_data_timestamp(data)\ndata[\"timestamp\"] must be present +\nnon-empty string + valid ISO-8601\n(fail loud at boundary \u2014 never silently dead-letter)" + label = "_validate_data_timestamp(data)\ndata[\"timestamp\"] must be present +\nnon-empty string + valid ISO-8601\n(fail loud at boundary -- never silently dead-letter)" fillcolor = "#F5F5F5" ] @@ -71,7 +71,7 @@ digraph durable_ingest_queue { ] IdempotencyHit [ - label = "Idempotency hit?\n(idempotency_key seen before)" + label = "idempotency_cache.seen(key)?\n(key is recorded ONLY after a durable append,\nso a hit proves the event is on disk)" shape = diamond fillcolor = "#FFF9C4" style = filled @@ -99,7 +99,20 @@ digraph durable_ingest_queue { ] QueueAppend [ - label = "QueueManager.append(worker_key, stamped_body)\nbody[\"created_by\"] = contributor_id (write-once, prevents spoofing)\nPERSIST-FIRST: append line to .log\n(zero-silent-loss window)" + label = "QueueManager.append(worker_key, stamped_body)\nbody[\"created_by\"] = contributor_id (write-once, prevents spoofing)\nPERSIST-FIRST: append line to .log\nSINGLE-WRITER, ATOMIC FRAMING: per-key file lock held by the\nWRITING THREAD (not the coroutine) for the whole record \--\ncancellation cannot release it mid-write\nPARTIAL WRITE \-> truncated away + OSError raised to the caller\n(never left as a torn/merged line)" + fillcolor = "#C8E6C9" + fontname = "Helvetica-Bold" + penwidth = 2.5 + ] + + AppendFailed [ + label = "append failed\npartial bytes DISCARDED, error surfaced\n(no idempotency key stored \-> a client retry\nis honoured, not falsely refused)" + shape = oval + fillcolor = "#FFCCBC" + ] + + IdempotencyStore [ + label = "idempotency_cache.store(key)\nAFTER the durable append succeeds \-- ordering is\nload-bearing: storing first would answer a retry\n\\\"duplicate\\\" for an event that is nowhere on disk" fillcolor = "#C8E6C9" fontname = "Helvetica-Bold" penwidth = 2.5 @@ -118,11 +131,11 @@ digraph durable_ingest_queue { } // ============================================================ - // PHASE 2 \u2014 Per-session drainer (single sticky async task) + // PHASE 2 -- Per-session drainer (single sticky async task) // ============================================================ subgraph cluster_drainer { graph [ - label = "Per-session drainer (registry.drain_worker)\n\u2014 single sticky async task per worker_key" + label = "Per-session drainer (registry.drain_worker)\n-- single sticky async task per worker_key" style = "filled" color = "#2E7D32" fillcolor = "#F1F8E9" @@ -136,12 +149,12 @@ digraph durable_ingest_queue { ] ProcessBatch [ - label = "process batch \u2192 build graph writes\n(handlers/enrichers \u2014 see diagram 01)" + label = "process batch -> build graph writes\n(handlers/enrichers -- see diagram 01)" fillcolor = "#F5F5F5" ] FlushBarrier [ - label = "_flush_barrier()\nasync with write_semaphore (cap=write_concurrency)\ngraph.flush() \u2192 Neo4j" + label = "_flush_barrier()\nasync with write_semaphore (cap=write_concurrency)\ngraph.flush() -> Neo4j" fillcolor = "#C8E6C9" fontname = "Helvetica-Bold" penwidth = 2.5 @@ -177,9 +190,16 @@ digraph durable_ingest_queue { ] DeadLetterOut [ - label = "\u2192 .dead.jsonl" + label = "POISON -> DEAD-LETTER -> CONTINUE\nthe unparseable/poison line goes to .dead.jsonl\nand the offset advances PAST it -- draining resumes\nat the next line; one bad record never halts the drain" fillcolor = "#F5F5F5" ] + + DrainSupervisor [ + label = "drain-task done-callback (the ONE supervision point)\ntask died with an exception?\n-> log 'drain_worker_died' at ERROR (session id + traceback)\n-> deregister the session, then close its store\nso the session is RECOVERABLE, never a silent orphan" + fillcolor = "#FFE0B2" + fontname = "Helvetica-Bold" + penwidth = 2.0 + ] } // ============================================================ @@ -220,12 +240,41 @@ digraph durable_ingest_queue { FileLog -> FileOffset -> FileDeadLetter [style = invis] } + // ============================================================ + // SELF-SHRINKING STORAGE + // the queue is a TRANSIENT BUFFER, not an archive + // ============================================================ + subgraph cluster_shrink { + graph [ + label = "Self-shrinking storage\n(the queue is a TRANSIENT BUFFER, not an archive)" + style = "filled" + color = "#00838F" + fillcolor = "#E0F7FA" + fontname = "Helvetica" + fontsize = 10 + ] + + Compaction [ + label = "compact_committed_prefix() (queue_compact_enabled, default ON)\nCONTINUOUS reclaim of a LIVE session's already-committed\nprefix [0, C) -- not just at session end. The .log then holds\nthe UNDRAINED TAIL, not the whole session history.\nfrequency bounded by queue_compact_min_prefix_bytes\nreclaims regardless of tail size\nrebase .offset to 0 FIRST, then swap the tail in\n(every crash window degrades to a bounded re-drive, never a loss)" + fillcolor = "#80DEEA" + fontname = "Helvetica-Bold" + penwidth = 2.0 + ] + + DeadLetterExpiry [ + label = "expire_dead_letters() (dead_letter_expiry_enabled, opt-in, default OFF)\nLOG-LESS ONLY: no .log exists (a structural proof that no\nboot mechanism can ever consult the file being expired)\nAND mtime older than dead_letter_retention_seconds (30d default)\nruns at boot (phase=expire) and on every sweep tick" + fillcolor = "#80DEEA" + ] + + Compaction -> DeadLetterExpiry [style = invis] + } + // ============================================================ // CRASH RECOVERY (startup, main.py lifespan) // ============================================================ subgraph cluster_recovery { graph [ - label = "Crash recovery (startup, main.py lifespan)" + label = "Phased boot reconciliation (main.py lifespan -> _boot_reconcile)\nspawned as a supervised BACKGROUND task, NOT awaited --\n/status and /version answer from the FIRST phase" style = "filled" color = "#6A1B9A" fillcolor = "#F3E5F5" @@ -233,19 +282,58 @@ digraph durable_ingest_queue { fontsize = 10 ] - ReconcileDead [ - label = "recovery_reconcile_dead()\n(advance past already-dead-lettered pending lines\n\u2014 closes the dead_letter\u2192commit crash window)" + PhaseRecovering [ + label = "phase = recovering\nboot_state.begin()\n(the DEFAULT phase from module import --\nnever 'ready' until proven)" fillcolor = "#E1BEE7" ] - SeedCounts [ - label = "recovery_seed_counts()\nseed accepted/written from disk\nso residual==0 at startup" + PhaseHeal [ + label = "phase = heal\nheal_torn_tails()\nquarantine any torn tail left by a hard kill" fillcolor = "#E1BEE7" ] - RecoverSpawn [ - label = "recover()\nrespawn a drainer per session\nwith unprocessed lines" + PhaseReclaim [ + label = "phase = reclaim\nCLASSIFY every pre-existing key:\nresumable | drained | unresumable | reset_offset\ndrained -> reclaimed automatically, always\nunresumable/reset_offset -> gated on reclaim_enabled\n(default False = dry-run: classify + audit-log only)\nunparseable offset -> always re-drain from byte 0, any size\nnegative/past-EOF offset -> re-drain below reclaim_redrain_max_bytes,\nelse delete" fillcolor = "#E1BEE7" + fontname = "Helvetica-Bold" + penwidth = 2.0 + ] + + PhaseExpire [ + label = "phase = expire\nexpire_dead_letters()\n(log-less + past dead_letter_retention_seconds)" + fillcolor = "#E1BEE7" + ] + + PhaseReconcile [ + label = "phase = reconcile\nrecovery_reconcile_dead()\n(advance past already-dead-lettered pending lines\n-- closes the dead_letter->commit crash window)" + fillcolor = "#E1BEE7" + ] + + PhaseSeed [ + label = "phase = seed\nrecovery_seed_counts()\nseed accepted/written from disk\nso residual==0 at startup" + fillcolor = "#E1BEE7" + ] + + PhaseTopup [ + label = "phase = topup\nrespawn a drainer per recovered session,\nCAPPED at crash_recovery_respawn_limit\ndeferred tail: untouched on disk, still recoverable" + fillcolor = "#E1BEE7" + ] + + PhaseSweep [ + label = "phase = sweep\nstart the periodic top-up loop\n(crash_recovery_sweep_interval_seconds)\nadvances the deferred tail at a BOUNDED rate" + fillcolor = "#E1BEE7" + ] + + PhaseReady [ + label = "phase = ready\nboot over: /status fills in metrics + spool" + shape = oval + fillcolor = "#C8E6C9" + ] + + PhaseFailed [ + label = "phase = failed\nANY step raised -> failed_step + error recorded,\nloud traceback logged, and THE SERVER KEEPS SERVING\n(/status stays HTTP 200 / status:ok at every phase --\nthe boot phase is informational, never a liveness signal)" + shape = oval + fillcolor = "#FFCCBC" ] WorkersGuard [ @@ -254,8 +342,16 @@ digraph durable_ingest_queue { style = filled ] - ReconcileDead -> SeedCounts -> RecoverSpawn - RecoverSpawn -> WorkersGuard [style = dotted, color = "#AAAAAA", arrowhead = none] + PhaseRecovering -> PhaseHeal -> PhaseReclaim -> PhaseExpire -> PhaseReconcile + PhaseReconcile -> PhaseSeed -> PhaseTopup -> PhaseSweep -> PhaseReady + PhaseReclaim -> PhaseFailed [ + label = "any phase raises" + style = dashed + color = "#C62828" + fontcolor = "#C62828" + constraint = false + ] + PhaseTopup -> WorkersGuard [style = dotted, color = "#AAAAAA", arrowhead = none] } // ============================================================ @@ -277,17 +373,31 @@ digraph durable_ingest_queue { ] StatusEndpoint [ - label = "GET /status .metrics\n(unauthenticated, aggregate-only)" + label = "GET /status (unauthenticated, aggregate-only)\nALWAYS present: .boot {phase, failed_step, error, counters}\n .writer_lease {mode, acquired, conflict, ...}\nWHILE BOOTING: .metrics = null AND .spool = null\n .status_detail = {reason: \"booting\"}\n (ZERO disk reads on this request path)\nONCE phase is ready|failed:\n .metrics + .spool populated" fillcolor = "#B3E5FC" + fontname = "Helvetica-Bold" + penwidth = 2.0 ] QueuesEndpoint [ - label = "routers/queues.py (authenticated: Bearer)\nGET /queues/dead-letter\nPOST .../{worker_key}/replay\nPOST .../{worker_key}/purge" + label = "routers/queues.py (authenticated: Bearer)\nGET /queues/dead-letter (read)\nPOST .../{worker_key}/replay (write)\nPOST .../{worker_key}/purge (write)" fillcolor = "#B3E5FC" ] + WriterLease [ + label = "WriterLease -- SINGLE-WRITER LEASE GUARD\n(writer_lease_mode, DEFAULT 'enforce')\nbackstop against an accidental second writer, NOT a\nrolling-deploy coordinator (deployment is single-replica).\nheartbeats /.writer.lease every\nwriter_lease_heartbeat_seconds (5s); stale after\nheartbeat x writer_lease_staleness_multiplier (15s)\n'enforce': REFUSES to boot against a LIVE foreign lease;\nreleases lease on clean shutdown; takes over a STALE\nforeign lease automatically.\n'detect': latches + surfaces a conflict on /status only,\nnever refuses to boot." + fillcolor = "#FFE0B2" + fontname = "Helvetica-Bold" + penwidth = 2.0 + ] + PipelineMetrics -> StatusEndpoint PipelineMetrics -> QueuesEndpoint [style = invis] + WriterLease -> StatusEndpoint [ + label = ".writer_lease\n(pure in-memory,\nevery phase)" + color = "#E65100" + fontcolor = "#E65100" + ] } // ============================================================ @@ -304,7 +414,9 @@ digraph durable_ingest_queue { IdempotencyHit -> WorkerKey [label = "no"] WorkerKey -> GetOrCreate GetOrCreate -> QueueAppend - QueueAppend -> RecordAccepted + QueueAppend -> AppendFailed [label = "OSError\n(partial discarded)"] + QueueAppend -> IdempotencyStore [label = "durably appended"] + IdempotencyStore -> RecordAccepted RecordAccepted -> Http202 // ============================================================ @@ -320,19 +432,28 @@ digraph durable_ingest_queue { MaxAttempts -> HandleExhausted [label = "yes"] WriteRetry -> FlushBarrier [label = "retry", constraint = false] HandleExhausted -> DeadLetterOut - DeadLetterOut -> ReadBatch [label = "resume", constraint = false] + DeadLetterOut -> ReadBatch [label = "continue past\nthe poison line", constraint = false] + + // Supervision: the drain task's done-callback observes ANY death. + ReadBatch -> DrainSupervisor [ + style = dotted + color = "#EF6C00" + fontcolor = "#EF6C00" + label = "drain task raised\n(done-callback)" + constraint = false + ] - // Recovery spawns drainer(s) at startup - RecoverSpawn -> ReadBatch [ + // The boot topup spawns drainer(s) at startup + PhaseTopup -> ReadBatch [ style = dashed color = "#6A1B9A" fontcolor = "#6A1B9A" - label = "spawns drainer\nper session" + label = "spawns drainer\nper recovered session\n(capped)" constraint = false ] // ============================================================ - // DURABLE FILE edges (dashed \u2014 annotate the side-effects) + // DURABLE FILE edges (dashed -- annotate the side-effects) // ============================================================ QueueAppend -> FileLog [ style = dashed @@ -348,6 +469,31 @@ digraph durable_ingest_queue { label = "commit\n(atomic)" constraint = false ] + + // Commit is what makes the prefix reclaimable: the committed offset is the + // single value the durability design already trusts. + CommitAndRecord -> Compaction [ + style = dashed + color = "#00838F" + fontcolor = "#00838F" + label = "committed offset C\ntriggers reclaim" + constraint = false + ] + Compaction -> FileLog [ + style = dashed + color = "#00838F" + fontcolor = "#00838F" + label = "rewrite: drop [0, C),\nkeep the undrained tail\n(.offset rebased to 0 first)" + constraint = false + ] + DeadLetterExpiry -> FileDeadLetter [ + style = dashed + color = "#00838F" + fontcolor = "#00838F" + label = "expire\n(log-less + aged out)" + arrowhead = odot + constraint = false + ] HandleExhausted -> FileDeadLetter [ style = dashed color = "#795548" @@ -384,6 +530,16 @@ digraph durable_ingest_queue { constraint = false ] + // The writer lease guards the DIRECTORY the append log lives in. + WriterLease -> FileLog [ + style = dotted + color = "#E65100" + fontcolor = "#E65100" + label = "guards /.writer.lease\nagainst a SECOND writer on this directory" + arrowhead = none + constraint = false + ] + // ============================================================ // LEGEND // ============================================================ @@ -429,7 +585,15 @@ digraph durable_ingest_queue { shape = oval fillcolor = "#FFCCBC" ] - + LegShrink [ + label = "self-shrinking storage\n(compaction / expiry)" + fillcolor = "#80DEEA" + ] + LegDetect [ + label = "supervision / detection\n(not prevention)" + fillcolor = "#FFE0B2" + ] LegGreen -> LegAuth -> LegDiamond -> LegFile -> LegOval -> LegError [style = invis] + LegError -> LegShrink -> LegDetect [style = invis] } } diff --git a/docs/architecture/05-durable-ingest-queue.png b/docs/architecture/05-durable-ingest-queue.png index c9926dfd..a259d501 100644 Binary files a/docs/architecture/05-durable-ingest-queue.png and b/docs/architecture/05-durable-ingest-queue.png differ diff --git a/docs/architecture/README.md b/docs/architecture/README.md index 8c58426f..5aca94e7 100644 --- a/docs/architecture/README.md +++ b/docs/architecture/README.md @@ -33,7 +33,7 @@ resolver and exempt-path set are wired at boot time. | 02 | [02-handler-architecture.dot](./02-handler-architecture.dot) | Handler class-level architecture | | 03 | [03-graph-model.dot](./03-graph-model.dot) | Neo4j property-graph schema | | 04 | [04-default-handler-flow.dot](./04-default-handler-flow.dot) | DefaultHandler internal decision flow | -| 05 | [05-durable-ingest-queue.dot](./05-durable-ingest-queue.dot) | Durable ingest queue + drain loop (incl. auth middleware entry) | +| 05 | [05-durable-ingest-queue.dot](./05-durable-ingest-queue.dot) | Durable ingest queue + drain loop — atomic per-key append, supervised drainers, phased boot, continuous prefix reclaim + dead-letter expiry, and the single-writer lease guard (incl. auth middleware entry) | | 06 | [06-auth-flow.dot](./06-auth-flow.dot) | **Per-request auth flow** — BearerTokenMiddleware → resolver dispatch → `/admin/*` branch or `post_events` | | 07 | [07-auth-startup.dot](./07-auth-startup.dot) | **Auth boot wiring** — mode selection, JWKS prefetch, fail-closed gate, exempt-path selection | | 08 | [08-identity-map-management.dot](./08-identity-map-management.dot) | **Runtime identity-map management** — admin API, `require_admin` gate, `IdentityStore` write-then-swap, live `flat_dict`, no-redeploy proof | @@ -209,14 +209,60 @@ Requests enter via `BearerTokenMiddleware` (auth gate — see diagram 06) and ar with `401`/`403` before reaching the route handler if credentials are missing or invalid. `POST /events` then validates `data.timestamp` (→ 400 on failure), stamps `created_by` from the verified contributor id, persists the event to a durable per-session append-log, -and returns `202` immediately (persist-then-202); an async single drainer per session -processes batches and flushes them to Neo4j under a global write semaphore, retrying -transient/deadlock failures and isolating poison events to a dead-letter file. Durable files -per session are `.log` (append-only raw events — `created_by`-stamped), -`.offset` (last committed byte position), and `.dead.jsonl` (poison -records). On startup the server replays unprocessed log lines and re-seeds counters from -disk (crash recovery). Live conservation metrics surface on `/status`, and authenticated -`/queues/dead-letter` endpoints support inspect, replay, and purge. +and returns `202` immediately (persist-then-202). Durable files per session are +`.log` (append-only raw events — `created_by`-stamped), `.offset` +(last committed byte position), and `.dead.jsonl` (poison records). + +**The append is atomic and single-writer-serialised.** Each worker key has its own file +lock, held by the *writing thread* for the duration of the write — not by the awaiting +coroutine — so cancellation cannot release it mid-record. One record therefore lands as one +contiguous, newline-terminated byte range or not at all: a write that fails part-way is +discarded (the partial bytes truncated away) and the failure raised to the caller, never +left on disk as a torn or merged line that a later append could fuse with. The idempotency +key is recorded **only after** that durable append returns successfully — burning the key +first would answer a client's retry "duplicate" for an event that is nowhere on disk, with +no recovery path. + +**Drainers are supervised.** An async single drainer per session reads batches, dispatches +them through the per-event spine, and flushes to Neo4j under one global write semaphore +(including the terminal `session:end` flush, which goes through the same barrier as every +other write — the semaphore is the single Neo4j-write boundary, with no bypass). Transient +and deadlock failures retry at the same offset; a poison line that survives +`max_delivery_attempts` is written to that key's dead-letter file and the offset advanced +past it, so draining continues rather than halting on one bad record. A drain task that +dies is not silent: a done-callback logs `drain_worker_died` at `ERROR` with the session id +and traceback, then deregisters the session and closes its store, so the session can be +recovered instead of lingering as an invisible orphan. + +**Boot is phased, fast, and never crash-looping.** `/status` and `/version` answer from the +very first phase: every share-reading recovery pass runs as a single supervised background +task rather than on the critical path to first request. The phases are `recovering → heal → +reclaim → expire → reconcile → seed → topup → sweep → ready`, terminating at `failed` if any +step raises — and `failed` keeps serving, since a boot hook must never restart-loop the +process it is reconciling for. The `reclaim` phase classifies pre-existing on-disk queue data +(resume vs. delete vs. bounded re-drain of a bad offset); classification is **dry-run by +default** (`reclaim_enabled` ships `false`), emitting the same audit line it would when live +while unlinking nothing, so an irreversible delete path's first contact with real data is +never an actual delete. Respawn of recovered drainers is capped per pass, with a periodic +sweep advancing the deferred tail — reclaiming already-drained data at a bounded rate rather +than all at once. + +**Storage shrinks itself.** A live session's +already-committed prefix is rewritten away continuously (compaction), so a `.log` tracks the +undrained tail rather than the whole session history — bounded in frequency by a minimum +committed-prefix threshold; there is no cap on the rewrite itself. A fully-drained log is +reclaimed automatically at boot regardless of `reclaim_enabled`. Dead-letter files with +no `.log` beside them expire on an mtime-based retention window. All three run automatically as +part of normal operation and require no operator action. + +Observability follows the same split. `/status` carries the always-present `boot` block (phase +plus reclaim/resume counters) and `writer_lease` block (the queue-directory single-writer +**lease guard** — a backstop against an accidental second writer, not a rolling-deploy +coordinator; default mode `enforce` refuses to boot against a live foreign lease). While the server is +still booting, `/status` performs zero disk reads on that request path: `metrics` and `spool` +are present but `null`, with `status_detail.reason == "booting"`, and both populate once boot +is over. Authenticated `/queues/dead-letter` endpoints still support inspect, replay, and +purge. --- diff --git a/docs/azure-deployment.md b/docs/azure-deployment.md index ec3cb784..74a81e31 100644 --- a/docs/azure-deployment.md +++ b/docs/azure-deployment.md @@ -503,7 +503,9 @@ services: volume: # persistent /data (identity store, queues, blobs, logs) mount_path: /data - size_gib: 16 + size_gib: 1024 # 1 TiB standing size — must match the live share quota + tier: premium # Premium_LRS SSD; Standard SMB per-op latency caps the + # write-heavy queue/blob throughput (premium min is 100 GiB) env: # Force Entra auth (server DEFAULTS to static — must override) @@ -758,26 +760,54 @@ guarantee is structural, not "be careful": - **Neo4j data** lives on the VM's **persistent managed data disk**, entirely outside the manifest (see Neo4j safety, above). Untouched by any server redeploy. +### `/data/queues` is self-shrinking + +The durable ingest queue is a **transient buffer, not an archive**. A live +session's already-committed prefix is reclaimed continuously (compaction), so a +`.log` holds the undrained tail rather than the whole session history; a +fully-drained log is removed at session finalize; and a dead-letter file with no +`.log` beside it expires on an mtime-based retention window. Steady-state +`/data/queues` growth therefore tracks **undrained** data, not total ingest. + +These mechanisms run automatically and require no operator action — including +reclamation of a fully-drained log, which happens at boot regardless of +`reclaim_enabled`. + ### The ONLY things that can lose `/data` — avoid during a version bump | Destructive action | Why it loses data | |--------------------|-------------------| | Remove/rename the `volume:` block | New revision has no `/data` mount → writes to ephemeral FS | | Change `mount_path` | App reads an empty path; share persists but is "gone" from the server's view | -| Change `size_gib` **or** add/change `tier` | Resize/tier change can **re-provision a new, empty share — no data migration** | +| **Shrink** `size_gib`, or change `tier` | Shrink/tier change can **re-provision a new, empty share — no data migration** | | `amplifier-online destroy` | Tears down per-project resources including the share | | Delete/re-provision the Azure Files share or storage account out-of-band | Removes the backing store | -`size_gib: 16` is safe to leave unchanged — leaving it is the safe path. **Treat -any `size_gib`/`tier` edit as potentially destructive.** A version bump changes -**only the image tag**. +**Treat any `size_gib`/`tier` edit as potentially destructive — but distinguish +two cases, because they are not the same risk:** + +- **Reconciling a GROW against an already-raised live quota is safe, and is what + the shipped manifest does.** The live share quota was raised out-of-band to + 1 TiB; the manifest declares `size_gib: 1024` so that manifest and actual + agree. A re-apply then grows-or-holds — it never drops below live usage, and + no re-provision occurs. Leaving the manifest *behind* the live quota is the + worse option: the drift is invisible until something reconciles it downward. +- **A SHRINK, or any `tier` change, is the destructive case.** Either can + re-provision a new, empty share with no data migration. Do not do it as part + of a version bump. + +A version bump should change **only the image tag**. If the volume block differs +from what is live, reconcile it deliberately and on its own — not bundled with a +release. ### Protect-the-durable-data checklist (before a version bump) ``` [ ] Manifest diff shows ONLY the image tag change (v6.6.6 → v6.7.0). - volume block byte-identical: mount_path: /data, size_gib: 16 - (no size_gib change, no tier added, no mount_path change). + volume block byte-identical: mount_path: /data, size_gib: 1024, tier: premium + (no size_gib change, no tier change, no mount_path change). + If size_gib/tier DOES differ from live: only a GROW to match an + already-raised live quota is safe. A shrink or tier change: STOP. [ ] `amplifier-online up --dry-run` reports ONLY the image/revision change — NO volume, share, or storage change. If it mentions volume/storage → STOP. [ ] (Extra safety) Snapshot the Azure Files share first: diff --git a/docs/operational-hardening.md b/docs/operational-hardening.md index 1229e057..71c1718e 100644 --- a/docs/operational-hardening.md +++ b/docs/operational-hardening.md @@ -234,17 +234,43 @@ A second host still running the default carries this hazard today. ## 5. Keep the spool out of the boot path `queues_path` holds one durable append-log per session (`.log` / `.offset` / -`.dead.jsonl`). It is meant to be transient — drained and deleted. When -draining stalls, it is not. - -There is **no size cap and no alert**. Boot cost scales with it: crash recovery -respawns one drainer per session with an undrained line and reads from disk to -rebuild the conservation baseline, so a large spool converts every restart into -a multi-minute, memory-heavy operation. That is what turned a recoverable stall -into a two-day loop. - -If a boot is already failing on an oversized spool, move it out of the way so -the server can start, then triage it separately: +`.dead.jsonl`). It is transient by design — drained, shrunk, and deleted. When +draining stalls, it stops shrinking. + +**The spool now shrinks itself, but it still has no hard size cap.** Three +mechanisms bound it: + +- **Compaction** reclaims a *live* session's already-committed prefix + continuously, not just at session end (`queue_compact_enabled`, on by default). + A `.log` therefore tracks the undrained tail, not the whole session history. + Frequency is bounded by `queue_compact_min_prefix_bytes`; it reclaims + regardless of tail size. +- **Dead-letter expiry** removes a `.dead.jsonl` that has no `.log` beside it + once it ages past `dead_letter_retention_seconds` (30 days by default), so + poison records cannot accumulate indefinitely. + +None of that bounds **undrained** data: a stalled drainer's tail is exactly the +data nothing may touch, and it grows for as long as the stall lasts. Boot cost +still scales with the spool — crash recovery respawns drainers and reads from +disk to rebuild the conservation baseline — but that work now runs in the +background after the server is already answering `/status` and `/version`, and +respawn is capped (`crash_recovery_respawn_limit`, default 8) with a periodic +top-up sweep (`crash_recovery_sweep_interval_seconds`, default 60s). Boot no +longer crash-loops on an oversized spool: a reconciliation failure sets the boot +phase to `failed` and the server keeps serving. + +**Reclaiming space.** In normal operation there is nothing to reclaim by hand: +compaction runs continuously and needs no operator action. Dead-letter expiry +is opt-in (`dead_letter_expiry_enabled`, default off — a dead-letter may be the +only surviving copy of an un-recovered event) and, once enabled, likewise runs +continuously. If an already-drained backlog ever does need manual reclamation, +that is an out-of-band admin/maintenance operation performed inside the +container — there is no API for it. + +### Last resort — archive the whole queue directory + +Only if the server cannot be made healthy any other way. Boot no longer +crash-loops on spool size, so this is rarely the right tool. ```bash DATA_DIR="$HOME/amplifier-context-intelligence-server-data-store" @@ -268,10 +294,33 @@ the event is durably appended, so **clients see success while nothing reaches the graph**. The only organic symptom is someone noticing the graph is stale. Check these directly. +> **Check `.boot.phase` and `.writer_lease` FIRST.** `/status` answers from the +> very first boot phase, and while the server is still booting the `metrics` and +> `spool` blocks are present but **`null`**, and `status_detail` is +> `{"reason": "booting"}`. Any one-liner below that indexes into `['metrics']` or +> `['spool']` will therefore raise `TypeError` on a booting server — that is a +> *booting* signal, not a broken server. Start with: +> +> ```bash +> curl -s http://localhost:8000/status | python3 -c "import json,sys; r=json.load(sys.stdin); print('boot:', r['boot']['phase'], r['boot'].get('failed_step'), r['boot'].get('error')); print('lease:', r['writer_lease']['mode'], 'conflict=', r['writer_lease']['conflict'])" +> ``` +> +> `phase` walks `recovering → heal → reclaim → expire → reconcile → seed → topup +> → sweep → ready`, or terminates at `failed`. **`failed` is not "down"** — the +> server keeps serving and keeps ingesting; it means one reconciliation step +> raised, and `failed_step` + `error` name which. `/status` stays HTTP `200` +> with `status: "ok"` at every phase, deliberately: the boot phase must never +> flip a liveness probe and restart the boot it is reporting on. `conflict: +> true` on the lease block is §8. + | Symptom | Command | What bad looks like | |---------|---------|---------------------| +| Still booting (check before anything below) | `curl -s http://localhost:8000/status \| python3 -c "import json,sys; r=json.load(sys.stdin); print(r['boot']['phase'], r.get('status_detail'))"` | anything other than `ready` — `metrics`/`spool` are `null` until then; `failed` names a step in `boot.failed_step` | +| Two writers on one queue directory | `curl -s http://localhost:8000/status \| python3 -c "import json,sys; print(json.load(sys.stdin)['writer_lease'])"` | `conflict: true` — see §8 | | Graph stopped updating | `curl -s http://localhost:8000/status \| python3 -c "import json,sys;[print(s['session_id'],s['last_successful_flush'],s['events_processed']) for s in json.load(sys.stdin)['sessions']]"` | `last_successful_flush` (unix seconds) hours old while `events_processed` keeps climbing | -| Backlog / loss accounting | `curl -s http://localhost:8000/status \| python3 -c "import json,sys; print(json.load(sys.stdin)['metrics'])"` | `in_queue_total` climbing and not falling; `degraded: true`; non-zero `dead_letter_total` | +| Backlog / loss accounting (**boot must be `ready`/`failed`**) | `curl -s http://localhost:8000/status \| python3 -c "import json,sys; print(json.load(sys.stdin)['metrics'])"` | `in_queue_total` climbing and not falling; `degraded: true`; non-zero `dead_letter_total`. `None` printed → still booting, not broken | +| Spool footprint from the server itself (**boot must be `ready`/`failed`**) | `curl -s http://localhost:8000/status \| python3 -c "import json,sys; print(json.load(sys.stdin)['spool'])"` | `spool_bytes_total` climbing and not falling. `None` printed → still booting | +| Drain workers dying | `journalctl --user -u context-intelligence-server --since '-1h' --no-pager \| grep drain_worker_died` | any hit — each line names the session id and carries the traceback; the session is recovered, but a repeat means something is reliably killing that drainer | | Dead drainers | `curl -s http://localhost:8000/status \| python3 -c "import json,sys; print(json.load(sys.stdin)['orphaned_sessions'])"` | non-zero — drain tasks that completed and stopped | | Spool growth | `du -sh ~/amplifier-context-intelligence-server-data-store/queues` | GB, not MB | | Spool file count / worst offender | `find ~/amplifier-context-intelligence-server-data-store/queues -name '*.log' \| wc -l` then `ls -lhS ~/amplifier-context-intelligence-server-data-store/queues/*.log \| head -5` | hundreds of files; any single file over ~1 GB | @@ -303,6 +352,8 @@ runs out of memory. It guarantees the failure is **bounded and attributable**: | `MemoryMax` / `MemoryHigh` | The server taking the host down with it; an unattributable kernel OOM kill | The server being killed — it just gets killed **in its own cgroup**, with `memory.events` naming it | | `StartLimitIntervalSec` / `StartLimitBurst` | An infinite restart loop burning CPU for days | The underlying failure — the unit stops in `failed` state and **stays down** until you fix it | | `write_concurrency: 4` | A backlog drain thundering Neo4j's transaction-memory ceiling | A backlog from forming in the first place | +| Compaction (on by default) + dead-letter expiry (opt-in, default off) | A *drained* session's bytes and orphaned poison records living on disk forever; a `.log` growing to the size of the whole session | **Undrained** data growing — a stalled drainer's tail is exactly what must not be touched; compaction shrinks nothing there | +| Writer-lease guard (§8) | A second process (accidental/misconfigured, not a rolling deploy — deployment is single-replica) silently writing the same queue directory | Nothing left to prevent by default: `enforce` refuses to boot against a live foreign lease. A stale lease (holder gone) is taken over automatically after the staleness window | The point of every one of them is the same: convert a silent, unbounded, self-perpetuating failure into a loud, bounded, one-shot one that shows up in @@ -317,6 +368,60 @@ restarts.** --- +## 8. Writer-lease guard + +The durable append log is correct because **exactly one process writes the +queue directory** (deployment is single-replica), with per-key locking inside +that process serializing every record to one contiguous, newline-terminated +line or not at all. The writer lease is a backstop against an accidental +second writer — a stray or misconfigured process pointed at the same `/data` +— not a rolling-deploy coordinator. + +**The default mode is `enforce`.** It refuses to boot against a **live** +foreign lease, so the second-writer overlap is prevented outright rather than +merely reported. It releases the lease on clean shutdown, so a restart +reacquires immediately, and it takes over a **stale** foreign lease +automatically once the holder's heartbeat ages past the staleness window +(heartbeat x multiplier — 15s at defaults) — an unclean exit recovers on its +own after that window instead of crash-looping. + +`detect` acquires the lease best-effort, heartbeats it, and only latches + +surfaces a conflict on `/status.writer_lease` — it never refuses to boot. +`off` disables the guard entirely. + +```bash +curl -s http://localhost:8000/status \ + | python3 -c "import json,sys; print(json.dumps(json.load(sys.stdin)['writer_lease'], indent=2))" +``` + +| Field | Read it as | +|-------|-----------| +| `mode` | `enforce` (default) / `detect` / `off` | +| `acquired` | whether this process currently believes it holds the lease | +| `conflict` | **the alarm.** `true` = a foreign writer was observed (under `enforce` this means boot was refused; under `detect` it is observe-only) | +| `conflict_source` | `boot` / `reacquire` / `runtime` — when it was seen (`runtime` = the lease was taken from a running process) | +| `observed_owner`, `observed_at` | who, and when | +| `took_over_stale`, `superseded_owner`, `superseded_age_seconds` | a *stale* lease was superseded — normal after an unclean exit, not a conflict | +| `error` | the guard is **not armed** for this process (share fault, hung mount). Evidence of nothing — never a conflict | +| `force_acquire` | the one-boot escape hatch is still set. Unset it | +| `heartbeat_seconds`, `staleness_seconds`, `lease_age_seconds` | heartbeat interval and staleness tolerance | + +The matching log lines are `writer_lease_conflict` (ERROR) and, on a stale +takeover, a `writer_lease` WARNING naming the superseded owner and its age. + +**On a conflict:** confirm you are not running two processes against one +queue directory (see the single-instance invariant in `AGENTS.md`), and get +down to one writer. A conflict latches — it stays visible after the overlap +ends — so treat it as "this happened", and correlate `observed_at` with your +timeline rather than assuming it is still happening. + +> `writer_lease_force_acquire` is a **one-boot** escape hatch for when you are +> certain the previous writer is gone but its lease hasn't gone stale yet. It +> logs a `WARNING` on every boot while set and shows on `/status`; unset it +> immediately afterwards. + +--- + ## See also - [docs/service-setup.md](service-setup.md) §5 — the systemd unit these guards live in. - [docs/service-setup.md](service-setup.md) §9 — troubleshooting table. diff --git a/pyproject.toml b/pyproject.toml index 2aaa7caf..e1fe9974 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.2" description = "Context Intelligence Server for Amplifier" requires-python = ">=3.11" dependencies = [ diff --git a/tests/conftest.py b/tests/conftest.py index ecfd94d9..3fddf0cf 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -10,16 +10,14 @@ "AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_ALLOW_UNAUTHENTICATED", "true" ) -from collections.abc import AsyncGenerator, Generator # noqa: E402 -from typing import Any # noqa: E402 +from collections.abc import AsyncGenerator, Generator +from typing import Any, Self -import httpx # noqa: E402 -import pytest # noqa: E402 - - -from context_intelligence_server.main import app, registry # noqa: E402 -from context_intelligence_server.services import HookStateService # noqa: E402 +import httpx +import pytest +from context_intelligence_server.main import app, registry +from context_intelligence_server.services import HookStateService # --------------------------------------------------------------------------- # Shared Neo4j mock helpers (used by POST /cypher tests) @@ -64,7 +62,7 @@ async def run(self, query: str, params: dict[str, Any]) -> MockNeo4jResult: raise self._exc return MockNeo4jResult(self._rows) - async def __aenter__(self) -> "MockNeo4jSession": + async def __aenter__(self) -> Self: return self async def __aexit__(self, *args: object) -> None: @@ -101,6 +99,7 @@ def anyio_backend() -> str: @pytest.fixture(autouse=True) def safe_settings(tmp_path: Any) -> Generator[None, None, None]: from unittest.mock import patch + from context_intelligence_server.config import Neo4jClientConfig from context_intelligence_server.config import Settings as _Settings @@ -122,6 +121,11 @@ class _SettingsProxy: neo4j_flush_chunk_rows: int = _real.neo4j_flush_chunk_rows neo4j_flush_chunk_bytes: int = _real.neo4j_flush_chunk_bytes neo4j_lock_timeout: float = _real.neo4j_lock_timeout + neo4j_max_connection_pool_size: int = _real.neo4j_max_connection_pool_size + neo4j_max_connection_lifetime: float = _real.neo4j_max_connection_lifetime + # Mirrors real Settings fields drain_worker reads via get_settings(). + queue_compact_enabled: bool = _real.queue_compact_enabled + queue_compact_min_prefix_bytes: int = _real.queue_compact_min_prefix_bytes # Neo4j two-client split (doc 12): SessionRegistry.get_or_create() calls # settings.resolve_neo4j_admin() directly, so this proxy (which stands @@ -150,6 +154,52 @@ def resolve_neo4j_query(self) -> Neo4jClientConfig: yield +@pytest.fixture(autouse=True) +def reset_boot_state() -> Generator[None, None, None]: + """Reset the module-level ``BootState`` singleton to ``"ready"`` around + each test, since ``/status``'s ``metrics``/``spool`` fields are + phase-gated and most tests never drive a real ``lifespan()`` boot.""" + from context_intelligence_server.status import boot_state as _boot_state + + def _reset() -> None: + _boot_state.phase = "ready" + _boot_state.started_at = 0.0 + _boot_state.completed_at = None + _boot_state.reclaimed = 0 + _boot_state.reclaimed_bytes = 0 + _boot_state.kept = 0 + _boot_state.failed = 0 + _boot_state.resumed = 0 + _boot_state.deferred = 0 + _boot_state.error = None + _boot_state.failed_step = None + _boot_state.fallback_workspace_byte0 = 0 + _boot_state.fallback_workspace_sentinel = 0 + _boot_state.reclaim_enabled = False + + _reset() + yield + _reset() + + +@pytest.fixture(autouse=True) +def _restore_lease_io() -> Generator[None, None, None]: + """Recreate the process-wide `_LEASE_IO` executor in ``writer_lease`` if + a prior test's real `lifespan()` shutdown already closed it, so later + tests aren't left without it.""" + yield + import concurrent.futures + + from context_intelligence_server import writer_lease as wl_module + + try: + wl_module._LEASE_IO.submit(lambda: None).result(timeout=1.0) + except RuntimeError: + wl_module._LEASE_IO = concurrent.futures.ThreadPoolExecutor( + max_workers=1, thread_name_prefix="writer-lease-io" + ) + + @pytest.fixture(autouse=True) def reset_registry() -> Generator[None, None, None]: """Ensure each test starts with a clean session registry.""" @@ -161,7 +211,7 @@ def reset_registry() -> Generator[None, None, None]: registry._queue_manager = None registry._write_semaphore = None # Zero the live pipeline-conservation counters on the shared singleton so - # each test starts from a clean conservation baseline (D2). + # each test starts from a clean conservation baseline. registry._accepted_total = 0 registry._written_total = 0 registry._replayed_total = 0 @@ -192,10 +242,10 @@ async def auth_client( monkeypatch: pytest.MonkeyPatch, ) -> AsyncGenerator[httpx.AsyncClient, None]: """Client routed through asgi_app (auth middleware applied) with a test API key set.""" - import hashlib # noqa: PLC0415 + import hashlib - from context_intelligence_server.auth import StaticKeyResolver # noqa: PLC0415 - from context_intelligence_server.main import asgi_app # noqa: PLC0415 + from context_intelligence_server.auth import StaticKeyResolver + from context_intelligence_server.main import asgi_app # Build a StaticKeyResolver that maps sha256("test-secret") → "owner" so existing # integration tests that send `Authorization: Bearer test-secret` continue to work. diff --git a/tests/neo4j/test_driver_leak_bounded.py b/tests/neo4j/test_driver_leak_bounded.py new file mode 100644 index 00000000..0487eaa5 --- /dev/null +++ b/tests/neo4j/test_driver_leak_bounded.py @@ -0,0 +1,146 @@ +"""Behavioral evidence that the per-session driver leak is gone. + +Drives many sessions through the real registry construction path +(``get_or_create`` -> shared ``Neo4jGraphStore``) against a live Neo4j, doing +a real write per session so bolt connections are actually opened, then queries +the server's own connection list to prove the open bolt connections stay +bounded by the pool (never scale with the session count) and are released on +driver close. +""" + +from __future__ import annotations + +import asyncio +from typing import Any + +import pytest +from context_intelligence_server.config import Settings +from context_intelligence_server.registry import SessionRegistry +from neo4j import AsyncGraphDatabase # type: ignore[attr-defined] + +pytestmark = pytest.mark.neo4j + +# Many more sessions than the pool can hold: if each session built its own +# driver (the leak), open bolt connections would scale with this number. +SESSION_COUNT = 30 +# Small, explicit pool so the bound is unmistakable in the observed count. +POOL_SIZE = 8 + +# The Python driver identifies itself with this user-agent prefix; the probe +# driver below uses a different one so it is excluded from the count. +_PY_DRIVER_UA_PREFIX = "neo4j-python" +_PROBE_UA = "leak-probe/1.0" + +_COUNT_QUERY = ( + "CALL dbms.listConnections() YIELD connector, userAgent " + f"WHERE connector = 'bolt' AND userAgent STARTS WITH '{_PY_DRIVER_UA_PREFIX}' " + "RETURN count(*) AS c" +) + + +async def _count_python_bolt_connections(probe_driver: Any) -> int: + """Return the number of open bolt connections opened by the Python driver. + + Excludes the probe driver itself (distinct user-agent) so the count + reflects only the registry's shared driver. + """ + async with probe_driver.session() as session: + result = await session.run(_COUNT_QUERY) + record = await result.single() + return int(record["c"]) + + +@pytest.mark.asyncio +async def test_bolt_connections_stay_bounded_and_release( + neo4j_container: dict[str, Any], + monkeypatch: pytest.MonkeyPatch, + tmp_path: Any, + capsys: Any, +) -> None: + bolt_url = neo4j_container["bolt_url"] + user = neo4j_container["user"] + password = neo4j_container["password"] + + # Real settings pointing at the container, with a small bounded pool and + # writable scratch paths. The registry resolves its shared driver from these. + settings = Settings( + neo4j_url=bolt_url, + neo4j_user=user, + neo4j_password=password, + neo4j_max_connection_pool_size=POOL_SIZE, + neo4j_max_connection_lifetime=3600.0, + blob_path=str(tmp_path / "blobs"), + queues_path=str(tmp_path / "queues"), + ) + monkeypatch.setattr( + "context_intelligence_server.registry.get_settings", + lambda: settings, + ) + + probe_driver = AsyncGraphDatabase.driver( + bolt_url, auth=(user, password), user_agent=_PROBE_UA + ) + + reg = SessionRegistry() + + baseline = await _count_python_bolt_connections(probe_driver) + + async def run_session(i: int) -> None: + # Exactly the construction path the leak came from: the registry builds + # (or reuses) its one shared driver and injects it into this session's + # store. The write forces a real bolt connection through the pool. + worker = reg.get_or_create(f"leak-session-{i}", f"/workspace/{i}") + graph = worker.services.graph + await graph.upsert_node( + f"node-{i}", {"label": "Event", "session": f"leak-session-{i}"} + ) + await graph.flush() + + await asyncio.gather(*(run_session(i) for i in range(SESSION_COUNT))) + + during_load = await _count_python_bolt_connections(probe_driver) + + # Stop the idle drain workers before tearing the driver down. + for worker in list(reg._workers.values()): + if worker.task is not None: + worker.task.cancel() + await asyncio.gather( + *(w.task for w in reg._workers.values() if w.task is not None), + return_exceptions=True, + ) + + # Reclaim: closing the one shared driver must release every bolt connection. + await reg.close_neo4j_driver() + + # Poll briefly for the server to observe the closed connections. + after_close = during_load + for _ in range(20): + after_close = await _count_python_bolt_connections(probe_driver) + if after_close == 0: + break + await asyncio.sleep(0.25) + + await probe_driver.close() + + with capsys.disabled(): + print( + f"\n[driver-leak evidence] sessions={SESSION_COUNT} pool_size={POOL_SIZE} " + f"baseline={baseline} during_load={during_load} after_close={after_close}" + ) + + # One shared driver was built for all sessions, not one per session. + assert reg._neo4j_driver is None # closed above + # The core proof: open bolt connections are bounded by the pool and do NOT + # scale with the session count. + assert during_load <= POOL_SIZE + 2, ( + f"open bolt connections ({during_load}) exceeded the pool bound " + f"({POOL_SIZE}); a per-session driver would scale with {SESSION_COUNT}" + ) + assert during_load < SESSION_COUNT, ( + f"open bolt connections ({during_load}) scaled with session count " + f"({SESSION_COUNT}) -- the leak is not fixed" + ) + # Reclaim proof: the shared driver's pool is fully released on close. + assert after_close == 0, ( + f"bolt connections not released after driver close (still {after_close})" + ) diff --git a/tests/neo4j/test_end_handler_buffer_shadow.py b/tests/neo4j/test_end_handler_buffer_shadow.py index 9f66fd38..dc835fe2 100644 --- a/tests/neo4j/test_end_handler_buffer_shadow.py +++ b/tests/neo4j/test_end_handler_buffer_shadow.py @@ -43,7 +43,6 @@ _current_type, ) - # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -60,6 +59,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 +283,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 +305,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: @@ -320,6 +349,9 @@ async def test_fork_flush_end_yields_single_terminal_label( "timestamp": "2026-01-01T10:05:00Z", }, ) + # _handle_end does not flush itself; the drainer flushes after the + # batch. Flush here so this test's read sees the persisted result. + await neo4j_services.graph.flush() final_labels = await _neo4j_labels(neo4j_services, child_id) 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..b416a8a2 100644 --- a/tests/neo4j/test_orphan_visibility.py +++ b/tests/neo4j/test_orphan_visibility.py @@ -43,11 +43,11 @@ import pytest from neo4j import AsyncGraphDatabase -from context_intelligence_server.status import build_status_response from context_intelligence_server.neo4j_store import Neo4jGraphStore from context_intelligence_server.queue_manager import QueueManager from context_intelligence_server.registry import SessionRegistry, SessionWorker from context_intelligence_server.services import HookStateService +from context_intelligence_server.status import build_status_response pytestmark = pytest.mark.neo4j @@ -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/neo4j/test_shared_driver_adverse_state.py b/tests/neo4j/test_shared_driver_adverse_state.py new file mode 100644 index 00000000..c817ac12 --- /dev/null +++ b/tests/neo4j/test_shared_driver_adverse_state.py @@ -0,0 +1,59 @@ +"""Adverse-state test: shared driver survives one session's close while +another session is mid-drain. + +Two Neo4jGraphStore instances share one driver (mirrors SessionRegistry's +per-session construction). Closing session A must not disturb session B's +in-flight write; the shared driver is closed exactly once, by its owner, +after both sessions are done with it. +""" + +from __future__ import annotations + +from typing import Any + +import pytest +from context_intelligence_server.neo4j_store import Neo4jGraphStore +from neo4j import AsyncGraphDatabase # type: ignore[attr-defined] + +pytestmark = pytest.mark.neo4j + + +@pytest.mark.asyncio +async def test_session_a_close_does_not_disrupt_session_b( + neo4j_container: dict[str, Any], +) -> None: + shared_driver = AsyncGraphDatabase.driver( + neo4j_container["bolt_url"], + auth=(neo4j_container["user"], neo4j_container["password"]), + ) + + store_a = Neo4jGraphStore( + uri=neo4j_container["bolt_url"], driver=shared_driver, workspace="session-a" + ) + store_b = Neo4jGraphStore( + uri=neo4j_container["bolt_url"], driver=shared_driver, workspace="session-b" + ) + + await store_b.upsert_node("node-b", {"label": "Event"}) + + # Session A finalizes and closes its store while B still has unflushed + # work buffered -- this is the adverse state: A's close must not touch + # the driver B is still using. + await store_a.close() + + # B's write still lands: the shared driver was never closed under it. + await store_b.flush() + fetched = await store_b.get_node("node-b") + assert fetched is not None + + await store_b.close() + + # The shared driver is still open after both stores are done with it -- + # neither store owned it. Its owner closes it exactly once, at shutdown. + async with shared_driver.session() as session: + result = await session.run("RETURN 1 AS one") + record = await result.single() + assert record is not None + assert record["one"] == 1 + + await shared_driver.close() diff --git a/tests/neo4j/test_steady_state_reclaim_neo4j.py b/tests/neo4j/test_steady_state_reclaim_neo4j.py new file mode 100644 index 00000000..0bb98a35 --- /dev/null +++ b/tests/neo4j/test_steady_state_reclaim_neo4j.py @@ -0,0 +1,310 @@ +"""Neo4j-backed tests for steady-state queue-log reclaim (compaction). + +Uses the real QueueManager, SessionRegistry/SessionWorker/drain_worker, and +Neo4jGraphStore against an isolated throwaway container. Covers: a drained-idle +no-terminal session's committed prefix is reclaimed with all events preserved; +concurrent ingest during repeated compaction never drops or reorders events; +and a boot-recovered session is compacted before it dry-exits. + + uv run pytest tests/neo4j/test_steady_state_reclaim_neo4j.py -q -m neo4j +""" + +from __future__ import annotations + +import asyncio +import json +import time +from pathlib import Path +from typing import Any + +import pytest +from neo4j import AsyncGraphDatabase + +from context_intelligence_server.neo4j_store import Neo4jGraphStore, ensure_neo4j_schema +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 = "reclaim" +_TS = "2026-08-21T00:00:00+00:00" + + +def _line(event: str, workspace: str, data: dict[str, Any]) -> bytes: + return json.dumps({"event": event, "workspace": workspace, "data": data}).encode( + "utf-8" + ) + + +def _small_event(i: int, sid: str) -> bytes: + return _line( + "tool:pre", + _WS, + { + "session_id": sid, + "timestamp": _TS, + "tool_call_id": f"call-{sid}-{i}", + "tool_name": "bash", + "tool_input": "x" * 128, + }, + ) + + +def _seq_event(i: int, sid: str) -> bytes: + # tool_input carries the append ordinal so a test can detect reordering, + # not just count loss. The check is int(tool_input), not string order. + return _line( + "tool:pre", + _WS, + { + "session_id": sid, + "timestamp": _TS, + "tool_call_id": f"call-{sid}-{i:06d}", + "tool_name": "bash", + "tool_input": f"{i:09d}", + }, + ) + + +def _build_registry(queues_dir: Path) -> SessionRegistry: + reg = SessionRegistry() + reg._queue_manager = QueueManager(queues_dir=queues_dir) + reg._write_semaphore = asyncio.Semaphore(8) + reg._max_delivery_attempts = 3 + return reg + + +def _build_worker( + container: dict[str, Any], sid: str, *, live_event_seen: bool = True +) -> 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, + live_event_seen=live_event_seen, + ) + + +async def _graph_event_count(container: dict[str, Any], sid: str) -> int: + driver = AsyncGraphDatabase.driver( + container["bolt_url"], auth=(container["user"], container["password"]) + ) + try: + async with driver.session() as session: + result = await session.run( + "MATCH (n:Event) WHERE n.session_id = $sid AND n.workspace = $ws " + "RETURN count(n) AS c", + sid=sid, + ws=_WS, + ) + record = await result.single() + return int(record["c"]) if record else 0 + finally: + await driver.close() + + +async def _current_backlog(qm: QueueManager, sid: str) -> int: + batch = await qm.read_batch(sid, max_items=1_000_000) + return len(batch.records) + + +async def _poll_until(predicate, *, timeout: float, interval: float = 0.05) -> bool: + deadline = time.monotonic() + timeout + while True: + if await predicate(): + return True + if time.monotonic() >= deadline: + return False + await asyncio.sleep(interval) + + +@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() + + +@pytest.mark.timeout(120) +async def test_drained_idle_no_terminal_session_prefix_reclaimed( + neo4j_container: dict[str, Any], tmp_path: Path +) -> None: + queues_dir = tmp_path / "queues" + reg = _build_registry(queues_dir) + qm = reg.queue_manager + + sid = "drained-idle" + n = 500 + for i in range(n): + await qm.append(sid, _small_event(i, sid)) + + log_path = queues_dir / f"{sid}.log" + 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 + + async def _drained() -> bool: + return await _current_backlog(qm, sid) == 0 + + assert await _poll_until(_drained, timeout=60.0), "drainer never reached backlog 0" + + async def _log_collapsed() -> bool: + return not log_path.exists() or log_path.stat().st_size == 0 + + assert await _poll_until(_log_collapsed, timeout=10.0), ( + "drained-idle .log did not collapse to its empty undrained tail" + ) + + assert await _current_backlog(qm, sid) == 0 + assert await _graph_event_count(neo4j_container, sid) == n + + # A manual invocation on the now-fully-compacted log reclaims nothing. + assert await qm.compact_committed_prefix(sid, 0) == 0 + + worker.task.cancel() + try: + await worker.task + except asyncio.CancelledError: + pass + + +@pytest.mark.timeout(180) +async def test_concurrent_ingest_during_repeated_compaction_no_loss( + neo4j_container: dict[str, Any], + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Verifies count conservation AND ordering: a spy wraps the real + # upsert_node (still calling through) to record the delivery order of each + # ToolCall's sequence marker, so a reorder that preserves the count is + # still caught. + import context_intelligence_server.config as cfg_module + import context_intelligence_server.registry as reg_module + from context_intelligence_server.neo4j_store import Neo4jGraphStore + + queues_dir = tmp_path / "queues" + reg = _build_registry(queues_dir) + qm = reg.queue_manager + + class _S: + queue_compact_enabled = True + queue_compact_min_prefix_bytes = 4096 # fires repeatedly mid-stream + stale_session_timeout = 3600.0 + + monkeypatch.setattr(cfg_module, "get_settings", lambda: _S()) + monkeypatch.setattr(reg_module, "get_settings", lambda: _S()) + + sid = "concurrent" + n_total = 400 + + observed_order: list[int] = [] + real_upsert_node = Neo4jGraphStore.upsert_node + + async def _spy_upsert_node(self: Any, node_id: str, data: dict[str, Any]) -> None: + if data.get("session_id") == sid and "ToolCall" in data.get("labels", []): + observed_order.append(int(data["tool_input"])) + await real_upsert_node(self, node_id, data) + + monkeypatch.setattr(Neo4jGraphStore, "upsert_node", _spy_upsert_node) + + worker = _build_worker(neo4j_container, sid) + reg._register_for_test(worker) + reg.start_drain(worker) + assert worker.task is not None + + for i in range(n_total): + await qm.append(sid, _seq_event(i, sid)) + reg.record_accepted(1) + if i % 23 == 0: + await asyncio.sleep(0.005) # interleave with the live drainer + + await qm.append( + sid, _line("session:end", _WS, {"session_id": sid, "timestamp": _TS}) + ) + reg.record_accepted(1) + await asyncio.wait_for(asyncio.shield(worker.task), timeout=120.0) + + # n_total tool:pre events + the session:end event, each its own :Event node. + assert await _graph_event_count(neo4j_container, sid) == n_total + 1 + assert len(await qm.read_dead_letters(sid)) == 0 + + assert len(observed_order) == n_total + assert observed_order == list(range(n_total)), ( + f"events delivered out of order: {observed_order[:10]}..." + ) + + metrics = await reg.pipeline_metrics() + assert metrics["residual"] == 0 + assert metrics["degraded"] is False + + +@pytest.mark.timeout(120) +async def test_boot_recovered_session_compacts_before_dry_exit( + neo4j_container: dict[str, Any], + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + # Drives the real get_or_create(..., recovered=True) entry point that + # crash-recovery and the sweep use, and asserts the log is compacted + # before the drainer's dry-exit (not after, not never). + import logging + + import context_intelligence_server.config as cfg_module + import context_intelligence_server.registry as reg_module + from context_intelligence_server.config import Settings + + queues_dir = tmp_path / "queues" + reg = SessionRegistry() + reg._queue_manager = QueueManager(queues_dir=queues_dir) + reg._write_semaphore = asyncio.Semaphore(8) + reg._max_delivery_attempts = 3 + qm = reg.queue_manager + + # Real Settings pointed at the test container, so get_or_create's own + # store/blob-store construction path runs for real. + settings = Settings() + settings.neo4j_url = neo4j_container["bolt_url"] + settings.neo4j_user = neo4j_container["user"] + settings.neo4j_password = neo4j_container["password"] + settings.blob_path = str(tmp_path / "blobs") + settings.queue_compact_enabled = True + settings.queue_compact_min_prefix_bytes = 0 + settings.stale_session_timeout = 3600.0 + + monkeypatch.setattr(cfg_module, "get_settings", lambda: settings) + monkeypatch.setattr(reg_module, "get_settings", lambda: settings) + + sid = "boot-recovered" + n = 300 + for i in range(n): + await qm.append(sid, _small_event(i, sid)) + + log_path = queues_dir / f"{sid}.log" + assert log_path.stat().st_size > 0 + + with caplog.at_level(logging.INFO): + worker = reg.get_or_create(sid, _WS, recovered=True) + assert worker.task is not None + await asyncio.wait_for(asyncio.shield(worker.task), timeout=60.0) + + assert any("recovered_drainer_exited" in r.getMessage() for r in caplog.records), ( + "expected the dry-exit log line; the worker never reached it" + ) + assert not log_path.exists() or log_path.stat().st_size == 0 + assert await _graph_event_count(neo4j_container, sid) == n diff --git a/tests/test_boot_safety.py b/tests/test_boot_safety.py new file mode 100644 index 00000000..60c1361d --- /dev/null +++ b/tests/test_boot_safety.py @@ -0,0 +1,1515 @@ +"""Boot-safety hardening tests. No real Neo4j is used anywhere in this file. + +Covers classify_session's decision table, log-then-delete ordering, the +anti-over-deletion guarantee, boot recovery of unparseable heads, lean +/status during boot, crash-loop guards, live-session ownership safety, +recovered-drainer bounds, sidecar retention, and clean shutdown ordering. +""" + +from __future__ import annotations + +import asyncio +import errno +import json +import logging +from pathlib import Path +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest +from neo4j.exceptions import ServiceUnavailable + +import context_intelligence_server.main as main_module +from context_intelligence_server.config import Settings +from context_intelligence_server.main import _head_is_resumable, lifespan +from context_intelligence_server.queue_manager import QueueManager, Verdict +from context_intelligence_server.registry import SessionRegistry, SessionWorker +from context_intelligence_server.services import HookStateService +from context_intelligence_server.status import boot_state + +pytestmark = pytest.mark.integration + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _line(workspace: str = "/ws", event: str = "tool_use", **data: Any) -> bytes: + """A newline-terminated record, matching QueueManager.append's own + on-disk framing exactly (append() always ensures a trailing "\n").""" + return ( + json.dumps({"event": event, "workspace": workspace, "data": data}) + "\n" + ).encode("utf-8") + + +async def _qm(tmp_path: Path) -> QueueManager: + return QueueManager(queues_dir=tmp_path) + + +def _seed_log(tmp_path: Path, key: str, content: bytes) -> Path: + path = tmp_path / f"{key}.log" + path.write_bytes(content) + return path + + +def _seed_offset(tmp_path: Path, key: str, content: str) -> Path: + path = tmp_path / f"{key}.offset" + path.write_text(content, encoding="utf-8") + return path + + +def _seed_dead(tmp_path: Path, key: str, content: str = "") -> Path: + path = tmp_path / f"{key}.dead.jsonl" + path.write_text(content, encoding="utf-8") + return path + + +class _NoOpGraph: + """A minimal graph store fake: flush/close are no-ops, buffer is empty.""" + + def __init__(self) -> None: + self.closed = False + self.workspace = "default" + self.created_by: str | None = None + + async def flush(self) -> None: + return None + + async def close(self) -> None: + self.closed = True + + def discard_buffer(self) -> None: + return None + + async def upsert_node(self, node_id: str, data: dict[str, Any]) -> None: + return None + + async def upsert_edge(self, src_id: str, dst_id: str, data: dict[str, Any]) -> None: + return None + + async def get_node(self, node_id: str) -> dict[str, Any] | None: + return None + + 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 _make_worker( + session_id: str, workspace: str = "/ws", *, live_event_seen: bool = True +) -> SessionWorker: + return SessionWorker( + session_id=session_id, + workspace=workspace, + services=HookStateService( + workspace=workspace, blob_store=MagicMock(), graph_store=_NoOpGraph() + ), + live_event_seen=live_event_seen, + ) + + +# --------------------------------------------------------------------------- +# classify_session decision-table (table-driven; side-effect-free) +# --------------------------------------------------------------------------- + + +async def test_classify_resumable_head_parses(tmp_path: Path) -> None: + qm = await _qm(tmp_path) + _seed_log(tmp_path, "k1", _line()) + c = await qm.classify_session("k1", _head_is_resumable) + assert c.verdict is Verdict.RESUMABLE + assert c.reason == "" + # side-effect-free: files untouched. + assert (tmp_path / "k1.log").exists() + + +async def test_classify_orphan_offset_is_reclaim_orphans_job_not_classify( + tmp_path: Path, +) -> None: + """classify iterates *.log stems only; a .log that vanished + between glob and classify (or was never there) is a benign race.""" + qm = await _qm(tmp_path) + c = await qm.classify_session("ghost", _head_is_resumable) + assert c.verdict is Verdict.UNREADABLE + assert c.reason == "log_vanished" + + +async def test_classify_unparseable_offset_small_resets(tmp_path: Path) -> None: + _seed_log(tmp_path, "k2", _line() * 3) + _seed_offset(tmp_path, "k2", "not-a-number") + _seed_dead(tmp_path, "k2", "") # empty -> dead_empty=True + qm = await _qm(tmp_path) + c = await qm.classify_session("k2", _head_is_resumable) + assert c.verdict is Verdict.RESET_OFFSET + assert c.reason == "unparseable_offset" + assert c.dead_empty is True + + +async def test_classify_unparseable_offset_with_dead_letters_kept( + tmp_path: Path, +) -> None: + """RESET_OFFSET requires an EMPTY .dead.jsonl; otherwise the reset would + re-dead-letter the same poison line on replay.""" + _seed_log(tmp_path, "k3", _line() * 3) + _seed_offset(tmp_path, "k3", "garbage") + _seed_dead( + tmp_path, "k3", json.dumps({"ts": 1, "error": "x", "payload": "p"}) + "\n" + ) + qm = await _qm(tmp_path) + c = await qm.classify_session("k3", _head_is_resumable) + assert c.verdict is Verdict.KEEP + assert c.reason == "bad_offset_with_dead" + + +async def test_classify_unparseable_offset_large_still_resets(tmp_path: Path) -> None: + """An unparseable .offset must never delete an intact .log, regardless + of size -- only genuinely bad parsed values (negative/past-eof) stay + size-gated (see test_classify_offset_past_eof_large_still_deletes).""" + big = _line() * 3 + _seed_log(tmp_path, "k4", big) + _seed_offset(tmp_path, "k4", "not-a-number") + qm = await _qm(tmp_path) + settings = Settings(reclaim_redrain_max_bytes=1) # would have forced "large" + with patch( + "context_intelligence_server.queue_manager.get_settings", + return_value=settings, + ): + c = await qm.classify_session("k4", _head_is_resumable) + assert c.verdict is Verdict.RESET_OFFSET + assert c.reason == "unparseable_offset" + + +async def test_classify_offset_past_eof_large_still_deletes(tmp_path: Path) -> None: + """Unlike an unparseable offset, a parsed-but-past-EOF offset stays + size-gated -- genuinely bad content, not an unreadable sidecar.""" + big = _line() * 3 + _seed_log(tmp_path, "k4b", big) + _seed_offset(tmp_path, "k4b", str(len(big) + 1000)) + qm = await _qm(tmp_path) + settings = Settings(reclaim_redrain_max_bytes=1) # force "large" + with patch( + "context_intelligence_server.queue_manager.get_settings", + return_value=settings, + ): + c = await qm.classify_session("k4b", _head_is_resumable) + assert c.verdict is Verdict.UNRESUMABLE + assert c.reason == "offset_past_eof" + + +async def test_classify_negative_offset(tmp_path: Path) -> None: + _seed_log(tmp_path, "k5", _line()) + _seed_offset(tmp_path, "k5", "-5") + qm = await _qm(tmp_path) + c = await qm.classify_session("k5", _head_is_resumable) + assert c.verdict is Verdict.RESET_OFFSET + assert c.reason == "negative_offset" + + +async def test_classify_offset_past_eof(tmp_path: Path) -> None: + line = _line() + _seed_log(tmp_path, "k6", line) + _seed_offset(tmp_path, "k6", str(len(line) + 1000)) + qm = await _qm(tmp_path) + c = await qm.classify_session("k6", _head_is_resumable) + assert c.verdict is Verdict.RESET_OFFSET + assert c.reason == "offset_past_eof" + + +async def test_classify_empty_log(tmp_path: Path) -> None: + _seed_log(tmp_path, "k7", b"") + qm = await _qm(tmp_path) + c = await qm.classify_session("k7", _head_is_resumable) + assert c.verdict is Verdict.UNRESUMABLE + assert c.reason == "empty_log" + + +async def test_classify_fully_drained(tmp_path: Path) -> None: + line = _line() + _seed_log(tmp_path, "k8", line) + _seed_offset(tmp_path, "k8", str(len(line))) + qm = await _qm(tmp_path) + c = await qm.classify_session("k8", _head_is_resumable) + assert c.verdict is Verdict.DRAINED + assert c.reason == "fully_drained" + + +async def test_classify_mid_line_offset_resumes(tmp_path: Path) -> None: + """Offset numerically sane, < size, mid-line -- RESUME (self-heals).""" + line = _line() + _seed_log(tmp_path, "k9", line) + _seed_offset(tmp_path, "k9", "1") # inside the line, not on a boundary + qm = await _qm(tmp_path) + c = await qm.classify_session("k9", _head_is_resumable) + assert c.verdict is Verdict.RESUMABLE + + +async def test_classify_merged_head_resumes_with_fallback_workspace_byte0( + tmp_path: Path, +) -> None: + """A merged/truncated first uncommitted line resumes with the workspace + read from byte 0 of the same file, rather than being deleted.""" + good_head = _line(workspace="/real-ws") + merged = ( + b'{"event":"tool_use","workspace":"/ws","data":{}' + b"\n" + ) # truncated JSON + _seed_log(tmp_path, "k10", good_head + merged) + _seed_offset(tmp_path, "k10", str(len(good_head))) # first uncommitted = merged + qm = await _qm(tmp_path) + c = await qm.classify_session("k10", _head_is_resumable) + assert c.verdict is Verdict.RESUMABLE + assert c.reason == "fallback_workspace" + assert c.fallback_source == "byte0" + + +async def test_classify_garbage_line_resumes_via_sentinel_not_deleted( + tmp_path: Path, +) -> None: + """A single-line non-JSON garbage `.log` (newline-terminated, <1MiB, no + siblings) must not classify UNRESUMABLE/DELETE -- the recovery path + sentinel-dispatches any complete first line regardless of JSON validity.""" + garbage = b"not json at all\n" + _seed_log(tmp_path, "k11", garbage) + _seed_offset(tmp_path, "k11", "0") + qm = await _qm(tmp_path) + c = await qm.classify_session("k11", _head_is_resumable) + assert c.verdict is Verdict.RESUMABLE + assert c.reason == "fallback_workspace" + assert c.fallback_source == "sentinel" + # side-effect-free AND not deleted. + assert (tmp_path / "k11.log").exists() + + +async def test_boot_reclaim_survives_garbage_log_and_real_drain_dead_letters_it( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Full-stack: with reclaim_enabled=True, a single-line garbage `.log` + survives the real `_boot_reclaim` pass; a subsequent real drain then + dead-letters the unparseable line and advances past it.""" + garbage = b"not json at all\n" + _seed_log(tmp_path, "garbage-key", garbage) + + qm = QueueManager(queues_dir=tmp_path) + monkeypatch.setattr(main_module.registry, "_queue_manager", qm) + monkeypatch.setattr(main_module._settings, "reclaim_enabled", True) + monkeypatch.setattr(main_module._settings, "crash_recovery_respawn_limit", 8) + + await main_module._boot_reclaim() + + # Survives the REAL reclaim pass -- NOT unlinked despite reclaim_enabled=True. + assert (tmp_path / "garbage-key.log").exists() + assert boot_state.fallback_workspace_sentinel >= 1 + + # A real drain (registry.drain_worker, the same code path a respawned + # recovery drainer runs) dead-letters the unparseable line and + # advances the offset -- no mocks of the unit under test. + reg = SessionRegistry() + reg._queue_manager = qm + worker = _make_worker("garbage-key", "unknown-recovered", live_event_seen=False) + reg._register_for_test(worker) + reg._ensure_infra() + reg._max_delivery_attempts = 1 # force immediate dead-letter, no retry delay + + await reg.drain_worker(worker) + + dead_content = (tmp_path / "garbage-key.dead.jsonl").read_text(encoding="utf-8") + assert "not json at all" in dead_content + # Offset advanced past the poison line (drained to EOF; dry-exit fires + # since there is nothing left and the worker was recovered). + assert "garbage-key" not in reg._workers + + +async def test_classify_unclassifiable_large_file_kept_not_deleted( + tmp_path: Path, +) -> None: + """DELETE only when the probe window covered the WHOLE file. + A large file whose first MiB has no parseable line is KEPT, not deleted.""" + garbage = b"not json at all\n" * 200_000 # > 1 MiB, no parseable line anywhere + _seed_log(tmp_path, "k12", garbage) + _seed_offset(tmp_path, "k12", "0") + qm = await _qm(tmp_path) + c = await qm.classify_session("k12", _head_is_resumable) + assert c.verdict is Verdict.KEEP + assert c.reason == "unclassifiable" + assert (tmp_path / "k12.log").exists() + + +async def test_classify_unreadable_offset_is_kept_never_deleted( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """A transient FS error reading the offset is NOT a corruption finding + and must never be laundered into a deletion.""" + _seed_log(tmp_path, "k13", _line()) + qm = await _qm(tmp_path) + + def _raise_oserror(_session_id: str) -> int: + raise OSError(errno.EIO, "simulated transient I/O error") + + monkeypatch.setattr(qm, "_read_committed_offset", _raise_oserror) + c = await qm.classify_session("k13", _head_is_resumable) + assert c.verdict is Verdict.UNREADABLE + assert c.reason == "unreadable_offset" + + +# --------------------------------------------------------------------------- +# log-then-delete: audit line BEFORE unlink, all closed-set files gone +# --------------------------------------------------------------------------- + + +async def test_reclaim_deletes_and_logs_before_unlink( + tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: + _seed_log(tmp_path, "d1", b"") # empty_log -> delete + _seed_dead(tmp_path, "d1", "kept-forever\n") + qm = await _qm(tmp_path) + c = await qm.classify_session("d1", _head_is_resumable) + assert c.verdict is Verdict.UNRESUMABLE + + with caplog.at_level(logging.WARNING): + ok = await qm.reclaim(c, lambda: False) + + assert ok is True + assert not (tmp_path / "d1.log").exists() + assert not (tmp_path / "d1.offset").exists() + # .dead.jsonl is NEVER deleted by the boot-safety classifier. + assert (tmp_path / "d1.dead.jsonl").exists() + warnings = [r for r in caplog.records if r.levelno == logging.WARNING] + assert any( + "boot_reclaimed" in r.getMessage() and "reason=empty_log" in r.getMessage() + for r in warnings + ) + + +async def test_reclaim_logs_before_failing_unlink( + tmp_path: Path, caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch +) -> None: + """The 'log first' policy is mechanically enforced, not just documented: + even when the unlink itself fails, the boot_reclaimed line already fired.""" + _seed_log(tmp_path, "d2", b"") + qm = await _qm(tmp_path) + c = await qm.classify_session("d2", _head_is_resumable) + + real_unlink = Path.unlink + + def _raise_unlink(self: Path, *a: Any, **kw: Any) -> None: + if self.name == "d2.log": + raise OSError(errno.EACCES, "simulated permission error") + return real_unlink(self, *a, **kw) + + with ( + caplog.at_level(logging.WARNING), + patch.object(Path, "unlink", _raise_unlink), + ): + ok = await qm.reclaim(c, lambda: False) + + assert ok is False + messages = [r.getMessage() for r in caplog.records] + reclaimed_idx = next(i for i, m in enumerate(messages) if "boot_reclaimed" in m) + failed_idx = next(i for i, m in enumerate(messages) if "boot_reclaim_failed" in m) + assert reclaimed_idx < failed_idx + + +# --------------------------------------------------------------------------- +# the anti-over-deletion test: resumable data survives reclaim +# --------------------------------------------------------------------------- + + +async def test_resumable_and_mid_line_offset_survive_reclaim(tmp_path: Path) -> None: + line = _line() + _seed_log(tmp_path, "r1", line) # scenario 3: stranded undrained tail + _seed_log(tmp_path, "r2", line) + _seed_offset(tmp_path, "r2", "1") # scenario 7c: mid-line offset + qm = await _qm(tmp_path) + + for key in ("r1", "r2"): + c = await qm.classify_session(key, _head_is_resumable) + assert c.verdict is Verdict.RESUMABLE + # RESUMABLE is never reclaimed -- the caller (main._boot_reclaim) + # simply never calls reclaim() for it. Files remain exactly as-is. + assert (tmp_path / f"{key}.log").exists() + + +# --------------------------------------------------------------------------- +# /status is LEAN + ZERO-DISK during boot +# --------------------------------------------------------------------------- + + +async def test_status_lean_during_boot_zero_disk_reads( + monkeypatch: pytest.MonkeyPatch, +) -> None: + boot_state.phase = "reclaim" # an ACTIVE boot phase (not ready/failed) + called = {"pipeline_metrics": False, "spool_stats": False} + + async def _spy_metrics() -> dict: + called["pipeline_metrics"] = True + return {} + + async def _spy_spool() -> dict: + called["spool_stats"] = True + return {} + + monkeypatch.setattr(main_module.registry, "pipeline_metrics", _spy_metrics) + monkeypatch.setattr(main_module.registry.queue_manager, "spool_stats", _spy_spool) + + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=main_module.app), base_url="http://test" + ) as c: + response = await c.get("/status") + + assert response.status_code == 200 + data = response.json() + assert data["status"] == "ok" + assert data["boot"]["phase"] == "reclaim" + assert data["metrics"] is None + assert data["spool"] is None + assert data["status_detail"] == {"reason": "booting"} + # Direct assertion that neither + # derive_all_stats' caller nor spool_stats is ever invoked while booting. + assert called["pipeline_metrics"] is False + assert called["spool_stats"] is False + + +async def test_status_populated_once_ready(monkeypatch: pytest.MonkeyPatch) -> None: + boot_state.phase = "ready" + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=main_module.app), base_url="http://test" + ) as c: + response = await c.get("/status") + data = response.json() + assert data["metrics"] is not None + assert data["spool"] is not None + assert "status_detail" not in data + + +async def test_status_populated_on_failed_boot_C1( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Gates on boot-is-over, not boot-succeeded: `failed` is terminal and + the server keeps ingesting, so metrics/spool must never be permanently + nulled by a reconcile failure.""" + boot_state.phase = "failed" + boot_state.error = "RuntimeError: simulated" + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=main_module.app), base_url="http://test" + ) as c: + response = await c.get("/status") + assert response.status_code == 200 + data = response.json() + assert data["status"] == "ok" # the boot phase is informational, never the probe + assert data["boot"]["phase"] == "failed" + assert data["metrics"] is not None + assert data["spool"] is not None + + +def test_boot_state_default_phase_is_recovering_never_ready() -> None: + """The pre-lifespan default must not be 'ready' -- the zero-disk-during-boot + guarantee depends on this from import time.""" + from context_intelligence_server.status import BootState + + assert BootState().phase == "recovering" + assert BootState().phase != "ready" + + +# --------------------------------------------------------------------------- +# _boot_reconcile is exception-safe (the boot done-callback analogue) +# --------------------------------------------------------------------------- + + +async def test_boot_reconcile_survives_a_failing_step_and_names_it( + monkeypatch: pytest.MonkeyPatch, +) -> None: + async def _raise(*_a: Any, **_kw: Any) -> Any: + raise OSError(errno.EIO, "simulated reconcile failure") + + monkeypatch.setattr( + main_module.registry.queue_manager, "recovery_reconcile_dead", _raise + ) + monkeypatch.setattr( + main_module.registry.queue_manager, + "heal_torn_tails", + AsyncMock(return_value={}), + ) + + await main_module._boot_reconcile() + + assert boot_state.phase == "failed" + assert boot_state.failed_step == "reconcile" + assert boot_state.error is not None and "OSError" in boot_state.error + + +async def test_boot_reconcile_success_reaches_ready( + monkeypatch: pytest.MonkeyPatch, +) -> None: + main_module.app.state.schema_ready = True # schema is out of scope here + monkeypatch.setattr( + main_module.registry.queue_manager, + "heal_torn_tails", + AsyncMock(return_value={}), + ) + monkeypatch.setattr( + main_module.registry.queue_manager, + "recovery_reconcile_dead", + AsyncMock(return_value=0), + ) + monkeypatch.setattr( + main_module.registry.queue_manager, + "recovery_seed_counts", + AsyncMock(return_value=(0, 0)), + ) + monkeypatch.setattr(main_module._settings, "crash_recovery_respawn_limit", 0) + + await main_module._boot_reconcile() + + assert boot_state.phase == "ready" + assert boot_state.completed_at is not None + assert boot_state.error is None + + +# --------------------------------------------------------------------------- +# crash-loop guards (test_g1..test_g4) -- the four crash triggers +# --------------------------------------------------------------------------- + + +async def test_g1_recover_skips_corrupt_offset_key(tmp_path: Path) -> None: + _seed_log(tmp_path, "good", _line()) + _seed_log(tmp_path, "bad", _line()) + _seed_offset(tmp_path, "bad", "\x00\x00\x00") # NUL-filled + qm = await _qm(tmp_path) + result = await qm.recover() # must not raise + assert "good" in result + assert "bad" not in result # skipped, not crash-looped + + +async def test_g2_reconcile_dead_skips_corrupt_key(tmp_path: Path) -> None: + _seed_log(tmp_path, "bad", _line()) + _seed_offset(tmp_path, "bad", "not-a-number") + _seed_dead( + tmp_path, "bad", json.dumps({"ts": 1, "error": "x", "payload": "p"}) + "\n" + ) + qm = await _qm(tmp_path) + total = await qm.recovery_reconcile_dead() # must not raise + assert total == 0 + + +async def test_g3_seed_counts_skips_corrupt_key(tmp_path: Path) -> None: + _seed_log(tmp_path, "good", _line()) + _seed_log(tmp_path, "bad", _line()) + _seed_offset(tmp_path, "bad", "negative-not-parseable") + qm = await _qm(tmp_path) + accepted, _written = await qm.recovery_seed_counts() # must not raise + assert accepted >= 1 # "good" still contributes + + +async def test_g4_topup_read_batch_guard_skips_corrupt_key( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr( + main_module.registry, "_queue_manager", QueueManager(queues_dir=tmp_path) + ) + qm = main_module.registry.queue_manager + _seed_log(tmp_path, "bad", _line()) + _seed_offset(tmp_path, "bad", "\x00") + + async def _raise_read_batch(session_id: str, max_items: int) -> Any: + raise ValueError("simulated corrupt offset in read_batch") + + monkeypatch.setattr(qm, "read_batch", _raise_read_batch) + spawned: list[str] = [] + monkeypatch.setattr( + main_module.registry, "get_or_create", lambda s, w, **kw: spawned.append(s) + ) + + result = await main_module._crash_recovery_topup(None) # must not raise + assert result.dispatched == 0 + + +async def test_g5_active_sessions_skips_corrupt_key(tmp_path: Path) -> None: + _seed_log(tmp_path, "good", _line()) + _seed_log(tmp_path, "bad", _line()) + _seed_offset(tmp_path, "bad", "garbage") + qm = await _qm(tmp_path) + result = await qm.active_sessions() # must not raise + assert "good" in result + assert "bad" not in result + + +async def test_all_four_crash_triggers_reach_ready_with_reclaim_disabled( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """The four crash-loop states, all seeded together; the whole + _boot_reconcile pass must reach phase=ready with reclaim disabled too, + proving the guards are independent defence, not merely masked by deletion.""" + _seed_log(tmp_path, "nul-offset", _line()) + _seed_offset(tmp_path, "nul-offset", "\x00\x00") + _seed_log(tmp_path, "neg-offset", _line()) + _seed_offset(tmp_path, "neg-offset", "-1") + _seed_dead(tmp_path, "bad-payload", json.dumps({"payload": 123}) + "\n") + _seed_log(tmp_path, "bad-payload", _line()) + + main_module.app.state.schema_ready = True # schema is out of scope here + monkeypatch.setattr( + main_module.registry, "_queue_manager", QueueManager(queues_dir=tmp_path) + ) + monkeypatch.setattr(main_module._settings, "reclaim_enabled", False) + monkeypatch.setattr(main_module._settings, "crash_recovery_respawn_limit", 8) + monkeypatch.setattr( + main_module.registry, "get_or_create", lambda *a, **kw: MagicMock() + ) + + await main_module._boot_reconcile() + + assert boot_state.phase == "ready" + + +# --------------------------------------------------------------------------- +# live-session ownership safety during boot reclaim +# --------------------------------------------------------------------------- + + +async def test_gate1_live_worker_key_never_reclaimed( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + line = _line() + _seed_log(tmp_path, "owned", line) + _seed_offset(tmp_path, "owned", str(len(line))) # fully_drained shape + + qm = QueueManager(queues_dir=tmp_path) + monkeypatch.setattr(main_module.registry, "_queue_manager", qm) + monkeypatch.setattr(main_module._settings, "reclaim_enabled", True) + monkeypatch.setattr(main_module._settings, "crash_recovery_respawn_limit", 8) + main_module.registry._register_for_test(_make_worker("owned")) + + await main_module._boot_reclaim() + + assert (tmp_path / "owned.log").exists() + assert boot_state.kept >= 1 + + +async def test_key_reclaimed_after_worker_removed_drains_from_zero( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + line = _line() + _seed_log(tmp_path, "was-owned", line) + _seed_offset(tmp_path, "was-owned", str(len(line))) + + qm = QueueManager(queues_dir=tmp_path) + monkeypatch.setattr(main_module.registry, "_queue_manager", qm) + monkeypatch.setattr(main_module._settings, "reclaim_enabled", True) + monkeypatch.setattr(main_module._settings, "crash_recovery_respawn_limit", 8) + + await main_module._boot_reclaim() + + assert not (tmp_path / "was-owned.log").exists() + # Re-append -> drains from byte 0, no stale .offset survives. + await qm.append("was-owned", _line(event="second")) + batch = await qm.read_batch("was-owned", max_items=10) + assert len(batch.records) == 1 + + +# --------------------------------------------------------------------------- +# the honest recovered-drainer bound -- recovered-only population + forward progress +# --------------------------------------------------------------------------- + + +async def test_recovered_drainer_population_bounded_and_makes_progress( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Seeds N sessions with no terminal record, ceiling < N. Across >=3 + sweep passes, the recovered-only population never exceeds the derived + bound, and cumulative distinct dispatches exceed the ceiling (forward progress).""" + ceiling = 3 + n_sessions = 9 + qm = QueueManager(queues_dir=tmp_path) + for i in range(n_sessions): + await qm.append(f"sess-{i}", _line(session_id=f"sess-{i}")) + + reg = SessionRegistry() + reg._queue_manager = qm + monkeypatch.setattr(main_module, "registry", reg) + + dispatched_all: set[str] = set() + + async def _drain_to_dry(worker: SessionWorker) -> None: + """A faithful stand-in for drain_worker's dry-exit: drains the one + record then exits immediately (no terminal record, so the real + drain loop would otherwise never finish).""" + batch = await qm.read_batch(worker.session_id, max_items=10) + if batch.records: + await qm.commit(worker.session_id, batch.end_offset) + reg._deregister(worker.session_id) + + def _get_or_create( + sid: str, workspace: str, created_by: Any = None, **kw: Any + ) -> Any: + recovered = kw.get("recovered", False) + worker = _make_worker(sid, workspace, live_event_seen=not recovered) + reg._register_for_test(worker) + dispatched_all.add(sid) + asyncio.ensure_future(_drain_to_dry(worker)) + return worker + + monkeypatch.setattr(reg, "get_or_create", _get_or_create) + + for _pass in range(4): + result = await main_module._crash_recovery_topup(ceiling) + # Bounded: recovered-only population never exceeds a generous + # derived bound (ceiling per pass; the fake drains+exits inline + # so live population settles back to 0 well within one pass). + recovered_only = [w for w in reg.workers() if not w.live_event_seen] + assert len(recovered_only) <= ceiling + await asyncio.sleep(0) # let the fire-and-forget drain tasks run + assert result.dispatched <= ceiling + + # Forward progress: cumulative distinct sessions dispatched exceeds the + # ceiling -- this is what distinguishes real drainage from a stall. + assert len(dispatched_all) > ceiling + + +# --------------------------------------------------------------------------- +# sidecar retention: a boot never destroys its own quarantine +# --------------------------------------------------------------------------- + + +async def test_reclaim_orphans_keeps_current_boots_sidecar(tmp_path: Path) -> None: + """before_ts is THIS PROCESS's start time. A sidecar older than that + (mtime < before_ts) came from a PRIOR boot -> reclaimed. A sidecar THIS + boot's heal just created (mtime >= before_ts, i.e. after process start) + must survive -- a boot must never destroy its own quarantine.""" + import os as _os + import time as _time + + before_ts = _time.time() # "process start" -- captured BEFORE either file exists + + old_sidecar = tmp_path / "s1.log.torn-111.bin" + old_sidecar.write_bytes(b"old") + _os.utime(old_sidecar, (before_ts - 100, before_ts - 100)) # from a PRIOR boot + + new_sidecar = tmp_path / "s2.log.torn-222.bin" + new_sidecar.write_bytes(b"new") # created just now, by THIS boot's heal + + qm = await _qm(tmp_path) + result = await qm.reclaim_orphans(before_ts=before_ts) + + assert not old_sidecar.exists() + assert new_sidecar.exists() + assert result["reclaimed"] == 1 + + +async def test_reclaim_orphans_removes_orphan_offset_and_tmp(tmp_path: Path) -> None: + (tmp_path / "orphan.offset").write_text("5", encoding="utf-8") + (tmp_path / "orphan.offset.tmp").write_text("5", encoding="utf-8") + qm = await _qm(tmp_path) + result = await qm.reclaim_orphans(before_ts=0.0) + assert not (tmp_path / "orphan.offset").exists() + assert not (tmp_path / "orphan.offset.tmp").exists() + assert result["reclaimed"] == 2 + + +async def test_reclaim_orphans_respects_reclaim_enabled_false_dry_run( + tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: + """With `reclaim_enabled=False`, `reclaim_orphans` must not unlink an + orphan `.offset` or a stale `.torn-*.bin` sidecar -- only classify and + log the would-delete (action=dry_run), same as the per-key `reclaim` path.""" + import os as _os + import time as _time + + before_ts = _time.time() + (tmp_path / "orphan.offset").write_text("5", encoding="utf-8") + old_sidecar = tmp_path / "s1.log.torn-111.bin" + old_sidecar.write_bytes(b"quarantined-bytes") + _os.utime(old_sidecar, (before_ts - 100, before_ts - 100)) # prior boot + + qm = await _qm(tmp_path) + with caplog.at_level(logging.WARNING): + result = await qm.reclaim_orphans(before_ts=before_ts, enabled=False) + + # Nothing unlinked -- both targets survive. + assert (tmp_path / "orphan.offset").exists() + assert old_sidecar.exists() + # A real deletion must never be invisible: the would-delete is named + # in the dry-run log, and NOT counted as an actual reclaim. + assert any( + "orphan_offset" in r.getMessage() and "action=dry_run" in r.getMessage() + for r in caplog.records + ) + assert any( + "torn_sidecar" in r.getMessage() and "action=dry_run" in r.getMessage() + for r in caplog.records + ) + assert result["reclaimed"] == 0 + + +async def test_reclaim_orphans_enabled_true_deletes_what_dry_run_named( + tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: + """The same targets, with `enabled=True`, are unlinked, logged + (action=delete), and counted -- a real deletion is never invisible.""" + import os as _os + import time as _time + + before_ts = _time.time() + (tmp_path / "orphan.offset").write_text("5", encoding="utf-8") + old_sidecar = tmp_path / "s1.log.torn-111.bin" + old_sidecar.write_bytes(b"quarantined-bytes") + _os.utime(old_sidecar, (before_ts - 100, before_ts - 100)) + + qm = await _qm(tmp_path) + with caplog.at_level(logging.WARNING): + result = await qm.reclaim_orphans(before_ts=before_ts, enabled=True) + + assert not (tmp_path / "orphan.offset").exists() + assert not old_sidecar.exists() + assert any( + "orphan_offset" in r.getMessage() and "action=delete" in r.getMessage() + for r in caplog.records + ) + assert any( + "torn_sidecar" in r.getMessage() and "action=delete" in r.getMessage() + for r in caplog.records + ) + assert result["reclaimed"] == 2 + assert result["reclaimed_bytes"] == len(b"5") + len(b"quarantined-bytes") + + +async def test_boot_reclaim_orphan_offset_survives_with_reclaim_disabled( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """End-to-end: a boot with `reclaim_enabled=False` must leave an orphan + `.offset` (no matching `.log`) untouched -- `reclaim_enabled` must reach + `reclaim_orphans` itself, not just gate a telemetry counter.""" + (tmp_path / "orphan.offset").write_text("5", encoding="utf-8") + + qm = QueueManager(queues_dir=tmp_path) + monkeypatch.setattr(main_module.registry, "_queue_manager", qm) + monkeypatch.setattr(main_module._settings, "reclaim_enabled", False) + monkeypatch.setattr(main_module._settings, "crash_recovery_respawn_limit", 8) + + await main_module._boot_reclaim() + + assert (tmp_path / "orphan.offset").exists() + + +# --------------------------------------------------------------------------- +# the dry-exit + the strand-after-await fix +# --------------------------------------------------------------------------- + + +async def test_dry_exit_fires_for_recovered_drainer_over_drained_log( + tmp_path: Path, +) -> None: + qm = QueueManager(queues_dir=tmp_path) + reg = SessionRegistry() + reg._queue_manager = qm + line = _line() + await qm.append("sess-x", line) + await qm.commit("sess-x", len(line) + 1) # fully drained, no terminal record + + worker = _make_worker("sess-x", live_event_seen=False) # recovered=True shape + reg._register_for_test(worker) + + await reg.drain_worker(worker) # must return promptly (dry-exit), not hang + + assert "sess-x" not in reg._workers + assert worker.services.graph.closed is True # type: ignore[attr-defined] + + +async def test_dry_exit_negative_control_live_created_worker_never_exits( + tmp_path: Path, +) -> None: + qm = QueueManager(queues_dir=tmp_path) + reg = SessionRegistry() + reg._queue_manager = qm + line = _line() + await qm.append("sess-live", line) + await qm.commit("sess-live", len(line) + 1) + + worker = _make_worker("sess-live", live_event_seen=True) # live path (default) + reg._register_for_test(worker) + + task = asyncio.ensure_future(reg.drain_worker(worker)) + await asyncio.sleep(0.05) + assert not task.done() # never exits on an empty batch when live + # drain_worker catches CancelledError internally (pre-existing, + # unrelated behaviour: it closes the store, deregisters, and RETURNS + # rather than re-raising) -- so cancel()+await completes normally here. + # This test's real assertion already happened above: `not task.done()`. + task.cancel() + await task + assert worker.services.graph.closed is True # type: ignore[attr-defined] + + +async def test_dry_exit_negative_control_recovered_that_drained_still_exits( + tmp_path: Path, +) -> None: + """The RED test proving the field is NOT last_event_time in disguise: + a recovered worker that DID process >=1 record must still exit once + it runs dry (v1.2's rejected fix excluded exactly this population).""" + qm = QueueManager(queues_dir=tmp_path) + reg = SessionRegistry() + reg._queue_manager = qm + line = _line() + await qm.append("sess-drained-recovered", line) + + worker = _make_worker("sess-drained-recovered", live_event_seen=False) + reg._register_for_test(worker) + worker.last_event_time = 1.0 # simulate having processed >=1 record + + await reg.drain_worker(worker) + + assert "sess-drained-recovered" not in reg._workers + + +async def test_d1_no_strand_after_recheck_await(tmp_path: Path) -> None: + """A POST landing during the recheck's await must not strand the drainer's + client -- the flag is re-read after the await, aborting the exit.""" + qm = QueueManager(queues_dir=tmp_path) + reg = SessionRegistry() + reg._queue_manager = qm + worker = _make_worker("sess-race", live_event_seen=False) + reg._register_for_test(worker) + + real_read_batch = qm.read_batch + call_count = {"n": 0} + + async def _read_batch_with_race(session_id: str, max_items: int) -> Any: + call_count["n"] += 1 + if call_count["n"] == 2: + # Simulate a POST's get_or_create landing during THIS await, + # before the recheck's result is used. + worker.live_event_seen = True + return await real_read_batch(session_id, max_items) + + qm.read_batch = _read_batch_with_race # type: ignore[method-assign] + + task = asyncio.ensure_future(reg.drain_worker(worker)) + for _ in range(50): + await asyncio.sleep(0.01) + if call_count["n"] >= 2: + break + await asyncio.sleep(0.05) + + # The drainer must NOT have exited/deregistered -- the flag flip during + # the recheck aborts the exit. + assert "sess-race" in reg._workers + assert not task.done() + # drain_worker swallows CancelledError internally; cancel()+await + # completes normally -- assertions above already covered the behavior + task.cancel() + await task + + +# --------------------------------------------------------------------------- +# reset-offset threshold + dead-empty + in-lock race +# --------------------------------------------------------------------------- + + +async def test_reset_offset_applies_and_drains_from_zero(tmp_path: Path) -> None: + line = _line() + _seed_log(tmp_path, "reset-me", line * 2) + _seed_offset(tmp_path, "reset-me", "not-a-number") + qm = await _qm(tmp_path) + c = await qm.classify_session("reset-me", _head_is_resumable) + assert c.verdict is Verdict.RESET_OFFSET + + ok = await qm.reclaim(c, lambda: False) + + assert ok is True + assert not (tmp_path / "reset-me.offset").exists() + assert (tmp_path / "reset-me.log").exists() + batch = await qm.read_batch("reset-me", max_items=10) + assert len(batch.records) == 2 # drains from byte 0 + + +async def test_reset_offset_refused_when_dead_letters_appear_after_classify( + tmp_path: Path, +) -> None: + """The .dead.jsonl-empty precondition is re-checked inside the guarded + body -- a live drain dirtying the dead file between classify and + reclaim must not let a reset apply and re-dead-letter the poison line.""" + line = _line() + _seed_log(tmp_path, "race-dead", line * 2) + _seed_offset(tmp_path, "race-dead", "garbage") + qm = await _qm(tmp_path) + c = await qm.classify_session("race-dead", _head_is_resumable) + assert c.verdict is Verdict.RESET_OFFSET # dead was empty AT CLASSIFY TIME + + # Simulate a concurrent live drain dead-lettering into this key's file + # BETWEEN classify and reclaim. + _seed_dead(tmp_path, "race-dead", json.dumps({"payload": "x"}) + "\n") + + ok = await qm.reclaim(c, lambda: False) + + assert ok is False + assert (tmp_path / "race-dead.offset").exists() # untouched + dead_content = (tmp_path / "race-dead.dead.jsonl").read_text(encoding="utf-8") + assert dead_content.count("\n") == 1 # NOT duplicated / re-dead-lettered + + +async def test_reclaim_size_drift_refuses_to_delete(tmp_path: Path) -> None: + """The delete window is closed: a live append growing the + `.log` between classify and reclaim must refuse the delete.""" + _seed_log(tmp_path, "grows", b"") # empty_log -> delete + qm = await _qm(tmp_path) + c = await qm.classify_session("grows", _head_is_resumable) + assert c.verdict is Verdict.UNRESUMABLE + + # Simulate a live append landing between classify and reclaim. + (tmp_path / "grows.log").write_bytes(_line()) + + ok = await qm.reclaim(c, lambda: False) + + assert ok is False + assert (tmp_path / "grows.log").exists() + + +async def test_reclaim_gate1_refuses_when_key_becomes_owned(tmp_path: Path) -> None: + _seed_log(tmp_path, "becomes-owned", b"") + qm = await _qm(tmp_path) + c = await qm.classify_session("becomes-owned", _head_is_resumable) + + ok = await qm.reclaim(c, lambda: True) # now owned + + assert ok is False + assert (tmp_path / "becomes-owned.log").exists() + + +# --------------------------------------------------------------------------- +# no phantom reclaim over a dead-file-only key +# --------------------------------------------------------------------------- + + +async def test_no_phantom_reclaim_for_dead_only_key( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """A correctly-finalized session retains only a .dead.jsonl; classify + must never see this key at all (it iterates *.log stems only).""" + _seed_dead(tmp_path, "finalized", "kept\n") + qm = QueueManager(queues_dir=tmp_path) + monkeypatch.setattr(main_module.registry, "_queue_manager", qm) + monkeypatch.setattr(main_module._settings, "reclaim_enabled", True) + monkeypatch.setattr(main_module._settings, "crash_recovery_respawn_limit", 8) + + await main_module._boot_reclaim() + await main_module._boot_reclaim() # second pass: still nothing to see + + assert (tmp_path / "finalized.dead.jsonl").exists() + + +# --------------------------------------------------------------------------- +# reclaim_enabled=False is a genuine dry run +# --------------------------------------------------------------------------- + + +async def test_reclaim_disabled_classifies_but_deletes_nothing( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + _seed_log(tmp_path, "would-delete", b"") + line = _line() + _seed_log(tmp_path, "would-resume", line) + + qm = QueueManager(queues_dir=tmp_path) + monkeypatch.setattr(main_module.registry, "_queue_manager", qm) + monkeypatch.setattr(main_module._settings, "reclaim_enabled", False) + monkeypatch.setattr(main_module._settings, "crash_recovery_respawn_limit", 8) + + with caplog.at_level(logging.WARNING): + await main_module._boot_reclaim() + + # Every seeded file still exists -- nothing was unlinked. + assert (tmp_path / "would-delete.log").exists() + assert (tmp_path / "would-resume.log").exists() + assert any("action=dry_run" in r.getMessage() for r in caplog.records) + assert boot_state.reclaim_enabled is False + + +async def test_reclaim_enabled_true_deletes_what_dry_run_named( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + _seed_log(tmp_path, "target", b"") + qm = QueueManager(queues_dir=tmp_path) + monkeypatch.setattr(main_module.registry, "_queue_manager", qm) + monkeypatch.setattr(main_module._settings, "crash_recovery_respawn_limit", 8) + + monkeypatch.setattr(main_module._settings, "reclaim_enabled", False) + await main_module._boot_reclaim() + assert (tmp_path / "target.log").exists() # dry run: nothing deleted + + boot_state.reclaimed = 0 + boot_state.kept = 0 + monkeypatch.setattr(main_module._settings, "reclaim_enabled", True) + await main_module._boot_reclaim() + assert not (tmp_path / "target.log").exists() # enabled: deleted for real + + +# --------------------------------------------------------------------------- +# DRAINED is auto-reclaimed at boot regardless of reclaim_enabled; the risky +# verdicts (UNRESUMABLE, RESET_OFFSET) stay gated behind it. +# --------------------------------------------------------------------------- + + +async def test_boot_reclaim_auto_reclaims_drained_log_under_shipped_defaults( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """A fully-drained log (committed == complete_end == size) is the same + evidence delete_drained already acts on unconditionally -- it must be + reclaimed at boot even under shipped defaults (reclaim_enabled=False).""" + qm = QueueManager(queues_dir=tmp_path) + line = _line() + await qm.append("drained-key", line) + await qm.commit("drained-key", len(line)) + + monkeypatch.setattr(main_module.registry, "_queue_manager", qm) + monkeypatch.setattr(main_module._settings, "reclaim_enabled", False) + monkeypatch.setattr(main_module._settings, "crash_recovery_respawn_limit", 8) + boot_state.reclaimed = 0 + boot_state.kept = 0 + + await main_module._boot_reclaim() + + assert not (tmp_path / "drained-key.log").exists() + assert not (tmp_path / "drained-key.offset").exists() + assert boot_state.reclaimed >= 1 + + +async def test_boot_reclaim_risky_verdicts_stay_gated_under_shipped_defaults( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """UNRESUMABLE and RESET_OFFSET must still be dry-run-only under + reclaim_enabled=False -- only DRAINED gets the auto-reclaim carve-out.""" + big = _line() * 3 + _seed_log(tmp_path, "unresumable-key", big) + # offset_past_eof (not unparseable_offset): genuinely bad parsed value, + # so it stays size-gated -- unaffected by the unparseable-offset fix. + _seed_offset(tmp_path, "unresumable-key", str(len(big) + 1000)) + + line = _line() + _seed_log(tmp_path, "reset-key", line * 2) + _seed_offset(tmp_path, "reset-key", "garbage") # small -> RESET_OFFSET + + qm = QueueManager(queues_dir=tmp_path) + monkeypatch.setattr(main_module.registry, "_queue_manager", qm) + monkeypatch.setattr(main_module._settings, "reclaim_enabled", False) + monkeypatch.setattr(main_module._settings, "crash_recovery_respawn_limit", 8) + + # Between reset-key's size (108) and unresumable-key's size (162): the + # former stays a bounded RESET_OFFSET, the latter tips into UNRESUMABLE. + # Patched directly on queue_manager's own get_settings, like + # test_classify_offset_past_eof_large_still_deletes does -- immune to any + # other test's get_settings.cache_clear() changing the shared singleton. + threshold_settings = Settings(reclaim_redrain_max_bytes=150) + with patch( + "context_intelligence_server.queue_manager.get_settings", + return_value=threshold_settings, + ): + # Confirm the verdicts are what this test claims before asserting. + c_unresumable = await qm.classify_session( + "unresumable-key", _head_is_resumable + ) + assert c_unresumable.verdict is Verdict.UNRESUMABLE + c_reset = await qm.classify_session("reset-key", _head_is_resumable) + assert c_reset.verdict is Verdict.RESET_OFFSET + + await main_module._boot_reclaim() + + assert (tmp_path / "unresumable-key.log").exists() + assert (tmp_path / "reset-key.log").exists() + assert (tmp_path / "reset-key.offset").exists() + + +async def test_boot_reclaim_large_unparseable_offset_never_deletes_log( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """A large intact .log with an unparseable .offset must survive a REAL + boot reclaim (reclaim_enabled=True) -- only the .offset resets, so the + session re-drains from byte 0 instead of losing its data.""" + big = _line() * 3 + _seed_log(tmp_path, "big-unparseable", big) + _seed_offset(tmp_path, "big-unparseable", "not-a-number") + + qm = QueueManager(queues_dir=tmp_path) + monkeypatch.setattr(main_module.registry, "_queue_manager", qm) + monkeypatch.setattr(main_module._settings, "reclaim_enabled", True) + monkeypatch.setattr(main_module._settings, "crash_recovery_respawn_limit", 8) + + threshold_settings = Settings(reclaim_redrain_max_bytes=1) # old "large" cliff + with patch( + "context_intelligence_server.queue_manager.get_settings", + return_value=threshold_settings, + ): + await main_module._boot_reclaim() + + assert (tmp_path / "big-unparseable.log").exists() + assert not (tmp_path / "big-unparseable.offset").exists() + batch = await qm.read_batch("big-unparseable", max_items=10) + assert len(batch.records) == 3 # re-drains from byte 0 + + +async def test_boot_reclaim_drained_log_with_live_worker_is_skipped( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """The has_worker guard runs BEFORE classify -- a live worker's log is + never touched by boot reclaim, DRAINED or not.""" + qm = QueueManager(queues_dir=tmp_path) + line = _line() + await qm.append("live-drained-key", line) + await qm.commit("live-drained-key", len(line)) + + reg = SessionRegistry() + reg._queue_manager = qm + worker = _make_worker("live-drained-key", live_event_seen=True) + reg._register_for_test(worker) + + monkeypatch.setattr(main_module, "registry", reg) + monkeypatch.setattr(main_module._settings, "reclaim_enabled", False) + monkeypatch.setattr(main_module._settings, "crash_recovery_respawn_limit", 8) + + await main_module._boot_reclaim() + + assert (tmp_path / "live-drained-key.log").exists() + + +# --------------------------------------------------------------------------- +# shutdown cancels every task before closing drivers +# --------------------------------------------------------------------------- + + +async def test_shutdown_cancels_sweep_and_boot_tasks_before_closing_drivers( + monkeypatch: pytest.MonkeyPatch, +) -> None: + mock_driver = MagicMock() + mock_driver.close = AsyncMock() + monkeypatch.setattr(main_module._settings, "crash_recovery_respawn_limit", 2) + monkeypatch.setattr( + main_module._settings, "crash_recovery_sweep_interval_seconds", 60 + ) + monkeypatch.setattr( + main_module.registry, "get_or_create", lambda *a, **kw: MagicMock() + ) + + with ( + patch("context_intelligence_server.main.setup_logging"), + patch( + "context_intelligence_server.neo4j_store.AsyncGraphDatabase.driver", + return_value=mock_driver, + ), + patch("context_intelligence_server.main.ensure_neo4j_schema", new=AsyncMock()), + ): + async with lifespan(main_module.app): + await main_module.app.state.boot_task + sweep_task = main_module.app.state.sweep_task + boot_task = main_module.app.state.boot_task + assert not sweep_task.done() # the sweep loop runs forever + + assert sweep_task.done() + assert boot_task.done() + assert mock_driver.close.await_count == 2 + + +# --------------------------------------------------------------------------- +# schema-init must never crash-loop the server +# --------------------------------------------------------------------------- + + +async def test_lifespan_survives_schema_init_neo4j_unreachable( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """RED before the fix: a connectivity failure during schema init must + not propagate out of lifespan -- the server still reaches yield (serves + /status), schema_ready stays False, and no drainers are started.""" + mock_driver = MagicMock() + mock_driver.close = AsyncMock() + monkeypatch.setattr(main_module._settings, "crash_recovery_respawn_limit", 8) + monkeypatch.setattr( + main_module._settings, "crash_recovery_sweep_interval_seconds", 60 + ) + schema_mock = AsyncMock(side_effect=ServiceUnavailable("simulated: unreachable")) + topup_mock = AsyncMock() + + with ( + patch("context_intelligence_server.main.setup_logging"), + patch( + "context_intelligence_server.neo4j_store.AsyncGraphDatabase.driver", + return_value=mock_driver, + ), + patch("context_intelligence_server.main.ensure_neo4j_schema", new=schema_mock), + patch( + "context_intelligence_server.main._crash_recovery_topup", new=topup_mock + ), + ): + async with lifespan(main_module.app): # must reach yield without raising + await main_module.app.state.boot_task + assert main_module.app.state.schema_ready is False + assert boot_state.phase not in ("failed", "ready") + + assert schema_mock.await_count >= 1 + topup_mock.assert_not_called() + + +async def test_boot_reconcile_proceeds_past_schema_once_neo4j_recovers( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Once schema init succeeds, the background reconcile proceeds past + the schema phase and starts drainers via the normal topup phase.""" + mock_driver = MagicMock() + mock_driver.close = AsyncMock() + monkeypatch.setattr(main_module._settings, "crash_recovery_respawn_limit", 8) + monkeypatch.setattr( + main_module._settings, "crash_recovery_sweep_interval_seconds", 60 + ) + topup_mock = AsyncMock(return_value=main_module.TopupResult(1, 1, 0)) + + with ( + patch("context_intelligence_server.main.setup_logging"), + patch( + "context_intelligence_server.neo4j_store.AsyncGraphDatabase.driver", + return_value=mock_driver, + ), + patch("context_intelligence_server.main.ensure_neo4j_schema", new=AsyncMock()), + patch( + "context_intelligence_server.main._crash_recovery_topup", new=topup_mock + ), + ): + async with lifespan(main_module.app): + await main_module.app.state.boot_task + assert main_module.app.state.schema_ready is True + assert boot_state.phase == "ready" + + topup_mock.assert_awaited_once() + + +async def test_lifespan_survives_schema_data_conflict_records_boot_fail( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Reachable graph, genuinely un-migrated (untagged>0): still refuses to + start drainers, but via boot_state.fail (visible on /status) -- never an + unguarded raise out of lifespan.""" + mock_driver = MagicMock() + mock_driver.close = AsyncMock() + monkeypatch.setattr(main_module._settings, "crash_recovery_respawn_limit", 8) + monkeypatch.setattr( + main_module._settings, "crash_recovery_sweep_interval_seconds", 60 + ) + topup_mock = AsyncMock() + + with ( + patch("context_intelligence_server.main.setup_logging"), + patch( + "context_intelligence_server.neo4j_store.AsyncGraphDatabase.driver", + return_value=mock_driver, + ), + patch("context_intelligence_server.main.ensure_neo4j_schema", new=AsyncMock()), + patch( + "context_intelligence_server.main.count_untagged_nodes", + new=AsyncMock(return_value=3), + ), + patch( + "context_intelligence_server.main._crash_recovery_topup", new=topup_mock + ), + ): + async with lifespan(main_module.app): # must reach yield without raising + await main_module.app.state.boot_task + assert main_module.app.state.schema_ready is False + assert boot_state.phase == "failed" + assert boot_state.failed_step == "schema" + + topup_mock.assert_not_called() + + +# --------------------------------------------------------------------------- +# a hung boot phase must not leave /status.spool/metrics null forever +# --------------------------------------------------------------------------- + + +async def test_hung_boot_phase_times_out_and_populates_status( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A phase that hangs beyond boot_phase_timeout_seconds must not leave + /status stuck with spool=null forever -- it ends the boot as a visible + failed phase (failed_step named) instead.""" + monkeypatch.setattr(main_module._settings, "boot_phase_timeout_seconds", 0.05) + main_module.app.state.schema_ready = True # skip schema phase for this test + + async def _hang(*a: Any, **kw: Any) -> Any: + await asyncio.sleep(2) + + monkeypatch.setattr(main_module.registry.queue_manager, "heal_torn_tails", _hang) + + await main_module._boot_reconcile() + + assert boot_state.phase == "failed" + assert boot_state.failed_step == "heal" + + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=main_module.app), base_url="http://test" + ) as c: + response = await c.get("/status") + data = response.json() + assert data["spool"] is not None + assert data["metrics"] is not None + + +def test_boot_phase_timeout_seconds_default_and_validator() -> None: + from context_intelligence_server.config import Settings + + assert Settings().boot_phase_timeout_seconds == 300.0 + assert Settings(boot_phase_timeout_seconds=0).boot_phase_timeout_seconds == 0 + assert Settings(boot_phase_timeout_seconds=-1).boot_phase_timeout_seconds == -1 + + +# --------------------------------------------------------------------------- +# end-to-end: merged-head corruption never causes a permanent boot stall +# --------------------------------------------------------------------------- + + +async def test_merged_head_session_resumes_and_drains_behind_the_poison_line( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """A session whose first uncommitted line is merged/truncated classifies + as fallback_workspace (kept, not deleted), and a second `_boot_reconcile` + pass classifies it identically -- recoverable every boot, not vanished.""" + merged = b'{"event":"tool_use","workspace":"/ws","dat' + b"\n" # truncated + good_head = _line(workspace="/w0") + _seed_log(tmp_path, "merged-head", good_head + merged) + _seed_offset( + tmp_path, "merged-head", str(len(good_head)) + ) # merged is first-uncommitted + + qm = QueueManager(queues_dir=tmp_path) + c1 = await qm.classify_session("merged-head", _head_is_resumable) + assert c1.verdict is Verdict.RESUMABLE + assert c1.reason == "fallback_workspace" + assert (tmp_path / "merged-head.log").exists() + + # A second pass over the SAME on-disk state classifies identically -- + # deterministic, not a one-shot escape hatch. + c2 = await qm.classify_session("merged-head", _head_is_resumable) + assert c2.verdict is Verdict.RESUMABLE + assert c2.reason == "fallback_workspace" + + +# --------------------------------------------------------------------------- +# config.py defaults -- smoke-checks the coupling from the boot-safety angle +# --------------------------------------------------------------------------- + + +def test_d8_config_defaults_are_coupled() -> None: + s = Settings() + assert s.crash_recovery_respawn_limit == 8 + assert s.crash_recovery_sweep_interval_seconds == 60 + assert s.reclaim_redrain_max_bytes == 64 * 1024 * 1024 + assert s.reclaim_enabled is False + + +def test_head_is_resumable_is_total_never_raises() -> None: + """Valid-but-non-dict JSON must not escape as AttributeError.""" + for raw in (b"123", b"null", b'"str"', b"[]", b"{}", b"not json", b""): + result = _head_is_resumable(raw) + assert isinstance(result, bool) + assert _head_is_resumable(json.dumps({"workspace": "/ws"}).encode()) is True + assert _head_is_resumable(json.dumps({"workspace": ""}).encode()) is False 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_config.py b/tests/test_config.py index 957d71a8..8a80b937 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -76,21 +76,27 @@ def test_settings_has_durable_queue_defaults(): # --------------------------------------------------------------------------- -def test_crash_recovery_respawn_limit_defaults_to_unbounded(): - """Default MUST preserve today's behaviour exactly: unbounded (None).""" +def test_crash_recovery_respawn_limit_defaults_to_8(): + """Default changes from unbounded (None) to a finite 8. + + An unbounded default was never actually protective (no sweep task was + ever started -- see the sweep-interval test below), so this makes the + ceiling finite by default: an operator is now protected out of the box + instead of only when they remember to opt in.""" from context_intelligence_server.config import Settings s = Settings() - assert s.crash_recovery_respawn_limit is None + assert s.crash_recovery_respawn_limit == 8 -def test_crash_recovery_sweep_interval_defaults_to_300(): - """A finite ceiling drains its deferred tail via a periodic sweep; the - default interval must be a sane positive value so a finite cap is safe - out of the box (not silently stranded).""" +def test_crash_recovery_sweep_interval_defaults_to_60(): + """Default changes from 300 to 60, coupled to the respawn-limit + default change above. A finite ceiling with no sweep task can never + drain a deferred tail; 60s makes throughput drain-bound rather than + timer-bound for a large recovered backlog.""" from context_intelligence_server.config import Settings - assert Settings().crash_recovery_sweep_interval_seconds == 300 + assert Settings().crash_recovery_sweep_interval_seconds == 60 def test_crash_recovery_sweep_interval_accepts_zero_and_positive(): @@ -435,7 +441,7 @@ def test_neo4j_flush_chunk_bytes_default(): class TestValidateApiKeys: """T8-T12 / T22: _validate_api_keys enforces the NESTED shape, fail-closed. - The NESTED form (design D4) maps a 64-char SHA-256 hex digest to a metadata + The NESTED form maps a 64-char SHA-256 hex digest to a metadata dict carrying at least ``id``:: api_keys: diff --git a/tests/test_drain_lifecycle_logging.py b/tests/test_drain_lifecycle_logging.py new file mode 100644 index 00000000..209512bc --- /dev/null +++ b/tests/test_drain_lifecycle_logging.py @@ -0,0 +1,537 @@ +"""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 +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Literal +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +import context_intelligence_server.main as main_module +from context_intelligence_server.queue_manager import QueueManager +from context_intelligence_server.registry import SessionRegistry +from context_intelligence_server.writer_lease import WriterLease, WriterLeaseConflict +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" + + +# --------------------------------------------------------------------------- +# reclaim_orphans()'s two swallowed mtime-stat OSErrors +# --------------------------------------------------------------------------- + + +class TestG6ReclaimOrphansMtimeStatFailure: + async def test_g6a_torn_sidecar_mtime_stat_failure_logged( + self, + tmp_path: Path, + caplog: pytest.LogCaptureFixture, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + qdir = tmp_path / "queues" + qdir.mkdir() + sid = "d3-g6-torn" + torn = qdir / f"{sid}.log.torn-123.bin" + torn.write_bytes(b"x") + + orig_stat = Path.stat + + def _fake_stat(self: Path, *a: object, **kw: object) -> object: + if self.name == torn.name: + raise OSError(errno.EIO, "Input/output error") + return orig_stat(self, *a, **kw) # type: ignore[misc] + + monkeypatch.setattr(Path, "stat", _fake_stat) + qm = QueueManager(queues_dir=qdir) + + with caplog.at_level(logging.WARNING, logger=LOGGER_NAME): + result = await qm.reclaim_orphans( + before_ts=time.time() + 3600, enabled=False + ) + + assert result["failed"] == 1 + matches = [ + r + for r in caplog.records + if r.levelno == logging.ERROR + and "boot_reclaim_failed" in r.getMessage() + and "torn_sidecar" in r.getMessage() + and torn.name in r.getMessage() + ] + assert matches, [r.getMessage() for r in caplog.records] + rec = matches[0] + assert rec.exc_info is not None, ( + "boot_reclaim_failed (torn_sidecar mtime stat) must carry exc_info" + ) + + async def test_g6b_compact_tmp_mtime_stat_failure_logged( + self, + tmp_path: Path, + caplog: pytest.LogCaptureFixture, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + qdir = tmp_path / "queues" + qdir.mkdir() + sid = "d3-g6-compact" + tmp_file = qdir / f"{sid}.log.compact.tmp" + tmp_file.write_bytes(b"x") + + orig_stat = Path.stat + + def _fake_stat(self: Path, *a: object, **kw: object) -> object: + if self.name == tmp_file.name: + raise OSError(errno.EIO, "Input/output error") + return orig_stat(self, *a, **kw) # type: ignore[misc] + + monkeypatch.setattr(Path, "stat", _fake_stat) + qm = QueueManager(queues_dir=qdir) + + with caplog.at_level(logging.WARNING, logger=LOGGER_NAME): + result = await qm.reclaim_orphans( + before_ts=time.time() + 3600, enabled=False + ) + + assert result["failed"] == 1 + matches = [ + r + for r in caplog.records + if r.levelno == logging.ERROR + and "boot_reclaim_failed" in r.getMessage() + and "orphan_compact_tmp" in r.getMessage() + and tmp_file.name in r.getMessage() + ] + assert matches, [r.getMessage() for r in caplog.records] + rec = matches[0] + assert rec.exc_info is not None, ( + "boot_reclaim_failed (orphan_compact_tmp mtime stat) must carry exc_info" + ) + + +# --------------------------------------------------------------------------- +# writer_lease enforce-mode boot refusal (two raise sites) +# --------------------------------------------------------------------------- + + +@dataclass +class _LeaseSettings: + """Minimal duck-typed settings stub -- mirrors + tests/test_writer_lease.py::_LeaseSettings exactly (the + six fields WriterLease.acquire reads).""" + + writer_lease_mode: Literal["off", "detect", "enforce"] = "detect" + writer_lease_heartbeat_seconds: float = 5.0 + writer_lease_staleness_multiplier: float = 3.0 + writer_lease_confirm_delay_seconds: float = 0.0 + writer_lease_acquire_timeout_seconds: float = 5.0 + writer_lease_force_acquire: bool = False + + +def _lease_settings(**overrides: object) -> _LeaseSettings: + return _LeaseSettings(**overrides) # type: ignore[arg-type] + + +class TestG7WriterLeaseRefusalLogged: + async def test_g7a_enforce_refuses_fresh_foreign_lease_logs_error( + self, tmp_path: Path, caplog: pytest.LogCaptureFixture + ) -> None: + peer = WriterLease() + await peer.acquire(_lease_settings(), lambda: tmp_path) + + lease = WriterLease() + with ( + caplog.at_level(logging.ERROR, logger=LOGGER_NAME), + pytest.raises(WriterLeaseConflict), + ): + await lease.acquire( + _lease_settings(writer_lease_mode="enforce"), lambda: tmp_path + ) + + assert any( + r.levelno == logging.ERROR + and "writer_lease_refused_boot" in r.getMessage() + and peer.owner in r.getMessage() + for r in caplog.records + ), [r.getMessage() for r in caplog.records] + + async def test_g7b_enforce_loses_confirm_race_logs_error( + self, tmp_path: Path, caplog: pytest.LogCaptureFixture + ) -> None: + settings = _lease_settings( + writer_lease_mode="enforce", writer_lease_confirm_delay_seconds=0.2 + ) + a = WriterLease() + b = WriterLease() + with caplog.at_level(logging.ERROR, logger=LOGGER_NAME): + results = await asyncio.gather( + a.acquire(settings, lambda: tmp_path), + b.acquire(settings, lambda: tmp_path), + return_exceptions=True, + ) + losers = [r for r in results if isinstance(r, Exception)] + assert len(losers) == 1, results + assert isinstance(losers[0], WriterLeaseConflict) + + assert any( + r.levelno == logging.ERROR + and "writer_lease_refused_boot" in r.getMessage() + and "lost the acquire race" in r.getMessage() + for r in caplog.records + ), [r.getMessage() for r in caplog.records] + + +# --------------------------------------------------------------------------- +# 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.dispatched == 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..6ec95fea --- /dev/null +++ b/tests/test_durable_append_framing.py @@ -0,0 +1,1059 @@ +"""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] + + +# --------------------------------------------------------------------------- +# heal_torn_tails truncates and quarantines +# --------------------------------------------------------------------------- + + +async def test_heal_torn_tails_truncates_and_quarantines(tmp_path: Path) -> None: + qm = QueueManager(tmp_path) + key = "torn-key" + log_path = qm._log_path(key) + good = b'{"event":"a"}\n' + torn = b'{"event":"b","data":"partial-fragment-no-terminator' + log_path.write_bytes(good + torn) + + result = await qm.heal_torn_tails() + + assert result["files_healed"] == 1 + assert result["bytes_discarded"] == len(torn) + assert result["files_failed"] == 0 + assert log_path.read_bytes() == good, ( + "file must end exactly at the last complete line" + ) + + sidecars = list(tmp_path.glob(f"{key}.log.torn-*.bin")) + assert len(sidecars) == 1 + assert sidecars[0].read_bytes() == torn, "quarantined bytes must be byte-identical" + + batch = await qm.read_batch(key, max_items=10) + assert len(batch.lines) == 1 + assert batch.lines[0] == good[:-1] + assert _parses(batch.lines[0]) + + +async def test_heal_torn_tails_on_empty_and_newline_free_files(tmp_path: Path) -> None: + qm = QueueManager(tmp_path) + empty_key = "empty-key" + nf_key = "newline-free-key" + qm._log_path(empty_key).write_bytes(b"") + nf_fragment = b'{"event":"no-newline-yet"' + qm._log_path(nf_key).write_bytes(nf_fragment) + + result = await qm.heal_torn_tails() + + assert qm._log_path(empty_key).read_bytes() == b"", "0-byte file must be untouched" + assert not list(tmp_path.glob(f"{empty_key}.log.torn-*.bin")), ( + "no sidecar for an empty file" + ) + + assert qm._log_path(nf_key).read_bytes() == b"", ( + "newline-free file truncates to empty" + ) + sidecars = list(tmp_path.glob(f"{nf_key}.log.torn-*.bin")) + assert len(sidecars) == 1 + assert sidecars[0].read_bytes() == nf_fragment + + assert result["files_failed"] == 0 + # A file that was never referenced simply does not appear -- no-op, no exception. + assert not (tmp_path / "missing-key.log").exists() + + +# --------------------------------------------------------------------------- +# 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 -- never truncated; heal_torn_tails removes it at boot + 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]) + + +# --------------------------------------------------------------------------- +# heal_torn_tails cannot crash boot +# --------------------------------------------------------------------------- + + +async def test_heal_torn_tails_survives_an_oserror_and_still_boots( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + qm = QueueManager(tmp_path) + good = b'{"event":"ok"}\n' + torn = b'{"event":"torn","data":"no-terminator-yet' + for name in ("a", "b", "c"): + qm._log_path(name).write_bytes(good + torn) + + real_os_open = os.open + real_os_truncate = os.truncate + + def _flaky_open( + path: Any, flags: int, mode: int = 0o777, *a: Any, **kw: Any + ) -> int: + if "a.log.torn-" in str(path): + raise OSError("simulated quarantine-copy failure for a") + return real_os_open(path, flags, mode, *a, **kw) + + def _flaky_truncate(path: Any, length: int) -> None: + if str(path).endswith("b.log"): + raise OSError("simulated truncate failure for b") + return real_os_truncate(path, length) + + monkeypatch.setattr(qm_module.os, "open", _flaky_open) + monkeypatch.setattr(qm_module.os, "truncate", _flaky_truncate) + + with caplog.at_level( + logging.ERROR, logger="context_intelligence_server.queue_manager" + ): + result = await qm.heal_torn_tails() # MUST NOT RAISE + + assert result["files_failed"] == 2 + assert result["files_healed"] == 1 + + # a: quarantine-copy raised -> file left EXACTLY as seeded, never truncated. + assert qm._log_path("a").read_bytes() == good + torn + assert not list(tmp_path.glob("a.log.torn-*.bin")) + + # b: copy succeeded, truncate raised -> STILL left exactly as seeded + # (never truncated on a failed truncate call). + assert qm._log_path("b").read_bytes() == good + torn + + # c: fully healed. + assert qm._log_path("c").read_bytes() == good + assert list(tmp_path.glob("c.log.torn-*.bin")) + + error_msgs = [r.message for r in caplog.records if r.levelno >= logging.ERROR] + assert sum("torn_tail_heal_failed" in m for m in error_msgs) == 2 + + +# --------------------------------------------------------------------------- +# 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..35198959 --- /dev/null +++ b/tests/test_finalize_delete_ordering.py @@ -0,0 +1,576 @@ +"""`_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" + ) + + +# --------------------------------------------------------------------------- +# No race with a prior compaction on the same key +# --------------------------------------------------------------------------- + + +async def test_finalize_retry_does_not_race_compaction_on_the_same_key() -> None: + """Drives the real drain_worker loop with compaction enabled, then a + late append lands inside the finalize window: the retry must still + retain-then-succeed after a prior compaction ran on the same key.""" + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d5-t5-no-compaction-race" + graph = _AccumGraph() + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + class _CompactAlwaysSettings: + queue_compact_enabled = True + queue_compact_min_prefix_bytes = 0 + + 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, + ), + patch( + "context_intelligence_server.registry.get_settings", + return_value=_CompactAlwaysSettings(), + ), + ): + await qm.append(sid, _line("e1", "/ws", {"session_id": sid})) + task = _start_supervised(reg, worker) + + for _ in range(500): + if "e1" in graph.flushed: + break + await asyncio.sleep(0.01) + assert "e1" in graph.flushed, ( + "precondition: e1 must drain (and Trigger H compact) BEFORE " + "session:end arrives, in its own separate non-terminal batch" + ) + + await qm.append(sid, _line("session:end", "/ws", {"session_id": sid})) + await asyncio.wait_for(task, timeout=10.0) + + assert graph.flushed == {"e1", "session:end", "late:event"}, ( + "every event persisted exactly once, no reorder, no exception" + ) + assert calls["count"] == 2, ( + "the late-append-during-finalize retry must still retain-then-succeed " + "even after a prior compaction ran on the same key" + ) + assert not qm._log_path(sid).exists() + assert not qm._offset_path(sid).exists() + + +# --------------------------------------------------------------------------- +# 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..5f8b7811 100644 --- a/tests/test_graph_store.py +++ b/tests/test_graph_store.py @@ -10,19 +10,36 @@ from context_intelligence_server.graph_store import GraphStore, QueryableStore - # --------------------------------------------------------------------------- # Minimal conforming implementations for isinstance checks # --------------------------------------------------------------------------- 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 +93,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"}) @@ -210,6 +242,7 @@ def test_queryable_store_exported(): def test_no_graph_forest_name_references(): """Verify graph_forest_name does not appear anywhere in graph_store.py.""" import inspect + import context_intelligence_server.graph_store as m source = inspect.getsource(m) diff --git a/tests/test_idempotency_store_on_success.py b/tests/test_idempotency_store_on_success.py new file mode 100644 index 00000000..f72710fa --- /dev/null +++ b/tests/test_idempotency_store_on_success.py @@ -0,0 +1,448 @@ +"""Idempotency key burned before the durable append. + +The idempotency cache must not burn a key until ``queue_manager.append`` +returns successfully: ``seen(key)`` (read-only duplicate check) runs before +the append; ``store(key)`` runs only after it succeeds. A failed append must +never poison a retry with a false "duplicate". +""" + +from __future__ import annotations + +import asyncio +import json +from pathlib import Path +from unittest.mock import MagicMock + +import httpx +import pytest + +import context_intelligence_server.main as main_module +from context_intelligence_server.idempotency import EventIdempotencyCache + +_TIMESTAMP = "2026-06-16T20:17:11.604690+00:00" + + +@pytest.fixture(autouse=True) +def _clear_idempotency_cache() -> None: + """Mirrors tests/test_main.py's autouse fixture -- this file is a + separate module, so it needs its own clear (module-level fixtures are + not shared across test files; the cache is a process-wide singleton). + """ + main_module.idempotency_cache.clear() + + +def _payload(session_id: str, idempotency_key: str) -> dict: + return { + "event": "tool_use", + "workspace": "/ws", + "idempotency_key": idempotency_key, + "data": { + "session_id": session_id, + "timestamp": _TIMESTAMP, + }, + } + + +# --------------------------------------------------------------------------- +# Append fails => key not burned => retry honoured, becomes durable +# --------------------------------------------------------------------------- + + +async def test_t1_append_failure_leaves_key_unburned_retry_honoured( + client: httpx.AsyncClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """POST #1 (append raises) must not burn the key; POST #2 with the same + key must be honoured (202 "queued", actually appended), not "duplicate".""" + monkeypatch.setattr( + main_module.registry, "get_or_create", lambda *a, **k: MagicMock() + ) + appended: list[tuple[str, bytes]] = [] + call_count = 0 + + async def _fake_append(worker_key: str, raw: bytes) -> None: + nonlocal call_count + call_count += 1 + if call_count == 1: + raise OSError("simulated durable-write failure") + appended.append((worker_key, raw)) + + monkeypatch.setattr(main_module.registry.queue_manager, "append", _fake_append) + + payload = _payload("sess-d7-t1", "aci-event-v1:d7-t1-key") + + # POST #1: append fails. No exception handler is registered in main.py, + # and the test client's ASGITransport uses raise_app_exceptions=True + # (the httpx default) -- so the OSError propagates OUT of + # `client.post(...)` itself. There is no response object to assert a + # status code against. + with pytest.raises(OSError): + await client.post("/events", json=payload) + + assert call_count == 1 + assert appended == [] + + # POST #2: same key. The key was never stored because append never + # returned successfully, so the retry is honoured and becomes durable. + second = await client.post("/events", json=payload) + assert second.status_code == 202 + assert second.json()["status"] == "queued" + assert call_count == 2 + assert len(appended) == 1 + worker_key, raw = appended[0] + assert worker_key == "sess-d7-t1" + assert json.loads(raw)["event"] == "tool_use" + + +# --------------------------------------------------------------------------- +# T1b -- the cache itself is untouched by a failed append +# --------------------------------------------------------------------------- + + +async def test_t1b_cache_untouched_by_failed_append( + client: httpx.AsyncClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Asserted directly on the cache object, not inferred from the response.""" + monkeypatch.setattr( + main_module.registry, "get_or_create", lambda *a, **k: MagicMock() + ) + + async def _always_fail_append(worker_key: str, raw: bytes) -> None: + raise OSError("simulated durable-write failure") + + monkeypatch.setattr( + main_module.registry.queue_manager, "append", _always_fail_append + ) + + key = "aci-event-v1:d7-t1b-key" + payload = _payload("sess-d7-t1b", key) + + with pytest.raises(OSError): + await client.post("/events", json=payload) + + # seen(key) must be False -- nothing was stored. + assert main_module.idempotency_cache.seen(key) is False + + +# --------------------------------------------------------------------------- +# Genuine duplicate still "duplicate" +# --------------------------------------------------------------------------- + + +async def test_t2_genuine_duplicate_still_duplicate_no_double_append( + client: httpx.AsyncClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Two sequential POSTs, append always succeeds -> first "queued", + second "duplicate", exactly one durable line.""" + monkeypatch.setattr( + main_module.registry, "get_or_create", lambda *a, **k: MagicMock() + ) + appended: list[tuple[str, bytes]] = [] + + async def _fake_append(worker_key: str, raw: bytes) -> None: + appended.append((worker_key, raw)) + + monkeypatch.setattr(main_module.registry.queue_manager, "append", _fake_append) + + payload = _payload("sess-d7-t2", "aci-event-v1:d7-t2-key") + + first = await client.post("/events", json=payload) + second = await client.post("/events", json=payload) + + assert first.status_code == 202 + assert first.json()["status"] == "queued" + assert second.status_code == 202 + assert second.json()["status"] == "duplicate" + assert len(appended) == 1 + + +# Two concurrent same-key POSTs are serialized by the per-key lock: exactly +# one durable append, the other answered "duplicate". + + +async def test_t3_concurrent_same_key_lock_serializes_exactly_one_append( + client: httpx.AsyncClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Two concurrent POSTs with the SAME idempotency_key must be serialized + by the per-key lock spanning seen()->append->store(): exactly ONE line + is durably appended; the other response is "duplicate". Before the fix, + both concurrent requests observed seen() == False and BOTH durably + appended (2 lines) -- this is the regression test for that defect.""" + monkeypatch.setattr( + main_module.registry, "get_or_create", lambda *a, **k: MagicMock() + ) + appended: list[tuple[str, bytes]] = [] + first_entered_append = asyncio.Event() + call_count = 0 + + async def _fake_append(worker_key: str, raw: bytes) -> None: + nonlocal call_count + call_count += 1 + first_entered_append.set() + # Bounded window for a genuinely concurrent second arrival to try + # to race in. Pre-fix (no lock spanning the sequence), the second + # arrival's own seen() check runs unlocked during this window and + # also reaches append -- the bug (both append). Post-fix, the + # second arrival cannot even reach its own seen() check until this + # request releases the per-key lock (after store()), so it never + # calls append at all -- this sleep simply elapses unobserved. + await asyncio.sleep(0.3) + appended.append((worker_key, raw)) + + monkeypatch.setattr(main_module.registry.queue_manager, "append", _fake_append) + + payload = _payload("sess-d7-t3", "aci-event-v1:d7-t3-key") + + async def _second() -> httpx.Response: + # Fire only once the first request is genuinely mid-append, so the + # two requests provably overlap in time. + await asyncio.wait_for(first_entered_append.wait(), timeout=5) + return await client.post("/events", json=payload) + + r1, r2 = await asyncio.wait_for( + asyncio.gather(client.post("/events", json=payload), _second()), + timeout=10, + ) + + assert r1.status_code == 202 + assert r2.status_code == 202 + statuses = {r1.json()["status"], r2.json()["status"]} + assert statuses == {"queued", "duplicate"}, ( + f"exactly one concurrent same-key POST may be honoured as new and " + f"the other must be a duplicate; got {statuses}" + ) + assert len(appended) == 1, ( + f"expected exactly one durable append under the per-key lock, got " + f"{len(appended)}" + ) + + +# --------------------------------------------------------------------------- +# Lock is released even when append raises -- no deadlock, next waiter with +# the same key is honoured once it acquires the (now-free) lock. +# --------------------------------------------------------------------------- + + +async def test_t3b_lock_released_after_append_failure_next_waiter_honoured( + client: httpx.AsyncClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A per-key lock held across a FAILING append must still be released + (never leaked/deadlocked): a second, concurrently-waiting request with + the same key acquires the lock once free, finds seen() still False (the + first request's key was never burned), and is durably honoured.""" + monkeypatch.setattr( + main_module.registry, "get_or_create", lambda *a, **k: MagicMock() + ) + appended: list[tuple[str, bytes]] = [] + first_attempted = asyncio.Event() + call_count = 0 + + async def _fake_append(worker_key: str, raw: bytes) -> None: + nonlocal call_count + call_count += 1 + first_attempted.set() + if call_count == 1: + await asyncio.sleep(0.1) # genuine checkpoint before failing + raise OSError("simulated durable-write failure") + appended.append((worker_key, raw)) + + monkeypatch.setattr(main_module.registry.queue_manager, "append", _fake_append) + + payload = _payload("sess-d7-t3b", "aci-event-v1:d7-t3b-key") + + async def _first() -> None: + with pytest.raises(OSError): + await client.post("/events", json=payload) + + async def _second() -> httpx.Response: + await asyncio.wait_for(first_attempted.wait(), timeout=5) + return await client.post("/events", json=payload) + + _, second_response = await asyncio.wait_for( + asyncio.gather(_first(), _second()), timeout=10 + ) + + assert second_response.status_code == 202 + assert second_response.json()["status"] == "queued" + assert call_count == 2 + assert len(appended) == 1 + + +# --------------------------------------------------------------------------- +# No idempotency_key: no lock is taken, concurrent no-key requests unaffected +# --------------------------------------------------------------------------- + + +async def test_t3c_no_dedup_key_concurrent_requests_unaffected_by_lock( + client: httpx.AsyncClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Requests carrying no idempotency_key take NO lock: two concurrent + no-key POSTs still both durably append, exactly as before the fix. This + reuses the classic mutual-release pattern -- if a lock were mistakenly + applied to the no-key path, this would deadlock and time out.""" + monkeypatch.setattr( + main_module.registry, "get_or_create", lambda *a, **k: MagicMock() + ) + appended: list[tuple[str, bytes]] = [] + release = asyncio.Event() + call_count = 0 + + async def _fake_append(worker_key: str, raw: bytes) -> None: + nonlocal call_count + call_count += 1 + if call_count == 1: + await asyncio.wait_for(release.wait(), timeout=5) + else: + release.set() + appended.append((worker_key, raw)) + + monkeypatch.setattr(main_module.registry.queue_manager, "append", _fake_append) + + payload = { + "event": "tool_use", + "workspace": "/ws", + "data": {"session_id": "sess-d7-t3c", "timestamp": _TIMESTAMP}, + } + + r1, r2 = await asyncio.wait_for( + asyncio.gather( + client.post("/events", json=payload), + client.post("/events", json=payload), + ), + timeout=10, + ) + + assert r1.status_code == 202 + assert r2.status_code == 202 + assert {r1.json()["status"], r2.json()["status"]} == {"queued"} + assert len(appended) == 2 + + +# --------------------------------------------------------------------------- +# Reservation leak: not applicable by construction +# --------------------------------------------------------------------------- + + +def test_t4_reservation_leak_not_applicable_by_construction() -> None: + """The store-on-success design has no reservation state analogous to + ``check_and_reserve``, so there is no reservation that can leak or a key + that can be permanently blocked by a lost ``release``.""" + assert True # documentation-only; see docstring + + +# --------------------------------------------------------------------------- +# Replay path unaffected (regression) +# --------------------------------------------------------------------------- + + +async def test_t5_replay_never_stores_key_non_replay_still_honoured( + client: httpx.AsyncClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """``?replay=true`` twice with key K -> both "queued", 2 appends; then a + non-replay POST with K -> "queued" (not "duplicate"), proving replay + never stored K.""" + monkeypatch.setattr( + main_module.registry, "get_or_create", lambda *a, **k: MagicMock() + ) + appended: list[tuple[str, bytes]] = [] + + async def _fake_append(worker_key: str, raw: bytes) -> None: + appended.append((worker_key, raw)) + + monkeypatch.setattr(main_module.registry.queue_manager, "append", _fake_append) + + payload = _payload("sess-d7-t5", "aci-event-v1:d7-t5-key") + + replay1 = await client.post("/events?replay=true", json=payload) + replay2 = await client.post("/events?replay=true", json=payload) + assert replay1.status_code == 202 + assert replay1.json()["status"] == "queued" + assert replay2.status_code == 202 + assert replay2.json()["status"] == "queued" + assert len(appended) == 2 + + non_replay = await client.post("/events", json=payload) + assert non_replay.status_code == 202 + assert non_replay.json()["status"] == "queued" # NOT "duplicate" + assert len(appended) == 3 + + +# --------------------------------------------------------------------------- +# EventIdempotencyCache.seen/store unit behaviour +# --------------------------------------------------------------------------- + + +class TestEventIdempotencyCacheSeenStore: + """Direct unit coverage of the seen()/store() split.""" + + def test_seen_false_then_store_then_seen_true(self) -> None: + cache = EventIdempotencyCache() + assert cache.seen("k1") is False + cache.store("k1") + assert cache.seen("k1") is True + + def test_seen_purges_past_ttl_entries(self) -> None: + cache = EventIdempotencyCache(ttl_seconds=10) + cache.store("k1", now=0.0) + assert cache.seen("k1", now=5.0) is True # still within ttl + assert cache.seen("k1", now=20.0) is False # purged -- past ttl + + def test_store_trims_at_max_entries(self) -> None: + cache = EventIdempotencyCache(max_entries=2) + cache.store("k1", now=1.0) + cache.store("k2", now=2.0) + cache.store("k3", now=3.0) # trims the oldest (k1) + assert cache.seen("k1", now=3.0) is False + assert cache.seen("k2", now=3.0) is True + assert cache.seen("k3", now=3.0) is True + + def test_seen_refreshes_lru_recency_on_hit(self) -> None: + cache = EventIdempotencyCache(max_entries=2) + cache.store("k1", now=1.0) + cache.store("k2", now=2.0) + assert cache.seen("k1", now=2.5) is True # touch k1 -> moves to end + cache.store("k3", now=3.0) # trims the now-LRU k2, not k1 + assert cache.seen("k1", now=3.0) is True + assert cache.seen("k2", now=3.0) is False + assert cache.seen("k3", now=3.0) is True + + def test_store_idempotent_under_repeat_calls(self) -> None: + cache = EventIdempotencyCache() + cache.store("k1", now=1.0) + cache.store("k1", now=2.0) # repeat store must not raise + assert cache.seen("k1", now=2.0) is True + + +# --------------------------------------------------------------------------- +# No stale check_and_store references anywhere in the repo +# --------------------------------------------------------------------------- + + +def test_t7_no_stale_check_and_store_references() -> None: + """``check_and_store`` must have no remaining call/def sites anywhere in + the repo. Scoped to the ``check_and_store(`` call/def form (not the bare + identifier) so explanatory prose mentioning the old name doesn't trip + the scan; this file is excluded from the scan for the same reason.""" + repo_root = Path(__file__).resolve().parents[1] + this_file = Path(__file__).resolve() + target = "check_and_store(" + hits: list[str] = [] + for py_file in repo_root.rglob("*.py"): + resolved = py_file.resolve() + if resolved == this_file: + continue + if any( + part in {".git", ".venv", "__pycache__", "node_modules"} + for part in py_file.parts + ): + continue + text = py_file.read_text(encoding="utf-8", errors="ignore") + if target in text: + hits.append(str(py_file)) + assert hits == [], f"stale check_and_store( call/def sites found in: {hits}" diff --git a/tests/test_large_event_tail_drop.py b/tests/test_large_event_tail_drop.py new file mode 100644 index 00000000..aca997b5 --- /dev/null +++ b/tests/test_large_event_tail_drop.py @@ -0,0 +1,340 @@ +"""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 +import logging +from unittest.mock import AsyncMock, 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 --- + assert task.cancelled() or task.exception() is None + + +# --------------------------------------------------------------------------- +# Case B -- DEFECT: same enqueue shape, but the unguarded seam inside +# _handle_exhausted_batch's except-clause (dead_letter) raises, reproducing +# the silent tail-drop symptom. +# --------------------------------------------------------------------------- + + +class TestOversizedEventDefectSilentlyDropsTail: + async def test_tail_after_oversized_event_is_silently_dropped( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """Same batch shape as Case A, but `dead_letter` itself raises while + isolating the oversized line, escaping past drain_worker's only + guard and killing the task. The prefix persists, the tail (oversized + + 2 events) is stranded uncommitted on disk, and the task dies with + an exception. The done-callback then logs it loud, closes the store, + and deregisters the worker so a respawn drains the exact stranded + suffix with no gap or duplicate. + """ + reg = SessionRegistry() + qm = reg.queue_manager + sid = "large-event-defect" + + 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) + + # the injected defect: dead_letter raises inside the unguarded + # except-clause seam; saved so it can be restored before the respawn + _real_dead_letter = qm.dead_letter + qm.dead_letter = AsyncMock( # type: ignore[method-assign] + side_effect=OSError("disk unavailable while writing dead-letter record") + ) + + 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, + ), + caplog.at_level(logging.ERROR, logger="context_intelligence_server"), + ): + # fire-and-forget accept semantics: every event is durably + # appended before the drain loop ever touches Neo4j + 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})) + assert len((await qm.read_batch(sid, 10)).lines) == 5, ( + "all 5 events durably accepted before any Neo4j write is attempted" + ) + + task = await _drive_drain_to_quiescence(reg, qm, worker, sid) + # Do NOT cancel: the defect kills the task on its own. If the + # task were somehow still alive, that itself falsifies the + # reproduction, so make that failure explicit rather than + # masking it with an unconditional cancel. + if not task.done(): + await _cancel_and_await(task) + pytest.fail( + "drain task did not die on its own -- the injected " + "dead_letter failure did not escape the loop as the " + "root-cause analysis predicted; investigate before " + "trusting this reproduction" + ) + # let the done-callback (scheduled via call_soon) run before asserting + for _ in range(5): + await asyncio.sleep(0) + + # the small prefix persisted + assert fake.flushed == ["small-1", "small-2"], ( + "expected exactly the prefix before the oversized event to have " + "flushed via the per-line isolation path" + ) + + # the tail after the oversized event did not persist + assert "tail-1" not in fake.flushed + assert "tail-2" not in fake.flushed + + # committed offset stuck at the prefix boundary; the tail bytes + # remain durably on disk, uncommitted + remaining = await qm.read_batch(sid, 10) + remaining_events = [json.loads(raw)["event"] for raw in remaining.lines] + assert remaining_events == ["oversized", "tail-1", "tail-2"], ( + "the tail must remain stranded, uncommitted, on disk -- this is " + "the silent data-loss window this defect produces (bytes are not " + "gone, but nothing will ever drain them without a fix)" + ) + + # the drain task is done() with a silent, unretrieved exception + assert task.done() + assert not task.cancelled() + exc = task.exception() + assert exc is not None, ( + "the drain task must have died from an unhandled exception, " + "escaping past drain_worker's sole guard (asyncio.CancelledError " + "only)" + ) + assert isinstance(exc, OSError) + + # the oversized event was never dead-lettered by the crashed task + # either -- resolved once the respawn below runs + dead = await qm.read_dead_letters(sid) + assert dead == [] + + # --- the death is loud and self-healing --- + + # drain_worker_died was logged at ERROR with this session's id + died_records = [ + r + for r in caplog.records + if r.levelno == logging.ERROR + and "drain_worker_died" in r.getMessage() + and getattr(r, "session_id", None) == sid + ] + assert died_records, ( + "expected a drain_worker_died ERROR with session_id=" + f"{sid!r} -- start_drain's done-callback (_on_drain_done) is " + "the supervisor that makes this death loud instead of silent" + ) + assert died_records[0].exc_info is not None + assert died_records[0].exc_info[0] is OSError + + # the store was closed by the callback -- no driver leak + assert worker.store_closed is True + assert fake.closed is True + + # the worker was deregistered -- a fresh get_or_create can revive it + assert sid not in reg.active_sessions(), ( + "the crashed worker must be deregistered by _on_drain_done, or " + "it is a spent worker (store_closed) that start_drain would " + "refuse forever" + ) + + # a fresh respawn drains the exact stranded suffix with no gap or + # duplicate: oversized is still genuinely poison, but dead_letter is + # restored to the real implementation, so this time it succeeds + qm.dead_letter = _real_dead_letter # type: ignore[method-assign] + fake2 = _FaultInjectableGraph(poison_event="oversized") + worker2 = SessionWorker( + session_id=sid, + workspace="/ws", + services=HookStateService(workspace="/ws"), + ) + worker2.services.graph = fake2 # type: ignore[assignment] + reg._register_for_test(worker2) + + async def _process2( + w: SessionWorker, event: str, data: object, h: object + ) -> None: + # route to the respawned worker's own fresh graph, not the crashed one's + fake2.buffer.add(event) + + with patch( + "context_intelligence_server.registry.process_event", + side_effect=_process2, + ): + task2 = await _drive_drain_to_quiescence(reg, qm, worker2, sid) + await _cancel_and_await(task2) + + assert fake2.flushed == ["tail-1", "tail-2"] + dead2 = await qm.read_dead_letters(sid) + assert len(dead2) == 1 + assert json.loads(dead2[0]["payload"])["event"] == "oversized" + assert (await qm.read_batch(sid, 10)).lines == [] + assert task2 is not task diff --git a/tests/test_main.py b/tests/test_main.py index 302fa1bf..5dd27c80 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -4,8 +4,8 @@ import contextlib import json import logging -from pathlib import Path from collections.abc import AsyncGenerator +from pathlib import Path from typing import Any from unittest.mock import AsyncMock, MagicMock, patch @@ -162,7 +162,7 @@ async def test_post_events_increments_accepted_counter( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: - """A durably-accepted event increments the registry accepted_total (D2).""" + """A durably-accepted event increments the registry accepted_total.""" from context_intelligence_server.queue_manager import QueueManager # Point the registry at a tmp queue dir so the durable append is isolated. @@ -241,9 +241,7 @@ async def test_drain_loop_processes_event( ) -> None: """A posted event is drained from the durable log by the sticky drainer. - Migrated off the vestigial in-memory worker.queue (Phase B2, Task 5): the - durable drain loop reads from the on-disk QueueManager log, so success is - observed by polling that log to empty rather than worker.queue.join(). + Success is observed by polling the on-disk QueueManager log to empty. """ from context_intelligence_server.neo4j_store import Neo4jGraphStore @@ -678,13 +676,15 @@ def test_main_doctor_dispatches_run_doctor_fix_false(self) -> None: with its return code.""" import context_intelligence_server.main as _main_mod - with patch( - "context_intelligence_server.doctor.run_doctor", - new_callable=AsyncMock, - return_value=0, - ) as mock_doctor: - with pytest.raises(SystemExit) as exc_info: - _main_mod.main(["doctor"]) + with ( + patch( + "context_intelligence_server.doctor.run_doctor", + new_callable=AsyncMock, + return_value=0, + ) as mock_doctor, + pytest.raises(SystemExit) as exc_info, + ): + _main_mod.main(["doctor"]) mock_doctor.assert_awaited_once_with(fix=False) assert exc_info.value.code == 0 @@ -693,13 +693,15 @@ def test_main_doctor_fix_dispatches_run_doctor_fix_true(self) -> None: """main(["doctor", "--fix"]) calls doctor.run_doctor(fix=True).""" import context_intelligence_server.main as _main_mod - with patch( - "context_intelligence_server.doctor.run_doctor", - new_callable=AsyncMock, - return_value=1, - ) as mock_doctor: - with pytest.raises(SystemExit) as exc_info: - _main_mod.main(["doctor", "--fix"]) + with ( + patch( + "context_intelligence_server.doctor.run_doctor", + new_callable=AsyncMock, + return_value=1, + ) as mock_doctor, + pytest.raises(SystemExit) as exc_info, + ): + _main_mod.main(["doctor", "--fix"]) mock_doctor.assert_awaited_once_with(fix=True) assert exc_info.value.code == 1 @@ -718,7 +720,7 @@ async def test_lifespan_creates_and_closes_driver( "context_intelligence_server.main.setup_logging", ) as mock_setup_logging, patch( - "context_intelligence_server.main.AsyncGraphDatabase.driver", + "context_intelligence_server.neo4j_store.AsyncGraphDatabase.driver", return_value=mock_driver, ) as mock_driver_factory, ): @@ -742,7 +744,7 @@ async def test_lifespan_creates_and_closes_driver( # --------------------------------------------------------------------------- -# Lifespan crash-recovery + workers==1 guard tests (Phase B2) +# Lifespan crash-recovery + workers==1 guard tests # --------------------------------------------------------------------------- @@ -780,7 +782,7 @@ async def test_lifespan_recovers_and_respawns_drainers( with ( patch("context_intelligence_server.main.setup_logging"), patch( - "context_intelligence_server.main.AsyncGraphDatabase.driver", + "context_intelligence_server.neo4j_store.AsyncGraphDatabase.driver", return_value=mock_driver, ), patch( @@ -789,7 +791,8 @@ async def test_lifespan_recovers_and_respawns_drainers( ), ): async with lifespan(main_module.app): - pass + # Recovery is backgrounded -- await it before shutdown cancels it. + await main_module.app.state.boot_task assert (sid, "/recovered-ws") in spawned @@ -797,9 +800,8 @@ 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 dispatched under + the `_RECOVERY_FALLBACK_WORKSPACE` sentinel, never under workspace=''.""" sid = "sess-empty-ws" qm = registry.queue_manager body = json.dumps( @@ -822,7 +824,7 @@ async def test_lifespan_skips_recovery_for_empty_workspace( with ( patch("context_intelligence_server.main.setup_logging"), patch( - "context_intelligence_server.main.AsyncGraphDatabase.driver", + "context_intelligence_server.neo4j_store.AsyncGraphDatabase.driver", return_value=mock_driver, ), patch( @@ -831,9 +833,15 @@ async def test_lifespan_skips_recovery_for_empty_workspace( ), ): async with lifespan(main_module.app): - pass + # Recovery is backgrounded -- await it before shutdown cancels it. + await main_module.app.state.boot_task - assert spawned == [] + # Dispatched under the sentinel fallback, never under workspace=''. + assert len(spawned) == 1 + dispatched_sid, dispatched_ws = spawned[0] + assert dispatched_sid == sid + assert dispatched_ws == main_module._RECOVERY_FALLBACK_WORKSPACE + assert dispatched_ws != "" # --------------------------------------------------------------------------- @@ -855,9 +863,9 @@ async def _seed_recoverable_session(qm: Any, sid: str, workspace: str) -> None: async def test_lifespan_default_respawns_all_recovered_sessions_unbounded( monkeypatch: pytest.MonkeyPatch, ) -> None: - """Default (crash_recovery_respawn_limit=None) MUST preserve today's - behaviour exactly: every recovered session is respawned on this boot, - no matter how many there are.""" + """An explicit unbounded ceiling (crash_recovery_respawn_limit=None) + respawns every recovered session on this boot, no matter how many + there are.""" qm = registry.queue_manager sids = [f"sess-unbounded-{i}" for i in range(10)] for sid in sids: @@ -867,19 +875,20 @@ async def test_lifespan_default_respawns_all_recovered_sessions_unbounded( monkeypatch.setattr( registry, "get_or_create", lambda s, w, **kw: spawned.append((s, w)) ) - assert main_module._settings.crash_recovery_respawn_limit is None + monkeypatch.setattr(main_module._settings, "crash_recovery_respawn_limit", None) mock_driver = _patched_lifespan_deps() with ( patch("context_intelligence_server.main.setup_logging"), patch( - "context_intelligence_server.main.AsyncGraphDatabase.driver", + "context_intelligence_server.neo4j_store.AsyncGraphDatabase.driver", return_value=mock_driver, ), patch("context_intelligence_server.main.ensure_neo4j_schema", new=AsyncMock()), ): async with lifespan(main_module.app): - pass + # Recovery is backgrounded -- await it before shutdown cancels it. + await main_module.app.state.boot_task assert {s for s, _w in spawned} == set(sids) @@ -906,14 +915,15 @@ async def test_lifespan_respawn_cap_defers_remainder_and_logs_warning( with ( patch("context_intelligence_server.main.setup_logging"), patch( - "context_intelligence_server.main.AsyncGraphDatabase.driver", + "context_intelligence_server.neo4j_store.AsyncGraphDatabase.driver", return_value=mock_driver, ), patch("context_intelligence_server.main.ensure_neo4j_schema", new=AsyncMock()), caplog.at_level(logging.WARNING, logger="context_intelligence_server"), ): async with lifespan(main_module.app): - pass + # Recovery is backgrounded -- await it before shutdown cancels it. + await main_module.app.state.boot_task # Exactly the cap's worth of sessions were respawned -- never more. assert len(spawned) == 2 @@ -952,13 +962,14 @@ async def test_lifespan_deferred_sessions_untouched_and_recoverable_next_boot( with ( patch("context_intelligence_server.main.setup_logging"), patch( - "context_intelligence_server.main.AsyncGraphDatabase.driver", + "context_intelligence_server.neo4j_store.AsyncGraphDatabase.driver", return_value=mock_driver, ), patch("context_intelligence_server.main.ensure_neo4j_schema", new=AsyncMock()), ): async with lifespan(main_module.app): - pass + # Recovery is backgrounded -- await it before shutdown cancels it. + await main_module.app.state.boot_task assert len(spawned_boot1) == 1 deferred_sids = set(sids) - {s for s, _w in spawned_boot1} @@ -993,13 +1004,14 @@ async def test_lifespan_respawn_cap_zero_defers_everything( with ( patch("context_intelligence_server.main.setup_logging"), patch( - "context_intelligence_server.main.AsyncGraphDatabase.driver", + "context_intelligence_server.neo4j_store.AsyncGraphDatabase.driver", return_value=mock_driver, ), patch("context_intelligence_server.main.ensure_neo4j_schema", new=AsyncMock()), ): async with lifespan(main_module.app): - pass + # Recovery is backgrounded -- await it before shutdown cancels it. + await main_module.app.state.boot_task assert spawned == [] recovered_again = await qm.recover() @@ -1016,10 +1028,10 @@ async def test_lifespan_respawn_cap_zero_defers_everything( async def test_crash_recovery_topup_drains_deferred_tail_across_passes( monkeypatch: pytest.MonkeyPatch, ) -> None: - """The deferred tail is not stranded: with ceiling=2, pass 1 dispatches the - 2 head sessions; once those finish draining (drop out of recover()), pass 2 - dispatches the previously-deferred 2. Live recovered drainers never exceed - the ceiling.""" + """With ceiling=2: pass 1 dispatches the 2 head sessions; once those + finish draining, pass 2 dispatches the previously-deferred 2. Tests + dispatch counting only -- see tests/test_boot_safety.py for the + live-drainer bound.""" qm = registry.queue_manager sids = sorted(f"sess-sweep-{i}" for i in range(4)) for sid in sids: @@ -1029,8 +1041,10 @@ async def test_crash_recovery_topup_drains_deferred_tail_across_passes( monkeypatch.setattr(registry, "get_or_create", lambda s, w, **kw: spawned.append(s)) # Pass 1: only the ceiling's worth (2) are dispatched; the tail is deferred. - dispatched = await main_module._crash_recovery_topup(2) - assert dispatched == 2 + result = await main_module._crash_recovery_topup(2) + assert result.dispatched == 2 + assert result.recovered == 4 + assert result.deferred == 2 assert set(spawned) == set(sids[:2]) # The 2 head sessions finish draining -> commit them to EOF so recover() @@ -1041,8 +1055,10 @@ async def test_crash_recovery_topup_drains_deferred_tail_across_passes( # Pass 2: the previously-DEFERRED tail is now dispatched -- not stranded. spawned.clear() - dispatched = await main_module._crash_recovery_topup(2) - assert dispatched == 2 + result = await main_module._crash_recovery_topup(2) + assert result.dispatched == 2 + assert result.recovered == 2 + assert result.deferred == 0 assert set(spawned) == set(sids[2:]) @@ -1063,14 +1079,16 @@ async def test_lifespan_enables_sweep_under_finite_limit( with ( patch("context_intelligence_server.main.setup_logging"), patch( - "context_intelligence_server.main.AsyncGraphDatabase.driver", + "context_intelligence_server.neo4j_store.AsyncGraphDatabase.driver", return_value=mock_driver, ), patch("context_intelligence_server.main.ensure_neo4j_schema", new=AsyncMock()), caplog.at_level(logging.INFO, logger="context_intelligence_server"), ): async with lifespan(main_module.app): - pass # task is created on entry and cancelled cleanly on exit + # Recovery (incl. sweep-task creation) is backgrounded -- + # await it before shutdown cancels it. + await main_module.app.state.boot_task assert any( "crash_recovery_sweep: enabled" in r.getMessage() for r in caplog.records @@ -1090,14 +1108,15 @@ async def test_lifespan_no_sweep_when_limit_unbounded( with ( patch("context_intelligence_server.main.setup_logging"), patch( - "context_intelligence_server.main.AsyncGraphDatabase.driver", + "context_intelligence_server.neo4j_store.AsyncGraphDatabase.driver", return_value=mock_driver, ), patch("context_intelligence_server.main.ensure_neo4j_schema", new=AsyncMock()), caplog.at_level(logging.INFO, logger="context_intelligence_server"), ): async with lifespan(main_module.app): - pass + # Recovery is backgrounded -- await it before shutdown cancels it. + await main_module.app.state.boot_task assert not any( "crash_recovery_sweep: enabled" in r.getMessage() for r in caplog.records @@ -1120,14 +1139,15 @@ async def test_lifespan_no_sweep_when_interval_zero( with ( patch("context_intelligence_server.main.setup_logging"), patch( - "context_intelligence_server.main.AsyncGraphDatabase.driver", + "context_intelligence_server.neo4j_store.AsyncGraphDatabase.driver", return_value=mock_driver, ), patch("context_intelligence_server.main.ensure_neo4j_schema", new=AsyncMock()), caplog.at_level(logging.INFO, logger="context_intelligence_server"), ): async with lifespan(main_module.app): - pass + # Recovery is backgrounded -- await it before shutdown cancels it. + await main_module.app.state.boot_task assert not any( "crash_recovery_sweep: enabled" in r.getMessage() for r in caplog.records @@ -1149,19 +1169,17 @@ async def test_lifespan_no_sweep_when_interval_zero( # --------------------------------------------------------------------------- -async def test_lifespan_calls_ensure_schema_with_fail_on_data_conflict() -> None: - """Cold start must call ensure_neo4j_schema with fail_on_data_conflict=True +async def test_ensure_schema_ready_calls_schema_init_with_fail_on_data_conflict() -> ( + None +): + """Schema init (now _boot_reconcile's first phase, not synchronous + lifespan) must call ensure_neo4j_schema with fail_on_data_conflict=True -- boot refuses to proceed on a genuine :Node constraint data conflict - (duplicate legacy nodes), mirroring run_repair's contract. Safe at cold - start (nothing flushed yet); the flush path keeps the opposite default.""" - mock_driver = _patched_lifespan_deps() + (duplicate legacy nodes), mirroring run_repair's contract.""" + main_module.app.state.schema_ready = False + main_module.app.state.neo4j_driver = MagicMock() mock_ensure_schema = AsyncMock(return_value=True) with ( - patch("context_intelligence_server.main.setup_logging"), - patch( - "context_intelligence_server.main.AsyncGraphDatabase.driver", - return_value=mock_driver, - ), patch( "context_intelligence_server.main.ensure_neo4j_schema", new=mock_ensure_schema, @@ -1171,29 +1189,27 @@ async def test_lifespan_calls_ensure_schema_with_fail_on_data_conflict() -> None new=AsyncMock(return_value=0), ), ): - async with lifespan(main_module.app): - pass + await main_module._ensure_schema_ready() mock_ensure_schema.assert_awaited_once() + assert mock_ensure_schema.await_args is not None _args, kwargs = mock_ensure_schema.await_args assert kwargs.get("fail_on_data_conflict") is True, ( - "lifespan must opt into fail_on_data_conflict=True -- cold start " + "schema init must opt into fail_on_data_conflict=True -- boot " "fails loud on a genuine data conflict (that contract now applies " "at boot too, not just to run_repair / `doctor --fix`)." ) + assert main_module.app.state.schema_ready is True -async def test_lifespan_raises_on_ensure_schema_data_conflict() -> None: +async def test_ensure_schema_ready_raises_on_data_conflict() -> None: """When ensure_neo4j_schema itself raises (a genuine :Node constraint - data conflict under fail_on_data_conflict=True), lifespan must propagate - the RuntimeError -- boot refuses to start.""" - mock_driver = _patched_lifespan_deps() + data conflict under fail_on_data_conflict=True), _ensure_schema_ready + propagates the RuntimeError -- caught by _boot_reconcile's own + exception-safe wrapper (boot_state.fail), never an ASGI-startup abort.""" + main_module.app.state.schema_ready = False + main_module.app.state.neo4j_driver = MagicMock() with ( - patch("context_intelligence_server.main.setup_logging"), - patch( - "context_intelligence_server.main.AsyncGraphDatabase.driver", - return_value=mock_driver, - ), patch( "context_intelligence_server.main.ensure_neo4j_schema", new=AsyncMock( @@ -1204,21 +1220,18 @@ async def test_lifespan_raises_on_ensure_schema_data_conflict() -> None: ), pytest.raises(RuntimeError, match="doctor --fix"), ): - async with lifespan(main_module.app): - pass + await main_module._ensure_schema_ready() + assert main_module.app.state.schema_ready is False -async def test_lifespan_raises_on_untagged_nodes() -> None: - """On an un-migrated graph (untagged :Node count > 0), startup MUST - raise a RuntimeError naming `doctor --fix` -- boot refuses to start - rather than silently risking write-path duplication.""" - mock_driver = _patched_lifespan_deps() + +async def test_ensure_schema_ready_raises_on_untagged_nodes() -> None: + """On an un-migrated graph (untagged :Node count > 0), schema init MUST + raise a RuntimeError naming `doctor --fix` -- boot refuses to mark + schema ready rather than silently risking write-path duplication.""" + main_module.app.state.schema_ready = False + main_module.app.state.neo4j_driver = MagicMock() with ( - patch("context_intelligence_server.main.setup_logging"), - patch( - "context_intelligence_server.main.AsyncGraphDatabase.driver", - return_value=mock_driver, - ), patch( "context_intelligence_server.main.ensure_neo4j_schema", new=AsyncMock(return_value=True), @@ -1229,25 +1242,20 @@ async def test_lifespan_raises_on_untagged_nodes() -> None: ), pytest.raises(RuntimeError, match="doctor --fix") as exc_info, ): - async with lifespan(main_module.app): - pass + await main_module._ensure_schema_ready() assert "42" in str(exc_info.value), ( f"Expected the untagged count in the error message, got: {exc_info.value}" ) + assert main_module.app.state.schema_ready is False -async def test_lifespan_does_not_raise_on_clean_graph() -> None: +async def test_ensure_schema_ready_does_not_raise_on_clean_graph() -> None: """On a fully-migrated graph (untagged count == 0, no constraint - conflict), startup does NOT raise -- the fail-loud guards are silent - when there is nothing to report.""" - mock_driver = _patched_lifespan_deps() + conflict), schema init does NOT raise and marks schema_ready True.""" + main_module.app.state.schema_ready = False + main_module.app.state.neo4j_driver = MagicMock() with ( - patch("context_intelligence_server.main.setup_logging"), - patch( - "context_intelligence_server.main.AsyncGraphDatabase.driver", - return_value=mock_driver, - ), patch( "context_intelligence_server.main.ensure_neo4j_schema", new=AsyncMock(return_value=True), @@ -1257,10 +1265,10 @@ async def test_lifespan_does_not_raise_on_clean_graph() -> None: new=AsyncMock(return_value=0), ) as mock_count, ): - async with lifespan(main_module.app): - pass + await main_module._ensure_schema_ready() mock_count.assert_awaited_once() + assert main_module.app.state.schema_ready is True async def test_lifespan_does_not_raise_when_health_check_itself_fails( @@ -1274,7 +1282,7 @@ async def test_lifespan_does_not_raise_when_health_check_itself_fails( with ( patch("context_intelligence_server.main.setup_logging"), patch( - "context_intelligence_server.main.AsyncGraphDatabase.driver", + "context_intelligence_server.neo4j_store.AsyncGraphDatabase.driver", return_value=mock_driver, ), patch( @@ -1468,7 +1476,7 @@ async def test_status_includes_neo4j_query_connected_false_when_no_driver( # --------------------------------------------------------------------------- -# /status pipeline metrics block (D3) +# /status pipeline metrics block # --------------------------------------------------------------------------- @@ -1666,7 +1674,7 @@ async def _fake_append(worker_key: str, raw: bytes) -> None: # Temporarily configure a StaticKeyResolver on asgi_app so auth injects contributor_id. # T2: middleware now uses resolver= seam; patch asgi_app.resolver, not asgi_app.keystore. - from context_intelligence_server.auth import StaticKeyResolver # noqa: PLC0415 + from context_intelligence_server.auth import StaticKeyResolver test_token = "test-secret" test_keystore = {hashlib.sha256(test_token.encode()).hexdigest(): "alice"} @@ -1716,7 +1724,7 @@ async def _fake_append(worker_key: str, raw: bytes) -> None: monkeypatch.setattr(main_module.registry.queue_manager, "append", _fake_append) # T2: middleware now uses resolver= seam; patch asgi_app.resolver, not asgi_app.keystore. - from context_intelligence_server.auth import StaticKeyResolver # noqa: PLC0415 + from context_intelligence_server.auth import StaticKeyResolver test_token = "test-secret" test_keystore = {hashlib.sha256(test_token.encode()).hexdigest(): "real-owner"} @@ -1812,6 +1820,7 @@ def _fake_get_or_create( session_id: str, workspace: str, created_by: str | None = None, + **_kwargs: Any, # Absorbs the new `recovered=` keyword, untested here ) -> MagicMock: calls.append( {"session_id": session_id, "workspace": workspace, "created_by": created_by} @@ -1854,6 +1863,7 @@ def _fake_get_or_create( session_id: str, workspace: str, created_by: str | None = None, + **_kwargs: Any, # Absorbs the new `recovered=` keyword, untested here ) -> MagicMock: calls.append( {"session_id": session_id, "workspace": workspace, "created_by": created_by} diff --git a/tests/test_neo4j_driver_sharing.py b/tests/test_neo4j_driver_sharing.py new file mode 100644 index 00000000..47218d4c --- /dev/null +++ b/tests/test_neo4j_driver_sharing.py @@ -0,0 +1,229 @@ +"""Tests for shared, pool-bounded Neo4j driver reuse across sessions. + +Covers: +- Neo4jGraphStore accepts an injected driver and never closes it (owns_driver). +- The self-built path (no injected driver) is unchanged: it owns and closes + its own driver. +- SessionRegistry hands the same driver instance to every per-session store + instead of building one per session. +""" + +from __future__ import annotations + +import asyncio +from unittest.mock import AsyncMock, patch + +import pytest +from context_intelligence_server.neo4j_store import Neo4jGraphStore +from context_intelligence_server.registry import SessionRegistry + +# --------------------------------------------------------------------------- +# Injected-driver seam (owns_driver) +# --------------------------------------------------------------------------- + + +def test_injected_driver_reports_owns_driver_false() -> None: + """A store built with driver= must report owns_driver False.""" + shared_driver = AsyncMock() + + store_a = Neo4jGraphStore(uri="bolt://unused:7687", driver=shared_driver) + store_b = Neo4jGraphStore(uri="bolt://unused:7687", driver=shared_driver) + + assert store_a.owns_driver is False + assert store_b.owns_driver is False + assert store_a._driver is shared_driver + assert store_b._driver is shared_driver + + +@pytest.mark.asyncio +async def test_close_on_injected_driver_does_not_close_it() -> None: + """Closing one store sharing an injected driver must not close the driver + out from under a second store still using it (the #489 safety property).""" + shared_driver = AsyncMock() + + store_a = Neo4jGraphStore(uri="bolt://unused:7687", driver=shared_driver) + store_b = Neo4jGraphStore(uri="bolt://unused:7687", driver=shared_driver) + + await store_a.close() + + shared_driver.close.assert_not_awaited() + # store_b's driver reference is untouched and still the live shared mock -- + # a real driver would still be open and usable by store_b at this point. + assert store_b._driver is shared_driver + + +@pytest.mark.asyncio +async def test_self_built_driver_still_owned_and_closed() -> None: + """With driver=None (default), behavior is unchanged: the store builds + and owns its driver, and close() closes it.""" + with patch( + "context_intelligence_server.neo4j_store.AsyncGraphDatabase" + ) as mock_adb: + mock_driver = AsyncMock() + mock_adb.driver.return_value = mock_driver + + store = Neo4jGraphStore(uri="bolt://localhost:7687", auth=("u", "p")) + assert store.owns_driver is True + + await store.close() + mock_driver.close.assert_awaited_once() + + +# --------------------------------------------------------------------------- +# Registry driver reuse +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_registry_shares_one_driver_across_sessions(monkeypatch) -> None: + """get_or_create must hand the SAME driver object to every per-session + Neo4jGraphStore -- reuse, not a new driver per session_id.""" + reg = SessionRegistry() + + built_drivers: list[object] = [] + + def fake_build_bounded_driver(config, **kwargs): + driver = AsyncMock() + built_drivers.append(driver) + return driver + + monkeypatch.setattr( + "context_intelligence_server.registry.build_bounded_neo4j_driver", + fake_build_bounded_driver, + ) + + worker_a = reg.get_or_create("session-a", "/workspace/a") + worker_b = reg.get_or_create("session-b", "/workspace/b") + + # Exactly one driver was ever built, and both sessions' stores share it. + assert len(built_drivers) == 1 + assert worker_a.services.graph._driver is built_drivers[0] + assert worker_b.services.graph._driver is built_drivers[0] + assert worker_a.services.graph.owns_driver is False + assert worker_b.services.graph.owns_driver is False + + for worker in (worker_a, worker_b): + if worker.task is not None: + worker.task.cancel() + + +def _spy_driver_factory(reg: SessionRegistry, monkeypatch) -> list[dict]: + """Patch the registry's driver factory to record every build call. + + Returns the list of recorded builds; each entry is + ``{"driver": , "kwargs": {...}}``. + """ + builds: list[dict] = [] + + def fake_build_bounded_driver(config, **kwargs): + driver = AsyncMock() + builds.append({"driver": driver, "kwargs": kwargs}) + return driver + + monkeypatch.setattr( + "context_intelligence_server.registry.build_bounded_neo4j_driver", + fake_build_bounded_driver, + ) + return builds + + +def _cancel_workers(reg: SessionRegistry) -> None: + """Cancel any real drain tasks started by get_or_create (test cleanup).""" + for worker in reg._workers.values(): + if worker.task is not None: + worker.task.cancel() + + +@pytest.mark.asyncio +async def test_n_sessions_build_exactly_one_driver(monkeypatch) -> None: + """The core leak-gone proof: N distinct sessions must build the driver + exactly ONCE, not once per session_id (which was the leak).""" + reg = SessionRegistry() + builds = _spy_driver_factory(reg, monkeypatch) + + n = 30 + workers = [reg.get_or_create(f"session-{i}", f"/workspace/{i}") for i in range(n)] + + assert len(builds) == 1, ( + f"expected exactly 1 driver build across {n} sessions, " + f"got {len(builds)} (a per-session build is the leak)" + ) + the_driver = builds[0]["driver"] + for worker in workers: + assert worker.services.graph._driver is the_driver + assert worker.services.graph.owns_driver is False + + _cancel_workers(reg) + + +@pytest.mark.asyncio +async def test_shared_driver_built_with_bounded_kwargs(monkeypatch) -> None: + """The single shared driver must be built WITH the bounded pool kwargs + (default pool size 50, lifetime 3600.0s) -- an unbounded build is the leak.""" + reg = SessionRegistry() + builds = _spy_driver_factory(reg, monkeypatch) + + reg.get_or_create("session-1", "/workspace/1") + + assert len(builds) == 1 + kwargs = builds[0]["kwargs"] + assert kwargs["max_connection_pool_size"] == 50 + assert kwargs["max_connection_lifetime"] == 3600.0 + + _cancel_workers(reg) + + +@pytest.mark.asyncio +async def test_concurrent_first_sessions_build_exactly_one_driver(monkeypatch) -> None: + """Racing many get_or_create calls as the FIRST sessions must still build + exactly one driver -- the lazy build is synchronous with no await between + the None-check and the assignment, so concurrent coroutines cannot double-build.""" + reg = SessionRegistry() + builds = _spy_driver_factory(reg, monkeypatch) + + async def make(i: int) -> None: + # Wrap the sync get_or_create so many run concurrently under gather. + reg.get_or_create(f"session-{i}", f"/workspace/{i}") + + await asyncio.gather(*(make(i) for i in range(40))) + + assert len(builds) == 1, ( + f"concurrent first-sessions raced into {len(builds)} driver builds; " + "the lazy build must be single-shot" + ) + + _cancel_workers(reg) + + +@pytest.mark.asyncio +async def test_close_neo4j_driver_reclaims_once_and_is_idempotent(monkeypatch) -> None: + """After sessions run, close_neo4j_driver() must close the shared driver + exactly once and clear it; a second call is a safe no-op.""" + reg = SessionRegistry() + builds = _spy_driver_factory(reg, monkeypatch) + + reg.get_or_create("session-1", "/workspace/1") + reg.get_or_create("session-2", "/workspace/2") + assert len(builds) == 1 + shared_driver = builds[0]["driver"] + + await reg.close_neo4j_driver() + shared_driver.close.assert_awaited_once() + assert reg._neo4j_driver is None + + # Second call: no driver left to close, must not raise or double-close. + await reg.close_neo4j_driver() + shared_driver.close.assert_awaited_once() + assert reg._neo4j_driver is None + + _cancel_workers(reg) + + +@pytest.mark.asyncio +async def test_close_neo4j_driver_none_safe_when_no_session_ran() -> None: + """close_neo4j_driver() must be safe when no session ever built a driver.""" + reg = SessionRegistry() + assert reg._neo4j_driver is None + # Must not raise. + await reg.close_neo4j_driver() + assert reg._neo4j_driver is None 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..d7e743e6 100644 --- a/tests/test_queue_manager.py +++ b/tests/test_queue_manager.py @@ -1,12 +1,16 @@ -"""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, + Verdict, +) @pytest.fixture @@ -22,13 +26,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 +235,82 @@ async def test_commit_is_atomic_no_temp_leftover(qm, tmp_path): assert list(qdir.glob("*.tmp")) == [] +# _read_committed_offset accepts the bare-int form and the legacy JSON offset +# document; commit() still writes bare int. 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_legacy_json_cursor(qm): + """A legacy JSON cursor document parses to its integer "offset" field.""" + qm._offset_path("s1").write_text( + '{"v":1,"offset":12345,"cursor":{"dl2":{"a":1},"dl3":{}}}', + encoding="utf-8", + ) + assert qm._read_committed_offset("s1") == 12345 + + +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_read_batch_drains_session_with_legacy_json_offset(qm): + """A session with a legacy JSON-cursor .offset drains via the normal + read path (read_batch) with no ValueError -- the fix must reach the + hot path, not just the private helper.""" + await qm.append("s1", b"a") + await qm.append("s1", b"b") + qm._offset_path("s1").write_text('{"v":1,"offset":2,"cursor":{}}', encoding="utf-8") + + batch = await qm.read_batch("s1", max_items=10) + + assert batch.start_offset == 2 + assert batch.lines == [b"b"] + + +async def test_classify_session_with_legacy_json_offset_is_not_corrupt(qm): + """A session with a legacy JSON-cursor .offset must not be classified + as an unreadable/corrupt offset -- it should classify the same as an + equivalent bare-int offset (drained, in this fully-committed case).""" + await qm.append("s1", b"a\n" * 0 + b"a") # single record "a" + line = b"a\n" + qm._offset_path("s1").write_text( + f'{{"v":1,"offset":{len(line)},"cursor":{{}}}}', encoding="utf-8" + ) + + classification = await qm.classify_session( + "s1", head_is_resumable=lambda _raw: True + ) + + assert classification.verdict == Verdict.DRAINED + + 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 +465,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 +541,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_steady_state_reclaim.py b/tests/test_steady_state_reclaim.py new file mode 100644 index 00000000..2d99bbce --- /dev/null +++ b/tests/test_steady_state_reclaim.py @@ -0,0 +1,798 @@ +"""Steady-state queue reclaim + dead-letter retention. + +Covers undrained-tail protection, crash-atomic compaction ordering, +dead-letter retention/expiry, /status non-blocking during compaction, +reclaim-on-finalize composing with compaction, and boot-vs-sweep expiry +accounting. No real Neo4j is used anywhere in this file (see +``tests/neo4j/test_steady_state_reclaim_neo4j.py`` for the Neo4j-backed cases). +""" + +from __future__ import annotations + +import asyncio +import json +import logging +import os +import time +from pathlib import Path +from typing import Any +from unittest.mock import AsyncMock + +import httpx +import pytest + +import context_intelligence_server.main as main_module +import context_intelligence_server.queue_manager as queue_manager_module +from context_intelligence_server.config import Settings +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 + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _fixed(i: int) -> bytes: + """A 9-byte, fixed-width record payload (10 bytes on disk incl. '\\n'). + + Fixed width makes byte-offset arithmetic in the crash-atomicity tests + exact and easy to reason about (event i occupies bytes [10*i, 10*i+10)). + """ + return f"{i:09d}".encode("ascii") + + +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" + ) + + +def _seed_dead(tmp_path: Path, key: str, content: str = "") -> Path: + path = tmp_path / f"{key}.dead.jsonl" + path.write_text(content, encoding="utf-8") + return path + + +def _seed_log(tmp_path: Path, key: str, content: bytes = b"") -> Path: + path = tmp_path / f"{key}.log" + path.write_bytes(content) + return path + + +_G_TS = "2026-08-21T00:00:00+00:00" + + +def _dead_record(payload: str, error: str = "boom") -> str: + return json.dumps({"ts": time.time(), "error": error, "payload": payload}) + "\n" + + +# --------------------------------------------------------------------------- +# (b) undrained tail is never prefix-reclaimed past the committed offset +# --------------------------------------------------------------------------- + + +async def test_b_undrained_tail_never_reclaimed_past_committed_c_less_than_tail( + tmp_path: Path, +) -> None: + """C < E-C: commit 40 of 100 events, compact, and prove events 41..100 + survive in order, untouched, with committed rebased to 0.""" + qm = QueueManager(queues_dir=tmp_path) + sid = "s-tail-c-lt-tail" + events = [_fixed(i) for i in range(100)] + for ev in events: + await qm.append(sid, ev) + + first_batch = await qm.read_batch(sid, max_items=40) + assert len(first_batch.records) == 40 + await qm.commit(sid, first_batch.end_offset) + c = first_batch.end_offset + log_path = tmp_path / f"{sid}.log" + e = log_path.stat().st_size + assert c < e - c # this sub-case: C < E-C + + reclaimed = await qm.compact_committed_prefix(sid, 0) + assert reclaimed == c + + assert log_path.stat().st_size == e - c + assert qm._read_committed_offset(sid) == 0 + + remaining = await qm.read_batch(sid, max_items=1000) + assert [r.raw for r in remaining.records] == events[40:] + + +async def test_b_undrained_tail_never_reclaimed_past_committed_c_greater_than_tail( + tmp_path: Path, +) -> None: + """C > E-C: commit 70 of 100 events -- the tail is now the SMALLER side.""" + qm = QueueManager(queues_dir=tmp_path) + sid = "s-tail-c-gt-tail" + events = [_fixed(i) for i in range(100)] + for ev in events: + await qm.append(sid, ev) + + first_batch = await qm.read_batch(sid, max_items=70) + assert len(first_batch.records) == 70 + await qm.commit(sid, first_batch.end_offset) + c = first_batch.end_offset + log_path = tmp_path / f"{sid}.log" + e = log_path.stat().st_size + assert c > e - c # this sub-case: C > E-C + + reclaimed = await qm.compact_committed_prefix(sid, 0) + assert reclaimed == c + assert log_path.stat().st_size == e - c + assert qm._read_committed_offset(sid) == 0 + + remaining = await qm.read_batch(sid, max_items=1000) + assert [r.raw for r in remaining.records] == events[70:] + + +async def test_b_below_min_prefix_bytes_is_a_noop(tmp_path: Path) -> None: + """C below min_prefix_bytes bails without touching the file at all.""" + qm = QueueManager(queues_dir=tmp_path) + sid = "s-below-threshold" + for i in range(10): + await qm.append(sid, _fixed(i)) + batch = await qm.read_batch(sid, max_items=5) + await qm.commit(sid, batch.end_offset) + log_path = tmp_path / f"{sid}.log" + before = log_path.read_bytes() + + reclaimed = await qm.compact_committed_prefix(sid, min_prefix_bytes=10_000) + assert reclaimed == 0 + assert log_path.read_bytes() == before + assert qm._read_committed_offset(sid) == batch.end_offset + + +# --------------------------------------------------------------------------- +# (c) crash atomicity +# --------------------------------------------------------------------------- + + +async def test_c_mid_copy_oserror_is_a_pure_noop( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """An OSError raised mid-copy (Precision 1) must not mutate anything and + must not escape -- it is caught and the method returns 0.""" + qm = QueueManager(queues_dir=tmp_path) + sid = "s-mid-copy-fault" + for i in range(9): + await qm.append(sid, _fixed(i)) + batch = await qm.read_batch(sid, max_items=3) + await qm.commit(sid, batch.end_offset) + + log_path = tmp_path / f"{sid}.log" + offset_path = tmp_path / f"{sid}.offset" + log_before = log_path.read_bytes() + offset_before = offset_path.read_text(encoding="utf-8") + assert offset_before.strip() == str(batch.end_offset) + + def _raise(fd: int, data: bytes) -> None: + raise OSError("simulated mid-copy failure") + + monkeypatch.setattr( + queue_manager_module.QueueManager, "_write_all", staticmethod(_raise) + ) + + reclaimed = await qm.compact_committed_prefix(sid, 0) # must not raise + + assert reclaimed == 0 + assert log_path.read_bytes() == log_before + assert offset_path.read_text(encoding="utf-8") == offset_before + + +async def test_c_window2_offset_rebased_before_log_replaced_bounded_redrive( + tmp_path: Path, +) -> None: + """Simulates a crash after the offset was rebased to 0 but before the log + was replaced. A reader resuming from this on-disk state must see a + bounded re-drive (the committed prefix duplicated) -- never a loss.""" + qm = QueueManager(queues_dir=tmp_path) + sid = "s-window2" + events = [_fixed(i) for i in range(9)] + for ev in events: + await qm.append(sid, ev) + batch = await qm.read_batch(sid, max_items=3) + await qm.commit(sid, batch.end_offset) # committed = 3 events (30 bytes) + + # Simulate the crash: offset already rebased to 0 (step 5 completed), + # log NOT yet replaced (step 6 never ran). + offset_path = tmp_path / f"{sid}.offset" + offset_path.write_text("0", encoding="utf-8") + + resumed = await qm.read_batch(sid, max_items=100) + resumed_raw = [r.raw for r in resumed.records] + + # Bounded re-drive: every original event is present -- zero loss. + assert resumed_raw == events + # the duplicated set is exactly the already-committed prefix + assert resumed_raw[:3] == events[:3] + + +async def test_c_control_rejected_log_then_offset_order_loses_data( + tmp_path: Path, +) -> None: + """CONTROL: applies the alternative log-then-offset ordering by hand and + stops mid-window (log replaced, offset not yet rewritten), proving that + order silently drops undrained data -- why offset-before-log is used.""" + qm = QueueManager(queues_dir=tmp_path) + sid = "s-rejected-order" + events = [_fixed(i) for i in range(9)] + for ev in events: + await qm.append(sid, ev) + batch = await qm.read_batch(sid, max_items=3) + await qm.commit(sid, batch.end_offset) # committed = 30 bytes (C == 30) + + log_path = tmp_path / f"{sid}.log" + c = qm._read_committed_offset(sid) + e = log_path.stat().st_size + assert c <= e - c # precondition for the table's "lands inside the tail" case + + # alternative order, applied by hand: replace the log with the tail first + tail_bytes = log_path.read_bytes()[c:e] + rejected_tmp = tmp_path / f"{sid}.log.rejected.tmp" + rejected_tmp.write_bytes(tail_bytes) + os.replace(rejected_tmp, log_path) + # crash here -- offset file still says C (unchanged) + + resumed = await qm.read_batch(sid, max_items=100) + resumed_raw = [r.raw for r in resumed.records] + + # silent loss: the drainer resumes at byte C inside the already-shortened + # file, skipping the first C bytes of real undrained data + skipped_events = events[3:6] + surviving_events = events[6:9] + assert resumed_raw == surviving_events + assert resumed_raw != events[3:9], ( + "control failed to reproduce the loss: the rejected ordering was " + "expected to silently drop the leading undrained events" + ) + for skipped in skipped_events: + assert skipped not in resumed_raw + + +# --------------------------------------------------------------------------- +# (i) os.replace failure restores the offset -- pure no-op, zero drift +# --------------------------------------------------------------------------- + + +async def test_i_replace_failure_restores_offset_zero_accounting_drift( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + qm = QueueManager(queues_dir=tmp_path) + sid = "s-replace-fails" + events = [_fixed(i) for i in range(9)] + for ev in events: + await qm.append(sid, ev) + batch = await qm.read_batch(sid, max_items=3) + await qm.commit(sid, batch.end_offset) + c = batch.end_offset + + log_path = tmp_path / f"{sid}.log" + offset_path = tmp_path / f"{sid}.offset" + log_before = log_path.read_bytes() + + real_replace = os.replace + + def _flaky_replace(src: Any, dst: Any) -> None: + # Fail ONLY the log replace (step 6); let the offset writes through. + if str(dst) == str(log_path): + raise OSError("simulated os.replace failure on the log") + real_replace(src, dst) + + monkeypatch.setattr( + queue_manager_module.os, "replace", _flaky_replace, raising=True + ) + + with caplog.at_level(logging.ERROR): + reclaimed = await qm.compact_committed_prefix(sid, 0) # must not raise + + assert reclaimed == 0 + # Offset restored to C -- a pure no-op, not a re-drive. + assert offset_path.read_text(encoding="utf-8").strip() == str(c) + # Log completely untouched. + assert log_path.read_bytes() == log_before + assert any( + "compact_replace_failed" in r.getMessage() + and "offset_restored" in r.getMessage() + for r in caplog.records + ) + + # Load-bearing proof of zero accounting drift: the NEXT read_batch + # resumes from the restored offset C, not from 0 -- no duplicate + # processing, hence no double record_written downstream. + resumed = await qm.read_batch(sid, max_items=100) + assert [r.raw for r in resumed.records] == events[3:] + + +async def test_i_double_replace_failure_logs_restore_failed_honestly( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + """If the RESTORE itself also fails, the honest (documented) fallback is + a logged `compact_restore_failed ... redrive_expected=true` -- never a + silent success claim.""" + qm = QueueManager(queues_dir=tmp_path) + sid = "s-double-fail" + for i in range(9): + await qm.append(sid, _fixed(i)) + batch = await qm.read_batch(sid, max_items=3) + await qm.commit(sid, batch.end_offset) + + log_path = tmp_path / f"{sid}.log" + + def _always_raise(src: Any, dst: Any) -> None: + # let the first offset rebase-to-0 write through, then fail both the + # log replace and the subsequent restore-to-C write + src_content = ( + Path(src).read_text(encoding="utf-8") if Path(src).exists() else "" + ) + if str(dst) == str(log_path): + raise OSError("simulated persistent log-replace failure") + if str(dst).endswith(".offset") and src_content != "0": + raise OSError("simulated persistent offset-restore failure") + + monkeypatch.setattr(queue_manager_module.os, "replace", _always_raise, raising=True) + + with caplog.at_level(logging.ERROR): + reclaimed = await qm.compact_committed_prefix(sid, 0) # must not raise + + assert reclaimed == 0 + assert any( + "compact_restore_failed" in r.getMessage() + and "redrive_expected=true" in r.getMessage() + for r in caplog.records + ) + + +# --------------------------------------------------------------------------- +# (j) a large tail no longer blocks reclaiming the committed prefix +# --------------------------------------------------------------------------- + + +async def test_j_large_tail_does_not_block_prefix_reclaim( + tmp_path: Path, +) -> None: + qm = QueueManager(queues_dir=tmp_path) + sid = "s-huge-tail" + # 20 events * 10 bytes = 200 bytes total. + for i in range(20): + await qm.append(sid, _fixed(i)) + batch = await qm.read_batch(sid, max_items=2) # C = 20 bytes + await qm.commit(sid, batch.end_offset) + log_path = tmp_path / f"{sid}.log" + assert batch.end_offset == 20 + + # Tail is 180 bytes -- well above min_prefix_bytes (10) and above what + # used to be a tail cap. The prefix must still be reclaimed. + reclaimed = await qm.compact_committed_prefix(sid, min_prefix_bytes=10) + + assert reclaimed == 20 + assert log_path.stat().st_size == 180 + assert qm._read_committed_offset(sid) == 0 + + # A concurrent append for the same key must complete promptly. + await asyncio.wait_for(qm.append(sid, _fixed(999)), timeout=2.0) + assert log_path.stat().st_size == 190 + + +# --------------------------------------------------------------------------- +# (e) dead-letter retention: log-less keys, whole-file mtime expiry +# --------------------------------------------------------------------------- + + +async def test_e_dead_letters_older_than_retention_expired_newer_kept( + tmp_path: Path, +) -> None: + qm = QueueManager(queues_dir=tmp_path) + now = time.time() + retention = 30 * 86400.0 + + old_path = _seed_dead(tmp_path, "old-log-less", _dead_record("p1")) + old_mtime = now - retention - 3600 # 1h past the window + os.utime(old_path, (old_mtime, old_mtime)) + + new_path = _seed_dead(tmp_path, "new-log-less", _dead_record("p2")) + new_mtime = now - 3600 # 1h old, well within the window + os.utime(new_path, (new_mtime, new_mtime)) + + # Old on age, but has a LIVE .log -- must NEVER be touched (boot-safety rule). + log_backed_dead = _seed_dead( + tmp_path, "old-but-log-present", _dead_record("p3") + _dead_record("p4") + ) + os.utime(log_backed_dead, (old_mtime, old_mtime)) + _seed_log(tmp_path, "old-but-log-present", b"") + + result = await qm.expire_dead_letters(now, retention, enabled=True) + + assert not old_path.exists() + assert new_path.exists() + assert log_backed_dead.exists() + assert result["expired_keys"] == 1 + assert result["expired_records"] == 1 # old-log-less had 1 record + assert result["failed"] == 0 + + +async def test_e_dry_run_deletes_nothing_but_still_classifies( + tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: + qm = QueueManager(queues_dir=tmp_path) + now = time.time() + retention = 30 * 86400.0 + old_path = _seed_dead(tmp_path, "would-expire", _dead_record("p")) + old_mtime = now - retention - 3600 + os.utime(old_path, (old_mtime, old_mtime)) + + with caplog.at_level(logging.WARNING): + result = await qm.expire_dead_letters(now, retention, enabled=False) + + assert old_path.exists() # nothing unlinked + assert result["expired_keys"] == 0 # dry-run never counts as "expired" + assert any( + "dead_letter_expired" in r.getMessage() and "action=dry_run" in r.getMessage() + for r in caplog.records + ) + + +async def test_e_retention_zero_disables_expiry(tmp_path: Path) -> None: + qm = QueueManager(queues_dir=tmp_path) + now = time.time() + old_path = _seed_dead(tmp_path, "ancient", _dead_record("p")) + os.utime(old_path, (now - 10_000_000, now - 10_000_000)) + + result = await qm.expire_dead_letters(now, retention_seconds=0, enabled=True) + + assert old_path.exists() + assert result == { + "expired_keys": 0, + "expired_records": 0, + "expired_bytes": 0, + "failed": 0, + } + + +# --------------------------------------------------------------------------- +# (k) dead-letter expiry works under SHIPPED DEFAULTS +# --------------------------------------------------------------------------- + + +async def test_k_expiry_is_opt_in_disabled_under_shipped_defaults( + tmp_path: Path, +) -> None: + """No overrides at all: `dead_letter_expiry_enabled` ships False -- an + un-recovered dead-letter's last surviving copy must never be silently + deleted out of the box. Explicit opt-in still deletes (mechanism + unchanged, only the default flipped).""" + settings = Settings() + assert settings.reclaim_enabled is False + assert settings.dead_letter_expiry_enabled is False + + qm = QueueManager(queues_dir=tmp_path) + now = time.time() + old_path = _seed_dead(tmp_path, "shipped-default-no-expire", _dead_record("p")) + old_mtime = now - settings.dead_letter_retention_seconds - 3600 + os.utime(old_path, (old_mtime, old_mtime)) + + result = await qm.expire_dead_letters( + now, settings.dead_letter_retention_seconds, settings.dead_letter_expiry_enabled + ) + + assert old_path.exists(), "shipped default must never auto-delete the last copy" + assert result["expired_keys"] == 0 # dry-run only, no override + + # Opt-in (enabled=True) still deletes -- the mechanism itself is untouched. + result = await qm.expire_dead_letters( + now, settings.dead_letter_retention_seconds, enabled=True + ) + assert not old_path.exists() + assert result["expired_keys"] == 1 + + +async def test_k_dead_letter_expiry_enabled_false_still_dry_runs( + tmp_path: Path, +) -> None: + qm = QueueManager(queues_dir=tmp_path) + now = time.time() + old_path = _seed_dead(tmp_path, "would-expire-2", _dead_record("p")) + os.utime(old_path, (now - (31 * 86400.0), now - (31 * 86400.0))) + + result = await qm.expire_dead_letters(now, 30 * 86400.0, enabled=False) + assert old_path.exists() + assert result["expired_keys"] == 0 + + +# --------------------------------------------------------------------------- +# (m) boot-vs-sweep dead-letter accounting: boot expiry runs before seed +# counting and must never call record_purged (nothing was counted yet); +# sweep expiry runs on live counters and must call it, or residual drifts. +# --------------------------------------------------------------------------- + + +async def test_m_boot_phase_expiry_never_calls_record_purged( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr( + main_module.registry, "_queue_manager", QueueManager(queues_dir=tmp_path) + ) + main_module.registry._accepted_total = 0 + main_module.registry._written_total = 0 + + monkeypatch.setattr(main_module._settings, "reclaim_enabled", False) + monkeypatch.setattr(main_module._settings, "dead_letter_expiry_enabled", True) + monkeypatch.setattr( + main_module._settings, "dead_letter_retention_seconds", 30 * 86400.0 + ) + # Prevent a background sweep task from starting at the end of + # _boot_reconcile -- it would otherwise run concurrently and race this + # test's own assertions (and outlive the test). + monkeypatch.setattr( + main_module._settings, "crash_recovery_sweep_interval_seconds", 0 + ) + + old_path = _seed_dead( + tmp_path, "boot-expire-me", _dead_record("p1") + _dead_record("p2") + ) + old_mtime = time.time() - main_module._settings.dead_letter_retention_seconds - 3600 + os.utime(old_path, (old_mtime, old_mtime)) + + purge_calls: list[int] = [] + monkeypatch.setattr( + main_module.registry, "record_purged", lambda n: purge_calls.append(n) + ) + + await main_module._boot_reconcile() + + assert not old_path.exists(), "boot phase never expired the dead-letter file" + assert purge_calls == [], ( + "boot's expire step must NEVER call record_purged -- it runs BEFORE " + "recovery_seed_counts, so expired lines are simply never counted " + "into accepted_seed in the first place (main.py's own comment at " + "the boot expire call site); calling record_purged here would " + "subtract records that were never added, driving residual negative" + ) + metrics = await main_module.registry.pipeline_metrics() + assert metrics["residual"] == 0 + assert metrics["degraded"] is False + + +async def test_m_sweep_tick_expiry_applies_record_purged( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr( + main_module.registry, "_queue_manager", QueueManager(queues_dir=tmp_path) + ) + main_module.registry._accepted_total = 0 + main_module.registry._written_total = 0 + + monkeypatch.setattr(main_module._settings, "dead_letter_expiry_enabled", True) + monkeypatch.setattr( + main_module._settings, "dead_letter_retention_seconds", 30 * 86400.0 + ) + + old_path = _seed_dead( + tmp_path, + "sweep-expire-me", + _dead_record("p1") + _dead_record("p2") + _dead_record("p3"), + ) + old_mtime = time.time() - main_module._settings.dead_letter_retention_seconds - 3600 + os.utime(old_path, (old_mtime, old_mtime)) + + real_record_purged = main_module.registry.record_purged + purge_calls: list[int] = [] + + def _spy_record_purged(n: int) -> None: + purge_calls.append(n) + real_record_purged(n) + + monkeypatch.setattr(main_module.registry, "record_purged", _spy_record_purged) + # Seed accepted_total > written_total so a real record_purged call is + # OBSERVABLE in the counters (not masked by record_purged's own + # accepted-can-never-fall-below-written clamp at a 0/0 baseline). + main_module.registry.record_accepted(3) + + sweep_task = asyncio.create_task(main_module._crash_recovery_sweep_loop(0, 100)) + try: + for _ in range(300): + if purge_calls: + break + await asyncio.sleep(0.02) + finally: + sweep_task.cancel() + try: + await sweep_task + except asyncio.CancelledError: + pass + + assert purge_calls == [3], ( + "sweep-tick expiry must call record_purged with the expired-record " + "count -- unlike boot, the sweep runs on LIVE counters and must " + "keep `accepted` conserved (main.py's own comment at the sweep " + "expire call site)" + ) + assert not old_path.exists() + assert main_module.registry.pipeline_counters()["accepted_total"] == 0 + metrics = await main_module.registry.pipeline_metrics() + assert metrics["residual"] == 0 + assert metrics["degraded"] is False + + +# --------------------------------------------------------------------------- +# (f) neither compaction nor expiry can block /status +# --------------------------------------------------------------------------- + + +async def test_f_status_not_blocked_by_an_in_progress_compaction( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Hold a key's file_lock (as an in-progress compaction would) from a + background thread; /status must still return promptly -- it never + acquires guard.file_lock for any key.""" + monkeypatch.setattr( + main_module.registry, "_queue_manager", QueueManager(queues_dir=tmp_path) + ) + qm = main_module.registry.queue_manager + sid = "s-status-lock" + await qm.append(sid, _fixed(0)) + await qm.commit(sid, 10) + + with qm._guard(sid) as guard: + loop = asyncio.get_event_loop() + held = asyncio.Event() + release = asyncio.Event() + + def _hold_lock() -> None: + guard.file_lock.acquire() + loop.call_soon_threadsafe(held.set) + while not release.is_set(): + time.sleep(0.01) + guard.file_lock.release() + + thread_task = loop.run_in_executor(None, _hold_lock) + await asyncio.wait_for(held.wait(), timeout=5.0) + + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=main_module.app), base_url="http://test" + ) as c: + response = await asyncio.wait_for(c.get("/status"), timeout=2.0) + + assert response.status_code == 200 + + release.set() + await thread_task + + +async def test_f_status_not_blocked_by_a_concurrent_dead_letter_expiry( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr( + main_module.registry, "_queue_manager", QueueManager(queues_dir=tmp_path) + ) + qm = main_module.registry.queue_manager + for i in range(50): + _seed_dead(tmp_path, f"dead-{i}", _dead_record("p")) + + expire_task = asyncio.create_task( + qm.expire_dead_letters(time.time(), 30 * 86400.0, enabled=True) + ) + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=main_module.app), base_url="http://test" + ) as c: + response = await asyncio.wait_for(c.get("/status"), timeout=2.0) + assert response.status_code == 200 + await expire_task + + +# --------------------------------------------------------------------------- +# (g) regression: reclaim-on-finalize unchanged, composes with compaction +# --------------------------------------------------------------------------- + + +@pytest.fixture +async def reg_qm(tmp_path: Path): + reg = SessionRegistry() + reg._queue_manager = QueueManager(queues_dir=tmp_path) + reg._write_semaphore = asyncio.Semaphore(8) + reg._max_delivery_attempts = 3 + yield reg, reg._queue_manager + for w in list(reg._workers.values()): + if w.task and not w.task.done(): + w.task.cancel() + tasks = [w.task for w in reg._workers.values() if w.task and not w.task.done()] + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + + +async def test_g_reclaim_on_finalize_still_removes_log_and_offset( + reg_qm: tuple[SessionRegistry, QueueManager], +) -> None: + reg, qm = reg_qm + sid = "s-finalize-regression" + worker = SessionWorker( + session_id=sid, workspace="/ws", services=HookStateService(workspace="/ws") + ) + worker.services.graph.flush = AsyncMock() # type: ignore[method-assign] + worker.services.graph.close = AsyncMock() # type: ignore[method-assign] + reg._register_for_test(worker) + + for i in range(5): + await qm.append( + sid, + _line("tool:pre", "/ws", {"session_id": sid, "i": i, "timestamp": _G_TS}), + ) + await qm.append( + sid, _line("session:end", "/ws", {"session_id": sid, "timestamp": _G_TS}) + ) + + reg.start_drain(worker) + assert worker.task is not None + await asyncio.wait_for(asyncio.shield(worker.task), timeout=10.0) + + assert not (qm.queues_dir / f"{sid}.log").exists() + assert not (qm.queues_dir / f"{sid}.offset").exists() + + +async def test_g_compaction_then_finalize_compose_correctly( + reg_qm: tuple[SessionRegistry, QueueManager], +) -> None: + """A compaction that already ran on an open session must not interfere + with a LATER finalize -- both paths must compose.""" + reg, qm = reg_qm + sid = "s-compact-then-finalize" + worker = SessionWorker( + session_id=sid, workspace="/ws", services=HookStateService(workspace="/ws") + ) + worker.services.graph.flush = AsyncMock() # type: ignore[method-assign] + worker.services.graph.close = AsyncMock() # type: ignore[method-assign] + reg._register_for_test(worker) + + for i in range(5): + await qm.append( + sid, + _line("tool:pre", "/ws", {"session_id": sid, "i": i, "timestamp": _G_TS}), + ) + + reg.start_drain(worker) + assert worker.task is not None + + async def _drained() -> bool: + b = await qm.read_batch(sid, max_items=1) + return not b.records + + for _ in range(200): + if await _drained(): + break + await asyncio.sleep(0.02) + assert await _drained() + + # The AUTOMATIC path (Trigger I, wired via reg_qm's real drain loop + + # the safe_settings proxy) has already been compacting on every idle + # poll tick since backlog hit 0 -- assert the composition it produced: + # the log has already collapsed to its (empty) undrained tail, entirely + # without a manual call. + log_path = qm.queues_dir / f"{sid}.log" + + async def _log_collapsed() -> bool: + return not log_path.exists() or log_path.stat().st_size == 0 + + for _ in range(100): + if await _log_collapsed(): + break + await asyncio.sleep(0.02) + assert await _log_collapsed() + + # A manual call on top finds nothing left -- idempotent (I7). + reclaimed = await qm.compact_committed_prefix(sid, 0) + assert reclaimed == 0 + assert not log_path.exists() or log_path.stat().st_size == 0 + + await qm.append( + sid, _line("session:end", "/ws", {"session_id": sid, "timestamp": _G_TS}) + ) + await asyncio.wait_for(asyncio.shield(worker.task), timeout=10.0) + + assert not (qm.queues_dir / f"{sid}.log").exists() + assert not (qm.queues_dir / f"{sid}.offset").exists() 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/tests/test_writer_lease.py b/tests/test_writer_lease.py new file mode 100644 index 00000000..760c1082 --- /dev/null +++ b/tests/test_writer_lease.py @@ -0,0 +1,954 @@ +"""Writer-lease detector tests. No real Neo4j is used anywhere in this file. + +The detector is a DETECTOR, not a mutex: it acquires a `.writer.lease` +sibling before boot recovery runs; `detect` mode latches + surfaces +conflicts without ever refusing boot; `enforce` mode additionally refuses a +fresh foreign lease. It never constructs a QueueManager, never crash-loops +on a share fault or hung mount, and bounds lease I/O to a private executor +so a stalled detector can never starve the append/commit path. +""" + +from __future__ import annotations + +import asyncio +import errno +import json +import threading +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Literal +from unittest.mock import AsyncMock, MagicMock, patch + +import context_intelligence_server.main as main_module +import httpx +import pytest +from context_intelligence_server.config import Settings +from context_intelligence_server.main import lifespan +from context_intelligence_server.queue_manager import QueueManager +from context_intelligence_server.status import boot_state +from context_intelligence_server.writer_lease import ( + WriterLease, + WriterLeaseBusy, + WriterLeaseConflict, +) + +pytestmark = pytest.mark.integration + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +@dataclass +class _LeaseSettings: + """A minimal duck-typed settings stub carrying only the six fields + ``WriterLease.acquire`` reads. Isolates the lease-mechanics tests from + the full ``Settings`` model (which is exercised directly by the config + validator test at the bottom of this file).""" + + writer_lease_mode: Literal["off", "detect", "enforce"] = "detect" + writer_lease_heartbeat_seconds: float = 5.0 + writer_lease_staleness_multiplier: float = 3.0 + writer_lease_confirm_delay_seconds: float = 0.0 + writer_lease_acquire_timeout_seconds: float = 5.0 + writer_lease_force_acquire: bool = False + + +def _settings(**overrides: object) -> _LeaseSettings: + return _LeaseSettings(**overrides) # type: ignore[arg-type] + + +def _read_lease(directory: Path) -> dict: + return json.loads((directory / ".writer.lease").read_text(encoding="utf-8")) + + +async def _clear_app_task(name: str) -> None: + if hasattr(main_module.app.state, name): + delattr(main_module.app.state, name) + + +class _NoOpGraph: + async def flush(self) -> None: + return None + + async def close(self) -> None: + return None + + def discard_buffer(self) -> None: + return None + + +def _reset_lease(lease: WriterLease) -> None: + """Reset a WriterLease instance's fields to fresh-`__init__` values IN + PLACE (never replaces the object identity, since ``main.py`` imported + ``writer_lease`` by name at module-import time).""" + fresh = WriterLease() + for name in vars(fresh): + setattr(lease, name, getattr(fresh, name)) + + +@pytest.fixture(autouse=True) +def _isolate_module_singleton(): + """The `main_module.writer_lease` singleton is process-wide, so tests + that exercise it directly (rather than a locally-constructed + `WriterLease()`) must not leak state into each other.""" + _reset_lease(main_module.writer_lease) + yield + _reset_lease(main_module.writer_lease) + + +# `_restore_lease_io` (the process-wide `_LEASE_IO` executor guard) now lives +# in tests/conftest.py as an autouse fixture, so it protects every test +# module regardless of collection order -- not just this one. + + +async def _drive_lifespan( + monkeypatch: pytest.MonkeyPatch, +) -> tuple[MagicMock, AsyncMock]: + """Return the patches needed to drive the real ``lifespan()`` without a + real Neo4j -- mirrors ``tests/test_boot_safety.py``'s own + pattern exactly, so this file's lifespan integration tests use the SAME + boot-driving convention the boot-safety suite already established.""" + mock_driver = MagicMock() + mock_driver.close = AsyncMock() + monkeypatch.setattr(main_module._settings, "crash_recovery_respawn_limit", 0) + monkeypatch.setattr( + main_module._settings, "crash_recovery_sweep_interval_seconds", 0 + ) + return mock_driver, AsyncMock() + + +# --------------------------------------------------------------------------- +# (a) Clean boot acquires +# --------------------------------------------------------------------------- + + +async def test_clean_boot_acquires(tmp_path: Path) -> None: + lease = WriterLease() + await lease.acquire(_settings(), lambda: tmp_path) + + assert lease.acquired is True + assert lease.ever_acquired is True + assert lease.conflict is False + assert lease.error is None + + path = tmp_path / ".writer.lease" + assert path.exists() + text = path.read_text(encoding="utf-8") + assert text.endswith("\n") + assert text.count("\n") == 1 # ONE newline-terminated line + + record = json.loads(text) + assert record["lease_version"] == 1 + assert record["owner"] == lease.owner + assert abs(record["heartbeat"] - time.time()) < 2.0 + assert record["host"] + assert isinstance(record["pid"], int) and record["pid"] > 0 + assert record["server_version"] + + +# --------------------------------------------------------------------------- +# (b) DETECT never refuses -- even against a FRESH foreign lease +# --------------------------------------------------------------------------- + + +async def test_detect_never_refuses_fresh_foreign_lease(tmp_path: Path) -> None: + """An unconditional raise here would deadlock every rolling deploy -- + `detect` mode must acquire over a fresh foreign lease instead.""" + peer = WriterLease() + await peer.acquire(_settings(), lambda: tmp_path) + + lease = WriterLease() + # Must NOT raise -- this is the deadlock guard. + await lease.acquire(_settings(writer_lease_mode="detect"), lambda: tmp_path) + + assert lease.acquired is True + assert lease.conflict is True + assert lease.conflict_source == "boot" + assert lease.observed_owner == peer.owner + + # The deliberate take-over: on-disk owner is now us. + record = _read_lease(tmp_path) + assert record["owner"] == lease.owner + + +# --------------------------------------------------------------------------- +# (c) ENFORCE refuses against a fresh foreign lease; boot_task never created +# --------------------------------------------------------------------------- + + +async def test_enforce_refuses_fresh_foreign_lease(tmp_path: Path) -> None: + peer = WriterLease() + await peer.acquire(_settings(), lambda: tmp_path) + + lease = WriterLease() + with pytest.raises(WriterLeaseConflict) as ei: + await lease.acquire(_settings(writer_lease_mode="enforce"), lambda: tmp_path) + assert peer.owner in str(ei.value) + + # A refusal must not overwrite the on-disk owner. + record = _read_lease(tmp_path) + assert record["owner"] == peer.owner + + +async def test_enforce_refusal_aborts_real_startup_before_boot_task( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + queues_dir = tmp_path / "queues" + queues_dir.mkdir() + peer = WriterLease() + await peer.acquire(_settings(), lambda: queues_dir) + + monkeypatch.setattr( + "context_intelligence_server.registry.get_settings", + lambda: type( + "S", + (), + { + "queues_path": str(queues_dir), + "write_concurrency": 8, + "max_delivery_attempts": 5, + }, + )(), + ) + monkeypatch.setattr(main_module._settings, "writer_lease_mode", "enforce") + main_module.registry._queue_manager = None + await _clear_app_task("boot_task") + await _clear_app_task("lease_task") + + mock_driver = MagicMock() + mock_driver.close = AsyncMock() + with ( + patch("context_intelligence_server.main.setup_logging"), + patch( + "context_intelligence_server.neo4j_store.AsyncGraphDatabase.driver", + return_value=mock_driver, + ), + patch("context_intelligence_server.main.ensure_neo4j_schema", new=AsyncMock()), + pytest.raises(WriterLeaseConflict), + ): + async with lifespan(main_module.app): + pytest.fail("lifespan must not reach yield in enforce+fresh-foreign") + + assert not hasattr(main_module.app.state, "boot_task") + + +# --------------------------------------------------------------------------- +# (d) Acquire over a STALE lease; took_over_stale surfaced +# --------------------------------------------------------------------------- + + +async def test_acquire_over_stale_lease_takes_over_and_latches( + tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: + peer = WriterLease() + await peer.acquire(_settings(), lambda: tmp_path) + # Rewrite the peer's heartbeat far in the past (data-driven staleness -- + # no sleep-and-hope). + stale_record = _read_lease(tmp_path) + stale_record["heartbeat"] = time.time() - 3600 + (tmp_path / ".writer.lease").write_text( + json.dumps(stale_record) + "\n", encoding="utf-8" + ) + + lease = WriterLease() + with caplog.at_level("WARNING"): + await lease.acquire(_settings(), lambda: tmp_path) + + assert lease.acquired is True + assert lease.took_over_stale is True + assert lease.superseded_owner == peer.owner + assert lease.superseded_age_seconds is not None + assert lease.superseded_age_seconds > lease.staleness_seconds # type: ignore[operator] + assert any("STALE" in r.message for r in caplog.records) + + record = _read_lease(tmp_path) + assert record["owner"] == lease.owner + + +async def test_enforce_takes_over_stale_lease_without_refusing( + tmp_path: Path, +) -> None: + """A single-replica restart after an unclean exit must not crash-loop: + `enforce` takes over a STALE foreign lease instead of refusing.""" + peer = WriterLease() + await peer.acquire(_settings(), lambda: tmp_path) + stale_record = _read_lease(tmp_path) + stale_record["heartbeat"] = time.time() - 3600 + (tmp_path / ".writer.lease").write_text( + json.dumps(stale_record) + "\n", encoding="utf-8" + ) + + lease = WriterLease() + await lease.acquire(_settings(writer_lease_mode="enforce"), lambda: tmp_path) + + assert lease.acquired is True + assert lease.took_over_stale is True + assert lease.conflict is False + assert lease.superseded_owner == peer.owner + + record = _read_lease(tmp_path) + assert record["owner"] == lease.owner + + +async def test_took_over_stale_surfaces_on_status( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + lease = main_module.writer_lease + peer = WriterLease() + await peer.acquire(_settings(), lambda: tmp_path) + stale_record = _read_lease(tmp_path) + stale_record["heartbeat"] = time.time() - 3600 + (tmp_path / ".writer.lease").write_text( + json.dumps(stale_record) + "\n", encoding="utf-8" + ) + + await lease.acquire(_settings(), lambda: tmp_path) + assert lease.took_over_stale is True + + boot_state.phase = "ready" + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=main_module.app), base_url="http://test" + ) as client: + response = await client.get("/status") + body = response.json() + assert body["writer_lease"]["took_over_stale"] is True + assert body["writer_lease"]["superseded_owner"] == peer.owner + + +# --------------------------------------------------------------------------- +# (e) Heartbeat renews +# --------------------------------------------------------------------------- + + +async def test_heartbeat_renews(tmp_path: Path) -> None: + lease = WriterLease() + await lease.acquire(_settings(), lambda: tmp_path) + h0 = _read_lease(tmp_path)["heartbeat"] + + await asyncio.sleep(0.01) + await lease.tick() + + h1 = _read_lease(tmp_path)["heartbeat"] + assert h1 > h0 + assert lease.conflict is False + assert lease.acquired is True + + +async def test_heartbeat_loop_calls_tick(tmp_path: Path) -> None: + lease = WriterLease() + await lease.acquire( + _settings(writer_lease_heartbeat_seconds=0.05), lambda: tmp_path + ) + evt = asyncio.Event() + orig_tick = lease.tick + + async def _wrapped_tick() -> None: + await orig_tick() + evt.set() + + lease.tick = _wrapped_tick # type: ignore[method-assign] + task = asyncio.create_task(lease.heartbeat_loop()) + try: + await asyncio.wait_for(evt.wait(), 2.0) + finally: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + +# --------------------------------------------------------------------------- +# (f) Foreign overwrite mid-run -> writer_lease_conflict on real /status +# within one heartbeat -- THE ACCEPTANCE TEST +# --------------------------------------------------------------------------- + + +async def test_foreign_overwrite_surfaces_conflict_on_status(tmp_path: Path) -> None: + lease = main_module.writer_lease + await lease.acquire(_settings(), lambda: tmp_path) + assert lease.acquired is True + + peer = WriterLease() + await peer.acquire(_settings(), lambda: tmp_path) # real second identity steals it + + await lease.tick() + assert lease.conflict is True + assert lease.conflict_source == "runtime" + assert lease.observed_owner == peer.owner + assert lease.acquired is False # lost it + assert lease.ever_acquired is True + + boot_state.phase = "ready" + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=main_module.app), base_url="http://test" + ) as client: + response = await client.get("/status") + body = response.json() + assert body["writer_lease"]["conflict"] is True + # The spool projection is UNTOUCHED by this change -- boot-verified shape. + assert set(body["spool"].keys()) == { + "pending_sessions", + "spool_bytes_total", + "corrupt_offsets", + } + + +async def test_conflict_latches_and_stops_renewing(tmp_path: Path) -> None: + lease = WriterLease() + await lease.acquire(_settings(), lambda: tmp_path) + peer = WriterLease() + await peer.acquire(_settings(), lambda: tmp_path) + await lease.tick() + assert lease.conflict is True + + (tmp_path / ".writer.lease").unlink() + await lease.tick() + await lease.tick() + + assert lease.conflict is True + assert not (tmp_path / ".writer.lease").exists() # no steal-back + + +# --------------------------------------------------------------------------- +# (g) Clean shutdown releases; never destroys a foreign lease +# --------------------------------------------------------------------------- + + +async def test_clean_shutdown_releases(tmp_path: Path) -> None: + lease = WriterLease() + await lease.acquire(_settings(), lambda: tmp_path) + await lease.release() + assert not (tmp_path / ".writer.lease").exists() + + +async def test_release_never_destroys_a_foreign_lease(tmp_path: Path) -> None: + lease = WriterLease() + await lease.acquire(_settings(), lambda: tmp_path) + peer = WriterLease() + await peer.acquire(_settings(), lambda: tmp_path) # peer steals it + + await lease.release() + assert (tmp_path / ".writer.lease").exists() + assert _read_lease(tmp_path)["owner"] == peer.owner + + +async def test_release_never_raises_when_lease_file_missing(tmp_path: Path) -> None: + lease = WriterLease() + await lease.acquire(_settings(), lambda: tmp_path) + (tmp_path / ".writer.lease").unlink() + + await lease.release() # must not raise FileNotFoundError + + +async def test_release_never_raises_when_dir_missing(tmp_path: Path) -> None: + missing_dir = tmp_path / "gone" + lease = WriterLease() + await lease.acquire(_settings(), lambda: missing_dir) # dir absent -> unarmed + + await lease.release() # must not raise + + +# --------------------------------------------------------------------------- +# (h) Acquire race: exactly one wins +# --------------------------------------------------------------------------- + + +async def test_acquire_race_exactly_one_wins(tmp_path: Path) -> None: + """Only `enforce` mode raises on the LOST confirm-race branch (`detect`'s + deliberate take-over means BOTH sides would otherwise report success, + which is correct `detect` behaviour, not a race failure -- see + ``test_detect_never_refuses_fresh_foreign_lease``).""" + for _ in range(5): + d = tmp_path / f"race-{_}" + d.mkdir() + a = WriterLease() + b = WriterLease() + settings = _settings( + writer_lease_mode="enforce", writer_lease_confirm_delay_seconds=0.2 + ) + results = await asyncio.gather( + a.acquire(settings, lambda directory=d: directory), + b.acquire(settings, lambda directory=d: directory), + return_exceptions=True, + ) + winners = [r for r in results if r is None] + losers = [r for r in results if isinstance(r, Exception)] + assert len(winners) == 1, results + assert len(losers) == 1, results + assert isinstance(losers[0], WriterLeaseConflict) + winner = a if results[0] is None else b + assert _read_lease(d)["owner"] == winner.owner + + +# --------------------------------------------------------------------------- +# (i) R1 -- the detector constructs NO QueueManager (the merged-line-corruption guard) +# --------------------------------------------------------------------------- + + +async def test_r1_no_queue_manager_constructed_by_d6( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """The detector's acquire/tick/release path must never trigger + QueueManager construction, or a concurrent construction race can + reproduce torn/merged-line append corruption.""" + queues_dir = tmp_path / "queues" + queues_dir.mkdir() + main_module.registry._queue_manager = None + + def _fail_init(self: object, queues_dir: Path) -> None: + pytest.fail("QueueManager.__init__ must never be called by the detector") + + monkeypatch.setattr(QueueManager, "__init__", _fail_init) + + lease = WriterLease() + await lease.acquire(_settings(), lambda: main_module.registry.queues_dir_path) + await lease.tick() + await lease.tick() + + assert main_module.registry._queue_manager is None + assert lease.acquired is True + assert (queues_dir / ".writer.lease").exists() + + await lease.release() + assert main_module.registry._queue_manager is None + assert not (queues_dir / ".writer.lease").exists() + + +async def test_r1_exactly_one_queue_manager_under_concurrent_construction( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + queues_dir = tmp_path / "queues" + monkeypatch.setattr( + "context_intelligence_server.registry.get_settings", + lambda: type( + "S", + (), + { + "queues_path": str(queues_dir), + "write_concurrency": 8, + "max_delivery_attempts": 5, + }, + )(), + ) + main_module.registry._queue_manager = None + + construct_count = 0 + real_init = QueueManager.__init__ + + def _counting_init(self: QueueManager, queues_dir: Path) -> None: + nonlocal construct_count + construct_count += 1 + real_init(self, queues_dir) + + monkeypatch.setattr(QueueManager, "__init__", _counting_init) + + lease = WriterLease() + + async def _boot_touch() -> None: + await lease.acquire(_settings(), lambda: main_module.registry.queues_dir_path) + + async def _reconcile_touch() -> None: + for _ in range(20): + _ = main_module.registry.queue_manager + await asyncio.sleep(0) + + await asyncio.gather(_boot_touch(), _reconcile_touch()) + + assert construct_count == 1 + main_module.registry._queue_manager = None + + +# --------------------------------------------------------------------------- +# (j) Cold boot: dir absent -> acquired False -> re-arms once dir exists +# --------------------------------------------------------------------------- + + +async def test_cold_boot_dir_absent_rearms_once_dir_exists(tmp_path: Path) -> None: + missing_dir = tmp_path / "does-not-exist-yet" + lease = WriterLease() + + await lease.acquire(_settings(), lambda: missing_dir) + assert lease.acquired is False + assert lease.conflict is False + assert lease.error is not None + assert not missing_dir.exists() # the detector never creates the directory + + missing_dir.mkdir(parents=True) # stands in for boot recovery's mkdir + await lease.tick() + + assert lease.acquired is True + assert lease.conflict is False + assert _read_lease(missing_dir)["owner"] == lease.owner + + +# --------------------------------------------------------------------------- +# (k) Share-fault OSError at boot -> continue in every mode +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("mode", ["off", "detect", "enforce"]) +async def test_share_fault_at_boot_continues_in_every_mode( + mode: str, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + lease = WriterLease() + + def _boom() -> None: + raise OSError(errno.ESTALE, "stale file handle") + + monkeypatch.setattr(lease, "_read", _boom) + + await lease.acquire(_settings(writer_lease_mode=mode), lambda: tmp_path) + + assert lease.acquired is False + if mode == "off": + assert lease.error is None + else: + assert lease.conflict is False + assert lease.error is not None + + +async def test_writer_lease_boot_wrapper_survives_any_exception( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """The main.py-level guard: only WriterLeaseConflict may escape.""" + await _clear_app_task("lease_task") + monkeypatch.setattr(main_module._settings, "writer_lease_mode", "detect") + + async def _raise_runtime(*_a: object, **_kw: object) -> None: + raise RuntimeError("boom") + + monkeypatch.setattr(main_module.writer_lease, "acquire", _raise_runtime) + await main_module._writer_lease_boot() + assert main_module.writer_lease.error is not None + assert main_module.app.state.lease_task is not None + main_module.app.state.lease_task.cancel() + with pytest.raises(asyncio.CancelledError): + await main_module.app.state.lease_task + + async def _raise_conflict(*_a: object, **_kw: object) -> None: + raise WriterLeaseConflict("boom") + + monkeypatch.setattr(main_module.writer_lease, "acquire", _raise_conflict) + with pytest.raises(WriterLeaseConflict): + await main_module._writer_lease_boot() + + +# --------------------------------------------------------------------------- +# (l) Hung mount -> wait_for times out -> lifespan/heartbeat always returns +# --------------------------------------------------------------------------- + + +async def test_hung_mount_acquire_times_out(tmp_path: Path) -> None: + lease = WriterLease() + blocker = threading.Event() + + def _hang() -> None: + blocker.wait(timeout=5.0) + + monkeypatch_target = lease + monkeypatch_target._read = _hang # type: ignore[method-assign] + + start = time.monotonic() + await lease.acquire( + _settings(writer_lease_acquire_timeout_seconds=0.1), lambda: tmp_path + ) + elapsed = time.monotonic() - start + + assert lease.acquired is False + assert lease.error is not None + assert "time" in lease.error.lower() + assert elapsed < 2.0 # upper bound only -- never a lower-bound sleep assertion + blocker.set() + + +async def test_hung_mount_does_not_starve_shared_pool(tmp_path: Path) -> None: + """F2's bound: a stalled detector can never starve the append/commit path, + which run on the SHARED default executor.""" + lease = WriterLease() + blocker = threading.Event() + + def _hang() -> None: + blocker.wait(timeout=3.0) + + lease._read = _hang # type: ignore[method-assign] + + await lease.acquire( + _settings( + writer_lease_acquire_timeout_seconds=0.1, + writer_lease_heartbeat_seconds=0.05, + ), + lambda: tmp_path, + ) + task = asyncio.create_task(lease.heartbeat_loop()) + try: + await asyncio.sleep(0.5) + + qm = QueueManager(queues_dir=tmp_path / "shared-pool-check") + start = time.monotonic() + await asyncio.wait_for(qm.append("sid-1", b"hello"), timeout=1.0) + elapsed = time.monotonic() - start + assert elapsed < 1.0 + finally: + blocker.set() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + +# --------------------------------------------------------------------------- +# (m) A failed submit() must release the gate (R3) +# --------------------------------------------------------------------------- + + +async def test_submit_fail_releases_gate(tmp_path: Path) -> None: + from context_intelligence_server import writer_lease as wl_module + + lease = WriterLease() + await lease.acquire(_settings(), lambda: tmp_path) + + real_submit = wl_module._LEASE_IO.submit + calls = {"n": 0} + + def _boom_once(*args: object, **kwargs: object) -> object: + calls["n"] += 1 + if calls["n"] == 1: + raise RuntimeError("shutdown") + return real_submit(*args, **kwargs) # type: ignore[arg-type] + + with ( + patch.object(wl_module._LEASE_IO, "submit", side_effect=_boom_once), + pytest.raises(WriterLeaseBusy), + ): + await lease._io(lease._read) + assert lease._io_inflight is False + + # Re-arm: unpatched, a subsequent op succeeds. + await lease.tick() + assert lease.conflict is False + + +# --------------------------------------------------------------------------- +# (n) Shutdown never raises, even with the gate busy +# --------------------------------------------------------------------------- + + +async def test_shutdown_never_raises_with_busy_gate( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Forces a genuinely busy lease-I/O gate (real outstanding future on the + single-worker executor, `_io_inflight` held True) at the exact moment + shutdown calls `release()`, so `release()` must swallow the resulting + `WriterLeaseBusy` rather than let it escape.""" + from context_intelligence_server import writer_lease as wl_module + + queues_dir = tmp_path / "queues" + queues_dir.mkdir() + + monkeypatch.setattr( + "context_intelligence_server.registry.get_settings", + lambda: type( + "S", + (), + { + "queues_path": str(queues_dir), + "write_concurrency": 8, + "max_delivery_attempts": 5, + }, + )(), + ) + main_module.registry._queue_manager = None + monkeypatch.setattr(main_module._settings, "writer_lease_mode", "detect") + # Keep the heartbeat from ticking during the forced-busy window so it + # cannot itself contend for the one-slot gate and mask the scenario. + monkeypatch.setattr(main_module._settings, "writer_lease_heartbeat_seconds", 999.0) + monkeypatch.setattr(main_module._settings, "crash_recovery_respawn_limit", 0) + monkeypatch.setattr( + main_module._settings, "crash_recovery_sweep_interval_seconds", 0 + ) + await _clear_app_task("boot_task") + await _clear_app_task("lease_task") + await _clear_app_task("sweep_task") + + mock_driver = MagicMock() + mock_driver.close = AsyncMock() + + blocker = threading.Event() + + def _hang() -> None: + blocker.wait(timeout=5.0) + + hung_future = None + try: + with ( + patch("context_intelligence_server.main.setup_logging"), + patch( + "context_intelligence_server.neo4j_store.AsyncGraphDatabase.driver", + return_value=mock_driver, + ), + patch( + "context_intelligence_server.main.ensure_neo4j_schema", + new=AsyncMock(), + ), + ): + async with lifespan(main_module.app): + await main_module.app.state.boot_task + assert main_module.writer_lease.acquired is True + + # Force the adverse state right before we fall out of this + # block (which runs lifespan's `finally`): a real + # outstanding future holding the executor's one worker + # thread, and the gate flag itself closed. + hung_future = wl_module._LEASE_IO.submit(_hang) + main_module.writer_lease._io_inflight = True + # Falling out of the `async with` here triggers shutdown: + # sweep/boot/lease tasks cancelled, then + # `writer_lease.release()` -- which must swallow the + # WriterLeaseBusy this forced state causes `_io()` to raise. + # Reaching here without an exception IS the assertion: shutdown + # never raises WriterLeaseBusy/OSError past `lifespan`'s finally, + # even with a genuinely busy gate at release() time. + finally: + blocker.set() + if hung_future is not None: + hung_future.result(timeout=5.0) + main_module.writer_lease._io_inflight = False + + assert ( + main_module.app.state.lease_task.cancelled() + or main_module.app.state.lease_task.done() + ) + + +# --------------------------------------------------------------------------- +# (o) /status writer_lease block present during boot, ZERO disk reads +# --------------------------------------------------------------------------- + + +async def test_status_writer_lease_present_during_boot_zero_disk_reads( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Instruments disk I/O directly (patches `Path.read_text`, filtered to + this lease's own path) rather than only asserting on the JSON body, so a + regression that reads the lease file from disk on every request is caught + even if `/status`'s JSON output looks unchanged.""" + lease = main_module.writer_lease + await lease.acquire(_settings(), lambda: tmp_path) + + lease_path = lease.path + read_calls = {"n": 0} + real_read_text = Path.read_text + + def _counting_read_text( + self: Path, encoding: str | None = None, errors: str | None = None + ) -> str: + if self == lease_path: + read_calls["n"] += 1 + return real_read_text(self, encoding=encoding, errors=errors) + + monkeypatch.setattr(Path, "read_text", _counting_read_text) + + boot_state.phase = "reclaim" # actively booting + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=main_module.app), base_url="http://test" + ) as client: + response = await client.get("/status") + + assert response.status_code == 200 + body = response.json() + assert body["spool"] is None # the boot-safety lean contract intact + assert "writer_lease" in body + assert body["writer_lease"]["acquired"] is True + assert body["writer_lease"]["owner"] == lease.owner + assert read_calls["n"] == 0, ( + "writer_lease.snapshot() (and /status while actively booting) must " + "be pure in-memory -- it must never read the on-disk lease file. " + f"Observed {read_calls['n']} disk read(s) of {lease_path}." + ) + + +async def test_status_writer_lease_never_500s_when_unacquired() -> None: + lease = main_module.writer_lease + saved = ( + lease.mode, + lease.acquired, + lease.heartbeat_seconds, + lease.staleness_seconds, + lease.last_renewed, + ) + lease.mode = None + lease.acquired = False + lease.heartbeat_seconds = None + lease.staleness_seconds = None + lease.last_renewed = None + try: + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=main_module.app), base_url="http://test" + ) as client: + response = await client.get("/status") + assert response.status_code == 200 + assert response.json()["writer_lease"]["acquired"] is False + finally: + ( + lease.mode, + lease.acquired, + lease.heartbeat_seconds, + lease.staleness_seconds, + lease.last_renewed, + ) = saved + + +# --------------------------------------------------------------------------- +# Config validators (writer-lease settings) +# --------------------------------------------------------------------------- + + +def test_writer_lease_config_defaults() -> None: + s = Settings() + assert s.writer_lease_mode == "enforce" + assert s.writer_lease_heartbeat_seconds == 5.0 + assert s.writer_lease_staleness_multiplier == 3.0 + assert s.writer_lease_confirm_delay_seconds == 1.0 + assert s.writer_lease_acquire_timeout_seconds == 5.0 + assert s.writer_lease_force_acquire is False + + +def test_writer_lease_config_validators_fail_loud() -> None: + from pydantic import ValidationError + + with pytest.raises(ValidationError): + Settings(writer_lease_mode="enforced") # pyright: ignore[reportArgumentType] -- typo, not a legal value + with pytest.raises(ValidationError): + Settings(writer_lease_staleness_multiplier=1.5) + with pytest.raises(ValidationError): + Settings(writer_lease_heartbeat_seconds=0) + with pytest.raises(ValidationError): + Settings(writer_lease_confirm_delay_seconds=-1.0) + with pytest.raises(ValidationError): + Settings( + writer_lease_acquire_timeout_seconds=0.5, + writer_lease_confirm_delay_seconds=1.0, + ) + + +# --------------------------------------------------------------------------- +# Collision guard: the writer-lease files never collide with other boot-path scans +# --------------------------------------------------------------------------- + + +async def test_no_collision_with_existing_session_scans(tmp_path: Path) -> None: + lease = WriterLease() + await lease.acquire(_settings(), lambda: tmp_path) + + qm = QueueManager(queues_dir=tmp_path) + (tmp_path / "sess-1.log").write_bytes(b'{"event":"x","workspace":"/w"}\n') + (tmp_path / "sess-1.offset").write_text("0", encoding="utf-8") + + keys = sorted(p.stem for p in qm.queues_dir.glob("*.log")) + assert "sess-1" in keys + assert ".writer" not in keys + + active = await qm.active_sessions() + assert ".writer.lease" not in active + assert (tmp_path / ".writer.lease").exists() diff --git a/uv.lock b/uv.lock index 06f8fbaf..a21b650a 100644 --- a/uv.lock +++ b/uv.lock @@ -233,7 +233,7 @@ wheels = [ [[package]] name = "context-intelligence-server" -version = "6.7.0" +version = "6.7.2" source = { editable = "." } dependencies = [ { name = "aiofiles" },