From 120ea7e9b523ba5436b0f100c2ab50da6dba4a26 Mon Sep 17 00:00:00 2001 From: colombod Date: Wed, 26 Aug 2026 18:26:17 +0000 Subject: [PATCH 1/4] fix(neo4j): reuse one bounded driver across per-session graph stores MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The drain path built a new AsyncGraphDatabase driver for every session_id with no pool bound, so bolt connections accumulated without limit until the server's bolt thread pool starved and ingest backpressured to clients. SessionRegistry now builds one shared, pool-bounded driver (lazily, on first session) and hands it to every Neo4jGraphStore it constructs. Neo4jGraphStore accepts an optional pre-built driver and tracks whether it owns it; close() only closes a driver it owns, so a per-session finalize can never take down the driver other live sessions are still using. The shared driver is closed exactly once, at lifespan shutdown. The pool-bounding kwargs (max_connection_pool_size, max_connection_lifetime) live in one helper in neo4j_store.py so the lifespan admin driver, the query driver, the doctor CLI, and the registry's shared driver can never diverge. Config gains neo4j_max_connection_pool_size (default 50) and neo4j_max_connection_lifetime (default 3600s). A live-Neo4j test proves the pool stays bounded across 30 sessions and returns to zero after close. 🤖 Generated with [Amplifier](https://github.com/microsoft/amplifier) Co-Authored-By: Amplifier <240397093+microsoft-amplifier@users.noreply.github.com> --- context_intelligence_server/config.py | 26 ++ context_intelligence_server/main.py | 29 ++- context_intelligence_server/neo4j_store.py | 93 +++++-- context_intelligence_server/registry.py | 45 +++- tests/conftest.py | 2 + tests/neo4j/test_driver_leak_bounded.py | 146 +++++++++++ .../neo4j/test_shared_driver_adverse_state.py | 59 +++++ tests/test_main.py | 30 +-- tests/test_neo4j_driver_sharing.py | 229 ++++++++++++++++++ 9 files changed, 617 insertions(+), 42 deletions(-) create mode 100644 tests/neo4j/test_driver_leak_bounded.py create mode 100644 tests/neo4j/test_shared_driver_adverse_state.py create mode 100644 tests/test_neo4j_driver_sharing.py diff --git a/context_intelligence_server/config.py b/context_intelligence_server/config.py index 48b20e97..095a8f8b 100644 --- a/context_intelligence_server/config.py +++ b/context_intelligence_server/config.py @@ -797,6 +797,32 @@ def resolve_neo4j_query(self) -> Neo4jClientConfig: access_mode="READ", ) + # Upper bound on concurrent bolt connections for a driver shared across many + # logical callers (the lifespan admin driver, the registry's per-session + # driver). Well under the server's default bolt thread-pool size so a + # driver leak can no longer starve it. + neo4j_max_connection_pool_size: int = 50 + + # Recycles a pooled connection after this many seconds, so a long-idle + # connection cannot accumulate indefinitely on the server side. + neo4j_max_connection_lifetime: float = 3600.0 + + @field_validator("neo4j_max_connection_pool_size") + @classmethod + def _validate_neo4j_max_connection_pool_size(cls, v: int) -> int: + """Fail loud on a non-positive pool size.""" + if v <= 0: + raise ValueError(f"neo4j_max_connection_pool_size must be > 0, got {v}") + return v + + @field_validator("neo4j_max_connection_lifetime") + @classmethod + def _validate_neo4j_max_connection_lifetime(cls, v: float) -> float: + """Fail loud on a non-positive lifetime (must be finite so idle connections recycle).""" + if v <= 0: + raise ValueError(f"neo4j_max_connection_lifetime must be > 0, got {v}") + return v + # ------------------------------------------------------------------------- # Storage paths # ------------------------------------------------------------------------- diff --git a/context_intelligence_server/main.py b/context_intelligence_server/main.py index 0bf0709f..83b97a61 100644 --- a/context_intelligence_server/main.py +++ b/context_intelligence_server/main.py @@ -16,7 +16,7 @@ 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 ( @@ -41,6 +41,7 @@ EventResponse, ) from context_intelligence_server.neo4j_store import ( + build_bounded_neo4j_driver, count_untagged_nodes, ensure_neo4j_schema, ) @@ -61,13 +62,19 @@ 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, + ) # --------------------------------------------------------------------------- @@ -221,9 +228,12 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: # 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. - 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. @@ -371,6 +381,11 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: 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( diff --git a/context_intelligence_server/neo4j_store.py b/context_intelligence_server/neo4j_store.py index 52ddc7d9..08b7cc3e 100644 --- a/context_intelligence_server/neo4j_store.py +++ b/context_intelligence_server/neo4j_store.py @@ -22,8 +22,32 @@ from neo4j import unit_of_work as _unit_of_work from neo4j.exceptions import DriverError, Neo4jError +from context_intelligence_server.config import Neo4jClientConfig + _LOG = logging.getLogger(__name__) + +def build_bounded_neo4j_driver( + config: Neo4jClientConfig, + *, + max_connection_pool_size: int, + max_connection_lifetime: float, +) -> Any: + """Construct an AsyncGraphDatabase driver with a bounded connection pool. + + Single source of truth for the pool-bounding kwargs applied to any driver + meant to be shared across many logical callers (the lifespan admin driver, + the registry's per-session driver). Both ``main.build_neo4j_driver`` and + ``SessionRegistry`` call this so the two never diverge. + """ + return AsyncGraphDatabase.driver( + config.url, + auth=config.auth, + max_connection_pool_size=max_connection_pool_size, + max_connection_lifetime=max_connection_lifetime, + ) + + # --------------------------------------------------------------------------- # Cypher identifier validation # --------------------------------------------------------------------------- @@ -1179,12 +1203,15 @@ def __init__( flush_chunk_rows: int = 100, flush_chunk_bytes: int = 4_194_304, neo4j_lock_timeout: float | None = None, + driver: Any | None = None, ) -> None: - """Initialise the store and create the async Neo4j driver. + """Initialise the store, reusing or creating the async Neo4j driver. Args: uri: Bolt/neo4j URI, e.g. ``bolt://localhost:7687``. + Ignored when ``driver`` is provided. auth: ``(username, password)`` tuple, or ``None`` for no-auth. + Ignored when ``driver`` is provided. database: Target Neo4j database name (default: ``"neo4j"``). workspace: Workspace to scope writes to. ``None`` resolves to ``"default"`` via the ``workspace`` property. @@ -1195,18 +1222,29 @@ def __init__( so a blocked flush raises ``Neo4jError`` instead of parking forever. ``None`` disables the timeout (default: no per-transaction limit). - Also sets ``connection_acquisition_timeout`` on - the driver to the same value so pool-exhaustion - failures also surface quickly. + When the store builds its own driver (``driver`` + not provided), this also sets + ``connection_acquisition_timeout`` on it to the + same value so pool-exhaustion failures surface + quickly. + driver: A pre-built async driver to reuse instead of + constructing a new one. When provided, this store + does not own the driver's lifecycle: ``close()`` + flushes and no-ops on the driver itself, leaving + it open for other stores sharing it. """ - # Explicit auto-retry budget for transient errors (e.g. deadlocks) so the - # managed-transaction retry window is deliberate and reviewable rather than - # relying on the driver default implicitly. 30.0s is a working default; - # design Open Question #3 — verify driver 6.1.0 backoff constants before tuning. - driver_kwargs: dict[str, Any] = {"max_transaction_retry_time": 30.0} - if neo4j_lock_timeout is not None and neo4j_lock_timeout > 0: - driver_kwargs["connection_acquisition_timeout"] = neo4j_lock_timeout - self._driver = AsyncGraphDatabase.driver(uri, auth=auth, **driver_kwargs) + if driver is not None: + self._driver = driver + self._owns_driver = False + else: + # Explicit auto-retry budget for transient errors (e.g. deadlocks) so + # the managed-transaction retry window is deliberate and reviewable + # rather than relying on the driver default implicitly. + driver_kwargs: dict[str, Any] = {"max_transaction_retry_time": 30.0} + if neo4j_lock_timeout is not None and neo4j_lock_timeout > 0: + driver_kwargs["connection_acquisition_timeout"] = neo4j_lock_timeout + self._driver = AsyncGraphDatabase.driver(uri, auth=auth, **driver_kwargs) + self._owns_driver = True self._database = database self._workspace = workspace self._created_by: str | None = None @@ -1224,6 +1262,19 @@ def __init__( else None ) + # ------------------------------------------------------------------ + # owns_driver property + # ------------------------------------------------------------------ + + @property + def owns_driver(self) -> bool: + """True when this store built its own driver; False when injected. + + Governs ``close()``: a store that does not own its driver must never + close it, since other stores may still be using it. + """ + return self._owns_driver + # ------------------------------------------------------------------ # workspace property # ------------------------------------------------------------------ @@ -1633,10 +1684,15 @@ async def _ensure_schema(self) -> None: # once Neo4j is reachable / duplicates are cleared by the dedup pass). async def close(self) -> None: - """Flush pending writes, await any background task, and close the driver. + """Flush pending writes and close the driver, if this store owns it. Handles event-loop mismatch gracefully when closing the driver from a different loop context. Sets ``_closed`` on completion. + + When the driver was injected (``owns_driver`` is False), the driver is + left open: it is shared with other stores/callers and closing it here + would break them out from under their own in-flight work. The shared + driver's owner is responsible for closing it exactly once. """ # Final flush to persist remaining buffer contents try: @@ -1646,11 +1702,12 @@ async def close(self) -> None: "Final flush failed during close; buffered writes may be lost" ) - # Close the driver, ignoring event-loop mismatch errors - try: - await self._driver.close() - except RuntimeError: - pass + if self._owns_driver: + # Close the driver, ignoring event-loop mismatch errors + try: + await self._driver.close() + except RuntimeError: + pass self._closed = True diff --git a/context_intelligence_server/registry.py b/context_intelligence_server/registry.py index ad078fee..38917837 100644 --- a/context_intelligence_server/registry.py +++ b/context_intelligence_server/registry.py @@ -11,11 +11,14 @@ from context_intelligence_server.blob_store import AsyncDiskBlobStore from context_intelligence_server.config import get_settings -from context_intelligence_server.status import EventRecord, ring_buffer -from context_intelligence_server.neo4j_store import Neo4jGraphStore +from context_intelligence_server.neo4j_store import ( + Neo4jGraphStore, + build_bounded_neo4j_driver, +) from context_intelligence_server.pipeline import process_event, setup_handlers from context_intelligence_server.queue_manager import Batch, QueueManager from context_intelligence_server.services import HookStateService +from context_intelligence_server.status import EventRecord, ring_buffer logger = logging.getLogger("context_intelligence_server") @@ -71,6 +74,11 @@ def __init__(self) -> None: self._queue_manager: QueueManager | None = None self._write_semaphore: asyncio.Semaphore | None = None self._max_delivery_attempts: int = 0 + # 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 pipeline-conservation counters (D2): make silently-dropped # events observable via /status. accepted = events admitted to the # log; written = events persisted to Neo4j; replayed = events @@ -106,6 +114,38 @@ def queue_manager(self) -> QueueManager: assert self._queue_manager is not None return self._queue_manager + 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.""" @@ -559,6 +599,7 @@ def get_or_create( neo4j_store = Neo4jGraphStore( uri=_admin.url, auth=_admin.auth, + driver=self.neo4j_driver, flush_chunk_rows=settings.neo4j_flush_chunk_rows, flush_chunk_bytes=settings.neo4j_flush_chunk_bytes, neo4j_lock_timeout=settings.neo4j_lock_timeout, diff --git a/tests/conftest.py b/tests/conftest.py index ecfd94d9..aafb3a58 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -122,6 +122,8 @@ 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 # Neo4j two-client split (doc 12): SessionRegistry.get_or_create() calls # settings.resolve_neo4j_admin() directly, so this proxy (which stands 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_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/test_main.py b/tests/test_main.py index 302fa1bf..11018ffa 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -718,7 +718,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, ): @@ -780,7 +780,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( @@ -822,7 +822,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( @@ -873,7 +873,7 @@ async def test_lifespan_default_respawns_all_recovered_sessions_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()), @@ -906,7 +906,7 @@ 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()), @@ -952,7 +952,7 @@ 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()), @@ -993,7 +993,7 @@ 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()), @@ -1063,7 +1063,7 @@ 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()), @@ -1090,7 +1090,7 @@ 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()), @@ -1120,7 +1120,7 @@ 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()), @@ -1159,7 +1159,7 @@ async def test_lifespan_calls_ensure_schema_with_fail_on_data_conflict() -> None 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( @@ -1191,7 +1191,7 @@ async def test_lifespan_raises_on_ensure_schema_data_conflict() -> None: 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( @@ -1216,7 +1216,7 @@ async def test_lifespan_raises_on_untagged_nodes() -> None: 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( @@ -1245,7 +1245,7 @@ async def test_lifespan_does_not_raise_on_clean_graph() -> None: 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( @@ -1274,7 +1274,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( 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 From a3c59629697d848ef3c9e6aaeb9ddb257f1f7ae2 Mon Sep 17 00:00:00 2001 From: sadlilas <11658960+sadlilas@users.noreply.github.com> Date: Tue, 1 Sep 2026 06:33:11 -0700 Subject: [PATCH 2/4] fix(neo4j): Harden shared-driver shutdown and restore driver kwargs - Quiesce drain workers before closing driver in lifespan shutdown to prevent data loss - Restore connection_acquisition_timeout and max_transaction_retry_time to shared driver - Remove no-op neo4j_max_connection_lifetime setting Fixes data loss (dead-lettering of queued events) at shutdown. Adds regression tests for shutdown ordering, driver kwarg parity, and live-Neo4j shutdown behavior. Hardens PR #91. Generated with [Amplifier](https://github.com/microsoft/amplifier) Co-Authored-By: Amplifier <240397093+microsoft-amplifier@users.noreply.github.com> --- context_intelligence_server/config.py | 24 ++-- context_intelligence_server/main.py | 13 +- context_intelligence_server/neo4j_store.py | 44 ++++-- context_intelligence_server/registry.py | 36 ++++- tests/conftest.py | 1 - tests/neo4j/test_driver_leak_bounded.py | 1 - tests/neo4j/test_shutdown_no_deadletter.py | 102 ++++++++++++++ tests/test_main.py | 44 ++++++ tests/test_neo4j_driver_sharing.py | 151 ++++++++++++++++++++- 9 files changed, 381 insertions(+), 35 deletions(-) create mode 100644 tests/neo4j/test_shutdown_no_deadletter.py diff --git a/context_intelligence_server/config.py b/context_intelligence_server/config.py index 095a8f8b..fd72680c 100644 --- a/context_intelligence_server/config.py +++ b/context_intelligence_server/config.py @@ -798,15 +798,17 @@ def resolve_neo4j_query(self) -> Neo4jClientConfig: ) # 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. + # logical callers (the lifespan admin driver, the lifespan query driver, the + # registry's shared per-session driver). Well under the server's default + # bolt thread-pool size so a driver leak can no longer starve it. The neo4j + # driver's own default is 100; 50 is a deliberate reduction, since every + # session now shares one pool instead of holding a private one. + # + # No companion max_connection_lifetime knob: the driver already recycles + # pooled connections at 3600 s by default, so a setting whose default equals + # the library default would change nothing. 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: @@ -815,14 +817,6 @@ def _validate_neo4j_max_connection_pool_size(cls, v: int) -> int: 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 # ------------------------------------------------------------------------- diff --git a/context_intelligence_server/main.py b/context_intelligence_server/main.py index 83b97a61..6d8f4deb 100644 --- a/context_intelligence_server/main.py +++ b/context_intelligence_server/main.py @@ -73,7 +73,6 @@ def build_neo4j_driver(config: Neo4jClientConfig) -> Any: return build_bounded_neo4j_driver( config, max_connection_pool_size=_settings.neo4j_max_connection_pool_size, - max_connection_lifetime=_settings.neo4j_max_connection_lifetime, ) @@ -233,7 +232,6 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: 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. @@ -378,6 +376,17 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: _sweep_task.cancel() with suppress(asyncio.CancelledError): await _sweep_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 + # spends its max_delivery_attempts budget in ~250 ms, and + # _handle_exhausted_batch dead-letters each line AND commits the offset + # past it -- discarding healthy events that merely happened to be + # queued at shutdown, with no replay on the next boot. Cancelling first + # routes each drainer through CancelledError -> _safe_close -> a final + # flush while the driver is still open. + logger.info("lifespan_shutdown: quiescing drain workers") + await registry.shutdown_workers() logger.info("lifespan_shutdown: closing Neo4j drivers") await app.state.neo4j_driver.close() await app.state.neo4j_query_driver.close() diff --git a/context_intelligence_server/neo4j_store.py b/context_intelligence_server/neo4j_store.py index 08b7cc3e..821db6c5 100644 --- a/context_intelligence_server/neo4j_store.py +++ b/context_intelligence_server/neo4j_store.py @@ -31,21 +31,47 @@ def build_bounded_neo4j_driver( config: Neo4jClientConfig, *, max_connection_pool_size: int, - max_connection_lifetime: float, + connection_acquisition_timeout: float | None = None, ) -> 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. + the lifespan query driver, the registry's shared per-session driver). All + three construct through here so they can never diverge. + + Args: + max_connection_pool_size: + Hard cap on concurrent bolt connections for this driver. + connection_acquisition_timeout: + How long a caller waits for a free pooled connection before + failing. ``None`` (the lifespan drivers) leaves the driver default + in place, matching what those two did before they were routed + through this helper. ``SessionRegistry`` passes + ``settings.neo4j_lock_timeout`` so the shared per-session driver + keeps the acquisition budget the per-session drivers it replaced + carried -- and it matters more now, not less: one bounded pool is + shared by every session, so acquisition can genuinely queue. + + ``max_connection_lifetime`` is deliberately not set: the neo4j driver + already recycles pooled connections at 3600 s by default, so passing it + would be a knob that changes nothing. """ - return AsyncGraphDatabase.driver( - config.url, - auth=config.auth, - max_connection_pool_size=max_connection_pool_size, - max_connection_lifetime=max_connection_lifetime, - ) + kwargs: dict[str, Any] = { + "max_connection_pool_size": max_connection_pool_size, + # 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. Carried over + # verbatim from the per-session driver construction this helper + # subsumed. + "max_transaction_retry_time": 30.0, + } + if ( + connection_acquisition_timeout is not None + and connection_acquisition_timeout > 0 + ): + kwargs["connection_acquisition_timeout"] = connection_acquisition_timeout + return AsyncGraphDatabase.driver(config.url, auth=config.auth, **kwargs) # --------------------------------------------------------------------------- diff --git a/context_intelligence_server/registry.py b/context_intelligence_server/registry.py index 38917837..e7f00acf 100644 --- a/context_intelligence_server/registry.py +++ b/context_intelligence_server/registry.py @@ -127,7 +127,12 @@ def _ensure_neo4j_driver(self) -> Any: 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, + # Parity with the per-session driver this one replaces: a + # blocked acquisition must surface on the SAME budget as a + # blocked transaction. Load-bearing now in a way it was not + # before -- every session shares this one bounded pool, so + # acquisition can actually queue. + connection_acquisition_timeout=settings.neo4j_lock_timeout, ) return self._neo4j_driver @@ -137,9 +142,38 @@ def neo4j_driver(self) -> Any: Neo4jGraphStore -- never closed by a per-session finalize.""" return self._ensure_neo4j_driver() + async def shutdown_workers(self) -> None: + """Quiesce every drain worker BEFORE the shared driver is closed. + + Ordering invariant (must run before ``close_neo4j_driver``): a live + drainer that meets a closed shared driver fails its batch, spends its + ``max_delivery_attempts`` budget in ~250 ms (5 attempts x the 50 ms + ``_DRAIN_POLL_INTERVAL`` backoff), and falls into + ``_handle_exhausted_batch`` -- which dead-letters each line AND commits + the offset past it. Those are healthy events that merely happened to be + queued at shutdown, and once dead-lettered they never replay. + + Cancelling instead routes each drainer through its ``CancelledError`` + handler -> ``_safe_close(worker)`` -> a final flush on a driver that is + still open. Anything left uncommitted stays in the durable queue and + replays on the next boot, which is the pre-shared-driver behaviour. + + Exceptions are collected, not raised: shutdown must not be derailed by + one failing worker. + """ + tasks = [w.task for w in self._workers.values() if w.task is not None] + if not tasks: + return + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + async def close_neo4j_driver(self) -> None: """Close the shared driver exactly once, at process shutdown. + Call ``shutdown_workers()`` first -- see its docstring for why closing + this driver under a live drainer dead-letters good events. + No-op if the driver was never built (no session has run yet). """ if self._neo4j_driver is not None: diff --git a/tests/conftest.py b/tests/conftest.py index aafb3a58..2c3b7c9f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -123,7 +123,6 @@ class _SettingsProxy: 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 # Neo4j two-client split (doc 12): SessionRegistry.get_or_create() calls # settings.resolve_neo4j_admin() directly, so this proxy (which stands diff --git a/tests/neo4j/test_driver_leak_bounded.py b/tests/neo4j/test_driver_leak_bounded.py index 0487eaa5..132ac3e0 100644 --- a/tests/neo4j/test_driver_leak_bounded.py +++ b/tests/neo4j/test_driver_leak_bounded.py @@ -68,7 +68,6 @@ async def test_bolt_connections_stay_bounded_and_release( 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"), ) diff --git a/tests/neo4j/test_shutdown_no_deadletter.py b/tests/neo4j/test_shutdown_no_deadletter.py new file mode 100644 index 00000000..39f84097 --- /dev/null +++ b/tests/neo4j/test_shutdown_no_deadletter.py @@ -0,0 +1,102 @@ +"""Behavioral evidence that shutdown does not discard queued events. + +The shared Neo4j driver made driver lifetime a cross-session concern: closing +it while a drain worker is still running is no longer "that session's driver +going away", it is *every* session's driver going away mid-flight. + +A drainer that meets a closed driver fails its batch, spends its +``max_delivery_attempts`` budget in ~250 ms (5 attempts x the 50 ms +``_DRAIN_POLL_INTERVAL`` backoff), and lands in ``_handle_exhausted_batch`` -- +which dead-letters each line AND commits the offset past it. Those are healthy +events that merely happened to be queued when the process stopped, and once +dead-lettered they never replay. + +This drives real events through the real registry against a live Neo4j and +asserts the shutdown sequence used by ``lifespan`` (quiesce the drainers, then +close the shared driver) dead-letters nothing. +""" + +from __future__ import annotations + +import asyncio +import json +from typing import Any + +import pytest +from context_intelligence_server.config import Settings +from context_intelligence_server.registry import SessionRegistry + +pytestmark = pytest.mark.neo4j + +# Enough to guarantee undrained lines are still queued when shutdown starts -- +# _DRAIN_MAX_BATCH is 100, so this is several batches deep. +EVENT_COUNT = 400 + + +def _event_line(session_id: str, i: int) -> bytes: + return json.dumps( + { + "event": "tool:pre", + "workspace": "/ws", + "data": { + "session_id": session_id, + "timestamp": "2024-01-01T00:00:00+00:00", + "tool_name": f"tool-{i}", + }, + } + ).encode("utf-8") + + +@pytest.mark.asyncio +async def test_shutdown_quiesce_then_close_deadletters_nothing( + neo4j_container: dict[str, Any], + monkeypatch: pytest.MonkeyPatch, + tmp_path: Any, +) -> None: + settings = Settings( + neo4j_url=neo4j_container["bolt_url"], + neo4j_user=neo4j_container["user"], + neo4j_password=neo4j_container["password"], + blob_path=str(tmp_path / "blobs"), + queues_path=str(tmp_path / "queues"), + ) + monkeypatch.setattr( + "context_intelligence_server.registry.get_settings", lambda: settings + ) + + reg = SessionRegistry() + session_id = "shutdown-session" + qm = reg.queue_manager + + dead_lettered: list[str] = [] + real_dead_letter = qm.dead_letter + + async def spy_dead_letter(sid: str, raw: bytes, error: str) -> None: + dead_lettered.append(error) + await real_dead_letter(sid, raw, error) + + monkeypatch.setattr(qm, "dead_letter", spy_dead_letter) + + for i in range(EVENT_COUNT): + await qm.append(session_id, _event_line(session_id, i)) + + reg.get_or_create(session_id, "/ws") + # Let the drainer get into its loop with work still queued behind it. + await asyncio.sleep(0.4) + + # The lifespan shutdown sequence, in order. Reversing these two lines is the + # regression this test exists to catch. + await reg.shutdown_workers() + await reg.close_neo4j_driver() + + # Stay on the loop as the server would during its shutdown window: a + # still-live drainer would burn its retry budget and dead-letter here. + await asyncio.sleep(2.0) + + assert dead_lettered == [], ( + f"{len(dead_lettered)} healthy queued events were dead-lettered during " + "shutdown; the drain workers must be quiesced before the shared driver " + "closes. First error: " + f"{dead_lettered[0] if dead_lettered else ''}" + ) + assert reg._neo4j_driver is None diff --git a/tests/test_main.py b/tests/test_main.py index 11018ffa..470aa09f 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -741,6 +741,50 @@ async def test_lifespan_creates_and_closes_driver( assert mock_driver.close.await_count == 2 +async def test_lifespan_quiesces_drain_workers_before_closing_shared_driver( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Shutdown MUST cancel the drain workers before closing the registry's + shared Neo4j driver. + + Every session's graph store shares one driver now, so closing it under a + live drainer is not a per-session concern: the drainer's batch fails, it + burns its max_delivery_attempts budget in ~250ms, and + _handle_exhausted_batch dead-letters each line AND commits the offset past + it -- silently discarding healthy events that merely happened to be queued + at shutdown, with no replay on the next boot. + """ + mock_driver = MagicMock() + mock_driver.close = AsyncMock() + + calls: list[str] = [] + + async def record_shutdown_workers() -> None: + calls.append("shutdown_workers") + + async def record_close_driver() -> None: + calls.append("close_neo4j_driver") + + monkeypatch.setattr( + main_module.registry, "shutdown_workers", record_shutdown_workers + ) + monkeypatch.setattr(main_module.registry, "close_neo4j_driver", record_close_driver) + + with ( + patch("context_intelligence_server.main.setup_logging"), + patch( + "context_intelligence_server.neo4j_store.AsyncGraphDatabase.driver", + return_value=mock_driver, + ), + ): + async with lifespan(main_module.app): + assert calls == [] + + assert calls == ["shutdown_workers", "close_neo4j_driver"], ( + f"drain workers must be quiesced BEFORE the shared driver closes; got {calls}" + ) + + # --------------------------------------------------------------------------- # Lifespan crash-recovery + workers==1 guard tests (Phase B2) # --------------------------------------------------------------------------- diff --git a/tests/test_neo4j_driver_sharing.py b/tests/test_neo4j_driver_sharing.py index 47218d4c..1977f9a9 100644 --- a/tests/test_neo4j_driver_sharing.py +++ b/tests/test_neo4j_driver_sharing.py @@ -6,16 +6,23 @@ its own driver. - SessionRegistry hands the same driver instance to every per-session store instead of building one per session. +- The shared driver keeps the pool cap AND the acquisition/retry budgets the + per-session drivers it replaced carried. +- shutdown_workers() quiesces every drainer before the shared driver closes. """ from __future__ import annotations import asyncio -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest -from context_intelligence_server.neo4j_store import Neo4jGraphStore -from context_intelligence_server.registry import SessionRegistry +from context_intelligence_server.config import Neo4jClientConfig, get_settings +from context_intelligence_server.neo4j_store import ( + Neo4jGraphStore, + build_bounded_neo4j_driver, +) +from context_intelligence_server.registry import SessionRegistry, SessionWorker # --------------------------------------------------------------------------- # Injected-driver seam (owns_driver) @@ -158,8 +165,9 @@ async def test_n_sessions_build_exactly_one_driver(monkeypatch) -> None: @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.""" + """The single shared driver must be built WITH the bounded pool cap + (default 50) -- an unbounded build is the leak -- AND with the acquisition + budget the per-session drivers it replaced carried.""" reg = SessionRegistry() builds = _spy_driver_factory(reg, monkeypatch) @@ -168,11 +176,64 @@ async def test_shared_driver_built_with_bounded_kwargs(monkeypatch) -> None: assert len(builds) == 1 kwargs = builds[0]["kwargs"] assert kwargs["max_connection_pool_size"] == 50 - assert kwargs["max_connection_lifetime"] == 3600.0 + # Parity with neo4j_lock_timeout (default 30.0): without this the driver + # silently falls back to its own 60.0s default, doubling how long a caller + # parks on an exhausted pool -- and the pool is now SHARED, so exhaustion + # is reachable in a way it never was with a private pool per session. + assert kwargs["connection_acquisition_timeout"] == get_settings().neo4j_lock_timeout _cancel_workers(reg) +# --------------------------------------------------------------------------- +# Driver-kwarg parity with the per-session driver this helper replaced +# --------------------------------------------------------------------------- + + +def test_build_bounded_driver_preserves_per_session_driver_kwargs() -> None: + """build_bounded_neo4j_driver must carry over BOTH kwargs the per-session + Neo4jGraphStore driver set, not just the new pool cap. + + Regression guard: extracting the construction into a shared helper silently + dropped ``connection_acquisition_timeout`` (30.0 -> the driver's 60.0 + default) and the explicit ``max_transaction_retry_time``. + """ + config = Neo4jClientConfig(url="bolt://unused:7687", username="u", password="p") + + with patch( + "context_intelligence_server.neo4j_store.AsyncGraphDatabase" + ) as mock_adb: + build_bounded_neo4j_driver( + config, + max_connection_pool_size=50, + connection_acquisition_timeout=30.0, + ) + + kwargs = mock_adb.driver.call_args.kwargs + assert kwargs["max_connection_pool_size"] == 50 + assert kwargs["connection_acquisition_timeout"] == 30.0 + assert kwargs["max_transaction_retry_time"] == 30.0 + # Deliberately absent: the driver already defaults to 3600 s, so setting it + # would be a knob that changes nothing. + assert "max_connection_lifetime" not in kwargs + + +def test_build_bounded_driver_omits_acquisition_timeout_when_not_given() -> None: + """The lifespan admin/query drivers pass no acquisition timeout, so the + helper must leave the driver default in place for them -- matching exactly + what those two did before they were routed through this helper.""" + config = Neo4jClientConfig(url="bolt://unused:7687", username="u", password="p") + + with patch( + "context_intelligence_server.neo4j_store.AsyncGraphDatabase" + ) as mock_adb: + build_bounded_neo4j_driver(config, max_connection_pool_size=50) + + kwargs = mock_adb.driver.call_args.kwargs + assert "connection_acquisition_timeout" not in kwargs + assert kwargs["max_transaction_retry_time"] == 30.0 + + @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 @@ -227,3 +288,81 @@ async def test_close_neo4j_driver_none_safe_when_no_session_ran() -> None: # Must not raise. await reg.close_neo4j_driver() assert reg._neo4j_driver is None + + +# --------------------------------------------------------------------------- +# Shutdown quiesce (shutdown_workers) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_shutdown_workers_cancels_and_awaits_every_drainer() -> None: + """shutdown_workers() must cancel AND await every live drain worker. + + This is the ordering guard for the shared driver: a drainer still running + when the shared driver closes fails its batch, exhausts its retry budget, + and dead-letters healthy events (committing the offset past them). + """ + reg = SessionRegistry() + + started = asyncio.Event() + + async def never_ending() -> None: + started.set() + await asyncio.sleep(3600) + + workers = [] + for i in range(3): + worker = SessionWorker( + session_id=f"s{i}", + workspace=f"/ws/{i}", + services=MagicMock(), + ) + worker.task = asyncio.create_task(never_ending()) + reg._workers[worker.session_id] = worker + workers.append(worker) + + await started.wait() + + await reg.shutdown_workers() + + for worker in workers: + assert worker.task is not None + assert worker.task.done(), "shutdown_workers must AWAIT, not just cancel" + assert worker.task.cancelled() + + +@pytest.mark.asyncio +async def test_shutdown_workers_no_op_with_no_workers() -> None: + """shutdown_workers() must be safe when no session ever ran.""" + reg = SessionRegistry() + await reg.shutdown_workers() # must not raise + + +@pytest.mark.asyncio +async def test_shutdown_workers_survives_a_failing_drainer() -> None: + """One worker raising during teardown must not abort the shutdown of the + others -- shutdown is not derailable by a single bad drainer.""" + reg = SessionRegistry() + + async def raises_on_cancel() -> None: + try: + await asyncio.sleep(3600) + except asyncio.CancelledError: + raise RuntimeError("teardown blew up") from None + + async def clean() -> None: + await asyncio.sleep(3600) + + bad = SessionWorker(session_id="bad", workspace="/ws", services=MagicMock()) + bad.task = asyncio.create_task(raises_on_cancel()) + good = SessionWorker(session_id="good", workspace="/ws", services=MagicMock()) + good.task = asyncio.create_task(clean()) + reg._workers["bad"] = bad + reg._workers["good"] = good + await asyncio.sleep(0) + + await reg.shutdown_workers() # must not raise + + assert bad.task.done() + assert good.task.done() From 8f874aa150b99d04d5538b2dce817c5965b71ac3 Mon Sep 17 00:00:00 2001 From: sadlilas <11658960+sadlilas@users.noreply.github.com> Date: Tue, 1 Sep 2026 06:34:02 -0700 Subject: [PATCH 3/4] docs(tests): drop a dangling cross-repo issue reference Issues are disabled on this repository, so the bare "#489" in the shared-driver test docstring could only resolve somewhere else -- a dangling pointer for anyone reading the code. Replaced with a self-describing statement of the property under test. Generated with [Amplifier](https://github.com/microsoft/amplifier) Co-Authored-By: Amplifier <240397093+microsoft-amplifier@users.noreply.github.com> --- tests/test_neo4j_driver_sharing.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_neo4j_driver_sharing.py b/tests/test_neo4j_driver_sharing.py index 1977f9a9..d9c49e7b 100644 --- a/tests/test_neo4j_driver_sharing.py +++ b/tests/test_neo4j_driver_sharing.py @@ -45,7 +45,8 @@ def test_injected_driver_reports_owns_driver_false() -> None: @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).""" + out from under a second store still using it -- the safety property the + whole shared-driver change rests on.""" shared_driver = AsyncMock() store_a = Neo4jGraphStore(uri="bolt://unused:7687", driver=shared_driver) From dd2a29effa600e1af66bd8a8c0a65a4abeb3b96f Mon Sep 17 00:00:00 2001 From: sadlilas <11658960+sadlilas@users.noreply.github.com> Date: Tue, 1 Sep 2026 06:46:13 -0700 Subject: [PATCH 4/4] docs(architecture): correct the Neo4j client topology for the shared session driver The topology section said the server uses two drivers and that the admin driver "carries the ingest path -- the drainer's batch flushes". Neither is true now: the drainer's flushes go through the registry's shared session driver, and /status probes only the admin and cypher_query drivers, so neo4j_connected: true no longer implies a healthy ingest path. An operator debugging ingest backpressure -- the exact failure this change fixes -- would have been pointed at the wrong driver. Documents all three drivers, which one carries ingest, the shared driver's lazy build and bounded pool, and the shutdown ordering requirement. Generated with [Amplifier](https://github.com/microsoft/amplifier) Co-Authored-By: Amplifier <240397093+microsoft-amplifier@users.noreply.github.com> --- docs/architecture/README.md | 48 ++++++++++++++++++++++++++----------- 1 file changed, 34 insertions(+), 14 deletions(-) diff --git a/docs/architecture/README.md b/docs/architecture/README.md index 8c58426f..965a3ac2 100644 --- a/docs/architecture/README.md +++ b/docs/architecture/README.md @@ -259,22 +259,42 @@ and their typed relationships (`HAS_EVENT`, `EMITTED`, `REFERENCES_BLOB`). ## Neo4j client topology -The server connects to Neo4j through **two drivers** rather than one. The **admin -driver** (`neo4j_driver`) is opened in **WRITE** access mode and carries the ingest -path — the drainer's batch flushes and schema work. The **cypher_query driver** -(`neo4j_query_driver`) is opened in **READ** access mode and serves `POST /cypher` -reads. Both drivers are created together in the dual-driver lifespan -at startup and closed together at shutdown. With legacy flat config -(`neo4j_url` / `neo4j_user` / `neo4j_password`), both drivers fall back to the same -endpoint and shared credentials, differing only by access-mode hint; a structured -`neo4j:` block lets the read driver take a separate credential and/or URL (e.g. a -read replica). Their connection health is reported independently on `/status` as -`neo4j_connected` (admin) and `neo4j_query_connected` (cypher_query). +The server connects to Neo4j through **three drivers**, each with a distinct job. + +Two are owned by the lifespan. The **admin driver** (`neo4j_driver`) is opened in +**WRITE** access mode and handles boot-time and operational work: schema creation, +the untagged-node integrity check, and the `/status` connectivity probe. The +**cypher_query driver** (`neo4j_query_driver`) is opened in **READ** access mode and +serves `POST /cypher` reads. Both are created together at startup and closed +together at shutdown. With legacy flat config (`neo4j_url` / `neo4j_user` / +`neo4j_password`), both fall back to the same endpoint and shared credentials, +differing only by access-mode hint; a structured `neo4j:` block lets the read driver +take a separate credential and/or URL (e.g. a read replica). Their connection health +is reported independently on `/status` as `neo4j_connected` (admin) and +`neo4j_query_connected` (cypher_query). + +The third is the **shared session driver**, owned by `SessionRegistry`, and it is the +one that carries the **ingest path** — every per-session `Neo4jGraphStore`'s batch +flushes. It is built lazily on the first session (from the same resolved admin client +config) with a bounded pool (`neo4j_max_connection_pool_size`, default 50) and +injected into every store, so a per-session finalize can never close the driver other +live sessions are still using. Before it existed, each session built and held its own +unshared driver, which is how bolt connections accumulated until the server's thread +pool starved. It is closed exactly once, at shutdown, and only *after* +`SessionRegistry.shutdown_workers()` has quiesced the drain workers — closing it +under a live drainer makes that drainer exhaust its retry budget and dead-letter +healthy queued events. + +Because `/status` probes only the admin and cypher_query drivers, +`neo4j_connected: true` does **not** by itself mean the ingest path is healthy. +Ingest health is visible through the pipeline-conservation counters in the `metrics` +block (`written_total`, `residual`, `degraded`). > **Note:** the existing `.dot` diagrams (e.g. `05-durable-ingest-queue`) show only -> the **write path** through the admin driver. The two-driver split is not yet -> rendered in any diagram — a dedicated topology diagram is a known follow-up. This -> prose subsection is the interim reference; no new `.dot` is authored in this pass. +> the **write path**, and label it as the admin driver — that write path is now the +> shared session driver. The three-driver split is not yet rendered in any diagram — +> a dedicated topology diagram is a known follow-up. This prose subsection is the +> interim reference; no new `.dot` is authored in this pass. ---