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/blob_processor.py b/context_intelligence_server/blob_processor.py index bd7d8c25..9a39ae41 100644 --- a/context_intelligence_server/blob_processor.py +++ b/context_intelligence_server/blob_processor.py @@ -10,6 +10,7 @@ from __future__ import annotations import logging +import re from typing import TYPE_CHECKING, Any if TYPE_CHECKING: @@ -27,6 +28,92 @@ ) +# --------------------------------------------------------------------------- +# Blob-ref carrier allowlist -- single source of truth for which property +# names may carry a ci-blob:// URI. +# --------------------------------------------------------------------------- +# +# Every ``ci-blob://`` URI minted below (``process_event_data``) is written +# into ``data``, which ``DefaultHandler`` always persists wholesale as the +# JSON-serialized ``data`` property on the Event node +# (handlers/data_layer_1/default.py). The blob-reclaim reference scan +# (``routers.admin._scan_referenced_uris``) enumerates every ``ci-blob://`` +# reference anywhere in the graph by walking a FIXED allowlist of node +# properties -- never an all-property/all-node scan (this codebase has scar +# tissue from a 1.3M-node AllNodesScan stall). A blob whose reference lives +# on a node property the scan doesn't know about is invisible to it and can +# be deleted as a false orphan. +# +# BLOB_REF_CARRIER_PROPERTIES is THE single source of truth for that +# allowlist, imported by ``routers.admin`` to build the scan's Cypher +# directly from this tuple (so the query text can never drift from it) and +# checked here, at the mint site, via :func:`assert_carrier_registered`. +# +# Adding a new carrier (a future field-lifter/enricher that promotes a +# blob-ref-shaped value onto a new node property) means adding its name +# here. Forgetting to is now a fail-closed error, not a silent GC hole. +BLOB_REF_CARRIER_PROPERTIES: tuple[str, ...] = ( + "data", + "tool_input", + "prompt", + "response", +) + +# Defensive validation, run once at import time: every carrier name must be +# a legal Cypher property identifier, because routers.admin interpolates +# these names directly into a Cypher query string. Guards against a future +# careless addition (e.g. containing a space or backtick) turning into a +# broken or injectable query rather than a loud, immediate import error. +_VALID_CARRIER_NAME_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") + + +def _validate_carrier_names(names: tuple[str, ...]) -> None: + for name in names: + if not _VALID_CARRIER_NAME_RE.match(name): + raise ValueError( + f"BLOB_REF_CARRIER_PROPERTIES entry {name!r} is not a valid " + "Cypher property identifier -- refusing to load (this tuple " + "is interpolated directly into a Cypher query by " + "routers.admin._scan_referenced_uris)" + ) + + +_validate_carrier_names(BLOB_REF_CARRIER_PROPERTIES) + + +class UnregisteredBlobCarrierError(RuntimeError): + """A ``ci-blob://`` reference is destined for a node property that is + not in :data:`BLOB_REF_CARRIER_PROPERTIES`. + + This converts a silent reclaim-GC hole + (a live blob deleted as an orphan because its carrier property was never + added to the allowlist) into a loud, immediate failure at the point the + omission is introduced -- not after a live blob is gone. + """ + + +def assert_carrier_registered(property_name: str) -> None: + """Fail loud if *property_name* is not a registered blob-ref carrier. + + Cheap (single tuple-membership check) and safe to call on every + ``process_event_data`` invocation. Raises + :class:`UnregisteredBlobCarrierError` -- deliberately NOT caught by the + per-field ``except Exception`` below, so it propagates out of + ``process_event_data``, through ``pipeline.process_event``'s outer + handler (which logs and re-raises), and the event is dead-lettered + instead of silently minting an unprotected blob reference. + """ + if property_name not in BLOB_REF_CARRIER_PROPERTIES: + raise UnregisteredBlobCarrierError( + f"ci-blob:// reference destined for node property {property_name!r} " + f"is not in BLOB_REF_CARRIER_PROPERTIES {BLOB_REF_CARRIER_PROPERTIES!r} " + "-- the blob-reclaim scan (context_intelligence_server.routers.admin) " + "will not see refs stored there and could delete this blob as an " + f"orphan. Add {property_name!r} to BLOB_REF_CARRIER_PROPERTIES " + "(context_intelligence_server/blob_processor.py) before shipping." + ) + + # --------------------------------------------------------------------------- # Internal helpers # --------------------------------------------------------------------------- @@ -88,6 +175,16 @@ async def process_event_data( """ _lift_raw_fields(data) + # Every ci-blob:// URI minted below lands in + # `data`, which DefaultHandler always persists wholesale onto the Event + # node's "data" property. Fail loud, BEFORE any blob is written, if that + # destination is ever missing from the allowlist the reclaim scan reads + # (see BLOB_REF_CARRIER_PROPERTIES above). Deliberately outside the + # per-field try/except below so it is never downgraded to a swallowed + # $blob_error -- it propagates out of process_event_data and dead-letters + # the event instead of silently minting an unprotected blob reference. + assert_carrier_registered("data") + for field_name in BLOB_FIELDS: value = data.get(field_name) if value is None: @@ -96,8 +193,8 @@ async def process_event_data( key = f"{node_id}__{field_name}" try: - uri = await blob_store.write(session_id, key, value) - data[field_name] = {"$blob_ref": uri} + ref = await blob_store.write(session_id, key, value) + data[field_name] = {"$blob_ref": ref.uri} except Exception as exc: # noqa: BLE001 logger.warning( "blob_offload_failed session=%s field=%s node=%s: %s", diff --git a/context_intelligence_server/blob_store.py b/context_intelligence_server/blob_store.py deleted file mode 100644 index 94511781..00000000 --- a/context_intelligence_server/blob_store.py +++ /dev/null @@ -1,237 +0,0 @@ -"""AsyncDiskBlobStore — async, disk-backed blob storage with ci-blob:// URIs. - -Disk layout: - //blobs/.json - -URI scheme: - ci-blob:/// - -All filesystem I/O is wrapped with ``asyncio.to_thread`` to keep the event -loop non-blocking. -""" - -from __future__ import annotations - -import asyncio -import json -import os -import shutil -import tempfile -from pathlib import Path -from typing import Any, Protocol, cast, runtime_checkable - -_SCHEME = "ci-blob://" - - -# --------------------------------------------------------------------------- -# BlobStore protocol -# --------------------------------------------------------------------------- - - -@runtime_checkable -class BlobStore(Protocol): - """Protocol for a session-scoped, URI-addressable blob store.""" - - async def write( - self, session_id: str, key: str, value: dict[str, Any] | list[Any] - ) -> str: - """Persist *value* as JSON and return a ``ci-blob://`` URI.""" - ... - - async def read(self, uri: str) -> dict[str, Any] | list[Any]: - """Resolve *uri* and return the stored value. - - Raises: - ValueError: If *uri* does not match the ``ci-blob://`` scheme. - FileNotFoundError: If no blob exists at the resolved path. - """ - ... - - async def list(self, session_id: str) -> list[str]: - """Return all blob URIs for *session_id*, sorted lexicographically.""" - ... - - async def dump(self, uri: str, dest_dir: Path | str | None = None) -> str: - """Copy the blob file addressed by *uri* to *dest_dir*. - - Args: - uri: ``ci-blob://`` URI identifying the blob to copy. - dest_dir: Destination directory. Defaults to - ``Path(tempfile.gettempdir()) / 'ci-blobs'``. - - Returns: - The destination file path as a string. - - Raises: - ValueError: If *uri* is not a valid ``ci-blob://`` URI. - FileNotFoundError: If no blob exists at the resolved path. - """ - ... - - -# --------------------------------------------------------------------------- -# AsyncDiskBlobStore -# --------------------------------------------------------------------------- - - -class AsyncDiskBlobStore: - """Async, disk-backed implementation of :class:`BlobStore`. - - Args: - root: Root directory under which all session blobs are stored. - """ - - def __init__(self, root: Path | str) -> None: - self._root = Path(root) - - # ------------------------------------------------------------------ - # Internal helpers - # ------------------------------------------------------------------ - - def _make_uri(self, session_id: str, key: str) -> str: - """Return the canonical ``ci-blob://`` URI for a session/key pair.""" - return f"{_SCHEME}{session_id}/{key}" - - def _parse_uri(self, uri: str) -> tuple[str, str]: - """Parse a ``ci-blob://`` URI into ``(session_id, key)``. - - Raises: - ValueError: If *uri* is not a valid ``ci-blob://`` URI. - """ - if not uri.startswith(_SCHEME): - raise ValueError( - f"Invalid URI scheme — expected '{_SCHEME}...', got: {uri!r}" - ) - remainder = uri[len(_SCHEME) :] - # remainder must be "/" — both parts non-empty - if "/" not in remainder: - raise ValueError(f"URI missing key component: {uri!r}") - session_id, _, key = remainder.partition("/") - if not session_id or not key: - raise ValueError(f"URI has empty session_id or key: {uri!r}") - return session_id, key - - def _blob_path(self, session_id: str, key: str) -> Path: - """Return the filesystem path for a given session/key blob.""" - return self._root / session_id / "blobs" / f"{key}.json" - - # ------------------------------------------------------------------ - # Public accessors (mirror of internal helpers for external callers) - # ------------------------------------------------------------------ - - def parse_uri(self, uri: str) -> tuple[str, str]: - """Public alias for :meth:`_parse_uri`.""" - return self._parse_uri(uri) - - def blob_path(self, session_id: str, key: str) -> Path: - """Public alias for :meth:`_blob_path`.""" - return self._blob_path(session_id, key) - - # ------------------------------------------------------------------ - # Async API - # ------------------------------------------------------------------ - - async def write( - self, session_id: str, key: str, value: dict[str, Any] | list[Any] - ) -> str: - """Persist *value* as JSON and return a ``ci-blob://`` URI. - - Creates the directory ``//blobs/`` if needed. - - Returns: - A ``ci-blob:///`` URI. - """ - path = self._blob_path(session_id, key) - - def _write() -> None: - path.parent.mkdir(parents=True, exist_ok=True) - data = json.dumps(value) - tmp_fd, tmp_name = tempfile.mkstemp( - dir=str(path.parent), prefix=f"{key}.", suffix=".tmp" - ) - try: - with os.fdopen(tmp_fd, "w", encoding="utf-8") as f: - f.write(data) - f.flush() - os.fsync(f.fileno()) - os.replace(tmp_name, path) - except BaseException: - try: - os.unlink(tmp_name) - except FileNotFoundError: - pass - raise - - await asyncio.to_thread(_write) - return self._make_uri(session_id, key) - - async def read(self, uri: str) -> dict[str, Any] | list[Any]: - """Return the blob addressed by *uri*. - - The session_id is resolved from the URI itself — callers do not - supply it separately (avoids the bundle footgun where the wrong - session_id is passed). - - Raises: - ValueError: If *uri* is not a valid ``ci-blob://`` URI. - FileNotFoundError: If no blob exists at the resolved path. - """ - session_id, key = self._parse_uri(uri) - path = self._blob_path(session_id, key) - - def _read() -> dict[str, Any] | list[Any]: - try: - return cast( - dict[str, Any] | list[Any], - json.loads(path.read_text(encoding="utf-8")), - ) - except FileNotFoundError: - raise FileNotFoundError(f"Blob not found: {uri!r} (path: {path})") - - return await asyncio.to_thread(_read) - - async def list(self, session_id: str) -> list[str]: - """Return all blob URIs for *session_id*, sorted lexicographically. - - Returns an empty list if the session directory does not exist. - """ - blobs_dir = self._root / session_id / "blobs" - - def _list() -> list[str]: - if not blobs_dir.exists(): - return [] - keys = sorted(p.stem for p in blobs_dir.glob("*.json")) - return [self._make_uri(session_id, key) for key in keys] - - return await asyncio.to_thread(_list) - - async def dump(self, uri: str, dest_dir: Path | str | None = None) -> str: - """Copy the blob file addressed by *uri* to *dest_dir*. - - Args: - uri: ``ci-blob://`` URI identifying the blob to copy. - dest_dir: Destination directory. Defaults to - ``Path(tempfile.gettempdir()) / 'ci-blobs'``. - - Returns: - The destination file path as a string. - - Raises: - ValueError: If *uri* is not a valid ``ci-blob://`` URI. - FileNotFoundError: If no blob exists at the resolved path. - """ - session_id, key = self._parse_uri(uri) - src = self._blob_path(session_id, key) - - if dest_dir is None: - dest_dir_path = Path(tempfile.gettempdir()) / "ci-blobs" - else: - dest_dir_path = Path(dest_dir) - - def _copy() -> str: - if not src.exists(): - raise FileNotFoundError(f"Blob not found: {uri!r}") - dest_dir_path.mkdir(parents=True, exist_ok=True) - return str(shutil.copy2(src, dest_dir_path)) - - return await asyncio.to_thread(_copy) diff --git a/context_intelligence_server/blob_store/__init__.py b/context_intelligence_server/blob_store/__init__.py new file mode 100644 index 00000000..0a9ecdac --- /dev/null +++ b/context_intelligence_server/blob_store/__init__.py @@ -0,0 +1,29 @@ +"""blob_store \u2014 session-scoped, URI-addressable blob storage. + +The public surface is the backend-neutral :class:`BlobStore` Protocol plus +:class:`BlobReference` / :class:`BlobNotFoundError` and the +:func:`create_blob_store` factory. Consumers should depend on these, never on +a concrete backend class. + +Package layout: + protocol.py BlobStore Protocol, BlobReference, BlobNotFoundError \u2014 the + backend-neutral seam (no filesystem imports). + filesystem.py FileSystemBlobStore \u2014 the disk-backed implementation. + factory.py create_blob_store(settings) \u2014 the ONLY place a backend is + selected and the ONLY place (besides config.py) that reads + settings.blob_path. +""" + +from __future__ import annotations + +from .factory import create_blob_store +from .filesystem import FileSystemBlobStore +from .protocol import BlobNotFoundError, BlobReference, BlobStore + +__all__ = [ + "BlobNotFoundError", + "BlobReference", + "BlobStore", + "FileSystemBlobStore", + "create_blob_store", +] diff --git a/context_intelligence_server/blob_store/factory.py b/context_intelligence_server/blob_store/factory.py new file mode 100644 index 00000000..adebde1b --- /dev/null +++ b/context_intelligence_server/blob_store/factory.py @@ -0,0 +1,46 @@ +"""Config-driven BlobStore factory \u2014 the ONLY place a blob-store backend is selected. + +This is the single seam through which the concrete backend is chosen. Adding +a new backend (e.g. Azure) means: one new module implementing +:class:`~.protocol.BlobStore`, one new branch here, and a config value \u2014 +zero changes to :mod:`context_intelligence_server.registry` or any consumer. + +This module (and :mod:`~context_intelligence_server.config`) are the only +places ``settings.blob_path`` is read \u2014 the on-disk root is a filesystem- +backend concern, resolved here and handed to the concrete backend at +construction time. Callers only ever see the :class:`~.protocol.BlobStore` +Protocol. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from .filesystem import FileSystemBlobStore +from .protocol import BlobStore + +if TYPE_CHECKING: + from context_intelligence_server.config import Settings + + +def create_blob_store(settings: Settings) -> BlobStore: + """Build the configured :class:`~.protocol.BlobStore` backend. + + Reads ``settings.blob_backend`` (default ``"filesystem"``) to select the + implementation: + + - ``"filesystem"``: :class:`~.filesystem.FileSystemBlobStore` rooted at + ``settings.blob_path``. + - ``"azure"``: not yet implemented. + - anything else: rejected as an unknown backend. + + Raises: + NotImplementedError: If ``blob_backend == "azure"`` (not yet built). + ValueError: If ``blob_backend`` names an unknown backend. + """ + backend = settings.blob_backend + if backend == "filesystem": + return FileSystemBlobStore(root=settings.blob_path) + if backend == "azure": + raise NotImplementedError("azure blob backend not yet implemented") + raise ValueError(f"Unknown blob_backend: {backend!r}") diff --git a/context_intelligence_server/blob_store/filesystem.py b/context_intelligence_server/blob_store/filesystem.py new file mode 100644 index 00000000..505508e2 --- /dev/null +++ b/context_intelligence_server/blob_store/filesystem.py @@ -0,0 +1,298 @@ +"""FileSystemBlobStore \u2014 async, disk-backed blob storage with ci-blob:// URIs. + +Disk layout: + //blobs/.json + +URI scheme: + ci-blob:/// + +All filesystem I/O is wrapped with ``asyncio.to_thread`` to keep the event +loop non-blocking. + +This is a concrete implementation of the :class:`~context_intelligence_server.blob_store.protocol.BlobStore` +Protocol. No ``Path``, on-disk layout, ``dest_dir``, or ``os.*`` detail +appears in the Protocol or in any value it returns \u2014 those details are +private to this class (and, later, an Azure equivalent). +""" + +from __future__ import annotations + +import asyncio +import json +import os +import shutil +import tempfile +from collections.abc import AsyncIterator +from pathlib import Path +from typing import Any, cast + +from .protocol import BlobNotFoundError, BlobReference + +_SCHEME = "ci-blob://" + + +class FileSystemBlobStore: + """Async, disk-backed implementation of :class:`~.protocol.BlobStore`. + + Args: + root: Root directory under which all session blobs are stored. + """ + + def __init__(self, root: Path | str) -> None: + self._root = Path(root) + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _make_uri(self, session_id: str, key: str) -> str: + """Return the canonical ``ci-blob://`` URI for a session/key pair.""" + return f"{_SCHEME}{session_id}/{key}" + + def _parse_uri(self, uri: str) -> tuple[str, str]: + """Parse a ``ci-blob://`` URI into ``(session_id, key)``. + + Raises: + ValueError: If *uri* is not a valid ``ci-blob://`` URI. + """ + if not uri.startswith(_SCHEME): + raise ValueError( + f"Invalid URI scheme \u2014 expected '{_SCHEME}...', got: {uri!r}" + ) + remainder = uri[len(_SCHEME) :] + # remainder must be "/" \u2014 both parts non-empty + if "/" not in remainder: + raise ValueError(f"URI missing key component: {uri!r}") + session_id, _, key = remainder.partition("/") + if not session_id or not key: + raise ValueError(f"URI has empty session_id or key: {uri!r}") + return session_id, key + + def _blob_path(self, session_id: str, key: str) -> Path: + """Return the filesystem path for a given session/key blob.""" + return self._root / session_id / "blobs" / f"{key}.json" + + # ------------------------------------------------------------------ + # Public accessors (mirror of internal helpers for external callers) + # ------------------------------------------------------------------ + + def parse_uri(self, uri: str) -> tuple[str, str]: + """Public alias for :meth:`_parse_uri`.""" + return self._parse_uri(uri) + + def blob_path(self, session_id: str, key: str) -> Path: + """Public alias for :meth:`_blob_path`.""" + return self._blob_path(session_id, key) + + # ------------------------------------------------------------------ + # Async API + # ------------------------------------------------------------------ + + async def write( + self, session_id: str, key: str, value: dict[str, Any] | list[Any] + ) -> BlobReference: + """Persist *value* as JSON and return a :class:`BlobReference`. + + Creates the directory ``//blobs/`` if needed. + ``last_modified`` is the storage mtime (from the same ``stat`` call + that produces ``size``) \u2014 never a writer-clock timestamp. + """ + path = self._blob_path(session_id, key) + + def _write() -> os.stat_result: + path.parent.mkdir(parents=True, exist_ok=True) + data = json.dumps(value) + tmp_fd, tmp_name = tempfile.mkstemp( + dir=str(path.parent), prefix=f"{key}.", suffix=".tmp" + ) + try: + with os.fdopen(tmp_fd, "w", encoding="utf-8") as f: + f.write(data) + f.flush() + os.fsync(f.fileno()) + os.replace(tmp_name, path) + except BaseException: + try: + os.unlink(tmp_name) + except FileNotFoundError: + pass + raise + return path.stat() + + st = await asyncio.to_thread(_write) + return BlobReference( + uri=self._make_uri(session_id, key), + session_id=session_id, + key=key, + size=st.st_size, + last_modified=st.st_mtime, + ) + + async def read(self, uri: str) -> dict[str, Any] | list[Any]: + """Return the blob addressed by *uri*. + + The session_id is resolved from the URI itself \u2014 callers do not + supply it separately (avoids the bundle footgun where the wrong + session_id is passed). + + Raises: + ValueError: If *uri* is not a valid ``ci-blob://`` URI. + BlobNotFoundError: If no blob exists for *uri*. Subclasses + ``FileNotFoundError`` for back-compat; the message carries + the URI only \u2014 never the on-disk path. + """ + session_id, key = self._parse_uri(uri) + path = self._blob_path(session_id, key) + + def _read() -> dict[str, Any] | list[Any]: + try: + return cast( + dict[str, Any] | list[Any], + json.loads(path.read_text(encoding="utf-8")), + ) + except FileNotFoundError: + raise BlobNotFoundError(f"Blob not found: {uri!r}") from None + + return await asyncio.to_thread(_read) + + async def list(self, session_id: str) -> AsyncIterator[BlobReference]: + """Stream all blob references for *session_id*. + + Yields nothing if the session's blobs directory does not exist. + The scandir + per-entry stat work is offloaded to a thread in small + units (per-entry), so the event loop stays responsive and references + stream out incrementally rather than blocking on the whole walk. + """ + blobs_dir = self._root / session_id / "blobs" + + def _list_entries() -> list[tuple[str, int, float]]: + if not blobs_dir.exists(): + return [] + entries: list[tuple[str, int, float]] = [] + with os.scandir(blobs_dir) as it: + for entry in it: + if not entry.name.endswith(".json"): + continue + st = entry.stat() + key = entry.name[: -len(".json")] + entries.append((key, st.st_size, st.st_mtime)) + entries.sort(key=lambda e: e[0]) + return entries + + entries = await asyncio.to_thread(_list_entries) + for key, size, last_modified in entries: + yield BlobReference( + uri=self._make_uri(session_id, key), + session_id=session_id, + key=key, + size=size, + last_modified=last_modified, + ) + + async def scan(self) -> AsyncIterator[BlobReference]: + """Stream all blob references across ALL sessions. + + Walks ``/*/blobs/*.json`` \u2014 session-dir enumeration and each + session's blob-dir scan are offloaded to a thread in small units + (never one giant ``to_thread`` for the whole tree), so references + stream out as they are discovered instead of materializing the + entire store in memory before yielding anything. + """ + + def _list_session_dirs() -> list[str]: + if not self._root.exists(): + return [] + with os.scandir(self._root) as it: + return sorted(entry.name for entry in it if entry.is_dir()) + + session_ids = await asyncio.to_thread(_list_session_dirs) + for session_id in session_ids: + async for ref in self.list(session_id): + yield ref + + async def delete( + self, uri: str, if_unmodified: BlobReference | None = None + ) -> bool: + """Delete the blob addressed by *uri*. + + Idempotent: returns ``False`` (never raises) if the blob is already + absent, ``True`` if it existed and was removed. + + Args: + uri: The ``ci-blob://`` URI to delete. + if_unmodified: When ``None`` (default), unconditional delete \u2014 + unlinks and returns ``True``, or ``False`` if already absent. + When provided, this is a **fenced compare-and-delete**: the + blob is re-``stat``'d (inside the same thread hop, right + before the unlink, to minimise the TOCTOU window) and the + delete only proceeds if ``st_mtime``/``st_size`` still match + *if_unmodified* \u2014 i.e. nothing rewrote the blob since it was + observed (e.g. by ``scan()``). If the blob is missing, or it + changed, the delete is refused and ``False`` is returned \u2014 + the blob is left untouched on disk. + """ + session_id, key = self._parse_uri(uri) + path = self._blob_path(session_id, key) + + def _delete() -> bool: + if if_unmodified is None: + try: + os.unlink(path) + return True + except FileNotFoundError: + return False + + # Fenced compare-and-delete: stat first, unlink only if unchanged. + try: + st = path.stat() + except FileNotFoundError: + return False + if ( + st.st_mtime != if_unmodified.last_modified + or st.st_size != if_unmodified.size + ): + # Blob was rewritten since it was observed \u2014 refuse to delete. + return False + try: + os.unlink(path) + return True + except FileNotFoundError: + # Deleted concurrently between our stat and unlink. + return False + + return await asyncio.to_thread(_delete) + + async def dump(self, uri: str, dest_dir: Path | str | None = None) -> str: + """Copy the blob file addressed by *uri* to *dest_dir*. + + Disk-only helper \u2014 NOT part of the :class:`~.protocol.BlobStore` + Protocol (no production caller; kept as a concrete convenience for + external tooling that needs a local export). + + Args: + uri: ``ci-blob://`` URI identifying the blob to copy. + dest_dir: Destination directory. Defaults to + ``Path(tempfile.gettempdir()) / 'ci-blobs'``. + + Returns: + The destination file path as a string. + + Raises: + ValueError: If *uri* is not a valid ``ci-blob://`` URI. + FileNotFoundError: If no blob exists at the resolved path. + """ + session_id, key = self._parse_uri(uri) + src = self._blob_path(session_id, key) + + if dest_dir is None: + dest_dir_path = Path(tempfile.gettempdir()) / "ci-blobs" + else: + dest_dir_path = Path(dest_dir) + + def _copy() -> str: + if not src.exists(): + raise FileNotFoundError(f"Blob not found: {uri!r}") + dest_dir_path.mkdir(parents=True, exist_ok=True) + return str(shutil.copy2(src, dest_dir_path)) + + return await asyncio.to_thread(_copy) diff --git a/context_intelligence_server/blob_store/protocol.py b/context_intelligence_server/blob_store/protocol.py new file mode 100644 index 00000000..a48c1eaa --- /dev/null +++ b/context_intelligence_server/blob_store/protocol.py @@ -0,0 +1,109 @@ +"""BlobStore Protocol \u2014 the backend-neutral seam. + +The only identity that crosses the boundary is the ``ci-blob:///`` +URI, carried by :class:`BlobReference`. No ``Path``, on-disk layout, ``dest_dir``, +or ``os.*`` detail appears here or in any value the Protocol returns \u2014 that is +private to a concrete backend (:class:`~context_intelligence_server.blob_store.filesystem.FileSystemBlobStore` +and, later, an Azure equivalent). + +This module is backend-neutral by construction: it imports nothing from +``os``, ``pathlib``, or any filesystem library. +""" + +from __future__ import annotations + +from collections.abc import AsyncIterator +from dataclasses import dataclass +from typing import Any, Protocol, runtime_checkable + +# --------------------------------------------------------------------------- +# BlobNotFoundError \u2014 backend-neutral missing-blob exception (guard #6) +# --------------------------------------------------------------------------- + + +class BlobNotFoundError(FileNotFoundError): + """Raised when a blob addressed by a ``ci-blob://`` URI does not exist. + + Subclasses :class:`FileNotFoundError` so existing ``except + FileNotFoundError`` callers keep working unchanged (zero caller churn). + The message carries the URI ONLY \u2014 never an on-disk path, container, or + account \u2014 so a future Azure backend can raise the same type/message + shape and no caller (or log line) ever learns which backend is in use. + """ + + +# --------------------------------------------------------------------------- +# BlobReference \u2014 cheap handle: identity + metadata, NO payload, NO Path +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class BlobReference: + """Cheap handle \u2014 identity + metadata, NO payload, NO Path. + + This is what ``scan()``/``list()`` return and what everything except a + payload read passes around. It is what gets serialized on the graph (as + its ``.uri``). + """ + + uri: str # ci-blob:/// \u2014 the ONLY address callers use + session_id: str + key: str + size: int # content length in bytes + last_modified: float # epoch seconds: disk st_mtime || azure Last-Modified + + +# --------------------------------------------------------------------------- +# BlobStore protocol +# --------------------------------------------------------------------------- + + +@runtime_checkable +class BlobStore(Protocol): + """Protocol for a session-scoped, URI-addressable blob store. + + 100% backend-neutral: the only identity that crosses the boundary is the + ``ci-blob://`` URI (carried by :class:`BlobReference`). No ``Path``, no + on-disk layout, no ``dest_dir``, no ``os.*`` \u2014 ever. + """ + + async def write( + self, session_id: str, key: str, value: dict[str, Any] | list[Any] + ) -> BlobReference: + """Persist *value* as JSON and return a :class:`BlobReference`.""" + ... + + def list(self, session_id: str) -> AsyncIterator[BlobReference]: + """Stream all blob references for *session_id* (one session).""" + ... + + def scan(self) -> AsyncIterator[BlobReference]: + """Stream all blob references across ALL sessions.""" + ... + + async def delete( + self, uri: str, if_unmodified: BlobReference | None = None + ) -> bool: + """Delete the blob addressed by *uri*. Idempotent: returns False if absent. + + Args: + uri: The ``ci-blob://`` URI to delete. + if_unmodified: When provided, this is a **fenced (compare-and-delete)** + delete \u2014 the store re-checks the blob's current metadata against + *if_unmodified* (disk: mtime + size; Azure: ``If-Match`` ETag) and + refuses (returns ``False``, does NOT delete) if the blob changed + since *if_unmodified* was observed (e.g. by a `scan()`). When + ``None`` (default), this is the unconditional idempotent delete. + """ + ... + + async def read(self, uri: str) -> dict[str, Any] | list[Any]: + """Resolve *uri* and return the stored value (the sole payload path). + + Raises: + ValueError: If *uri* does not match the ``ci-blob://`` scheme. + BlobNotFoundError: If no blob exists for *uri*. Subclasses + ``FileNotFoundError`` for back-compat; the message carries the + URI only \u2014 never an on-disk path, container, or account. + """ + ... diff --git a/context_intelligence_server/config.py b/context_intelligence_server/config.py index 48b20e97..82837447 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,56 +595,81 @@ 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 # ------------------------------------------------------------------------- + # blob_backend selects the BlobStore implementation via + # context_intelligence_server.blob_store.create_blob_store(); "filesystem" + # is the only backend implemented today. blob_path is that filesystem + # backend's own root -- read only by the factory (and here, at + # declaration) and otherwise meaningless to any other backend. + blob_backend: str = "filesystem" blob_path: str = "/data/blobs" queues_path: str = "/data/queues" # ------------------------------------------------------------------------- # 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 +682,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 +700,158 @@ 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 + + # Maintenance mode: the live schema-health gate + /admin/maintenance repair. + maintenance_probe_ttl_seconds: float = 5.0 # :Node constraint probe cache TTL + maintenance_retry_after_seconds: int = 30 # Retry-After on the maintenance 503 + # Bounded pre-repair quiesce so ordinary in-flight flushes settle before + # run_repair (drain poll interval is 0.05s); a flush outliving it is a + # residual risk the constraint create then fails loud on. + maintenance_quiesce_seconds: float = 2.0 + + @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/content_block.py b/context_intelligence_server/handlers/data_layer_2/content_block.py index c0238f79..627ee5bc 100644 --- a/context_intelligence_server/handlers/data_layer_2/content_block.py +++ b/context_intelligence_server/handlers/data_layer_2/content_block.py @@ -49,11 +49,13 @@ async def __call__(self, event: str, data: dict[str, Any]) -> HookResult: return HookResult(action="continue") block_index = data.get("block_index") + # Key the block off the FULL active_iteration_id so a block inherits the + # iteration's run scope (including the run tiebreaker). Using only the + # trailing iteration number would let two runs that share an iteration + # number collide on the same block_node_id and MERGE-overwrite. iteration_id = self.services.data_layer_2.active_iteration_id - # ID format is "{session_id}::iteration::{n}"; [-1] extracts the iteration number. - # If the cursor format ever changes, this extraction must be updated to match. - iteration_n = iteration_id.split("::")[-1] if iteration_id else "0" - block_node_id = f"{session_id}::block::{iteration_n}::{block_index}" + iteration_key = iteration_id if iteration_id else f"{session_id}::iteration::0" + block_node_id = f"{iteration_key}::block::{block_index}" if event == "content_block:start": await self._handle_start(session_id, block_node_id, block_index, data) diff --git a/context_intelligence_server/handlers/data_layer_2/iteration.py b/context_intelligence_server/handlers/data_layer_2/iteration.py index 4c10d1e8..7088261b 100644 --- a/context_intelligence_server/handlers/data_layer_2/iteration.py +++ b/context_intelligence_server/handlers/data_layer_2/iteration.py @@ -36,6 +36,18 @@ class IterationHandler: def __init__(self, services: HookStateService) -> None: self.services = services + def _current_iteration_scope(self) -> str: + """The 'run' | 'unscoped' discriminator, sourced from the SAME cursor + field (``execution_start_ts``) used to decide the node_id shape, so all + three upsert_node call sites (provider:request, llm:request, + llm:response) always agree. + """ + return ( + "run" + if self.services.data_layer_2.execution_start_ts is not None + else "unscoped" + ) + async def __call__(self, event: str, data: dict[str, Any]) -> HookResult: """Dispatch to the appropriate sub-handler. @@ -64,22 +76,51 @@ async def _handle_provider_request( ) -> None: """Create Iteration node and set active_iteration_id cursor. - - Computes iteration_id as '{session_id}::iteration::{iteration_number}' + - Computes iteration_id as run-scoped, reusing the active run's full + orch_run_id (which carries the run tiebreaker): + '{orch_run_id}::iteration::{iteration_number}' when a run is active, + falling back to the bare '{session_id}::iteration::{iteration_number}' + shape when no run is active. - Sets active_iteration_id cursor on DataLayer2State - Creates Iteration:SST_EVENT node with session_id, iteration_number, started_at - - Conditionally creates E06: OrchestratorRun -[:HAS_PART {sst_semantic: 'CONTAINS'}]-> - Iteration when execution_start_ts cursor is set + - Conditionally creates the OrchestratorRun -[:HAS_PART]-> Iteration edge + when a run is active. + + Without run-scoping, iteration_number alone (a counter scoped to the whole + session, not the run) can repeat across orchestrator runs -- e.g. after a + drainer restart resets the in-memory counter -- causing distinct runs' + Iteration nodes to MERGE onto the same node_id and their usage figures to + clobber each other. """ # Increment counter to get the next iteration number self.services.data_layer_2.iteration_count += 1 iteration_number = self.services.data_layer_2.iteration_count timestamp: str = data.get("timestamp", "") - iteration_id = f"{session_id}::iteration::{iteration_number}" + + orch_run_id = self.services.data_layer_2.active_orch_run_id + if orch_run_id is not None: + iteration_id = f"{orch_run_id}::iteration::{iteration_number}" + else: + iteration_id = f"{session_id}::iteration::{iteration_number}" # Set cursor so llm:request and llm:response can find this iteration self.services.data_layer_2.active_iteration_id = iteration_id + # A queryable discriminator between a run-scoped iteration and a + # legitimate loop-basic session with no active orchestrator run. + iteration_scope = self._current_iteration_scope() + if iteration_scope == "unscoped": + # INFO not WARNING: a loop-basic session with no execution:start is + # a normal case, not an alert-worthy anomaly. + logger.info( + "unscoped_iteration_emitted session=%s iteration_number=%d iteration_id=%s", + session_id, + iteration_number, + iteration_id, + extra={"session_id": session_id}, + ) + # Create the Iteration node await self.services.graph.upsert_node( iteration_id, @@ -88,13 +129,12 @@ async def _handle_provider_request( "session_id": session_id, "iteration_number": iteration_number, "started_at": timestamp, + "iteration_scope": iteration_scope, }, ) # E06 (conditional): OrchestratorRun -[:HAS_PART {sst_semantic: 'CONTAINS'}]-> Iteration - execution_start_ts = self.services.data_layer_2.execution_start_ts - if execution_start_ts is not None: - orch_run_id = f"{session_id}::orch_run::{execution_start_ts}" + if orch_run_id is not None: await self.services.graph.upsert_edge( orch_run_id, iteration_id, @@ -127,6 +167,12 @@ async def _handle_llm_request(self, data: dict[str, Any]) -> None: "model": data.get("model"), "message_count": data.get("message_count"), "has_system": data.get("has_system"), + # Stamp independently of provider:request's own write -- an + # Iteration node must never be created/updated without a scope + # value, even if this write is the first one + # to ever reach the node (e.g. a dead-lettered provider:request + # whose cursor mutation nonetheless survived). + "iteration_scope": self._current_iteration_scope(), }, ) @@ -157,6 +203,9 @@ async def _handle_llm_response(self, data: dict[str, Any]) -> None: "usage_input": usage.get("input_tokens"), "usage_output": usage.get("output_tokens"), "usage_cache_write": usage.get("cache_creation_input_tokens"), + # See _handle_llm_request -- same completeness rationale applies + # to this, the third of the three sites. + "iteration_scope": self._current_iteration_scope(), }, ) diff --git a/context_intelligence_server/handlers/data_layer_2/orchestrator_run.py b/context_intelligence_server/handlers/data_layer_2/orchestrator_run.py index 45cb7bb3..7c26ce35 100644 --- a/context_intelligence_server/handlers/data_layer_2/orchestrator_run.py +++ b/context_intelligence_server/handlers/data_layer_2/orchestrator_run.py @@ -67,7 +67,7 @@ async def _handle_execution_start( ) -> None: """Create OrchestratorRun node and wire E01 (and optionally E14). - - Computes orch_run_id as '{session_id}::orch_run::{timestamp}' + - Computes orch_run_id as '{session_id}::orch_run::{timestamp}::{seq}' - Sets execution_start_ts cursor on DataLayer2State - Creates OrchestratorRun:SST_EVENT node with session_id + started_at - Creates E01: Session -[:HAS_EXECUTION {sst_semantic: 'CONTAINS'}]-> OrchestratorRun @@ -75,10 +75,16 @@ async def _handle_execution_start( OrchestratorRun when last_prompt_id cursor is set """ timestamp: str = data.get("timestamp", "") - orch_run_id = f"{session_id}::orch_run::{timestamp}" - - # Store cursor so execution:end and orchestrator:complete can find this run + # Tiebreaker: two runs sharing an identical timestamp (coarse clock or a + # replayed execution:start) must not collide on the same orch_run_id. + self.services.data_layer_2.orch_run_seq += 1 + seq = self.services.data_layer_2.orch_run_seq + orch_run_id = f"{session_id}::orch_run::{timestamp}::{seq}" + + # Store cursors so execution:end, orchestrator:complete, and the + # Iteration/ContentBlock handlers all reuse this exact id. self.services.data_layer_2.execution_start_ts = timestamp + self.services.data_layer_2.active_orch_run_id = orch_run_id # Create the OrchestratorRun node await self.services.graph.upsert_node( @@ -130,7 +136,10 @@ async def _handle_execution_end( if ts is None: return - orch_run_id = f"{session_id}::orch_run::{ts}" + orch_run_id = ( + self.services.data_layer_2.active_orch_run_id + or f"{session_id}::orch_run::{ts}" + ) timestamp: str = data.get("timestamp", "") node_data: dict[str, Any] = { @@ -167,7 +176,10 @@ async def _handle_orchestrator_complete( if ts is None: return - orch_run_id = f"{session_id}::orch_run::{ts}" + orch_run_id = ( + self.services.data_layer_2.active_orch_run_id + or f"{session_id}::orch_run::{ts}" + ) timestamp: str = data.get("timestamp", "") orchestrator: str = data.get("orchestrator", "") turn_count: Any = data.get("turn_count") @@ -213,3 +225,4 @@ async def _handle_orchestrator_complete( # Update cursors self.services.data_layer_2.last_completed_orch_run_id = orch_run_id self.services.data_layer_2.execution_start_ts = None + self.services.data_layer_2.active_orch_run_id = None diff --git a/context_intelligence_server/handlers/data_layer_2/session.py b/context_intelligence_server/handlers/data_layer_2/session.py index a39118f1..b0ccb7f3 100644 --- a/context_intelligence_server/handlers/data_layer_2/session.py +++ b/context_intelligence_server/handlers/data_layer_2/session.py @@ -68,72 +68,84 @@ class SessionLabelStateMachine: ForkedSession > SubSession > RootSession in specificity (terminal ordering). """ + @staticmethod + def _heal_forward(transition: LabelTransition) -> LabelTransition: + """Strip a stale IncompleteSession marker on a start/fork transition. + + IncompleteSession is only correct at session:end on a bare session whose + start/fork was lost. Its co-occurrence with a start/fork is always the + out-of-order case -- a forked sub-session's session:end drained before + its session:start/fork, stamping the marker first. Removing an absent + label is a no-op, so applying this to every start/fork result is safe and + makes the invariant impossible to miss when a branch is added later. + """ + if "IncompleteSession" in transition.remove: + return transition + return LabelTransition( + add=transition.add, remove=[*transition.remove, "IncompleteSession"] + ) + 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). Start/fork + # results are passed through _heal_forward so a stale IncompleteSession + # marker is stripped on the transition (it is only correct at end). if event == "start": - if current_type in ("ForkedSession", "SubSession"): - return LabelTransition() - if current_type == "RootSession": - if has_parent: - return LabelTransition( - add=["SubSession", "SST_EVENT"], - remove=["RootSession", "StubSession"], - ) + return self._heal_forward(self._classify_start(current_type, has_parent)) + + if event == "fork": + return self._heal_forward(self._classify_fork(current_type, has_parent)) + + if event == "end": + if current_type is not None: return LabelTransition() - # bare session (current_type is None) - if has_parent: - return LabelTransition( - add=["Session", "SubSession", "SST_EVENT"], - remove=["StubSession"], - ) + # Bare session: start/fork was permanently lost. Mark + # IncompleteSession rather than fabricating a real terminal. return LabelTransition( - add=["RootSession", "Session", "SST_EVENT"], - remove=["StubSession"], + add=["IncompleteSession", "SST_EVENT"], remove=["StubSession"] ) - if event == "fork": - if current_type == "ForkedSession": - return LabelTransition() - if current_type in ("RootSession", "SubSession"): + raise ValueError(f"classify() received unknown event: {event!r}") + + @staticmethod + def _classify_start(current_type: str | None, has_parent: bool) -> LabelTransition: + if current_type in ("ForkedSession", "SubSession"): + return LabelTransition() + if current_type == "RootSession": + if has_parent: return LabelTransition( - add=["ForkedSession", "SST_EVENT"], - remove=[current_type, "StubSession"], + add=["SubSession", "SST_EVENT"], + remove=["RootSession", "StubSession"], ) - # bare session (current_type is None) + return LabelTransition() + # bare session (current_type is None) + if has_parent: return LabelTransition( - add=["Session", "ForkedSession", "SST_EVENT"], + add=["Session", "SubSession", "SST_EVENT"], remove=["StubSession"], ) + return LabelTransition( + add=["RootSession", "Session", "SST_EVENT"], + remove=["StubSession"], + ) - 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. + @staticmethod + def _classify_fork(current_type: str | None, has_parent: bool) -> LabelTransition: + if current_type == "ForkedSession": + return LabelTransition() + if current_type in ("RootSession", "SubSession"): return LabelTransition( - add=["IncompleteSession", "SST_EVENT"], remove=["StubSession"] + add=["ForkedSession", "SST_EVENT"], + remove=[current_type, "StubSession"], ) - - raise ValueError(f"classify() received unknown event: {event!r}") + # bare session (current_type is None) + return LabelTransition( + add=["Session", "ForkedSession", "SST_EVENT"], + remove=["StubSession"], + ) class SessionHandler: @@ -184,9 +196,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 +214,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 +224,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 +258,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 +285,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 +307,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 +334,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 +352,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 +368,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/handlers/data_layer_2/state.py b/context_intelligence_server/handlers/data_layer_2/state.py index 47b363c2..37b8ec1a 100644 --- a/context_intelligence_server/handlers/data_layer_2/state.py +++ b/context_intelligence_server/handlers/data_layer_2/state.py @@ -16,6 +16,19 @@ class DataLayer2State: # OrchestratorRun identity execution_start_ts: str | None = None + # Monotonic per-session run counter, incremented on each execution:start and + # folded into orch_run_id. Two runs that share an identical execution_start_ts + # (coarse clock, or a replayed execution:start) would otherwise collide on the + # same orch_run_id and MERGE-overwrite each other's nodes. Lives here so the + # durable cursor persists it across a worker rebuild. + orch_run_seq: int = 0 + + # The full orch_run_id of the active run, including the tiebreaker. Set at + # execution:start and read back by execution:end/orchestrator:complete and + # the Iteration/ContentBlock handlers, so every node in one run agrees on the + # same id without any site recomputing it from execution_start_ts alone. + active_orch_run_id: str | None = None + # Iteration cursor read by ContentBlockHandler + ToolCallHandler active_iteration_id: str | None = None @@ -28,6 +41,12 @@ class DataLayer2State: # E15 OrchestratorRun→Prompt turn-flow cursor last_completed_orch_run_id: str | None = None - # Iteration counter — incremented on each provider:request; used to compute - # iteration_id as '{session_id}::iteration::{iteration_count}' + # Iteration counter — incremented on each provider:request; combined with + # execution_start_ts (via IterationHandler) to compute the run-scoped + # iteration_id '{session_id}::orch_run::{execution_start_ts}::{seq}::iteration::{iteration_count}' + # (falls back to the bare '{session_id}::iteration::{iteration_count}' shape when no + # orchestrator run is active). Scoped per-session, not reset per run: uniqueness across + # runs comes from the orch_run_id prefix, not from this counter — + # resetting it would collide with ContentBlockHandler's block_node_id derivation, which + # keys solely off this counter's value, not the run. iteration_count: int = 0 diff --git a/context_intelligence_server/handlers/data_layer_3/delegation.py b/context_intelligence_server/handlers/data_layer_3/delegation.py index efeda5c8..e837d0ad 100644 --- a/context_intelligence_server/handlers/data_layer_3/delegation.py +++ b/context_intelligence_server/handlers/data_layer_3/delegation.py @@ -52,12 +52,17 @@ def _discriminate_root_vs_unresolved(parent_session: dict[str, Any] | None) -> s -> "unresolved" (fails loud; monitored; never mistaken for a real answer) CRITICAL: branches on the TERMINAL label ONLY, never on - ``IncompleteSession``. Live graph data shows ``IncompleteSession`` - co-labels a terminal label ~41% of the time (a session can reach - session:end with session:start/fork permanently missed, out of order -- - see SessionHandler._handle_end). Treating ``IncompleteSession`` as a - discriminator would mis-flag hundreds of genuine root/forked sessions as - unresolved. + ``IncompleteSession``. Historically ``IncompleteSession`` co-labeled a + terminal label ~41% of the time (a session's session:end drained before + its session:start/fork, out of order -- see SessionHandler._handle_end). + As of the heal-forward fix, SessionLabelStateMachine.classify() now + strips ``IncompleteSession`` the moment a real start/fork is processed, + so co-occurrence is + no longer expected for newly-processed events -- only pending a one-off + backfill for historical nodes written before the fix. This function + still ignores ``IncompleteSession`` as a discriminator regardless: it was + never a reliable signal and using it would mis-flag genuine root/forked + sessions as unresolved. """ labels: list[str] = (parent_session or {}).get("labels", []) if "RootSession" in labels: 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/identity_store/__init__.py b/context_intelligence_server/identity_store/__init__.py new file mode 100644 index 00000000..1cf81b75 --- /dev/null +++ b/context_intelligence_server/identity_store/__init__.py @@ -0,0 +1,29 @@ +"""identity_store -- durable, write-through key -> contributor-identity map. + +The public surface is the backend-neutral :class:`IdentityStore` Protocol plus +the :func:`create_identity_store` factory. Consumers should depend on these, +never on a concrete backend class. + +Package layout: + protocol.py IdentityStore Protocol -- the backend-neutral seam (no + filesystem imports). AUTH-CRITICAL commit-order and + fail-closed-load guarantees are documented here. + filesystem.py FileSystemIdentityStore -- the JSON-file-backed + implementation. + factory.py create_identity_store(settings, kind) -- the ONLY place a + backend is selected and the ONLY place (besides config.py) + that reads settings.entra_identities_store_path / + settings.api_keys_store_path. +""" + +from __future__ import annotations + +from .factory import create_identity_store +from .filesystem import FileSystemIdentityStore +from .protocol import IdentityStore + +__all__ = [ + "FileSystemIdentityStore", + "IdentityStore", + "create_identity_store", +] diff --git a/context_intelligence_server/identity_store/factory.py b/context_intelligence_server/identity_store/factory.py new file mode 100644 index 00000000..e2d64c66 --- /dev/null +++ b/context_intelligence_server/identity_store/factory.py @@ -0,0 +1,49 @@ +"""Config-driven IdentityStore factory -- the ONLY place a backend is selected. + +This is the single seam through which the concrete backend is chosen. Adding +a new backend (e.g. Azure) means: one new module implementing +:class:`~.protocol.IdentityStore`, one new branch here -- zero changes to +:mod:`context_intelligence_server.main` or any other consumer. + +This module (and :mod:`~context_intelligence_server.config`) are the only +places ``settings.entra_identities_store_path`` / ``settings.api_keys_store_path`` +are read -- the on-disk location is a filesystem-backend concern, resolved +here and handed to the concrete backend at construction time. Callers only +ever see the :class:`~.protocol.IdentityStore` Protocol. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import TYPE_CHECKING + +from .filesystem import FileSystemIdentityStore +from .protocol import IdentityStore + +if TYPE_CHECKING: + from context_intelligence_server.config import Settings + + +def create_identity_store(settings: Settings, kind: str) -> IdentityStore: + """Build an :class:`~.protocol.IdentityStore` rooted at the configured path for *kind*. + + Construction only -- callers are responsible for ``load()`` / ``seed()`` + and any auth-mode wiring (this mirrors + ``blob_store.factory.create_blob_store``: a single-backend, config-reading + seam that keeps the store paths out of consumers such as ``main.py``). + + Args: + settings: The active ``Settings``. + kind: ``"entra"`` for the OID identity map, ``"api_key"`` for the + SHA-256 digest keystore. + + Raises: + ValueError: If *kind* is neither ``"entra"`` nor ``"api_key"``. + """ + if kind == "entra": + path = settings.entra_identities_store_path + elif kind == "api_key": + path = settings.api_keys_store_path + else: + raise ValueError(f"Unknown identity store kind: {kind!r}") + return FileSystemIdentityStore(Path(path)) diff --git a/context_intelligence_server/identity_store.py b/context_intelligence_server/identity_store/filesystem.py similarity index 71% rename from context_intelligence_server/identity_store.py rename to context_intelligence_server/identity_store/filesystem.py index a7da96e7..c18d92d6 100644 --- a/context_intelligence_server/identity_store.py +++ b/context_intelligence_server/identity_store/filesystem.py @@ -1,19 +1,29 @@ -"""Durable identity-map store for the Context Intelligence Server. +"""FileSystemIdentityStore -- disk-backed identity map with atomic writes. -Each ``IdentityStore`` wraps ONE JSON file and keeps an in-process dict -(``_data``) that IS the live source of truth for the single-replica process. -A second derived dict, ``flat_dict``, exposes ``{key: contributor_id}`` and is -kept in-sync with ``_data`` via in-place mutations so that any object holding a -reference to ``flat_dict`` always sees the latest state without a restart. +Each ``FileSystemIdentityStore`` wraps ONE JSON file and keeps an in-process +dict (``_data``) that IS the live source of truth for the single-replica +process. A second derived dict, ``flat_dict``, exposes ``{key: contributor_id}`` +and is kept in-sync with ``_data`` via in-place mutations so that any object +holding a reference to ``flat_dict`` always sees the latest state without a +restart. -**Commit order (ROB F2 — NON-NEGOTIABLE)** +This is a concrete implementation of the +:class:`~context_intelligence_server.identity_store.protocol.IdentityStore` +Protocol. No ``Path``, on-disk layout, or ``os.*`` detail appears in the +Protocol or in any value it returns -- those details are private to this +class (and, later, an Azure equivalent). + +**Commit order (ROB F2 -- NON-NEGOTIABLE)** On every mutation (put / delete): 1. Build the new data dict (do NOT touch ``_data`` yet). -2. Serialize and write to a tempfile **in the same directory** as the target file. -3. ``os.replace()`` the tempfile onto the target (atomic rename on POSIX / Azure Files). -4. **ONLY IF the above succeeds**: update ``_data`` and ``flat_dict`` in-place. +2. Serialize and write to a tempfile **in the same directory** as the target + file. +3. ``os.replace()`` the tempfile onto the target (atomic rename on POSIX / + Azure Files). +4. **ONLY IF the above succeeds**: update ``_data`` and ``flat_dict`` + in-place. If the file write raises for any reason, ``_data`` and ``flat_dict`` are **unchanged** and the exception propagates to the caller (who returns 5xx). @@ -23,13 +33,13 @@ On ``load()``: -- Missing file → empty dict (normal first boot). No log, no raise. -- Corrupt / torn / partial / invalid-JSON file → **empty dict + a LOUD - ``logger.error`` / ``logger.critical``**. The server MUST NOT crash-loop on - a bad store file. An empty map means "nobody is bound yet" — every auth +- Missing file -> empty dict (normal first boot). No log, no raise. +- Corrupt / torn / partial / invalid-JSON file -> **empty dict + a LOUD + ``logger.error`` / ``logger.critical``**. The server MUST NOT crash-loop on + a bad store file. An empty map means "nobody is bound yet" -- every auth attempt then fails normally until an admin re-populates via the /admin API. -File format (both modes share the same abstraction):: +File format (both modes share the same on-disk shape):: # api-keys.json { @@ -42,6 +52,8 @@ } """ +from __future__ import annotations + import json import logging import os @@ -51,18 +63,19 @@ logger = logging.getLogger(__name__) -class IdentityStore: +class FileSystemIdentityStore: """Durable, write-through identity map backed by a single JSON file. - See module docstring for the commit-order contract and fail-closed guarantees. + See module docstring for the commit-order contract and fail-closed + guarantees. Args: - path: Absolute path to the JSON store file. The parent directory is + path: Absolute path to the JSON store file. The parent directory is created automatically on the first write. """ def __init__(self, path: Path) -> None: - self.path = path + self._path = path # Rich format: {key: {id: ..., display_name?: ...}} self._data: dict[str, dict[str, str]] = {} # Flat derived cache: {key: contributor_id}. @@ -77,23 +90,23 @@ def __init__(self, path: Path) -> None: def load(self) -> None: """Read the file and populate the in-process map. - Missing file → empty dict (normal first boot, no log). - Corrupt / non-dict → empty dict + LOUD error log, never raise. + Missing file -> empty dict (normal first boot, no log). + Corrupt / non-dict -> empty dict + LOUD error log, never raise. """ - if not self.path.exists(): - # Normal first boot — the file hasn't been written yet. + if not self._path.exists(): + # Normal first boot -- the file hasn't been written yet. self._data = {} self._rebuild_flat() return raw: object try: - raw = json.loads(self.path.read_text(encoding="utf-8")) + raw = json.loads(self._path.read_text(encoding="utf-8")) except (json.JSONDecodeError, UnicodeDecodeError, OSError) as exc: logger.error( - "identity_store.load CORRUPT FILE path=%s error=%r — " + "identity_store.load CORRUPT FILE path=%s error=%r -- " "failing CLOSED to empty map. Re-populate via /admin API.", - self.path, + self._path, exc, ) self._data = {} @@ -102,9 +115,9 @@ def load(self) -> None: if not isinstance(raw, dict): logger.critical( - "identity_store.load INVALID FORMAT path=%s got=%r — " + "identity_store.load INVALID FORMAT path=%s got=%r -- " "expected a JSON object at top level. Failing CLOSED to empty map.", - self.path, + self._path, type(raw).__name__, ) self._data = {} @@ -118,15 +131,15 @@ def load(self) -> None: self._rebuild_flat() def put(self, key: str, value: dict[str, str]) -> None: - """Upsert *key* → *value*. + """Upsert *key* -> *value*. - Commit order (F2): write tempfile → os.replace → update in-process. + Commit order (F2): write tempfile -> os.replace -> update in-process. Raises on file-write failure; in-process state is UNCHANGED. """ new_data = dict(self._data) new_data[key] = value self._write_atomic(new_data) - # File write succeeded — now update in-process state. + # File write succeeded -- now update in-process state. self._data[key] = value contributor_id = value.get("id", "") if contributor_id: @@ -137,13 +150,13 @@ def put(self, key: str, value: dict[str, str]) -> None: def delete(self, key: str) -> None: """Remove *key* from the store. - Commit order (F2): write tempfile → os.replace → update in-process. + Commit order (F2): write tempfile -> os.replace -> update in-process. Raises on file-write failure; in-process state is UNCHANGED. No-op if *key* is not present. """ new_data = {k: v for k, v in self._data.items() if k != key} self._write_atomic(new_data) - # File write succeeded — now update in-process state. + # File write succeeded -- now update in-process state. self._data.pop(key, None) self.flat_dict.pop(key, None) @@ -166,12 +179,12 @@ def seed(self, data: dict[str, dict[str, str]]) -> None: except Exception as exc: logger.warning( "identity_store.seed: could not write seed to %s: %r " - "— in-memory map is live but the file is not yet persisted. " + "-- in-memory map is live but the file is not yet persisted. " "The next mutation via /admin API will persist the file.", - self.path, + self._path, exc, ) - # Update in-memory regardless — data is from durable config. + # Update in-memory regardless -- data is from durable config. self._data = dict(data) self._rebuild_flat() @@ -186,6 +199,10 @@ def items(self): # type: ignore[override] def __len__(self) -> int: return len(self._data) + def exists(self) -> bool: + """Whether the store has ever been persisted to its backing file.""" + return self._path.exists() + # ------------------------------------------------------------------ # Internals # ------------------------------------------------------------------ @@ -203,7 +220,7 @@ def _rebuild_flat(self) -> None: self.flat_dict[key] = contributor_id def _write_atomic(self, data: dict[str, dict[str, str]]) -> None: - """Write *data* atomically to ``self.path``. + """Write *data* atomically to ``self._path``. Steps: 1. Create parent directory (parents=True, exist_ok=True). @@ -216,15 +233,15 @@ def _write_atomic(self, data: dict[str, dict[str, str]]) -> None: Raises the underlying OS/IO exception so the caller (put/delete) knows the write failed and leaves in-process state unchanged. """ - self.path.parent.mkdir(parents=True, exist_ok=True) - tmp_fd, tmp_str = tempfile.mkstemp(dir=str(self.path.parent), suffix=".tmp") + self._path.parent.mkdir(parents=True, exist_ok=True) + tmp_fd, tmp_str = tempfile.mkstemp(dir=str(self._path.parent), suffix=".tmp") tmp_path = Path(tmp_str) try: with os.fdopen(tmp_fd, "w", encoding="utf-8") as fh: json.dump(data, fh, indent=2, ensure_ascii=False) fh.flush() os.fsync(fh.fileno()) - os.replace(str(tmp_path), str(self.path)) + os.replace(str(tmp_path), str(self._path)) except Exception: # Best-effort cleanup of the tempfile before propagating. try: diff --git a/context_intelligence_server/identity_store/protocol.py b/context_intelligence_server/identity_store/protocol.py new file mode 100644 index 00000000..3182eb69 --- /dev/null +++ b/context_intelligence_server/identity_store/protocol.py @@ -0,0 +1,117 @@ +"""IdentityStore Protocol -- the backend-neutral seam. + +Each ``IdentityStore`` wraps a durable ``key -> {"id": contributor_id, ...}`` +map and keeps an in-process, live-mutated view of it (:attr:`flat_dict`) so +that any object holding a reference (e.g. an auth resolver's keystore) always +sees the latest state without a restart. No ``Path``, on-disk layout, or +``os.*`` detail appears here or in any value the Protocol returns -- that is +private to a concrete backend +(:class:`~context_intelligence_server.identity_store.filesystem.FileSystemIdentityStore` +and, later, an Azure equivalent). + +This is AUTH-CRITICAL surface. The following guarantees are part of the +Protocol's contract and every backend MUST uphold them: + +**Commit order (ROB F2 -- NON-NEGOTIABLE)** + +On every mutation (``put`` / ``delete``): + +1. The durable store is written FIRST (whatever "durable" means for the + backend -- a file, a blob, etc.). +2. **ONLY IF** that write succeeds does in-process state (``flat_dict`` and + any internal map) get updated, IN-PLACE, so existing references to + ``flat_dict`` observe the change immediately. +3. If the durable write fails, in-process state is **UNCHANGED** and the + exception propagates to the caller (who returns 5xx). The durable store + and in-process memory are never out of sync. + +**Fail-CLOSED load()** + +On ``load()``: + +- Missing / never-persisted store -> empty map (normal first boot). No log, + no raise. +- Corrupt / unreadable store -> empty map + a LOUD error log. The server + MUST NOT crash-loop on a bad store. An empty map means "nobody is bound + yet" -- every auth attempt then fails normally until an admin re-populates + via the /admin API. + +**flat_dict is a live, shared view** + +``flat_dict`` is ``{key: contributor_id}``, mutated IN-PLACE on every +``put`` / ``delete`` / ``seed`` / ``load`` so any object holding a reference +to it always sees the latest state. + +This module is backend-neutral by construction: it imports nothing from +``os``, ``pathlib``, or any filesystem library. +""" + +from __future__ import annotations + +from collections.abc import ItemsView +from typing import Protocol, runtime_checkable + + +@runtime_checkable +class IdentityStore(Protocol): + """Protocol for a durable, write-through identity map. + + 100% backend-neutral: no ``Path``, no on-disk layout, no ``os.*`` -- + ever. + """ + + flat_dict: dict[str, str] + """Live derived view: ``{key: contributor_id}``. Shared BY REFERENCE with + consumers (e.g. auth resolvers) so mutations are visible with no restart.""" + + def load(self) -> None: + """Read the durable store and populate the in-process map. + + Fail-closed: missing store -> empty map (silent, normal first boot); + corrupt store -> empty map + a loud error log. Never raises. + """ + ... + + def put(self, key: str, value: dict[str, str]) -> None: + """Upsert ``key`` -> ``value``. + + Commit order (F2): durable write -> in-process update. Raises on + durable-write failure; in-process state is left unchanged. + """ + ... + + def delete(self, key: str) -> None: + """Remove ``key`` from the store. + + Commit order (F2): durable write -> in-process update. Raises on + durable-write failure; in-process state is left unchanged. No-op if + ``key`` is not present. + """ + ... + + def seed(self, data: dict[str, dict[str, str]]) -> None: + """Bulk-seed from config on first boot. + + Unlike ``put()`` (which enforces F2 write-before-memory strictly), + in-process state is updated even if the durable write fails -- the + data came from durable config, so memory-ahead-of-durable-store is + safe (a restart re-seeds from config again). A warning is logged if + the durable write fails. + """ + ... + + def get(self, key: str) -> dict[str, str] | None: + """Return the value for ``key``, or ``None`` if not present.""" + ... + + def items(self) -> ItemsView[str, dict[str, str]]: + """Iterate over ``(key, value)`` pairs in the store.""" + ... + + def __len__(self) -> int: + """Return the number of entries currently in the store.""" + ... + + def exists(self) -> bool: + """Whether the store has ever been persisted to its backing store.""" + ... diff --git a/context_intelligence_server/lease_store/__init__.py b/context_intelligence_server/lease_store/__init__.py new file mode 100644 index 00000000..ca5fe090 --- /dev/null +++ b/context_intelligence_server/lease_store/__init__.py @@ -0,0 +1,8 @@ +"""Writer-lease persistence behind a backend-neutral Protocol.""" + +from __future__ import annotations + +from context_intelligence_server.lease_store.factory import create_lease_store +from context_intelligence_server.lease_store.protocol import LeaseRecord, LeaseStore + +__all__ = ["LeaseRecord", "LeaseStore", "create_lease_store"] diff --git a/context_intelligence_server/lease_store/factory.py b/context_intelligence_server/lease_store/factory.py new file mode 100644 index 00000000..0f55411d --- /dev/null +++ b/context_intelligence_server/lease_store/factory.py @@ -0,0 +1,19 @@ +"""Backend selection for the writer-lease store. + +One backend today (filesystem). A future backend is added here and nowhere +else -- the detector never learns which one it got. +""" + +from __future__ import annotations + +from collections.abc import Callable +from pathlib import Path + +from context_intelligence_server.lease_store.filesystem import FileSystemLeaseStore +from context_intelligence_server.lease_store.protocol import LeaseStore + + +def create_lease_store(dir_source: Callable[[], Path]) -> LeaseStore: + """Build the lease store. *dir_source* is resolved lazily per operation, so + nothing is constructed and no path is read at build time.""" + return FileSystemLeaseStore(dir_source) diff --git a/context_intelligence_server/lease_store/filesystem.py b/context_intelligence_server/lease_store/filesystem.py new file mode 100644 index 00000000..efdd4001 --- /dev/null +++ b/context_intelligence_server/lease_store/filesystem.py @@ -0,0 +1,89 @@ +"""Filesystem-backed writer-lease store. + +The lease is one atomically-replaced ``.writer.lease`` file in a directory +resolved lazily via *dir_source* -- resolved per operation (a cheap attribute +read, zero syscalls) so the store, like the detector it serves, constructs +nothing at build time and reflects a directory the tests may re-point. +""" + +from __future__ import annotations + +import json +import os +from collections.abc import Callable +from pathlib import Path + +from context_intelligence_server.lease_store.protocol import LeaseRecord + +LEASE_FILENAME = ".writer.lease" +LEASE_TMP_FILENAME = ".writer.lease.tmp" +_LEASE_VERSION = 1 + + +class FileSystemLeaseStore: + """A ``LeaseStore`` backed by a single atomically-written file on disk.""" + + def __init__(self, dir_source: Callable[[], Path]) -> None: + self._dir_source = dir_source + + def _path(self) -> Path: + return self._dir_source() / LEASE_FILENAME + + def read(self) -> LeaseRecord | 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, record: LeaseRecord) -> None: + directory = self._dir_source() + payload = { + "lease_version": record.lease_version, + "owner": record.owner, + "host": record.host, + "pid": record.pid, + "started_at": record.started_at, + "heartbeat": record.heartbeat, + "revision": record.revision, + "server_version": record.server_version, + } + tmp = directory / LEASE_TMP_FILENAME + tmp.write_text( + json.dumps(payload, separators=(",", ":")) + "\n", encoding="utf-8" + ) + os.replace(tmp, directory / LEASE_FILENAME) + + def delete_if_owned(self, owner: str) -> None: + rec = self.read() + if rec is not None and not rec.unreadable and rec.owner == owner: + try: + self._path().unlink() + except FileNotFoundError: + pass diff --git a/context_intelligence_server/lease_store/protocol.py b/context_intelligence_server/lease_store/protocol.py new file mode 100644 index 00000000..8281a555 --- /dev/null +++ b/context_intelligence_server/lease_store/protocol.py @@ -0,0 +1,66 @@ +"""The writer-lease persistence boundary. + +A single named lease record (owner, heartbeat, identity) persisted somewhere +durable. The writer-lease DETECTOR (``writer_lease.py``) owns all policy -- +staleness, conflict, the bounded-thread I/O executor -- and reaches the lease +only through this Protocol, so the same detector runs unchanged against any +backend (a filesystem file today, a blob lease or a row tomorrow). + +The methods are synchronous by contract: the detector runs each one on its own +private single-thread executor to bound a hung mount to a single leaked thread, +which ``asyncio.to_thread`` (shared pool) cannot guarantee. A backend whose I/O +is natively async wraps itself to satisfy this sync surface. +""" + +from __future__ import annotations + +import dataclasses +from typing import Protocol + + +@dataclasses.dataclass +class LeaseRecord: + """Parsed view of one persisted lease record. + + ``unreadable=True`` marks a synthetic record standing in for a torn or + hand-mangled lease (decode error / 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 + + +class LeaseStore(Protocol): + """Persistence for exactly one writer-lease record. + + All three operations may raise ``OSError`` (a share fault); the detector + absorbs that as "not armed", never as a conflict. + """ + + def read(self) -> LeaseRecord | None: + """Return the current lease record, or ``None`` when no lease exists + (a free directory). A torn/malformed record returns a ``LeaseRecord`` + with ``unreadable=True`` rather than ``None``.""" + ... + + def write(self, record: LeaseRecord) -> None: + """Persist *record* as the current lease, atomically (a reader never + observes a half-written record).""" + ... + + def delete_if_owned(self, owner: str) -> None: + """Delete the lease only if it is still owned by *owner*. + + Never deletes a foreign lease: if a peer took it over, removing theirs + would actively hand the directory to a third writer. Best-effort: a + lease already gone is not an error.""" + ... diff --git a/context_intelligence_server/main.py b/context_intelligence_server/main.py index 0bf0709f..186943e5 100644 --- a/context_intelligence_server/main.py +++ b/context_intelligence_server/main.py @@ -9,14 +9,15 @@ 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 pathlib import Path +from functools import partial 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 ( @@ -30,25 +31,48 @@ require_read, require_write, ) -from context_intelligence_server.blob_store import AsyncDiskBlobStore +from context_intelligence_server.blob_store import create_blob_store from context_intelligence_server.config import Neo4jClientConfig, Settings, get_settings -from context_intelligence_server.idempotency import EventIdempotencyCache -from context_intelligence_server.identity_store import IdentityStore +from context_intelligence_server.idempotency import ( + EventIdempotencyCache, + KeyedAsyncLocks, +) +from context_intelligence_server.identity_store import ( + IdentityStore, + create_identity_store, +) from context_intelligence_server.logging_config import setup_logging +from context_intelligence_server.maintenance import ( + MAINTENANCE_ALLOW_LIST, + coordinator, + maintenance_gate_middleware, +) +from context_intelligence_server.maintenance_ops import run_maintenance_operation from context_intelligence_server.models import ( CypherRequest, EventRequest, EventResponse, ) from context_intelligence_server.neo4j_store import ( + build_bounded_neo4j_driver, count_untagged_nodes, ensure_neo4j_schema, + read_graph_schema_version, ) from context_intelligence_server.registry import SessionRegistry 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 ( + SCHEMA_VERSION, + boot_state, + build_status_response, +) +from context_intelligence_server.writer_lease import ( + WriterLeaseConflict, + shutdown_lease_io, + writer_lease, +) _settings = get_settings() @@ -61,23 +85,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 +126,561 @@ 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 - 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. + 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 - 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. + logger.warning( + "recovery_skipped session=%s: torn or empty workspace in first line", + sid, + ) + return False - Returns: - True – drainer was (re)spawned via *get_or_create*. - False – session skipped (empty/torn workspace, or malformed JSON line). + +@dataclass +class TopupResult: + """Result of one ``_crash_recovery_topup`` pass. + + ``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 _schema_ready_or_auto_repair() -> None: + """Run the schema check; on un-migrated data, attempt a lease-armed repair. + + ``_ensure_schema_ready`` refuses (raises ``RuntimeError``) on un-migrated + data. Rather than let that abort boot, repair it in place when it is safe to: + + CROSS-REPLICA SAFETY: ``run_repair`` mutates the graph, so it fires ONLY when + ``writer_lease.acquired`` is True. The lease is a single-writer DETECTOR, not + a mutex -- a storage fault leaves it unarmed and boot proceeds -- so an + unarmed lease means we cannot prove we are the sole writer and MUST NOT + auto-mutate; the server stays gated and an operator repairs via + ``POST /admin/maintenance``. The repair runs under the coordinator's + single-flight (``try_begin_op``), so boot-repair and a concurrent + ``/admin/maintenance`` on this instance can never overlap. On success the + schema check is re-run to re-arm the global ``app.state.schema_ready`` gate. + + Never raises: an unrepaired graph leaves ``schema_ready`` False and a + ``degraded_reason`` set, and the caller proceeds (the periodic sweep retries). + """ + try: + await _ensure_schema_ready() + return + except RuntimeError: + pass # un-migrated data -- fall through to the repair attempt + + if not writer_lease.acquired: + reason = ( + "un-migrated graph data and the writer lease is not armed -- " + "refusing to auto-repair (cannot confirm single writer); repair via " + "POST /admin/maintenance" + ) + boot_state.degrade(reason) + logger.warning("schema_auto_repair_skipped reason=lease_unarmed") + return + + run_id = coordinator.try_begin_op() + if run_id is None: + # A maintenance op is already running (e.g. /admin/maintenance); let it + # finish rather than double-running the repair. + logger.info("schema_auto_repair_skipped reason=op_already_running") + return + + logger.info("schema_auto_repair_started run_id=%s", run_id) + # Records success/failure on the coordinator; never raises. + await run_maintenance_operation( + app.state.neo4j_driver, + run_id, + quiesce_seconds=_settings.maintenance_quiesce_seconds, + ) + try: + await _ensure_schema_ready() + except RuntimeError as exc: + boot_state.degrade(f"auto-repair ran but schema still not ready: {exc}") + logger.error("schema_auto_repair_incomplete run_id=%s", run_id) + return + if app.state.schema_ready: + boot_state.clear_degraded() + logger.info("schema_auto_repair_succeeded run_id=%s", run_id) 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 + # Enumerate through the QueueManager so the sweep is backend-agnostic -- + # the queue owns its own storage layout, this loop never reaches into it. + keys = await qm.session_keys() + 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, + qm.queues_dir / 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" + # On un-migrated data this attempts a lease-armed auto-repair instead of + # raising -- so a schema failure never aborts boot and skips the sweep + # scheduling below (the sweep is this pass's retry mechanism). + await _schema_ready_or_auto_repair() + + 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 +695,82 @@ 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), + # Wire the live admin driver into the maintenance coordinator so its gate/ + # status probe can run. Left unbound in tests, the probe returns None and the + # gate stays open -- the no-regression property for the existing suite. + coordinator.bind_driver( + app.state.neo4j_driver, + probe_ttl_seconds=_settings.maintenance_probe_ttl_seconds, ) - # 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", @@ -387,16 +778,22 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: app.include_router(admin_router) app.include_router(version_router) app.include_router(queues_router) +# Allow-list gate: refuses non-allow-listed paths while a maintenance op runs. +# Registered on `app` (not the auth-wrapped ASGI app) so it cannot be bypassed +# by the bare `app` entrypoint; auth still runs first via the outer wrapper. +app.middleware("http")(maintenance_gate_middleware) _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 +804,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 +824,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 @@ -462,15 +851,26 @@ def _assert_admin_not_exempt() -> None: ) +def _assert_maintenance_endpoint_allow_listed() -> None: + """/admin/maintenance, /status, and /version must be on the maintenance + allow-list. Without this, /admin/maintenance could be gated by maintenance + mode -- 503ing at the exact moment it exists to unblock. + """ + required = {"/admin/maintenance", "/status", "/version"} + missing = required - MAINTENANCE_ALLOW_LIST + if missing: + raise RuntimeError( + f"Availability invariant violated: {sorted(missing)!r} missing from " + f"maintenance.MAINTENANCE_ALLOW_LIST -- these paths must never be " + f"gated by maintenance mode." + ) + + 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 +890,46 @@ 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() + # /admin/maintenance, /status, /version must never be gated by maintenance + # mode -- assert before middleware construction so the failure is immediate. + _assert_maintenance_endpoint_allow_listed() 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( @@ -572,12 +948,11 @@ def create_asgi_app( if s.auth_mode == "entra": # Build and load the entra identity store. - entra_store = IdentityStore(Path(s.entra_identities_store_path)) + entra_store = create_identity_store(s, "entra") 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. + if not entra_store.exists(): + # 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,27 +960,20 @@ 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 " "is UP and serving, but every delegated (human) token will " "receive 403 until identities are onboarded. Bind the first user " - "with an IdentityAdmin-role token via PUT /admin/identities/{oid} " - "(store=%s). This is expected on a fresh /data volume.", - s.entra_identities_store_path, + "with an IdentityAdmin-role token via PUT /admin/identities/{oid}. " + "This is expected on a fresh /data volume." ) - # 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 +986,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 = create_identity_store(s, "api_key") 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. + if not key_store.exists(): + # 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,18 +1015,15 @@ 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( "static keystore is EMPTY at startup (0 bound keys) — server " "is UP but fail-CLOSED; every request will 401 until keys are " "onboarded. Add the first key with the admin token via " - "PUT /admin/keys/{sha256hash} (store=%s). Expected on a fresh " - "/data volume.", - s.api_keys_store_path, + "PUT /admin/keys/{sha256hash}. Expected on a fresh /data volume." ) else: logger.warning( @@ -670,18 +1033,15 @@ def create_asgi_app( "/admin API is unreachable without an admin key: every token " "401s at the middleware before require_admin runs). Set " "admin_api_key/admin_api_key_sha256 to enable runtime " - "onboarding, or add api_keys in config and restart. (store=%s)", - s.api_keys_store_path, + "onboarding, or add api_keys in config and restart." ) # Pass key_store.flat_dict (the LIVE dict) so the resolver sees any # 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 +1051,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 +1070,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 +1084,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 +1113,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 +1123,56 @@ 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() + # Schema-version drift: the compiled model version vs the graph's stored + # :SchemaMeta version. read_graph_schema_version returns None when the + # graph is unreachable or the baseline was never written; an absent + # driver is treated the same way -- drift=None (unknown), never a false + # "in sync" and never a 500 (/status must never raise). Gated with the + # disk reads above so /status stays graph-read-free while booting. + _schema_driver = getattr(request.app.state, "neo4j_driver", None) + graph_schema_version = ( + await read_graph_schema_version(_schema_driver) + if _schema_driver is not None + else None + ) + response["schema_version"] = SCHEMA_VERSION + response["graph_schema_version"] = graph_schema_version + response["schema_version_current"] = ( + None + if graph_schema_version is None + else graph_schema_version == SCHEMA_VERSION + ) + 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["schema_version"] = SCHEMA_VERSION + response["graph_schema_version"] = None + response["schema_version_current"] = 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 +1187,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 +1220,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,46 +1241,56 @@ 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 + # Lift the optional top-level working_dir envelope field into data so the + # Session-node write sees it; absent/empty leaves Session.working_dir null + # for a later event to populate. + if request.working_dir and isinstance(body_obj.get("data"), dict): + body_obj["data"]["working_dir"] = request.working_dir + 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)]) async def list_blobs(session_id: str) -> JSONResponse: - blob_store = AsyncDiskBlobStore(root=_settings.blob_path) - uris = await blob_store.list(session_id) + blob_store = create_blob_store(_settings) + uris = [ref.uri async for ref in blob_store.list(session_id)] return JSONResponse(content={"session_id": session_id, "blobs": uris}) @app.get("/blobs/{session_id}/{key}", dependencies=[Depends(require_read)]) async def get_blob(session_id: str, key: str) -> JSONResponse: - blob_store = AsyncDiskBlobStore(root=_settings.blob_path) + blob_store = create_blob_store(_settings) uri = f"ci-blob://{session_id}/{key}" try: content = await blob_store.read(uri) @@ -955,23 +1317,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 +1353,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 +1404,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 +1421,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/maintenance.py b/context_intelligence_server/maintenance.py new file mode 100644 index 00000000..470a6e69 --- /dev/null +++ b/context_intelligence_server/maintenance.py @@ -0,0 +1,527 @@ +"""Maintenance-mode coordinator -- the single seam for gate + status + op state. + +This module owns ALL maintenance-mode state: the tri-state live constraint +probe, the current-mode derivation, the single-flight maintenance-operation +record, and the structured 503 response. It is deliberately the ONLY place +this state lives so that the HTTP gate (``maintenance_gate_middleware``), the +drain-loop gate (``registry.drain_worker``), and the ``/status`` / +``GET /admin/maintenance`` surfaces can never drift from one another. + +Constraints (do not violate): +- This module MUST NOT import ``registry`` -- ``registry`` imports this + module, never the reverse. +- This module MUST NOT contain graph-mutation logic. Migration/repair logic + lives in ``neo4j_store.run_repair``; this module only reads (the + constraint-presence probe) and tracks in-process op state. + +Swappability (multi-replica deferred): every call site touches +only ``gate_closed`` / ``status`` / ``try_begin_op`` / ``finish_op`` / +``current_op``. Replacing the in-process op record with a graph-store-backed +lock node later means reimplementing the body of those five methods -- +zero call-site changes. + +No driver bound -- the load-bearing no-regression property: tests that never +run ``lifespan`` never call ``bind_driver``, so the probe always returns +``None`` (unknown) and the gate is never closed. This keeps the ENTIRE +existing unit-test suite's behavior unchanged. +""" + +from __future__ import annotations + +import asyncio +import logging +import time +import uuid +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from datetime import UTC, datetime +from typing import Any, Literal + +from fastapi import Request +from fastapi.responses import JSONResponse, Response + +from context_intelligence_server.config import get_settings +from context_intelligence_server.neo4j_store import count_untagged_nodes + +logger = logging.getLogger("context_intelligence_server") + +# --------------------------------------------------------------------------- +# Types +# --------------------------------------------------------------------------- + +MaintenanceMode = Literal["healthy", "maintenance", "degraded", "unknown"] +OpState = Literal["unknown", "running", "succeeded", "failed"] + + +@dataclass(frozen=True) +class OpRecord: + """Snapshot of the (per-process) maintenance operation state. + + ``state`` initializes to "unknown" so "never ran" is + distinguishable from "ran, record lost to a crash" (which would also + read as a false "succeeded" if initialized there instead). + """ + + state: OpState + run_id: str | None # uuid4 hex; freshness marker AND future fencing token + started_at: str | None # ISO-8601 UTC + completed_at: str | None # ISO-8601 UTC; set ONLY on the genuine-execution path + records_affected: int | None + error: str | None # human-readable; persists across "failed" + + +@dataclass(frozen=True) +class MaintenanceStatus: + """Snapshot of the current maintenance mode, returned by ``status()``.""" + + mode: MaintenanceMode + constraint_present: bool | None # None == probe could not answer + reason: str | None # human-readable cause of the current mode + started_at: str | None # when the CURRENT maintenance window opened + elapsed_seconds: float | None # None when not in maintenance + op: OpRecord + # Live (TTL-cached) untagged-node count backing the `degraded` term. + # Defaulted so existing MaintenanceStatus(...) constructions in tests keep + # working; None == the probe could not answer (or no driver bound). + untagged_nodes: int | None = None + + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +_CONSTRAINT_NAME = "node_node_id_workspace_unique" +_PROBE_CYPHER = ( + "SHOW CONSTRAINTS YIELD name " + f"WHERE name = '{_CONSTRAINT_NAME}' RETURN count(*) AS c" +) +_PROBE_TTL_SECONDS = 5.0 # default; overridable per-instance via bind_driver() +_RETRY_AFTER_DEFAULT = ( + 30 # fallback; live value is settings.maintenance_retry_after_seconds +) + +# Allow-list, not deny-list: a deny-list is scatter-and-miss. +# Any route added later to the app is blocked-by-default -- the safe direction. +MAINTENANCE_ALLOW_LIST: frozenset[str] = frozenset( + {"/status", "/version", "/admin/maintenance", "/docs", "/openapi.json"} +) + + +def _now_iso() -> str: + return datetime.now(UTC).isoformat() + + +# --------------------------------------------------------------------------- +# The seam +# --------------------------------------------------------------------------- + + +class MaintenanceCoordinator: + """The seam. See module docstring for the swappability contract.""" + + def __init__(self) -> None: + self._driver: Any = None + self._boot_untagged: int | None = None + self._probe_ttl_seconds: float = _PROBE_TTL_SECONDS + + # TTL-cached, single-flight probe state (constraint presence). + self._probe_lock = asyncio.Lock() + self._cache_populated: bool = False + self._cached_present: bool | None = None + self._cache_expires_at: float = 0.0 # monotonic clock + + # TTL-cached, single-flight probe state for the untagged-node count -- + # a SEPARATE cache from the constraint probe above (independent fate: + # a count failure must not poison the constraint signal, and vice + # versa). This is what de-latches the degraded/untagged half of the + # health signal: it self-clears within one TTL after an out-of-band + # repair, exactly like the constraint probe, instead of staying pinned + # to the boot-time snapshot until a restart. count_untagged_nodes is + # O(1) via Neo4j's counts store (see its docstring), so probing it on + # the cached path is as cheap as the constraint catalog read. + self._untagged_lock = asyncio.Lock() + self._untagged_cache_populated: bool = False + self._cached_untagged: int | None = None + self._untagged_cache_expires_at: float = 0.0 # monotonic clock + + # Op state -- init "unknown": never-run != crash-lost. + self._op = OpRecord( + state="unknown", + run_id=None, + started_at=None, + completed_at=None, + records_affected=None, + error=None, + ) + + # Maintenance-window bookkeeping (for /status + transition logging). + self._window_started_at: str | None = None + self._window_started_monotonic: float | None = None + + # Strong references to in-flight maintenance-op background + # tasks. asyncio only holds a WEAK reference to a task created via + # ``asyncio.create_task`` -- without an external strong ref the task + # can be garbage-collected mid-run. ``retain_task`` holds it here and + # a completion callback discards it so this set never grows unbounded. + self._background_tasks: set[asyncio.Task[Any]] = set() + + # -- binding -------------------------------------------------------- + + def bind_driver( + self, + driver: Any, + *, + untagged: int | None = None, + probe_ttl_seconds: float | None = None, + ) -> None: + """Wire the live admin Neo4j driver (called from ``lifespan``). + + ``untagged`` is the boot-time untagged-node count already computed by + ``main._record_schema_health`` -- injected here so this module never + imports ``main`` (would be circular). ``probe_ttl_seconds`` overrides + the default TTL (normally sourced from + ``settings.maintenance_probe_ttl_seconds`` by the caller). + + Tests that never call this leave the driver unbound: the probe then + always returns ``None`` (unknown) and the gate stays OPEN -- the + load-bearing no-regression property for the existing suite. + """ + self._driver = driver + self._boot_untagged = untagged + if probe_ttl_seconds is not None: + self._probe_ttl_seconds = probe_ttl_seconds + # Rebinding invalidates any cached constraint probe result. + self._cache_populated = False + self._cached_present = None + self._cache_expires_at = 0.0 + # Seed the untagged cache with the boot-time count so the FIRST + # /status (before any live re-probe) reports the same value the old + # boot snapshot did -- then let it expire after one TTL so the live + # probe takes over and de-latches it. Seeding as populated (not cold) + # keeps the existing TTL-cache tests' hit-counts unchanged for the + # first in-window call. + self._cached_untagged = untagged + self._untagged_cache_populated = True + self._untagged_cache_expires_at = time.monotonic() + self._probe_ttl_seconds + + # -- the constraint probe -------------------------------------------- + + async def _run_probe(self) -> bool | None: + """One live catalog read. Tri-state: True / False / None (unknown). + + ``SHOW CONSTRAINTS`` is a catalog read -- no data scan, no dependence + on graph size. Unlike ``count_untagged_nodes``/``count_duplicate_nodes`` + (neo4j_store.py), this MUST be cheap enough to run on the request + path; it must never become an ``AllNodesScan``. + """ + if self._driver is None: + return None + try: + async with self._driver.session() as session: + result = await session.run(_PROBE_CYPHER) + count = 0 + async for record in result: + count = record["c"] + return count > 0 + except Exception as exc: # noqa: BLE001 -- connectivity probe, not confirmed bad state + logger.warning("maintenance_probe_failed error=%s", exc) + return None + + async def _probe_constraint_present(self) -> bool | None: + """TTL-cached, single-flight wrapper around ``_run_probe``. + + Double-checked locking: the fast path (warm cache) never touches the + lock, so N concurrent callers with a warm cache never contend on it. + Only callers that observe an expired/empty cache take the lock, and + the re-check immediately after acquiring collapses concurrent + expiry-time callers into exactly one live probe. + """ + now = time.monotonic() + if self._cache_populated and now < self._cache_expires_at: + return self._cached_present + async with self._probe_lock: + now = time.monotonic() + if self._cache_populated and now < self._cache_expires_at: + return self._cached_present + result = await self._run_probe() + self._cached_present = result + self._cache_populated = True + self._cache_expires_at = time.monotonic() + self._probe_ttl_seconds + return result + + # -- the untagged-node probe (de-latches the degraded half) ---------- + + async def _run_untagged_probe(self) -> int | None: + """One live untagged-node count. Tri-state: int / None (unknown). + + ``count_untagged_nodes`` is O(1) via Neo4j's counts store (total minus + :Node count), NOT the ``WHERE NOT n:Node`` AllNodesScan -- safe on the + cached request path. A probe failure returns None (unknown, no + ``degraded`` signal fabricated) exactly as the constraint probe does; + it is caught here so it can never poison the independent constraint + signal. + """ + if self._driver is None: + return None + try: + return await count_untagged_nodes(self._driver) + except Exception as exc: # noqa: BLE001 -- connectivity probe, not confirmed bad state + logger.warning("maintenance_untagged_probe_failed error=%s", exc) + return None + + async def _probe_untagged(self) -> int | None: + """TTL-cached, single-flight wrapper around ``_run_untagged_probe``. + + Same double-checked-locking shape as ``_probe_constraint_present`` but + with its OWN cache/lock so the two probes never share fate. Seeded at + ``bind_driver`` with the boot count, then live-refreshed each TTL -- + this is what lets an out-of-band repair clear ``degraded`` without a + restart (the latch this fixes). + """ + now = time.monotonic() + if self._untagged_cache_populated and now < self._untagged_cache_expires_at: + return self._cached_untagged + async with self._untagged_lock: + now = time.monotonic() + if self._untagged_cache_populated and now < self._untagged_cache_expires_at: + return self._cached_untagged + result = await self._run_untagged_probe() + self._cached_untagged = result + self._untagged_cache_populated = True + self._untagged_cache_expires_at = time.monotonic() + self._probe_ttl_seconds + return result + + # -- mode derivation (single source of truth for /status + the gate) -- + + def _handle_transition( + self, mode: MaintenanceMode, reason: str | None, run_id: str | None + ) -> None: + """Detect + log open<->closed transitions exactly once each. + + No ``await`` anywhere in this method: it cannot be preempted + mid-execution by another coroutine, so two concurrent callers + observing the same transition can never both log it (whichever runs + first flips ``_window_started_at``; the second then sees the + already-updated state and no-ops). + """ + is_maintenance = mode == "maintenance" + was_maintenance = self._window_started_at is not None + if is_maintenance and not was_maintenance: + self._window_started_at = _now_iso() + self._window_started_monotonic = time.monotonic() + logger.info( + "maintenance_entered", + extra={ + "reason": reason, + "run_id": run_id, + "trigger": "op" if self._op.state == "running" else "constraint", + }, + ) + elif not is_maintenance and was_maintenance: + duration = ( + time.monotonic() - self._window_started_monotonic + if self._window_started_monotonic is not None + else None + ) + logger.info( + "maintenance_completed", + extra={ + "reason": reason, + "run_id": run_id, + "duration_seconds": duration, + }, + ) + self._window_started_at = None + self._window_started_monotonic = None + + async def _derive_mode( + self, + ) -> tuple[MaintenanceMode, str | None, bool | None, int | None]: + """The ONE place mode is computed. Both ``gate_closed`` and + ``status`` call this so a transition is caught no matter which + surface is being polled.""" + op = self._op + constraint_present = await self._probe_constraint_present() + # untagged is only load-bearing for the `degraded` term, which only + # applies when the constraint IS present. Probing it only in that + # branch keeps the absent/unknown/op-running paths (and their existing + # TTL-cache hit-count tests) untouched, and avoids a needless count on + # a graph we already know is in maintenance. + untagged: int | None = None + + if op.state == "running": + mode: MaintenanceMode = "maintenance" + reason = "maintenance operation in progress" + elif constraint_present is False: + mode = "maintenance" + reason = ":Node uniqueness constraint absent -- migration required" + elif constraint_present is None: + mode = "unknown" + reason = "constraint probe could not determine graph state" + else: + # constraint present: consult the LIVE (TTL-cached) untagged count + # so an out-of-band repair de-latches degraded->healthy with no + # restart. None (probe could not answer) is NOT coerced to + # degraded -- no evidence, so healthy stands. + untagged = await self._probe_untagged() + if untagged is not None and untagged > 0: + mode = "degraded" + reason = f"{untagged} node(s) lacking the :Node label" + else: + mode = "healthy" + reason = None + + self._handle_transition(mode, reason, op.run_id) + return mode, reason, constraint_present, untagged + + # -- public seam ------------------------------------------------------ + + async def gate_closed(self) -> bool: + """True iff ingest/query must be refused right now. + + ``gate_closed() == op_running (live) OR constraint_absent (TTL-cached)``. + ``degraded`` and ``unknown`` do NOT close the gate -- only + ``mode == "maintenance"`` does. + """ + mode, _reason, _constraint_present, _untagged = await self._derive_mode() + return mode == "maintenance" + + async def status(self) -> MaintenanceStatus: + """Full snapshot for ``/status`` and ``GET /admin/maintenance``.""" + mode, reason, constraint_present, untagged = await self._derive_mode() + elapsed = ( + time.monotonic() - self._window_started_monotonic + if self._window_started_monotonic is not None + else None + ) + return MaintenanceStatus( + mode=mode, + constraint_present=constraint_present, + reason=reason, + started_at=self._window_started_at, + elapsed_seconds=elapsed, + op=self._op, + untagged_nodes=untagged, + ) + + def try_begin_op(self) -> str | None: + """Synchronous single-flight CAS -- begin an op iff none is running. + + No ``await`` between the check and the set: in asyncio this makes + the check-and-set atomic. Per-process only; + ``run_id`` is the future multi-replica fencing token. + """ + if self._op.state == "running": + return None + run_id = uuid.uuid4().hex + self._op = OpRecord( + state="running", + run_id=run_id, + started_at=_now_iso(), + completed_at=None, + records_affected=None, + error=None, + ) + return run_id + + def finish_op( + self, run_id: str, *, records_affected: int | None, error: str | None + ) -> None: + """Record the outcome of the op started by ``try_begin_op``. + + Sets ``completed_at`` -- the ONLY place it is written, which is what + makes ``completed_at`` a genuine freshness marker. A + run_id mismatch (stale/foreign completion signal) is logged and + ignored rather than corrupting the current op record. + """ + if self._op.run_id != run_id: + logger.warning( + "maintenance_finish_op_run_id_mismatch expected=%s got=%s", + self._op.run_id, + run_id, + ) + return + self._op = OpRecord( + state="failed" if error else "succeeded", + run_id=run_id, + started_at=self._op.started_at, + completed_at=_now_iso(), + records_affected=records_affected, + error=error, + ) + + def current_op(self) -> OpRecord: + return self._op + + def retain_task(self, task: asyncio.Task[Any]) -> None: + """Hold a strong reference to an in-flight maintenance-op task. + + ``asyncio.create_task`` only returns a task the event loop tracks + weakly; with no other strong reference, the task object can be + garbage-collected mid-run (a well-known asyncio footgun -- see the + "Important" note in the stdlib ``asyncio.create_task`` docs). The + caller (``routers/admin.py``) MUST call this immediately after + creating the task. The completion callback discards the reference + once the task finishes, so this set never grows unbounded. + """ + self._background_tasks.add(task) + task.add_done_callback(self._background_tasks.discard) + + +# Module singleton -- the ONE coordinator instance shared by the HTTP gate, +# the drain-loop gate, and /status. +coordinator: MaintenanceCoordinator = MaintenanceCoordinator() + + +# --------------------------------------------------------------------------- +# The structured 503 (one producer) +# --------------------------------------------------------------------------- + + +def maintenance_response(status: MaintenanceStatus, retry_after: int) -> JSONResponse: + """The ONE producer of the maintenance 503. + + Deliberately NOT ``HTTPException(503, detail=...)`` -- that renders + ``{"detail": ...}``, the wrong contract for this response. + """ + schema_health = "unknown" if status.constraint_present is None else "degraded" + return JSONResponse( + status_code=503, + content={ + "status": "maintenance", + "reason": status.reason, + "retry_after": retry_after, + "schema_health": schema_health, + "maintenance_started_at": status.started_at, + }, + headers={"Retry-After": str(retry_after)}, + ) + + +# --------------------------------------------------------------------------- +# HTTP gate middleware +# --------------------------------------------------------------------------- + + +async def maintenance_gate_middleware( + request: Request, call_next: Callable[[Request], Awaitable[Response]] +) -> Response: + """Allow-list middleware: refuse every non-allow-listed path while the + coordinator reports ``mode == "maintenance"``. + + Registered on ``app`` itself (not the auth-wrapped ASGI app) so it cannot + be bypassed by the bare ``app`` entrypoint. + Blocks (non-exhaustive; see docs/maintenance-mode.md once shipped): + ``POST /events``, ``POST /cypher``, ``GET /blobs/*``, ``GET/POST + /queues/*`` (including dead-letter replay), and all ``/admin/*`` except + ``/admin/maintenance``. + """ + if request.url.path in MAINTENANCE_ALLOW_LIST: + return await call_next(request) + st = await coordinator.status() + if st.mode == "maintenance": + retry_after = getattr( + get_settings(), "maintenance_retry_after_seconds", _RETRY_AFTER_DEFAULT + ) + return maintenance_response(st, retry_after) + return await call_next(request) diff --git a/context_intelligence_server/maintenance_ops.py b/context_intelligence_server/maintenance_ops.py new file mode 100644 index 00000000..8a942f10 --- /dev/null +++ b/context_intelligence_server/maintenance_ops.py @@ -0,0 +1,79 @@ +"""Maintenance operation execution -- the ONE shared logic home for running +the maintenance repair operation, used by both ``POST /admin/maintenance`` +and, later, the standalone out-of-band migration script +sec 5.4, not built in this change). + +This module writes NO new dedup/repair algorithm: +it wraps the existing ``neo4j_store.run_repair`` with exactly the two pieces +of bookkeeping an *in-process, gated* caller needs that a standalone script +would not: + +- a bounded pre-op quiesce sleep to let ordinary in-flight + flushes land before the schema DDL runs, and +- recording the outcome via the ``MaintenanceCoordinator`` seam + (``finish_op``), so ``GET /admin/maintenance`` and ``/status`` observe it. + +Kept standalone-friendly on purpose: ``run_maintenance_operation`` takes the +driver and coordinator explicitly (no FastAPI ``Request``, no import of +``main`` or ``routers.admin``), so a future standalone script can reuse it +directly if it ever wants coordinator-aware bookkeeping. No +bypass mechanism exists here or anywhere else -- this function reaches Neo4j +only through the driver it is given (the admin driver, ``app.state.neo4j_driver`` +at the HTTP call site) and never consults ``coordinator.gate_closed()``. +""" + +from __future__ import annotations + +import asyncio +import logging +from typing import Any + +from context_intelligence_server.maintenance import MaintenanceCoordinator, coordinator +from context_intelligence_server.neo4j_store import run_repair + +logger = logging.getLogger("context_intelligence_server") + + +async def run_maintenance_operation( + driver: Any, + run_id: str, + *, + quiesce_seconds: float, + database: str = "neo4j", + coord: MaintenanceCoordinator = coordinator, +) -> None: + """Run one maintenance operation to completion and record its outcome. + + Args: + driver: The Neo4j admin driver (``app.state.neo4j_driver`` at the + HTTP call site). Never the query-only driver. + run_id: The run id returned by ``coord.try_begin_op()`` -- the + caller MUST have already won the CAS before scheduling this. + quiesce_seconds: Seconds to sleep before calling ``run_repair`` (spec + sec 5.3). Pass 0 to skip (e.g. direct unit testing). + database: Neo4j database name, forwarded to ``run_repair``. + coord: The coordinator instance to report back to. Defaults to the + process-wide singleton; overridable for tests. + + On success, calls ``coord.finish_op(run_id, records_affected=n, error=None)`` + with ``n = duplicates_removed + nodes_tagged``. On any + exception, calls ``coord.finish_op(run_id, records_affected=None, + error=str(exc))`` -- the op is recorded ``failed``, never silently lost, + and the gate stays closed (``op_running`` only clears via ``finish_op``) + until an operator retries. + """ + if quiesce_seconds > 0: + logger.info("maintenance_quiesce run_id=%s seconds=%s", run_id, quiesce_seconds) + await asyncio.sleep(quiesce_seconds) + try: + result = await run_repair(driver, database=database) + records_affected = result["duplicates_removed"] + result["nodes_tagged"] + coord.finish_op(run_id, records_affected=records_affected, error=None) + logger.info( + "maintenance_op_succeeded run_id=%s records_affected=%s", + run_id, + records_affected, + ) + except Exception as exc: + logger.exception("maintenance_op_failed run_id=%s", run_id) + coord.finish_op(run_id, records_affected=None, error=str(exc)) diff --git a/context_intelligence_server/models.py b/context_intelligence_server/models.py index 03cf0793..6a8bc52b 100644 --- a/context_intelligence_server/models.py +++ b/context_intelligence_server/models.py @@ -14,10 +14,19 @@ class EventRequest(BaseModel): The Amplifier client must always supply workspace on every event. Events without workspace (e.g. an incorrectly configured hook) are rejected at the endpoint with HTTP 422. + + working_dir is OPTIONAL — the bundle hook emits it + as a top-level envelope field alongside workspace, but older clients/events + won't have it. Absent/empty is fine and leaves the Session node's + working_dir property null. Populate-if-missing: the Session node's + working_dir is filled in by the first subsequent event (including a + re-import via the upload CLI) that carries a non-empty value, but an + already-populated value is never overwritten. """ event: str workspace: str + working_dir: str | None = None idempotency_key: str | None = None data: dict[str, Any] @@ -29,6 +38,20 @@ def workspace_must_not_be_empty(cls, v: str) -> str: raise ValueError("workspace must not be empty") return v + @field_validator("working_dir") + @classmethod + def working_dir_must_not_be_blank(cls, v: str | None) -> str | None: + """Allow ``None`` (working_dir is optional, unlike workspace) but reject + blank/whitespace-only strings. + + Mirrors ``workspace_must_not_be_empty``'s normalize-or-reject stance: a + whitespace-only value (e.g. ``" "``) is never a legitimate path and must + not write through to the Session node verbatim. + """ + if v is not None and not v.strip(): + raise ValueError("working_dir must not be blank") + return v + class EventResponse(BaseModel): """Response returned after an event is accepted.""" diff --git a/context_intelligence_server/neo4j_store.py b/context_intelligence_server/neo4j_store.py index 52ddc7d9..d16d6941 100644 --- a/context_intelligence_server/neo4j_store.py +++ b/context_intelligence_server/neo4j_store.py @@ -15,15 +15,40 @@ import logging import re from collections.abc import Generator -from datetime import datetime +from datetime import UTC, datetime from typing import Any, LiteralString, cast from neo4j import AsyncGraphDatabase from neo4j import unit_of_work as _unit_of_work from neo4j.exceptions import DriverError, Neo4jError +from context_intelligence_server.config import Neo4jClientConfig +from context_intelligence_server.status import SCHEMA_VERSION + _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 # --------------------------------------------------------------------------- @@ -903,6 +928,115 @@ async def _create_constraint( return fully_established +async def ensure_schema_version_baseline( + driver: Any, + *, + database: str = "neo4j", +) -> None: + """Create the :SchemaMeta uniqueness constraint and baseline singleton. + + Baseline only -- create-if-absent, no comparison/migration. Call this + exactly once, from the lifespan startup handler, AFTER ``ensure_neo4j_schema`` + has established the rest of the schema (indexes + uniqueness constraints). + + Kept out of ``ensure_neo4j_schema`` on purpose: that runs on every + ``Neo4jGraphStore``'s first flush (once per worker, concurrently, on cold + start) and from ``run_repair``/``doctor --fix``. A SchemaMeta baseline write + is a single-writer, startup-only concern that must not fire per worker. + + Ordering matters: the ``(:SchemaMeta).id`` uniqueness constraint is created + FIRST, then the singleton MERGE -- without the constraint, two concurrent + MERGEs on a fresh database can each create a ``{id: 'singleton'}`` node. + + ``ON CREATE SET`` only: an existing node is left untouched. Reconciling a + stored ``schema_version`` against the running server's value is deferred; + the read path (``read_graph_schema_version``) and this write path stay + structurally separate so comparison/upgrade logic cannot creep in here. + + O(1): one constraint DDL plus a MERGE on a fixed key -- never a scan. Any + ``Neo4jError``/``DriverError`` is logged and swallowed: a transient failure + on this passive data point must never crash boot. + """ + try: + async with driver.session(database=database) as session: + try: + await session.run( + "CREATE CONSTRAINT schemameta_id_unique IF NOT EXISTS " + "FOR (m:SchemaMeta) REQUIRE m.id IS UNIQUE" + ) + except (Neo4jError, DriverError) as exc: + if isinstance(exc, Neo4jError) and exc.code in _BENIGN_SCHEMA_CODES: + _LOG.debug( + "ensure_schema_version_baseline: SchemaMeta " + "uniqueness constraint already present (benign " + "concurrent-schema race, code=%s)", + exc.code, + ) + else: + _LOG.warning( + "ensure_schema_version_baseline: could not create " + "SchemaMeta uniqueness constraint; continuing " + "without it: %s", + exc, + ) + + await session.run( + "MERGE (m:SchemaMeta {id: 'singleton'}) " + "ON CREATE SET m.schema_version = $schema_version, " + "m.last_updated = $now", + schema_version=SCHEMA_VERSION, + now=datetime.now(UTC).isoformat(), + ) + except (Neo4jError, DriverError) as exc: + _LOG.warning( + "ensure_schema_version_baseline: could not write SchemaMeta " + "baseline singleton (connectivity error); continuing without " + "it: %s", + exc, + ) + + +async def read_graph_schema_version( + driver: Any, + *, + database: str = "neo4j", +) -> int | None: + """Read-only: the STORED ``:SchemaMeta{id:'singleton'}.schema_version``. + + Advisory drift signal only, not a guard. Read-only companion to + ``ensure_schema_version_baseline``, kept structurally apart from that write + path. It does NOT import or compare against ``SCHEMA_VERSION`` -- it only + reads back whatever is stored. ``GET /status`` compares the returned value + against ``status.SCHEMA_VERSION`` itself so a server/graph mismatch is + detectable; this function performs no comparison, gating, or migration. + + Returns ``None`` when the singleton is absent (startup never ran the + baseline against this graph) or when the read fails for any reason (treated + as "unknown", never an error), mirroring the never-500-``/status`` contract. + + O(1): a point lookup by the unique ``id`` key. Reads via ``async for`` + (not ``.single()``) so it works against both the real async driver and the + test suite's mock session, which only implements async iteration. + """ + try: + async with driver.session(database=database) as session: + result = await session.run( + "MATCH (m:SchemaMeta {id: 'singleton'}) " + "RETURN m.schema_version AS schema_version" + ) + async for record in result: + value = record["schema_version"] + return int(value) if value is not None else None + return None + except Exception as exc: # noqa: BLE001 - defensive: /status must never 500 + _LOG.warning( + "read_graph_schema_version: could not read SchemaMeta singleton " + "(connectivity error?); returning None: %s", + exc, + ) + return None + + def _serialized_row_size(value: Any) -> int: """Return a cheap conservative proxy for the serialized byte size of *value*. @@ -985,12 +1119,17 @@ def _chunk_list( def _build_node_props(data: dict[str, Any], workspace: str) -> dict[str, Any]: """Assemble the sanitized props dict for a single node row. - ``labels`` and ``created_by`` are excluded from the returned dict: + ``labels``, ``created_by``, and ``working_dir`` are excluded from the returned dict: - ``labels`` are applied separately via ``SET n:Label`` statements (not stored as props). - ``created_by`` travels ONLY as the ``$created_by`` query param (never a node property) so it cannot affect node identity or clobber the ``ON CREATE SET n.created_by`` stamp. + - ``working_dir`` is excluded from the blind ``SET n += row.props`` so the + Session-node write can apply ``coalesce(n.working_dir, row.working_dir)`` + instead of last-write-wins -- an already-set working_dir is never clobbered. """ - raw = {k: v for k, v in data.items() if k not in ("labels", "created_by")} + raw = { + k: v for k, v in data.items() if k not in ("labels", "created_by", "working_dir") + } _convert_temporal_props(raw) # ISO str -> datetime, in place props = Neo4jGraphStore._sanitize_properties(raw) props["workspace"] = workspace @@ -1032,6 +1171,13 @@ async def _write_batch( row: dict[str, Any] = {"node_id": node_id, "props": props} if "Session" in labels: + # working_dir is a Session-only property carried as a separate + # top-level row key (not in props) so the MERGE below can coalesce + # it rather than blindly overwrite. Omitted when empty/absent, so + # coalesce(n.working_dir, null) is a no-op for such rows. + working_dir_value = data.get("working_dir") + if working_dir_value: + row["working_dir"] = working_dir_value session_rows.append(row) else: other_rows.append(row) @@ -1061,7 +1207,12 @@ async def _write_batch( f"MERGE (n:{_UNIVERSAL_NODE_LABEL} " "{node_id: row.node_id, workspace: row.props.workspace}) " "ON CREATE SET n.created_by = $created_by " - "SET n += row.props, n:Session", + "SET n += row.props, n:Session " + # Populate-if-missing: fill working_dir only when it is still null, + # never clobber an already-set value (a concurrent/replica writer or + # an earlier event may have populated it). Rows without a working_dir + # carry a null row key, so this is a no-op for them. + "SET n.working_dir = coalesce(n.working_dir, row.working_dir)", rows=session_rows, created_by=created_by, ) @@ -1179,12 +1330,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 +1349,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 +1389,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 +1811,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 +1829,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/pipeline.py b/context_intelligence_server/pipeline.py index 7310fd9c..d114cb13 100644 --- a/context_intelligence_server/pipeline.py +++ b/context_intelligence_server/pipeline.py @@ -173,7 +173,16 @@ async def process_event( data.get("timestamp") if isinstance(data, dict) else None ) if session_id and timestamp and worker.services.blob_store: - node_id = make_node_id(session_id, event, timestamp) + # The blob-key node_id MUST match handlers/data_layer_1/default.py's + # event_node_id (same session_id + event + timestamp + tool_call_id), + # otherwise two distinct same-millisecond events (e.g. parallel tool + # calls in the same batch) collide on an identical blob key and the + # second write silently overwrites the first, while the first Event + # node's $blob_ref still points at that URI (now holding the wrong + # payload). tool_call_id is present on ALL event types that carry it, + # not just tool:*, so it is safe to read unconditionally here. + disambiguator = data.get("tool_call_id") if isinstance(data, dict) else None + node_id = make_node_id(session_id, event, timestamp, disambiguator) await process_event_data( data, worker.services.blob_store, session_id, node_id ) diff --git a/context_intelligence_server/queue_manager.py b/context_intelligence_server/queue_manager.py deleted file mode 100644 index 6fba186e..00000000 --- a/context_intelligence_server/queue_manager.py +++ /dev/null @@ -1,731 +0,0 @@ -"""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. -""" - -from __future__ import annotations - -import asyncio -import base64 -import json -import os -import time -from dataclasses import dataclass -from pathlib import Path -from typing import Any - -# Fixed buffer size for streaming scans over a session ``.log`` (last-newline -# search and newline counting). Bounds boot-time and /status memory to O(chunk) -# instead of O(file): a durable log can be multi-GB (4.9 GB in the incident), -# and loading one into RAM just to count newlines is what drove ~44 GB RSS at -# startup. 1 MiB balances syscall count against per-scan memory. -_SCAN_CHUNK_BYTES = 1 << 20 - - -@dataclass(frozen=True) -class Batch: - """A contiguous batch of log lines 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. - 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 - are available, ``end_offset == start_offset``. - """ - - session_id: str - lines: list[bytes] - start_offset: int - end_offset: int - - -class QueueManager: - """Manages per-session append-only queues on disk.""" - - def __init__(self, queues_dir: Path): - self._dir = Path(queues_dir) - self._dir.mkdir(parents=True, exist_ok=True) - self._stats_cache: dict[str, Any] | None = None - self._stats_cache_at: float = 0.0 - self._stats_cache_ttl: float = 1.0 - # Separate cache for spool_stats() (Change 2 / /status spool block). - # A longer TTL than _stats_cache_ttl is fine here: spool_stats() is an - # operator-facing "is the backlog growing" signal, not a - # correctness-sensitive value, so a few extra seconds of staleness is - # an acceptable trade for fewer directory scans under frequent - # /status polling. - self._spool_cache: dict[str, int] | None = None - self._spool_cache_at: float = 0.0 - self._spool_cache_ttl: float = 5.0 - - def _log_path(self, session_id: str) -> Path: - return self._dir / f"{session_id}.log" - - def _offset_path(self, session_id: str) -> Path: - return self._dir / f"{session_id}.offset" - - def _dead_path(self, session_id: str) -> Path: - return self._dir / f"{session_id}.dead.jsonl" - - def _read_committed_offset(self, session_id: str) -> int: - 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 - - def _complete_data_end(self, session_id: str) -> int: - """Byte position after the last complete (newline-terminated) line. - - 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. - - 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. - """ - path = self._log_path(session_id) - try: - with open(path, "rb") as f: - f.seek(0, os.SEEK_END) - pos = f.tell() - while pos > 0: - read_size = min(_SCAN_CHUNK_BYTES, pos) - pos -= read_size - f.seek(pos) - buf = f.read(read_size) - idx = buf.rfind(b"\n") - if idx != -1: - return pos + idx + 1 - return 0 - except FileNotFoundError: - return 0 - - @staticmethod - def _stream_newlines(path: Path, start: int = 0, end: int | None = None) -> int: - """Count ``\\n`` bytes in ``path``'s byte range ``[start, end)`` -- streamed. - - ``end=None`` counts to EOF. Reads the range in fixed-size chunks - (O(chunk) memory) instead of materialising the whole file (or a slice - copy of it) in RAM, which is what ``read_bytes()`` + - ``data[a:b].count(b"\\n")`` did on multi-GB spool files. Numerically - identical to that slice-count for any range; a missing file counts 0. - - Path-based (not session-id-based) so it serves both ``.log`` scans - (``_count_newlines``) and the whole-file ``.dead.jsonl`` count - (``_count_dead``). - """ - if end is not None and end <= start: - return 0 - try: - with open(path, "rb") as f: - f.seek(start) - remaining = None if end is None else end - start - count = 0 - while True: - to_read = ( - _SCAN_CHUNK_BYTES - if remaining is None - else min(_SCAN_CHUNK_BYTES, remaining) - ) - if to_read <= 0: - break - buf = f.read(to_read) - if not buf: - break - count += buf.count(b"\n") - if remaining is not None: - remaining -= len(buf) - return count - except FileNotFoundError: - return 0 - - def _count_newlines( - self, session_id: str, start: int = 0, end: int | None = None - ) -> int: - """Streamed newline count over a session ``.log``'s ``[start, end)``.""" - return self._stream_newlines(self._log_path(session_id), start, end) - - @staticmethod - def _validate_session_id(session_id: str) -> None: - if ( - not session_id - or "/" in session_id - or "\\" in session_id - or "\0" in session_id - ): - raise ValueError(f"Invalid session_id: {session_id!r}") - - async def append(self, session_id: str, raw: bytes) -> None: - self._validate_session_id(session_id) - line = raw if raw.endswith(b"\n") else raw + b"\n" - path = self._log_path(session_id) - - def _append() -> None: - with open(path, "ab") as f: - f.write(line) - - await asyncio.to_thread(_append) - - async def read_batch(self, session_id: str, max_items: int) -> Batch: - self._validate_session_id(session_id) - path = self._log_path(session_id) - - def _read() -> Batch: - start = self._read_committed_offset(session_id) - lines: list[bytes] = [] - consumed = 0 - try: - with open(path, "rb") as f: - f.seek(start) - while len(lines) < 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]) - consumed += len(raw) - except FileNotFoundError: - pass - return Batch(session_id, lines, start, start + consumed) - - return await asyncio.to_thread(_read) - - async def commit(self, session_id: str, new_offset: int) -> None: - """Atomically and durably persist ``new_offset`` (the ack). - - 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). - """ - self._validate_session_id(session_id) - final = self._offset_path(session_id) - tmp = self._dir / f"{session_id}.offset.tmp" - - def _commit() -> None: - tmp.write_text(str(new_offset), encoding="utf-8") - os.replace(tmp, final) - - await asyncio.to_thread(_commit) - - 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. - """ - self._validate_session_id(session_id) - payload = raw[:-1] if raw.endswith(b"\n") else raw - record: dict = {"ts": time.time(), "error": error} - try: - record["payload"] = payload.decode("utf-8") - except UnicodeDecodeError: - 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. - """ - self._validate_session_id(session_id) - - def _delete() -> None: - for p in (self._log_path(session_id), self._offset_path(session_id)): - try: - p.unlink() - except FileNotFoundError: - pass - - await asyncio.to_thread(_delete) - - 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. - """ - self._validate_session_id(session_id) - - def _read() -> list[dict]: - try: - text = self._dead_path(session_id).read_text(encoding="utf-8") - except FileNotFoundError: - return [] - return [json.loads(ln) for ln in text.splitlines() if ln.strip()] - - return await asyncio.to_thread(_read) - - async def active_sessions(self) -> list[str]: - """Return sorted session_ids with undrained data. - - A session is "active" when its committed offset is strictly less than - the byte length of its ``.log`` file (i.e. there are appended bytes - that have not yet been committed). Fully-committed sessions are - excluded. The result is sorted by session_id. - """ - - 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) - return result - - return await asyncio.to_thread(_scan) - - async def recover(self) -> list[str]: - """Return sorted session_ids that have a complete unprocessed line. - - A session is recoverable when its committed offset is strictly less - than the end of its complete (newline-terminated) data, i.e. at least - one whole line remains to be processed. A torn trailing line (bytes - after the final newline) is ignored, so a session whose only remaining - data is a partial line is NOT reported. - - This method is idempotent, safe on an empty directory, and performs no - drainer logic; respawning drainers for the reported sessions is Phase - B2. - """ - - 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) - return result - - return await asyncio.to_thread(_scan) - - def _count_dead(self, worker_key: str) -> int: - """Count complete (newline-terminated) dead-letter lines for a key. - - Returns 0 when no dead-letter file exists. Dead-letter records are - always written newline-terminated, so counting newlines yields the - number of complete records. - - Streamed (bounded memory), not ``read_bytes()``: a .dead.jsonl is - usually small but is NOT bounded -- a systematically-failing session - dead-letters every line -- and this is called on the same boot and - polled-/status paths as the .log scans. - """ - return self._stream_newlines(self._dead_path(worker_key)) - - def _all_worker_keys(self) -> list[str]: - """Return the sorted union of ``.log`` and ``.dead.jsonl`` stems. - - ``Path.stem`` only strips the final suffix, so for ``s1.dead.jsonl`` it - returns ``s1.dead``; the ``.dead.jsonl`` suffix is sliced explicitly to - recover the bare worker key. - """ - keys: set[str] = set() - for log in self._dir.glob("*.log"): - keys.add(log.stem) - for dead in self._dir.glob("*.dead.jsonl"): - keys.add(dead.name[: -len(".dead.jsonl")]) - return sorted(keys) - - 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. - """ - now = time.monotonic() - if ( - self._stats_cache is not None - and (now - self._stats_cache_at) < self._stats_cache_ttl - ): - return self._stats_cache - - def _all() -> dict[str, Any]: - per_key: list[dict[str, Any]] = [] - in_queue_total = 0 - dead_total = 0 - for worker_key in self._all_worker_keys(): - try: - committed = self._read_committed_offset(worker_key) - except (OSError, ValueError): - # /status calls this (via pipeline_metrics); a corrupt or - # transiently-unreadable .offset must NOT 500 the health - # probe. Degrade to 0 for this key's stats -- mirroring the - # existing missing-file->0 convention in - # _read_committed_offset, and tending the conservation - # residual negative (benign, never a false `degraded`). - # Deliberately NO logging here: /status is polled, and a - # per-scan warning on a persistently-corrupt offset would - # flood the hot path. The visibility signal is the aggregate - # `spool.corrupt_offsets` field (see spool_stats()). - committed = 0 - # Streamed count of complete lines from committed -> EOF. - # Equivalent to the old f.read() + data[:last_nl+1].count(b"\n") - # (every b"\n" lies at or before the last one), but without - # materialising the undrained tail -- which can be gigabytes - # under a large backlog on this (polled) /status path. - in_queue = self._count_newlines(worker_key, committed) - dead = self._count_dead(worker_key) - per_key.append( - {"worker_key": worker_key, "in_queue": in_queue, "dead": dead} - ) - in_queue_total += in_queue - dead_total += dead - return { - "per_key": per_key, - "in_queue_total": in_queue_total, - "dead_total": dead_total, - } - - stats = await asyncio.to_thread(_all) - self._stats_cache = stats - self._stats_cache_at = now - return stats - - 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. - """ - now = time.monotonic() - if ( - self._spool_cache is not None - and (now - self._spool_cache_at) < self._spool_cache_ttl - ): - return self._spool_cache - - def _scan() -> dict[str, int]: - spool_bytes_total = 0 - pending_sessions = 0 - corrupt_offsets = 0 - for entry in self._dir.iterdir(): - if not entry.is_file(): - continue - try: - size = entry.stat().st_size - except FileNotFoundError: - # Raced with a concurrent delete_drained()/purge; the - # entry no longer exists -- simply exclude it, don't fail - # a cheap, best-effort aggregate over a live directory. - continue - spool_bytes_total += size - if entry.suffix == ".log": - try: - committed = self._read_committed_offset(entry.stem) - except ValueError: - # The .offset exists but is not a valid integer -- a - # GENUINELY corrupt offset. This is the one visibility - # signal for it (no logging anywhere, to avoid flooding - # the polled health path): surface it as an aggregate - # count on /status so `spool.corrupt_offsets > 0` is the - # operator's alarm. Count this file's bytes; skip its - # pending calc. - corrupt_offsets += 1 - continue - except OSError: - # A transient/racing FS error reading the offset (NOT - # corruption): count bytes, skip pending calc, and do - # NOT inflate corrupt_offsets with a non-corruption cause. - continue - if committed < size: - pending_sessions += 1 - return { - "pending_sessions": pending_sessions, - "spool_bytes_total": spool_bytes_total, - "corrupt_offsets": corrupt_offsets, - } - - try: - stats = await asyncio.to_thread(_scan) - except (OSError, ValueError): - # Queue dir missing/unavailable (e.g. Azure Files SMB remount) or a - # transient FS error mid-scan. /status is the health probe and MUST - # return 200 -- degrade to a sentinel and DO NOT cache it, so the - # next poll retries immediately once the filesystem recovers. All - # three fields are -1 = "temporarily unavailable" (distinct from a - # real 0, and from a real corrupt_offsets count). - return { - "pending_sessions": -1, - "spool_bytes_total": -1, - "corrupt_offsets": -1, - } - - self._spool_cache = stats - self._spool_cache_at = now - return stats - - async def dead_letter_keys(self) -> list[str]: - """Return sorted worker keys that have a ``.dead.jsonl`` file. - - Keys with only main-log data (no dead-letter file) are excluded. - ``Path.name`` is sliced by the ``.dead.jsonl`` suffix to recover the - bare worker key (``Path.stem`` would only strip ``.jsonl``). - """ - - def _scan() -> list[str]: - return sorted( - dead.name[: -len(".dead.jsonl")] - for dead in self._dir.glob("*.dead.jsonl") - ) - - return await asyncio.to_thread(_scan) - - async def purge_dead_letters(self, worker_key: str) -> int: - """Delete the dead-letter file for ``worker_key`` and return the count. - - Counts the dead-letter records via ``_count_dead``, then unlinks the - ``.dead.jsonl`` file. Returns the number of records removed (0 when no - dead-letter file exists). Raises ``ValueError`` for an unsafe key. - - Deletion is routed exclusively through this method: callers must never - touch the filesystem directly. - """ - self._validate_session_id(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: - - - ``C`` = complete lines below the committed offset - - ``P`` = complete lines between the committed offset and the end of - complete data (== ``in_queue``) - - ``D`` = dead-letter records - - Formula:: - - written_seed = max(0, C - D) - accepted_seed = written_seed + P + D - - The ``max(0, ...)`` clamp is load-bearing. In a crash/replay window a - dead-but-pending line (dead-lettered, but whose commit has not yet - advanced past it) makes ``C - D`` go negative. The naive formula - ``accepted = C + P`` / ``written = C - D`` yields a negative written - count -- residual ``-1``, a false DEGRADED. Clamping written to zero - and counting the line in BOTH ``P`` and ``D`` absorbs it into - ``accepted_seed`` so the residual stays exactly zero. - - Ordering is load-bearing: this MUST run AFTER ``recovery_reconcile_dead`` - in the lifespan so the dead-letter counts it reads are already settled. - """ - - 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) - written_seed = max(0, before - dead) - accepted += written_seed + pending + dead - written += written_seed - return accepted, written - - return await asyncio.to_thread(_seed) - - def _dead_payload_set(self, worker_key: str) -> set[bytes]: - """Return the set of original raw line bytes recorded as dead-letters. - - Each dead-letter record stores the original line (sans trailing - newline) either as ``payload`` (a UTF-8 string) or ``payload_b64`` - (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. - """ - try: - text = self._dead_path(worker_key).read_text(encoding="utf-8") - except FileNotFoundError: - return set() - payloads: set[bytes] = set() - 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"])) - 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. - """ - - 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 - 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) - self._stats_cache = None - return total_skipped - - return await asyncio.to_thread(_reconcile) diff --git a/context_intelligence_server/queue_manager/__init__.py b/context_intelligence_server/queue_manager/__init__.py new file mode 100644 index 00000000..3c0ce8a2 --- /dev/null +++ b/context_intelligence_server/queue_manager/__init__.py @@ -0,0 +1,31 @@ +"""queue_manager — durable, per-session append-only queue for the event-write pipeline. + +The public surface is the backend-neutral :class:`QueueManager` Protocol plus +the :class:`Batch` / :class:`Record` value types and the +:func:`create_queue_manager` factory. Consumers depend on these, never on a +concrete backend class. + +Package layout: + protocol.py QueueManager Protocol, Batch, Record — the backend-neutral + seam (no filesystem imports). + filesystem.py FileSystemQueueManager — the on-disk implementation, plus + the on-disk-only Verdict classification enum. + factory.py create_queue_manager(settings) — the ONLY place a backend is + selected and the ONLY place (besides config.py) that reads + settings.queues_path. +""" + +from __future__ import annotations + +from .factory import create_queue_manager +from .filesystem import FileSystemQueueManager, Verdict +from .protocol import Batch, QueueManager, Record + +__all__ = [ + "Batch", + "FileSystemQueueManager", + "QueueManager", + "Record", + "Verdict", + "create_queue_manager", +] diff --git a/context_intelligence_server/queue_manager/factory.py b/context_intelligence_server/queue_manager/factory.py new file mode 100644 index 00000000..fadf960e --- /dev/null +++ b/context_intelligence_server/queue_manager/factory.py @@ -0,0 +1,35 @@ +"""Config-driven QueueManager factory — the ONLY place a queue backend is selected. + +This is the single seam through which the concrete backend is chosen. Adding +a new backend (e.g. Azure) means: one new module implementing +:class:`~.protocol.QueueManager`, one new branch here, and a config value — +zero changes to :mod:`context_intelligence_server.registry` or any consumer. + +This module (and :mod:`~context_intelligence_server.config`) are the only +places ``settings.queues_path`` is read — the on-disk root is a filesystem- +backend concern, resolved here and handed to the concrete backend at +construction time. Callers only ever see the :class:`~.protocol.QueueManager` +Protocol. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import TYPE_CHECKING + +from .filesystem import FileSystemQueueManager +from .protocol import QueueManager + +if TYPE_CHECKING: + from context_intelligence_server.config import Settings + + +def create_queue_manager(settings: Settings) -> QueueManager: + """Build the durable ``QueueManager`` from config. + + Single backend today (on-disk), so this is a thin config-reading seam + rather than a multi-backend dispatcher — but it keeps ``settings.queues_path`` + out of consumers (mirrors ``blob_store.factory.create_blob_store`` and + ``identity_store.factory.create_identity_store``). + """ + return FileSystemQueueManager(queues_dir=Path(settings.queues_path)) diff --git a/context_intelligence_server/queue_manager/filesystem.py b/context_intelligence_server/queue_manager/filesystem.py new file mode 100644 index 00000000..b044ead9 --- /dev/null +++ b/context_intelligence_server/queue_manager/filesystem.py @@ -0,0 +1,1950 @@ +"""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, TypeVar + +from context_intelligence_server.config import get_settings +from context_intelligence_server.queue_manager.protocol import Batch, Record + +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) +# instead of O(file): a durable log can be multi-GB (4.9 GB in the incident), +# and loading one into RAM just to count newlines is what drove ~44 GB RSS at +# startup. 1 MiB balances syscall count against per-scan memory. +_SCAN_CHUNK_BYTES = 1 << 20 + + +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 FileSystemQueueManager: + """Manages per-session append-only queues on disk. + + Implements :class:`~context_intelligence_server.queue_manager.protocol.QueueManager`. + """ + + def __init__(self, queues_dir: Path): + self._dir = Path(queues_dir) + self._dir.mkdir(parents=True, exist_ok=True) + self._stats_cache: dict[str, Any] | None = None + self._stats_cache_at: float = 0.0 + self._stats_cache_ttl: float = 1.0 + # Separate cache for spool_stats() (Change 2 / /status spool block). + # A longer TTL than _stats_cache_ttl is fine here: spool_stats() is an + # operator-facing "is the backlog growing" signal, not a + # correctness-sensitive value, so a few extra seconds of staleness is + # an acceptable trade for fewer directory scans under frequent + # /status polling. + 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" + + def _offset_path(self, session_id: str) -> Path: + return self._dir / f"{session_id}.offset" + + 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 _write_offset_record( + self, session_id: str, offset: int, cursor: dict[str, Any] | None + ) -> None: + """Sole writer of the ``.offset`` file: one atomic JSON record. + + Writes ``{"v": 1, "offset": offset, "cursor": cursor}`` to a temp file + and ``os.replace``s it in, so a reader never sees a torn record. Folding + the cursor into the SAME record as the offset (not a sidecar file) is + what keeps the two from ever skewing: a crash loses both together, never + one without the other. No ``fsync`` -- process-crash-durable, not + power-durable, matching the rest of this module. + """ + if cursor is not None and not isinstance(cursor, dict): + raise TypeError( + f"cursor must be a dict or None, got {type(cursor).__name__}" + ) + final = self._offset_path(session_id) + tmp = self._dir / f"{session_id}.offset.tmp" + record = {"v": 1, "offset": offset, "cursor": cursor} + tmp.write_text(json.dumps(record, separators=(",", ":")), encoding="utf-8") + try: + os.replace(tmp, final) + except OSError: + # A failed rename leaves the prior committed record intact; drop the + # staged temp so a crashed write leaves nothing behind. + tmp.unlink(missing_ok=True) + raise + + def _read_offset_record(self, session_id: str) -> tuple[int, dict[str, Any] | None]: + """Read the ``.offset`` file, returning ``(offset, cursor)``. + + Accepts the current JSON record and the legacy bare-integer shape; a + missing or empty file yields ``(0, None)``. An unreadable cursor + degrades to ``None`` (an unknown ``v`` or non-dict ``cursor`` keeps the + offset but drops the cursor -- a corrupt cursor must not crash boot). An + unreadable OFFSET is NOT degraded: a malformed record raises, because + silently resetting a corrupt offset to 0 would replay the whole log and + manufacture duplicate nodes -- a worse, quieter failure than a loud one. + """ + try: + text = self._offset_path(session_id).read_text("utf-8") + except FileNotFoundError: + return 0, None + text = text.strip() + if not text: + return 0, None + if text[0] == "{": + try: + rec = json.loads(text) + offset = int(rec["offset"]) + except (json.JSONDecodeError, KeyError, TypeError, ValueError): + raise ValueError( + f"unparseable legacy offset document for session {session_id!r}" + ) from None + cursor = rec.get("cursor") + if rec.get("v") == 1 and isinstance(cursor, dict): + return offset, cursor + return offset, None + return int(text), None + + def _read_committed_offset(self, session_id: str) -> int: + """Committed byte offset; reads envelope, bare-int, and legacy JSON.""" + return self._read_offset_record(session_id)[0] + + @staticmethod + def _last_complete_end(path: Path) -> int: + """Byte position after the last complete line in ``path`` (0 if none). + + 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). 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``). + """ + try: + with open(path, "rb") as f: + f.seek(0, os.SEEK_END) + pos = f.tell() + while pos > 0: + read_size = min(_SCAN_CHUNK_BYTES, pos) + pos -= read_size + f.seek(pos) + buf = f.read(read_size) + idx = buf.rfind(b"\n") + if idx != -1: + return pos + idx + 1 + return 0 + 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. + + ``end=None`` counts to EOF. Reads the range in fixed-size chunks + (O(chunk) memory) instead of materialising the whole file (or a slice + copy of it) in RAM, which is what ``read_bytes()`` + + ``data[a:b].count(b"\\n")`` did on multi-GB spool files. Numerically + identical to that slice-count for any range; a missing file counts 0. + + Path-based (not session-id-based) so it serves both ``.log`` scans + (``_count_newlines``) and the whole-file ``.dead.jsonl`` count + (``_count_dead``). + """ + if end is not None and end <= start: + return 0 + try: + with open(path, "rb") as f: + f.seek(start) + remaining = None if end is None else end - start + count = 0 + while True: + to_read = ( + _SCAN_CHUNK_BYTES + if remaining is None + else min(_SCAN_CHUNK_BYTES, remaining) + ) + if to_read <= 0: + break + buf = f.read(to_read) + if not buf: + break + count += buf.count(b"\n") + if remaining is not None: + remaining -= len(buf) + return count + except FileNotFoundError: + return 0 + + def _count_newlines( + self, session_id: str, start: int = 0, end: int | None = None + ) -> int: + """Streamed newline count over a session ``.log``'s ``[start, end)``.""" + return self._stream_newlines(self._log_path(session_id), start, end) + + @staticmethod + def _validate_session_id(session_id: str) -> None: + if ( + not session_id + or "/" in session_id + or "\\" in session_id + or "\0" in session_id + ): + 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: + FileSystemQueueManager._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 = FileSystemQueueManager._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: + FileSystemQueueManager._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" # unchanged (was :186) + path = self._log_path(session_id) + # ``_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) + path = self._log_path(session_id) + + def _read() -> Batch: + start = self._read_committed_offset(session_id) + records: list[Record] = [] + consumed = 0 + try: + with open(path, "rb") as f: + f.seek(start) + while len(records) < max_items: + raw = f.readline() + if not raw or not raw.endswith(b"\n"): + # EOF, or a torn trailing line with no newline yet: + # ignore the partial line and stop on a line boundary. + break + rec_start = start + consumed + consumed += len(raw) + rec_end = start + consumed + records.append(Record(raw[:-1], rec_start, rec_end)) + except FileNotFoundError: + pass + return Batch(session_id, records, start, start + consumed) + + return await asyncio.to_thread(_read) + + async def commit( + self, session_id: str, new_offset: int, cursor: dict[str, Any] | None + ) -> None: + """Atomically and durably persist ``new_offset`` (the ack) + ``cursor``. + + ``cursor`` has NO default: every call site must pass it explicitly so it + is never silently omitted -- a missed site (or a plain rolling deploy + against an older signature) would null the cursor and reset the + cross-handler counters, manufacturing duplicate nodes on the next run. + Offset and cursor are written in the SAME atomic record, so they can + never skew. No ``fsync``: process-crash-durable, not power-durable. + """ + self._validate_session_id(session_id) + await asyncio.to_thread( + self._write_offset_record, session_id, new_offset, cursor + ) + + async def read_cursor(self, session_id: str) -> dict[str, Any] | None: + """Return the persisted cursor for ``session_id``, or ``None``. + + ``None`` covers: no offset file, a legacy bare-integer offset file, or a + JSON record whose cursor is absent/unreadable. This is the read side a + rebuilt worker uses to restore its cross-handler counters. + """ + self._validate_session_id(session_id) + return await asyncio.to_thread(lambda: self._read_offset_record(session_id)[1]) + + async def dead_letter(self, session_id: str, raw: bytes, error: str) -> None: + """Append one dead-letter record for an unprocessable batch line. + + 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 + record: dict = {"ts": time.time(), "error": error} + try: + record["payload"] = payload.decode("utf-8") + except UnicodeDecodeError: + record["payload_b64"] = base64.b64encode(payload).decode("ascii") + line = (json.dumps(record) + "\n").encode("utf-8") + path = self._dead_path(session_id) + 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 + + try: + 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 + + # Step 6: replace the log with the verified tail copy. + try: + os.replace(tmp, log) + except OSError: + # Restore the offset to the committed value so a failed + # compaction is a pure no-op, never an in-process re-drive + # (which would double-count `written` and drive the + # residual negative). + 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. 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) + + def _read() -> list[dict]: + try: + text = self._dead_path(session_id).read_text(encoding="utf-8") + except FileNotFoundError: + return [] + records: list[dict] = [] + skipped = 0 + for ln in text.splitlines(): + if not ln.strip(): + continue + 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. + + A session is "active" when its committed offset is strictly less than + the byte length of its ``.log`` file (i.e. there are appended bytes + that have not yet been committed). Fully-committed sessions are + excluded. The result is sorted by session_id. + """ + + def _scan() -> list[str]: + result: list[str] = [] + for log in sorted(self._dir.glob("*.log")): + session_id = log.stem + # 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) + + async def session_keys(self) -> list[str]: + """Return every persisted session key (sorted). + + A session key is present whenever the backend holds a queue log for it, + regardless of drain state or live-worker status. Callers that need to + sweep the whole queue (e.g. boot reclaim) enumerate here rather than + reaching into the backend's on-disk layout, so the sweep works + unchanged against any backend. + """ + + def _scan() -> list[str]: + return sorted(log.stem for log in self._dir.glob("*.log")) + + return await asyncio.to_thread(_scan) + + async def is_fully_drained(self, session_id: str) -> bool: + """True iff the session has no undrained log data left. + + Compares the committed offset against the end of complete (newline- + terminated) data, read straight from disk -- so it is independent of + in-memory worker liveness and stays correct across a crash + restart. A + session with no ``.log`` reads as drained (0 >= 0): it never had data. + """ + self._validate_session_id(session_id) + + def _check() -> bool: + return self._read_committed_offset(session_id) >= self._complete_data_end( + session_id + ) + + return await asyncio.to_thread(_check) + + async def recover(self) -> list[str]: + """Return sorted session_ids that have a complete unprocessed line. + + A session is recoverable when its committed offset is strictly less + than the end of its complete (newline-terminated) data, i.e. at least + one whole line remains to be processed. A torn trailing line (bytes + after the final newline) is ignored, so a session whose only remaining + data is a partial line is NOT reported. + + This method is idempotent, safe on an empty directory, and performs no + drainer logic; respawning drainers for the reported sessions is Phase + B2. + """ + + def _scan() -> list[str]: + result: list[str] = [] + for log in sorted(self._dir.glob("*.log")): + session_id = log.stem + # 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) + + def _count_dead(self, worker_key: str) -> int: + """Count complete (newline-terminated) dead-letter lines for a key. + + Returns 0 when no dead-letter file exists. Dead-letter records are + always written newline-terminated, so counting newlines yields the + number of complete records. + + Streamed (bounded memory), not ``read_bytes()``: a .dead.jsonl is + usually small but is NOT bounded -- a systematically-failing session + dead-letters every line -- and this is called on the same boot and + polled-/status paths as the .log scans. + """ + return self._stream_newlines(self._dead_path(worker_key)) + + def _all_worker_keys(self) -> list[str]: + """Return the sorted union of ``.log`` and ``.dead.jsonl`` stems. + + ``Path.stem`` only strips the final suffix, so for ``s1.dead.jsonl`` it + returns ``s1.dead``; the ``.dead.jsonl`` suffix is sliced explicitly to + recover the bare worker key. + """ + keys: set[str] = set() + for log in self._dir.glob("*.log"): + keys.add(log.stem) + for dead in self._dir.glob("*.dead.jsonl"): + keys.add(dead.name[: -len(".dead.jsonl")]) + return sorted(keys) + + async def derive_all_stats(self) -> dict[str, Any]: + """Derive live queue stats purely from disk, with a short TTL cache. + + 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 ( + self._stats_cache is not None + and (now - self._stats_cache_at) < self._stats_cache_ttl + ): + return self._stats_cache + + def _all() -> dict[str, Any]: + per_key: list[dict[str, Any]] = [] + in_queue_total = 0 + dead_total = 0 + for worker_key in self._all_worker_keys(): + try: + committed = self._read_committed_offset(worker_key) + except (OSError, ValueError): + # /status calls this (via pipeline_metrics); a corrupt or + # transiently-unreadable .offset must NOT 500 the health + # probe. Degrade to 0 for this key's stats -- mirroring the + # existing missing-file->0 convention in + # _read_committed_offset, and tending the conservation + # residual negative (benign, never a false `degraded`). + # Deliberately NO logging here: /status is polled, and a + # per-scan warning on a persistently-corrupt offset would + # flood the hot path. The visibility signal is the aggregate + # `spool.corrupt_offsets` field (see spool_stats()). + committed = 0 + # Streamed count of complete lines from committed -> EOF. + # Equivalent to the old f.read() + data[:last_nl+1].count(b"\n") + # (every b"\n" lies at or before the last one), but without + # materialising the undrained tail -- which can be gigabytes + # under a large backlog on this (polled) /status path. + in_queue = self._count_newlines(worker_key, committed) + dead = self._count_dead(worker_key) + per_key.append( + {"worker_key": worker_key, "in_queue": in_queue, "dead": dead} + ) + in_queue_total += in_queue + dead_total += dead + return { + "per_key": per_key, + "in_queue_total": in_queue_total, + "dead_total": dead_total, + } + + stats = await asyncio.to_thread(_all) + self._stats_cache = stats + self._stats_cache_at = now + return stats + + async def spool_stats(self) -> dict[str, int]: + """Cheap, aggregate-only spool footprint for the unauthenticated /status. + + 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 ( + self._spool_cache is not None + and (now - self._spool_cache_at) < self._spool_cache_ttl + ): + return self._spool_cache + + def _scan() -> dict[str, int]: + spool_bytes_total = 0 + pending_sessions = 0 + corrupt_offsets = 0 + for entry in self._dir.iterdir(): + if not entry.is_file(): + continue + try: + size = entry.stat().st_size + except FileNotFoundError: + # Raced with a concurrent delete_drained()/purge; the + # entry no longer exists -- simply exclude it, don't fail + # a cheap, best-effort aggregate over a live directory. + continue + spool_bytes_total += size + if entry.suffix == ".log": + try: + committed = self._read_committed_offset(entry.stem) + except ValueError: + # The .offset exists but is not a valid integer -- a + # GENUINELY corrupt offset. This is the one visibility + # signal for it (no logging anywhere, to avoid flooding + # the polled health path): surface it as an aggregate + # count on /status so `spool.corrupt_offsets > 0` is the + # operator's alarm. Count this file's bytes; skip its + # pending calc. + corrupt_offsets += 1 + continue + except OSError: + # A transient/racing FS error reading the offset (NOT + # corruption): count bytes, skip pending calc, and do + # NOT inflate corrupt_offsets with a non-corruption cause. + continue + if committed < size: + pending_sessions += 1 + return { + "pending_sessions": pending_sessions, + "spool_bytes_total": spool_bytes_total, + "corrupt_offsets": corrupt_offsets, + } + + try: + stats = await asyncio.to_thread(_scan) + except (OSError, ValueError): + # Queue dir missing/unavailable (e.g. Azure Files SMB remount) or a + # transient FS error mid-scan. /status is the health probe and MUST + # return 200 -- degrade to a sentinel and DO NOT cache it, so the + # next poll retries immediately once the filesystem recovers. All + # three fields are -1 = "temporarily unavailable" (distinct from a + # real 0, and from a real corrupt_offsets count). + return { + "pending_sessions": -1, + "spool_bytes_total": -1, + "corrupt_offsets": -1, + } + + self._spool_cache = stats + self._spool_cache_at = now + return stats + + async def dead_letter_keys(self) -> list[str]: + """Return sorted worker keys that have a ``.dead.jsonl`` file. + + Keys with only main-log data (no dead-letter file) are excluded. + ``Path.name`` is sliced by the ``.dead.jsonl`` suffix to recover the + bare worker key (``Path.stem`` would only strip ``.jsonl``). + """ + + def _scan() -> list[str]: + return sorted( + dead.name[: -len(".dead.jsonl")] + for dead in self._dir.glob("*.dead.jsonl") + ) + + return await asyncio.to_thread(_scan) + + async def purge_dead_letters(self, worker_key: str) -> int: + """Delete the dead-letter file for ``worker_key`` and return the count. + + Counts the dead-letter records via ``_count_dead``, then unlinks the + ``.dead.jsonl`` file. Returns the number of records removed (0 when no + dead-letter file exists). Raises ``ValueError`` for an unsafe key. + + 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: + with guard.file_lock: + count = self._count_dead(worker_key) + try: + path.unlink() + except FileNotFoundError: + pass + return count + + 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 + + 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 + + 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 + + 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, + } + + 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(): + # 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 + return accepted, written + + return await asyncio.to_thread(_seed) + + def _dead_payload_set(self, worker_key: str) -> set[bytes]: + """Return the set of original raw line bytes recorded as dead-letters. + + Each dead-letter record stores the original line (sans trailing + newline) either as ``payload`` (a UTF-8 string) or ``payload_b64`` + (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 + 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: 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(): + # Q-6: check `.log` existence BEFORE reading the whole + # `.dead.jsonl` into RAM (+ a payload set at ~3.6x its size). + # A key with only a `.dead.jsonl` (the common shape left by + # `delete_drained`) previously paid that read for nothing, + # every boot, forever. Free fix. + log_path = self._log_path(key) + if not log_path.exists(): + continue + # Fault-isolate this key -- a corrupt/unreadable + # dead-payload set or offset for ONE key must not abort the + # reconcile pass for every other key. The boot + # hook this feeds must never crash-loop the share it reads. + try: + dead_payloads = self._dead_payload_set(key) + if not dead_payloads: + continue + committed = self._read_committed_offset(key) + complete_end = self._complete_data_end(key) + pos = committed + with open(log_path, "rb") as f: + f.seek(committed) + while pos < complete_end: + raw = f.readline() + if not raw or not raw.endswith(b"\n"): + break + if raw[:-1] not in dead_payloads: + break + pos += len(raw) + total_skipped += 1 + if pos > committed: + final = self._offset_path(key) + tmp = self._dir / f"{key}.offset.tmp" + tmp.write_text(str(pos), encoding="utf-8") + os.replace(tmp, final) + except (OSError, ValueError): + logger.exception("recovery_reconcile_dead_key_failed key=%s", key) + continue + self._stats_cache = None + return total_skipped + + return await asyncio.to_thread(_reconcile) diff --git a/context_intelligence_server/queue_manager/protocol.py b/context_intelligence_server/queue_manager/protocol.py new file mode 100644 index 00000000..14ccc519 --- /dev/null +++ b/context_intelligence_server/queue_manager/protocol.py @@ -0,0 +1,127 @@ +"""QueueManager Protocol — the backend-neutral seam. + +A ``QueueManager`` manages a durable, per-session append-only queue for the +event-write pipeline. Consumers depend on this Protocol plus the ``Batch`` / +``Record`` value types, never on a concrete backend class. No ``Path``, +on-disk layout, or ``os.*`` detail appears here or in any value the Protocol +returns — those are private to a concrete backend +(:class:`~context_intelligence_server.queue_manager.filesystem.FileSystemQueueManager`, +and, later, an Azure equivalent). +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any, Protocol, runtime_checkable + + +@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 records read from a session's append-only log. + + Attributes: + 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 record. + This is the value passed to ``commit``. When no complete records + are available, ``end_offset == start_offset``. + """ + + session_id: str + 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] + + +@runtime_checkable +class QueueManager(Protocol): + """Durable, per-session append-only queue. + + The method set mirrors the on-disk backend's public surface. A backend + reports its own queue root via ``queues_dir``; every other on-disk detail + stays private to the implementation. + """ + + @property + def queues_dir(self) -> Any: ... + + async def heal_torn_tails(self) -> dict[str, int]: ... + + async def append(self, session_id: str, raw: bytes) -> None: ... + + async def read_batch(self, session_id: str, max_items: int) -> Batch: ... + + async def commit(self, session_id: str, new_offset: int) -> None: ... + + async def dead_letter(self, session_id: str, raw: bytes, error: str) -> None: ... + + async def delete_drained(self, session_id: str) -> bool: ... + + async def compact_committed_prefix( + self, session_id: str, min_prefix_bytes: int = 0 + ) -> int: ... + + async def read_dead_letters(self, session_id: str) -> list[dict]: ... + + async def read_first_line(self, key: str) -> bytes | None: ... + + async def classify_session( + self, key: str, head_is_resumable: Callable[[bytes], bool] + ) -> Any: ... + + async def reclaim(self, c: Any, is_owned: Callable[[], bool]) -> bool: ... + + async def reclaim_orphans( + self, before_ts: float, enabled: bool = True + ) -> dict[str, int]: ... + + async def active_sessions(self) -> list[str]: ... + + async def session_keys(self) -> list[str]: ... + + async def recover(self) -> list[str]: ... + + async def derive_all_stats(self) -> dict[str, Any]: ... + + async def spool_stats(self) -> dict[str, int]: ... + + async def dead_letter_keys(self) -> list[str]: ... + + async def purge_dead_letters(self, worker_key: str) -> int: ... + + async def expire_dead_letters( + self, now: float, retention_seconds: float, enabled: bool + ) -> dict[str, int]: ... + + async def recovery_seed_counts(self) -> tuple[int, int]: ... + + async def recovery_reconcile_dead(self) -> int: ... diff --git a/context_intelligence_server/registry.py b/context_intelligence_server/registry.py index ad078fee..e478a6c2 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 @@ -9,24 +10,32 @@ from pathlib import Path from typing import Any -from context_intelligence_server.blob_store import AsyncDiskBlobStore +from context_intelligence_server.blob_store import create_blob_store 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.queue_manager import ( + Batch, + QueueManager, + create_queue_manager, +) 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 +50,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 +81,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: @@ -95,7 +113,7 @@ def _ensure_infra(self) -> None: """ if self._queue_manager is None: settings = get_settings() - self._queue_manager = QueueManager(queues_dir=Path(settings.queues_path)) + self._queue_manager = create_queue_manager(settings) self._write_semaphore = asyncio.Semaphore(settings.write_concurrency) self._max_delivery_attempts = settings.max_delivery_attempts @@ -106,6 +124,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 +191,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 +234,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 +303,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 +317,16 @@ async def _process_one( ) async def _flush_barrier(self, worker: SessionWorker) -> None: - """The ONE Neo4j-write boundary: a semaphore-gated, awaited flush. - - Acquiring self.write_semaphore caps the number of concurrent Neo4j - write transactions across ALL session drainers (the starvation guard). - The offset must only ever advance AFTER this returns successfully. + """The one Neo4j-write boundary: a semaphore-gated, awaited flush. - 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 +334,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 @@ -323,12 +353,44 @@ async def drain_worker( poll_interval = min(flush_timeout, _DRAIN_POLL_INTERVAL) idle_elapsed = 0.0 attempts = 0 + # A (re)built worker starts with empty cross-handler counters. Restore the + # cursor the last commit persisted -- lazily, just before the FIRST real + # batch is processed -- so a rebuild resumes those counters instead of + # restarting from zero (which would remint node ids and duplicate them). + # Kept off the idle path so an idle worker never does this read. + cursor_restored = False + # Cursor state as it was BEFORE the first attempt on the current batch. + # A failed attempt leaves the cross-handler counters advanced; replaying + # the same batch from that advanced state mints fresh node ids and + # duplicates the Iteration. Restoring this snapshot before each retry + # makes the replay reproduce the SAME ids (idempotent MERGE). + pre_batch_cursor: dict[str, Any] | None = None while True: try: batch = await qm.read_batch(session_id, max_items=_DRAIN_MAX_BATCH) - 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 +402,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) @@ -351,21 +414,45 @@ async def drain_worker( idle_elapsed = 0.0 + # Restore the durable cursor once, before this worker's FIRST + # real batch is processed (and before the pre-attempt snapshot + # below, so that snapshot captures the restored state). A + # brand-new session has no committed cursor (read_cursor -> None + # -> no-op). + if not cursor_restored: + worker.services.restore_cursor(await qm.read_cursor(session_id)) + cursor_restored = True + + # Snapshot once, before the first attempt on this batch, so a + # retry can roll the counters back (see pre_batch_cursor above). + if attempts == 0: + pre_batch_cursor = worker.services.snapshot_cursor() + # --- dispatch + durable write barrier, one error path --- try: - 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 +476,66 @@ 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: roll the cross-handler counters back to + # their pre-attempt state so the replay reproduces the same + # node ids (idempotent MERGE) instead of duplicating them, + # then back off before re-reading the same offset. + worker.services.restore_cursor(pre_batch_cursor) await asyncio.sleep(poll_interval) continue 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, worker.services.snapshot_cursor() + ) + 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 +546,162 @@ 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, worker.services.snapshot_cursor() + ) # 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, worker.services.snapshot_cursor() + ) # 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 - await qm.commit(session_id, tail.end_offset) - self.record_written(len(tail.lines)) + return False # NOT finalized: keep worker alive, tail uncommitted + await qm.commit( + session_id, tail.end_offset, worker.services.snapshot_cursor() + ) + self.record_written(len(tail.records)) logger.debug( "batch_committed events=%d offset=%d", - 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 +714,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 +756,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,14 +847,23 @@ 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) + blob_store = create_blob_store(settings) _admin = settings.resolve_neo4j_admin() 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 +877,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 +886,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 +917,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 +946,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/admin.py b/context_intelligence_server/routers/admin.py index 4606a303..48d7f612 100644 --- a/context_intelligence_server/routers/admin.py +++ b/context_intelligence_server/routers/admin.py @@ -2,7 +2,7 @@ Endpoints manage the entra-identities and api-keys stores at runtime with no server restart. Mutations flow through ``IdentityStore.put`` / ``delete``, which -use the ROB-F2 commit order (write-file-then-swap-memory), so the persistent +use the write-file-then-swap-memory commit order, so the persistent file and the in-process dict are always in sync. **Store access pattern** @@ -13,10 +13,10 @@ the router and the main module. If the relevant store is ``None`` for the current ``auth_mode``, the endpoint returns **503**. -**Auth seam — T4 placeholder** +**Auth seam — enforcement placeholder** ``require_admin`` is a NO-OP dependency applied to the whole ``/admin`` router -via ``APIRouter(dependencies=[Depends(require_admin)])``. T5 replaces the +via ``APIRouter(dependencies=[Depends(require_admin)])``. This replaces the function body with real enforcement (static: admin_api_key; entra: IdentityAdmin App Role) — the routes and tests do not change. Tests override it via:: @@ -33,14 +33,26 @@ from __future__ import annotations +import asyncio +import json import logging import re +import time +from typing import Any from fastapi import APIRouter, Depends, HTTPException, Request -from pydantic import BaseModel, field_validator +from fastapi.responses import JSONResponse +from neo4j import READ_ACCESS, WRITE_ACCESS +from pydantic import BaseModel, Field, field_validator -from context_intelligence_server.config import _ALL_ZEROS_GUID, _GUID_RE +from context_intelligence_server.blob_processor import ( + BLOB_REF_CARRIER_PROPERTIES as _BLOB_REF_CARRIER_PROPERTIES, +) +from context_intelligence_server.blob_store import BlobReference, create_blob_store +from context_intelligence_server.config import _ALL_ZEROS_GUID, _GUID_RE, get_settings from context_intelligence_server.identity_store import IdentityStore +from context_intelligence_server.maintenance import coordinator +from context_intelligence_server.maintenance_ops import run_maintenance_operation # --------------------------------------------------------------------------- # Validation constants @@ -49,10 +61,32 @@ # Mirrors the 64-hex check in config._validate_api_keys (config.py:172-179). _HASH_RE = re.compile(r"^[0-9a-f]{64}$") -# Maximum contributor id length (TB-12). Matches the cap implied by config +# Maximum contributor id length. Matches the cap implied by config # (non-empty, non-whitespace, sane upper bound for an identifier string). _MAX_CONTRIBUTOR_LEN = 256 +# --------------------------------------------------------------------------- +# Blob-reclaim constants. +# --------------------------------------------------------------------------- + +# Hard mtime-floor safety net, defense-in-depth behind +# the durable undrained-queue gate. min_age_minutes below this is rejected +# (422) rather than silently raised -- a caller passing 0 must not be able to +# disable the age gate entirely. +_MIN_AGE_FLOOR_MINUTES = 15 + +# "sample" is bounded to keep the response small; totals (orphans_found, +# reclaimable_bytes) remain authoritative even when the sample is truncated. +_MAX_SAMPLE = 50 + +# Single-flight guard for the DESTRUCTIVE reclaim apply. Two concurrent applies +# would each honour max_delete independently, so together they could delete +# twice the operator's intended blast radius; a second overlapping apply is +# refused (409) rather than admitted. Per-process only. Flipped False->True with +# no await between the check and the set, so the check-and-set is atomic under +# asyncio. Dry-run never takes it -- only the delete phase is serialized. +_reclaim_apply_inflight = False + # --------------------------------------------------------------------------- # Module-level audit logger # --------------------------------------------------------------------------- @@ -61,7 +95,7 @@ # --------------------------------------------------------------------------- -# Auth seam — real enforcement (T5) +# Auth seam — real enforcement # --------------------------------------------------------------------------- @@ -70,9 +104,9 @@ def require_admin(request: Request) -> None: Applied router-wide via ``APIRouter(dependencies=[Depends(require_admin)])``. - Security model (design §6, T5): + Security model: - The middleware (BearerTokenMiddleware) ALWAYS enforces authentication on - /admin/* paths (they are never in an exempt set — TB-07 startup assertion). + /admin/* paths (they are never in an exempt set — enforced by a startup assertion). A missing/invalid token → 401 before this function is ever reached. - This dependency enforces *authorization* (not authentication): the request has already been authenticated; here we check whether the authenticated @@ -95,11 +129,11 @@ def require_admin(request: Request) -> None: Notes: - Tests override this with ``app.dependency_overrides[require_admin] = lambda: None`` - to bypass enforcement in T4 route tests (this is the standard FastAPI override + to bypass enforcement in route-level tests (this is the standard FastAPI override mechanism; the lambda's signature must satisfy ITS OWN declared parameters — FastAPI injects based on the override's signature, not the original's). - The ``roles`` check reads ONLY the ``roles`` claim — never ``groups``. - Group membership in the token cannot grant admin access (TB-09). + Group membership in the token cannot grant admin access. """ auth_mode: str = getattr(request.app.state, "auth_mode", "static") # Read auth metadata from scope state (set by BearerTokenMiddleware). @@ -146,7 +180,7 @@ def require_admin(request: Request) -> None: # 403 when the token's `roles` claim does not contain the required role. # ONLY the `roles` claim is checked — `groups` is intentionally excluded - # so group membership cannot grant admin access (TB-09). + # so group membership cannot grant admin access. roles: list[str] = scope_state.get("roles", []) if entra_admin_role not in roles: raise HTTPException( @@ -174,7 +208,7 @@ class IdentityBody(BaseModel): @field_validator("id") @classmethod def id_must_be_valid(cls, v: str) -> str: - """Validate contributor id: non-empty, non-whitespace, bounded, no null bytes (TB-12).""" + """Validate contributor id: non-empty, non-whitespace, bounded, no null bytes.""" if not v.strip(): raise ValueError("id must be a non-empty, non-whitespace string") if len(v) > _MAX_CONTRIBUTOR_LEN: @@ -194,7 +228,7 @@ class KeyBody(BaseModel): @field_validator("id") @classmethod def id_must_be_valid(cls, v: str) -> str: - """Validate contributor id: non-empty, non-whitespace, bounded, no null bytes (TB-12).""" + """Validate contributor id: non-empty, non-whitespace, bounded, no null bytes.""" if not v.strip(): raise ValueError("id must be a non-empty, non-whitespace string") if len(v) > _MAX_CONTRIBUTOR_LEN: @@ -207,7 +241,7 @@ def id_must_be_valid(cls, v: str) -> str: # --------------------------------------------------------------------------- -# Path-param validation helpers (TB-10) +# Path-param validation helpers # --------------------------------------------------------------------------- @@ -276,14 +310,14 @@ def _audit_put( """Emit one structured audit log line for a PUT (upsert) mutation. When *old_contributor* is set and differs from *contributor*, the entry - records an OVERWRITE event with old → new (TB-11). Otherwise it records + records an OVERWRITE event with old → new. Otherwise it records a normal insert/same-contributor upsert. NEVER logs raw keys — only the *target* (oid or hash) is recorded. """ who = _admin_who(request) if old_contributor is not None and old_contributor != contributor: - # Overwrite: different contributor (TB-11 explicit old→new audit). + # Overwrite: different contributor (explicit old→new audit). logger.info( "admin.audit action=put target=%s old_contributor=%r new_contributor=%r who=%s", target, @@ -309,6 +343,318 @@ def _audit_delete(request: Request, *, target: str) -> None: ) +def _audit_blob_reclaim_delete(request: Request, *, uri: str) -> None: + """Emit one structured audit log line per successfully-deleted blob. + + NEVER logs blob contents -- only the ``ci-blob://`` URI is recorded. + """ + logger.info( + "admin.audit action=blob_reclaim target=%s who=%s", + uri, + _admin_who(request), + ) + + +# --------------------------------------------------------------------------- +# Blob reclaim -- orphaned-blob GC. +# --------------------------------------------------------------------------- + + +def _access_mode_const(mode: str) -> str: + """Map the configured query access-mode string ("READ"/"WRITE") to the + driver's access-mode constant. + + Deliberately duplicated (not imported) from ``main._neo4j_access_const``: + importing from ``main`` here would create a circular import (``main`` + already imports ``routers.admin`` at module load time). Two lines of + duplication is cheaper than that coupling. + """ + return READ_ACCESS if mode == "READ" else WRITE_ACCESS + + +def _collect_blob_refs(obj: Any, out: set[str]) -> None: + """Recursively walk a decoded JSON value collecting ``$blob_ref`` URIs. + + This is STRUCTURAL extraction + over the parsed object, never a regex over the serialized string. A + regex anchored on ``ci-blob://`` truncates at the first unescaped + special character (e.g. a literal ``"`` in a session_id, which + ``queue_manager._validate_session_id`` explicitly permits -- it only + rejects ``/ \\ \\0``), silently misclassifying a genuinely-referenced + blob as orphan. ``json.loads`` has already resolved all escaping by the + time this function runs, so any character in a URI (quotes, non-ASCII) + is handled correctly -- there is no APOC path and no regex path. + """ + if isinstance(obj, dict): + ref = obj.get("$blob_ref") + if isinstance(ref, str): + out.add(ref) + for v in obj.values(): + _collect_blob_refs(v, out) + elif isinstance(obj, list): + for item in obj: + _collect_blob_refs(item, out) + + +# Reference-scan hardening: the known node properties that can carry a +# "ci-blob://" token anywhere in the graph, keyed to the property name +# Cypher matches on. +# +# ASSUMPTION -- make it explicit + greppable: the reference scan is scoped to +# EXACTLY these properties, not an all-property-all-node walk, for +# performance (a per-property UNION lets Neo4j evaluate ONE property per row +# instead of toString()-ing every key of every node -- this codebase has +# scar tissue from a 1.3M-node AllNodesScan stall). Adding a NEW ci-blob +# carrier property in the future (a new field-lifter, a new enricher +# property, etc.) REQUIRES adding it to this tuple -- otherwise the reclaim +# GC will not see refs stored there and could misclassify a live blob as an +# orphan. Covers every carrier in the hardening doc's table: +# data -- *Event.data (all 14 event types) +# tool_input -- ToolPreEvent/ToolPostEvent.tool_input (L1), ToolCall.tool_input (L2) +# prompt -- PromptSubmitEvent/PromptCompleteEvent.prompt (L1), Prompt.prompt (L2) +# response -- OrchestratorRun.response (L2) +# NOTE: the canonical allowlist lives in blob_processor as BLOB_REF_CARRIER_PROPERTIES +# and is imported (aliased to _BLOB_REF_CARRIER_PROPERTIES) at the top of this module -- +# a single source of truth shared by the mint site (blob_processor) and this reclaim +# scan. Do NOT re-declare it here. + +# Fallback extraction for carrier values that are NOT valid JSON (a bare +# string property, e.g. a plain-string tool_input/prompt that itself +# contains a ci-blob:// URI rather than the {"$blob_ref": "..."} wrapper). +# Safe here -- unlike a regex over a JSON-*serialized* string (see the +# docstring above) -- because these values are already fully-decoded Neo4j +# property strings with no JSON escaping left to trip over. +_BARE_BLOB_URI_RE = re.compile(r'ci-blob://[^"\s]+') + + +def _extract_blob_refs_from_value(val: str, out: set[str]) -> None: + """Extract every ``ci-blob://`` URI referenced by one carrier property value. + + Two extraction paths, unioned into *out*: + + 1. **Structural (preferred):** ``json.loads(val)`` then recurse with + :func:`_collect_blob_refs`. Handles the ``{"$blob_ref": "..."}`` + JSON-string carriers -- ``neo4j_store._sanitize_properties`` + JSON-serializes any dict/list property value on write, so + ``Event.data``, ``ToolCall.tool_input``, ``Prompt.prompt``, and + ``OrchestratorRun.response`` all round-trip through this path when + they hold a dict/list. + 2. **Regex fallback (only on JSON-parse failure):** a plain-string + carrier is written through VERBATIM (``_sanitize_properties`` only + JSON-serializes dict/list values -- a bare string is stored as-is), + so it is never valid JSON and always lands here. Extracts every bare + ``ci-blob://[^"\\s]+`` token directly from the decoded string -- + covers a lifted ``*.tool_input``/``*.prompt`` property that is a + plain string mentioning a blob URI. + + A value that is neither valid JSON nor contains a bare token contributes + nothing -- this can only ever fail to positively assert a reference, + never falsely assert one, matching the conservative-skip contract of + :func:`_scan_referenced_uris`. + """ + try: + obj = json.loads(val) + except (TypeError, ValueError): + out.update(_BARE_BLOB_URI_RE.findall(val)) + return + _collect_blob_refs(obj, out) + + +def _carrier_scan_clause(prop: str) -> str: + """Build one ``UNION ALL`` branch of the reclaim reference-scan query + for carrier property *prop*. + + ``data`` is special-cased: ``Event.data`` is always a JSON string + (``DefaultHandler`` writes ``json.dumps(data)`` -- see + ``handlers/data_layer_1/default.py``) and the scan for it is scoped to + ``:Event``, matching this carrier's scope in the pre-hardening scan. Every + other registered carrier is an unrestricted-label match with + ``toString()``, since the property may be lifted onto any node type as a + dict, list, or bare string. + """ + if prop == "data": + return "MATCH (n:Event) WHERE n.data CONTAINS 'ci-blob://' RETURN n.data AS val" + return ( + f"MATCH (n) WHERE n.{prop} IS NOT NULL " + f"AND toString(n.{prop}) CONTAINS 'ci-blob://' " + f"RETURN toString(n.{prop}) AS val" + ) + + +# Generated FROM _BLOB_REF_CARRIER_PROPERTIES (imported from +# blob_processor.BLOB_REF_CARRIER_PROPERTIES) -- not hand-duplicated -- so the +# query text can never drift from the allowlist. See +# tests/test_blob_processor.py for the regression test that locks this +# agreement structurally (extracts every `n.` reference out of the +# built query and diffs it against the allowlist tuple). +_BLOB_REF_SCAN_QUERY = " UNION ALL ".join( + _carrier_scan_clause(prop) for prop in _BLOB_REF_CARRIER_PROPERTIES +) + + +async def _scan_referenced_uris(request: Request) -> set[str]: + """Enumerate every ``ci-blob://`` URI referenced anywhere in the graph. + + Step 2 of the design -- deliberately GLOBAL, never workspace-filtered + (hazard #1): blobs are session_id-scoped while nodes are + (node_id, workspace)-scoped, so a per-workspace scan could delete another + workspace's live data. + + Reference-scan hardening: the referenced set is computed GRAPH-WIDE over the + known ``ci-blob://`` carrier properties (:data:`_BLOB_REF_CARRIER_PROPERTIES` + -- ``data``, ``tool_input``, ``prompt``, ``response``), not ``Event.data`` + only. This makes the scan correct BY CONSTRUCTION -- a strict superset of + the prior ``Event.data``-only scan -- rather than resting on the + (empirically true today, but unenforced) pipeline-ordering invariant that + ``DefaultHandler`` always persists every ref onto ``Event.data`` before any + field-lifter/enricher can strip or promote it elsewhere. Widening the scan + can only ever *protect* more blobs, never delete more: any URI the old + scan found is still found here (still scanned via the ``data`` branch), + plus any URI that lives ONLY on ``tool_input``/``prompt``/``response`` is + now ALSO found. See + ``tests/neo4j/test_blob_reclaim.py::test_b1_event_data_carries_every_blob_ref`` + for the regression test that continues to pin the pipeline-ordering + invariant (belt-and-suspenders now, not the sole safety net), and the + ``test_*_only_referenced_via_*`` tests alongside it that pin the new + per-property carriers directly. + + Query shape (performance-critical -- see :data:`_BLOB_REF_CARRIER_PROPERTIES`): + a ``UNION ALL`` of four single-property predicates, each touching exactly + ONE property per row (``data`` restricted to ``:Event``, matching the + prior scan's scope for that carrier; the other three unrestricted across + labels since ToolCall/Prompt/OrchestratorRun are ordinary nodes). This + avoids a pathological ``MATCH (n) ... [k IN keys(n) WHERE toString(n[k]) + ...]`` all-property-all-node walk, which would toString() every key of + every node in the graph -- this codebase has scar tissue from exactly + that shape of full scan (a 1.3M-node AllNodesScan stall). ``UNION ALL`` + (not plain ``UNION``) is deliberate: plain ``UNION``'s implicit DISTINCT + would force Neo4j to materialize and dedupe every row before returning + the first one, defeating streaming; the Python ``set`` below already + dedupes, so ``UNION ALL`` costs nothing and preserves the stream. + + No APOC. Each returned value is + extracted via :func:`_extract_blob_refs_from_value` (structural + ``json.loads`` + recursive walk, falling back to a bare-token regex only + on JSON-parse failure -- see that function's docstring for why the regex + path is safe here and was NOT safe for the original ``Event.data`` scan). + A malformed/unparseable value is skipped conservatively -- it can never + positively assert an orphan, only fail to positively assert a reference. + + Streams rows via the async driver iterator rather than materializing all + carrier-property strings at once (only the resulting, much smaller, URI + set is retained), per the design's "bound its own work" cost note. + """ + driver = request.app.state.neo4j_query_driver + access_mode = _access_mode_const(request.app.state.neo4j_query_access_mode) + referenced: set[str] = set() + async with driver.session(default_access_mode=access_mode) as session: + # Graph-wide, per-property UNION ALL -- see the docstring above for + # why this shape (not an all-property walk, not plain UNION). The + # query text is generated from _BLOB_REF_CARRIER_PROPERTIES + # (_BLOB_REF_SCAN_QUERY, module level) -- not hand-duplicated here -- + # so it can never drift from the allowlist. + result = await session.run(_BLOB_REF_SCAN_QUERY) + async for record in result: + val = record["val"] + if not isinstance(val, str): + # toString() on a non-null value is always a str; this guards + # conservatively against an unexpected driver type mapping. + continue + _extract_blob_refs_from_value(val, referenced) + return referenced + + +async def _select_orphans(request: Request, *, min_age_minutes: int) -> dict[str, Any]: + """The ONE selection path shared by dry-run and apply. + + Returns a dict with every response field EXCEPT ``dry_run``/``sample``/ + ``rescanned``/``deleted``/``deleted_bytes`` (the caller fills those in), + plus a ``candidates`` key (list[BlobReference], sorted by uri for + deterministic sampling/capping) that the caller pops before returning the + response and uses to actually delete in apply mode. + + Every blob is reached only through the BlobStore protocol -- ``scan()`` + streams a ``BlobReference`` (uri + size + last_modified) per blob; no + filesystem path or on-disk layout crosses into this router. + + Safety gates applied to every blob not in the referenced set: + 1. Undrained-queue gate (primary, durable): skipped when the + session has a live worker (``registry.active_sessions()``) OR its + queue is not fully drained (``QueueManager.is_fully_drained``, + durable across restarts). Counted as ``skipped_pending_session``. + 2. Age floor (defense-in-depth): skipped when younger than + ``min_age_minutes`` (already clamped >= ``_MIN_AGE_FLOOR_MINUTES`` by + the request body validator), measured by ``BlobReference.last_modified``. + Counted as ``skipped_recent``. + """ + settings = get_settings() + blob_store = create_blob_store(settings) + + referenced = await _scan_referenced_uris(request) + + registry = request.app.state.registry + queue_manager = registry.queue_manager + live_workers = set(registry.active_sessions()) + + now = time.time() + age_cutoff_seconds = min_age_minutes * 60 + + scanned = 0 + candidates: list[BlobReference] = [] + skipped_recent = 0 + skipped_pending_session = 0 + reclaimable_bytes = 0 + + async for ref in blob_store.scan(): + scanned += 1 + if ref.uri in referenced: + continue + if ref.session_id in live_workers or not await queue_manager.is_fully_drained( + ref.session_id + ): + skipped_pending_session += 1 + continue + if now - ref.last_modified < age_cutoff_seconds: + skipped_recent += 1 + continue + candidates.append(ref) + reclaimable_bytes += ref.size + + candidates.sort(key=lambda b: b.uri) + + return { + "scanned_disk_blobs": scanned, + "referenced_uris": len(referenced), + "orphans_found": len(candidates), + "reclaimable_bytes": reclaimable_bytes, + "skipped_recent": skipped_recent, + "skipped_pending_session": skipped_pending_session, + "candidates": candidates, + } + + +class BlobReclaimBody(BaseModel): + """Body for POST /admin/blobs/reclaim.""" + + dry_run: bool = True + min_age_minutes: int = 60 + max_delete: int | None = Field(default=None, ge=1) + + @field_validator("min_age_minutes") + @classmethod + def _min_age_at_least_floor(cls, v: int) -> int: + """Reject (422) below the hard safety floor + rather than silently raising -- a caller passing 0 must not be able + to disable the age gate. + """ + if v < _MIN_AGE_FLOOR_MINUTES: + raise ValueError( + f"min_age_minutes must be >= {_MIN_AGE_FLOOR_MINUTES} " + f"(hard safety floor); got {v}" + ) + return v + + # --------------------------------------------------------------------------- # Per-request store dependencies (via app.state — no circular import) # --------------------------------------------------------------------------- @@ -375,14 +721,14 @@ def put_identity( """Upsert an entra identity (OID → contributor). Path-param guard: ``oid`` must be a valid GUID in lowercase hex and must - not be the all-zeros sentinel → **422** on violation (TB-10). + not be the all-zeros sentinel → **422** on violation. Write-through via ``IdentityStore.put``: the persistent file is updated atomically and the in-process map (shared with ``EntraResolver``) is updated immediately — no server restart required. An overwrite (existing oid with a different contributor) emits an explicit - audit line recording old → new contributor (TB-11). + audit line recording old → new contributor. Returns the stored record as ``{oid, id[, display_name]}``. """ @@ -447,10 +793,10 @@ def put_key( """Upsert a static API-key entry (sha256 hex → contributor). Path-param guard: ``sha256hash`` must be exactly 64 lowercase hex chars - → **422** on violation (TB-10). + → **422** on violation. Admin-key guard: the hash of the configured ``admin_api_key`` cannot be - shadow-bound via this endpoint → **409** (TB-05). The admin key lives in + shadow-bound via this endpoint → **409**. The admin key lives in config, not in the data keystore; rebinding its hash here would be confusing and is explicitly rejected. @@ -464,7 +810,7 @@ def put_key( """ _validate_hash(sha256hash) - # Admin-key un-shadowable guard (TB-05): reject PUT targeting the admin + # Admin-key un-shadowable guard: reject PUT targeting the admin # key's hash. The admin key is a config credential, not a data store # entry; allowing it to be shadowed here would silently rebind it. admin_digest: str | None = getattr(request.app.state, "admin_api_key_digest", None) @@ -501,10 +847,10 @@ def delete_key( """Delete a static API-key entry. Path-param guard: ``sha256hash`` must be exactly 64 lowercase hex chars - → **422** on violation (TB-10). + → **422** on violation. Admin-key guard: the hash of the configured ``admin_api_key`` cannot be - deleted via this endpoint → **409** (TB-05). The admin key is the + deleted via this endpoint → **409**. The admin key is the bootstrap floor — deleting it via the API must not be possible. Returns 200 on success, 404 when the hash is not present. @@ -512,7 +858,7 @@ def delete_key( """ _validate_hash(sha256hash) - # Admin-key un-deletable guard (TB-05): reject DELETE targeting the admin + # Admin-key un-deletable guard: reject DELETE targeting the admin # key's hash. The admin key is the emergency-bootstrap credential; it must # never be deletable via the API (operators would lock themselves out). admin_digest: str | None = getattr(request.app.state, "admin_api_key_digest", None) @@ -544,3 +890,201 @@ def list_keys( return { "keys": [{"hash": h, "id": record.get("id", "")} for h, record in store.items()] } + + +# --- Blob reclaim (orphaned-blob GC) ---------------------------------------- + + +@router.post("/blobs/reclaim", status_code=200) +async def reclaim_blobs(body: BlobReclaimBody, request: Request) -> dict[str, Any]: + """Preview (dry-run) or apply reclamation of orphaned blob files. + + An orphan is an on-disk blob (``//blobs/.json``) + whose ``ci-blob://`` URI is referenced by NO ``:Event.data`` anywhere in + the graph (scanned globally, across ALL workspaces), and whose session is + both fully drained (durable ``QueueManager`` state, not in-memory worker + liveness) and older than ``min_age_minutes``. See + the safety amendments (dry-run default, required max_delete cap, + graph-wide reference scan). + + ``dry_run=true`` (default) computes and reports the candidate set without + deleting anything. ``dry_run=false`` requires ``max_delete`` (422 + otherwise -- a conscious blast-radius opt-in for an irreversible, + cross-workspace delete) and performs its OWN fresh, authoritative + ``_select_orphans`` scan at delete time (``rescanned: true`` in the + response) -- it never deletes a URI that is referenced or pending at the + moment of deletion, independent of any earlier dry-run preview. + + Deletion goes through ``BlobStore.delete(uri, if_unmodified=ref)`` -- a + fenced compare-and-delete that refuses (leaves the blob untouched) if the + blob was rewritten since ``scan()`` observed it, and is idempotent (a blob + already gone counts as not-deleted, never an error). Capped at + ``max_delete``; ``orphans_found`` and ``reclaimable_bytes`` always reflect + the FULL candidate set even when ``max_delete`` caps how many are actually + removed. One structured audit log line is emitted per successful delete; + blob CONTENTS are never logged, only the ``ci-blob://`` URI. + """ + if not body.dry_run and body.max_delete is None: + raise HTTPException( + status_code=422, + detail=( + "max_delete is required when dry_run=false -- a conscious " + "blast-radius opt-in for an irreversible, cross-workspace " + "delete. Omit dry_run (or set it true) to preview first." + ), + ) + + if body.dry_run: + # Preview only: compute the candidate set and report it, delete nothing. + selection = await _select_orphans(request, min_age_minutes=body.min_age_minutes) + candidates: list[BlobReference] = selection.pop("candidates") + return { + "dry_run": True, + **selection, + "sample": [b.uri for b in candidates[:_MAX_SAMPLE]], + "rescanned": False, + "deleted": 0, + "deleted_bytes": 0, + } + + assert body.max_delete is not None # guaranteed by the 422 guard above + + # Acquire the destructive-apply single-flight BEFORE the authoritative scan, + # so a rejected concurrent apply fails fast and never even scans. + global _reclaim_apply_inflight + if _reclaim_apply_inflight: + raise HTTPException( + status_code=409, + detail=( + "a blob-reclaim apply is already in progress -- concurrent " + "applies would each honour max_delete independently and " + "together exceed the intended blast radius. Retry once it " + "completes." + ), + ) + _reclaim_apply_inflight = True + try: + # The apply's OWN fresh, authoritative scan at delete time -- it never + # deletes a URI that is referenced or pending at the moment of deletion, + # independent of any earlier dry-run preview. + selection = await _select_orphans(request, min_age_minutes=body.min_age_minutes) + candidates = selection.pop("candidates") + response: dict[str, Any] = { + "dry_run": False, + **selection, + "sample": [b.uri for b in candidates[:_MAX_SAMPLE]], + "rescanned": True, + "deleted": 0, + "deleted_bytes": 0, + } + + blob_store = create_blob_store(get_settings()) + deleted = 0 + deleted_bytes = 0 + # Each delete is fenced against the reference just observed by the scan + # above (delete refuses if the blob was rewritten since), so a blob that + # was concurrently re-minted or modified is left intact and counted as + # not-deleted rather than destroyed. + for ref in candidates[: body.max_delete]: + if not await blob_store.delete(ref.uri, if_unmodified=ref): + continue # absent or changed since scan -- left untouched + deleted += 1 + deleted_bytes += ref.size + _audit_blob_reclaim_delete(request, uri=ref.uri) + finally: + _reclaim_apply_inflight = False + + response["deleted"] = deleted + response["deleted_bytes"] = deleted_bytes + return response + + +# --- Maintenance operation -------------------------------------------------- +# +# This endpoint is on ``maintenance.MAINTENANCE_ALLOW_LIST`` (enforced +# structurally by ``main._assert_maintenance_endpoint_allow_listed`` at +# startup) so it stays reachable even while the gate is closed -- otherwise +# it would 503 at exactly the moment it exists to unblock. + + +@router.post("/maintenance", status_code=202) +async def post_maintenance(request: Request) -> JSONResponse: + """Trigger the maintenance repair operation (dedup -> :Node backfill -> + schema DDL), reusing the existing ``neo4j_store.run_repair`` -- see + ``maintenance_ops.run_maintenance_operation``. + + Single-flight: ``coordinator.try_begin_op()`` is a + synchronous compare-and-swap with no ``await`` between check and set, so + two concurrent POSTs can never both start an op. A second POST while one + is running gets **409** with the in-progress ``run_id`` -- it never + starts a second op. + + Returns promptly: the operation runs as a background + task; this handler returns **202** immediately rather than blocking for + the op's full duration (which includes the quiesce sleep + an + O(graph-size) dedup pass). + + Idempotent honest re-scan: POST against an already-clean graph + still performs a genuine ``run_repair`` call and returns a **fresh** + ``run_id``/``completed_at`` with ``records_affected: 0`` -- it never + short-circuits, which would make ``0`` indistinguishable from "did not + run" and defeat the freshness marker. + """ + run_id = coordinator.try_begin_op() + if run_id is None: + current = coordinator.current_op() + return JSONResponse( + status_code=409, + content={ + "detail": "maintenance operation already running", + "run_id": current.run_id, + "state": current.state, + }, + ) + + settings = get_settings() + # Defensive getattr (matches _check_driver_connected's convention above): + # lets this route degrade gracefully rather than 500 if ever hit before + # lifespan has bound a driver -- run_repair() will raise on a None + # driver, which finish_op() records as a normal `failed` outcome. + driver = getattr(request.app.state, "neo4j_driver", None) + task = asyncio.create_task( + run_maintenance_operation( + driver, + run_id, + quiesce_seconds=settings.maintenance_quiesce_seconds, + ) + ) + coordinator.retain_task( + task + ) # strong ref -- see MaintenanceCoordinator.retain_task + + op = coordinator.current_op() + return JSONResponse( + status_code=202, + content={"run_id": run_id, "state": op.state, "started_at": op.started_at}, + ) + + +@router.get("/maintenance", status_code=200) +async def get_maintenance() -> dict[str, Any]: + """Report maintenance-operation progress. + + Separate, admin-authenticated route, and never + folded into ``/status`` (``/status`` is unauthenticated; folding these + fields in would be an auth-boundary leak). ``state`` initializes to + ``"unknown"`` on boot so "never ran" is distinguishable from "ran, record + lost to a crash". + """ + st = await coordinator.status() + op = st.op + return { + "mode": st.mode, + "state": op.state, + "run_id": op.run_id, + "started_at": op.started_at, + "completed_at": op.completed_at, + "elapsed_seconds": st.elapsed_seconds, + "records_affected": op.records_affected, + "error": op.error, + } 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/routers/version.py b/context_intelligence_server/routers/version.py index 1273273f..97377136 100644 --- a/context_intelligence_server/routers/version.py +++ b/context_intelligence_server/routers/version.py @@ -4,19 +4,23 @@ from fastapi import APIRouter -from context_intelligence_server.status import SERVER_VERSION +from context_intelligence_server.status import SCHEMA_VERSION, SERVER_VERSION router = APIRouter() @router.get("/version") -async def get_version() -> dict[str, str]: - """Return the running server version. +async def get_version() -> dict[str, str | int]: + """Return the running server version and expected graph schema version. This endpoint is intentionally unauthenticated so clients can check server compatibility without credentials. + ``schema_version`` is a read-only baseline data point (no comparison or + upgrade logic lives here — see ``SCHEMA_VERSION`` in ``status.py``). + Returns: - JSON object with a single ``version`` key, e.g. ``{"version": "2.0.0"}``. + JSON object with ``version`` and ``schema_version`` keys, e.g. + ``{"version": "2.0.0", "schema_version": 1}``. """ - return {"version": SERVER_VERSION} + return {"version": SERVER_VERSION, "schema_version": SCHEMA_VERSION} diff --git a/context_intelligence_server/services.py b/context_intelligence_server/services.py index b5f0230b..ccb154fa 100644 --- a/context_intelligence_server/services.py +++ b/context_intelligence_server/services.py @@ -7,11 +7,14 @@ from __future__ import annotations +import dataclasses import fnmatch import logging +from dataclasses import asdict 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 +157,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 +221,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: @@ -244,66 +243,88 @@ def __init__( self.data_layer_3 = DataLayer3State() # ------------------------------------------------------------------ - # Session node management + # Durable cursor (survives a worker rebuild) # ------------------------------------------------------------------ - 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*. + def snapshot_cursor(self) -> dict[str, Any]: + """Return a JSON-safe snapshot of cross-handler cursor state. - Uses a two-tier lookup for replay resilience: + Snapshots the whole ``data_layer_2``/``data_layer_3`` dataclasses via + ``asdict`` (every field is JSON-native) rather than a hand-picked + allowlist that would silently miss a newly added field. + """ + return { + "dl2": asdict(self.data_layer_2), + "dl3": asdict(self.data_layer_3), + } - 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'``. + def restore_cursor(self, record: dict[str, Any] | None) -> None: + """Restore cross-handler cursor state from a persisted snapshot. - 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. + No-op on ``record is None`` (a brand-new session, or a legacy + ``.offset`` with no cursor). Otherwise only field NAMES present on the + current dataclass are assigned -- unknown/renamed keys are dropped and a + field absent from the record keeps its default, so the persisted format + tolerates dataclass evolution in both directions without a version bump. - Only caches session_id after a successful write to ensure retry - resilience on write failure. + Any failure is caught and logged, leaving the dataclasses at their + defaults: a corrupt or unexpected cursor must never crash boot. + """ + if record is None: + return + try: + for key, target in ( + ("dl2", self.data_layer_2), + ("dl3", self.data_layer_3), + ): + value = record.get(key) + if not isinstance(value, dict): + continue + valid_fields = {f.name for f in dataclasses.fields(type(target))} + for field_name, field_value in value.items(): + if field_name in valid_fields: + setattr(target, field_name, field_value) + except Exception: + logger.warning("cursor_restore_failed", exc_info=True) + + # ------------------------------------------------------------------ + # Session node management + # ------------------------------------------------------------------ + + 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*. + + 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. - await self.graph.upsert_node( - session_id, - {"labels": ["Session"], "status": "running", "session_id": session_id}, - ) + # Upsert a stub so this worker's own flush uses MERGE (idempotent) + # instead of racing a second worker into creating a duplicate node. + stub_data: dict[str, Any] = { + "labels": ["Session"], + "status": "running", + "session_id": session_id, + } + # Populate-if-missing: backfill working_dir on a node created before + # working_dir was known (or by a bare reference). Only when the event + # supplies one AND the node still lacks it; the DB-level coalesce + # guarantees an already-set value is never clobbered. + if data.get("working_dir") and not existing.get("working_dir"): + stub_data["working_dir"] = data["working_dir"] + await self.graph.upsert_node(session_id, stub_data) 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", @@ -316,6 +337,10 @@ async def ensure_session_node(self, session_id: str, data: dict[str, Any]) -> No node_data["started_at"] = _ts if "agent" in data: node_data["agent"] = data["agent"] + # Lift working_dir onto the Session node when the event carries one; + # absent/empty leaves it null for a later event to populate. + if data.get("working_dir"): + node_data["working_dir"] = data["working_dir"] await self.graph.upsert_node(session_id, node_data) self._seen_sessions.add(session_id) # only cache after successful write @@ -323,28 +348,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..6fdc8080 100644 --- a/context_intelligence_server/status.py +++ b/context_intelligence_server/status.py @@ -13,6 +13,11 @@ # Resolved once at import time — never changes within a process lifetime. SERVER_VERSION: str = _pkg_version("context-intelligence-server") +# Graph data-model version, distinct from the server release version above. +# Bumped only when the node/edge schema changes; /status compares it against +# the graph's stored :SchemaMeta.schema_version to surface drift. +SCHEMA_VERSION: int = 1 + if TYPE_CHECKING: from context_intelligence_server.registry import SessionRegistry @@ -61,6 +66,100 @@ 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 + # Why the graph-wide schema gate is currently closed, or None when open. + # Single global value: the un-migrated-data condition is graph-wide, not + # per-session. Set when boot detects/repairs un-migrated data, cleared once + # the gate re-opens; surfaced on /status so an operator can see the cause. + degraded_reason: str | None = None + + 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 degrade(self, reason: str) -> None: + """Record why the graph-wide schema gate is closed.""" + self.degraded_reason = reason + + def clear_degraded(self) -> None: + """Clear the reason once the schema gate re-opens.""" + self.degraded_reason = None + + 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 +184,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..b58d4e38 --- /dev/null +++ b/context_intelligence_server/writer_lease.py @@ -0,0 +1,508 @@ +"""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 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.lease_store import ( + LeaseRecord, + LeaseStore, + create_lease_store, +) +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_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.""" + + +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._store: LeaseStore | 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 + + def _build_record(self, heartbeat: float) -> LeaseRecord: + """The lease record this process would write at *heartbeat* -- pure + identity, no I/O. The store persists it; the detector owns it.""" + return LeaseRecord( + 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, + lease_version=_LEASE_VERSION, + ) + + # ----------------------------------------------------------------- + # 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. + # The store resolves the directory lazily per op, so this constructs + # nothing and reads no path here. + self._dir_source = dir_source + self._store = create_lease_store(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 + assert self._store is not None + + rec = await self._io(self._store.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() + store = self._store + await self._io(lambda: store.write(self._build_record(heartbeat))) + await asyncio.sleep(self._confirm_delay) + rec2 = await self._io(self._store.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_source 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_source()}\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: + assert self._store is not None + rec = await self._io(self._store.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() + store = self._store + await self._io(lambda: store.write(self._build_record(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 + assert self._store is not None + try: + rec = await asyncio.wait_for( + self._io(self._store.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._store is None: + return + store = self._store + timeout = self._acquire_timeout if self._acquire_timeout is not None else 5.0 + try: + await asyncio.wait_for( + self._io(lambda: store.delete_if_owned(self.owner)), + 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 bbca231d..518c92a3 100644 --- a/docs/azure-deployment.md +++ b/docs/azure-deployment.md @@ -494,7 +494,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) @@ -749,26 +751,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/maintenance-mode.md b/docs/maintenance-mode.md new file mode 100644 index 00000000..82104d11 --- /dev/null +++ b/docs/maintenance-mode.md @@ -0,0 +1,371 @@ +# Maintenance Mode + +Runtime, no-restart gate that refuses ingest and query while the graph's +identity constraint is absent or a repair operation is in flight — the +live counterpart to the cold-start refusal described in the [README's +"Cold Start No Longer Auto-Migrates"](../README.md) section. All state +lives in one seam: `context_intelligence_server/maintenance.py`'s +`MaintenanceCoordinator` (module singleton `coordinator`). The HTTP gate +(`maintenance_gate_middleware`), the ingest drain loop +(`registry.drain_worker`), and `GET /status` all read from this one +coordinator, so they can never disagree about whether the server is in +maintenance. + +## Why this exists + +`node_node_id_workspace_unique` is the `:Node` uniqueness constraint that +makes every write path's `MERGE` safe. If that constraint is missing — +because the graph has duplicate legacy nodes it conflicts with — the +schema is in a state where **both writes and reads are unsafe**: a write +can silently create a duplicate node instead of merging onto the existing +one, and a read (`/cypher`) can return results computed against that +duplicated, inconsistent shape. Likewise, while an operator-triggered +repair (dedup + `:Node` backfill + constraint creation) is actually +executing, the graph is mid-mutation — reads and writes during that +window would race the repair. + +Maintenance mode is the server's answer to both cases: refuse the request +with a `503` and a `Retry-After` header instead of serving a request +against schema-unsafe or mid-repair data. + +## The mode state machine + +`MaintenanceCoordinator._derive_mode()` is the **one** place mode is +computed; every caller (`gate_closed()`, `status()`) goes through it, so a +mode transition is always caught regardless of which surface polls first. +`MaintenanceMode` is one of four values: + +| Mode | When | Gates ingest/query? | +|---|---|---| +| `healthy` | Constraint present, no op running, boot-time untagged-node count is 0 | No | +| `degraded` | Constraint present, no op running, but the **boot-time** untagged-node probe found `> 0` nodes lacking the `:Node` label | No | +| `unknown` | The live constraint probe could not run (Neo4j unreachable, credentials rejected, etc.) | No — an unreachable graph can't serve writes/reads anyway, so gating here would only buy an outage with no safety benefit | +| `maintenance` | **Either** a maintenance op is currently running (`op.state == "running"`), **or** the live probe found the `:Node` constraint absent | **Yes** — this is the only mode `gate_closed()` treats as closed | + +Exact derivation order (`maintenance.py:267-291`), evaluated in this +priority: + +1. `op.state == "running"` → `mode="maintenance"`, `reason="maintenance operation in progress"` +2. else `constraint_present is False` → `mode="maintenance"`, `reason=":Node uniqueness constraint absent -- migration required"` +3. else `constraint_present is None` → `mode="unknown"`, `reason="constraint probe could not determine graph state"` +4. else boot-time `untagged_nodes > 0` → `mode="degraded"`, `reason="{n} node(s) lacking the :Node label"` +5. else → `mode="healthy"`, `reason=None` + +`gate_closed()` is exactly `mode == "maintenance"` — `degraded` and +`unknown` never gate (`maintenance.py:295-303`). + +### The live probe + +The constraint check is a cheap catalog read, not a data scan: + +```cypher +SHOW CONSTRAINTS YIELD name WHERE name = 'node_node_id_workspace_unique' RETURN count(*) AS c +``` + +It is **TTL-cached and single-flight** (`maintenance_probe_ttl_seconds`, +default `5.0` seconds): the fast path never takes a lock, and concurrent +callers that see an expired cache collapse into exactly one live probe. +This is what makes maintenance mode **self-clearing with no server +restart** — once the constraint is repaired out-of-band, the next probe +(within the TTL window) sees `constraint_present=True` and the mode flips +back on its own. + +## What happens while gated + +While `mode == "maintenance"`: + +- **Ingest is gated in the drain loop, not at spawn.** The check is the + first statement inside `drain_worker`'s `while True:` loop + (`registry.py:373-388`), before `read_batch`. This placement (rather + than refusing at `get_or_create → start_drain`) is deliberate: gating + *spawn* would be a **latch** — refuse once, never retried — the exact + defect class this feature exists to prevent. Drainers still spawn + during maintenance; they idle in a `_GATED_POLL_INTERVAL` sleep loop + instead. Because everything downstream of the gate check (`read_batch`, + `process`, `flush`, `commit`) is unreachable while gated, the on-disk + offset **never advances** during a gated window — a crash mid-maintenance + and restart replays from the last successfully-flushed offset, exactly + once. +- **`POST /events`** (including dead-letter `?replay=true` re-ingestion, + since it flows through the same drain loop) is refused via the HTTP + gate below before it can even reach the queue-append step. +- **`POST /cypher`** is refused via the same HTTP gate. +- **`GET /status` and `GET /version` stay up.** They are explicitly on the + allow-list (see below) — this is how an operator/automation observes + maintenance mode in the first place. + +### The structured 503 + +`maintenance_gate_middleware` (registered on `app` itself, via +`app.middleware("http")(maintenance_gate_middleware)` at `main.py:466`) +runs on every request whose path is not on the allow-list. When the +coordinator reports `mode == "maintenance"`, it returns the **one** +producer of this response, `maintenance_response()` — deliberately not a +FastAPI `HTTPException` (which would render `{"detail": ...}`, the wrong +contract here): + +```json +{ + "status": "maintenance", + "reason": ":Node uniqueness constraint absent -- migration required", + "retry_after": 30, + "schema_health": "degraded", + "maintenance_started_at": "2026-08-13T12:00:00.123456+00:00" +} +``` + +with an HTTP header `Retry-After: 30` (the numeric value mirrors the +`retry_after` field). `retry_after` is sourced from +`maintenance_retry_after_seconds` (default `30`). + +`reason` is the same human-readable string from the mode derivation above +(e.g. `"maintenance operation in progress"` when an op is running). +`schema_health` in this body is a **coarser, two-state** field than the +one on `/status` (below) — it is `"unknown"` only when the live probe +returned `None`, and `"degraded"` for every other case that reaches this +response, **including** the "op running, constraint actually present" +case. Operators should not read `"degraded"` here as "the constraint is +absent" — check `reason` for the actual cause. + +## `GET /status` fields + +`GET /status` is unauthenticated and always reachable (it's on the +allow-list). The maintenance-relevant fields, set in `main.py`'s +`get_status` handler (`main.py:897-933`): + +| Field | Source | Meaning | +|---|---|---| +| `mode` | live, via `coordinator.status()` | `"healthy"` \| `"degraded"` \| `"unknown"` \| `"maintenance"` — de-latched: reflects the current probe, not a boot-time snapshot | +| `maintenance_started_at` | live | ISO-8601 UTC timestamp the **current** maintenance window opened; `null` when not in maintenance | +| `maintenance_elapsed_seconds` | live | Seconds since that window opened; `null` when not in maintenance | +| `schema_health` | live (computed independently in `main.py`, a **3-state** version, not reused from the 503's 2-state ternary) | `"unknown"` if the probe returned `None`; `"degraded"` if the constraint is absent **or** the boot-time untagged count was `> 0`; else `"healthy"` | +| `untagged_nodes` | **boot-time snapshot**, not live | Count of `:Node`-label-missing nodes at the last server boot. Documented as boot-time — it is not a gate input; the live gate/mode signal is the constraint probe above | +| `schema_checked_at` | `datetime.now(UTC)` at the moment this `/status` call is served | Not the exact instant the constraint probe last ran — the probe result may be up to `maintenance_probe_ttl_seconds` stale | +| `degraded_reason` | **boot-time snapshot**, set once by `_record_schema_health` at startup | Human-readable cause captured at boot (e.g. `":Node uniqueness constraint absent (data conflict)"`, `"N node(s) lacking the :Node label"`, or the probe-unreachable message); **not** recomputed on every call | + +Example (in maintenance, constraint absent): + +```json +{ + "mode": "maintenance", + "maintenance_started_at": "2026-08-13T12:00:00.123456+00:00", + "maintenance_elapsed_seconds": 42.7, + "schema_health": "degraded", + "untagged_nodes": 0, + "schema_checked_at": "2026-08-13T12:00:42.831112+00:00", + "degraded_reason": null +} +``` + +(`degraded_reason` is `null` here deliberately — it is a boot-time field, +and this example assumes the constraint was present at boot and only went +absent later; a boot-time degradation would populate it.) + +Example (healthy): + +```json +{ + "mode": "healthy", + "maintenance_started_at": null, + "maintenance_elapsed_seconds": null, + "schema_health": "healthy", + "untagged_nodes": 0, + "schema_checked_at": "2026-08-13T12:05:00.001112+00:00", + "degraded_reason": null +} +``` + +**Note:** `/status` does not surface the coordinator's internal op record +(`run_id` / `state` / `records_affected` / `error`). That detail is +deliberately kept off the unauthenticated `/status` surface and lives on +the admin-authenticated [`GET /admin/maintenance`](#postget-adminmaintenance--the-live-endpoint-contract) +instead, plus the transition log lines (below). + +## Allow-list — blast radius + +`MAINTENANCE_ALLOW_LIST` (`maintenance.py:100-102`) is an **allow-list, +not a deny-list**, matched on exact path (`request.url.path in +MAINTENANCE_ALLOW_LIST`) — deliberately, so any route added to the app +later is blocked-by-default rather than silently exempt: + +```python +MAINTENANCE_ALLOW_LIST: frozenset[str] = frozenset( + {"/status", "/version", "/admin/maintenance", "/docs", "/openapi.json"} +) +``` + +| Stays reachable during maintenance | Returns `503` during maintenance | +|---|---| +| `GET /status` | `POST /events` (including dead-letter replay) | +| `GET /version` | `POST /cypher` | +| `POST`/`GET /admin/maintenance` (live -- the one `/admin/*` route reachable during maintenance; see [the endpoint contract below](#postget-adminmaintenance--the-live-endpoint-contract)) | `GET /blobs/{session_id}`, `GET /blobs/{session_id}/{key}` | +| `GET /docs`, `GET /openapi.json` (Swagger UI / OpenAPI schema) | `GET /queues/dead-letter` and the other `/queues/*` routes | +| | `GET /admin/identities`, `GET /admin/keys`, `POST /admin/blobs/reclaim`, and every other `/admin/*` route | + +This is **intentional**, not a bug: a blob-reclaim scan +(`POST /admin/blobs/reclaim`) walks the graph to decide which blobs are +still referenced, and a mid-dedup or mid-backfill graph would misclassify +live blobs as orphaned. Every other `/admin/*` route and all data-plane +routes are unsafe against a schema-unsafe or mid-repair graph for the +same underlying reason described in [Why this exists](#why-this-exists). + +The middleware is registered on `app` itself (`main.py:466`), **not** the +auth-wrapped `asgi_app` — so it cannot be bypassed by the bare +`main:app` entrypoint. `BearerTokenMiddleware` wraps `app`, so +authentication still runs first; an unauthenticated request 401s before +it ever reaches this gate. A startup assertion, +`_assert_maintenance_endpoint_allow_listed()` (`main.py:542-559`), fails +loud if `/admin/maintenance`, `/status`, or `/version` are ever missing +from the allow-list — this prevents `/admin/maintenance` from ever +502-ing itself out of existence. + +## Log events to grep for + +| Event | Level | Emitted by | Fields | +|---|---|---|---| +| `maintenance_entered` | INFO | `maintenance.py` (`_handle_transition`) | `reason`, `run_id`, `trigger` (`"op"` or `"constraint"`) | +| `maintenance_completed` | INFO | `maintenance.py` (`_handle_transition`) | `reason`, `run_id`, `duration_seconds` | +| `maintenance_probe_failed` | WARNING | `maintenance.py` (`_run_probe`) | `error` — the live probe raised | +| `maintenance_finish_op_run_id_mismatch` | WARNING | `maintenance.py` (`finish_op`) | `expected`, `got` — a stale/foreign completion signal was ignored | +| `maintenance_quiesce` | INFO | `maintenance_ops.py` (`run_maintenance_operation`) | `run_id`, `seconds` | +| `maintenance_op_succeeded` | INFO | `maintenance_ops.py` | `run_id`, `records_affected` | +| `maintenance_op_failed` | ERROR (exception) | `maintenance_ops.py` | `run_id` (+ traceback) | +| `schema_degraded` | ERROR | `main.py` (`_record_schema_health`, boot only) | reason string | +| `startup_degraded` | ERROR (exception) | `main.py` (`lifespan`, boot only) | the exception that made the whole startup boundary degrade | + +`_handle_transition` logs each open↔closed transition **exactly once**: +it runs with no `await` in its body, so two concurrent callers observing +the same transition can never both log it. + +The last three of the `maintenance_ops.py` events +(`maintenance_quiesce` / `maintenance_op_succeeded` / `maintenance_op_failed`) +fire on every `POST /admin/maintenance` call, since that endpoint schedules +`run_maintenance_operation()` as its background task (see below). They do +**not** fire for `migrations/run.py --apply` or `doctor --fix` — both call +`neo4j_store.run_repair()` directly and never go through +`run_maintenance_operation` (or the coordinator's `try_begin_op`/`finish_op` +bookkeeping). Their effect is still observable via the constraint-probe +transition lines (`maintenance_entered` / `maintenance_completed`) once the +next live probe notices the restored constraint. + +## Clearing maintenance + +The coordinator self-clears the moment its live probe sees the constraint +restored — no restart required, on any of the channels below — but +something has to actually restore it first. Two supported channels exist +today, both ultimately calling the same `neo4j_store.run_repair` (dedup → +`:Node` backfill → constraint create): + +- **`POST /admin/maintenance`** — the network-reachable channel (e.g. + cloud/ACA deployments where the private Neo4j is not directly + reachable). Admin-authenticated HTTP call; schedules + `maintenance_ops.run_maintenance_operation()` as a background task, + which wraps `run_repair` with a quiesce sleep and reports the outcome + back through the coordinator (`try_begin_op`/`finish_op`). Poll + `GET /admin/maintenance` for progress. See + [the endpoint contract below](#postget-adminmaintenance--the-live-endpoint-contract) + for the full request/response shape. +- **`migrations/run.py --apply`** — the local/VM/direct-Neo4j channel + (`migrations/run.py`), for operators who can reach Neo4j directly. + Standalone script, run out-of-band from the server process; it calls + `run_repair` directly with no coordinator/HTTP/admin-auth involvement. + `migrations/run.py --status` gives a read-only report first (constraint + presence, untagged/duplicate counts) without writing anything. + +Both channels run the identical repair logic — which one to use is purely +a matter of which one a given deployment can reach. `doctor --fix` +(`context-intelligence-server doctor --fix`) is an older, still-available +entry point to the same `run_repair` function (`doctor.py`); it remains an +equivalent low-level tool, but it is not the primary lever documented here +for clearing a *live* maintenance window — it predates this endpoint and +was originally documented (see the [README](../README.md)) for the +cold-start-refuses-to-boot case. + +Once `run_repair` successfully re-creates the +`node_node_id_workspace_unique` constraint — via any of the three tools +above — the **next** probe (at most `maintenance_probe_ttl_seconds` later, +whether triggered by an ingest drainer's gate check, a `GET /status` poll, +or a `GET /admin/maintenance` poll) observes `constraint_present=True` and +`mode` flips back to `healthy`/`degraded` on its own. No server restart is +required. + +If the gate is closed because an **op is running** (`op.state == +"running"`) rather than because the constraint is absent, that sub-state +clears only via `coordinator.finish_op(...)`, which +`run_maintenance_operation` calls automatically once the background task +completes (success or failure) — see the endpoint contract below. +`migrations/run.py --apply` and `doctor --fix` never go through the +coordinator, so they can neither open nor close the "op running" +sub-state; they only affect the constraint-presence sub-state. + +## `POST`/`GET /admin/maintenance` — the live endpoint contract + +Both routes are defined in `context_intelligence_server/routers/admin.py` +and are on `maintenance.MAINTENANCE_ALLOW_LIST`, so they stay reachable +even while the gate is closed (enforced structurally at startup by +`main._assert_maintenance_endpoint_allow_listed()` — otherwise the +endpoint would 503 at exactly the moment it exists to unblock). + +- **`POST /admin/maintenance`** — admin-authenticated (`require_admin`, + same 401/403 matrix as every other `/admin/*` route). Single-flight: + `coordinator.try_begin_op()` is a synchronous compare-and-swap with no + `await` between check and set, so two concurrent POSTs can never both + start an op. On winning the CAS, it schedules + `maintenance_ops.run_maintenance_operation()` as a background + `asyncio.Task` (retained via `coordinator.retain_task(...)` so it can't + be garbage-collected mid-run) and returns **promptly** — it does not + await the op's duration (quiesce sleep + the O(graph-size) dedup pass): + - **202** on winning the CAS: `{"run_id": "...", "state": "running", "started_at": "..."}` + - **409** if an op is already running: `{"detail": "maintenance operation already running", "run_id": "", "state": "running"}` + - Re-running on an already-clean graph is a genuine re-scan, not a + short-circuit: each `POST` performs a real `run_repair` call and gets + a fresh `run_id`/`completed_at`, even when `records_affected` comes + back `0`. +- **`GET /admin/maintenance`** — admin-authenticated, same matrix. Returns: + + ```json + { + "mode": "maintenance", + "state": "running", + "run_id": "3f9e...", + "started_at": "2026-08-13T12:00:00.123456+00:00", + "completed_at": null, + "elapsed_seconds": 4.2, + "records_affected": null, + "error": null + } + ``` + + | Field | Meaning | + |---|---| + | `mode` | The coordinator's current mode (`"healthy"` \| `"degraded"` \| `"unknown"` \| `"maintenance"`) — the same value `/status` reports | + | `state` | The op record's state: `"unknown"` \| `"running"` \| `"succeeded"` \| `"failed"`. `"unknown"` with every other field `null` means no op has run yet in this process | + | `run_id` | The op's id (uuid4 hex), `null` if none has run | + | `started_at` / `completed_at` | ISO-8601 UTC timestamps for the op; `completed_at` is `null` while `state == "running"` | + | `elapsed_seconds` | Seconds since the **current maintenance window** opened (the same value as `/status`'s `maintenance_elapsed_seconds`) — not strictly the op's own runtime; `null` when not in maintenance | + | `records_affected` | `duplicates_removed + nodes_tagged` from the last completed `run_repair` call, `null` until one completes | + | `error` | Human-readable exception string if the last op failed, else `null` | + +Neither route is exposed in `/openapi.json` or `/docs` (the whole +`/admin` router is mounted with `include_in_schema=False`) — this affects +schema visibility only, not routing or auth. + +## Do not wire this to a liveness/readiness probe + +The same written prohibition that applies to boot-time `schema_health` +(see the [Azure deployment guide](azure-deployment.md)) applies here, +verbatim, per the code comments in both `main.py` (`_record_schema_health` +docstring and the `/status` handler) and `maintenance.py`: + +> `schema_health`/`mode` MUST NOT be wired to a Kubernetes/ACA liveness or +> readiness probe. Doing so would recreate the exact crash-loop the +> deploy-safe-boot fix removes, one layer up. + +`mode`, `schema_health`, `maintenance_started_at`, and +`maintenance_elapsed_seconds` are **data-migration signals** for an +operator or automation to *read* — not a gate to wire into container +orchestration health checks. `/status` (and `/version`) stay reachable +and return `200` throughout a maintenance window specifically so a plain +HTTP-200 liveness check against either of those paths keeps working +un-interrupted; use that for liveness if one is needed, and treat the +maintenance fields as a separate, human/automation-facing signal. 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/migrations/__init__.py b/migrations/__init__.py new file mode 100644 index 00000000..cda7f6c8 --- /dev/null +++ b/migrations/__init__.py @@ -0,0 +1,6 @@ +"""Out-of-band graph rectification tools. + +Nothing in this package is imported by the server at startup or on the +request path -- see ``run.py`` for the CORE PRINCIPLE this package exists to +honor (migrations run OUT-OF-BAND, never in the server's critical path). +""" diff --git a/migrations/manifest.yaml b/migrations/manifest.yaml new file mode 100644 index 00000000..d6605f7e --- /dev/null +++ b/migrations/manifest.yaml @@ -0,0 +1,29 @@ +# Machine-readable upgrade/migration manifest. +# +# The lean upgrade mechanism: one ordered entry per server version that requires +# out-of-band action on the graph. No rollback machinery -- these are +# forward-only structural rectifications, not reversible schema migrations. +# +# Consumed by: operators, and Amplifier deployment automation deciding +# whether an upgrade needs an out-of-band step before clients resume. +entries: + - id: "6.7.3-maintenance-mode" + server_version: "6.7.3" + schema_version: 1 + schema_affecting: false + summary: > + Adds maintenance mode (gates ingest + query while the :Node uniqueness + constraint is absent, structured 503 + Retry-After, live re-probe + self-clear, /admin/maintenance execution channel). Does not change any + stored node/edge shape -- schema_version stays 1. A deployment with + pre-existing un-migrated/duplicate nodes will boot into maintenance + mode under 6.7.3 (previously it booted degraded but kept accepting + writes); that graph must be rectified once, out-of-band. + scripts: + - "migrations/run.py --apply" + gating: > + If GET /status reports mode=maintenance or mode=degraded, run the + rectification before clients resume sending events. A healthy graph + (schema_health=healthy) needs no action -- boots normally under 6.8.0 + with zero behavior change. + verify: "GET /status shows schema_health=healthy (equivalently mode=healthy)" diff --git a/migrations/run.py b/migrations/run.py new file mode 100644 index 00000000..7d1a812d --- /dev/null +++ b/migrations/run.py @@ -0,0 +1,271 @@ +"""Standalone, OUT-OF-BAND graph rectification tool. + +``migrations/run.py`` is the local/VM/direct-Neo4j sibling of +``POST /admin/maintenance`` (the network-reachable channel for cloud +deployments). Both call the SAME underlying repair logic +(``neo4j_store.run_repair`` / ``diagnose``) -- this script adds no new +algorithm, only a CLI around the existing functions. + +WHAT THIS DOES + Rectifies an already-degraded graph: dedups duplicate legacy nodes, + backfills the universal ``:Node`` label, and (re-)creates the + ``:Node`` uniqueness constraint that the running server's maintenance + gate (``context_intelligence_server/maintenance.py``) probes for. + This is NOT a schema-version migration -- no stored node/edge shape + changes, so ``SCHEMA_VERSION`` (``status.py``) stays ``1``. + It is structural rectification of pre-existing degraded/un-migrated + data, made necessary by upgrading the server past 6.7.x. + +WHY IT EXISTS -- OUT-OF-BAND, NEVER AT STARTUP + CORE PRINCIPLE (AGENTS.md, workspace root): migrations run OUT-OF-BAND + against a live Neo4j instance -- never inside the server's + startup/critical path. This script is NOT imported by the server and + is NEVER invoked from ``lifespan()`` or any request handler. Run it by + hand (local/VM) or trigger the equivalent in-server channel, + ``POST /admin/maintenance`` (cloud, where a human cannot reach the + private Neo4j directly) -- see the README "Upgrading" section. + +IDEMPOTENT + ``--status`` never writes. ``--apply`` calls ``run_repair``, which is + idempotent (dedup -> :Node backfill -> constraint create, all + ``IF NOT EXISTS`` / MERGE-based) -- re-running against an + already-clean graph is a safe no-op. + +USAGE + python migrations/run.py --status + uv run python migrations/run.py --apply + python -m migrations.run --status + + Connection is resolved the same way the server resolves it: via + ``config.Settings`` (``server-config.yaml`` then + ``AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_*`` env vars), using the + SAME driver constructor the server and ``doctor`` CLI use + (``main.build_neo4j_driver``) so this tool can never connect + differently than the process it is rectifying data for. Pass + ``--neo4j-url`` / ``--neo4j-user`` / ``--neo4j-password`` to override + any of the three independently (e.g. pointing at a different Neo4j + than the local ``server-config.yaml`` without editing it). +""" + +from __future__ import annotations + +import argparse +import asyncio +import sys + +from context_intelligence_server.config import Neo4jClientConfig, get_settings +from context_intelligence_server.main import build_neo4j_driver +from context_intelligence_server.maintenance import _CONSTRAINT_NAME, _PROBE_CYPHER +from context_intelligence_server.neo4j_store import diagnose, run_repair +from context_intelligence_server.status import SCHEMA_VERSION, SERVER_VERSION + +# Self-declared from -> to. This is a server-version / structural- +# rectification step, NOT a schema_version bump -- no stored node/edge +# shape changes, so schema_version stays 1 -> 1. +FROM_SERVER_VERSION = "6.7.x" +TO_SERVER_VERSION = "6.8.0" +FROM_SCHEMA_VERSION = 1 +TO_SCHEMA_VERSION = 1 + +_OK = "\033[32m\u2713\033[0m" # green check +_FAIL = "\033[31m\u2717\033[0m" # red x +_WARN = "\033[33m!\033[0m" # yellow warning + + +def _print_banner() -> None: + print("=" * 72) + print("context-intelligence-server -- out-of-band graph rectification") + print(f" server version : {FROM_SERVER_VERSION} -> {TO_SERVER_VERSION}") + print( + f" schema_version : {FROM_SCHEMA_VERSION} -> {TO_SCHEMA_VERSION} (unchanged -- structural only)" + ) + print( + f" running server reports: version={SERVER_VERSION} schema_version={SCHEMA_VERSION}" + ) + print(" OUT-OF-BAND: never runs at server startup or on the request path.") + print(" IDEMPOTENT: safe to re-run; --apply is a no-op on an already-clean graph.") + print("=" * 72) + + +def _resolve_config(args: argparse.Namespace) -> Neo4jClientConfig: + """Same resolution the server uses (``Settings``), with CLI overrides. + + Base config comes from ``get_settings().resolve_neo4j_admin()`` -- the + identical call ``doctor.py`` and ``lifespan()`` make. ``--neo4j-*`` + flags override individual fields without requiring a + ``server-config.yaml`` edit. + """ + admin = get_settings().resolve_neo4j_admin() + return admin.model_copy( + update={ + k: v + for k, v in ( + ("url", args.neo4j_url), + ("username", args.neo4j_user), + ("password", args.neo4j_password), + ) + if v is not None + } + ) + + +async def _constraint_present(driver: object) -> bool | None: + """Tri-state ``:Node`` uniqueness constraint check. + + Reuses the exact catalog-read query the running server's maintenance + gate probes with (``maintenance._PROBE_CYPHER`` / ``_CONSTRAINT_NAME``) + -- one source of truth for "is the constraint present", never a + second hardcoded copy of the constraint name. + """ + try: + async with driver.session() as session: # type: ignore[attr-defined] + result = await session.run(_PROBE_CYPHER) + count = 0 + async for record in result: + count = record["c"] + return count > 0 + except Exception: # noqa: BLE001 -- connectivity/catalog probe, report unknown + return None + + +def _print_diagnosis( + diagnosis: dict[str, int], constraint_present: bool | None +) -> None: + mark = ( + _OK if constraint_present else (_WARN if constraint_present is False else _FAIL) + ) + label = {True: "present", False: "ABSENT", None: "unknown (probe failed)"}[ + constraint_present + ] + print(f" {mark} :Node uniqueness constraint ({_CONSTRAINT_NAME}): {label}") + + untagged = diagnosis["untagged_nodes"] + mark = _OK if untagged == 0 else _WARN + print(f" {mark} Untagged :Node count: {untagged}") + + duplicates = diagnosis["duplicate_nodes"] + mark = _OK if duplicates == 0 else _WARN + print(f" {mark} Duplicate node count: {duplicates}") + + +async def _run_status(driver: object) -> int: + """Read-only report. Never writes. Exits 0 once Neo4j is reachable and + the report has been printed, regardless of whether rectification is + needed -- ``--status`` reports, it does not judge. + """ + try: + await driver.verify_connectivity() # type: ignore[attr-defined] + except Exception as exc: # noqa: BLE001 + print(f" {_FAIL} Neo4j reachable -- {exc}") + return 1 + print(f" {_OK} Neo4j reachable") + + constraint_present = await _constraint_present(driver) + diagnosis = await diagnose(driver) + _print_diagnosis(diagnosis, constraint_present) + + needs_rectification = ( + constraint_present is not True + or diagnosis["untagged_nodes"] + or diagnosis["duplicate_nodes"] + ) + if needs_rectification: + print( + f" {_WARN} Rectification needed -- re-run with --apply (or POST /admin/maintenance)." + ) + else: + print(f" {_OK} Graph is healthy -- no rectification needed.") + return 0 + + +async def _run_apply(driver: object) -> int: + """Rectify: dedup -> :Node backfill -> constraint create, via the + SAME ``run_repair`` the server's ``/admin/maintenance`` endpoint and + ``doctor --fix`` call. Prints before/after counts. + """ + try: + await driver.verify_connectivity() # type: ignore[attr-defined] + except Exception as exc: # noqa: BLE001 + print(f" {_FAIL} Neo4j reachable -- {exc}") + return 1 + print(f" {_OK} Neo4j reachable") + + print("-- before --") + before = await diagnose(driver) + _print_diagnosis(before, await _constraint_present(driver)) + + print("Rectifying (dedup + :Node backfill + constraint create)...") + result = await run_repair(driver) + print( + f" {_OK} {result['duplicates_removed']} duplicate(s) removed, {result['nodes_tagged']} node(s) tagged :Node." + ) + + print("-- after --") + after = await diagnose(driver) + constraint_present = await _constraint_present(driver) + _print_diagnosis(after, constraint_present) + + if ( + constraint_present is True + and after["untagged_nodes"] == 0 + and after["duplicate_nodes"] == 0 + ): + print(f" {_OK} Graph is healthy after rectification.") + return 0 + print(f" {_FAIL} Graph still has issues after rectification -- see counts above.") + return 1 + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="migrations/run.py", + description="Out-of-band Neo4j graph rectification (never at server startup).", + ) + mode = parser.add_mutually_exclusive_group(required=True) + mode.add_argument( + "--status", + action="store_true", + help="Read-only report (constraint presence, untagged/duplicate counts). Exit 0.", + ) + mode.add_argument( + "--apply", + action="store_true", + help="Rectify: dedup + :Node backfill + constraint create. Idempotent.", + ) + parser.add_argument( + "--neo4j-url", + default=None, + help="Override the Neo4j bolt URL (default: from Settings).", + ) + parser.add_argument( + "--neo4j-user", + default=None, + help="Override the Neo4j username (default: from Settings).", + ) + parser.add_argument( + "--neo4j-password", + default=None, + help="Override the Neo4j password (default: from Settings).", + ) + return parser + + +async def _amain(args: argparse.Namespace) -> int: + config = _resolve_config(args) + driver = build_neo4j_driver(config) + try: + if args.apply: + return await _run_apply(driver) + return await _run_status(driver) + finally: + await driver.close() + + +def main(argv: list[str] | None = None) -> int: + args = build_parser().parse_args(argv) + _print_banner() + return asyncio.run(_amain(args)) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/pyproject.toml b/pyproject.toml index 2aaa7caf..0e032375 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.3" description = "Context Intelligence Server for Amplifier" requires-python = ">=3.11" dependencies = [ diff --git a/scripts/relabel_incomplete_sessions.py b/scripts/relabel_incomplete_sessions.py new file mode 100644 index 00000000..a6dae673 --- /dev/null +++ b/scripts/relabel_incomplete_sessions.py @@ -0,0 +1,667 @@ +#!/usr/bin/env python +"""Maintenance script: one-off backfill removing the stale IncompleteSession +false-positive marker from historical Session nodes (the one-off backfill +half of the IncompleteSession relabel fix). + +NOT product code -- integration-tested in +tests/neo4j/test_relabel_incomplete_sessions.py. Never execute this file as +part of the normal application lifecycle. + +Background +---------- +:IncompleteSession is stamped at session:end only when a Session node has no +type label yet. Forked sub-sessions drain in independent, concurrently +draining per-session queues, so a child's session:end can be processed +BEFORE its session:fork/session:start -- stamping the marker before the real +terminal type arrives. The heal-forward classify() fix +(SessionLabelStateMachine.classify(), see handlers/data_layer_2/session.py) +heals this FORWARD: every start/fork transition now strips a stale +IncompleteSession marker the moment it is processed. That fix handles +new/future events. This script is the one-off backfill for nodes that were +ALREADY mislabeled before the classify() fix shipped and will never see +another start/fork event to heal them forward. + +POST-DEPLOY GATE +---------------- +Run this script ONLY after the heal-forward classify() fix is deployed AND +verified live. Running --apply before that fix is deployed races fresh +mislabeling: new out-of-order nodes can still be stamped IncompleteSession by +the old code while this script is removing the stale marker from historical +nodes, so a subsequent read could see the marker reappear on a node this +script "fixed." The heal-forward fix must be live first so the false-positive +population stops growing before this one-off cleans up the existing backlog. + +Selector direction correction (found while writing the Neo4j integration +tests for this script -- NOT what an earlier draft's literal Cypher assumed) +-------------------------------------------------------------------------- +An earlier draft's Cypher fragments write the linked-event check as +``(s)<-[:SOURCED_FROM]-(:SessionStartEvent)`` -- i.e. SessionStartEvent as +the edge SOURCE, Session as the TARGET. Verifying against the shipping +write path shows the real graph is the other way around: + +* ``SessionHandler._handle_start``/``_handle_fork`` + (handlers/data_layer_2/session.py) call + ``graph.upsert_edge(session_id, data_layer_1_node_id, {"type": "SOURCED_FROM"})`` + -- ``session_id`` is the FIRST (src) argument. +* ``Neo4jGraphStore._edge_merge_cypher`` (neo4j_store.py) emits + ``MERGE (src)-[r:{edge_type}]->(dst)`` with src/dst bound to the first/second + ``upsert_edge`` arguments respectively. +* Every other SOURCED_FROM bridge in the codebase (ToolCall -> ToolPreEvent, + Prompt -> PromptSubmitEvent, etc. -- see tests/handlers/data_layer_1/test_default.py + and tests/integration/test_event_pipeline.py) follows the same convention: + the domain entity is the edge SOURCE, the data_layer_1 event node is the + TARGET. + +So the real graph shape is ``(Session)-[:SOURCED_FROM]->(SessionStartEvent)``, +NOT the reverse. This script's selectors use the code-verified direction +``(s)-[:SOURCED_FROM]->(:SessionStartEvent)`` / +``(s)-[:SOURCED_FROM]->(:SessionForkEvent)`` throughout. A selector built on +the earlier draft's literal (reversed) arrow would never match a single real +node and would silently do nothing. + +Concrete execution shape +------------------------------- +A single-statement, server-side batched REMOVE using +``CALL { ... } IN TRANSACTIONS OF $batch_size ROWS`` (the pattern proven in +scripts/tag_legacy_pooled_iterations.py -- NOT the per-node Python loop of +scripts/repair_dual_labels.py). The false-positive selector +(_FALSE_POSITIVE_MATCH below) is the ONE module constant reused by the +diagnostic, count, sample, collect, and apply queries -- one definition of +"false positive." ``apply_relabel()`` additionally collects the node_ids it +actually touched (for the printed "removed N node(s)" report) via ``RETURN +collect(s.node_id)`` after the batched CALL block -- the same "outer +variable survives the per-batch commits" technique +``tag_legacy_pooled_iterations.py`` uses for its ``RETURN count(i)``, just +projecting ids instead of a count. This is a REPORTING count only: the +undo-log (below) is sourced from a separate, EARLIER, read-only collection -- +see below for why. + +Read-only reconciliation diagnostic, REQUIRED before --apply +-------------------------------------------------------------------- +The selector's core assumption -- "a linked SessionStartEvent/SessionForkEvent +implies the terminal label is already set" -- holds for today's handler code +but is UNVERIFIED for historical nodes written before this fix existed. So +--apply is hard-gated on a read-only diagnostic (``diagnostic_report()``) +that this script always runs first: + +* ``linked_but_untyped`` -- IncompleteSession nodes with a linked start/fork + event but NO terminal label. This is the assumption's blind spot. If this + is > 0, the assumption is FALSIFIED on this DB and --apply REFUSES (exit 1, + no write), printing the count and up to 20 sample node_ids. The operator + must reconcile (or narrow the selector to the terminal-label clause only) + before re-running --apply. +* ``typed_but_unlinked`` -- IncompleteSession nodes with a terminal label but + no linked start/fork event at all. Informational only; does not gate. + +(The POST-DEPLOY GATE banner above is the third hardening measure this script +applies.) + +Lightweight rollback +--------------------------- +--apply writes a plain JSON undo-log (--undo-log PATH, default a timestamped +path in the current directory) BEFORE it mutates the graph. The log is +sourced from ``collect_false_positive_ids()`` -- a read-only, pre-mutation +collection of the FULL false-positive candidate set (same _FALSE_POSITIVE_MATCH +selector as everything else) -- NOT from the ids ``apply_relabel()`` reports +as touched. This ordering matters: ``apply_relabel()``'s batched +``CALL { ... } IN TRANSACTIONS OF N ROWS`` COMMITS PER BATCH, so a mid-run +crash can leave some REMOVEs already durable in Neo4j. Writing the undo-log +from the pre-mutation candidate set means that even a crash after batch 1 of +N still leaves a COMPLETE undo-log on disk -- there is no auditability +window where a committed removal has no undo record anywhere. See +``run_apply()`` for the inline safety argument for why logging the (superset) +candidate set instead of the (subset) touched set is safe. The file's +header records the Neo4j host and an ISO-8601 generation timestamp. +--restore PATH re-adds :IncompleteSession to exactly those node_ids via the +same batched ``CALL { ... } IN TRANSACTIONS OF N ROWS`` shape. No +edge/temporal snapshot is needed: this is one non-destructive label REMOVE, +fully recomputable from the node_id list alone. + +Acceptance check +------------------------ +Both --dry-run and --apply print a population summary before (and --apply +also after) any write: total :IncompleteSession count, and the +genuine-incomplete count (:IncompleteSession AND NOT linked to any +start/fork event at all -- the true, un-mislabeled data-loss population this +script must never touch). This lets the operator see the population move. + +"Only once" guarantees (both hold, no version machinery) +---------------------------------------------------------- +1. Out-of-band -- standalone script, run manually post-upgrade; NEVER at + server startup, and NOT triggered by any product code path. +2. Idempotent -- the false-positive selector's WHERE clause leads with + ``s:IncompleteSession``, so once a node's marker is removed it no longer + matches; re-running --apply against an already-healed graph matches (and + touches) zero rows. Exactly-once degrades safely to at-least-once. + +Leaves untouched the genuine ~0.5% (no linked start/fork event at all -- +these carry no evidence that the labeling race, rather than a real data-loss +event, produced the marker). + +SCHEMA_VERSION is unchanged (this is a data backfill, not a schema change). + +Modes +----- +--dry-run (default) + Read-only report: population summary, the reconciliation diagnostic + (informational here), and the count + a bounded sample (first 20 + node_ids) of what --apply would touch. Writes NOTHING. Always exits 0 + (health-check, not a gate). + +--apply + Runs the reconciliation diagnostic first. If ``linked_but_untyped`` > 0, + REFUSES (exit 1, no write). Otherwise runs the batched REMOVE + (``apply_relabel``), writes the undo-log, and prints the before/after + population summary. + +--restore PATH + Reads an undo-log written by a previous --apply and re-adds + :IncompleteSession to exactly those node_ids. + +Connection +---------- +Connection details come from :func:`context_intelligence_server.config.get_settings`. +Resolution order (highest first): + +1. Environment variables with prefix ``AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_``. +2. ``server-config.yaml`` in the current working directory. +3. Built-in defaults. + +--neo4j-url/--neo4j-user/--neo4j-password override the resolved settings. +This script is standalone: it connects to Neo4j directly via the ``neo4j`` +driver and does NOT import or start the FastAPI application. + +Exit codes +---------- +* 0 -- --dry-run (always); --apply completed (whether or not any node + needed healing); --restore completed. +* 1 -- --apply refused by the reconciliation-diagnostic gate + (``linked_but_untyped`` > 0); or --restore could not read the given + undo-log file. +""" + +from __future__ import annotations + +import argparse +import json +import sys +from datetime import datetime, timezone +from typing import Any +from urllib.parse import urlparse + +from context_intelligence_server.config import get_settings +from neo4j import GraphDatabase + +DEFAULT_BATCH_SIZE = 500 +DEFAULT_SAMPLE_LIMIT = 20 + +# --------------------------------------------------------------------------- +# Selector fragments -- the ONE place each predicate is defined. Every +# diagnostic / count / sample / apply query below is built from these same +# fragments, so there is exactly one definition of "has a terminal label" and +# "has a linked start/fork event." Both fragments are only ever embedded +# inside a WHERE clause (standard Cypher pattern-predicate usage) -- never +# used as a bare expression inside RETURN/CASE. +# --------------------------------------------------------------------------- + +# A Session node carries a real terminal type. +_HAS_TERMINAL_LABEL = "(s:RootSession OR s:SubSession OR s:ForkedSession)" + +# A Session node has its own linked SessionStartEvent/SessionForkEvent. +# Direction verified against the shipping write path -- see "Selector +# direction correction" in the module docstring above. +_HAS_LINKED_START_OR_FORK = ( + "( (s)-[:SOURCED_FROM]->(:SessionStartEvent) " + "OR (s)-[:SOURCED_FROM]->(:SessionForkEvent) )" +) + +# The ONE false-positive selector. An IncompleteSession node is a +# false positive if it already carries a real terminal label, OR has its own +# linked start/fork event. This is the exact set --apply removes the stale +# marker from. ``WHERE s:IncompleteSession`` leads the predicate, which is +# also what makes the operation idempotent: once removed, a node no longer +# matches (see "Only once guarantees" above). +_FALSE_POSITIVE_MATCH = ( + "MATCH (s:Session) WHERE s:IncompleteSession " + f"AND ( {_HAS_TERMINAL_LABEL} OR {_HAS_LINKED_START_OR_FORK} )" +) + + +# --------------------------------------------------------------------------- +# Read-only helpers +# --------------------------------------------------------------------------- + + +def population_summary(session) -> dict[str, int]: + """Acceptance check: total IncompleteSession count + genuine-incomplete count. + + genuine_incomplete = IncompleteSession AND NOT linked to any + SessionStartEvent/SessionForkEvent at all -- the true data-loss + population this fix must never touch. + """ + total = session.run( + "MATCH (s:Session) WHERE s:IncompleteSession RETURN count(s) AS n" + ).single()["n"] + + genuine = session.run( + "MATCH (s:Session) WHERE s:IncompleteSession " + f"AND NOT {_HAS_LINKED_START_OR_FORK} " + "RETURN count(s) AS n" + ).single()["n"] + + return {"total": total or 0, "genuine_incomplete": genuine or 0} + + +def diagnostic_report(session) -> dict[str, Any]: + """Read-only reconciliation diagnostic. No writes. + + Returns: + { + "linked_but_untyped": int, + "linked_but_untyped_samples": list[str], # up to 20 node_ids + "typed_but_unlinked": int, + } + + linked_but_untyped is the assumption's blind spot: an IncompleteSession + node with a linked start/fork event but NO terminal label yet. If this + is > 0 anywhere in the graph, the selector's core assumption is + FALSIFIED on this DB and --apply must refuse. + + typed_but_unlinked is informational only (never gates --apply): an + IncompleteSession node with a terminal label but no linked start/fork + event at all. + """ + linked_but_untyped_rows = session.run( + "MATCH (s:Session) WHERE s:IncompleteSession " + f"AND {_HAS_LINKED_START_OR_FORK} AND NOT {_HAS_TERMINAL_LABEL} " + "RETURN s.node_id AS node_id" + ) + linked_but_untyped_ids = [row["node_id"] for row in linked_but_untyped_rows] + + typed_but_unlinked = ( + session.run( + "MATCH (s:Session) WHERE s:IncompleteSession " + f"AND {_HAS_TERMINAL_LABEL} AND NOT {_HAS_LINKED_START_OR_FORK} " + "RETURN count(s) AS n" + ).single()["n"] + or 0 + ) + + return { + "linked_but_untyped": len(linked_but_untyped_ids), + "linked_but_untyped_samples": linked_but_untyped_ids[:DEFAULT_SAMPLE_LIMIT], + "typed_but_unlinked": typed_but_unlinked, + } + + +def count_false_positives(session) -> int: + """Count of IncompleteSession nodes the false-positive selector matches + (i.e. how many --apply would touch).""" + return ( + session.run(f"{_FALSE_POSITIVE_MATCH} RETURN count(s) AS n").single()["n"] or 0 + ) + + +def sample_false_positives(session, limit: int = DEFAULT_SAMPLE_LIMIT) -> list[str]: + """Bounded sample of node_ids the false-positive selector matches. + + ``limit`` is always an internal int constant (never user-controlled + text), so it is safe to interpolate directly into the LIMIT clause. + """ + rows = session.run( + f"{_FALSE_POSITIVE_MATCH} RETURN s.node_id AS node_id LIMIT {limit}" + ) + return [row["node_id"] for row in rows] + + +def collect_false_positive_ids(session) -> list[str]: + """Read-only collection of the FULL (unbounded) set of node_ids the + false-positive selector matches -- i.e. every candidate --apply would + touch. + + This is the pre-mutation read ``run_apply`` uses to source the + undo-log BEFORE calling ``apply_relabel`` (see "undo-log-before-mutation" + in ``run_apply``'s docstring for the ordering rationale). Unlike + ``sample_false_positives`` (bounded by ``DEFAULT_SAMPLE_LIMIT``, for + human-readable reporting), this has no LIMIT -- it must capture every + candidate so the undo-log is complete even if a subsequent crash + prevents ``apply_relabel`` from reporting its own touched set. + """ + result = session.run(f"{_FALSE_POSITIVE_MATCH} RETURN collect(s.node_id) AS ids") + return list(result.single()["ids"] or []) + + +# --------------------------------------------------------------------------- +# Mutating operations +# --------------------------------------------------------------------------- + + +def apply_relabel(session, batch_size: int = DEFAULT_BATCH_SIZE) -> list[str]: + """The single-statement, server-side batched REMOVE. + + Strips :IncompleteSession from every node matched by + _FALSE_POSITIVE_MATCH, batched via + ``CALL { ... } IN TRANSACTIONS OF $batch_size ROWS`` (the pattern in + scripts/tag_legacy_pooled_iterations.py). Returns the node_ids actually + touched -- needed for the undo-log. + + Idempotent: because the outer MATCH re-requires ``s:IncompleteSession``, + a node healed by a previous run no longer matches; re-running this + function against an already-healed graph returns []. + + This function performs NO gate check. Callers that must honor the + reconciliation gate use ``run_apply()``, not this function directly. + """ + result = session.run( + f"{_FALSE_POSITIVE_MATCH} " + "CALL { WITH s REMOVE s:IncompleteSession } IN TRANSACTIONS OF $batch_size ROWS " + "RETURN collect(s.node_id) AS touched_ids", + batch_size=batch_size, + ) + return list(result.single()["touched_ids"] or []) + + +def restore_ids( + session, node_ids: list[str], batch_size: int = DEFAULT_BATCH_SIZE +) -> int: + """--restore: re-add :IncompleteSession to exactly the given node_ids. + + Uses the same batched ``CALL { ... } IN TRANSACTIONS OF N ROWS`` shape as + apply_relabel. Returns the number of nodes actually restored (node_ids + that still exist in the graph). + """ + result = session.run( + "UNWIND $node_ids AS nid " + "MATCH (s:Session {node_id: nid}) " + "CALL { WITH s SET s:IncompleteSession } IN TRANSACTIONS OF $batch_size ROWS " + "RETURN count(s) AS restored", + node_ids=node_ids, + batch_size=batch_size, + ) + return result.single()["restored"] or 0 + + +# --------------------------------------------------------------------------- +# Undo-log +# --------------------------------------------------------------------------- + + +def _neo4j_host(url: str) -> str: + """Return a host-only fragment of a Neo4j URL for the undo-log header.""" + return urlparse(url).hostname or url + + +def default_undo_log_path() -> str: + """Default timestamped undo-log path, used when --undo-log is omitted.""" + ts = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + return f"relabel_incomplete_sessions_undo_{ts}.json" + + +def write_undo_log(path: str, neo4j_url: str, node_ids: list[str]) -> None: + """Write an undo-log: header (Neo4j host + ISO timestamp) + touched node_ids.""" + payload = { + "neo4j_host": _neo4j_host(neo4j_url), + "generated_at": datetime.now(timezone.utc).isoformat(), + "node_ids": node_ids, + } + with open(path, "w", encoding="utf-8") as fh: + json.dump(payload, fh, indent=2) + + +def read_undo_log(path: str) -> dict[str, Any]: + """Read an undo-log written by write_undo_log.""" + with open(path, encoding="utf-8") as fh: + return json.load(fh) + + +# --------------------------------------------------------------------------- +# Reporting / orchestration +# --------------------------------------------------------------------------- + + +def run_dry_run(session) -> int: + """Read-only report. Writes nothing. Always returns 0 (health-check, + not a gate).""" + before = population_summary(session) + print("DRY RUN -- relabel_incomplete_sessions.py (no writes)\n") + print("POPULATION SUMMARY:") + print(f" total IncompleteSession : {before['total']}") + print(f" genuine_incomplete (no link) : {before['genuine_incomplete']}\n") + + diag = diagnostic_report(session) + print("RECONCILIATION DIAGNOSTIC (gates --apply; informational here):") + print(f" linked_but_untyped : {diag['linked_but_untyped']}") + print(f" typed_but_unlinked : {diag['typed_but_unlinked']} (informational only)") + if diag["linked_but_untyped"]: + print("\n --apply would REFUSE (see module docstring). Sample node_ids:") + for node_id in diag["linked_but_untyped_samples"]: + print(f" - {node_id}") + print() + + would_touch = count_false_positives(session) + sample = sample_false_positives(session) + print( + f"WOULD APPLY -- {would_touch} node(s) would have :IncompleteSession removed." + ) + if sample: + print(f"Sample node_ids (up to {DEFAULT_SAMPLE_LIMIT}):") + for node_id in sample: + print(f" - {node_id}") + return 0 + + +def run_apply(session, batch_size: int, undo_log_path: str, neo4j_url: str) -> int: + """Gated --apply: diagnostic first, write only if the gate is clear. + + 1. Print the BEFORE population summary. + 2. Run the reconciliation diagnostic. If linked_but_untyped > 0, + REFUSE (no write): print the count and up to 20 sample node_ids, + return 1. + 3. Otherwise: collect the FULL false-positive candidate set (read-only), + write the undo-log from THAT set, THEN run apply_relabel() to + mutate, print the AFTER population summary, return 0. + + undo-log-before-mutation + ----------------------------------- + Step 3 writes the undo-log BEFORE calling apply_relabel(), not after. + apply_relabel()'s batched ``CALL { ... } IN TRANSACTIONS OF N ROWS`` + COMMITS PER BATCH -- each batch's REMOVE is durable in Neo4j the moment + that batch completes, well before apply_relabel() returns. If the + process crashed mid-run under the OLD ordering (mutate first, log + after), some REMOVEs would already be committed with NO undo record + anywhere -- an auditability hole. Logging the pre-mutation candidate set + first closes that window: even a crash after the very first batch still + leaves a COMPLETE undo-log on disk. + + Why logging the (pre-mutation) candidate set instead of apply_relabel's + (post-mutation) touched set is safe: + - candidate_ids is collected via the exact same _FALSE_POSITIVE_MATCH + selector apply_relabel() uses, moments before apply_relabel() runs + -- so candidate_ids is a SUPERSET of (on a clean run, EQUAL to) + whatever apply_relabel() actually removes. + - restore_ids() (the --restore consumer of this log) is idempotent + over a superset: it SETs :IncompleteSession on each node_id that + still exists. For a node that WAS removed, this correctly restores + the marker. For a candidate that -- for any reason -- was NOT + removed (still carries :IncompleteSession), the SET is a no-op: the + label is already present. So restoring from the candidate superset + is always safe -- it never incorrectly re-adds a marker that + shouldn't exist, and it never fails to restore a node that WAS + removed. + - On a clean run (the common case, and the only case this script's + idempotent selector allows in practice) candidate_ids == the ids + apply_relabel() reports touched, so the printed "removed N node(s)" + count is UNCHANGED by this reordering -- only the undo-log's source + and its write timing move earlier. + """ + before = population_summary(session) + print("POPULATION SUMMARY (before):") + print(f" total IncompleteSession : {before['total']}") + print(f" genuine_incomplete (no link) : {before['genuine_incomplete']}\n") + + diag = diagnostic_report(session) + print("RECONCILIATION DIAGNOSTIC:") + print(f" linked_but_untyped : {diag['linked_but_untyped']}") + print(f" typed_but_unlinked : {diag['typed_but_unlinked']} (informational only)\n") + + if diag["linked_but_untyped"] > 0: + print( + "APPLY REFUSED -- the selector's assumption (a linked " + "SessionStartEvent/SessionForkEvent implies the terminal label is " + f"already set) does NOT hold on this DB: {diag['linked_but_untyped']} " + "node(s) are linked but carry no terminal label.\n" + ) + print("Sample node_ids (up to 20):") + for node_id in diag["linked_but_untyped_samples"]: + print(f" - {node_id}") + print( + "\nReconcile these nodes (or narrow the selector to the " + "terminal-label clause only) before re-running --apply." + ) + return 1 + + # Read-only pre-mutation capture, THEN write the undo-log, THEN mutate. + # See the "undo-log-before-mutation" note in this function's + # docstring for why this order (and logging the candidate superset + # rather than apply_relabel's touched subset) is safe. + candidate_ids = collect_false_positive_ids(session) + write_undo_log(undo_log_path, neo4j_url, candidate_ids) + print( + f"Undo log written to {undo_log_path} ({len(candidate_ids)} node_id(s), " + "pre-mutation candidate set).\n" + ) + + touched_ids = apply_relabel(session, batch_size) + print(f"APPLIED -- removed :IncompleteSession from {len(touched_ids)} node(s).\n") + + after = population_summary(session) + print("POPULATION SUMMARY (after):") + print(f" total IncompleteSession : {after['total']}") + print(f" genuine_incomplete (no link) : {after['genuine_incomplete']}") + return 0 + + +def run_restore(session, path: str, batch_size: int) -> int: + """--restore: re-add :IncompleteSession to exactly the ids in *path*.""" + try: + snap = read_undo_log(path) + except (OSError, json.JSONDecodeError) as exc: + print(f"RESTORE FAILED -- could not read undo-log {path}: {exc}") + return 1 + + node_ids = snap.get("node_ids", []) + restored = restore_ids(session, node_ids, batch_size) + print( + f"RESTORED -- re-added :IncompleteSession to {restored}/{len(node_ids)} " + f"node(s) from {path} (generated_at={snap.get('generated_at')})." + ) + return 0 + + +# --------------------------------------------------------------------------- +# Entry point +# --------------------------------------------------------------------------- + + +def main() -> int: + """Parse arguments, connect to Neo4j, and run the requested mode.""" + parser = argparse.ArgumentParser( + prog="relabel_incomplete_sessions.py", + description=( + "One-off maintenance tool (the backfill half of the IncompleteSession " + "relabel fix): remove the stale :IncompleteSession false-positive " + "marker from historical Session nodes that already carry a real " + "terminal label or a linked SessionStartEvent/SessionForkEvent. " + "POST-DEPLOY GATE -- run only after the heal-forward classify() fix " + "is deployed and verified live. --apply is itself gated on a " + "read-only reconciliation diagnostic; see module docstring." + ), + ) + parser.add_argument( + "--dry-run", + action="store_true", + help=( + "Read-only report: population summary, reconciliation diagnostic, and the " + "count + bounded sample of what --apply would touch. Writes " + "nothing. This is the default mode." + ), + ) + parser.add_argument( + "--apply", + action="store_true", + help=( + "Run the reconciliation diagnostic; if clear, remove :IncompleteSession from " + "the false-positive set (batched, idempotent), write an " + "undo-log, and print the before/after population summary." + ), + ) + parser.add_argument( + "--restore", + metavar="PATH", + default=None, + help="Re-add :IncompleteSession to exactly the node_ids in an undo-log file.", + ) + parser.add_argument( + "--undo-log", + metavar="PATH", + default=None, + help=( + "Path to write the undo-log (used with --apply). Defaults to " + "a timestamped path in the current directory." + ), + ) + parser.add_argument( + "--batch-size", + type=int, + default=DEFAULT_BATCH_SIZE, + metavar="N", + help=f"Rows per transaction for --apply/--restore (default: {DEFAULT_BATCH_SIZE}).", + ) + parser.add_argument( + "--neo4j-url", + metavar="URL", + default=None, + help="Neo4j Bolt URL (overrides server-config.yaml / env var)", + ) + parser.add_argument( + "--neo4j-user", + metavar="USER", + default=None, + help="Neo4j username (overrides server-config.yaml / env var)", + ) + parser.add_argument( + "--neo4j-password", + metavar="PW", + default=None, + help="Neo4j password (overrides server-config.yaml / env var)", + ) + args = parser.parse_args() + + modes_selected = sum(bool(m) for m in (args.dry_run, args.apply, args.restore)) + if modes_selected > 1: + parser.error("--dry-run, --apply, and --restore are mutually exclusive") + + settings = get_settings() + neo4j_url = args.neo4j_url or settings.neo4j_url + neo4j_user = args.neo4j_user or settings.neo4j_user + neo4j_password = args.neo4j_password or settings.neo4j_password + + print(f"Connecting to Neo4j at {neo4j_url} as {neo4j_user}\n") + driver = GraphDatabase.driver(neo4j_url, auth=(neo4j_user, neo4j_password)) + try: + with driver.session() as neo4j_session: + if args.restore: + return run_restore(neo4j_session, args.restore, args.batch_size) + if args.apply: + undo_log_path = args.undo_log or default_undo_log_path() + return run_apply( + neo4j_session, args.batch_size, undo_log_path, neo4j_url + ) + # --dry-run is the default: no writes unless --apply is explicit. + return run_dry_run(neo4j_session) + finally: + driver.close() + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/tag_legacy_pooled_iterations.py b/scripts/tag_legacy_pooled_iterations.py new file mode 100644 index 00000000..21be1b01 --- /dev/null +++ b/scripts/tag_legacy_pooled_iterations.py @@ -0,0 +1,375 @@ +#!/usr/bin/env python +"""Maintenance tool: TAG (non-destructive) legacy Iteration nodes that were +confirmed MERGEd across >=2 distinct OrchestratorRuns before the run-scoped-id fix. + +**NOT product code -- no unit tests.** This is a standalone, one-off +maintenance script. The regression that PREVENTS new corruption is covered +by ``tests/handlers/data_layer_2/test_iteration.py`` +(``TestIterationRunScopingP21``); this script only deals with historical data +already in a live graph. + +Background +--------------------------------------- +Before the run-scoped-id fix, ``Iteration.node_id`` was the bare shape +``{session_id}::iteration::{N}``, with ``N`` a per-*session* (not per-*run*) +counter. Because the counter could restart (e.g. drainer restart/replay), +two DIFFERENT ``OrchestratorRun``s could reuse the same ``N`` and MERGE onto +the SAME bare-id Iteration node, clobbering its usage/message properties +(last-write-wins). After the run-scoped-id fix, new Iteration node_ids are run-scoped: +``{session_id}::orch_run::{ts}::iteration::{N}`` -- these can never collide +across runs and are excluded from consideration by this script entirely. + +A **bare-id node_id does NOT mean the node is corrupt.** A design review + +live-data verification established that only nodes whose bare id is reached +by ``HAS_PART`` from **two or more distinct** ``OrchestratorRun`` nodes are +provably corrupt (pooled). The rest of the bare-id population (the large +majority -- ~93-95% of bare-id nodes in live data) are clean, single-run +nodes created before the fix and must NOT be touched: a run=1 bare-id node is +exactly what a healthy pre-fix Iteration looked like. + +What this script does +---------------------- +This script TAGS -- and only tags -- the confirmed-corrupt subset with a +non-destructive marker property ``data_quality = 'legacy_pooled_pre_fix'``. +It never deletes or restructures anything. The destructive cleanup (e.g. +splitting or deleting pooled nodes) is a SEPARATE, gated follow-up and is +explicitly OUT OF SCOPE here. + +CONFIRMED-CORRUPT SELECTOR (the ONLY nodes this script ever touches):: + + MATCH (run:OrchestratorRun)-[:HAS_PART]->(i:Iteration) + WHERE NOT i.node_id CONTAINS '::orch_run::' + WITH i, count(DISTINCT run) AS runs + WHERE runs >= 2 + +Selector-sanity note (does not over-claim) +------------------------------------------- +``runs >= 2`` is a **confirmed lower bound** on corruption, not an exhaustive +enumeration of it. A bare-id node with exactly one surviving ``HAS_PART`` +parent (``runs == 1``) could -- in principle -- have been pooled across MORE +runs whose ``OrchestratorRun`` node or edge was later pruned/expired, leaving +only one parent behind; such a node would be indistinguishable, from current +graph state, from a genuinely single-run node. This script deliberately does +**not** attempt to catch that case: it tags only what is provably corroborated +by graph structure today (>=2 live distinct parents), and explicitly leaves +``runs == 1`` (``confirmed_clean_single_run``) and ``runs == 0`` +(``no_run_edge`` -- no ``HAS_PART`` parent survives at all, so corruption +cannot be confirmed or denied) untouched. + +Nodes this script explicitly does NOT touch +-------------------------------------------- +* ``runs == 1`` -- confirmed clean single-run bare-id node. Leave alone. +* ``runs == 0`` -- no surviving ``OrchestratorRun`` parent edge at all, so + pooling cannot be confirmed. Leave alone (see selector-sanity note above). +* Any run-scoped node (``node_id`` contains ``'::orch_run::'``) -- these are + post-fix and structurally cannot collide across runs. + +DEPLOYMENT GATE +--------------- +This script is **not a deployment gate**. It performs a single additive, +idempotent ``SET`` of one marker property on a provably-corrupt subset of +historical nodes; it does not race the run-scoped-id write-path fix (which +governs only newly-created Iteration nodes) and carries no risk of +mixed-type or mixed-shape state. It may be run at any time, before or after +the run-scoped-id fix is deployed to a given server, and re-run as often as +desired (see Idempotency below). Running it is a maintenance convenience, +not a precondition for shipping the run-scoped-id fix. + +Connection +---------- +Connection details come from :func:`context_intelligence_server.config.get_settings`. +Resolution order (highest first): + +1. Environment variables with prefix ``AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_``. +2. ``server-config.yaml`` in the current working directory. +3. Built-in defaults. + +To target a DTU, override env vars:: + + AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_NEO4J_URL=bolt://localhost:7688 \\ + AMPLIFIER_CONTEXT_INTELLIGENCE_SERVER_NEO4J_PASSWORD=testpassword \\ + uv run python scripts/tag_legacy_pooled_iterations.py --dry-run + +This script is standalone: it connects to Neo4j directly via the ``neo4j`` +driver and does NOT import or start the FastAPI application. + +Idempotency +----------- +The write query only touches nodes where +``i.data_quality IS NULL OR i.data_quality <> 'legacy_pooled_pre_fix'``, so a +second ``--apply`` run against an already-tagged graph matches zero rows and +performs zero writes. The tag itself is never removed by this script (it is +a permanent, non-destructive marker); un-tagging is out of scope. + +Modes +----- +--dry-run (default) + Read-only report: total bare-id Iteration count, broken down into + ``confirmed_corrupt`` (runs >= 2), ``confirmed_clean_single_run`` + (runs == 1), ``no_run_edge`` (runs == 0), and how many of the + ``confirmed_corrupt`` set are NOT yet tagged (i.e. would be written by + ``--apply``). Writes NOTHING. Always exits 0 (this is a health-check / + reporting mode, not a gate). + +--apply + ``SET i.data_quality = 'legacy_pooled_pre_fix'`` on the confirmed-corrupt + set, batched via ``CALL { ... } IN TRANSACTIONS OF N ROWS``. Only writes + rows not already carrying the tag (idempotent). Prints the number of rows + tagged, then re-runs the untagged-count verification query and reports + it (must be 0 after a successful apply). + +Exit codes +---------- +* 0 -- success (dry-run report printed; or apply completed and verification + found zero confirmed-corrupt nodes remaining untagged). +* 1 -- apply completed but verification found untagged confirmed-corrupt + nodes remaining (should not happen; signals a bug or concurrent writer). +""" + +from __future__ import annotations + +import argparse +import sys + +from context_intelligence_server.config import get_settings +from neo4j import GraphDatabase + +DEFAULT_BATCH_SIZE = 500 + +TAG_VALUE = "legacy_pooled_pre_fix" + +# --------------------------------------------------------------------------- +# Selector -- the ONLY nodes this script ever touches. +# Shared verbatim (as a fragment) across classify/apply/verify so there is +# exactly one place that defines "confirmed corrupt". +# --------------------------------------------------------------------------- +_CONFIRMED_CORRUPT_MATCH = ( + "MATCH (run:OrchestratorRun)-[:HAS_PART]->(i:Iteration) " + "WHERE NOT i.node_id CONTAINS '::orch_run::' " + "WITH i, count(DISTINCT run) AS runs " + "WHERE runs >= 2" +) + +_UNTAGGED_GUARD = ( + "WITH DISTINCT i WHERE i.data_quality IS NULL OR i.data_quality <> $tag_value" +) + + +# --------------------------------------------------------------------------- +# Read-only helpers +# --------------------------------------------------------------------------- + + +def classify(session) -> dict[str, int]: + """Return bucketed counts over ALL bare-id Iteration nodes. No writes. + + Buckets: + - total: every bare-id (pre-fix-shaped) Iteration node + - confirmed_corrupt: runs >= 2 (the ONLY set this script will ever tag) + - confirmed_clean_single_run: runs == 1 (leave alone) + - no_run_edge: runs == 0, i.e. no surviving OrchestratorRun HAS_PART + parent at all (leave alone -- corruption cannot be confirmed) + - confirmed_corrupt_untagged: of confirmed_corrupt, how many are NOT yet + carrying the tag (i.e. how many --apply would write) + """ + result = session.run( + "MATCH (i:Iteration) " + "WHERE NOT i.node_id CONTAINS '::orch_run::' " + "OPTIONAL MATCH (run:OrchestratorRun)-[:HAS_PART]->(i) " + "WITH i, count(DISTINCT run) AS runs " + "RETURN " + " count(i) AS total, " + " sum(CASE WHEN runs >= 2 THEN 1 ELSE 0 END) AS confirmed_corrupt, " + " sum(CASE WHEN runs = 1 THEN 1 ELSE 0 END) AS confirmed_clean_single_run, " + " sum(CASE WHEN runs = 0 THEN 1 ELSE 0 END) AS no_run_edge" + ) + row = result.single() + counts = { + "total": row["total"] or 0, + "confirmed_corrupt": row["confirmed_corrupt"] or 0, + "confirmed_clean_single_run": row["confirmed_clean_single_run"] or 0, + "no_run_edge": row["no_run_edge"] or 0, + } + counts["confirmed_corrupt_untagged"] = count_untagged(session) + return counts + + +def count_untagged(session) -> int: + """Return the count of confirmed-corrupt Iteration nodes NOT yet tagged. + + Standalone re-runnable verification query: after a successful --apply + this must return 0. + """ + result = session.run( + f"{_CONFIRMED_CORRUPT_MATCH} {_UNTAGGED_GUARD} RETURN count(i) AS untagged", + tag_value=TAG_VALUE, + ) + return result.single()["untagged"] or 0 + + +# --------------------------------------------------------------------------- +# Mutating operation +# --------------------------------------------------------------------------- + + +def tag_confirmed_corrupt(session, batch_size: int) -> int: + """SET data_quality='legacy_pooled_pre_fix' on the confirmed-corrupt, + not-yet-tagged subset, batched via CALL { ... } IN TRANSACTIONS OF N ROWS. + + Idempotent: rows already carrying the tag are excluded by the guard + before the CALL, so re-running matches (and writes) zero rows. + + Returns the number of rows tagged. + """ + result = session.run( + f"{_CONFIRMED_CORRUPT_MATCH} {_UNTAGGED_GUARD} " + "CALL { " + " WITH i " + " SET i.data_quality = $tag_value " + "} IN TRANSACTIONS OF $batch_size ROWS " + "RETURN count(i) AS tagged", + tag_value=TAG_VALUE, + batch_size=batch_size, + ) + return result.single()["tagged"] or 0 + + +# --------------------------------------------------------------------------- +# Reporting +# --------------------------------------------------------------------------- + + +def run_dry_run(session) -> int: + """Print classification report. Always returns 0 (health-check, not a gate).""" + counts = classify(session) + print("DRY RUN -- confirmed-corrupt Iteration nodes (TAG ONLY, no writes):\n") + print(f"{'Bucket':<35} {'Count':>10}") + print("-" * 46) + print(f"{'total (all bare-id Iterations)':<35} {counts['total']:>10}") + print(f"{'confirmed_corrupt (runs >= 2)':<35} {counts['confirmed_corrupt']:>10}") + print( + f"{'confirmed_clean_single_run (runs == 1)':<35} " + f"{counts['confirmed_clean_single_run']:>10}" + ) + print(f"{'no_run_edge (runs == 0)':<35} {counts['no_run_edge']:>10}") + print("-" * 46) + print( + f"{'confirmed_corrupt NOT yet tagged':<35} " + f"{counts['confirmed_corrupt_untagged']:>10}" + ) + print( + f"\n --apply would tag {counts['confirmed_corrupt_untagged']} node(s). " + f"{counts['confirmed_clean_single_run'] + counts['no_run_edge']} bare-id " + "node(s) are left untouched (not confirmed corrupt).\n" + ) + return 0 + + +def run_apply(session, batch_size: int) -> int: + """Tag the confirmed-corrupt set, then verify zero remain untagged. + + Returns the process exit code (0 = verified clean, 1 = residual untagged). + """ + print("APPLY -- tagging confirmed-corrupt Iteration nodes:\n") + tagged = tag_confirmed_corrupt(session, batch_size) + print(f" Tagged {tagged} node(s) with data_quality='{TAG_VALUE}'.\n") + + remaining = count_untagged(session) + print("VERIFICATION -- confirmed-corrupt nodes still missing the tag:\n") + print(f" remaining untagged: {remaining}") + if remaining: + print( + "\nAPPLY INCOMPLETE -- " + f"{remaining} confirmed-corrupt node(s) remain untagged after apply. " + "This should not happen; investigate for a concurrent writer or a bug." + ) + return 1 + print("\nAPPLY COMPLETE -- all confirmed-corrupt nodes are tagged.") + return 0 + + +# --------------------------------------------------------------------------- +# Entry point +# --------------------------------------------------------------------------- + + +def main() -> int: + """Parse arguments, connect to Neo4j, and run dry-run or apply.""" + parser = argparse.ArgumentParser( + prog="tag_legacy_pooled_iterations.py", + description=( + "Maintenance tool: non-destructively TAG (data_quality=" + f"'{TAG_VALUE}') legacy bare-id Iteration nodes confirmed pooled " + "across >=2 distinct OrchestratorRuns before the run-scoped-id fix. Never " + "touches run-scoped nodes, single-run bare-id nodes, or bare-id " + "nodes with no surviving run parent. The destructive cleanup is " + "a separate, gated follow-up -- out of scope here." + ), + ) + parser.add_argument( + "--dry-run", + action="store_true", + help=( + "Report bucketed counts (confirmed_corrupt / " + "confirmed_clean_single_run / no_run_edge) and how many " + "confirmed-corrupt nodes would be tagged. Writes nothing. " + "This is the default mode." + ), + ) + parser.add_argument( + "--apply", + action="store_true", + help=( + "Tag the confirmed-corrupt set (batched, idempotent), then " + "verify zero remain untagged." + ), + ) + parser.add_argument( + "--batch-size", + type=int, + default=DEFAULT_BATCH_SIZE, + metavar="N", + help=f"Rows per transaction for --apply (default: {DEFAULT_BATCH_SIZE}).", + ) + parser.add_argument( + "--neo4j-url", + metavar="URL", + default=None, + help="Neo4j Bolt URL (overrides server-config.yaml / env var)", + ) + parser.add_argument( + "--neo4j-user", + metavar="USER", + default=None, + help="Neo4j username (overrides server-config.yaml / env var)", + ) + parser.add_argument( + "--neo4j-password", + metavar="PW", + default=None, + help="Neo4j password (overrides server-config.yaml / env var)", + ) + args = parser.parse_args() + + if args.apply and args.dry_run: + parser.error("--dry-run and --apply are mutually exclusive") + + settings = get_settings() + neo4j_url = args.neo4j_url or settings.neo4j_url + neo4j_user = args.neo4j_user or settings.neo4j_user + neo4j_password = args.neo4j_password or settings.neo4j_password + + print(f"Connecting to Neo4j at {neo4j_url} as {neo4j_user}\n") + driver = GraphDatabase.driver(neo4j_url, auth=(neo4j_user, neo4j_password)) + try: + with driver.session() as neo_session: + if args.apply: + return run_apply(neo_session, args.batch_size) + # --dry-run is the default: no writes unless --apply is explicit. + return run_dry_run(neo_session) + finally: + driver.close() + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/conftest.py b/tests/conftest.py index ecfd94d9..11835162 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,12 +99,14 @@ 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 _real = _Settings() class _SettingsProxy: + blob_backend: str = _real.blob_backend blob_path: str = _real.blob_path queues_path: str = str(tmp_path / "queues") # Redirect identity-store paths so the registry proxy never touches the @@ -122,6 +122,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 +155,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 +212,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 @@ -178,6 +229,29 @@ def reset_registry() -> Generator[None, None, None]: registry._write_semaphore = None +@pytest.fixture(autouse=True) +def reset_maintenance_coordinator() -> Generator[None, None, None]: + """Each test starts with (and leaves) a pristine MaintenanceCoordinator. + + ``maintenance.coordinator`` is a process-wide singleton shared by the + drain-loop gate and the HTTP gate/status. A test that runs the real + ``lifespan()`` calls ``coordinator.bind_driver(...)`` against its own mock + driver; without this reset that leaks into the coordinator's TTL-cached + probe/op state and can spuriously close the gate for every other test. + Reset by copying a fresh instance's attributes in place, since other modules + hold a direct reference to THIS object. + """ + from context_intelligence_server.maintenance import MaintenanceCoordinator + from context_intelligence_server.maintenance import coordinator as _coordinator + + def _reset() -> None: + _coordinator.__dict__.update(vars(MaintenanceCoordinator())) + + _reset() + yield + _reset() + + @pytest.fixture async def client() -> AsyncGenerator[httpx.AsyncClient, None]: async with httpx.AsyncClient( @@ -192,10 +266,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/handlers/data_layer_2/test_content_block.py b/tests/handlers/data_layer_2/test_content_block.py index 53746570..32baefef 100644 --- a/tests/handlers/data_layer_2/test_content_block.py +++ b/tests/handlers/data_layer_2/test_content_block.py @@ -3,7 +3,7 @@ Covers: - handled_events == frozenset({'content_block:start', 'content_block:end'}) - content_block:start creates ContentBlock:SST_EVENT node keyed as - '{session_id}::block::{iteration_n}::{block_index}' with session_id, block_index, + '{iteration_id}::block::{block_index}' with session_id, block_index, started_at; iteration_n extracted from active_iteration_id cursor (split('::')[-1]) - E07: Iteration -[:HAS_PART {sst_semantic: 'CONTAINS'}]-> ContentBlock created when active_iteration_id is set; NOT created when no active iteration (zero edges) @@ -53,7 +53,7 @@ class TestContentBlockStartCreatesNode: async def test_node_created_with_correct_compound_key( self, services: HookStateService ) -> None: - """content_block:start must create node at '{session_id}::block::{iteration_n}::{block_index}'.""" + """content_block:start must create node at '{iteration_id}::block::{block_index}'.""" services.data_layer_2.active_iteration_id = "s1::iteration::1" handler = ContentBlockHandler(services) await handler( @@ -64,7 +64,7 @@ async def test_node_created_with_correct_compound_key( "block_index": 0, }, ) - node_id = "s1::block::1::0" + node_id = "s1::iteration::1::block::0" node = await services.graph.get_node(node_id) assert node is not None, f"content_block:start must create node at '{node_id}'" @@ -82,7 +82,7 @@ async def test_node_has_content_block_and_sst_event_labels( "block_index": 0, }, ) - node = await services.graph.get_node("s1::block::1::0") + node = await services.graph.get_node("s1::iteration::1::block::0") assert node is not None assert "ContentBlock" in node["labels"], ( f"ContentBlock label missing. Got: {node['labels']}" @@ -105,7 +105,7 @@ async def test_node_has_session_id_property( "block_index": 0, }, ) - node = await services.graph.get_node("s1::block::1::0") + node = await services.graph.get_node("s1::iteration::1::block::0") assert node is not None assert node.get("session_id") == "s1", ( f"session_id property missing or wrong. Got: {node!r}" @@ -125,7 +125,7 @@ async def test_node_has_block_index_and_started_at_properties( "block_index": 2, }, ) - node = await services.graph.get_node("s1::block::1::2") + node = await services.graph.get_node("s1::iteration::1::block::2") assert node is not None assert node.get("block_index") == 2, ( f"block_index property missing or wrong. Got: {node!r}" @@ -158,7 +158,7 @@ async def test_e07_edge_created_when_active_iteration_id_is_set( }, ) iteration_id = "s1::iteration::1" - block_id = "s1::block::1::0" + block_id = "s1::iteration::1::block::0" edge = await services.graph.get_edge(iteration_id, block_id) assert edge is not None, ( f"E07 HAS_PART edge from '{iteration_id}' to '{block_id}' must exist " @@ -228,7 +228,7 @@ async def test_content_block_end_sets_block_type( "block": {"type": "text"}, }, ) - node = await services.graph.get_node("s1::block::1::0") + node = await services.graph.get_node("s1::iteration::1::block::0") assert node is not None assert node.get("block_type") == "text", ( f"content_block:end must set block_type from block.type. Got: {node!r}" @@ -258,7 +258,7 @@ async def test_content_block_end_sets_ended_at( "block": {"type": "text"}, }, ) - node = await services.graph.get_node("s1::block::1::0") + node = await services.graph.get_node("s1::iteration::1::block::0") assert node is not None assert node.get("ended_at") == "2026-01-01T00:01:00Z", ( f"content_block:end must set ended_at. Got: {node!r}" @@ -302,9 +302,9 @@ async def test_tool_call_block_with_id_is_cached( ) assert ( services.data_layer_2.pending_tool_block_ids["tool-block-abc"] - == "s1::block::1::0" + == "s1::iteration::1::block::0" ), ( - "pending_tool_block_ids['tool-block-abc'] must map to the block node id 's1::block::1::0'" + "pending_tool_block_ids['tool-block-abc'] must map to the block node id 's1::iteration::1::block::0'" ) async def test_text_block_not_cached(self, services: HookStateService) -> None: @@ -448,7 +448,7 @@ async def test_content_block_start_creates_sourced_from_edge( "block_index": 0, }, ) - block_node_id = "s1::block::1::0" + block_node_id = "s1::iteration::1::block::0" data_layer_1_node_id = make_node_id("s1", "content_block:start", timestamp) edge = await services.graph.get_edge(block_node_id, data_layer_1_node_id) assert edge is not None, ( @@ -484,7 +484,7 @@ async def test_content_block_end_creates_sourced_from_edge( "block": {"type": "text"}, }, ) - block_node_id = "s1::block::1::0" + block_node_id = "s1::iteration::1::block::0" data_layer_1_node_id = make_node_id("s1", "content_block:end", end_timestamp) edge = await services.graph.get_edge(block_node_id, data_layer_1_node_id) assert edge is not None, ( diff --git a/tests/handlers/data_layer_2/test_iteration.py b/tests/handlers/data_layer_2/test_iteration.py index 77512c4f..1e511025 100644 --- a/tests/handlers/data_layer_2/test_iteration.py +++ b/tests/handlers/data_layer_2/test_iteration.py @@ -3,10 +3,15 @@ Covers: - handled_events == frozenset({'provider:request', 'llm:request', 'llm:response'}) - provider:request creates Iteration:SST_EVENT node keyed as - '{session_id}::iteration::{iteration_number}' with session_id, iteration_number, - and started_at; sets active_iteration_id cursor + '{session_id}::iteration::{iteration_number}' (no active orchestrator run) or + '{session_id}::orch_run::{execution_start_ts}::iteration::{iteration_number}' + (run-scoped when a run is active) with session_id, + iteration_number, and started_at; sets active_iteration_id cursor - E06: OrchestratorRun -[:HAS_PART {sst_semantic: 'CONTAINS'}]-> Iteration - created when execution_start_ts cursor is set; NOT created when None + created when execution_start_ts cursor is set (target is the run-scoped + iteration_id); NOT created when None +- two iterations sharing the same iteration_number under DIFFERENT + orchestrator runs get DISTINCT node_ids (no cross-run collision) - llm:request enriches active Iteration with provider, model, message_count, has_system; noop when active_iteration_id is None - llm:response enriches active Iteration with usage_input, usage_output, usage_cache_write; @@ -16,11 +21,15 @@ from __future__ import annotations +import logging + from context_intelligence_server.handlers.data_layer_2.iteration import IterationHandler +from context_intelligence_server.handlers.data_layer_2.orchestrator_run import ( + OrchestratorRunHandler, +) from context_intelligence_server.services import HookStateService from context_intelligence_server.utils import make_node_id - # --------------------------------------------------------------------------- # 1. TestIterationHandlerHandledEvents # --------------------------------------------------------------------------- @@ -172,10 +181,14 @@ class TestE06HasPartEdge: async def test_e06_has_part_edge_created_when_execution_start_ts_is_set( self, services: HookStateService ) -> None: - """E06 edge must be created when execution_start_ts cursor is set before provider:request.""" + """E06 edge must be created when execution_start_ts cursor is set before provider:request. + + The edge target is the run-scoped iteration_id, not the bare shape. + """ handler = IterationHandler(services) - # Simulate that execution:start previously fired and set the cursor + # Simulate that execution:start previously fired and set both cursors. services.data_layer_2.execution_start_ts = "2026-01-01T00:00:00Z" + services.data_layer_2.active_orch_run_id = "s1::orch_run::2026-01-01T00:00:00Z::1" await handler( "provider:request", @@ -184,8 +197,8 @@ async def test_e06_has_part_edge_created_when_execution_start_ts_is_set( "timestamp": "2026-01-01T00:00:01Z", }, ) - orch_run_id = "s1::orch_run::2026-01-01T00:00:00Z" - iteration_id = "s1::iteration::1" + orch_run_id = "s1::orch_run::2026-01-01T00:00:00Z::1" + iteration_id = "s1::orch_run::2026-01-01T00:00:00Z::1::iteration::1" edge = await services.graph.get_edge(orch_run_id, iteration_id) assert edge is not None, ( f"E06 HAS_PART edge from '{orch_run_id}' to '{iteration_id}' must exist " @@ -197,6 +210,9 @@ async def test_e06_has_part_edge_created_when_execution_start_ts_is_set( assert edge.get("sst_semantic") == "CONTAINS", ( f"E06 edge must have sst_semantic='CONTAINS'. Got: {edge.get('sst_semantic')}" ) + assert services.data_layer_2.active_iteration_id == iteration_id, ( + "active_iteration_id cursor must be set to the run-scoped iteration_id" + ) async def test_e06_not_created_when_execution_start_ts_is_none( self, services: HookStateService @@ -218,6 +234,118 @@ async def test_e06_not_created_when_execution_start_ts_is_none( f"Only SOURCED_FROM edge should exist when execution_start_ts is None. " f"Got {len(services.graph._edges)} edges: {list(services.graph._edges.keys())}" ) + # Falls back to the bare (pre-fix-shaped) id when no orchestrator run is active + assert services.data_layer_2.active_iteration_id == "s1::iteration::1" + + +# --------------------------------------------------------------------------- +# 3b. TestIterationRunScopingP21 +# --------------------------------------------------------------------------- + + +class TestIterationRunScopingP21: + """Iteration node_id is run-scoped and does not collide across runs.""" + + async def test_same_iteration_number_under_different_runs_gets_distinct_node_ids( + self, services: HookStateService + ) -> None: + """Two iterations sharing iteration_number=1 under DIFFERENT orchestrator runs + must produce DISTINCT node_ids and both nodes must independently exist. + + This reproduces the real-world collision: iteration_count is a per-session + (not per-run) counter, so after a drainer restart/replay recreates + DataLayer2State the counter can restart from zero and reproduce a prior + run's iteration_number under a NEW orchestrator run. Before run-scoping, + both runs' first iteration would MERGE onto the bare + 's1::iteration::1' node_id. After the fix, each run's iteration_number=1 is + prefixed with its own orch_run_id and the two nodes stay distinct. + """ + orch = OrchestratorRunHandler(services) + handler = IterationHandler(services) + + # --- Run 1: a real execution:start mints the run + its tiebreaker (seq 1). + await orch( + "execution:start", + {"session_id": "s1", "timestamp": "2026-01-01T00:00:00Z"}, + ) + await handler( + "provider:request", + {"session_id": "s1", "timestamp": "2026-01-01T00:00:01Z"}, + ) + run1_iteration_id = services.data_layer_2.active_iteration_id + await orch( + "orchestrator:complete", + {"session_id": "s1", "timestamp": "2026-01-01T00:00:02Z"}, + ) + + # --- Simulate a drainer restart: a fresh DataLayer2State restarts + # iteration_count from 0. The durable orch_run_seq is what a restored + # cursor carries; leave it as-is so a genuinely NEW run keeps advancing it. + services.data_layer_2.iteration_count = 0 + # A SECOND run that shares run 1's identical timestamp (coarse clock / + # replayed execution:start) -- the exact collision the tiebreaker exists for. + await orch( + "execution:start", + {"session_id": "s1", "timestamp": "2026-01-01T00:00:00Z"}, + ) + await handler( + "provider:request", + {"session_id": "s1", "timestamp": "2026-01-01T00:00:03Z"}, + ) + run2_iteration_id = services.data_layer_2.active_iteration_id + + # Same timestamp, different seq -> distinct ids (tiebreaker load-bearing). + assert run1_iteration_id == "s1::orch_run::2026-01-01T00:00:00Z::1::iteration::1" + assert run2_iteration_id == "s1::orch_run::2026-01-01T00:00:00Z::2::iteration::1" + assert run1_iteration_id != run2_iteration_id, ( + "Two runs sharing an identical execution_start_ts must still get " + "distinct Iteration ids (the run tiebreaker)." + ) + + node1 = await services.graph.get_node(run1_iteration_id) + node2 = await services.graph.get_node(run2_iteration_id) + assert node1 is not None, ( + f"Run 1's Iteration node '{run1_iteration_id}' must exist" + ) + assert node2 is not None, ( + f"Run 2's Iteration node '{run2_iteration_id}' must exist" + ) + assert node1.get("iteration_number") == 1 + assert node2.get("iteration_number") == 1 + assert node1.get("started_at") == "2026-01-01T00:00:01Z", ( + "Run 1's Iteration node must retain its own started_at, not run 2's " + "(proves the two nodes were never merged together)" + ) + assert node2.get("started_at") == "2026-01-01T00:00:03Z" + + # --- Regression guard: exactly ONE distinct HAS_PART parent per Iteration. + run1_orch_run_id = "s1::orch_run::2026-01-01T00:00:00Z::1" + run2_orch_run_id = "s1::orch_run::2026-01-01T00:00:00Z::2" + + def has_part_parents(iteration_id: str) -> list[str]: + """Distinct HAS_PART parent ids pointing at *iteration_id* in the fake graph.""" + return [ + src + for (src, dst), data in services.graph._edges.items() + if dst == iteration_id and data.get("type") == "HAS_PART" + ] + + run1_parents = has_part_parents(run1_iteration_id) + run2_parents = has_part_parents(run2_iteration_id) + + assert run1_parents == [run1_orch_run_id], ( + f"Run 1's Iteration node '{run1_iteration_id}' must have exactly ONE " + f"HAS_PART parent (its own OrchestratorRun). Got: {run1_parents!r}" + ) + assert run2_parents == [run2_orch_run_id], ( + f"Run 2's Iteration node '{run2_iteration_id}' must have exactly ONE " + f"HAS_PART parent (its own OrchestratorRun). Got: {run2_parents!r}" + ) + assert set(run1_parents).isdisjoint(run2_parents), ( + "No Iteration node may be shared (MERGEd) across the two " + "OrchestratorRuns -- each run's iterations must be distinct nodes " + "with distinct, non-overlapping HAS_PART parents." + ) # --------------------------------------------------------------------------- @@ -597,3 +725,175 @@ async def test_llm_response_creates_sourced_from_edge( assert edge.get("type") == "SOURCED_FROM", ( f"Edge type must be 'SOURCED_FROM'. Got: {edge.get('type')!r}" ) + + +# --------------------------------------------------------------------------- +# iteration_scope completeness +# +# The Iteration node is upsert_node'd from THREE sites: provider:request, +# llm:request, llm:response. Each site stamps the additive 'iteration_scope' +# ('run' | 'unscoped') property independently, sourced from the SAME cursor +# field (execution_start_ts) -- so a node can never be created/updated +# without a scope value, regardless of which of the three sites happens to +# be the one that actually writes it (e.g. a dead-lettered provider:request +# whose active_iteration_id mutation nonetheless survives to a later +# llm:request/llm:response). +# --------------------------------------------------------------------------- + + +class TestIterationScopeCompleteness: + """iteration_scope stamped at ALL THREE upsert_node call sites.""" + + async def test_provider_request_unscoped_with_no_execution_start( + self, services: HookStateService + ) -> None: + """provider:request with NO preceding execution:start -> + iteration_scope == 'unscoped'.""" + assert services.data_layer_2.execution_start_ts is None + handler = IterationHandler(services) + await handler( + "provider:request", + {"session_id": "s1", "timestamp": "2026-01-01T00:00:00Z"}, + ) + node = await services.graph.get_node("s1::iteration::1") + assert node is not None + assert node.get("iteration_scope") == "unscoped", ( + f"Expected iteration_scope='unscoped'. Got: {node!r}" + ) + + async def test_provider_request_run_scoped_with_execution_start( + self, services: HookStateService + ) -> None: + """A normal (execution:start already seen) run stamps 'run'.""" + services.data_layer_2.execution_start_ts = "2026-01-01T00:00:00Z" + services.data_layer_2.active_orch_run_id = "s1::orch_run::2026-01-01T00:00:00Z::1" + handler = IterationHandler(services) + await handler( + "provider:request", + {"session_id": "s1", "timestamp": "2026-01-01T00:00:01Z"}, + ) + node_id = "s1::orch_run::2026-01-01T00:00:00Z::1::iteration::1" + node = await services.graph.get_node(node_id) + assert node is not None + assert node.get("iteration_scope") == "run", ( + f"Expected iteration_scope='run'. Got: {node!r}" + ) + + async def test_unscoped_emitted_log_is_info_not_warning( + self, services: HookStateService, caplog + ) -> None: + """The log line for an unscoped iteration must be + INFO (or rate-limited), NEVER WARNING -- a loop-basic unscoped + session is a normal case, not an alert-worthy anomaly.""" + handler = IterationHandler(services) + with caplog.at_level( + logging.INFO, + logger="context_intelligence_server.handlers.data_layer_2.iteration", + ): + await handler( + "provider:request", + {"session_id": "s1", "timestamp": "2026-01-01T00:00:00Z"}, + ) + unscoped_records = [ + r for r in caplog.records if "unscoped_iteration_emitted" in r.message + ] + assert unscoped_records, ( + "expected an 'unscoped_iteration_emitted' log record to be emitted" + ) + assert all(r.levelno == logging.INFO for r in unscoped_records), ( + "unscoped_iteration_emitted must log at INFO, got levels: " + f"{[r.levelname for r in unscoped_records]}" + ) + assert not any(r.levelno >= logging.WARNING for r in caplog.records), ( + "no WARNING (or higher) should be emitted for a normal unscoped iteration" + ) + + async def test_llm_request_stamps_unscoped_independent_of_provider_request( + self, services: HookStateService + ) -> None: + """llm:request must stamp iteration_scope + on its OWN upsert_node call, even when the node was never created by + provider:request (e.g. a dead-lettered provider:request whose + active_iteration_id cursor mutation nonetheless survives). Simulated + here by setting the cursor directly, bypassing provider:request.""" + services.data_layer_2.active_iteration_id = "s1::iteration::99" + handler = IterationHandler(services) + await handler( + "llm:request", + { + "session_id": "s1", + "timestamp": "2026-01-01T00:00:00Z", + "provider": "anthropic", + "model": "claude", + }, + ) + node = await services.graph.get_node("s1::iteration::99") + assert node is not None + assert node.get("iteration_scope") == "unscoped", ( + f"llm:request must stamp iteration_scope='unscoped'. Got: {node!r}" + ) + + async def test_llm_request_stamps_run_scope_when_execution_start_ts_set( + self, services: HookStateService + ) -> None: + services.data_layer_2.execution_start_ts = "2026-01-01T00:00:00Z" + services.data_layer_2.active_orch_run_id = "s1::orch_run::2026-01-01T00:00:00Z::1" + node_id = "s1::orch_run::2026-01-01T00:00:00Z::1::iteration::1" + services.data_layer_2.active_iteration_id = node_id + handler = IterationHandler(services) + await handler( + "llm:request", + { + "session_id": "s1", + "timestamp": "2026-01-01T00:00:01Z", + "provider": "anthropic", + "model": "claude", + }, + ) + node = await services.graph.get_node(node_id) + assert node is not None + assert node.get("iteration_scope") == "run", ( + f"llm:request must stamp iteration_scope='run'. Got: {node!r}" + ) + + async def test_llm_response_stamps_unscoped_independent_of_provider_request( + self, services: HookStateService + ) -> None: + """Same completeness gap as llm:request, for the third site.""" + services.data_layer_2.active_iteration_id = "s1::iteration::99" + handler = IterationHandler(services) + await handler( + "llm:response", + { + "session_id": "s1", + "timestamp": "2026-01-01T00:00:00Z", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + node = await services.graph.get_node("s1::iteration::99") + assert node is not None + assert node.get("iteration_scope") == "unscoped", ( + f"llm:response must stamp iteration_scope='unscoped'. Got: {node!r}" + ) + + async def test_llm_response_stamps_run_scope_when_execution_start_ts_set( + self, services: HookStateService + ) -> None: + services.data_layer_2.execution_start_ts = "2026-01-01T00:00:00Z" + services.data_layer_2.active_orch_run_id = "s1::orch_run::2026-01-01T00:00:00Z::1" + node_id = "s1::orch_run::2026-01-01T00:00:00Z::1::iteration::1" + services.data_layer_2.active_iteration_id = node_id + handler = IterationHandler(services) + await handler( + "llm:response", + { + "session_id": "s1", + "timestamp": "2026-01-01T00:00:01Z", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + node = await services.graph.get_node(node_id) + assert node is not None + assert node.get("iteration_scope") == "run", ( + f"llm:response must stamp iteration_scope='run'. Got: {node!r}" + ) diff --git a/tests/handlers/data_layer_2/test_orchestrator_run.py b/tests/handlers/data_layer_2/test_orchestrator_run.py index 44bd54f1..a7d04662 100644 --- a/tests/handlers/data_layer_2/test_orchestrator_run.py +++ b/tests/handlers/data_layer_2/test_orchestrator_run.py @@ -66,7 +66,7 @@ async def test_node_created_with_correct_compound_key( "timestamp": "2026-01-01T00:00:00Z", }, ) - node_id = "s1::orch_run::2026-01-01T00:00:00Z" + node_id = "s1::orch_run::2026-01-01T00:00:00Z::1" node = await services.graph.get_node(node_id) assert node is not None, f"execution:start must create node at '{node_id}'" @@ -82,7 +82,7 @@ async def test_node_has_orchestrator_run_and_sst_event_labels( "timestamp": "2026-01-01T00:00:00Z", }, ) - node = await services.graph.get_node("s1::orch_run::2026-01-01T00:00:00Z") + node = await services.graph.get_node("s1::orch_run::2026-01-01T00:00:00Z::1") assert node is not None assert "OrchestratorRun" in node["labels"], ( f"OrchestratorRun label missing. Got: {node['labels']}" @@ -103,7 +103,7 @@ async def test_node_has_session_id_and_started_at( "timestamp": "2026-01-01T00:00:00Z", }, ) - node = await services.graph.get_node("s1::orch_run::2026-01-01T00:00:00Z") + node = await services.graph.get_node("s1::orch_run::2026-01-01T00:00:00Z::1") assert node is not None assert node.get("session_id") == "s1", ( f"session_id property missing or wrong. Got: {node!r}" @@ -124,7 +124,7 @@ async def test_e01_has_execution_edge_created( "timestamp": "2026-01-01T00:00:00Z", }, ) - orch_run_id = "s1::orch_run::2026-01-01T00:00:00Z" + orch_run_id = "s1::orch_run::2026-01-01T00:00:00Z::1" edge = await services.graph.get_edge("s1", orch_run_id) assert edge is not None, ( f"E01 HAS_EXECUTION edge from 's1' to '{orch_run_id}' must exist" @@ -191,7 +191,7 @@ async def test_execution_end_sets_ended_at( "status": "completed", }, ) - node = await services.graph.get_node("s1::orch_run::2026-01-01T00:00:00Z") + node = await services.graph.get_node("s1::orch_run::2026-01-01T00:00:00Z::1") assert node is not None assert node.get("ended_at") == "2026-01-01T00:01:00Z", ( f"ended_at must be set by execution:end. Got: {node!r}" @@ -212,7 +212,7 @@ async def test_execution_end_sets_status(self, services: HookStateService) -> No "status": "completed", }, ) - node = await services.graph.get_node("s1::orch_run::2026-01-01T00:00:00Z") + node = await services.graph.get_node("s1::orch_run::2026-01-01T00:00:00Z::1") assert node is not None assert node.get("status") == "completed", ( f"status must be set by execution:end. Got: {node!r}" @@ -236,7 +236,7 @@ async def test_execution_end_sets_response_properties( "response": "final answer text", }, ) - node = await services.graph.get_node("s1::orch_run::2026-01-01T00:00:00Z") + node = await services.graph.get_node("s1::orch_run::2026-01-01T00:00:00Z::1") assert node is not None assert node.get("response") == "final answer text", ( f"response must be set by execution:end. Got: {node!r}" @@ -269,7 +269,7 @@ async def test_orchestrator_complete_enriches_name_turn_count_completed_at( "turn_count": 3, }, ) - node = await services.graph.get_node("s1::orch_run::2026-01-01T00:00:00Z") + node = await services.graph.get_node("s1::orch_run::2026-01-01T00:00:00Z::1") assert node is not None assert node.get("orchestrator_name") == "my-orchestrator", ( f"orchestrator_name must be set by orchestrator:complete. Got: {node!r}" @@ -366,7 +366,7 @@ async def test_orchestrator_complete_sets_last_completed_orch_run_id( "turn_count": 1, }, ) - expected_run_id = "s1::orch_run::2026-01-01T00:00:00Z" + expected_run_id = "s1::orch_run::2026-01-01T00:00:00Z::1" assert services.data_layer_2.last_completed_orch_run_id == expected_run_id, ( f"last_completed_orch_run_id must be set to '{expected_run_id}'. " f"Got: {services.data_layer_2.last_completed_orch_run_id!r}" @@ -420,7 +420,7 @@ async def test_e14_prompt_triggers_orchestrator_run_edge_created( "timestamp": "2026-01-01T00:01:00Z", }, ) - orch_run_id = "s1::orch_run::2026-01-01T00:01:00Z" + orch_run_id = "s1::orch_run::2026-01-01T00:01:00Z::1" prompt_id = "s1::prompt::2026-01-01T00:00:00Z" edge = await services.graph.get_edge(prompt_id, orch_run_id) assert edge is not None, ( @@ -509,7 +509,7 @@ async def test_execution_start_creates_sourced_from_edge( "timestamp": "2026-01-01T00:00:00Z", }, ) - orch_run_id = "s1::orch_run::2026-01-01T00:00:00Z" + orch_run_id = "s1::orch_run::2026-01-01T00:00:00Z::1" data_layer_1_node_id = make_node_id( "s1", "execution:start", "2026-01-01T00:00:00Z" ) @@ -539,7 +539,7 @@ async def test_execution_end_creates_sourced_from_edge( "status": "completed", }, ) - orch_run_id = "s1::orch_run::2026-01-01T00:00:00Z" + orch_run_id = "s1::orch_run::2026-01-01T00:00:00Z::1" data_layer_1_node_id = make_node_id( "s1", "execution:end", "2026-01-01T00:01:00Z" ) @@ -570,7 +570,7 @@ async def test_orchestrator_complete_creates_sourced_from_edge( "turn_count": 3, }, ) - orch_run_id = "s1::orch_run::2026-01-01T00:00:00Z" + orch_run_id = "s1::orch_run::2026-01-01T00:00:00Z::1" data_layer_1_node_id = make_node_id( "s1", "orchestrator:complete", "2026-01-01T00:02:00Z" ) diff --git a/tests/handlers/data_layer_2/test_session.py b/tests/handlers/data_layer_2/test_session.py index 19983caf..7f49a9fb 100644 --- a/tests/handlers/data_layer_2/test_session.py +++ b/tests/handlers/data_layer_2/test_session.py @@ -2078,6 +2078,14 @@ class TestClassifyMatrix: remove, because by the time a node reaches one of those states its StubSession marker has already been cleared by the transition that got it there. + + "IncompleteSession" heal-forward: EVERY start/fork cell (including the + former pure no-ops) now also removes "IncompleteSession" — see + _heal_forward() in session.py. A node may carry a stale IncompleteSession + marker from an earlier out-of-order session:end regardless of which + start/fork branch it reaches, so every start/fork transition strips it. + The end/None cells are unchanged — session:end still stamps + IncompleteSession when genuinely no type is known at end. """ CASES: list[tuple[str, str | None, bool, list[str], list[str]]] = [ @@ -2085,78 +2093,78 @@ class TestClassifyMatrix: # ------------------------------------------------------------------ # event=start # ------------------------------------------------------------------ - ("start", "ForkedSession", True, [], []), - ("start", "ForkedSession", False, [], []), - ("start", "SubSession", True, [], []), - ("start", "SubSession", False, [], []), - ("start", "RootSession", False, [], []), + ("start", "ForkedSession", True, [], ["IncompleteSession"]), + ("start", "ForkedSession", False, [], ["IncompleteSession"]), + ("start", "SubSession", True, [], ["IncompleteSession"]), + ("start", "SubSession", False, [], ["IncompleteSession"]), + ("start", "RootSession", False, [], ["IncompleteSession"]), ( "start", "RootSession", True, ["SubSession", "SST_EVENT"], - ["RootSession", "StubSession"], + ["RootSession", "StubSession", "IncompleteSession"], ), ( "start", None, True, ["Session", "SubSession", "SST_EVENT"], - ["StubSession"], + ["StubSession", "IncompleteSession"], ), ( "start", None, False, ["RootSession", "Session", "SST_EVENT"], - ["StubSession"], + ["StubSession", "IncompleteSession"], ), # ------------------------------------------------------------------ # event=fork # ------------------------------------------------------------------ - ("fork", "ForkedSession", True, [], []), - ("fork", "ForkedSession", False, [], []), + ("fork", "ForkedSession", True, [], ["IncompleteSession"]), + ("fork", "ForkedSession", False, [], ["IncompleteSession"]), ( "fork", "RootSession", True, ["ForkedSession", "SST_EVENT"], - ["RootSession", "StubSession"], + ["RootSession", "StubSession", "IncompleteSession"], ), ( "fork", "RootSession", False, ["ForkedSession", "SST_EVENT"], - ["RootSession", "StubSession"], + ["RootSession", "StubSession", "IncompleteSession"], ), ( "fork", "SubSession", True, ["ForkedSession", "SST_EVENT"], - ["SubSession", "StubSession"], + ["SubSession", "StubSession", "IncompleteSession"], ), ( "fork", "SubSession", False, ["ForkedSession", "SST_EVENT"], - ["SubSession", "StubSession"], + ["SubSession", "StubSession", "IncompleteSession"], ), ( "fork", None, True, ["Session", "ForkedSession", "SST_EVENT"], - ["StubSession"], + ["StubSession", "IncompleteSession"], ), ( "fork", None, False, ["Session", "ForkedSession", "SST_EVENT"], - ["StubSession"], + ["StubSession", "IncompleteSession"], ), # ------------------------------------------------------------------ # event=end @@ -2165,7 +2173,8 @@ class TestClassifyMatrix: # not a fabricated Root/Sub terminal. has_parent is irrelevant here — # the server never guesses; it marks and surfaces the health signal. # IncompleteSession is a confirmed (if incomplete) terminal, so - # StubSession is cleared here too. + # StubSession is cleared here too. UNCHANGED by heal-forward: end is + # not a start/fork transition. ("end", None, True, ["IncompleteSession", "SST_EVENT"], ["StubSession"]), ("end", None, False, ["IncompleteSession", "SST_EVENT"], ["StubSession"]), ("end", "RootSession", True, [], []), @@ -2241,7 +2250,7 @@ async def test_ensure_session_node_existing_branch_does_not_add_stub_session( "already-enriched", {"labels": ["Session", "RootSession"], "status": "running"}, ) - # _seen_sessions cache is empty, so this call hits Tier 2 (graph query) + # _seen_sessions cache is empty, so this call hits the graph query # and takes the "existing is not None" branch. await services.ensure_session_node("already-enriched", {}) @@ -2882,3 +2891,126 @@ async def test_current_type_returns_none_for_incomplete_session_node( "_current_type must ignore IncompleteSession and return None, " "so a late start/fork can still classify the session normally" ) + + # ----------------------------------------------------------------------- + # Heal-forward: out-of-order end -> fork/start race + # + # A forked sub-session's session:end can drain, in an independent queue, + # BEFORE its session:fork/session:start. classify() sees current_type=None + # at end -> stamps IncompleteSession. When the real fork/start is then + # processed, it must strip the stale marker (heal-forward), leaving the + # node with ONLY the real terminal label + # (out-of-order end-before-start/fork heal-forward). + # ----------------------------------------------------------------------- + + async def test_out_of_order_end_then_fork_heals_incomplete_session( + self, services: HookStateService + ) -> None: + """end processed BEFORE fork: stale IncompleteSession must be stripped + the moment the real session:fork arrives, leaving ForkedSession only. + """ + handler = SessionHandler(services) + + # session:end drains first (simulating the cross-queue race) — stamps + # IncompleteSession on the bare node. + await handler( + "session:end", + {"session_id": "s-race-fork", "timestamp": "2026-01-01T01:00:00Z"}, + ) + node = await services.graph.get_node("s-race-fork") + assert node is not None + assert "IncompleteSession" in node["labels"], ( + "Precondition: out-of-order end must stamp the stale marker" + ) + + # The real session:fork arrives late. + await handler( + "session:fork", + { + "session_id": "s-race-fork", + "parent_id": "p-race", + "timestamp": "2026-01-01T00:00:00Z", + }, + ) + + node = await services.graph.get_node("s-race-fork") + assert node is not None + labels = node["labels"] + assert "ForkedSession" in labels, "Real terminal must be assigned by fork" + assert "IncompleteSession" not in labels, ( + "Heal-forward: session:fork must strip the stale IncompleteSession " + "marker left by the out-of-order session:end" + ) + + async def test_out_of_order_end_then_start_heals_incomplete_session( + self, services: HookStateService + ) -> None: + """end processed BEFORE start (no parent): stale IncompleteSession must + be stripped the moment the real session:start arrives, leaving + RootSession only. + """ + handler = SessionHandler(services) + + await handler( + "session:end", + {"session_id": "s-race-root", "timestamp": "2026-01-01T01:00:00Z"}, + ) + node = await services.graph.get_node("s-race-root") + assert node is not None + assert "IncompleteSession" in node["labels"], ( + "Precondition: out-of-order end must stamp the stale marker" + ) + + await handler( + "session:start", + {"session_id": "s-race-root", "timestamp": "2026-01-01T00:00:00Z"}, + ) + + node = await services.graph.get_node("s-race-root") + assert node is not None + labels = node["labels"] + assert "RootSession" in labels, "Real terminal must be assigned by start" + assert "IncompleteSession" not in labels, ( + "Heal-forward: session:start must strip the stale IncompleteSession " + "marker left by the out-of-order session:end" + ) + + async def test_out_of_order_end_then_start_with_parent_heals_to_subsession( + self, services: HookStateService + ) -> None: + """end processed BEFORE start (with parent): stale IncompleteSession + must be stripped, leaving SubSession only. + """ + handler = SessionHandler(services) + + await handler( + "session:end", + { + "session_id": "s-race-sub", + "parent_id": "p-race-sub", + "timestamp": "2026-01-01T01:00:00Z", + }, + ) + node = await services.graph.get_node("s-race-sub") + assert node is not None + assert "IncompleteSession" in node["labels"], ( + "Precondition: out-of-order end must stamp the stale marker" + ) + + await handler( + "session:start", + { + "session_id": "s-race-sub", + "parent_id": "p-race-sub", + "timestamp": "2026-01-01T00:00:00Z", + }, + ) + + node = await services.graph.get_node("s-race-sub") + assert node is not None + labels = node["labels"] + assert "SubSession" in labels, "Real terminal must be assigned by start" + assert "IncompleteSession" not in labels, ( + "Heal-forward: session:start must strip the stale IncompleteSession " + "marker left by the out-of-order session:end" + ) diff --git a/tests/integration/test_crash_recovery.py b/tests/integration/test_crash_recovery.py index 61fb544a..aeae3f1a 100644 --- a/tests/integration/test_crash_recovery.py +++ b/tests/integration/test_crash_recovery.py @@ -9,7 +9,7 @@ from unittest.mock import AsyncMock, patch from context_intelligence_server import registry as registry_module -from context_intelligence_server.queue_manager import QueueManager +from context_intelligence_server.queue_manager import FileSystemQueueManager, QueueManager from context_intelligence_server.registry import SessionRegistry, SessionWorker from context_intelligence_server.services import HookStateService @@ -23,7 +23,7 @@ async def test_no_loss_after_crash_mid_drain() -> None: registry_module.get_settings() ) # queues_path patched to tmp_path by safe_settings sid = "crash-sess" - qm = QueueManager(queues_dir=Path(settings.queues_path)) + qm = FileSystemQueueManager(queues_dir=Path(settings.queues_path)) K = 5 for i in range(K): @@ -94,7 +94,7 @@ async def test_offset_never_advances_over_undurable_data() -> None: buffer-restore-on-failure (neo4j_store.py:686-696).""" settings = registry_module.get_settings() sid = "flush-fail-sess" - qm = QueueManager(queues_dir=Path(settings.queues_path)) + qm = FileSystemQueueManager(queues_dir=Path(settings.queues_path)) await qm.append(sid, _line("e0", "/ws", {"session_id": sid})) reg = SessionRegistry() diff --git a/tests/integration/test_data_layer_3_delegation_skill.py b/tests/integration/test_data_layer_3_delegation_skill.py index ad56d110..14883482 100644 --- a/tests/integration/test_data_layer_3_delegation_skill.py +++ b/tests/integration/test_data_layer_3_delegation_skill.py @@ -236,8 +236,9 @@ async def test_skill_loaded_during_iteration_creates_e05(self) -> None: handlers, ) - # iteration_id = SESSION_ID::iteration::1 (first iteration) - iteration_id = f"{SESSION_ID}::iteration::1" + # A run is active (execution:start fired), so the iteration id is + # run-scoped and carries the run tiebreaker (seq 1 for the first run). + iteration_id = f"{SESSION_ID}::orch_run::{T1}::1::iteration::1" skill_load_id = f"{SESSION_ID}::skill::{SKILL_NAME}::{T2}" # Verify active_iteration_id was set by IterationHandler diff --git a/tests/integration/test_turn_chain.py b/tests/integration/test_turn_chain.py index 06ad16f8..04b2d4fd 100644 --- a/tests/integration/test_turn_chain.py +++ b/tests/integration/test_turn_chain.py @@ -48,9 +48,9 @@ # Computed node IDs matching the handlers' key conventions PROMPT_1_ID = f"{SESSION_ID}::prompt::{T1}" -RUN_1_ID = f"{SESSION_ID}::orch_run::{T2}" +RUN_1_ID = f"{SESSION_ID}::orch_run::{T2}::1" PROMPT_2_ID = f"{SESSION_ID}::prompt::{T8}" -RUN_2_ID = f"{SESSION_ID}::orch_run::{T9}" +RUN_2_ID = f"{SESSION_ID}::orch_run::{T9}::2" PROMPT_3_ID = f"{SESSION_ID}::prompt::{T11}" diff --git a/tests/neo4j/test_blob_reclaim_e2e.py b/tests/neo4j/test_blob_reclaim_e2e.py new file mode 100644 index 00000000..a27a4d81 --- /dev/null +++ b/tests/neo4j/test_blob_reclaim_e2e.py @@ -0,0 +1,144 @@ +"""End-to-end blob-reclaim against a REAL Neo4j graph and a REAL filesystem +blob store -- no mocks, no stubs. + +Proves the reshaped reclaim GC for real: the reference scan runs as Cypher +against a live graph, orphan selection uses the real QueueManager drain state, +and deletion goes through the real fenced BlobStore.delete -- the orphan file +actually disappears from disk while the referenced blob survives. + + uv run pytest tests/neo4j/test_blob_reclaim_e2e.py -q -m neo4j +""" + +from __future__ import annotations + +import asyncio +import json +import os +import time +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest +from neo4j import READ_ACCESS, AsyncGraphDatabase + +from context_intelligence_server.blob_store import create_blob_store +from context_intelligence_server.config import Settings +from context_intelligence_server.neo4j_store import ensure_neo4j_schema +from context_intelligence_server.queue_manager import FileSystemQueueManager +from context_intelligence_server.registry import SessionRegistry +from context_intelligence_server.routers import admin +from context_intelligence_server.routers.admin import BlobReclaimBody, reclaim_blobs + +pytestmark = pytest.mark.neo4j + +_WS = "reclaim_e2e" + + +def _settings(tmp_path: Path) -> Settings: + # The reclaim path reaches Neo4j via app.state.neo4j_query_driver, so only + # the filesystem roots matter here. + s = Settings() + s.blob_path = str(tmp_path / "blobs") + s.queues_path = str(tmp_path / "queues") + return s + + +def _build_registry(queues_dir: Path) -> SessionRegistry: + reg = SessionRegistry() + reg._queue_manager = FileSystemQueueManager(queues_dir=queues_dir) + reg._write_semaphore = asyncio.Semaphore(8) + reg._max_delivery_attempts = 3 + return reg + + +def _backdate(path: Path, minutes: int) -> None: + """Age a blob file past the reclaim min-age floor (>= 15 min).""" + old = time.time() - minutes * 60 + os.utime(path, (old, old)) + + +def _blob_file(blob_root: Path, session_id: str, key: str) -> Path: + return blob_root / session_id / "blobs" / f"{key}.json" + + +@pytest.fixture +async def _driver(neo4j_container: dict[str, Any]): + driver = AsyncGraphDatabase.driver( + neo4j_container["bolt_url"], + auth=(neo4j_container["user"], neo4j_container["password"]), + ) + await ensure_neo4j_schema(driver) + async with driver.session() as session: + await session.run("MATCH (n) DETACH DELETE n") + yield driver + async with driver.session() as session: + await session.run("MATCH (n) DETACH DELETE n") + await driver.close() + + +async def test_reclaim_deletes_orphan_keeps_referenced_e2e( + tmp_path: Path, neo4j_container: dict[str, Any], _driver, monkeypatch +) -> None: + settings = _settings(tmp_path) + monkeypatch.setattr(admin, "get_settings", lambda: settings) + + # Real filesystem blob store: write two real blobs. + store = create_blob_store(settings) + ref_kept = await store.write("sess_ref", "kept", {"payload": "referenced"}) + ref_orphan = await store.write("sess_orphan", "orphan", {"payload": "dangling"}) + + blob_root = Path(settings.blob_path) + kept_file = _blob_file(blob_root, "sess_ref", "kept") + orphan_file = _blob_file(blob_root, "sess_orphan", "orphan") + assert kept_file.exists() and orphan_file.exists() + + # Age both past the 15-minute floor so neither is skipped_recent. + _backdate(kept_file, 60) + _backdate(orphan_file, 60) + + # Reference ONLY the kept blob in the real graph, via a carrier property. + async with _driver.session() as session: + await session.run( + "CREATE (:Event {node_id: 'e1', workspace: $ws, session_id: 'sess_ref', " + "data: $data})", + ws=_WS, + data=json.dumps({"$blob_ref": ref_kept.uri}), + ) + + # Real registry + real queue manager; both sessions fully drained (no logs). + registry = _build_registry(tmp_path / "queues") + + request = SimpleNamespace( + # request.scope["state"]["contributor_id"] is read by the audit logger. + scope={"state": {"contributor_id": "admin"}}, + app=SimpleNamespace( + state=SimpleNamespace( + registry=registry, + neo4j_query_driver=_driver, + neo4j_query_access_mode="READ", + ) + ), + ) + # Sanity: the exact access-mode constant the endpoint will use. + assert admin._access_mode_const("READ") is READ_ACCESS + + # 1) DRY-RUN: finds exactly the orphan, deletes nothing. + admin._reclaim_apply_inflight = False + dry = await reclaim_blobs( + BlobReclaimBody(dry_run=True, min_age_minutes=15), request + ) + assert dry["dry_run"] is True + assert dry["orphans_found"] == 1, dry + assert ref_orphan.uri in dry["sample"] + assert ref_kept.uri not in dry["sample"] + assert orphan_file.exists() # nothing deleted in dry-run + + # 2) APPLY: the orphan file actually disappears; the referenced one survives. + applied = await reclaim_blobs( + BlobReclaimBody(dry_run=False, min_age_minutes=15, max_delete=10), request + ) + assert applied["deleted"] == 1, applied + assert not orphan_file.exists(), "orphan blob must be gone from disk" + assert kept_file.exists(), "referenced blob must survive" + assert admin._reclaim_apply_inflight is False # single-flight released diff --git a/tests/neo4j/test_concurrent_flush.py b/tests/neo4j/test_concurrent_flush.py index 470aa83d..d130a6a7 100644 --- a/tests/neo4j/test_concurrent_flush.py +++ b/tests/neo4j/test_concurrent_flush.py @@ -25,7 +25,7 @@ Neo4jGraphStore, ensure_neo4j_schema, ) -from context_intelligence_server.queue_manager import QueueManager +from context_intelligence_server.queue_manager import FileSystemQueueManager, QueueManager from context_intelligence_server.registry import SessionRegistry, SessionWorker from context_intelligence_server.services import HookStateService from context_intelligence_server.utils import make_node_id @@ -226,7 +226,7 @@ async def test_durable_drain_multi_writer_zero_loss( await driver.close() reg = SessionRegistry() - reg._queue_manager = QueueManager(queues_dir=tmp_path / "queues") + reg._queue_manager = FileSystemQueueManager(queues_dir=tmp_path / "queues") reg._write_semaphore = asyncio.Semaphore(2) reg._max_delivery_attempts = 5 qm = reg._queue_manager @@ -327,7 +327,7 @@ async def test_durable_poison_isolation_no_contamination( await driver.close() reg = SessionRegistry() - reg._queue_manager = QueueManager(queues_dir=tmp_path / "queues") + reg._queue_manager = FileSystemQueueManager(queues_dir=tmp_path / "queues") reg._write_semaphore = asyncio.Semaphore(2) reg._max_delivery_attempts = 2 qm = reg._queue_manager 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..6da04c9a --- /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 FileSystemQueueManager, 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 = FileSystemQueueManager(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_incomplete_session_heal_forward.py b/tests/neo4j/test_incomplete_session_heal_forward.py new file mode 100644 index 00000000..065a63f7 --- /dev/null +++ b/tests/neo4j/test_incomplete_session_heal_forward.py @@ -0,0 +1,191 @@ +"""Neo4j integration proof for IncompleteSession heal-forward. + +Encodes the out-of-order race that produces the ~99% false-positive +IncompleteSession population: a forked sub-session's session:end drains +(independent per-session queue) BEFORE its session:fork/session:start. +SessionLabelStateMachine.classify() must strip the stale IncompleteSession +marker the moment the real start/fork is processed +(out-of-order end-before-start/fork). + +This closes a real coverage gap: no prior real-Neo4j test exercised +set_labels() with a non-empty remove_labels list. It runs the real +SessionHandler (not just the pure classify() unit) against a live Neo4j +container, flushes, and reads back with a fresh Cypher query (bypassing the +in-memory buffer) to prove the label is PHYSICALLY removed from the node. + +It also proves non-interaction with the terminal-label lattice +normalization (_LATTICE_NORMALIZATION in neo4j_store.py): removing +IncompleteSession must not disturb the RootSession/SubSession/ForkedSession +convergence guarantee, since IncompleteSession is not a member of +_TERMINAL_LABELS. + +Run: uv run pytest tests/neo4j/test_incomplete_session_heal_forward.py -v -m neo4j +""" + +from __future__ import annotations + +from typing import Any + +import pytest + +from context_intelligence_server.handlers.data_layer_2.session import SessionHandler + +pytestmark = pytest.mark.neo4j + + +async def _neo4j_labels(services: Any, node_id: str) -> list[str]: + """Read labels directly from Neo4j (bypasses the in-memory buffer).""" + 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 [] + + +@pytest.mark.neo4j +class TestIncompleteSessionHealForward: + """Out-of-order end -> fork/start must physically strip IncompleteSession.""" + + async def test_out_of_order_end_then_fork_strips_incomplete_session_in_neo4j( + self, neo4j_services: Any + ) -> None: + """end (bare, stamps IncompleteSession) -> flush -> fork -> flush: + the real Neo4j node must end with ForkedSession only, no + IncompleteSession, proving REMOVE n:IncompleteSession actually ran. + """ + handler = SessionHandler(neo4j_services) + session_id = "child-heal-fork-001" + + # session:end drains first (out-of-order race) -- stamps IncompleteSession + await handler( + "session:end", + {"session_id": session_id, "timestamp": "2026-01-01T10:00:00Z"}, + ) + await neo4j_services.graph.flush() + + mid_labels = await _neo4j_labels(neo4j_services, session_id) + assert "IncompleteSession" in mid_labels, ( + f"precondition: out-of-order end must stamp IncompleteSession in " + f"Neo4j; got {mid_labels}" + ) + + # The real session:fork arrives late, in its own flush cycle. + await handler( + "session:fork", + { + "session_id": session_id, + "parent_id": "parent-heal-fork-001", + "timestamp": "2026-01-01T09:59:59Z", + }, + ) + await neo4j_services.graph.flush() + + final_labels = await _neo4j_labels(neo4j_services, session_id) + assert "ForkedSession" in final_labels, ( + f"expected ForkedSession in {final_labels}" + ) + assert "IncompleteSession" not in final_labels, ( + f"heal-forward must physically REMOVE n:IncompleteSession at " + f"flush; still present in Neo4j: {final_labels}" + ) + + async def test_out_of_order_end_then_start_strips_incomplete_session_in_neo4j( + self, neo4j_services: Any + ) -> None: + """end (bare, stamps IncompleteSession) -> flush -> start (no parent) + -> flush: the real Neo4j node must end with RootSession only. + """ + handler = SessionHandler(neo4j_services) + session_id = "root-heal-start-001" + + await handler( + "session:end", + {"session_id": session_id, "timestamp": "2026-01-01T10:00:00Z"}, + ) + await neo4j_services.graph.flush() + + mid_labels = await _neo4j_labels(neo4j_services, session_id) + assert "IncompleteSession" in mid_labels, ( + f"precondition: out-of-order end must stamp IncompleteSession in " + f"Neo4j; got {mid_labels}" + ) + + await handler( + "session:start", + {"session_id": session_id, "timestamp": "2026-01-01T09:59:59Z"}, + ) + await neo4j_services.graph.flush() + + final_labels = await _neo4j_labels(neo4j_services, session_id) + assert "RootSession" in final_labels, f"expected RootSession in {final_labels}" + assert "IncompleteSession" not in final_labels, ( + f"heal-forward must physically REMOVE n:IncompleteSession at " + f"flush; still present in Neo4j: {final_labels}" + ) + + async def test_heal_forward_does_not_disturb_terminal_lattice_normalization( + self, neo4j_services: Any + ) -> None: + """Healing IncompleteSession must not interact with, or break, the + RootSession/SubSession/ForkedSession lattice-normalization guarantee. + + Drives: end (bare, stamps IncompleteSession) -> flush -> start WITH a + parent (assigns SubSession) -> flush -> fork (reclassifies to + ForkedSession, the lattice's specificity ordering) -> flush. The node + must converge to exactly ONE terminal label (ForkedSession) with + IncompleteSession and the stale SubSession both absent -- proving the + IncompleteSession REMOVE and the terminal-lattice REMOVE/SET both ran + correctly and did not clobber each other. + """ + handler = SessionHandler(neo4j_services) + session_id = "lattice-heal-001" + parent_id = "lattice-heal-parent-001" + + await handler( + "session:end", + {"session_id": session_id, "timestamp": "2026-01-01T10:00:00Z"}, + ) + await neo4j_services.graph.flush() + + await handler( + "session:start", + { + "session_id": session_id, + "parent_id": parent_id, + "timestamp": "2026-01-01T09:59:58Z", + }, + ) + await neo4j_services.graph.flush() + + mid_labels = await _neo4j_labels(neo4j_services, session_id) + assert "SubSession" in mid_labels, f"expected SubSession in {mid_labels}" + assert "IncompleteSession" not in mid_labels, ( + f"IncompleteSession must already be healed after start: {mid_labels}" + ) + + await handler( + "session:fork", + { + "session_id": session_id, + "parent_id": parent_id, + "timestamp": "2026-01-01T09:59:59Z", + }, + ) + await neo4j_services.graph.flush() + + final_labels = await _neo4j_labels(neo4j_services, session_id) + terminals = [ + lbl + for lbl in final_labels + if lbl in ("RootSession", "SubSession", "ForkedSession") + ] + assert terminals == ["ForkedSession"], ( + f"lattice must converge to exactly one terminal (ForkedSession); " + f"got {terminals} in {final_labels}" + ) + assert "IncompleteSession" not in final_labels, ( + f"IncompleteSession must remain healed through the lattice " + f"reclassification: {final_labels}" + ) diff --git a/tests/neo4j/test_migrations_cli_e2e.py b/tests/neo4j/test_migrations_cli_e2e.py new file mode 100644 index 00000000..3577e3cf --- /dev/null +++ b/tests/neo4j/test_migrations_cli_e2e.py @@ -0,0 +1,122 @@ +"""End-to-end migration CLI against a REAL Neo4j -- no mocked driver. + +The unit test (tests/test_migrations_run.py) drives ``_amain`` with a mocked +driver and mocked ``run_repair``; this reconciles those assumptions with +reality: the CLI builds its OWN real driver from ``--neo4j-*`` flags, runs +``--status`` (read-only) then ``--apply`` against a genuinely dirty graph, and +the graph is actually rectified. Re-running ``--apply`` is a real no-op. + + uv run pytest tests/neo4j/test_migrations_cli_e2e.py -q -m neo4j +""" + +from __future__ import annotations + +from typing import Any + +import migrations.run as run_module +import pytest +from neo4j import GraphDatabase + +pytestmark = pytest.mark.neo4j + +_WS = "migrate_e2e" + + +def _wipe(container: dict[str, Any]) -> None: + driver = GraphDatabase.driver( + container["bolt_url"], auth=(container["user"], container["password"]) + ) + try: + with driver.session() as s: + s.run("MATCH (n) DETACH DELETE n") + for rec in list(s.run("SHOW CONSTRAINTS YIELD name RETURN name")): + s.run(f"DROP CONSTRAINT {rec['name']} IF EXISTS") + for rec in list(s.run("SHOW INDEXES YIELD name RETURN name")): + try: + s.run(f"DROP INDEX {rec['name']} IF EXISTS") + except Exception: # noqa: BLE001 - constraint-backed/lookup indexes + pass + finally: + driver.close() + + +def _seed_dirty(container: dict[str, Any]) -> None: + """Two duplicate untagged :Event nodes + one legacy :Session (no :Node).""" + driver = GraphDatabase.driver( + container["bolt_url"], auth=(container["user"], container["password"]) + ) + try: + with driver.session() as s: + s.run("CREATE (:Event {node_id: 'dup-1', workspace: $ws, v: 1})", ws=_WS) + s.run("CREATE (:Event {node_id: 'dup-1', workspace: $ws, v: 2})", ws=_WS) + s.run("CREATE (:Session {node_id: 'legacy-sess', workspace: $ws})", ws=_WS) + finally: + driver.close() + + +def _graph_state(container: dict[str, Any]) -> dict[str, int]: + driver = GraphDatabase.driver( + container["bolt_url"], auth=(container["user"], container["password"]) + ) + try: + with driver.session() as s: + untagged = s.run( + "MATCH (n) WHERE n.node_id IS NOT NULL AND NOT n:Node RETURN count(n) AS c" + ).single()["c"] + dup1 = s.run( + "MATCH (n {node_id: 'dup-1', workspace: $ws}) RETURN count(n) AS c", + ws=_WS, + ).single()["c"] + constraint = len( + list(s.run("SHOW CONSTRAINTS YIELD name RETURN name")) + ) + return {"untagged": untagged, "dup1": dup1, "constraints": constraint} + finally: + driver.close() + + +def _args(container: dict[str, Any], flag: str): + return run_module.build_parser().parse_args( + [ + flag, + "--neo4j-url", + container["bolt_url"], + "--neo4j-user", + container["user"], + "--neo4j-password", + container["password"], + ] + ) + + +async def test_migration_cli_status_then_apply_rectifies_real_graph( + neo4j_container: dict[str, Any], +) -> None: + _wipe(neo4j_container) + try: + _seed_dirty(neo4j_container) + + before = _graph_state(neo4j_container) + assert before["untagged"] >= 3, before + assert before["dup1"] == 2, before + + # --status is read-only: reports, never mutates. + code = await run_module._amain(_args(neo4j_container, "--status")) + assert code == 0 + after_status = _graph_state(neo4j_container) + assert after_status == before, "status must not touch the graph" + + # --apply rectifies for real. + code = await run_module._amain(_args(neo4j_container, "--apply")) + assert code == 0 + after_apply = _graph_state(neo4j_container) + assert after_apply["untagged"] == 0, after_apply + assert after_apply["dup1"] == 1, after_apply + assert after_apply["constraints"] >= 1, after_apply + + # Re-running --apply is a real idempotent no-op (still clean, still 0). + code = await run_module._amain(_args(neo4j_container, "--apply")) + assert code == 0 + assert _graph_state(neo4j_container) == after_apply + finally: + _wipe(neo4j_container) diff --git a/tests/neo4j/test_oom_regression.py b/tests/neo4j/test_oom_regression.py index b7fd2e1b..94f7b0ea 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, @@ -306,7 +310,7 @@ async def test_finalization_path_freezes_then_restart_then_drains( from unittest.mock import AsyncMock from context_intelligence_server.pipeline import setup_handlers - from context_intelligence_server.queue_manager import QueueManager + from context_intelligence_server.queue_manager import FileSystemQueueManager, QueueManager from context_intelligence_server.registry import SessionRegistry, SessionWorker from context_intelligence_server.services import HookStateService @@ -320,7 +324,7 @@ async def test_finalization_path_freezes_then_restart_then_drains( # fixture) patches into get_settings().queues_path for every new # SessionRegistry() instance constructed in this test. queues_dir = tmp_path / "queues" - qm = QueueManager(queues_dir=queues_dir) + qm = FileSystemQueueManager(queues_dir=queues_dir) async def _make( rows: int, byts: int diff --git a/tests/neo4j/test_orphan_visibility.py b/tests/neo4j/test_orphan_visibility.py index a2ced1e0..188da47b 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.queue_manager import FileSystemQueueManager, 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 @@ -145,7 +145,7 @@ async def test_finalization_orphan_surfaces_on_status( # guarantees this QueueManager and the registry's queue_manager point at # the same on-disk queue. queues_dir = tmp_path / "queues" - qm = QueueManager(queues_dir=queues_dir) + qm = FileSystemQueueManager(queues_dir=queues_dir) # ----------------------------------------------------------------------- # Seed the durable queue in exact order (line counts are load-bearing). @@ -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_queues_actions.py b/tests/neo4j/test_queues_actions.py index 6276bbc1..59e7c41a 100644 --- a/tests/neo4j/test_queues_actions.py +++ b/tests/neo4j/test_queues_actions.py @@ -34,7 +34,7 @@ Neo4jGraphStore, ensure_neo4j_schema, ) -from context_intelligence_server.queue_manager import QueueManager +from context_intelligence_server.queue_manager import FileSystemQueueManager, QueueManager from context_intelligence_server.registry import SessionRegistry, SessionWorker from context_intelligence_server.services import HookStateService @@ -96,7 +96,7 @@ async def test_replay_rewrites_through_real_drainer( # worker is pre-registered, so get_or_create never builds a settings-derived # store). reg = SessionRegistry() - reg._queue_manager = QueueManager(queues_dir=tmp_path / "queues") + reg._queue_manager = FileSystemQueueManager(queues_dir=tmp_path / "queues") reg._write_semaphore = asyncio.Semaphore(1) reg._max_delivery_attempts = 3 qm = reg._queue_manager diff --git a/tests/neo4j/test_relabel_incomplete_sessions.py b/tests/neo4j/test_relabel_incomplete_sessions.py new file mode 100644 index 00000000..02aba124 --- /dev/null +++ b/tests/neo4j/test_relabel_incomplete_sessions.py @@ -0,0 +1,488 @@ +"""Neo4j integration test -- IncompleteSession backfill. + +Seeds real Session/SessionStartEvent/SessionForkEvent nodes and SOURCED_FROM +edges to exercise scripts/relabel_incomplete_sessions.py against a live +Neo4j container. Covers six scenarios: + +(a) forked node with a real terminal label + stale marker -> cleared by the + gated --apply path (run_apply). +(b) node with a linked SessionStartEvent (no terminal label yet) + stale + marker -> cleared by the raw selector/write function (apply_relabel), + proving the raw selector's mechanics work on this node shape in isolation. +(c) genuine no-start/no-fork node -> RETAINED. +(d) idempotence: a second apply_relabel() call matches zero rows. +(e) the reconciliation gate: seed a linked_but_untyped node and assert + the gated --apply path (run_apply) REFUSES and makes NO write. +(f) --restore re-adds :IncompleteSession to exactly the touched ids. + +Note: (b) and (e) seed the SAME node shape (IncompleteSession + a linked +start/fork event + no terminal label -- exactly "linked_but_untyped"). They +are deliberately not a contradiction: (b) proves the underlying raw selector +mechanically heals that shape when invoked directly; (e) proves the +reconciliation safety gate refuses the FULL --apply workflow when such a node +exists anywhere in the graph (the gate exists precisely because the +assumption is unverified on historical data, independent of whether the +mechanics work). + +Run: uv run pytest tests/neo4j/test_relabel_incomplete_sessions.py -v -m neo4j +""" + +from __future__ import annotations + +from typing import Any + +import pytest +from neo4j import GraphDatabase +from scripts import relabel_incomplete_sessions as relabel + +WORKSPACE = "test" + +pytestmark = pytest.mark.neo4j + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _driver(neo4j_container: dict) -> GraphDatabase: + """Return a synchronous Neo4j driver for the test container.""" + return GraphDatabase.driver( + neo4j_container["bolt_url"], + auth=(neo4j_container["user"], neo4j_container["password"]), + ) + + +def _seed_session(session, node_id: str, extra_labels: list[str] | None = None) -> None: + """Seed a bare :Session:IncompleteSession node, optionally with extra labels.""" + labels = "".join(f":{lbl}" for lbl in (extra_labels or [])) + session.run( + f"MERGE (s:Session:IncompleteSession{labels} " + "{node_id: $node_id, workspace: $workspace})", + node_id=node_id, + workspace=WORKSPACE, + ) + + +def _seed_linked_event( + session, session_node_id: str, event_label: str, event_node_id: str +) -> None: + """Seed a SessionStartEvent/SessionForkEvent node and the real + (Session)-[:SOURCED_FROM]->(Event) edge direction (see the module + docstring's "Selector direction correction" for why this direction, not + the reverse, is what the shipping write path actually produces). + """ + session.run( + f"MATCH (s:Session {{node_id: $session_node_id, workspace: $workspace}}) " + f"MERGE (e:{event_label} {{node_id: $event_node_id, workspace: $workspace}}) " + "MERGE (s)-[:SOURCED_FROM]->(e)", + session_node_id=session_node_id, + event_node_id=event_node_id, + workspace=WORKSPACE, + ) + + +def _labels(session, node_id: str) -> list[str]: + """Return the labels of a node matched by node_id + workspace, or [].""" + result = session.run( + "MATCH (n {node_id: $node_id, workspace: $workspace}) RETURN labels(n) AS lbls", + node_id=node_id, + workspace=WORKSPACE, + ) + record = result.single() + if record is None: + return [] + return list(record["lbls"]) + + +# --------------------------------------------------------------------------- +# Per-test isolation +# --------------------------------------------------------------------------- + + +@pytest.fixture(autouse=True) +def _clear_neo4j(neo4j_container: dict) -> None: # type: ignore[return] + """Wipe the container clean before each test for complete isolation.""" + driver = GraphDatabase.driver( + neo4j_container["bolt_url"], + auth=(neo4j_container["user"], neo4j_container["password"]), + ) + try: + with driver.session() as s: + s.run("MATCH (n) DETACH DELETE n") + finally: + driver.close() + + +# --------------------------------------------------------------------------- +# (a) forked node with real terminal + stale marker -> cleared by --apply +# --------------------------------------------------------------------------- + + +@pytest.mark.neo4j +class TestApplyClearsTerminalLabeledNode: + def test_forked_node_with_terminal_label_is_healed_by_gated_apply( + self, neo4j_container: dict[str, Any], tmp_path: Any + ) -> None: + driver = _driver(neo4j_container) + try: + with driver.session() as s: + _seed_session(s, "sess-a", extra_labels=["ForkedSession"]) + + exit_code = relabel.run_apply( + s, + batch_size=relabel.DEFAULT_BATCH_SIZE, + undo_log_path=str(tmp_path / "undo-a.json"), + neo4j_url=neo4j_container["bolt_url"], + ) + + assert exit_code == 0 + labels = _labels(s, "sess-a") + assert "ForkedSession" in labels + assert "IncompleteSession" not in labels + finally: + driver.close() + + +# --------------------------------------------------------------------------- +# (b) linked SessionStartEvent, no terminal label -> cleared by the raw +# selector/write function (apply_relabel), proving the raw selector +# mechanics. +# --------------------------------------------------------------------------- + + +@pytest.mark.neo4j +class TestApplyRelabelClearsLinkedUntypedNode: + def test_linked_start_event_no_terminal_label_healed_by_raw_selector( + self, neo4j_container: dict[str, Any] + ) -> None: + driver = _driver(neo4j_container) + try: + with driver.session() as s: + _seed_session(s, "sess-b") + _seed_linked_event(s, "sess-b", "SessionStartEvent", "sess-b::start") + + touched = relabel.apply_relabel( + s, batch_size=relabel.DEFAULT_BATCH_SIZE + ) + + assert touched == ["sess-b"] + labels = _labels(s, "sess-b") + assert "IncompleteSession" not in labels + assert "Session" in labels + finally: + driver.close() + + +# --------------------------------------------------------------------------- +# (c) genuine no-start/no-fork node -> RETAINED +# --------------------------------------------------------------------------- + + +@pytest.mark.neo4j +class TestGenuineIncompleteRetained: + def test_genuine_node_with_no_linked_event_and_no_terminal_is_retained( + self, neo4j_container: dict[str, Any], tmp_path: Any + ) -> None: + driver = _driver(neo4j_container) + try: + with driver.session() as s: + _seed_session(s, "sess-c") + + exit_code = relabel.run_apply( + s, + batch_size=relabel.DEFAULT_BATCH_SIZE, + undo_log_path=str(tmp_path / "undo-c.json"), + neo4j_url=neo4j_container["bolt_url"], + ) + + assert exit_code == 0 + labels = _labels(s, "sess-c") + assert "IncompleteSession" in labels, ( + "genuine incomplete session (no linked start/fork event) " + f"must be retained; got {labels}" + ) + finally: + driver.close() + + +# --------------------------------------------------------------------------- +# (d) idempotence: second apply_relabel() matches zero rows +# --------------------------------------------------------------------------- + + +@pytest.mark.neo4j +class TestIdempotence: + def test_second_apply_relabel_touches_zero_rows( + self, neo4j_container: dict[str, Any] + ) -> None: + driver = _driver(neo4j_container) + try: + with driver.session() as s: + _seed_session(s, "sess-d", extra_labels=["ForkedSession"]) + + first = relabel.apply_relabel(s, batch_size=relabel.DEFAULT_BATCH_SIZE) + second = relabel.apply_relabel(s, batch_size=relabel.DEFAULT_BATCH_SIZE) + + assert first == ["sess-d"] + assert second == [] + labels = _labels(s, "sess-d") + assert "IncompleteSession" not in labels + assert "ForkedSession" in labels + finally: + driver.close() + + +# --------------------------------------------------------------------------- +# (e) reconciliation gate: linked_but_untyped node -> --apply REFUSES +# --------------------------------------------------------------------------- + + +@pytest.mark.neo4j +class TestReconciliationGateRefusesApply: + def test_linked_but_untyped_node_causes_gated_apply_to_refuse( + self, neo4j_container: dict[str, Any], tmp_path: Any + ) -> None: + driver = _driver(neo4j_container) + try: + with driver.session() as s: + # Same shape as (b) -- IncompleteSession + linked start event, + # no terminal label -- but exercised through the GATED + # run_apply() path instead of the raw apply_relabel(). + _seed_session(s, "sess-e") + _seed_linked_event(s, "sess-e", "SessionStartEvent", "sess-e::start") + + diag = relabel.diagnostic_report(s) + assert diag["linked_but_untyped"] == 1 + assert diag["linked_but_untyped_samples"] == ["sess-e"] + + exit_code = relabel.run_apply( + s, + batch_size=relabel.DEFAULT_BATCH_SIZE, + undo_log_path=str(tmp_path / "undo-e.json"), + neo4j_url=neo4j_container["bolt_url"], + ) + + assert exit_code == 1, ( + "run_apply must REFUSE (non-zero exit) when the " + "reconciliation diagnostic finds a linked_but_untyped node" + ) + # No write must have occurred: the marker is still present. + labels = _labels(s, "sess-e") + assert "IncompleteSession" in labels, ( + "gate refusal must not mutate the graph; " + f"IncompleteSession missing from {labels}" + ) + # The undo-log must not have been written either. + assert not (tmp_path / "undo-e.json").exists() + finally: + driver.close() + + +# --------------------------------------------------------------------------- +# (f) --restore re-adds :IncompleteSession to exactly the touched ids +# --------------------------------------------------------------------------- + + +@pytest.mark.neo4j +class TestRestore: + def test_restore_re_adds_marker_to_exactly_the_touched_ids( + self, neo4j_container: dict[str, Any], tmp_path: Any + ) -> None: + driver = _driver(neo4j_container) + try: + with driver.session() as s: + _seed_session(s, "sess-f", extra_labels=["ForkedSession"]) + undo_log_path = tmp_path / "undo-f.json" + + exit_code = relabel.run_apply( + s, + batch_size=relabel.DEFAULT_BATCH_SIZE, + undo_log_path=str(undo_log_path), + neo4j_url=neo4j_container["bolt_url"], + ) + assert exit_code == 0 + assert undo_log_path.exists() + + healed_labels = _labels(s, "sess-f") + assert "IncompleteSession" not in healed_labels + assert "ForkedSession" in healed_labels + + snap = relabel.read_undo_log(str(undo_log_path)) + assert snap["node_ids"] == ["sess-f"] + assert "generated_at" in snap + assert "neo4j_host" in snap + + restore_exit_code = relabel.run_restore( + s, str(undo_log_path), batch_size=relabel.DEFAULT_BATCH_SIZE + ) + assert restore_exit_code == 0 + + restored_labels = _labels(s, "sess-f") + assert "IncompleteSession" in restored_labels + # The terminal label must never have been touched. + assert "ForkedSession" in restored_labels + finally: + driver.close() + + +# --------------------------------------------------------------------------- +# (g) undo-log is written from a read-only PRE-MUTATION candidate +# collection, EXACTLY matching the full false-positive set, and restore +# from it fully re-adds the label to every candidate. +# --------------------------------------------------------------------------- + + +@pytest.mark.neo4j +class TestUndoLogCapturesFullPreMutationCandidateSet: + """apply_relabel()'s batched CALL {...} IN TRANSACTIONS OF N ROWS + commits per batch, so the undo-log must exist -- complete -- the moment + any batch has committed, not only after apply_relabel() returns. This + test seeds a NON-VACUOUS candidate set (two distinct false-positive + shapes that both clear the reconciliation gate: a terminal-labeled node, and a + terminal-labeled node that ALSO has a linked start event) and asserts + the undo-log written by the gated --apply path (run_apply) contains + EXACTLY that candidate set, and that restoring from it fully re-adds + :IncompleteSession to both. + """ + + def test_undo_log_exactly_matches_candidates_and_restore_heals_all( + self, neo4j_container: dict[str, Any], tmp_path: Any + ) -> None: + driver = _driver(neo4j_container) + try: + with driver.session() as s: + # Non-vacuity: two distinct false-positive shapes, so the + # candidate set has more than one member and is not a + # degenerate single-node case. Both carry a terminal label + # (so NEITHER trips the linked_but_untyped gate); the + # second additionally has a linked start event, to prove the + # selector's OR-shape doesn't affect completeness. + _seed_session(s, "sess-g1", extra_labels=["ForkedSession"]) + _seed_session(s, "sess-g2", extra_labels=["SubSession"]) + _seed_linked_event(s, "sess-g2", "SessionStartEvent", "sess-g2::start") + + # Ground truth: the pre-mutation candidate set, read + # independently via the same read-only helper run_apply + # uses internally, BEFORE run_apply does anything. + expected_candidates = sorted(relabel.collect_false_positive_ids(s)) + assert expected_candidates == ["sess-g1", "sess-g2"], ( + "test setup must be non-vacuous: exactly two " + f"false-positive candidates expected, got {expected_candidates}" + ) + + undo_log_path = tmp_path / "undo-g.json" + exit_code = relabel.run_apply( + s, + batch_size=relabel.DEFAULT_BATCH_SIZE, + undo_log_path=str(undo_log_path), + neo4j_url=neo4j_container["bolt_url"], + ) + assert exit_code == 0 + assert undo_log_path.exists(), "undo-log file must exist after --apply" + + snap = relabel.read_undo_log(str(undo_log_path)) + logged_ids = sorted(snap["node_ids"]) + # (a) The undo log holds EXACTLY the candidate set -- no + # more, no less. + assert logged_ids == expected_candidates == ["sess-g1", "sess-g2"] + + # Both were actually mutated. + assert "IncompleteSession" not in _labels(s, "sess-g1") + assert "IncompleteSession" not in _labels(s, "sess-g2") + + # (b) Restoring from the undo log fully re-adds the marker + # to every logged node_id. + restore_exit_code = relabel.run_restore( + s, str(undo_log_path), batch_size=relabel.DEFAULT_BATCH_SIZE + ) + assert restore_exit_code == 0 + assert "IncompleteSession" in _labels(s, "sess-g1") + assert "IncompleteSession" in _labels(s, "sess-g2") + finally: + driver.close() + + +@pytest.mark.neo4j +class TestUndoLogSourcedFromPreMutationReadNotApplyReturnValue: + """The undo-log must be sourced from a read-only, pre-mutation + collection of the candidate set (``collect_false_positive_ids``) -- NOT + derived from whatever ``apply_relabel()`` happens to report as touched. + + Proven by monkeypatching ``apply_relabel`` to still perform the REAL + mutation (so the graph state and restore path stay meaningful) but + report a deliberately WRONG, truncated touched set. Under the pre-fix + ordering (log written from apply_relabel's return value, AFTER the + mutation), the undo-log would contain only the truncated set. Under the + fix, run_apply captures the full candidate set via a separate read + BEFORE calling apply_relabel at all, so the log is unaffected by + whatever apply_relabel reports. + """ + + def test_undo_log_unaffected_by_apply_relabels_return_value( + self, + neo4j_container: dict[str, Any], + tmp_path: Any, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + driver = _driver(neo4j_container) + try: + with driver.session() as s: + _seed_session(s, "sess-h1", extra_labels=["ForkedSession"]) + _seed_session(s, "sess-h2", extra_labels=["RootSession"]) + + expected_candidates = sorted(relabel.collect_false_positive_ids(s)) + assert expected_candidates == ["sess-h1", "sess-h2"], ( + "test setup must be non-vacuous: exactly two " + f"false-positive candidates expected, got {expected_candidates}" + ) + + real_apply_relabel = relabel.apply_relabel + + def _fake_apply_relabel(session, batch_size): + # Perform the REAL mutation (so the restore-path + # assertions below stay meaningful) but report a + # deliberately WRONG, truncated touched set -- the + # shape of divergence a mid-run crash would produce if + # the undo-log were sourced from this return value. + real_apply_relabel(session, batch_size) + return ["sess-h1"] # deliberately omits sess-h2 + + monkeypatch.setattr(relabel, "apply_relabel", _fake_apply_relabel) + + undo_log_path = tmp_path / "undo-h.json" + exit_code = relabel.run_apply( + s, + batch_size=relabel.DEFAULT_BATCH_SIZE, + undo_log_path=str(undo_log_path), + neo4j_url=neo4j_container["bolt_url"], + ) + assert exit_code == 0 + assert undo_log_path.exists() + + snap = relabel.read_undo_log(str(undo_log_path)) + logged_ids = sorted(snap["node_ids"]) + + # THE core assertion: the log holds the FULL pre-mutation + # candidate set, not apply_relabel's (faked, truncated) + # return value. If run_apply sourced the log from + # touched_ids (the pre-fix ordering), logged_ids would + # equal ["sess-h1"] here. + assert logged_ids == expected_candidates == ["sess-h1", "sess-h2"] + assert logged_ids != ["sess-h1"], ( + "undo log must not be derived from apply_relabel's " + "reported touched set" + ) + + # Restoring from the (correctly complete) log heals BOTH + # nodes, including "sess-h2" which the faked touched-set + # omitted -- proving the log's completeness has real + # recovery value, not just structural equality. + restore_exit_code = relabel.run_restore( + s, str(undo_log_path), batch_size=relabel.DEFAULT_BATCH_SIZE + ) + assert restore_exit_code == 0 + assert "IncompleteSession" in _labels(s, "sess-h1") + assert "IncompleteSession" in _labels(s, "sess-h2") + finally: + driver.close() diff --git a/tests/neo4j/test_schema_version_baseline.py b/tests/neo4j/test_schema_version_baseline.py new file mode 100644 index 00000000..d26f5b36 --- /dev/null +++ b/tests/neo4j/test_schema_version_baseline.py @@ -0,0 +1,198 @@ +"""Neo4j integration proof for the SchemaMeta baseline singleton. + +BASELINE DATA POINTS ONLY: this proves ``ensure_schema_version_baseline``'s +``:SchemaMeta {id: 'singleton'}`` write is create-if-absent (``ON CREATE SET`` +only, no ``ON MATCH SET``) and O(1) -- calling it twice must leave exactly one +node in place with an UNCHANGED `last_updated`, proving the second call did +not clobber it. It also proves the uniqueness constraint on +``(:SchemaMeta).id`` is actually created, and that the constraint is what +makes the singleton race-free under real concurrency (N concurrent callers +still leave exactly one node) -- the whole point of this hardening +follow-up: moving the write off the per-worker-flush path onto a +single-writer startup call, backed by a uniqueness constraint so even a +violation of that single-writer invariant could never create a duplicate. + +Also proves ``ensure_neo4j_schema`` itself no longer writes SchemaMeta at +all -- that responsibility now belongs exclusively to +``ensure_schema_version_baseline``, called once from the lifespan startup +handler, never from the per-flush / doctor-repair paths that call +``ensure_neo4j_schema``. + +No comparison/upgrade/migration logic is exercised here; there is none to +exercise -- that is the point of this test. + +Run: uv run pytest tests/neo4j/test_schema_version_baseline.py -v -m neo4j +""" + +from __future__ import annotations + +import asyncio +from typing import Any + +import pytest +from context_intelligence_server.neo4j_store import ( + ensure_neo4j_schema, + ensure_schema_version_baseline, +) +from context_intelligence_server.status import SCHEMA_VERSION +from neo4j import AsyncGraphDatabase + + +async def _schema_meta_rows(driver: Any) -> list[Any]: + """Return all :SchemaMeta{id:'singleton'} rows (schema_version, last_updated).""" + async with driver.session() as session: + result = await session.run( + "MATCH (m:SchemaMeta {id: 'singleton'}) " + "RETURN m.schema_version AS schema_version, " + "m.last_updated AS last_updated" + ) + return [record async for record in result] + + +@pytest.mark.neo4j +class TestSchemaMetaBaselineIdempotence: + """``ensure_schema_version_baseline`` writes the singleton create-if-absent only.""" + + async def test_second_call_does_not_clobber_first( + self, neo4j_container: dict[str, Any] + ) -> None: + auth = (neo4j_container["user"], neo4j_container["password"]) + bolt = neo4j_container["bolt_url"] + + driver = AsyncGraphDatabase.driver(bolt, auth=auth) + try: + # First call: creates the singleton (and the uniqueness constraint). + await ensure_schema_version_baseline(driver) + + rows = await _schema_meta_rows(driver) + assert len(rows) == 1, ( + f"expected exactly one :SchemaMeta singleton after first call, " + f"got {len(rows)}" + ) + assert rows[0]["schema_version"] == SCHEMA_VERSION + first_last_updated = rows[0]["last_updated"] + assert first_last_updated is not None + + # Second call: must be a no-op on the existing node (ON CREATE only). + await ensure_schema_version_baseline(driver) + + rows_after = await _schema_meta_rows(driver) + assert len(rows_after) == 1, ( + f"expected exactly one :SchemaMeta node after second call " + f"(no duplicate created), got {len(rows_after)}" + ) + assert rows_after[0]["last_updated"] == first_last_updated, ( + "last_updated changed on the second call -- ON MATCH SET must " + "not be present; the singleton must be left untouched once it " + "exists" + ) + assert rows_after[0]["schema_version"] == SCHEMA_VERSION + finally: + await driver.close() + + async def test_uniqueness_constraint_exists( + self, neo4j_container: dict[str, Any] + ) -> None: + """The (:SchemaMeta).id uniqueness constraint is created and enforced.""" + auth = (neo4j_container["user"], neo4j_container["password"]) + bolt = neo4j_container["bolt_url"] + + driver = AsyncGraphDatabase.driver(bolt, auth=auth) + try: + await ensure_schema_version_baseline(driver) + + async with driver.session() as session: + result = await session.run("SHOW CONSTRAINTS") + constraints = [record async for record in result] + + schema_meta_constraints = [ + c + for c in constraints + if "SchemaMeta" in (c.get("labelsOrTypes") or []) + and "id" in (c.get("properties") or []) + ] + assert schema_meta_constraints, ( + "expected a uniqueness constraint on (:SchemaMeta).id to exist " + f"after ensure_schema_version_baseline; SHOW CONSTRAINTS returned: " + f"{constraints}" + ) + + # Belt-and-suspenders: attempting to create a second singleton node + # directly (bypassing MERGE) must be rejected by the constraint. + with pytest.raises(Exception): # noqa: B017 - Neo4jError subtype + async with driver.session() as session: + await session.run("CREATE (m:SchemaMeta {id: 'singleton'})") + finally: + await driver.close() + + async def test_concurrent_calls_create_exactly_one_node( + self, neo4j_container: dict[str, Any] + ) -> None: + """N concurrent baseline calls against the same DB leave exactly one node. + + This is the whole point of the uniqueness-constraint hardening: without + it, concurrent MERGEs on a fresh singleton key can each pass the + existence check and create divergent duplicate nodes. Fire many + concurrent calls (each on its own driver, mirroring independent + SessionWorker stores) and assert the constraint prevents any + duplication. + """ + auth = (neo4j_container["user"], neo4j_container["password"]) + bolt = neo4j_container["bolt_url"] + + n_concurrent = 20 + drivers = [ + AsyncGraphDatabase.driver(bolt, auth=auth) for _ in range(n_concurrent) + ] + try: + await asyncio.gather(*(ensure_schema_version_baseline(d) for d in drivers)) + + rows = await _schema_meta_rows(drivers[0]) + assert len(rows) == 1, ( + f"expected exactly one :SchemaMeta singleton after " + f"{n_concurrent} concurrent calls, got {len(rows)} -- the " + "uniqueness constraint should make concurrent creation race-free" + ) + assert rows[0]["schema_version"] == SCHEMA_VERSION + finally: + for d in drivers: + await d.close() + + +@pytest.mark.neo4j +class TestEnsureNeo4jSchemaNoLongerWritesSchemaMeta: + """``ensure_neo4j_schema`` must not touch :SchemaMeta at all (hardening follow-up). + + That responsibility moved exclusively to ``ensure_schema_version_baseline``, + called once from the lifespan startup handler -- NOT from the per-flush / + doctor-repair paths that call ``ensure_neo4j_schema``. If ``ensure_neo4j_schema`` + still created the singleton, it would fire redundantly (and concurrently) on + every SessionWorker's first flush. + """ + + async def test_ensure_neo4j_schema_does_not_create_schema_meta( + self, neo4j_container: dict[str, Any] + ) -> None: + auth = (neo4j_container["user"], neo4j_container["password"]) + bolt = neo4j_container["bolt_url"] + + driver = AsyncGraphDatabase.driver(bolt, auth=auth) + try: + # Remove any pre-existing singleton so this test is unambiguous + # regardless of what earlier tests in this (session-scoped) + # container have already written. + async with driver.session() as session: + await session.run( + "MATCH (m:SchemaMeta {id: 'singleton'}) DETACH DELETE m" + ) + + await ensure_neo4j_schema(driver) + + rows = await _schema_meta_rows(driver) + assert rows == [], ( + "ensure_neo4j_schema must not write the :SchemaMeta singleton " + f"-- that is ensure_schema_version_baseline's job now, but " + f"found {len(rows)} node(s)" + ) + finally: + await driver.close() 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..5c07ce0e --- /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 FileSystemQueueManager, 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 = FileSystemQueueManager(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 = FileSystemQueueManager(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/neo4j/test_tag_legacy_pooled_iterations.py b/tests/neo4j/test_tag_legacy_pooled_iterations.py new file mode 100644 index 00000000..53468d09 --- /dev/null +++ b/tests/neo4j/test_tag_legacy_pooled_iterations.py @@ -0,0 +1,254 @@ +"""Live Neo4j test coverage for scripts/tag_legacy_pooled_iterations.py. + +The script's own module docstring states it has **NO test coverage** ("NOT +product code -- no unit tests"), despite mutating a live graph (it ``SET``s a +non-destructive marker property on a confirmed-corrupt subset of :Iteration +nodes). This module closes that gap against a REAL Neo4j test container (see +tests/neo4j/conftest.py), mirroring the direct-function-import pattern used +by ``test_relabel_incomplete_sessions.py`` -- both scripts share the same +``CALL { ... } IN TRANSACTIONS OF N ROWS`` batching shape and both expose +their session-taking functions (``run_dry_run``, ``run_apply``, +``tag_confirmed_corrupt``) as directly importable/testable seams, so no +subprocess/CLI invocation is needed. + +Covers: + +(a) Selector precision -- an :Iteration node reached by ``HAS_PART`` from + ``>=2`` distinct :OrchestratorRun nodes (the confirmed-corrupt / legacy + pooled shape) IS tagged, and a node with exactly ONE run parent (the + clean, single-run shape that is the large majority of bare-id nodes in + live data per the script's own docstring) is NEVER tagged. Both shapes + are seeded in the SAME graph so the selector's precision is proven, not + just its recall. +(b) The --apply gate -- ``run_dry_run()`` (the default / non-apply + invocation) performs ZERO writes: the pooled node is left untagged. + Only ``run_apply()`` (the --apply path) sets the marker. +(c) Idempotence -- a second ``tag_confirmed_corrupt``/``run_apply`` call + against an already-tagged graph tags/matches zero additional rows, and + a subsequent ``run_apply()`` still reports success (exit 0). + +Run: uv run pytest tests/neo4j/test_tag_legacy_pooled_iterations.py -v -m neo4j +""" + +from __future__ import annotations + +from typing import Any + +import pytest +from neo4j import GraphDatabase +from scripts import tag_legacy_pooled_iterations as tagger + +WORKSPACE = "test" + +pytestmark = pytest.mark.neo4j + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _driver(neo4j_container: dict) -> GraphDatabase: + """Return a synchronous Neo4j driver for the test container.""" + return GraphDatabase.driver( + neo4j_container["bolt_url"], + auth=(neo4j_container["user"], neo4j_container["password"]), + ) + + +def _seed_run(session, run_id: str) -> None: + """Seed a bare :OrchestratorRun node.""" + session.run( + "MERGE (r:OrchestratorRun {node_id: $run_id, workspace: $workspace})", + run_id=run_id, + workspace=WORKSPACE, + ) + + +def _seed_bare_iteration(session, iter_id: str) -> None: + """Seed a bare-id :Iteration node. + + ``iter_id`` must NOT contain ``'::orch_run::'`` -- that is the pre-fix, + per-session-counter shape the script's ``_CONFIRMED_CORRUPT_MATCH`` + selector considers at all. Run-scoped (post-fix) node_ids are excluded by + the selector's ``WHERE NOT i.node_id CONTAINS '::orch_run::'`` clause. + """ + session.run( + "MERGE (i:Iteration {node_id: $iter_id, workspace: $workspace})", + iter_id=iter_id, + workspace=WORKSPACE, + ) + + +def _link(session, run_id: str, iter_id: str) -> None: + """Seed the real ``(OrchestratorRun)-[:HAS_PART]->(Iteration)`` edge + shape the script's selector matches (``MATCH (run:OrchestratorRun) + -[:HAS_PART]->(i:Iteration)``).""" + session.run( + "MATCH (r:OrchestratorRun {node_id: $run_id, workspace: $workspace}) " + "MATCH (i:Iteration {node_id: $iter_id, workspace: $workspace}) " + "MERGE (r)-[:HAS_PART]->(i)", + run_id=run_id, + iter_id=iter_id, + workspace=WORKSPACE, + ) + + +def _data_quality(session, iter_id: str) -> str | None: + """Return the ``data_quality`` property of an :Iteration node, or None.""" + result = session.run( + "MATCH (i:Iteration {node_id: $iter_id, workspace: $workspace}) " + "RETURN i.data_quality AS dq", + iter_id=iter_id, + workspace=WORKSPACE, + ) + record = result.single() + if record is None: + return None + return record["dq"] + + +def _seed_pooled_pair(session) -> None: + """Seed one pooled (>=2 run parents) and one clean (1 run parent) + bare-id Iteration node -- the shared fixture used by every test below.""" + _seed_run(session, "run-1") + _seed_run(session, "run-2") + _seed_bare_iteration(session, "sess-pooled::iteration::1") + _link(session, "run-1", "sess-pooled::iteration::1") + _link(session, "run-2", "sess-pooled::iteration::1") + + _seed_run(session, "run-3") + _seed_bare_iteration(session, "sess-clean::iteration::1") + _link(session, "run-3", "sess-clean::iteration::1") + + +# --------------------------------------------------------------------------- +# Per-test isolation +# --------------------------------------------------------------------------- + + +@pytest.fixture(autouse=True) +def _clear_neo4j(neo4j_container: dict) -> None: # type: ignore[return] + """Wipe the container clean before each test for complete isolation.""" + driver = GraphDatabase.driver( + neo4j_container["bolt_url"], + auth=(neo4j_container["user"], neo4j_container["password"]), + ) + try: + with driver.session() as s: + s.run("MATCH (n) DETACH DELETE n") + finally: + driver.close() + + +# --------------------------------------------------------------------------- +# (a) selector precision: >=2 parents tagged, exactly-1 parent NEVER tagged +# --------------------------------------------------------------------------- + + +@pytest.mark.neo4j +class TestSelectorPrecision: + def test_pooled_node_tagged_clean_node_untouched( + self, neo4j_container: dict[str, Any] + ) -> None: + driver = _driver(neo4j_container) + try: + with driver.session() as s: + _seed_pooled_pair(s) + + exit_code = tagger.run_apply(s, batch_size=tagger.DEFAULT_BATCH_SIZE) + + assert exit_code == 0 + assert ( + _data_quality(s, "sess-pooled::iteration::1") == tagger.TAG_VALUE + ), "the >=2-distinct-run-parent node must be tagged" + assert _data_quality(s, "sess-clean::iteration::1") is None, ( + "the exactly-1-run-parent (clean, single-run) node must " + "NEVER be tagged -- this is exactly the shape the " + "script's own docstring says is the majority of live " + "bare-id nodes and must be left untouched" + ) + finally: + driver.close() + + +# --------------------------------------------------------------------------- +# (b) --apply gate: dry-run makes NO mutation; only --apply writes +# --------------------------------------------------------------------------- + + +@pytest.mark.neo4j +class TestApplyGate: + def test_dry_run_makes_no_mutation_only_apply_writes( + self, neo4j_container: dict[str, Any] + ) -> None: + driver = _driver(neo4j_container) + try: + with driver.session() as s: + _seed_pooled_pair(s) + + # The default / non-apply invocation: run_dry_run() only + # calls classify() (read-only counts, no SET anywhere in its + # call graph). It must leave the confirmed-corrupt node + # completely untagged. + dry_exit_code = tagger.run_dry_run(s) + assert dry_exit_code == 0 + assert _data_quality(s, "sess-pooled::iteration::1") is None, ( + "run_dry_run must make ZERO writes -- the confirmed-" + "corrupt pooled node must remain untagged after a " + "dry-run invocation" + ) + assert _data_quality(s, "sess-clean::iteration::1") is None + + # Only --apply (run_apply) is permitted to write the tag. + apply_exit_code = tagger.run_apply( + s, batch_size=tagger.DEFAULT_BATCH_SIZE + ) + assert apply_exit_code == 0 + assert ( + _data_quality(s, "sess-pooled::iteration::1") == tagger.TAG_VALUE + ), "run_apply must tag the pooled node once explicitly invoked" + finally: + driver.close() + + +# --------------------------------------------------------------------------- +# (c) idempotence: second apply tags/matches zero additional rows +# --------------------------------------------------------------------------- + + +@pytest.mark.neo4j +class TestIdempotence: + def test_second_apply_tags_zero_additional_rows( + self, neo4j_container: dict[str, Any] + ) -> None: + driver = _driver(neo4j_container) + try: + with driver.session() as s: + _seed_pooled_pair(s) + + first_tagged = tagger.tag_confirmed_corrupt( + s, batch_size=tagger.DEFAULT_BATCH_SIZE + ) + second_tagged = tagger.tag_confirmed_corrupt( + s, batch_size=tagger.DEFAULT_BATCH_SIZE + ) + + assert first_tagged == 1, ( + "expected exactly 1 node tagged on the first apply " + f"(the pooled node only), got {first_tagged}" + ) + assert second_tagged == 0, ( + "a second apply against an already-tagged graph must " + f"match and write zero rows (idempotent), got {second_tagged}" + ) + assert _data_quality(s, "sess-pooled::iteration::1") == tagger.TAG_VALUE + assert _data_quality(s, "sess-clean::iteration::1") is None + + # The full run_apply() gate (not just the raw write helper) + # must also report success and find nothing outstanding on + # a third, still-idempotent re-run -- no error, no double-tag. + exit_code = tagger.run_apply(s, batch_size=tagger.DEFAULT_BATCH_SIZE) + assert exit_code == 0 + finally: + driver.close() diff --git a/tests/neo4j/test_working_dir_non_overwrite.py b/tests/neo4j/test_working_dir_non_overwrite.py new file mode 100644 index 00000000..1b1549e8 --- /dev/null +++ b/tests/neo4j/test_working_dir_non_overwrite.py @@ -0,0 +1,158 @@ +"""Live E2E tests: working_dir is never clobbered at the DB level. + +Root cause being guarded here: prior to this fix, the ONLY guarantee that an +already-populated ``working_dir`` is never overwritten lived in the Python layer +(``services.py``'s ``if data.get("working_dir") and not existing.get("working_dir")`` +populate-if-missing check). That check reads the buffered/graph node BEFORE the +write is issued, so it cannot protect against a cross-writer or replica race: a +second concurrent flush that read the node before the first write committed would +still see no working_dir, and its ``SET n += row.props`` would blindly overwrite +whatever the first writer just set. + +The fix adds a genuine DB-level guarantee: the Session-node MERGE in +``_write_batch`` excludes working_dir from the blind ``+=`` merge and instead +applies ``SET n.working_dir = coalesce(n.working_dir, row.working_dir)`` -- a +non-overwrite rule enforced by Neo4j itself, at the same MERGE lock hold, immune +to read-then-write races between writers. + +Requires Docker and the docker Python package. Skip-if-absent via the +``neo4j_container`` fixture in tests/neo4j/conftest.py. + +Run explicitly: + cd amplifier-context-intelligence + uv run pytest tests/neo4j/test_working_dir_non_overwrite.py -v -m neo4j +""" + +from __future__ import annotations + +from typing import Any + +import pytest +from context_intelligence_server.neo4j_store import Neo4jGraphStore +from neo4j import GraphDatabase + +pytestmark = pytest.mark.neo4j + + +async def _flush_session_node( + container: dict[str, Any], node_id: str, data: dict[str, Any] +) -> None: + """Drive a single Session node through the real flush path via a FRESH store. + + A fresh ``Neo4jGraphStore`` per call simulates independent writers (e.g. two + drainer workers, or a writer racing a replica) rather than two writes queued + through the same in-process buffer -- the scenario the Python-layer + populate-if-missing check cannot see. + """ + store = Neo4jGraphStore( + uri=container["bolt_url"], + auth=(container["user"], container["password"]), + workspace="test", + ) + try: + await store.upsert_node(node_id, {"labels": ["Session"], **data}) + await store.flush() + finally: + await store.close() + + +def _read_working_dir(container: dict[str, Any], node_id: str) -> Any: + driver = GraphDatabase.driver( + container["bolt_url"], auth=(container["user"], container["password"]) + ) + try: + with driver.session() as session: + rec = session.run( + "MATCH (n:Session {node_id: $nid, workspace: $ws}) " + "RETURN n.working_dir AS wd", + nid=node_id, + ws="test", + ).single() + assert rec is not None, f"Session node {node_id!r} was not written" + return rec["wd"] + finally: + driver.close() + + +async def test_existing_working_dir_survives_conflicting_later_write( + neo4j_container: dict[str, Any], +) -> None: + """DB-level guarantee: an already-set working_dir is never overwritten, + even by a second independent writer (simulating a cross-writer/replica race). + + RED (unfixed): ``SET n += row.props`` blindly overwrites -> working_dir + becomes "/y" -> this assertion fails. + GREEN (fixed): ``coalesce(n.working_dir, row.working_dir)`` keeps the + existing value -> working_dir stays "/x". + """ + node_id = "sess-wd-no-clobber-live" + + # Writer 1: establishes working_dir="/x". + await _flush_session_node( + neo4j_container, + node_id, + {"status": "running", "working_dir": "/x"}, + ) + assert _read_working_dir(neo4j_container, node_id) == "/x" + + # Writer 2 (independent store instance): tries to write a DIFFERENT + # working_dir for the SAME node -- must be rejected at the DB level. + await _flush_session_node( + neo4j_container, + node_id, + {"status": "running", "working_dir": "/y"}, + ) + + assert _read_working_dir(neo4j_container, node_id) == "/x", ( + "An already-populated working_dir must never be overwritten by a " + "later/concurrent writer -- DB-level coalesce guarantee failed" + ) + + +async def test_working_dir_fills_gap_at_db_level_when_previously_absent( + neo4j_container: dict[str, Any], +) -> None: + """DB-level populate-if-missing: a node with no working_dir gets filled in + by a later write that supplies one (coalesce(null, value) -> value). + + Mirrors the Python-layer guarantee (services.py) but proves it also holds + purely at the DB level, independent of the in-process buffer. + """ + node_id = "sess-wd-fill-gap-live" + + # Writer 1: no working_dir supplied. + await _flush_session_node(neo4j_container, node_id, {"status": "running"}) + assert _read_working_dir(neo4j_container, node_id) is None + + # Writer 2: supplies working_dir for the first time. + await _flush_session_node( + neo4j_container, + node_id, + {"status": "running", "working_dir": "/first-value"}, + ) + + assert _read_working_dir(neo4j_container, node_id) == "/first-value", ( + "working_dir must be filled in at the DB level once a writer supplies " + "a value for a node that previously had none" + ) + + +async def test_working_dir_absent_write_does_not_clear_existing_value( + neo4j_container: dict[str, Any], +) -> None: + """A later write that omits working_dir entirely must not null out an + already-set value (coalesce(n.working_dir, null) -> unchanged). + """ + node_id = "sess-wd-absent-no-clear-live" + + await _flush_session_node( + neo4j_container, node_id, {"status": "running", "working_dir": "/keep-me"} + ) + assert _read_working_dir(neo4j_container, node_id) == "/keep-me" + + # Second write carries no working_dir key at all (e.g. a touch/heartbeat event). + await _flush_session_node(neo4j_container, node_id, {"status": "still-running"}) + + assert _read_working_dir(neo4j_container, node_id) == "/keep-me", ( + "A write that omits working_dir must never null out an already-set value" + ) diff --git a/tests/routers/test_queues.py b/tests/routers/test_queues.py index c8909e57..8a218891 100644 --- a/tests/routers/test_queues.py +++ b/tests/routers/test_queues.py @@ -9,7 +9,7 @@ import pytest from context_intelligence_server.main import registry -from context_intelligence_server.queue_manager import QueueManager +from context_intelligence_server.queue_manager import FileSystemQueueManager, QueueManager def _point_registry_at(tmp_path: Path) -> QueueManager: @@ -17,7 +17,7 @@ def _point_registry_at(tmp_path: Path) -> QueueManager: Returns the QueueManager so tests can seed dead-letter records directly. """ - qm = QueueManager(queues_dir=tmp_path / "queues") + qm = FileSystemQueueManager(queues_dir=tmp_path / "queues") registry._queue_manager = qm registry._write_semaphore = asyncio.Semaphore(2) registry._max_delivery_attempts = 5 diff --git a/tests/test_blob_carrier_allowlist.py b/tests/test_blob_carrier_allowlist.py new file mode 100644 index 00000000..168445ce --- /dev/null +++ b/tests/test_blob_carrier_allowlist.py @@ -0,0 +1,160 @@ +"""Tests for the blob-carrier allowlist runtime tripwire. + +Covers the single source of truth shared by the mint path +(``blob_processor.BLOB_REF_CARRIER_PROPERTIES`` / +``blob_processor.assert_carrier_registered``) and the reclaim-scan path +(``routers.admin._BLOB_REF_CARRIER_PROPERTIES`` / ``_BLOB_REF_SCAN_QUERY``): + +1. The allowlist is exactly the current 4-item tuple (regression lock). +2. ``routers.admin`` imports the SAME object -- no local re-declaration to + drift out of sync. +3. The generated Cypher scan query references exactly the allowlist's + properties -- no more, no less (mint/scan agreement, structurally). +4. ``assert_carrier_registered`` is non-vacuous: it raises for an + unregistered property and is a no-op for a registered one. +5. ``process_event_data`` (the real mint call path) propagates the + tripwire's exception -- fail-closed, BEFORE any blob is written -- when + its destination carrier ("data") is not registered, and completes + normally when it is. +""" + +from __future__ import annotations + +import re +from typing import Any +from unittest.mock import AsyncMock + +import pytest +from context_intelligence_server.blob_processor import ( + BLOB_REF_CARRIER_PROPERTIES, + UnregisteredBlobCarrierError, + assert_carrier_registered, + process_event_data, +) +from context_intelligence_server.blob_store import BlobReference + + +def _ref(uri: str) -> BlobReference: + """A BlobReference for a mocked write() return (only .uri is read here).""" + session_id, _, key = uri.removeprefix("ci-blob://").partition("/") + return BlobReference( + uri=uri, session_id=session_id, key=key, size=0, last_modified=0.0 + ) +from context_intelligence_server.routers.admin import ( + _BLOB_REF_CARRIER_PROPERTIES as admin_carrier_properties, +) +from context_intelligence_server.routers.admin import ( + _BLOB_REF_SCAN_QUERY, +) + +# --------------------------------------------------------------------------- +# 1. Regression lock -- current 4-item allowlist +# --------------------------------------------------------------------------- + + +def test_carrier_properties_locked() -> None: + """BLOB_REF_CARRIER_PROPERTIES is exactly the specified 4-item tuple.""" + assert BLOB_REF_CARRIER_PROPERTIES == ("data", "tool_input", "prompt", "response") + + +# --------------------------------------------------------------------------- +# 2. admin.py imports the SAME allowlist -- no local re-declaration +# --------------------------------------------------------------------------- + + +def test_admin_imports_same_allowlist_object() -> None: + """routers.admin re-exports blob_processor's tuple by identity, not a + hand-copied duplicate -- proves there is exactly ONE allowlist object.""" + assert admin_carrier_properties is BLOB_REF_CARRIER_PROPERTIES + + +# --------------------------------------------------------------------------- +# 3. Generated scan query references exactly the allowlist's properties +# --------------------------------------------------------------------------- + + +def test_scan_query_matches_allowlist_exactly() -> None: + """The Cypher query built for the reclaim scan mentions exactly the + properties in BLOB_REF_CARRIER_PROPERTIES -- no more, no less. + + This is the structural lock that makes mint/scan drift impossible: if a + property is ever added to (or removed from) the allowlist without the + query being regenerated from it, this test fails. + """ + referenced_props = set(re.findall(r"n\.(\w+)", _BLOB_REF_SCAN_QUERY)) + assert referenced_props == set(BLOB_REF_CARRIER_PROPERTIES) + + +# --------------------------------------------------------------------------- +# 4. assert_carrier_registered -- non-vacuous tripwire +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("carrier", list(BLOB_REF_CARRIER_PROPERTIES)) +def test_assert_carrier_registered_passes_for_registered_carrier( + carrier: str, +) -> None: + """A registered carrier (each of the current 4) does not trip the guard.""" + assert_carrier_registered(carrier) # must not raise + + +def test_assert_carrier_registered_raises_for_unregistered_carrier() -> None: + """An unregistered carrier property trips the guard immediately. + + Proves the tripwire is non-vacuous: it actually fires. Simulates a + plausible future scenario -- a new field-lifter/enricher promoting a + value onto a brand-new node property ("artifact_content") + that nobody added to BLOB_REF_CARRIER_PROPERTIES. + """ + with pytest.raises(UnregisteredBlobCarrierError, match="artifact_content"): + assert_carrier_registered("artifact_content") + + +# --------------------------------------------------------------------------- +# 5. process_event_data -- the real mint call path +# --------------------------------------------------------------------------- + + +async def test_process_event_data_fails_closed_when_data_carrier_unregistered( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """If "data" is (hypothetically) removed from the allowlist, the mint + call fails loud BEFORE writing any blob -- not after. + + Regression target: this is the exact failure mode the allowlist tripwire + exists to catch -- a carrier property silently dropping out while the + mint path keeps writing to it. blob_store.write must never be called: + the guard fires before any blob is persisted, matching "fail loud at the + source, not after a live blob is deleted." + """ + import context_intelligence_server.blob_processor as blob_processor_module + + monkeypatch.setattr( + blob_processor_module, + "BLOB_REF_CARRIER_PROPERTIES", + ("tool_input", "prompt", "response"), # "data" removed + ) + + data: dict[str, Any] = {"result": {"answer": 42}} + blob_store = AsyncMock() + blob_store.write = AsyncMock(return_value=_ref("ci-blob://sess/node__result")) + + with pytest.raises(UnregisteredBlobCarrierError, match="'data'"): + await process_event_data(data, blob_store, "sess", "node") + + blob_store.write.assert_not_called() + # data must be untouched -- the guard fired before any mutation/write + assert data == {"result": {"answer": 42}} + + +async def test_process_event_data_succeeds_when_data_carrier_registered() -> None: + """Sanity/non-regression: with the real (unmodified) allowlist, the mint + path completes normally -- the guard does not false-trip on the + ordinary, correctly-registered path.""" + data: dict[str, Any] = {"result": {"answer": 42}} + blob_store = AsyncMock() + blob_store.write = AsyncMock(return_value=_ref("ci-blob://sess/node__result")) + + await process_event_data(data, blob_store, "sess", "node") + + assert data["result"] == {"$blob_ref": "ci-blob://sess/node__result"} diff --git a/tests/test_blob_processor.py b/tests/test_blob_processor.py index 82ee196c..43ae8074 100644 --- a/tests/test_blob_processor.py +++ b/tests/test_blob_processor.py @@ -29,6 +29,15 @@ _lift_raw_fields, process_event_data, ) +from context_intelligence_server.blob_store import BlobReference + + +def _ref(uri: str) -> BlobReference: + """A BlobReference for a mocked write() return (only .uri is read here).""" + session_id, _, key = uri.removeprefix("ci-blob://").partition("/") + return BlobReference( + uri=uri, session_id=session_id, key=key, size=0, last_modified=0.0 + ) # --------------------------------------------------------------------------- @@ -55,7 +64,7 @@ async def test_process_event_data_mutates_in_place() -> None: original_id = id(data) blob_store = AsyncMock() - blob_store.write = AsyncMock(return_value="ci-blob://sess/node__raw") + blob_store.write = AsyncMock(return_value=_ref("ci-blob://sess/node__raw")) await process_event_data(data, blob_store, "sess", "node") @@ -74,7 +83,7 @@ async def test_process_event_data_returns_none() -> None: """process_event_data returns None.""" data: dict[str, Any] = {"result": {"answer": 42}} blob_store = AsyncMock() - blob_store.write = AsyncMock(return_value="ci-blob://sess/node__result") + blob_store.write = AsyncMock(return_value=_ref("ci-blob://sess/node__result")) result = await process_event_data(data, blob_store, "sess", "node") @@ -96,8 +105,8 @@ async def test_blob_ref_substitution_on_successful_write() -> None: # Use a function-based side_effect so the returned URI always matches # the actual key argument, regardless of BLOB_FIELDS frozenset iteration order. - async def _write(session_id: str, key: str, value: object) -> str: - return f"ci-blob://{session_id}/{key}" + async def _write(session_id: str, key: str, value: object) -> BlobReference: + return _ref(f"ci-blob://{session_id}/{key}") blob_store.write = AsyncMock(side_effect=_write) @@ -134,7 +143,7 @@ async def test_absent_fields_are_skipped() -> None: """Fields in BLOB_FIELDS that are absent from data are not added.""" data: dict[str, Any] = {"other_field": "untouched"} blob_store = AsyncMock() - blob_store.write = AsyncMock(return_value="ci-blob://sess/node__something") + blob_store.write = AsyncMock(return_value=_ref("ci-blob://sess/node__something")) await process_event_data(data, blob_store, "sess", "node") @@ -158,7 +167,7 @@ async def test_none_fields_are_skipped() -> None: "messages": None, } blob_store = AsyncMock() - blob_store.write = AsyncMock(return_value="ci-blob://sess/node__x") + blob_store.write = AsyncMock(return_value=_ref("ci-blob://sess/node__x")) await process_event_data(data, blob_store, "sess", "node") diff --git a/tests/test_blob_reclaim_endpoint.py b/tests/test_blob_reclaim_endpoint.py new file mode 100644 index 00000000..73363369 --- /dev/null +++ b/tests/test_blob_reclaim_endpoint.py @@ -0,0 +1,187 @@ +"""The POST /admin/blobs/reclaim orchestration: dry-run vs apply, the +blast-radius cap, the destructive-apply single-flight, and fenced deletion. + +The *selection* logic (which blobs are orphans) is covered against a real +graph elsewhere; these tests pin the endpoint's own contract by stubbing the +one selection call, so they run without neo4j. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import pytest +from context_intelligence_server.blob_store import BlobReference +from context_intelligence_server.routers import admin +from context_intelligence_server.routers.admin import BlobReclaimBody, reclaim_blobs +from fastapi import HTTPException + +pytestmark = pytest.mark.integration + + +def _ref(uri: str, size: int = 10) -> BlobReference: + session_id, _, key = uri.removeprefix("ci-blob://").partition("/") + return BlobReference( + uri=uri, session_id=session_id, key=key, size=size, last_modified=1.0 + ) + + +def _request() -> Any: + return SimpleNamespace(app=SimpleNamespace(state=SimpleNamespace())) + + +def _stub_selection(monkeypatch, candidates: list[BlobReference]) -> None: + async def _fake_select(_request: Any, *, min_age_minutes: int) -> dict[str, Any]: + return { + "scanned_disk_blobs": len(candidates), + "referenced_uris": 0, + "orphans_found": len(candidates), + "reclaimable_bytes": sum(c.size for c in candidates), + "skipped_recent": 0, + "skipped_pending_session": 0, + "candidates": list(candidates), + } + + monkeypatch.setattr(admin, "_select_orphans", _fake_select) + + +class _FakeStore: + """A blob store whose delete() honours a per-uri fence verdict.""" + + def __init__(self, deletable: set[str]) -> None: + self._deletable = deletable + self.deleted: list[str] = [] + + async def delete(self, uri: str, if_unmodified: BlobReference | None = None) -> bool: + if uri in self._deletable: + self.deleted.append(uri) + return True + return False # absent or changed since scan -- fenced out + + +@pytest.fixture(autouse=True) +def _reset_single_flight(): + admin._reclaim_apply_inflight = False + yield + admin._reclaim_apply_inflight = False + + +async def test_apply_without_max_delete_is_422(monkeypatch) -> None: + """F2: a destructive apply must name its blast radius.""" + with pytest.raises(HTTPException) as exc: + await reclaim_blobs( + BlobReclaimBody(dry_run=False, max_delete=None), _request() + ) + assert exc.value.status_code == 422 + + +async def test_dry_run_reports_without_deleting(monkeypatch) -> None: + """F1: dry-run (the default) previews and deletes nothing.""" + cands = [_ref("ci-blob://s1/a"), _ref("ci-blob://s1/b")] + _stub_selection(monkeypatch, cands) + store = _FakeStore({"ci-blob://s1/a", "ci-blob://s1/b"}) + monkeypatch.setattr(admin, "create_blob_store", lambda _s: store) + + resp = await reclaim_blobs(BlobReclaimBody(dry_run=True), _request()) + + assert resp["dry_run"] is True + assert resp["rescanned"] is False + assert resp["orphans_found"] == 2 + assert resp["deleted"] == 0 + assert store.deleted == [] # nothing touched + + +async def test_apply_deletes_through_fenced_protocol(monkeypatch) -> None: + """Happy-path apply: fresh scan (rescanned), fenced delete, audit per delete.""" + cands = [_ref("ci-blob://s1/a"), _ref("ci-blob://s1/b")] + _stub_selection(monkeypatch, cands) + store = _FakeStore({"ci-blob://s1/a", "ci-blob://s1/b"}) + monkeypatch.setattr(admin, "create_blob_store", lambda _s: store) + audited: list[str] = [] + monkeypatch.setattr( + admin, "_audit_blob_reclaim_delete", lambda _r, *, uri: audited.append(uri) + ) + + resp = await reclaim_blobs( + BlobReclaimBody(dry_run=False, max_delete=10), _request() + ) + + assert resp["rescanned"] is True + assert resp["deleted"] == 2 + assert sorted(store.deleted) == ["ci-blob://s1/a", "ci-blob://s1/b"] + assert sorted(audited) == ["ci-blob://s1/a", "ci-blob://s1/b"] + assert admin._reclaim_apply_inflight is False # released + + +async def test_fenced_delete_refusal_is_not_counted(monkeypatch) -> None: + """R3/R4: a blob changed/re-referenced since the scan is fenced out -- + delete() returns False, it stays on disk and is NOT counted as deleted.""" + cands = [_ref("ci-blob://s1/a"), _ref("ci-blob://s1/b")] + _stub_selection(monkeypatch, cands) + # Only 'a' is still deletable; 'b' was re-minted since the scan. + store = _FakeStore({"ci-blob://s1/a"}) + monkeypatch.setattr(admin, "create_blob_store", lambda _s: store) + monkeypatch.setattr( + admin, "_audit_blob_reclaim_delete", lambda _r, *, uri: None + ) + + resp = await reclaim_blobs( + BlobReclaimBody(dry_run=False, max_delete=10), _request() + ) + + assert resp["deleted"] == 1 + assert store.deleted == ["ci-blob://s1/a"] # 'b' left intact + + +async def test_max_delete_caps_blast_radius(monkeypatch) -> None: + """F4: orphans_found reflects the FULL set; only max_delete are removed.""" + cands = [_ref(f"ci-blob://s1/{k}") for k in "abcde"] + _stub_selection(monkeypatch, cands) + store = _FakeStore({c.uri for c in cands}) + monkeypatch.setattr(admin, "create_blob_store", lambda _s: store) + monkeypatch.setattr( + admin, "_audit_blob_reclaim_delete", lambda _r, *, uri: None + ) + + resp = await reclaim_blobs( + BlobReclaimBody(dry_run=False, max_delete=2), _request() + ) + + assert resp["orphans_found"] == 5 # full candidate set + assert resp["deleted"] == 2 # capped + assert len(store.deleted) == 2 + + +async def test_concurrent_apply_is_single_flighted(monkeypatch) -> None: + """E1: a second apply while one is in flight is refused (409) before it + even scans -- two applies would each honour max_delete and jointly exceed + the operator's intended blast radius.""" + admin._reclaim_apply_inflight = True # simulate an apply already running + scanned = False + + async def _should_not_run(_request: Any, *, min_age_minutes: int) -> dict[str, Any]: + nonlocal scanned + scanned = True + return {"candidates": []} + + monkeypatch.setattr(admin, "_select_orphans", _should_not_run) + + with pytest.raises(HTTPException) as exc: + await reclaim_blobs( + BlobReclaimBody(dry_run=False, max_delete=1), _request() + ) + + assert exc.value.status_code == 409 + assert scanned is False # fail-fast: rejected before the authoritative scan + + +async def test_dry_run_is_never_single_flighted(monkeypatch) -> None: + """A preview must never be blocked by an in-flight apply.""" + admin._reclaim_apply_inflight = True + _stub_selection(monkeypatch, [_ref("ci-blob://s1/a")]) + + resp = await reclaim_blobs(BlobReclaimBody(dry_run=True), _request()) + + assert resp["dry_run"] is True + assert resp["orphans_found"] == 1 diff --git a/tests/test_blob_store.py b/tests/test_blob_store.py index 4a3a31ad..13884795 100644 --- a/tests/test_blob_store.py +++ b/tests/test_blob_store.py @@ -1,14 +1,14 @@ -"""Tests for AsyncDiskBlobStore — Write, Read, List, Dump. +"""Tests for FileSystemBlobStore — Write, Read, List, Scan, Delete, Dump. -15 tests covering: +Covers: 1. write/read roundtrip -2. URI format +2. BlobReference.uri format 3. directory structure creation 4. URI-based session_id resolution 5. missing blob raises FileNotFoundError 6. invalid URI raises ValueError 7. empty list for missing session -8. correct URI listing +8. correct BlobReference listing (async iterator) 9. session isolation 10. asyncio.to_thread delegation verification 11. dump() copies blob to specified dest_dir @@ -16,19 +16,26 @@ 13. dump() missing blob raises FileNotFoundError 14. dump() delegates copy2 via asyncio.to_thread 15. BlobStore protocol conformance +16. scan() yields BlobReference across multiple sessions +17. delete() is idempotent (True then False) and removes the blob +18. list()/write() BlobReference has correct uri/size/last_modified """ from __future__ import annotations import asyncio import json +import os from pathlib import Path from unittest.mock import patch import pytest - -from context_intelligence_server.blob_store import AsyncDiskBlobStore, BlobStore - +from context_intelligence_server.blob_store import ( + BlobNotFoundError, + BlobReference, + BlobStore, + FileSystemBlobStore, +) # --------------------------------------------------------------------------- # Fixtures @@ -36,9 +43,13 @@ @pytest.fixture -def store(tmp_path: Path) -> AsyncDiskBlobStore: - """Return a fresh AsyncDiskBlobStore rooted at a temporary directory.""" - return AsyncDiskBlobStore(root=tmp_path) +def store(tmp_path: Path) -> FileSystemBlobStore: + """Return a fresh FileSystemBlobStore rooted at a temporary directory.""" + return FileSystemBlobStore(root=tmp_path) + + +async def _list_uris(store: FileSystemBlobStore, session_id: str) -> list[str]: + return [ref.uri async for ref in store.list(session_id)] # --------------------------------------------------------------------------- @@ -46,23 +57,28 @@ def store(tmp_path: Path) -> AsyncDiskBlobStore: # --------------------------------------------------------------------------- -async def test_write_read_roundtrip(store: AsyncDiskBlobStore) -> None: +async def test_write_read_roundtrip(store: FileSystemBlobStore) -> None: """Data written can be read back unchanged.""" payload = {"event": "tool_call", "tool": "bash", "args": ["ls"]} - uri = await store.write("session-abc", "tool_call_01", payload) - result = await store.read(uri) + ref = await store.write("session-abc", "tool_call_01", payload) + result = await store.read(ref.uri) assert result == payload # --------------------------------------------------------------------------- -# 2. URI format +# 2. BlobReference.uri format # --------------------------------------------------------------------------- -async def test_uri_format(store: AsyncDiskBlobStore) -> None: - """write() returns a ci-blob:/// URI.""" - uri = await store.write("session-xyz", "my_key", {"x": 1}) - assert uri == "ci-blob://session-xyz/my_key" +async def test_uri_format(store: FileSystemBlobStore) -> None: + """write() returns a BlobReference whose .uri is ci-blob:///.""" + ref = await store.write("session-xyz", "my_key", {"x": 1}) + assert isinstance(ref, BlobReference) + assert ref.uri == "ci-blob://session-xyz/my_key" + assert ref.session_id == "session-xyz" + assert ref.key == "my_key" + assert ref.size > 0 + assert ref.last_modified > 0 # --------------------------------------------------------------------------- @@ -71,7 +87,7 @@ async def test_uri_format(store: AsyncDiskBlobStore) -> None: async def test_directory_structure_creation( - store: AsyncDiskBlobStore, tmp_path: Path + store: FileSystemBlobStore, tmp_path: Path ) -> None: """write() creates //blobs/.json on disk.""" await store.write("session-123", "blob_key", {"data": "value"}) @@ -87,17 +103,17 @@ async def test_directory_structure_creation( async def test_uri_based_session_id_resolution( - store: AsyncDiskBlobStore, tmp_path: Path + store: FileSystemBlobStore, tmp_path: Path ) -> None: """read() resolves the session_id from the URI, not from a parameter.""" session_id = "session-uri-resolve" key = "my_blob" payload = {"resolved": True} - uri = await store.write(session_id, key, payload) + ref = await store.write(session_id, key, payload) # Confirm URI contains session_id - assert session_id in uri + assert session_id in ref.uri # read must successfully resolve session_id from URI - result = await store.read(uri) + result = await store.read(ref.uri) assert result == payload @@ -106,7 +122,7 @@ async def test_uri_based_session_id_resolution( # --------------------------------------------------------------------------- -async def test_missing_blob_raises_file_not_found(store: AsyncDiskBlobStore) -> None: +async def test_missing_blob_raises_file_not_found(store: FileSystemBlobStore) -> None: """read() raises FileNotFoundError for a URI pointing to a non-existent blob.""" uri = "ci-blob://session-missing/nonexistent_key" with pytest.raises(FileNotFoundError): @@ -118,7 +134,7 @@ async def test_missing_blob_raises_file_not_found(store: AsyncDiskBlobStore) -> # --------------------------------------------------------------------------- -async def test_invalid_uri_raises_value_error(store: AsyncDiskBlobStore) -> None: +async def test_invalid_uri_raises_value_error(store: FileSystemBlobStore) -> None: """read() raises ValueError for URIs that don't match the ci-blob:// scheme.""" with pytest.raises(ValueError): await store.read("not-a-ci-blob-uri") @@ -135,30 +151,35 @@ async def test_invalid_uri_raises_value_error(store: AsyncDiskBlobStore) -> None # --------------------------------------------------------------------------- -async def test_empty_list_for_missing_session(store: AsyncDiskBlobStore) -> None: - """list() returns an empty list when no blobs exist for the session.""" - result = await store.list("session-does-not-exist") +async def test_empty_list_for_missing_session(store: FileSystemBlobStore) -> None: + """list() yields nothing when no blobs exist for the session.""" + result = await _list_uris(store, "session-does-not-exist") assert result == [] # --------------------------------------------------------------------------- -# 8. Correct URI listing +# 8. Correct BlobReference listing (async iterator) # --------------------------------------------------------------------------- -async def test_correct_uri_listing(store: AsyncDiskBlobStore) -> None: - """list() returns all blob URIs for a session, sorted.""" +async def test_correct_uri_listing(store: FileSystemBlobStore) -> None: + """list() yields all blob references for a session, sorted by key.""" session_id = "session-list" await store.write(session_id, "key_b", {"b": 2}) await store.write(session_id, "key_a", {"a": 1}) await store.write(session_id, "key_c", {"c": 3}) - uris = await store.list(session_id) - assert uris == [ + refs = [ref async for ref in store.list(session_id)] + assert [r.uri for r in refs] == [ "ci-blob://session-list/key_a", "ci-blob://session-list/key_b", "ci-blob://session-list/key_c", ] + for r in refs: + assert isinstance(r, BlobReference) + assert r.session_id == session_id + assert r.size > 0 + assert r.last_modified > 0 # --------------------------------------------------------------------------- @@ -166,14 +187,14 @@ async def test_correct_uri_listing(store: AsyncDiskBlobStore) -> None: # --------------------------------------------------------------------------- -async def test_session_isolation(store: AsyncDiskBlobStore) -> None: - """list() only returns URIs for the requested session, not other sessions.""" +async def test_session_isolation(store: FileSystemBlobStore) -> None: + """list() only returns references for the requested session, not other sessions.""" await store.write("session-alpha", "blob_1", {"alpha": True}) await store.write("session-beta", "blob_2", {"beta": True}) await store.write("session-alpha", "blob_3", {"alpha2": True}) - alpha_uris = await store.list("session-alpha") - beta_uris = await store.list("session-beta") + alpha_uris = await _list_uris(store, "session-alpha") + beta_uris = await _list_uris(store, "session-beta") assert all("session-alpha" in u for u in alpha_uris) assert all("session-beta" in u for u in beta_uris) @@ -188,7 +209,7 @@ async def test_session_isolation(store: AsyncDiskBlobStore) -> None: async def test_asyncio_to_thread_delegation(tmp_path: Path) -> None: """All filesystem I/O is delegated to asyncio.to_thread for non-blocking I/O.""" - store = AsyncDiskBlobStore(root=tmp_path) + store = FileSystemBlobStore(root=tmp_path) to_thread_calls: list[str] = [] original_to_thread = asyncio.to_thread @@ -200,7 +221,8 @@ async def tracking_to_thread(func, *args, **kwargs): # type: ignore[no-untyped- with patch("asyncio.to_thread", side_effect=tracking_to_thread): await store.write("sess", "k", {"v": 1}) await store.read("ci-blob://sess/k") - await store.list("sess") + async for _ in store.list("sess"): + pass assert len(to_thread_calls) >= 3, ( f"Expected at least 3 asyncio.to_thread calls (write, read, list), " @@ -214,16 +236,16 @@ async def tracking_to_thread(func, *args, **kwargs): # type: ignore[no-untyped- async def test_dump_copy_to_specified_dest_dir( - store: AsyncDiskBlobStore, tmp_path: Path + store: FileSystemBlobStore, tmp_path: Path ) -> None: """dump() copies the blob file to the specified dest_dir and returns the path.""" session_id = "session-dump-copy" key = "blob_to_copy" payload = {"copy": "me"} - uri = await store.write(session_id, key, payload) + ref = await store.write(session_id, key, payload) dest_dir = tmp_path / "my_dest" - result = await store.dump(uri, dest_dir=dest_dir) + result = await store.dump(ref.uri, dest_dir=dest_dir) result_path = Path(result) assert result_path.exists() @@ -236,15 +258,15 @@ async def test_dump_copy_to_specified_dest_dir( # --------------------------------------------------------------------------- -async def test_dump_default_dest_dir(store: AsyncDiskBlobStore) -> None: +async def test_dump_default_dest_dir(store: FileSystemBlobStore) -> None: """dump() uses Path(tempfile.gettempdir()) / 'ci-blobs' when dest_dir is None.""" import tempfile session_id = "session-dump-default" key = "default_blob" - uri = await store.write(session_id, key, {"default": True}) + ref = await store.write(session_id, key, {"default": True}) - result = await store.dump(uri) + result = await store.dump(ref.uri) expected_dir = Path(tempfile.gettempdir()) / "ci-blobs" result_path = Path(result) @@ -258,7 +280,7 @@ async def test_dump_default_dest_dir(store: AsyncDiskBlobStore) -> None: async def test_dump_missing_blob_raises_file_not_found( - store: AsyncDiskBlobStore, + store: FileSystemBlobStore, ) -> None: """dump() raises FileNotFoundError with 'Blob not found' message for missing blob.""" uri = "ci-blob://session-nonexistent/missing_blob" @@ -272,12 +294,12 @@ async def test_dump_missing_blob_raises_file_not_found( async def test_dump_uses_asyncio_to_thread_for_copy2( - store: AsyncDiskBlobStore, tmp_path: Path + store: FileSystemBlobStore, tmp_path: Path ) -> None: """dump() delegates shutil.copy2 to asyncio.to_thread for non-blocking I/O.""" session_id = "session-dump-thread" key = "thread_blob" - uri = await store.write(session_id, key, {"thread": True}) + ref = await store.write(session_id, key, {"thread": True}) dest_dir = tmp_path / "thread_dest" to_thread_calls: list[str] = [] @@ -288,7 +310,7 @@ async def tracking_to_thread(func, *args, **kwargs): # type: ignore[no-untyped- return await original_to_thread(func, *args, **kwargs) with patch("asyncio.to_thread", side_effect=tracking_to_thread): - await store.dump(uri, dest_dir=dest_dir) + await store.dump(ref.uri, dest_dir=dest_dir) assert len(to_thread_calls) >= 1, ( f"Expected at least 1 asyncio.to_thread call for dump(), " @@ -301,8 +323,8 @@ async def tracking_to_thread(func, *args, **kwargs): # type: ignore[no-untyped- # --------------------------------------------------------------------------- -def test_blob_store_protocol_conformance(store: AsyncDiskBlobStore) -> None: - """AsyncDiskBlobStore conforms to the BlobStore protocol.""" +def test_blob_store_protocol_conformance(store: FileSystemBlobStore) -> None: + """FileSystemBlobStore conforms to the BlobStore protocol.""" assert isinstance(store, BlobStore) @@ -312,18 +334,20 @@ def test_blob_store_protocol_conformance(store: AsyncDiskBlobStore) -> None: async def test_write_is_atomic_no_torn_file_on_failure( - store: AsyncDiskBlobStore, tmp_path: Path + store: FileSystemBlobStore, tmp_path: Path ) -> None: """A failure during os.replace leaves no torn final file and no temp siblings.""" session_id = "sess-atomic" key = "k1" - with patch( - "context_intelligence_server.blob_store.os.replace", - side_effect=OSError("simulated replace failure"), + with ( + patch( + "context_intelligence_server.blob_store.filesystem.os.replace", + side_effect=OSError("simulated replace failure"), + ), + pytest.raises(OSError), ): - with pytest.raises(OSError): - await store.write(session_id, key, {"v": 1}) + await store.write(session_id, key, {"v": 1}) final_path = store.blob_path(session_id, key) # No torn file observable at the final path. @@ -335,16 +359,180 @@ async def test_write_is_atomic_no_torn_file_on_failure( async def test_write_replaces_atomically_on_success( - store: AsyncDiskBlobStore, + store: FileSystemBlobStore, ) -> None: """On success the final file has the exact JSON, no temp remains, URI is correct.""" session_id = "sess-atomic" key = "k2" - uri = await store.write(session_id, key, {"v": 1}) + ref = await store.write(session_id, key, {"v": 1}) - assert uri == "ci-blob://sess-atomic/k2" + assert ref.uri == "ci-blob://sess-atomic/k2" final_path = store.blob_path(session_id, key) assert final_path.read_text(encoding="utf-8") == '{"v": 1}' # No leftover temp files. assert list(final_path.parent.glob("*.tmp")) == [] + + +# --------------------------------------------------------------------------- +# 16. scan() yields BlobReference across multiple sessions +# --------------------------------------------------------------------------- + + +async def test_scan_yields_references_across_sessions( + store: FileSystemBlobStore, +) -> None: + """scan() streams a BlobReference for every blob across ALL sessions.""" + await store.write("session-scan-a", "k1", {"a": 1}) + await store.write("session-scan-a", "k2", {"a": 2}) + await store.write("session-scan-b", "k1", {"b": 1}) + + refs = [ref async for ref in store.scan()] + uris = {r.uri for r in refs} + assert uris == { + "ci-blob://session-scan-a/k1", + "ci-blob://session-scan-a/k2", + "ci-blob://session-scan-b/k1", + } + for r in refs: + assert isinstance(r, BlobReference) + assert r.size > 0 + assert r.last_modified > 0 + + +async def test_scan_empty_store_yields_nothing(store: FileSystemBlobStore) -> None: + """scan() over an empty store yields no references.""" + refs = [ref async for ref in store.scan()] + assert refs == [] + + +# --------------------------------------------------------------------------- +# 17. delete() is idempotent and removes the blob +# --------------------------------------------------------------------------- + + +async def test_delete_idempotent_true_then_false(store: FileSystemBlobStore) -> None: + """delete() returns True the first time (blob existed), False thereafter.""" + ref = await store.write("session-delete", "to_delete", {"gone": "soon"}) + + first = await store.delete(ref.uri) + assert first is True + + # The blob is actually removed from disk. + with pytest.raises(FileNotFoundError): + await store.read(ref.uri) + + second = await store.delete(ref.uri) + assert second is False + + +async def test_delete_missing_blob_returns_false(store: FileSystemBlobStore) -> None: + """delete() on a never-written blob returns False, never raises.""" + result = await store.delete("ci-blob://never-existed/nope") + assert result is False + + +# --------------------------------------------------------------------------- +# BlobNotFoundError — neutral missing-blob error (guard #6) +# --------------------------------------------------------------------------- + + +async def test_missing_blob_raises_blob_not_found_error_no_path_leak( + store: FileSystemBlobStore, tmp_path: Path +) -> None: + """read() of a missing uri raises BlobNotFoundError (a FileNotFoundError + subclass, for back-compat) whose message carries the uri only — never the + on-disk path/root. + """ + uri = "ci-blob://session-missing/nonexistent_key" + + with pytest.raises(BlobNotFoundError) as exc_info: + await store.read(uri) + + # Back-compat: existing `except FileNotFoundError` callers still catch it. + assert isinstance(exc_info.value, FileNotFoundError) + + message = str(exc_info.value) + assert uri in message + # No on-disk path fragment or root leaks into the message. + assert "path" not in message.lower() + assert str(tmp_path) not in message + + +# --------------------------------------------------------------------------- +# Fenced (compare-and-delete) delete — guard #1 +# --------------------------------------------------------------------------- + + +async def test_fenced_delete_succeeds_when_unchanged( + store: FileSystemBlobStore, +) -> None: + """delete(uri, if_unmodified=ref) removes the blob when it has not + changed since ref was observed (e.g. by scan()/list()).""" + ref = await store.write("session-fence", "unchanged_key", {"v": 1}) + + result = await store.delete(ref.uri, if_unmodified=ref) + assert result is True + + with pytest.raises(BlobNotFoundError): + await store.read(ref.uri) + + +async def test_fenced_delete_refuses_when_rewritten( + store: FileSystemBlobStore, +) -> None: + """delete(uri, if_unmodified=stale_ref) returns False and leaves the + (new) blob on disk when the blob was rewritten after stale_ref was + observed. + + Deterministic (not sleep-based): the rewrite uses a longer JSON payload + so the size differs regardless of filesystem mtime granularity, and the + file's mtime is also forced forward via os.utime so both the size AND + mtime comparisons independently detect the change. + """ + session_id, key = "session-fence-stale", "rewritten_key" + + stale_ref = await store.write(session_id, key, {"v": 1}) + + # Rewrite with a longer payload -> different size, independent of mtime + # resolution/granularity on the filesystem. + new_ref = await store.write(session_id, key, {"v": 1, "extra": "x" * 64}) + assert new_ref.size != stale_ref.size + + # Force the mtime to be unambiguously different too (belt-and-suspenders + # against any filesystem where sizes could coincidentally collide). + path = store.blob_path(session_id, key) + new_mtime = stale_ref.last_modified + 100.0 + os.utime(path, (new_mtime, new_mtime)) + + result = await store.delete(stale_ref.uri, if_unmodified=stale_ref) + assert result is False + + # The (new) blob survives on disk, untouched. + survived = await store.read(stale_ref.uri) + assert survived == {"v": 1, "extra": "x" * 64} + + +async def test_fenced_delete_missing_blob_returns_false( + store: FileSystemBlobStore, +) -> None: + """delete(uri, if_unmodified=ref) on an already-absent blob returns False.""" + ref = await store.write("session-fence-missing", "gone_key", {"v": 1}) + assert await store.delete(ref.uri) is True # unconditional delete first + + result = await store.delete(ref.uri, if_unmodified=ref) + assert result is False + + +async def test_unconditional_delete_still_idempotent( + store: FileSystemBlobStore, +) -> None: + """Unconditional delete(uri) (if_unmodified=None, the default) is + unchanged: True then False, idempotent.""" + ref = await store.write("session-fence-uncond", "plain_key", {"v": 1}) + + first = await store.delete(ref.uri) + assert first is True + + second = await store.delete(ref.uri) + assert second is False diff --git a/tests/test_boot_safety.py b/tests/test_boot_safety.py new file mode 100644 index 00000000..90516420 --- /dev/null +++ b/tests/test_boot_safety.py @@ -0,0 +1,1520 @@ +"""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 FileSystemQueueManager, 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 FileSystemQueueManager(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.filesystem.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.filesystem.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 = FileSystemQueueManager(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", FileSystemQueueManager(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", FileSystemQueueManager(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 = FileSystemQueueManager(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 = FileSystemQueueManager(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 = FileSystemQueueManager(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, None) + 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 = FileSystemQueueManager(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 = FileSystemQueueManager(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, None) # 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 = FileSystemQueueManager(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, None) + + 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 = FileSystemQueueManager(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 = FileSystemQueueManager(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 = FileSystemQueueManager(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 = FileSystemQueueManager(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 = FileSystemQueueManager(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 = FileSystemQueueManager(queues_dir=tmp_path) + line = _line() + await qm.append("drained-key", line) + await qm.commit("drained-key", len(line), None) + + monkeypatch.setattr(main_module.registry, "_queue_manager", qm) + monkeypatch.setattr(main_module._settings, "reclaim_enabled", False) + 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 = FileSystemQueueManager(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.filesystem.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 = FileSystemQueueManager(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.filesystem.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 = FileSystemQueueManager(queues_dir=tmp_path) + line = _line() + await qm.append("live-drained-key", line) + await qm.commit("live-drained-key", len(line), None) + + 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_stays_gated_and_schedules_sweep( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Reachable graph, genuinely un-migrated (untagged>0), unarmed lease: it + refuses to start drainers and does NOT auto-mutate, but must stay gated via + ``awaiting_schema`` (not ``failed``) so the retry sweep is still scheduled. + Aborting at ``failed`` would skip sweep scheduling and strand the graph.""" + 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), + ), + # Unarmed lease: auto-repair must NOT run; the server stays gated. + patch.object(main_module.writer_lease, "acquired", False), + 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 + # Stays gated (not "failed") so the sweep loop is scheduled to retry. + assert boot_state.phase == "awaiting_schema" + assert boot_state.degraded_reason is not None + assert main_module.app.state.sweep_task is not None + + 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 = FileSystemQueueManager(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..406559d0 --- /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 FileSystemQueueManager, 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 = FileSystemQueueManager(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 = FileSystemQueueManager(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..1821436e 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -71,26 +71,45 @@ def test_settings_has_durable_queue_defaults(): assert s.max_delivery_attempts == 5 +def test_maintenance_knob_defaults(): + """Maintenance-mode knobs exist without disturbing PR#78 defaults.""" + from context_intelligence_server.config import Settings + + s = Settings() + assert s.maintenance_probe_ttl_seconds == 5.0 + assert s.maintenance_retry_after_seconds == 30 + assert s.maintenance_quiesce_seconds == 2.0 + # Guard the two PR#78 defaults the maintenance work must not change. + assert s.writer_lease_mode == "enforce" + assert s.dead_letter_expiry_enabled is False + + # --------------------------------------------------------------------------- # crash_recovery_respawn_limit (Change 1) # --------------------------------------------------------------------------- -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 +454,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..f3ae3f76 --- /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 FileSystemQueueManager, 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 = FileSystemQueueManager(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 = FileSystemQueueManager(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 = FileSystemQueueManager(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..a14f5d29 --- /dev/null +++ b/tests/test_drain_supervision.py @@ -0,0 +1,1197 @@ +"""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" + + +class TestRetryDoesNotDuplicateIteration: + """In-place retry rolls the cross-handler counter back to its pre-attempt + state, so a replayed batch reproduces the SAME node ids instead of + duplicating them.""" + + async def test_in_place_retry_yields_exactly_one_iteration(self) -> None: + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d1-retry-dedup" + # Fail the first multi-record flush exactly once, then succeed (in-place + # retry, no respawn). flushed is a SET: a duplicate id shows as an extra + # member, and a non-rolled-back counter yields iter::3/iter::4 too. + flush_calls = {"n": 0} + + def _fail_first_multi(buf: set[str]) -> bool: + if len(buf) > 1: + flush_calls["n"] += 1 + return flush_calls["n"] == 1 + return False + + graph = _FlakyGraph(fail_when=_fail_first_multi) + worker = _make_worker(sid, graph) + reg._register_for_test(worker) + + async def _advance_and_buffer( + w: SessionWorker, event: str, data: object, handlers: object + ) -> None: + # Mimic a real enricher: bump the iteration counter, then buffer the + # id it produces. Without the pre-retry rollback the counter keeps + # climbing across attempts and the replay emits a DIFFERENT id. + w.services.data_layer_2.iteration_count += 1 + w.services.graph.buffer.add( + f"iter::{w.services.data_layer_2.iteration_count}" + ) + + with patch( + "context_intelligence_server.registry.process_event", + side_effect=_advance_and_buffer, + ): + await qm.append(sid, _line("e1", "/ws", {"session_id": sid})) + await qm.append(sid, _line("e2", "/ws", {"session_id": sid})) + reg.start_drain(worker) + await _drain_until_idle(reg, qm, worker, sid) + await _cancel_and_await(worker.task) + + # Two records => ids iter::1, iter::2 written exactly once each. A missing + # rollback would additionally leave iter::3/iter::4 from the failed attempt. + assert graph.flushed == {"iter::1", "iter::2"} + assert (await qm.read_batch(sid, 10)).lines == [] + + +class TestCursorRestoredOnWorkerRebuild: + """A rebuilt worker restores the durable cursor the last commit persisted, + so cross-handler counters resume instead of restarting from zero (which + would remint node ids and duplicate them).""" + + async def test_orch_run_seq_is_restored_before_a_rebuilt_worker_processes( + self, + ) -> None: + reg = SessionRegistry() + qm = reg.queue_manager + sid = "d1-cursor-rebuild" + + # A prior worker committed a cursor carrying orch_run_seq=5, leaving one + # event still undrained (offset 0) for the rebuilt worker to pick up. + await qm.append(sid, _line("e1", "/ws", {"session_id": sid})) + await qm.commit(sid, 0, {"dl2": {"orch_run_seq": 5}, "dl3": {}}) + assert (await qm.read_cursor(sid))["dl2"]["orch_run_seq"] == 5 + + # A REBUILT worker: a brand-new SessionWorker whose DataLayer2State starts + # at orch_run_seq=0. drain_worker must restore the persisted cursor before + # processing, so the counter is 5 (not 0) by the time an event is handled + # -- without the restore the next run would remint the already-used seq. + graph = _FlakyGraph() + worker = _make_worker(sid, graph) + assert worker.services.data_layer_2.orch_run_seq == 0 + reg._register_for_test(worker) + + captured: dict[str, int] = {} + + async def _capture(w: SessionWorker, event, data, handlers) -> None: + captured["seq_at_process"] = w.services.data_layer_2.orch_run_seq + w.services.graph.buffer.add("e1") + + with patch( + "context_intelligence_server.registry.process_event", + side_effect=_capture, + ): + reg.start_drain(worker) + await _drain_until_idle(reg, qm, worker, sid) + await _cancel_and_await(worker.task) + + assert captured.get("seq_at_process") == 5, ( + "drain_worker must restore the durable cursor (orch_run_seq=5) before " + "processing any event; a rebuilt worker that starts from 0 would " + "remint already-used run ids" + ) diff --git a/tests/test_durable_append_framing.py b/tests/test_durable_append_framing.py new file mode 100644 index 00000000..2e41bb88 --- /dev/null +++ b/tests/test_durable_append_framing.py @@ -0,0 +1,1060 @@ +"""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.queue_manager import filesystem as qm_module +from context_intelligence_server.queue_manager import FileSystemQueueManager, QueueManager +from context_intelligence_server.queue_manager.filesystem import _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 = FileSystemQueueManager(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 = FileSystemQueueManager(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 = FileSystemQueueManager(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 = FileSystemQueueManager(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 = FileSystemQueueManager(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 = FileSystemQueueManager(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 = FileSystemQueueManager(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 = FileSystemQueueManager(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 = FileSystemQueueManager(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 = FileSystemQueueManager(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 = FileSystemQueueManager(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 = FileSystemQueueManager(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 = FileSystemQueueManager(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, None) # 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 = FileSystemQueueManager(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, None) + 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, None) + 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 = FileSystemQueueManager(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 = FileSystemQueueManager(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 = FileSystemQueueManager(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, None) # 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 = FileSystemQueueManager(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 = FileSystemQueueManager(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_identity_store.py b/tests/test_identity_store.py index 61733c91..bc8cd6e9 100644 --- a/tests/test_identity_store.py +++ b/tests/test_identity_store.py @@ -14,7 +14,7 @@ import pytest -from context_intelligence_server.identity_store import IdentityStore +from context_intelligence_server.identity_store import FileSystemIdentityStore # --------------------------------------------------------------------------- @@ -42,7 +42,7 @@ def _bob_entry() -> dict[str, str]: class TestPutGetRoundtrip: def test_put_then_get_returns_value(self, tmp_path: Path) -> None: """put(key, value) → get(key) returns that value immediately.""" - store = IdentityStore(path=tmp_path / "store.json") + store = FileSystemIdentityStore(path=tmp_path / "store.json") store.load() store.put(FAKE_HASH_A, _alice_entry()) @@ -53,25 +53,25 @@ def test_put_then_get_returns_value(self, tmp_path: Path) -> None: def test_put_persists_to_file_and_loads_fresh(self, tmp_path: Path) -> None: """Round-trip: put → create new store → load → value is present.""" store_path = tmp_path / "store.json" - store = IdentityStore(path=store_path) + store = FileSystemIdentityStore(path=store_path) store.load() store.put(FAKE_HASH_A, _alice_entry()) # New store instance reads from disk - store2 = IdentityStore(path=store_path) + store2 = FileSystemIdentityStore(path=store_path) store2.load() assert store2.get(FAKE_HASH_A) == _alice_entry() def test_get_missing_key_returns_none(self, tmp_path: Path) -> None: """get() returns None for a key that was never put.""" - store = IdentityStore(path=tmp_path / "store.json") + store = FileSystemIdentityStore(path=tmp_path / "store.json") store.load() assert store.get(FAKE_HASH_A) is None def test_delete_removes_key(self, tmp_path: Path) -> None: """delete(key) removes the entry from in-process dict AND file.""" store_path = tmp_path / "store.json" - store = IdentityStore(path=store_path) + store = FileSystemIdentityStore(path=store_path) store.load() store.put(FAKE_HASH_A, _alice_entry()) store.delete(FAKE_HASH_A) @@ -79,13 +79,13 @@ def test_delete_removes_key(self, tmp_path: Path) -> None: assert store.get(FAKE_HASH_A) is None # Verify file is also updated - store2 = IdentityStore(path=store_path) + store2 = FileSystemIdentityStore(path=store_path) store2.load() assert store2.get(FAKE_HASH_A) is None def test_items_returns_all_entries(self, tmp_path: Path) -> None: """items() yields all key-value pairs currently in the store.""" - store = IdentityStore(path=tmp_path / "store.json") + store = FileSystemIdentityStore(path=tmp_path / "store.json") store.load() store.put(FAKE_HASH_A, _alice_entry()) store.put(FAKE_HASH_B, _bob_entry()) @@ -95,7 +95,7 @@ def test_items_returns_all_entries(self, tmp_path: Path) -> None: def test_upsert_overwrites_existing(self, tmp_path: Path) -> None: """put() on an existing key overwrites the value.""" - store = IdentityStore(path=tmp_path / "store.json") + store = FileSystemIdentityStore(path=tmp_path / "store.json") store.load() store.put(FAKE_HASH_A, _alice_entry()) store.put(FAKE_HASH_A, {"id": "alice-updated"}) @@ -105,12 +105,12 @@ def test_upsert_overwrites_existing(self, tmp_path: Path) -> None: def test_sequential_puts_all_persist(self, tmp_path: Path) -> None: """Multiple sequential puts all persist correctly (each write is the full map).""" store_path = tmp_path / "store.json" - store = IdentityStore(path=store_path) + store = FileSystemIdentityStore(path=store_path) store.load() store.put(FAKE_HASH_A, _alice_entry()) store.put(FAKE_HASH_B, _bob_entry()) - store2 = IdentityStore(path=store_path) + store2 = FileSystemIdentityStore(path=store_path) store2.load() assert store2.get(FAKE_HASH_A) == _alice_entry() assert store2.get(FAKE_HASH_B) == _bob_entry() @@ -124,7 +124,7 @@ def test_sequential_puts_all_persist(self, tmp_path: Path) -> None: class TestLoadFailClosed: def test_missing_file_yields_empty_dict(self, tmp_path: Path) -> None: """Missing store file → load() yields empty dict (normal first boot), no raise.""" - store = IdentityStore(path=tmp_path / "nonexistent.json") + store = FileSystemIdentityStore(path=tmp_path / "nonexistent.json") store.load() # must not raise assert store.get(FAKE_HASH_A) is None assert list(store.items()) == [] @@ -136,7 +136,7 @@ def test_corrupt_json_yields_empty_dict_and_logs_error( store_path = tmp_path / "store.json" store_path.write_text("{{{{not valid json at all}}}}", encoding="utf-8") - store = IdentityStore(path=store_path) + store = FileSystemIdentityStore(path=store_path) with caplog.at_level(logging.ERROR): store.load() # must NOT raise @@ -156,7 +156,7 @@ def test_valid_json_but_not_dict_yields_empty_and_logs( json.dumps([{"id": "alice"}]), encoding="utf-8" ) # list, not dict - store = IdentityStore(path=store_path) + store = FileSystemIdentityStore(path=store_path) with caplog.at_level(logging.ERROR): store.load() # must NOT raise @@ -171,7 +171,7 @@ def test_partial_write_torn_file_loads_empty( store_path = tmp_path / "store.json" store_path.write_bytes(b'{"aaa": {"id": "al') # truncated mid-write - store = IdentityStore(path=store_path) + store = FileSystemIdentityStore(path=store_path) with caplog.at_level(logging.ERROR): store.load() # must NOT raise @@ -188,7 +188,7 @@ def test_partial_write_torn_file_loads_empty( class TestWriteFileThenSwapMemory: def test_put_write_failure_leaves_dict_unchanged(self, tmp_path: Path) -> None: """If os.replace raises, the in-process dict is UNCHANGED (F2 contract).""" - store = IdentityStore(path=tmp_path / "store.json") + store = FileSystemIdentityStore(path=tmp_path / "store.json") store.load() # Establish an existing entry store.put(FAKE_HASH_A, _alice_entry()) @@ -208,7 +208,7 @@ def test_put_write_failure_leaves_dict_unchanged(self, tmp_path: Path) -> None: def test_delete_write_failure_leaves_dict_unchanged(self, tmp_path: Path) -> None: """If delete's file write fails, the in-process dict is UNCHANGED.""" - store = IdentityStore(path=tmp_path / "store.json") + store = FileSystemIdentityStore(path=tmp_path / "store.json") store.load() store.put(FAKE_HASH_A, _alice_entry()) @@ -221,7 +221,7 @@ def test_delete_write_failure_leaves_dict_unchanged(self, tmp_path: Path) -> Non def test_failed_write_leaves_no_torn_tempfile(self, tmp_path: Path) -> None: """A failed os.replace must clean up the tempfile — no orphaned .tmp files.""" - store = IdentityStore(path=tmp_path / "store.json") + store = FileSystemIdentityStore(path=tmp_path / "store.json") store.load() with patch("os.replace", side_effect=OSError("simulated disk full")): @@ -235,7 +235,7 @@ def test_failed_write_leaves_no_torn_tempfile(self, tmp_path: Path) -> None: def test_atomic_write_uses_tempfile_in_same_dir(self, tmp_path: Path) -> None: """Writes use a temp file in the same directory (then os.replace).""" store_path = tmp_path / "store.json" - store = IdentityStore(path=store_path) + store = FileSystemIdentityStore(path=store_path) store.load() replaced_from: list[str] = [] @@ -266,13 +266,13 @@ class TestFlatDictLiveReference: def test_flat_dict_empty_on_new_store(self, tmp_path: Path) -> None: """flat_dict is empty after load() with no file.""" - store = IdentityStore(path=tmp_path / "store.json") + store = FileSystemIdentityStore(path=tmp_path / "store.json") store.load() assert store.flat_dict == {} def test_flat_dict_updated_after_put(self, tmp_path: Path) -> None: """flat_dict is updated immediately after put().""" - store = IdentityStore(path=tmp_path / "store.json") + store = FileSystemIdentityStore(path=tmp_path / "store.json") store.load() store.put(FAKE_HASH_A, {"id": "alice"}) @@ -280,7 +280,7 @@ def test_flat_dict_updated_after_put(self, tmp_path: Path) -> None: def test_flat_dict_updated_after_delete(self, tmp_path: Path) -> None: """flat_dict removes key immediately after delete().""" - store = IdentityStore(path=tmp_path / "store.json") + store = FileSystemIdentityStore(path=tmp_path / "store.json") store.load() store.put(FAKE_HASH_A, {"id": "alice"}) store.delete(FAKE_HASH_A) @@ -290,7 +290,7 @@ def test_flat_dict_updated_after_delete(self, tmp_path: Path) -> None: def test_flat_dict_is_same_object_across_puts(self, tmp_path: Path) -> None: """flat_dict is the SAME dict object before and after put() (so a shared reference to flat_dict stays live).""" - store = IdentityStore(path=tmp_path / "store.json") + store = FileSystemIdentityStore(path=tmp_path / "store.json") store.load() flat_ref = store.flat_dict # capture the reference @@ -309,7 +309,7 @@ def test_flat_dict_populated_from_file_on_load(self, tmp_path: Path) -> None: json.dumps({FAKE_HASH_A: {"id": "alice"}, FAKE_HASH_B: {"id": "bob"}}), encoding="utf-8", ) - store = IdentityStore(path=store_path) + store = FileSystemIdentityStore(path=store_path) store.load() assert store.flat_dict[FAKE_HASH_A] == "alice" 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_lease_store.py b/tests/test_lease_store.py new file mode 100644 index 00000000..cf19fc1d --- /dev/null +++ b/tests/test_lease_store.py @@ -0,0 +1,99 @@ +"""Filesystem lease-store: the writer-lease persistence backend. + +The writer-lease detector (``writer_lease.py``) reaches the lease only through +this store, so these tests pin the persistence contract independently of the +detector's policy. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +from context_intelligence_server.lease_store import LeaseRecord, create_lease_store +from context_intelligence_server.lease_store.filesystem import ( + LEASE_FILENAME, + FileSystemLeaseStore, +) + +pytestmark = pytest.mark.integration + + +def _store(directory: Path) -> FileSystemLeaseStore: + store = create_lease_store(lambda: directory) + assert isinstance(store, FileSystemLeaseStore) + return store + + +def _record(owner: str = "me", heartbeat: float = 100.0) -> LeaseRecord: + return LeaseRecord( + owner=owner, + host="h", + pid=7, + started_at=1.0, + heartbeat=heartbeat, + revision="rev", + server_version="6.7.3", + lease_version=1, + ) + + +def test_read_missing_is_none(tmp_path: Path) -> None: + """A missing lease reads as None (a free directory), never an error.""" + assert _store(tmp_path).read() is None + + +def test_write_then_read_roundtrips(tmp_path: Path) -> None: + store = _store(tmp_path) + store.write(_record(owner="alice", heartbeat=42.0)) + got = store.read() + assert got is not None + assert got.owner == "alice" + assert got.heartbeat == 42.0 + assert got.unreadable is False + + +def test_write_is_atomic_no_tmp_left(tmp_path: Path) -> None: + store = _store(tmp_path) + store.write(_record()) + assert (tmp_path / LEASE_FILENAME).exists() + assert list(tmp_path.glob("*.tmp")) == [] + + +def test_torn_lease_reads_unreadable(tmp_path: Path) -> None: + """A hand-mangled lease is a synthetic unreadable record (fresh-foreign + strength), not None and not a crash.""" + (tmp_path / LEASE_FILENAME).write_text("{not json", encoding="utf-8") + got = _store(tmp_path).read() + assert got is not None + assert got.unreadable is True + + +def test_delete_if_owned_only_deletes_own(tmp_path: Path) -> None: + store = _store(tmp_path) + store.write(_record(owner="mine")) + + # A foreign lease is never deleted -- deleting it would hand the directory + # to a third writer. + store.delete_if_owned("someone_else") + assert store.read() is not None + + store.delete_if_owned("mine") + assert store.read() is None + + +def test_delete_if_owned_absent_is_noop(tmp_path: Path) -> None: + """Best-effort: deleting an already-absent lease is not an error.""" + _store(tmp_path).delete_if_owned("mine") # no raise + + +def test_dir_source_resolved_lazily(tmp_path: Path) -> None: + """The store constructs nothing and reads no path at build time -- the + directory is resolved per operation, so a store built before its directory + exists still works once it does.""" + target = tmp_path / "queues" + store = create_lease_store(lambda: target) + target.mkdir() # created AFTER the store was built + store.write(_record(owner="late")) + got = store.read() + assert got is not None and got.owner == "late" diff --git a/tests/test_m2_service_auth.py b/tests/test_m2_service_auth.py index a6f9781f..024454e0 100644 --- a/tests/test_m2_service_auth.py +++ b/tests/test_m2_service_auth.py @@ -123,18 +123,20 @@ def _service_claims(roles: list[str], appid: str = FAKE_APPID) -> dict[str, Any] class _MockBlobStore: - """Mock for AsyncDiskBlobStore — returns empty list, never touches filesystem.""" + """Mock BlobStore — empty list, never touches filesystem.""" - def __init__(self, root: Any) -> None: - pass - - async def list(self, session_id: str) -> list[str]: - return [] + async def list(self, session_id: str) -> AsyncGenerator[Any, None]: + return + yield # unreachable: makes list() an async generator async def read(self, uri: str) -> Any: raise FileNotFoundError(f"mock blob store: not found: {uri}") +def _mock_blob_store_factory(settings: Any) -> _MockBlobStore: + return _MockBlobStore() + + class _MockNeo4jResult: """Async-iterable result mock that yields a fixed list of rows.""" @@ -377,7 +379,7 @@ async def test_cap_sr_r_reader_read_capable( private_key, asgi = service_asgi token = _sign_jwt(private_key, _service_claims(roles=["Reader"])) - monkeypatch.setattr(main_module, "AsyncDiskBlobStore", _MockBlobStore) + monkeypatch.setattr(main_module, "create_blob_store", _mock_blob_store_factory) async with _make_client(asgi) as c: resp = await c.get( @@ -456,7 +458,7 @@ async def test_cap_sc_r_contributor_read_capable( private_key, asgi = service_asgi token = _sign_jwt(private_key, _service_claims(roles=["Contributor"])) - monkeypatch.setattr(main_module, "AsyncDiskBlobStore", _MockBlobStore) + monkeypatch.setattr(main_module, "create_blob_store", _mock_blob_store_factory) async with _make_client(asgi) as c: resp = await c.get( diff --git a/tests/test_main.py b/tests/test_main.py index 302fa1bf..192f2e8d 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 @@ -37,6 +37,37 @@ async def test_status_body(client: httpx.AsyncClient) -> None: assert data["active_sessions"] == 0 +async def test_status_exposes_schema_version_drift_fields( + client: httpx.AsyncClient, +) -> None: + """/status carries the compiled schema_version, the graph's stored version, + and a tri-state drift flag (True/False/None) -- distinct from server_version + and never a false "in sync".""" + from context_intelligence_server.status import SCHEMA_VERSION + + data = (await client.get("/status")).json() + assert data["schema_version"] == SCHEMA_VERSION + assert "graph_schema_version" in data + # No live graph in this client fixture -> unknown drift, never a false sync. + assert data["schema_version_current"] is None + # Distinct key from PR#78's server_version (no collision). + assert data["server_version"] != data["schema_version"] + + +async def test_status_carries_degraded_reason_via_boot( + client: httpx.AsyncClient, +) -> None: + """The single global degraded_reason surfaces under /status boot.""" + from context_intelligence_server.status import boot_state + + boot_state.degrade("3 node(s) lacking the :Node label") + try: + data = (await client.get("/status")).json() + assert data["boot"]["degraded_reason"] == "3 node(s) lacking the :Node label" + finally: + boot_state.clear_degraded() + + async def test_post_events_returns_202(client: httpx.AsyncClient) -> None: response = await client.post( "/events", @@ -162,12 +193,12 @@ 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).""" - from context_intelligence_server.queue_manager import QueueManager + """A durably-accepted event increments the registry accepted_total.""" + from context_intelligence_server.queue_manager import FileSystemQueueManager, QueueManager # Point the registry at a tmp queue dir so the durable append is isolated. monkeypatch.setattr( - main_module.registry, "_queue_manager", QueueManager(queues_dir=tmp_path) + main_module.registry, "_queue_manager", FileSystemQueueManager(queues_dir=tmp_path) ) monkeypatch.setattr( main_module.registry, "get_or_create", lambda *args, **kwargs: MagicMock() @@ -241,9 +272,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 +707,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 +724,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 +751,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 +775,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 +813,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 +822,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 +831,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 +855,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 +864,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 +894,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 +906,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 +946,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 +993,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 +1035,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 +1059,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,20 +1072,24 @@ 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() # stops reporting them (exactly what a real drainer does on completion). for sid in sids[:2]: batch = await qm.read_batch(sid, max_items=10) - await qm.commit(sid, batch.end_offset) + await qm.commit(sid, batch.end_offset, None) # Pass 2: the previously-DEFERRED tail is now dispatched -- not stranded. spawned.clear() - 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 +1110,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 +1139,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 +1170,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 +1200,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 +1220,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 +1251,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 +1273,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 +1296,125 @@ 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 + + +class TestSchemaAutoRepair: + """_schema_ready_or_auto_repair: lease-armed auto-repair on un-migrated data.""" + + @staticmethod + def _unmigrated_then_clean() -> AsyncMock: + # First schema check refuses (un-migrated), second succeeds (repaired). + return AsyncMock( + side_effect=[ + RuntimeError("42 node(s) lacking the :Node label -- doctor --fix"), + None, + ] + ) + + async def test_armed_lease_auto_repairs_and_rearms_gate(self) -> None: + main_module.app.state.schema_ready = False + main_module.app.state.neo4j_driver = MagicMock() + # Start DEGRADED so a passing test proves clear_degraded() actually ran, + # not that degraded_reason happened to already be None. + main_module.boot_state.degrade("42 node(s) lacking the :Node label") + run_op = AsyncMock() + + # First schema check refuses (un-migrated); the SECOND call -- the + # re-arm after repair -- is what must flip schema_ready. Model that by + # having the second call set it (as the real _ensure_schema_ready does), + # so a deleted re-arm block leaves schema_ready False and fails here. + call_count = {"n": 0} + + async def _schema_check() -> None: + call_count["n"] += 1 + if call_count["n"] == 1: + raise RuntimeError("42 node(s) lacking :Node -- doctor --fix") + main_module.app.state.schema_ready = True + + with ( + patch.object(main_module.writer_lease, "acquired", True), + patch( + "context_intelligence_server.main._ensure_schema_ready", + new=AsyncMock(side_effect=_schema_check), + ), + patch( + "context_intelligence_server.main.run_maintenance_operation", + new=run_op, + ), + ): + await main_module._schema_ready_or_auto_repair() + + run_op.assert_awaited_once() # repair ran (armed lease) + assert call_count["n"] == 2 # the re-arm re-check actually ran + assert main_module.app.state.schema_ready is True # gate re-armed + assert main_module.boot_state.degraded_reason is None # clear_degraded ran + + async def test_unarmed_lease_does_not_repair_and_stays_gated(self) -> None: + main_module.app.state.schema_ready = False + main_module.app.state.neo4j_driver = MagicMock() + main_module.boot_state.clear_degraded() + run_op = AsyncMock() + + with ( + patch.object(main_module.writer_lease, "acquired", False), + patch( + "context_intelligence_server.main._ensure_schema_ready", + new=AsyncMock( + side_effect=RuntimeError("42 node(s) lacking :Node -- doctor --fix") + ), + ), + patch( + "context_intelligence_server.main.run_maintenance_operation", + new=run_op, + ), + ): + await main_module._schema_ready_or_auto_repair() + + run_op.assert_not_awaited() # cross-replica safety: NO mutation unarmed + assert main_module.app.state.schema_ready is False # stays gated + assert main_module.boot_state.degraded_reason is not None + + async def test_never_raises_so_boot_continues_to_sweep(self) -> None: + # An unrepairable graph must not raise -- boot proceeds so the periodic + # sweep is still scheduled (the retry mechanism). + main_module.app.state.schema_ready = False + main_module.app.state.neo4j_driver = MagicMock() + with ( + patch.object(main_module.writer_lease, "acquired", False), + patch( + "context_intelligence_server.main._ensure_schema_ready", + new=AsyncMock(side_effect=RuntimeError("un-migrated")), + ), + ): + await main_module._schema_ready_or_auto_repair() # MUST NOT raise + + async def test_concurrent_op_running_skips_double_repair(self) -> None: + # If a maintenance op is already running (e.g. /admin/maintenance won the + # coordinator CAS), boot auto-repair must NOT double-run. + main_module.app.state.schema_ready = False + main_module.app.state.neo4j_driver = MagicMock() + run_op = AsyncMock() + with ( + patch.object(main_module.writer_lease, "acquired", True), + patch( + "context_intelligence_server.main._ensure_schema_ready", + new=AsyncMock(side_effect=RuntimeError("un-migrated")), + ), + patch.object( + main_module.coordinator, "try_begin_op", return_value=None + ), + patch( + "context_intelligence_server.main.run_maintenance_operation", + new=run_op, + ), + ): + await main_module._schema_ready_or_auto_repair() + + run_op.assert_not_awaited() # coordinator single-flight excluded it async def test_lifespan_does_not_raise_when_health_check_itself_fails( @@ -1274,7 +1428,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( @@ -1305,22 +1459,22 @@ async def test_lifespan_seeds_counters_from_disk(tmp_path: Path) -> None: zero residual: 1 committed + 1 pending line yields accepted=2, written=1, in_queue=1, residual=0 after reconcile -> seed_counts -> seed_counters. """ - from context_intelligence_server.queue_manager import QueueManager + from context_intelligence_server.queue_manager import FileSystemQueueManager, QueueManager from context_intelligence_server.registry import SessionRegistry # Seed a queue dir with one committed line and one still-pending line. - seed_qm = QueueManager(queues_dir=tmp_path) + seed_qm = FileSystemQueueManager(queues_dir=tmp_path) sid = "sess-seed" line1 = json.dumps({"event": "a", "workspace": "/ws", "data": {}}).encode("utf-8") line2 = json.dumps({"event": "b", "workspace": "/ws", "data": {}}).encode("utf-8") await seed_qm.append(sid, line1) await seed_qm.append(sid, line2) committed = len(line1) + 1 # +1 for the appended trailing newline - await seed_qm.commit(sid, committed) + await seed_qm.commit(sid, committed, None) # Fresh registry reusing the same on-disk queue dir. reg = SessionRegistry() - reg._queue_manager = QueueManager(queues_dir=tmp_path) + reg._queue_manager = FileSystemQueueManager(queues_dir=tmp_path) # Production order: reconcile dead lines BEFORE seeding the counts. await reg.queue_manager.recovery_reconcile_dead() @@ -1468,7 +1622,7 @@ async def test_status_includes_neo4j_query_connected_false_when_no_driver( # --------------------------------------------------------------------------- -# /status pipeline metrics block (D3) +# /status pipeline metrics block # --------------------------------------------------------------------------- @@ -1666,7 +1820,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 +1870,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 +1966,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 +2009,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_maintenance.py b/tests/test_maintenance.py new file mode 100644 index 00000000..f835d6f7 --- /dev/null +++ b/tests/test_maintenance.py @@ -0,0 +1,443 @@ +"""Tests for the MaintenanceCoordinator seam. + +Covers: +probe tri-state + fail-open, TTL caching, single-flight, the latch-defect +regression, the op-running/constraint-present interaction, and transition +logging. Allow-list/gate/status HTTP-surface tests live in test_main.py; +drain-loop offset-ordering tests live in test_registry.py; the docs +tripwire lives in test_docs_entrypoint.py; the maintenance-mode admin +surface also has coverage in test_main.py. +""" + +from __future__ import annotations + +import asyncio +import logging +from typing import Any, Self + +import pytest +from context_intelligence_server.maintenance import ( + MAINTENANCE_ALLOW_LIST, + MaintenanceCoordinator, +) + +# NOTE: asyncio_mode = "auto" (pyproject.toml) runs async tests automatically +# -- no pytest.mark.asyncio needed. A blanket `pytestmark` would incorrectly +# tag the sync tests below (TestOpSeam, test_allow_list_contains_required_paths). + + +# --------------------------------------------------------------------------- +# Fake Neo4j driver -- controls the constraint-probe result and counts real +# "catalog reads" so tests can assert single-flight/caching behavior. +# --------------------------------------------------------------------------- + + +class _FakeConstraintResult: + def __init__(self, count: int) -> None: + self._count = count + self._yielded = False + + def __aiter__(self) -> _FakeConstraintResult: + return self + + async def __anext__(self) -> dict[str, int]: + if self._yielded: + raise StopAsyncIteration + self._yielded = True + return {"c": self._count} + + +class _FakeConstraintSession: + def __init__(self, driver: _FakeConstraintDriver) -> None: + self._driver = driver + + async def __aenter__(self) -> Self: + return self + + async def __aexit__(self, *exc: object) -> None: + return None + + async def run(self, cypher: str) -> _FakeConstraintResult: + self._driver.call_count += 1 + if self._driver.raise_exc is not None: + raise self._driver.raise_exc + return _FakeConstraintResult(1 if self._driver.present else 0) + + +class _FakeConstraintDriver: + """Test double: counts real probe hits, lets tests flip the result.""" + + def __init__( + self, present: bool = True, raise_exc: Exception | None = None + ) -> None: + self.present = present + self.raise_exc = raise_exc + self.call_count = 0 + + def session(self, **kwargs: Any) -> _FakeConstraintSession: + return _FakeConstraintSession(self) + + +# --------------------------------------------------------------------------- +# Tri-state probe + fail-open rule +# --------------------------------------------------------------------------- + + +class TestProbeTriState: + async def test_constraint_absent_closes_gate(self) -> None: + coord = MaintenanceCoordinator() + coord.bind_driver(_FakeConstraintDriver(present=False)) + assert await coord.gate_closed() is True + + async def test_constraint_present_opens_gate(self) -> None: + coord = MaintenanceCoordinator() + coord.bind_driver(_FakeConstraintDriver(present=True)) + assert await coord.gate_closed() is False + + async def test_probe_raises_yields_unknown_and_gate_open(self) -> None: + coord = MaintenanceCoordinator() + coord.bind_driver( + _FakeConstraintDriver(raise_exc=RuntimeError("neo4j unreachable")) + ) + st = await coord.status() + assert st.constraint_present is None + assert st.mode == "unknown" + assert await coord.gate_closed() is False # unknown must NOT close the gate + + +# --------------------------------------------------------------------------- +# TTL cache works +# --------------------------------------------------------------------------- + + +class TestProbeTtlCache: + async def test_two_calls_within_ttl_hit_driver_once(self) -> None: + driver = _FakeConstraintDriver(present=True) + coord = MaintenanceCoordinator() + coord.bind_driver(driver, probe_ttl_seconds=10.0) + + await coord.gate_closed() + await coord.gate_closed() + + assert driver.call_count == 1 + + async def test_call_after_ttl_expiry_hits_driver_again(self) -> None: + driver = _FakeConstraintDriver(present=True) + coord = MaintenanceCoordinator() + coord.bind_driver(driver, probe_ttl_seconds=0.05) + + await coord.gate_closed() + # First in-window call: constraint catalog read (cold) + untagged + # count (seeded warm at bind_driver, so 0 reads here) == 1 read. + hits_after_first = driver.call_count + assert hits_after_first == 1 + await asyncio.sleep(0.1) # both caches expire + await coord.gate_closed() + + # After expiry a FRESH probe runs -- the load-bearing property is that + # the cache expired and re-probed, not the exact read count (which now + # also includes the live O(1) untagged count, present-constraint path). + assert driver.call_count > hits_after_first + + +# --------------------------------------------------------------------------- +# Single-flight: N concurrent callers at cache expiry -> ONE probe +# --------------------------------------------------------------------------- + + +class TestProbeSingleFlight: + async def test_fifty_concurrent_calls_at_expiry_yield_one_catalog_read( + self, + ) -> None: + driver = _FakeConstraintDriver(present=True) + coord = MaintenanceCoordinator() + coord.bind_driver(driver, probe_ttl_seconds=0.02) + + # Prime + expire the cache once, deterministically. + await coord.gate_closed() + await asyncio.sleep(0.05) + hits_before = driver.call_count + assert hits_before == 1 + + results = await asyncio.gather(*[coord.gate_closed() for _ in range(50)]) + + # Single-flight: all 50 callers at the expiry boundary collapse to ONE + # probe window, NOT 50. That window now reads the constraint catalog + # AND the live O(1) untagged count (constraint-present path), so the + # increment is a small constant (<= 3), never proportional to the 50 + # callers -- which is the property this test guards. + new_reads = driver.call_count - hits_before + assert new_reads <= 3 + assert all(r is False for r in results) # constraint present -> gate open + + +# --------------------------------------------------------------------------- +# THE LATCH DEFECT, DIRECTLY: probe flips False->True with NO restart +# --------------------------------------------------------------------------- + + +class TestNoRestartSelfClear: + async def test_gate_opens_within_ttl_with_no_restart(self) -> None: + """This is the exact bug this spec exists to fix: a boot-latched + schema_health never re-probed until the next restart. Here the SAME + coordinator instance (no restart) observes the constraint appear.""" + driver = _FakeConstraintDriver(present=False) + coord = MaintenanceCoordinator() + coord.bind_driver(driver, probe_ttl_seconds=0.05) + + assert await coord.gate_closed() is True # migration required + + # Out-of-band repair happens (e.g. `doctor --fix` in another + # process) -- constraint now exists. NO restart of this process. + driver.present = True + await asyncio.sleep(0.1) # let the TTL expire + + assert await coord.gate_closed() is False # self-cleared, no restart + + async def test_non_vacuous_without_ttl_expiry_gate_stays_stale(self) -> None: + """Non-vacuity proof: if we DON'T wait out the TTL, + the cached (stale) answer is returned -- proving the test above is + actually exercising the cache-expiry path, not a tautology.""" + driver = _FakeConstraintDriver(present=False) + coord = MaintenanceCoordinator() + coord.bind_driver(driver, probe_ttl_seconds=10.0) # long TTL + + assert await coord.gate_closed() is True + + driver.present = True # flips immediately, but cache is still warm + # NO sleep -- still within the 10s TTL window. + assert await coord.gate_closed() is True # stale cached answer + + +# --------------------------------------------------------------------------- +# A4b -- THE UNTAGGED LATCH: degraded->healthy self-clears with NO restart +# --------------------------------------------------------------------------- + + +class _FakeUntaggedResult: + def __init__(self, count: int) -> None: + self._count = count + self._yielded = False + + def __aiter__(self) -> _FakeUntaggedResult: + return self + + async def __anext__(self) -> dict[str, int]: + if self._yielded: + raise StopAsyncIteration + self._yielded = True + return {"c": self._count} + + +class _FakeUntaggedSession: + def __init__(self, driver: _FakeUntaggedDriver) -> None: + self._driver = driver + + async def __aenter__(self) -> Self: + return self + + async def __aexit__(self, *exc: object) -> None: + return None + + async def run(self, cypher: str, *args: Any, **kwargs: Any) -> _FakeUntaggedResult: + self._driver.call_count += 1 + if "SHOW CONSTRAINTS" in cypher: + return _FakeUntaggedResult(1) # constraint present + if ":Node)" in cypher: + return _FakeUntaggedResult(self._driver.tagged) # tagged-node count + return _FakeUntaggedResult(self._driver.total) # total-node count + + +class _FakeUntaggedDriver: + """Test double: constraint always present; untagged = total - tagged, and + both are flippable at runtime so a test can simulate an out-of-band repair + tagging the last untagged node WITHOUT a restart.""" + + def __init__(self, total: int, tagged: int) -> None: + self.total = total + self.tagged = tagged + self.call_count = 0 + + def session(self, **kwargs: Any) -> _FakeUntaggedSession: + return _FakeUntaggedSession(self) + + +class TestUntaggedNoRestartSelfClear: + async def test_degraded_self_clears_within_ttl_no_restart(self) -> None: + """The untagged half of the latch: schema_health/mode report `degraded` + (1 node lacking :Node) at boot; an out-of-band repair tags it; the SAME + coordinator (no restart) reports `healthy` within one TTL.""" + driver = _FakeUntaggedDriver(total=1, tagged=0) # 1 untagged -> degraded + coord = MaintenanceCoordinator() + coord.bind_driver(driver, untagged=1, probe_ttl_seconds=0.05) + + st = await coord.status() + assert st.mode == "degraded" + assert st.untagged_nodes == 1 + assert await coord.gate_closed() is False # degraded NEVER closes gate + + # Out-of-band repair tags the node -- untagged now 0. NO restart. + driver.tagged = 1 + await asyncio.sleep(0.1) # let the untagged-probe TTL expire + + st2 = await coord.status() + assert st2.mode == "healthy" # self-cleared, no restart + assert st2.untagged_nodes == 0 + assert st2.reason is None + + async def test_non_vacuous_within_ttl_stays_degraded(self) -> None: + """Non-vacuity: without waiting out the TTL the seeded/cached count is + returned -- proving the test above exercises the expiry path, not a + tautology.""" + driver = _FakeUntaggedDriver(total=1, tagged=0) + coord = MaintenanceCoordinator() + coord.bind_driver(driver, untagged=1, probe_ttl_seconds=10.0) # long TTL + + st = await coord.status() + assert st.mode == "degraded" + + driver.tagged = 1 # repaired immediately, but cache still warm + # NO sleep -- still within the 10s TTL window. + st2 = await coord.status() + assert st2.mode == "degraded" # stale cached count + assert st2.untagged_nodes == 1 + + +# --------------------------------------------------------------------------- +# op_running holds the gate closed even when constraint IS present +# --------------------------------------------------------------------------- + + +class TestOpRunningKeepsGateClosed: + async def test_op_running_with_constraint_present_still_closed(self) -> None: + coord = MaintenanceCoordinator() + coord.bind_driver(_FakeConstraintDriver(present=True)) + assert await coord.gate_closed() is False # sanity: open beforehand + + run_id = coord.try_begin_op() + assert run_id is not None + + assert await coord.gate_closed() is True # op_running term + + coord.finish_op(run_id, records_affected=0, error=None) + assert await coord.gate_closed() is False # reopens once finished + + +# --------------------------------------------------------------------------- +# Transition logging fires EXACTLY once per transition +# --------------------------------------------------------------------------- + + +class TestTransitionLogging: + async def test_closed_open_closed_emits_exactly_one_each( + self, caplog: pytest.LogCaptureFixture + ) -> None: + driver = _FakeConstraintDriver( + present=False + ) # start closed (migration required) + coord = MaintenanceCoordinator() + coord.bind_driver(driver, probe_ttl_seconds=0.02) + + with caplog.at_level(logging.INFO, logger="context_intelligence_server"): + # Already "maintenance" from bind (closed=migration required). + await coord.status() + await coord.status() # repeat call: must not double-log + + # Flip open. + driver.present = True + await asyncio.sleep(0.05) + await coord.status() + await coord.status() # repeat: must not double-log + + # Flip closed again. + driver.present = False + await asyncio.sleep(0.05) + await coord.status() + + entered = [r for r in caplog.records if r.getMessage() == "maintenance_entered"] + completed = [ + r for r in caplog.records if r.getMessage() == "maintenance_completed" + ] + # Two "entered" (initial closed state, then re-closed) and one + # "completed" (the middle open window). + assert len(entered) == 2 + assert len(completed) == 1 + + +# --------------------------------------------------------------------------- +# No driver bound => gate OPEN (the existing-suite no-regression rule) +# --------------------------------------------------------------------------- + + +class TestNoDriverBound: + async def test_unbound_coordinator_gate_is_open(self) -> None: + coord = MaintenanceCoordinator() # bind_driver() never called + assert await coord.gate_closed() is False + st = await coord.status() + assert st.mode == "unknown" + assert st.constraint_present is None + + +# --------------------------------------------------------------------------- +# try_begin_op / finish_op / current_op -- the CAS seam itself +# --------------------------------------------------------------------------- + + +class TestOpSeam: + def test_try_begin_op_is_synchronous_cas(self) -> None: + coord = MaintenanceCoordinator() + run_id_1 = coord.try_begin_op() + assert run_id_1 is not None + run_id_2 = coord.try_begin_op() + assert run_id_2 is None # already running -- CAS refuses a second op + + def test_finish_op_sets_completed_at_and_state(self) -> None: + coord = MaintenanceCoordinator() + run_id = coord.try_begin_op() + assert run_id is not None + assert coord.current_op().completed_at is None + + coord.finish_op(run_id, records_affected=3, error=None) + + op = coord.current_op() + assert op.state == "succeeded" + assert op.completed_at is not None + assert op.records_affected == 3 + + def test_finish_op_with_error_marks_failed(self) -> None: + coord = MaintenanceCoordinator() + run_id = coord.try_begin_op() + assert run_id is not None + + coord.finish_op(run_id, records_affected=None, error="boom") + + op = coord.current_op() + assert op.state == "failed" + assert op.error == "boom" + + def test_op_state_initializes_unknown_not_succeeded(self) -> None: + """Council D4: never-run must not read as a false 'succeeded'.""" + coord = MaintenanceCoordinator() + assert coord.current_op().state == "unknown" + + def test_finish_op_stale_run_id_is_ignored(self) -> None: + coord = MaintenanceCoordinator() + run_id = coord.try_begin_op() + assert run_id is not None + + coord.finish_op("not-the-real-run-id", records_affected=1, error=None) + + # The real op is untouched -- still running, not clobbered by a + # foreign/stale completion signal. + assert coord.current_op().state == "running" + + +# --------------------------------------------------------------------------- +# Allow-list sanity (used by the startup assertion + the middleware) +# --------------------------------------------------------------------------- + + +def test_allow_list_contains_required_paths() -> None: + assert "/admin/maintenance" in MAINTENANCE_ALLOW_LIST + assert "/status" in MAINTENANCE_ALLOW_LIST + assert "/version" in MAINTENANCE_ALLOW_LIST diff --git a/tests/test_migrations_run.py b/tests/test_migrations_run.py new file mode 100644 index 00000000..36863d9b --- /dev/null +++ b/tests/test_migrations_run.py @@ -0,0 +1,203 @@ +"""Tests for migrations/run.py -- the standalone, out-of-band graph +rectification CLI. Unit-level only (mocked driver/config); real +rectification against a live Neo4j is DTU-validated separately. +""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from migrations import run as run_module + + +def _fake_driver() -> MagicMock: + driver = MagicMock() + driver.verify_connectivity = AsyncMock(return_value=None) + driver.close = AsyncMock(return_value=None) + session_cm = MagicMock() + session_cm.__aenter__ = AsyncMock(return_value=session_cm) + session_cm.__aexit__ = AsyncMock(return_value=False) + driver.session = MagicMock(return_value=session_cm) + return driver + + +@pytest.fixture(autouse=True) +def _patch_settings_and_driver(): + """Every test patches get_settings/build_neo4j_driver so no real config + or Neo4j connection is required (mirrors tests/test_doctor.py).""" + with ( + patch.object(run_module, "get_settings") as mock_get_settings, + patch.object(run_module, "build_neo4j_driver") as mock_build_driver, + ): + admin_config = MagicMock() + admin_config.model_copy.return_value = "admin-cfg" + mock_get_settings.return_value.resolve_neo4j_admin.return_value = admin_config + mock_build_driver.return_value = _fake_driver() + yield mock_build_driver.return_value + + +def test_module_imports_cleanly() -> None: + """The module must import without touching Neo4j (no side effects at + import time) -- required for the DTU update script's `grep` probe.""" + import migrations.run # noqa: F401 -- import-cleanliness check + + +def test_status_flag_exists_and_parses() -> None: + args = run_module.build_parser().parse_args(["--status"]) + assert args.status is True + assert args.apply is False + + +def test_apply_flag_exists_and_parses() -> None: + args = run_module.build_parser().parse_args(["--apply"]) + assert args.apply is True + assert args.status is False + + +def test_status_and_apply_are_mutually_exclusive() -> None: + with pytest.raises(SystemExit): + run_module.build_parser().parse_args(["--status", "--apply"]) + + +def test_one_of_status_or_apply_is_required() -> None: + with pytest.raises(SystemExit): + run_module.build_parser().parse_args([]) + + +def test_neo4j_overrides_are_optional_flags() -> None: + args = run_module.build_parser().parse_args( + [ + "--status", + "--neo4j-url", + "bolt://example:7687", + "--neo4j-user", + "u", + "--neo4j-password", + "p", + ] + ) + assert args.neo4j_url == "bolt://example:7687" + assert args.neo4j_user == "u" + assert args.neo4j_password == "p" + + +def test_banner_declares_from_to_and_out_of_band( + capsys: pytest.CaptureFixture[str], +) -> None: + run_module._print_banner() + out = capsys.readouterr().out + assert run_module.FROM_SERVER_VERSION in out + assert run_module.TO_SERVER_VERSION in out + assert str(run_module.FROM_SCHEMA_VERSION) in out + assert str(run_module.TO_SCHEMA_VERSION) in out + assert "OUT-OF-BAND" in out + assert "IDEMPOTENT" in out + assert "never runs at server startup" in out + + +async def test_status_mode_does_not_call_run_repair( + _patch_settings_and_driver: MagicMock, +) -> None: + with ( + patch.object( + run_module, + "diagnose", + AsyncMock(return_value={"untagged_nodes": 0, "duplicate_nodes": 0}), + ), + patch.object(run_module, "run_repair", AsyncMock()) as repair_mock, + patch.object(run_module, "_constraint_present", AsyncMock(return_value=True)), + ): + args = run_module.build_parser().parse_args(["--status"]) + code = await run_module._amain(args) + + assert code == 0 + repair_mock.assert_not_awaited() + + +async def test_apply_mode_calls_run_repair( + _patch_settings_and_driver: MagicMock, +) -> None: + diagnose_mock = AsyncMock( + side_effect=[ + {"untagged_nodes": 5, "duplicate_nodes": 2}, # before + {"untagged_nodes": 0, "duplicate_nodes": 0}, # after + ] + ) + repair_mock = AsyncMock(return_value={"duplicates_removed": 2, "nodes_tagged": 5}) + with ( + patch.object(run_module, "diagnose", diagnose_mock), + patch.object(run_module, "run_repair", repair_mock), + patch.object(run_module, "_constraint_present", AsyncMock(return_value=True)), + ): + args = run_module.build_parser().parse_args(["--apply"]) + code = await run_module._amain(args) + + assert code == 0 + repair_mock.assert_awaited_once() + assert diagnose_mock.await_count == 2 + + +async def test_apply_is_idempotent_noop_on_already_clean_graph( + _patch_settings_and_driver: MagicMock, +) -> None: + """--apply still calls run_repair (it is itself idempotent/no-op-safe), + but a clean before-state must still report healthy after.""" + diagnose_mock = AsyncMock(return_value={"untagged_nodes": 0, "duplicate_nodes": 0}) + repair_mock = AsyncMock(return_value={"duplicates_removed": 0, "nodes_tagged": 0}) + with ( + patch.object(run_module, "diagnose", diagnose_mock), + patch.object(run_module, "run_repair", repair_mock), + patch.object(run_module, "_constraint_present", AsyncMock(return_value=True)), + ): + args = run_module.build_parser().parse_args(["--apply"]) + code = await run_module._amain(args) + + assert code == 0 + repair_mock.assert_awaited_once() + + +async def test_status_unreachable_neo4j_fails_gracefully( + _patch_settings_and_driver: MagicMock, +) -> None: + """Pointed at an unreachable URL, --status must not crash/traceback -- + it reports the connectivity failure and returns a non-zero exit code.""" + driver = _patch_settings_and_driver + driver.verify_connectivity = AsyncMock( + side_effect=RuntimeError("connection refused") + ) + + args = run_module.build_parser().parse_args(["--status"]) + code = await run_module._amain(args) + + assert code != 0 + driver.close.assert_awaited_once() + + +async def test_status_unreachable_neo4j_never_touches_diagnose_or_repair( + _patch_settings_and_driver: MagicMock, +) -> None: + driver = _patch_settings_and_driver + driver.verify_connectivity = AsyncMock( + side_effect=RuntimeError("connection refused") + ) + with ( + patch.object(run_module, "diagnose", AsyncMock()) as diagnose_mock, + patch.object(run_module, "run_repair", AsyncMock()) as repair_mock, + ): + args = run_module.build_parser().parse_args(["--status"]) + await run_module._amain(args) + + diagnose_mock.assert_not_awaited() + repair_mock.assert_not_awaited() + + +async def test_constraint_present_returns_none_on_probe_failure( + _patch_settings_and_driver: MagicMock, +) -> None: + driver = _patch_settings_and_driver + driver.session.side_effect = RuntimeError("catalog read failed") + + result = await run_module._constraint_present(driver) + + assert result is None 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..36c97d97 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -488,8 +488,12 @@ async def test_blob_processing_called_when_all_conditions_met( ) as mock_node_id, ): await process_event(worker, "session:start", data, pipeline_handlers) + # The blob-key node_id must be computed with the SAME disambiguator + # (tool_call_id) that handlers/data_layer_1/default.py uses for the + # Event node id -- otherwise parallel same-millisecond events collide + # on the blob key (see test_blob_node_id_matches_default_handler_event_node_id). mock_node_id.assert_called_once_with( - "sess-123", "session:start", "2024-01-01T00:00:00Z" + "sess-123", "session:start", "2024-01-01T00:00:00Z", None ) mock_process.assert_called_once_with( data, worker.services.blob_store, "sess-123", "test-node-id" @@ -566,6 +570,123 @@ async def test_blob_skip_missing_timestamp_logs_warning( assert "missing timestamp" in caplog.text +# =========================================================================== +# Blob-key collision regression (data integrity) +# +# Root cause: pipeline.py computed the blob-key node_id WITHOUT the +# tool_call_id disambiguator that handlers/data_layer_1/default.py uses for +# the Event node id. Two distinct same-session, same-event, same-millisecond +# events with DIFFERENT tool_call_id (e.g. parallel tool calls) therefore +# minted the SAME blob key, and the second write silently clobbered the +# first via FileSystemBlobStore's os.replace -- while the first Event node's +# $blob_ref still pointed at that (now-overwritten) URI. +# =========================================================================== + + +async def test_parallel_same_millisecond_events_do_not_collide_on_blob_key( + pipeline_handlers: Any, + tmp_path: Any, +) -> None: + """Two events sharing session_id + event name + timestamp (same epoch-ms) + but with DIFFERENT tool_call_id must mint DISTINCT ci-blob:// URIs, and + each blob must read back its own payload -- no silent overwrite.""" + from context_intelligence_server.blob_store import FileSystemBlobStore + from context_intelligence_server.pipeline import process_event + + blob_store = FileSystemBlobStore(root=tmp_path) + + worker = MagicMock() + worker.services.ensure_session_node = AsyncMock() + worker.services.touch_session = AsyncMock() + worker.services.graph = MagicMock() + worker.services.graph.flush = AsyncMock() + worker.services.blob_store = blob_store + + session_id = "sess-parallel" + event_name = "tool_call:end" + timestamp = "2024-06-01T12:00:00.000Z" # fixed -- identical epoch-ms for both + + data_a: dict[str, Any] = { + "session_id": session_id, + "timestamp": timestamp, + "tool_call_id": "call-A", + "result": {"payload": "result-from-call-A"}, + } + data_b: dict[str, Any] = { + "session_id": session_id, + "timestamp": timestamp, + "tool_call_id": "call-B", + "result": {"payload": "result-from-call-B"}, + } + + await process_event(worker, event_name, data_a, pipeline_handlers) + await process_event(worker, event_name, data_b, pipeline_handlers) + + # (c) each event's data[field] == {"$blob_ref": } + assert "$blob_ref" in data_a["result"] + assert "$blob_ref" in data_b["result"] + uri_a = data_a["result"]["$blob_ref"] + uri_b = data_b["result"]["$blob_ref"] + + # (a) distinct URIs -- no collision + assert uri_a != uri_b, ( + f"Blob key collision: both events minted the same URI {uri_a!r} -- " + "the second write silently overwrote the first's blob." + ) + + # (b) both blobs exist and each reads back its OWN distinct payload + read_a = await blob_store.read(uri_a) + read_b = await blob_store.read(uri_b) + assert read_a == {"payload": "result-from-call-A"} + assert read_b == {"payload": "result-from-call-B"} + + +async def test_blob_node_id_matches_default_handler_event_node_id( + pipeline_handlers: Any, +) -> None: + """Pins the invariant: the blob-key node_id pipeline.process_event computes + must EQUAL the event_node_id handlers/data_layer_1/default.py computes for + the same event (make_node_id(session_id, event, timestamp, + data.get("tool_call_id"))). If these ever drift apart again, the + collision this test suite guards against reappears.""" + from context_intelligence_server.pipeline import process_event + from context_intelligence_server.utils import make_node_id + + worker = MagicMock() + worker.services.ensure_session_node = AsyncMock() + worker.services.touch_session = AsyncMock() + worker.services.graph = MagicMock() + worker.services.graph.flush = AsyncMock() + worker.services.blob_store = MagicMock() # truthy blob_store + + session_id = "sess-invariant" + event_name = "tool_call:end" + timestamp = "2024-06-01T12:00:00.000Z" + tool_call_id = "call-invariant" + + data = { + "session_id": session_id, + "timestamp": timestamp, + "tool_call_id": tool_call_id, + } + + expected_event_node_id = make_node_id( + session_id, event_name, timestamp, tool_call_id + ) + + with patch( + "context_intelligence_server.pipeline.process_event_data", + new_callable=AsyncMock, + ) as mock_process: + await process_event(worker, event_name, data, pipeline_handlers) + actual_blob_node_id = mock_process.call_args.args[3] + + assert actual_blob_node_id == expected_event_node_id, ( + "Blob-key node_id has drifted from default.py's event_node_id -- " + "this reintroduces the same-millisecond blob-key collision." + ) + + # =========================================================================== # process_event — touch_session call site # =========================================================================== diff --git a/tests/test_queue_manager.py b/tests/test_queue_manager.py index bd270601..541145af 100644 --- a/tests/test_queue_manager.py +++ b/tests/test_queue_manager.py @@ -1,34 +1,106 @@ -"""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, + FileSystemQueueManager, + Record, + Verdict, +) @pytest.fixture def qm(tmp_path): - return QueueManager(queues_dir=tmp_path / "queues") + return FileSystemQueueManager(queues_dir=tmp_path / "queues") def test_constructor_creates_queues_dir(tmp_path): target = tmp_path / "nested" / "queues" assert not target.exists() - QueueManager(queues_dir=target) + FileSystemQueueManager(queues_dir=target) assert target.is_dir() 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, None) + + 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" @@ -136,7 +208,7 @@ async def test_commit_advances_offset(qm): await qm.append("s1", b"a") await qm.append("s1", b"b") first = await qm.read_batch("s1", max_items=1) - await qm.commit("s1", first.end_offset) + await qm.commit("s1", first.end_offset, None) await qm.append("s1", b"c") second = await qm.read_batch("s1", max_items=10) assert second.lines == [b"b", b"c"] @@ -145,33 +217,135 @@ async def test_commit_advances_offset(qm): async def test_commit_persists_across_a_new_instance(tmp_path): qdir = tmp_path / "queues" - qm1 = QueueManager(queues_dir=qdir) + qm1 = FileSystemQueueManager(queues_dir=qdir) await qm1.append("s1", b"a") await qm1.append("s1", b"b") batch = await qm1.read_batch("s1", max_items=1) - await qm1.commit("s1", batch.end_offset) - qm2 = QueueManager(queues_dir=qdir) # simulate restart + await qm1.commit("s1", batch.end_offset, None) + qm2 = FileSystemQueueManager(queues_dir=qdir) # simulate restart resumed = await qm2.read_batch("s1", max_items=10) assert resumed.lines == [b"b"] async def test_commit_is_atomic_no_temp_leftover(qm, tmp_path): await qm.append("s1", b"a") - await qm.commit("s1", 2) + await qm.commit("s1", 2, None) qdir = tmp_path / "queues" - assert (qdir / "s1.offset").read_text("utf-8") == "2" + assert (qdir / "s1.offset").read_text("utf-8") == '{"v":1,"offset":2,"cursor":null}' assert list(qdir.glob("*.tmp")) == [] +# _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") done = await qm.read_batch("s_done", max_items=10) - await qm.commit("s_done", done.end_offset) # drained + await qm.commit("s_done", done.end_offset, None) # drained active = await qm.active_sessions() assert active == ["s_active"] +async def test_is_fully_drained_true_for_unknown_session(qm): + assert await qm.is_fully_drained("never_seen") is True + + +async def test_is_fully_drained_false_while_uncommitted(qm): + await qm.append("s1", b"x") + assert await qm.is_fully_drained("s1") is False + + +async def test_is_fully_drained_true_after_commit(qm): + await qm.append("s1", b"x") + batch = await qm.read_batch("s1", max_items=10) + await qm.commit("s1", batch.end_offset, None) + assert await qm.is_fully_drained("s1") is True + + +async def test_is_fully_drained_ignores_torn_trailing_fragment(qm, tmp_path): + # A torn tail (bytes after the final newline) is not complete data, so a + # session whose complete lines are all committed reads as drained. + log = tmp_path / "queues" / "s1.log" + log.write_bytes(b"a\nb\nTORN_PARTIAL") + batch = await qm.read_batch("s1", max_items=10) + await qm.commit("s1", batch.end_offset, None) + assert await qm.is_fully_drained("s1") is True + + async def test_recover_empty_dir_is_safe(qm): assert await qm.recover() == [] @@ -180,7 +354,7 @@ async def test_recover_reports_session_with_uncommitted_complete_line(qm, tmp_pa log = tmp_path / "queues" / "s1.log" log.write_bytes(b"a\nb\nTORN") # two complete lines + torn tail assert await qm.recover() == ["s1"] - await qm.commit("s1", 4) # past 'a\nb\n' == 4 bytes + await qm.commit("s1", 4, None) # past 'a\nb\n' == 4 bytes assert await qm.recover() == [] # only torn tail remains -> not recoverable @@ -211,7 +385,7 @@ async def test_read_batch_rejects_unsafe_session_id(qm, bad_id): @pytest.mark.parametrize("bad_id", ["", "a/b", "a\\b", "a\x00b"]) async def test_commit_rejects_unsafe_session_id(qm, bad_id): with pytest.raises(ValueError): - await qm.commit(bad_id, 0) + await qm.commit(bad_id, 0, None) @pytest.mark.parametrize("bad_id", ["", "a/b", "a\\b", "a\x00b"]) @@ -227,11 +401,11 @@ async def test_read_dead_letters_rejects_unsafe_session_id(qm, bad_id): async def test_delete_drained_removes_log_and_offset_keeps_dead(tmp_path) -> None: - from context_intelligence_server.queue_manager import QueueManager + from context_intelligence_server.queue_manager import FileSystemQueueManager - qm = QueueManager(queues_dir=tmp_path) + qm = FileSystemQueueManager(queues_dir=tmp_path) await qm.append("s", b"line") - await qm.commit("s", 5) + await qm.commit("s", 5, None) await qm.dead_letter("s", b"bad\n", "boom") await qm.delete_drained("s") @@ -317,7 +491,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): @@ -325,7 +499,7 @@ async def test_recovery_seed_counts_pending_and_committed(qm): await qm.append("s1", b"a") await qm.append("s1", b"b") await qm.append("s1", b"c") - await qm.commit("s1", 4) # commit the first two complete lines + await qm.commit("s1", 4, None) # commit the first two complete lines accepted, written = await qm.recovery_seed_counts() @@ -336,7 +510,7 @@ async def test_recovery_seed_counts_pending_and_committed(qm): async def test_recovery_seed_counts_committed_includes_dead(qm): # C=1 committed, P=0 pending, D=1 dead. before-dead == 0. await qm.append("s2", b"a") - await qm.commit("s2", 2) + await qm.commit("s2", 2, None) await qm.dead_letter("s2", b"a", error="boom") accepted, written = await qm.recovery_seed_counts() @@ -360,10 +534,10 @@ async def test_recovery_seed_counts_residual_is_zero_mixed_shape(qm): await qm.append("a", b"1") await qm.append("a", b"2") await qm.append("a", b"3") - await qm.commit("a", 4) + await qm.commit("a", 4, None) # Key B: 1 committed + 1 dead. await qm.append("b", b"x") - await qm.commit("b", 2) + await qm.commit("b", 2, None) await qm.dead_letter("b", b"x", error="boom") # Key C: dead-only (log reclaimed). await qm.dead_letter("c", b"poison", error="boom") @@ -393,7 +567,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): @@ -453,7 +627,7 @@ async def test_recovery_seed_counts_replay_window_residual_zero(qm): # log = [line0 committed][line0 re-appended pending]. C=1, P=1, D=1. # The re-appended line is absorbed into accepted_seed (counted in P and D). await qm.append("s6", b"a") - await qm.commit("s6", 2) + await qm.commit("s6", 2, None) await qm.dead_letter("s6", b"a", error="boom") await qm.append("s6", b"a") # re-append the dead line for replay @@ -493,7 +667,7 @@ async def test_spool_stats_fully_committed_session_not_pending(qm): (spool_bytes_total still reflects them).""" await qm.append("s1", b"a") line = b"a\n" - await qm.commit("s1", len(line)) + await qm.commit("s1", len(line), None) stats = await qm.spool_stats() @@ -518,7 +692,7 @@ async def test_spool_stats_multiple_sessions_aggregate(qm): await qm.append("s1", b"a") # pending await qm.append("s2", b"b") line = b"b\n" - await qm.commit("s2", len(line)) # fully committed, not pending + await qm.commit("s2", len(line), None) # fully committed, not pending await qm.append("s3", b"c") # pending stats = await qm.spool_stats() @@ -613,7 +787,7 @@ def test_complete_data_end_newline_on_chunk_boundary(qm, tmp_path, monkeypatch): """The backward scan reads fixed non-overlapping windows; a newline landing exactly on a chunk boundary must still be found (regression guard for the streaming rewrite).""" - import context_intelligence_server.queue_manager as qm_mod + import context_intelligence_server.queue_manager.filesystem as qm_mod monkeypatch.setattr(qm_mod, "_SCAN_CHUNK_BYTES", 8) log = tmp_path / "queues" / "s1.log" @@ -628,7 +802,7 @@ def test_complete_data_end_newline_on_chunk_boundary(qm, tmp_path, monkeypatch): def test_count_newlines_matches_naive_across_ranges(qm, tmp_path, monkeypatch): """_count_newlines(start,end) == data[start:end].count(b'\\n') for arbitrary ranges, including across a small chunk size (multi-chunk streaming).""" - import context_intelligence_server.queue_manager as qm_mod + import context_intelligence_server.queue_manager.filesystem as qm_mod monkeypatch.setattr(qm_mod, "_SCAN_CHUNK_BYTES", 4) log = tmp_path / "queues" / "s1.log" @@ -652,7 +826,7 @@ def test_count_newlines_missing_and_empty_range(qm): def test_count_dead_matches_naive_and_streams(qm, tmp_path, monkeypatch): """_count_dead == old data.count(b'\\n') for empty / multi-record / missing, including a newline on a chunk boundary (streamed, not read_bytes).""" - import context_intelligence_server.queue_manager as qm_mod + import context_intelligence_server.queue_manager.filesystem as qm_mod monkeypatch.setattr(qm_mod, "_SCAN_CHUNK_BYTES", 8) dead = tmp_path / "queues" / "s1.dead.jsonl" @@ -677,7 +851,7 @@ async def test_recovery_seed_counts_unchanged_under_streaming(qm): await qm.append("s1", b"a") await qm.append("s1", b"bb") line1 = b"a\n" - await qm.commit("s1", len(line1)) # 1 written, 1 still pending + await qm.commit("s1", len(line1), None) # 1 written, 1 still pending accepted, written = await qm.recovery_seed_counts() @@ -779,9 +953,91 @@ async def test_spool_stats_healthy_offsets_report_zero_corrupt(qm): fire on the normal committed-offset path).""" await qm.append("s1", b"a") line = b"a\n" - await qm.commit("s1", len(line)) # writes a valid numeric .offset + await qm.commit("s1", len(line), None) # writes a valid numeric .offset qm._spool_cache = None stats = await qm.spool_stats() assert stats["corrupt_offsets"] == 0 + + +# --------------------------------------------------------------------------- +# Durable cursor folded into the atomic offset write +# --------------------------------------------------------------------------- + + +async def test_commit_requires_cursor_argument(qm): + """cursor has no default: omitting it fails loudly at call time. + + A default would silently null the cursor on a missed migration site or a + rolling deploy against an older signature -- the exact silent-loss class. + """ + with pytest.raises(TypeError): + await qm.commit("s1", 0) # type: ignore[call-arg] + + +async def test_commit_rejects_wrong_typed_cursor(qm): + """A non-dict, non-None cursor fails loud at write time. + + Without this, a wrong-typed cursor is written verbatim and silently reads + back as None -- the same silent cross-handler-counter reset the required + arg exists to prevent. + """ + await qm.append("s1", b"a") + with pytest.raises(TypeError): + await qm.commit("s1", 2, "not-a-dict") # type: ignore[arg-type] + + +async def test_commit_persists_cursor_in_same_record_as_offset(qm): + cursor = {"dl2": {"iteration_count": 4}, "dl3": {}} + await qm.append("s1", b"a") + await qm.commit("s1", 2, cursor) + assert await qm.read_cursor("s1") == cursor + assert qm._read_committed_offset("s1") == 2 # offset + cursor never skew + + +async def test_read_cursor_is_none_for_bare_int_and_missing(qm): + assert await qm.read_cursor("never_seen") is None + qm._offset_path("s1").write_text("42", encoding="utf-8") # legacy bare int + assert await qm.read_cursor("s1") is None + assert qm._read_committed_offset("s1") == 42 # bare-int still parses + + +async def test_rolling_upgrade_bare_int_then_envelope(qm): + # An old worker wrote a bare-int offset; the new worker commits an envelope + # on top of the same session -- both are readable, the cursor now persists. + qm._offset_path("s1").write_text("10", encoding="utf-8") + assert await qm.read_cursor("s1") is None + await qm.commit("s1", 20, {"dl2": {"iteration_count": 9}, "dl3": {}}) + assert qm._read_committed_offset("s1") == 20 + assert await qm.read_cursor("s1") == {"dl2": {"iteration_count": 9}, "dl3": {}} + + +async def test_corrupt_offset_record_raises_not_silently_zero(qm): + # A malformed record must raise, never degrade to 0 (0 replays the whole + # log and manufactures duplicate nodes -- a worse, quieter failure). + qm._offset_path("s1").write_text('{"v":1,"offset":"NaN"}', encoding="utf-8") + with pytest.raises(ValueError): + qm._read_committed_offset("s1") + + +async def test_commit_is_atomic_across_a_mid_write_crash(qm, tmp_path, monkeypatch): + # Simulate os.replace failing mid-commit: the previously committed record + # must survive intact (no torn/partial offset file), and no .tmp leaks. + await qm.append("s1", b"a") + await qm.commit("s1", 2, {"dl2": {"iteration_count": 1}, "dl3": {}}) + + import context_intelligence_server.queue_manager.filesystem as qmmod + + def _boom(src, dst): + raise OSError("simulated crash during os.replace") + + monkeypatch.setattr(qmmod.os, "replace", _boom) + with pytest.raises(OSError): + await qm.commit("s1", 99, {"dl2": {"iteration_count": 2}, "dl3": {}}) + monkeypatch.undo() + + # The pre-crash record is intact; the crashed write left nothing behind. + assert qm._read_committed_offset("s1") == 2 + assert await qm.read_cursor("s1") == {"dl2": {"iteration_count": 1}, "dl3": {}} + assert list((tmp_path / "queues").glob("*.tmp")) == [] diff --git a/tests/test_registry.py b/tests/test_registry.py index c31efeb4..f707d973 100644 --- a/tests/test_registry.py +++ b/tests/test_registry.py @@ -15,9 +15,9 @@ import pytest import context_intelligence_server.registry as registry_module -from context_intelligence_server.blob_store import AsyncDiskBlobStore +from context_intelligence_server.blob_store import FileSystemBlobStore 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 # --------------------------------------------------------------------------- @@ -248,10 +247,10 @@ async def test_worker_has_workspace_attribute( async def test_worker_services_blob_store_is_async_disk_blob_store( self, registry: SessionRegistry ) -> None: - """worker.services.blob_store is an AsyncDiskBlobStore instance.""" + """worker.services.blob_store is an FileSystemBlobStore instance.""" worker = registry.get_or_create("session-1", "/workspace/test") - assert isinstance(worker.services.blob_store, AsyncDiskBlobStore) + assert isinstance(worker.services.blob_store, FileSystemBlobStore) @pytest.mark.asyncio async def test_worker_services_blob_store_root_matches_settings_blob_path( @@ -262,7 +261,7 @@ async def test_worker_services_blob_store_root_matches_settings_blob_path( settings = get_settings() blob_store = worker.services.blob_store - assert isinstance(blob_store, AsyncDiskBlobStore) + assert isinstance(blob_store, FileSystemBlobStore) assert blob_store._root == Path(settings.blob_path) @@ -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] @@ -1988,7 +1990,7 @@ def test_get_or_create_accepts_created_by_kwarg(self) -> None: with ( patch("context_intelligence_server.registry.Neo4jGraphStore") as MockStore, patch( - "context_intelligence_server.registry.AsyncDiskBlobStore" + "context_intelligence_server.registry.create_blob_store" ) as MockBlob, patch( "context_intelligence_server.registry.HookStateService" @@ -2019,7 +2021,7 @@ def test_get_or_create_default_created_by_is_none(self) -> None: with ( patch("context_intelligence_server.registry.Neo4jGraphStore") as MockStore, patch( - "context_intelligence_server.registry.AsyncDiskBlobStore" + "context_intelligence_server.registry.create_blob_store" ) as MockBlob, patch( "context_intelligence_server.registry.HookStateService" @@ -2045,7 +2047,7 @@ class TestSessionOwnershipInvariant: - log nothing at ERROR when the same (or None) created_by arrives; - log an ERROR and preserve the bound id when a different created_by arrives. - Mocking strategy: patch Neo4jGraphStore / AsyncDiskBlobStore / HookStateService + Mocking strategy: patch Neo4jGraphStore / FileSystemBlobStore / HookStateService exactly as TestGetOrCreateCreatedBy does. After the first get_or_create call (which stores a worker whose .services is the mock_svc MagicMock), we manually set mock_svc.graph.created_by = "alice" to simulate the bound state — the real @@ -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.create_blob_store" + ) 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.create_blob_store" + ) 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.create_blob_store" + ) 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_services.py b/tests/test_services.py index 7b37014b..601d19b5 100644 --- a/tests/test_services.py +++ b/tests/test_services.py @@ -710,3 +710,41 @@ def test_created_by_propagated_to_custom_graph_store(self) -> None: store = GraphState() svc = HookStateService(workspace="/ws", graph_store=store, created_by="carol") assert svc.graph.created_by == "carol" + + +class TestDurableCursor: + def test_snapshot_restore_round_trips_cursor_state(self) -> None: + src = HookStateService(workspace="/ws") + src.data_layer_2.iteration_count = 7 + src.data_layer_2.execution_start_ts = "2026-01-01T00:00:02+00:00" + src.data_layer_2.active_iteration_id = "sid::iteration::7" + src.data_layer_3.active_recipe_run_stack = ["run-1", "run-2"] + + snapshot = src.snapshot_cursor() + + # A fresh worker (empty in-process state) restores from the snapshot. + dst = HookStateService(workspace="/ws") + dst.restore_cursor(snapshot) + assert dst.data_layer_2.iteration_count == 7 + assert dst.data_layer_2.execution_start_ts == "2026-01-01T00:00:02+00:00" + assert dst.data_layer_2.active_iteration_id == "sid::iteration::7" + assert dst.data_layer_3.active_recipe_run_stack == ["run-1", "run-2"] + + def test_restore_none_is_a_noop(self) -> None: + svc = HookStateService(workspace="/ws") + svc.data_layer_2.iteration_count = 3 + svc.restore_cursor(None) # legacy .offset with no cursor + assert svc.data_layer_2.iteration_count == 3 + + def test_restore_drops_unknown_keys_and_keeps_defaults(self) -> None: + svc = HookStateService(workspace="/ws") + svc.restore_cursor({"dl2": {"iteration_count": 5, "gone_field": "x"}}) + assert svc.data_layer_2.iteration_count == 5 + assert not hasattr(svc.data_layer_2, "gone_field") + # A field absent from the record keeps its dataclass default. + assert svc.data_layer_2.execution_start_ts is None + + def test_corrupt_record_never_raises(self) -> None: + svc = HookStateService(workspace="/ws") + svc.restore_cursor({"dl2": "not-a-dict", "dl3": 123}) # type: ignore[dict-item] + assert svc.data_layer_2.iteration_count == 0 diff --git a/tests/test_status.py b/tests/test_status.py index 572267fc..6e102792 100644 --- a/tests/test_status.py +++ b/tests/test_status.py @@ -491,3 +491,31 @@ def test_last_successful_flush_present_in_session(self) -> None: sess = response["sessions"][0] assert "last_successful_flush" in sess assert sess["last_successful_flush"] == expected_flush + + +class TestBootStateDegradedReason: + def test_degraded_reason_defaults_none_and_surfaces_on_snapshot(self) -> None: + from context_intelligence_server.status import BootState + + bs = BootState() + snap = bs.snapshot() + assert "degraded_reason" in snap + assert snap["degraded_reason"] is None + + def test_degrade_and_clear_round_trip(self) -> None: + from context_intelligence_server.status import BootState + + bs = BootState() + bs.degrade("3 node(s) lacking the :Node label") + assert bs.snapshot()["degraded_reason"] == "3 node(s) lacking the :Node label" + bs.clear_degraded() + assert bs.snapshot()["degraded_reason"] is None + + def test_no_parallel_schema_health_subsystem(self) -> None: + # degraded_reason lives on the existing BootState; no separate + # schema_health/schema_untagged fields were introduced. + from context_intelligence_server.status import BootState + + fields = set(BootState().snapshot()) + assert "schema_health" not in fields + assert "schema_untagged_nodes" not in fields diff --git a/tests/test_steady_state_reclaim.py b/tests/test_steady_state_reclaim.py new file mode 100644 index 00000000..5bc14029 --- /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.filesystem as queue_manager_module +from context_intelligence_server.config import Settings +from context_intelligence_server.queue_manager import FileSystemQueueManager, 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 = FileSystemQueueManager(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, None) + 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 = FileSystemQueueManager(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, None) + 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 = FileSystemQueueManager(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, None) + 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 = FileSystemQueueManager(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, None) + + log_path = tmp_path / f"{sid}.log" + offset_path = tmp_path / f"{sid}.offset" + log_before = log_path.read_bytes() + offset_before = offset_path.read_text(encoding="utf-8") + assert offset_before.strip() == f'{{"v":1,"offset":{batch.end_offset},"cursor":null}}' + + def _raise(fd: int, data: bytes) -> None: + raise OSError("simulated mid-copy failure") + + monkeypatch.setattr( + queue_manager_module.FileSystemQueueManager, "_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 = FileSystemQueueManager(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, None) # committed = 3 events (30 bytes) + + # Simulate the crash: offset already rebased to 0 (step 5 completed), + # log NOT yet replaced (step 6 never ran). + 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 = FileSystemQueueManager(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, None) # 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 = FileSystemQueueManager(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, None) + c = batch.end_offset + + log_path = tmp_path / f"{sid}.log" + offset_path = tmp_path / f"{sid}.offset" + log_before = log_path.read_bytes() + + real_replace = os.replace + + 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 = FileSystemQueueManager(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, None) + + log_path = tmp_path / f"{sid}.log" + + def _always_raise(src: Any, dst: Any) -> None: + # let the first offset rebase-to-0 write through, then fail both the + # log replace and the subsequent restore-to-C write + 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 = FileSystemQueueManager(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, None) + 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 = FileSystemQueueManager(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 = FileSystemQueueManager(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 = FileSystemQueueManager(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 = FileSystemQueueManager(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 = FileSystemQueueManager(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", FileSystemQueueManager(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", FileSystemQueueManager(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", FileSystemQueueManager(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, None) + + 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", FileSystemQueueManager(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 = FileSystemQueueManager(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_storage_boundary_guard.py b/tests/test_storage_boundary_guard.py new file mode 100644 index 00000000..1cc0b9d7 --- /dev/null +++ b/tests/test_storage_boundary_guard.py @@ -0,0 +1,125 @@ +"""Best-effort AST tripwire for the storage-agnosticism boundary. + +Storage artifacts (blobs, durable queues, identity stores) are reached ONLY +through their backend Protocols. No module OUTSIDE the three backend packages +may enumerate, stat, or unlink a storage artifact, nor read a storage root +path from settings -- otherwise a second backend (e.g. Azure) could not be +dropped in without editing consumers. + +This guard walks the AST of every non-storage module and fails on the file +operations and settings reads that would reach around a Protocol. It is a +TRIPWIRE, not a proof: a determined caller can defeat any static check +(dynamic import, getattr, os.system, a C-extension). The real guarantee is +that each consumer is positively verified protocol-only by reading. This test +exists to catch the accidental reintroduction of a KNOWN leak shape, and to +fail loudly the moment one lands. +""" + +from __future__ import annotations + +import ast +from pathlib import Path + +# The three storage backend packages -- the ONLY place a raw file operation on +# a storage artifact is allowed to live. +_STORAGE_PACKAGES = {"blob_store", "queue_manager", "identity_store", "lease_store"} + +# Attribute calls that mutate or enumerate a storage artifact on disk. +_BANNED_CALLS = { + ("os", "unlink"), + ("os", "remove"), + ("os", "removedirs"), + ("os", "rmdir"), + ("os", "scandir"), + ("os", "listdir"), + ("os", "walk"), +} +# Attribute names that are banned regardless of the receiver (glob on any Path, +# any shutil operation, Path.unlink()). +_BANNED_METHODS = {"glob", "rglob", "iterdir"} +_BANNED_MODULE_PREFIXES = {"shutil"} + +# settings.*_path reads that leak a storage root into a consumer. The factory +# and config own these; nobody else reads them. +_BANNED_SETTINGS_ATTRS = {"blob_path", "queues_path"} + +# The single human-approved exception (workspace AGENTS.md): the WriterLease +# boot detector resolves the queue directory WITHOUT constructing a +# QueueManager, so registry.queues_dir_path reads settings.queues_path. +_APPROVED_EXCEPTIONS = { + ("context_intelligence_server/registry.py", "queues_path"), +} + +_SERVER_ROOT = Path(__file__).resolve().parent.parent / "context_intelligence_server" +_SCRIPTS_ROOT = Path(__file__).resolve().parent.parent / "scripts" + + +def _iter_guarded_files() -> list[Path]: + files: list[Path] = [] + for root in (_SERVER_ROOT, _SCRIPTS_ROOT): + if not root.exists(): + continue + for path in root.rglob("*.py"): + parts = set(path.relative_to(root.parent).parts) + if parts & _STORAGE_PACKAGES: + continue # backend packages are the sanctioned home + files.append(path) + return files + + +def _rel(path: Path) -> str: + try: + return str(path.relative_to(_SERVER_ROOT.parent)) + except ValueError: + # A path outside the repo (e.g. the planted-leak self-test's tmp file). + return str(path) + + +def _violations(path: Path) -> list[str]: + tree = ast.parse(path.read_text("utf-8"), filename=str(path)) + rel = _rel(path) + found: list[str] = [] + + for node in ast.walk(tree): + # os.unlink(...) / os.walk(...) / shutil.rmtree(...) etc. + if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute): + attr = node.func.attr + recv = node.func.value + if isinstance(recv, ast.Name): + if (recv.id, attr) in _BANNED_CALLS: + found.append(f"{rel}:{node.lineno} {recv.id}.{attr}(...)") + if recv.id in _BANNED_MODULE_PREFIXES: + found.append(f"{rel}:{node.lineno} {recv.id}.{attr}(...)") + if attr in _BANNED_METHODS: + found.append(f"{rel}:{node.lineno} .{attr}(...)") + if attr == "unlink": # Path(...).unlink() + found.append(f"{rel}:{node.lineno} .unlink(...)") + + # settings.blob_path / settings.queues_path reads + if isinstance(node, ast.Attribute) and node.attr in _BANNED_SETTINGS_ATTRS: + if (rel, node.attr) in _APPROVED_EXCEPTIONS: + continue + found.append(f"{rel}:{node.lineno} settings.{node.attr}") + + return found + + +def test_no_storage_file_ops_outside_backend_packages() -> None: + """No consumer reaches around a storage Protocol with a raw file op.""" + violations: list[str] = [] + for path in _iter_guarded_files(): + violations.extend(_violations(path)) + + assert not violations, ( + "storage-agnosticism boundary breached -- storage artifacts must be " + "reached only through their backend Protocol. Offending sites:\n " + + "\n ".join(sorted(violations)) + ) + + +def test_guard_actually_detects_a_planted_leak(tmp_path: Path) -> None: + """The tripwire is armed: a planted os.unlink is caught (red-on-violation).""" + leak = tmp_path / "context_intelligence_server" / "routers" / "leaky.py" + leak.parent.mkdir(parents=True) + leak.write_text("import os\n\n\ndef f(p):\n os.unlink(p)\n", "utf-8") + assert _violations(leak), "guard failed to detect a planted os.unlink leak" 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..83da004b --- /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 FileSystemQueueManager, QueueManager +from context_intelligence_server.status import boot_state +from context_intelligence_server.lease_store.filesystem import FileSystemLeaseStore +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(FileSystemQueueManager, "__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 = FileSystemQueueManager.__init__ + + def _counting_init(self: FileSystemQueueManager, queues_dir: Path) -> None: + nonlocal construct_count + construct_count += 1 + real_init(self, queues_dir) + + monkeypatch.setattr(FileSystemQueueManager, "__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(_self: object) -> None: + raise OSError(errno.ESTALE, "stale file handle") + + monkeypatch.setattr(FileSystemLeaseStore, "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(_self: object) -> None: + blocker.wait(timeout=5.0) + + with patch.object(FileSystemLeaseStore, "read", _hang): + 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 = FileSystemQueueManager(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), + ): + assert lease._store is not None + await lease._io(lease._store.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 = tmp_path / ".writer.lease" + 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 = FileSystemQueueManager(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..ee1380f4 100644 --- a/uv.lock +++ b/uv.lock @@ -233,7 +233,7 @@ wheels = [ [[package]] name = "context-intelligence-server" -version = "6.7.0" +version = "6.7.3" source = { editable = "." } dependencies = [ { name = "aiofiles" },