Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
293 changes: 182 additions & 111 deletions context_intelligence_server/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,6 @@
)
from context_intelligence_server.neo4j_store import (
build_bounded_neo4j_driver,
count_untagged_nodes,
ensure_neo4j_schema,
mark_schema_ready,
)
Expand Down Expand Up @@ -245,105 +244,41 @@ async def _crash_recovery_sweep_loop(interval: int, respawn_limit: int) -> None:
logger.warning("crash_recovery_sweep: tick failed, will retry: %s", exc)


@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
"""Manage application lifespan: configure logging and create shared Neo4j driver."""
setup_logging()
_admin = _settings.resolve_neo4j_admin()
_query = _settings.resolve_neo4j_query()
logger.info(
"lifespan_startup: creating Neo4j drivers admin_url=%s query_url=%s query_access_mode=%s",
_admin.url,
_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.
app.state.neo4j_driver = build_neo4j_driver(_admin)
# 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,
)
# 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.
schema_fully_established = await ensure_neo4j_schema(
app.state.neo4j_driver, fail_on_data_conflict=True
)
# Seed the PROCESS-wide schema latch, but ONLY on a fully-established pass.
#
# fail_on_data_conflict=True makes this call fail closed on a :Node
# constraint DATA conflict -- but a CONNECTIVITY failure on any individual
# index/constraint is deliberately swallowed and reported through the
# return value instead (see ensure_neo4j_schema's docstring). Latching
# unconditionally would therefore mark a HALF-BUILT schema as ready and
# permanently disable the per-flush self-heal for the whole process --
# exactly the "constraint created once, never retried" gap
# Neo4jGraphStore._ensure_schema exists to close.
#
# On the happy path this seed is what stops every per-session store from
# re-running the same ~11-statement catalog pass on its first flush -- and,
# whenever that pass cannot complete, on EVERY subsequent flush -- competing
# for the very bolt pool it needs. See neo4j_store._SCHEMA_READY.
if schema_fully_established:
mark_schema_ready()
logger.info("lifespan_startup: Neo4j schema initialized")
else:
logger.warning(
"lifespan_startup: Neo4j schema NOT fully established (indexes or "
"constraints missing); leaving the process-wide latch unset so the "
"flush path retries schema init (rate-limited by "
"neo4j_store._SCHEMA_RETRY_BACKOFF_SECONDS)."
)
# 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.
async def _startup_recovery(app: FastAPI) -> None:
"""Crash-recovery pass + deferred-backlog sweep, off the startup path.

Runs as a background task created by ``lifespan`` so that startup can
complete -- and the HTTP surface open -- in ~1s regardless of how large
the durable spool is. Sets ``app.state.recovery_complete`` when the
one-shot recovery pass is done, so ``/status`` can say so honestly
rather than reporting un-seeded counters as though they were settled.

Thin wrapper around ``_startup_recovery_body``: as a BACKGROUND task, an
unhandled exception here would otherwise be swallowed until garbage
collection ("Task exception was never retrieved") -- silently leaving the
counters unseeded and no drainers respawned, with a server that looks
perfectly healthy. On the old inline path the same failure aborted boot
loudly. Log it, record it for ``/status``, and always set the event so
``/status`` cannot report "recovery in progress" forever.
"""
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"
await _startup_recovery_body(app)
except asyncio.CancelledError:
raise # shutdown -- not a failure, and must propagate
except Exception as exc: # noqa: BLE001 - background task boundary
app.state.recovery_error = repr(exc)
logger.exception(
"startup_recovery: FAILED -- conservation counters may be unseeded "
"and recovered drainers may not have respawned. The server is still "
"serving; a later boot's recover() reports the same sessions again, "
"and a new event for a session spawns its drainer via get_or_create()."
)
finally:
app.state.recovery_complete.set()


async def _startup_recovery_body(app: FastAPI) -> None:
"""The actual recovery pass. See ``_startup_recovery`` for why it is split."""
# Crash recovery: on startup, respawn one drainer per
# session that still has an undrained, complete line. The workspace is
# parsed from that session's FIRST log line so the respawned worker is
Expand Down Expand Up @@ -405,7 +340,7 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
# silent, un-discoverable fact. Names the exact counts and the
# setting to raise.
logger.warning(
"lifespan_startup: crash-recovery respawn cap reached "
"startup_recovery: 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 "
Expand All @@ -416,36 +351,163 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
deferred_count,
)
logger.info(
"lifespan_startup: crash recovery respawned %d/%d drainers",
"startup_recovery: crash recovery respawned %d/%d drainers",
respawned,
len(recovered),
)
# The one-shot pass is done: counters are seeded and this boot's drainers
# are up. /status can now report settled numbers rather than mid-recovery
# ones. Set BEFORE the sweep loop below, which never returns. (The wrapper
# also sets it in a finally, so a failure cannot pin /status at
# "in progress" forever -- but on the success path it must be set HERE, or
# the sweep loop would delay it indefinitely.)
app.state.recovery_complete.set()
# 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
# (respawn_limit is None) there is no deferred tail, so the loop is not
# entered -- 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
#
# AWAITED rather than wrapped in its own create_task: this coroutine is
# already a background task that lifespan cancels on shutdown, so awaiting
# the sweep here means one task to cancel instead of two -- and no way for
# the sweep to outlive its parent.
_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,
)
await _crash_recovery_sweep_loop(_sweep_interval, respawn_limit)


@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
"""Manage application lifespan: configure logging and create shared Neo4j driver."""
setup_logging()
_admin = _settings.resolve_neo4j_admin()
_query = _settings.resolve_neo4j_query()
logger.info(
"lifespan_startup: creating Neo4j drivers admin_url=%s query_url=%s query_access_mode=%s",
_admin.url,
_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.
app.state.neo4j_driver = build_neo4j_driver(_admin)
# 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,
)
# 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.
schema_fully_established = await ensure_neo4j_schema(
app.state.neo4j_driver, fail_on_data_conflict=True
)
# Seed the PROCESS-wide schema latch, but ONLY on a fully-established pass.
#
# fail_on_data_conflict=True makes this call fail closed on a :Node
# constraint DATA conflict -- but a CONNECTIVITY failure on any individual
# index/constraint is deliberately swallowed and reported through the
# return value instead (see ensure_neo4j_schema's docstring). Latching
# unconditionally would therefore mark a HALF-BUILT schema as ready and
# permanently disable the per-flush self-heal for the whole process --
# exactly the "constraint created once, never retried" gap
# Neo4jGraphStore._ensure_schema exists to close.
#
# On the happy path this seed is what stops every per-session store from
# re-running the same ~11-statement catalog pass on its first flush -- and,
# whenever that pass cannot complete, on EVERY subsequent flush -- competing
# for the very bolt pool it needs. See neo4j_store._SCHEMA_READY.
if schema_fully_established:
mark_schema_ready()
logger.info("lifespan_startup: Neo4j schema initialized")
else:
logger.warning(
"lifespan_startup: Neo4j schema NOT fully established (indexes or "
"constraints missing); leaving the process-wide latch unset so the "
"flush path retries schema init (rate-limited by "
"neo4j_store._SCHEMA_RETRY_BACKOFF_SECONDS)."
)
# NOTE: the O(1) untagged-:Node boot guard that used to sit here has been
# REMOVED. It refused to boot when any node lacked the universal :Node
# label. Two things made it dead weight:
#
# 1. Nothing in this service can produce an untagged node. Every
# node-creating statement in neo4j_store is label-scoped -- the node
# MERGE, both edge-endpoint MERGEs, and the session write path all
# MERGE (n:Node {node_id, workspace}).
# 2. Since get_node()/get_edge() became :Node-scoped, an untagged node is
# also INVISIBLE to every read. Read and write paths now agree, so a
# stray untagged node is inert dead data, not a correctness hazard --
# and refusing to serve over inert data is disproportionate.
#
# The capability is not gone: `context-intelligence-server doctor` still
# reports untagged nodes via diagnose()/count_untagged_nodes, and
# `doctor --fix` still repairs them. Detection moved to the operator path,
# where it belongs; it is no longer a boot gate. (Historically this check
# was itself an AllNodesScan and caused a 25-30s boot stall until PR #67
# made it O(1) -- a second reason not to keep it on the startup path.)

# Crash recovery runs as a BACKGROUND task -- it MUST NOT gate serving.
#
# uvicorn runs the ASGI lifespan to completion BEFORE it handles a single
# request (uvicorn/server.py: `await self.startup()` then
# `await self.main_loop()`), and gunicorn binds the socket before that. So
# any work done here is work during which the process ACCEPTS connections
# and answers none -- /version and /status included, however cheap they
# are. Incident 2026-09-09: a ~5000-file spool on Azure Files kept this
# loop busy for 9+ minutes; ACA's tcpSocket probes saw an open port and
# reported the replica Healthy while every request timed out at the APIM
# gateway (231s), and clients dropped events from full buffers.
#
# Moving it behind create_task lets startup complete in ~1s, so the HTTP
# surface is up regardless of spool size. This is safe because respawn is
# idempotent by construction -- see _crash_recovery_topup, which is already
# documented as safe to call repeatedly on a LIVE server, and which the
# sweep loop has always called against a serving process.
app.state.recovery_complete = asyncio.Event()
app.state.recovery_error = None
_recovery_task: asyncio.Task[None] = asyncio.create_task(_startup_recovery(app))

try:
yield
finally:
if _sweep_task is not None:
_sweep_task.cancel()
with suppress(asyncio.CancelledError):
await _sweep_task
# Cancel background recovery FIRST: it spawns drainers, and shutdown
# below quiesces them. Letting it keep spawning into a shutting-down
# registry would race the quiesce and dead-letter healthy events.
_recovery_task.cancel()
with suppress(asyncio.CancelledError):
await _recovery_task
# ORDER IS LOAD-BEARING. Quiesce the drainers FIRST. Every session's
# graph store now shares ONE driver, so closing it under a live drainer
# is no longer a per-session concern: the drainer's batch fails, it
Expand Down Expand Up @@ -929,6 +991,15 @@ async def get_status(request: Request) -> dict[str, Any]:
# above. Cheap by construction (stat-only, short-TTL cached); see
# QueueManager.spool_stats() for why this holds even under a huge spool.
response["spool"] = await registry.queue_manager.spool_stats()
# Crash recovery now runs in the BACKGROUND so the HTTP surface opens
# immediately (see _startup_recovery). While it is still running, the
# conservation counters above have not been seeded yet, so say so rather
# than letting a caller read mid-recovery numbers as settled ones.
_recovery_evt = getattr(request.app.state, "recovery_complete", None)
response["recovery_complete"] = bool(_recovery_evt and _recovery_evt.is_set())
# "complete" alone would be a lie by omission when recovery FAILED: the
# event is set either way (see _startup_recovery's finally).
response["recovery_error"] = getattr(request.app.state, "recovery_error", None)
# 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
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ build-backend = "hatchling.build"

[project]
name = "context-intelligence-server"
version = "6.7.2"
version = "6.7.3"
description = "Context Intelligence Server for Amplifier"
requires-python = ">=3.11"
dependencies = [
Expand Down
Loading
Loading