diff --git a/context_intelligence_server/main.py b/context_intelligence_server/main.py index 6ef8965..c2faf78 100644 --- a/context_intelligence_server/main.py +++ b/context_intelligence_server/main.py @@ -44,6 +44,7 @@ build_bounded_neo4j_driver, count_untagged_nodes, ensure_neo4j_schema, + mark_schema_ready, ) from context_intelligence_server.registry import SessionRegistry from context_intelligence_server.routers.admin import router as admin_router @@ -289,8 +290,34 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: # 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") + 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 -- diff --git a/context_intelligence_server/neo4j_store.py b/context_intelligence_server/neo4j_store.py index fbfd159..bb72149 100644 --- a/context_intelligence_server/neo4j_store.py +++ b/context_intelligence_server/neo4j_store.py @@ -14,6 +14,7 @@ import json import logging import re +import time from collections.abc import Generator from datetime import datetime from typing import Any, LiteralString, cast @@ -188,6 +189,55 @@ def build_bounded_neo4j_driver( f"MATCH (n:{_UNIVERSAL_NODE_LABEL} {{node_id: $node_id, workspace: $workspace}})" ) +# Single-node read-by-identity, built on the shared _NODE_MATCH_BY_ID prefix so +# it can never drift back to a label-free MATCH. This is the get_node() Neo4j +# fallback, which runs on the hot per-event path (touch_session -> get_node), so +# it MUST plan as a NodeUniqueIndexSeek: as a label-free MATCH it planned as an +# AllNodesScan over the whole graph and, at 12.09M nodes, took ~60s per call +# while holding its pooled connection for the duration. +_NODE_GET_BY_ID_CYPHER = ( + f"{_NODE_MATCH_BY_ID} RETURN properties(n) AS props, labels(n) AS lbls" +) + +# Single-edge read-by-endpoints, anchored on BOTH :Node endpoints so the +# planner seeks the :Node(node_id, workspace) unique index twice and then walks +# only that node's own relationships. +# +# The previous form was ``MATCH ()-[r]->() WHERE r.src_id = ... AND r.dst_id = +# ... AND r.workspace = ...``: relationship property indexes in Neo4j are +# TYPE-scoped, and this pattern names no type, so no index could back it and it +# planned as an AllRelationshipsScan -- the same defect class as the label-free +# get_node() MATCH, over a set that is larger still (a graph has more +# relationships than nodes). +# +# The relationship-property predicates are deliberately RETAINED after the +# anchor. They are now evaluated against the handful of relationships incident +# to one node rather than every relationship in the graph, so they cost nothing +# -- and keeping them makes this a strictly NARROWING change: any edge the old +# form would have returned and the new one would not is one whose endpoints do +# not exist as :Node, which the edge writer (_edge_merge_cypher, which MERGEs +# both :Node endpoints) structurally cannot produce. +_EDGE_GET_BY_ENDPOINTS_CYPHER = ( + f"MATCH (src:{_UNIVERSAL_NODE_LABEL} " + "{node_id: $src_id, workspace: $workspace})-[r]->" + f"(dst:{_UNIVERSAL_NODE_LABEL} " + "{node_id: $dst_id, workspace: $workspace}) " + "WHERE r.src_id = $src_id AND r.dst_id = $dst_id " + "AND r.workspace = $workspace " + "RETURN properties(r) AS props" +) + +# Delegation lookup by the session it spawned. Named here (rather than inlined +# at the call site) so it sits beside the composite index that backs it -- +# idx_delegation_sub_session on :Delegation(sub_session_id, workspace), created +# in ensure_neo4j_schema. The query shape was already correct; before that index +# existed it simply had nothing to seek and planned as a NodeByLabelScan across +# every :Delegation node. +_DELEGATION_BY_SUB_SESSION_CYPHER = ( + "MATCH (d:Delegation {sub_session_id: $sid, workspace: $workspace}) " + "RETURN properties(d) AS props" +) + def _edge_merge_cypher(edge_type: str) -> str: """Return the UNWIND edge-MERGE query for *edge_type* — self-healing endpoints. @@ -619,6 +669,96 @@ async def run_repair(driver: Any, database: str = "neo4j") -> dict[str, int]: } +# --------------------------------------------------------------------------- +# Process-wide schema latch. +# +# ``ensure_neo4j_schema`` has two callers whose SCOPES differ, and conflating +# them is what turned a connectivity blip into a per-flush DDL storm: +# +# - The FastAPI lifespan (``main.py``) runs it ONCE PER PROCESS, before the +# server accepts a single request, with ``fail_on_data_conflict=True`` -- +# it refuses to boot unless the schema is fully established. +# - ``Neo4jGraphStore._ensure_schema`` runs it from inside ``_flush_body``, +# latched on ``self._schema_initialized`` -- a PER-STORE flag, and a store +# is constructed PER SESSION (``registry.get_or_create``). +# +# So the per-store latch made every new session pay a full ~11-statement DDL +# pass on its first flush, and -- whenever that pass could not complete -- pay +# it again on EVERY subsequent flush of EVERY session, with no backoff. Each +# attempt opens a session and issues ~11 sequential catalog statements, every +# one of which must acquire a pooled bolt connection; under pool starvation +# each acquisition burns the full acquisition timeout before failing. That is +# a large, futile, self-reinforcing load applied at exactly the moment the +# pool is already exhausted, and it runs BEFORE any data is written +# (``_flush_body`` awaits ``_ensure_schema`` ahead of Phase 1). +# +# Two changes fix that without weakening the self-heal: +# +# 1. ``_SCHEMA_READY`` is PROCESS-wide. The lifespan seeds it via +# ``mark_schema_ready()`` once its own fail-closed pass succeeds, so in +# the server the per-flush call becomes a single boolean read forever -- +# it can never discover anything cold start did not already establish. +# Contexts with no lifespan (tests, CLI tools, direct store use) never +# seed it and keep today's behaviour: the first store to fully establish +# the schema sets it for the process. +# 2. A failed pass no longer retries on the very next flush. +# ``_SCHEMA_RETRY_BACKOFF_SECONDS`` puts a floor between attempts, so a +# connectivity blip can no longer become a per-flush storm. The self-heal +# survives -- it just stops competing with the writes it is meant to +# enable. +# --------------------------------------------------------------------------- +_SCHEMA_RETRY_BACKOFF_SECONDS = 30.0 + +_SCHEMA_READY: bool = False +_SCHEMA_LAST_ATTEMPT: float | None = None + + +def mark_schema_ready() -> None: + """Record that the Neo4j schema is fully established for THIS process. + + Called by the FastAPI lifespan after its own fail-closed + ``ensure_neo4j_schema(..., fail_on_data_conflict=True)`` pass succeeds, and + by any store whose own pass fully establishes the schema. Once set, every + per-flush ``_ensure_schema`` call short-circuits to a boolean read. + """ + global _SCHEMA_READY + _SCHEMA_READY = True + + +def schema_ready() -> bool: + """Whether the schema has been fully established in this process.""" + return _SCHEMA_READY + + +def reset_schema_state() -> None: + """Clear the process-wide schema latch and backoff clock. + + Exists for tests: the latch is deliberately process-global, so without an + explicit reset one test that establishes the schema would silently suppress + the schema path in every test that ran after it. + """ + global _SCHEMA_READY, _SCHEMA_LAST_ATTEMPT + _SCHEMA_READY = False + _SCHEMA_LAST_ATTEMPT = None + + +def _schema_retry_allowed() -> bool: + """Whether enough time has passed since the last un-latched schema attempt. + + Records the attempt time as a side effect when it returns True, so two + concurrent flushes cannot both decide to run a full DDL pass. + """ + global _SCHEMA_LAST_ATTEMPT + now = time.monotonic() + if ( + _SCHEMA_LAST_ATTEMPT is not None + and now - _SCHEMA_LAST_ATTEMPT < _SCHEMA_RETRY_BACKOFF_SECONDS + ): + return False + _SCHEMA_LAST_ATTEMPT = now + return True + + async def ensure_neo4j_schema( driver: Any, database: str = "neo4j", @@ -777,6 +917,20 @@ async def _create_index(statement: str) -> bool: ) and fully_established ) + # Backs find_delegation_by_sub_session's lookup + # (_DELEGATION_BY_SUB_SESSION_CYPHER). Without it that MATCH names a + # label but no indexed property, so it plans as a NodeByLabelScan over + # EVERY :Delegation node in the graph. It is not a rare path -- the + # self-delegation resolver hits it on live ingest whenever a parent + # Delegation has already flushed out of the in-memory buffer -- and it + # grows without bound as delegation volume grows. + fully_established = ( + await _create_index( + "CREATE INDEX idx_delegation_sub_session IF NOT EXISTS " + "FOR (n:Delegation) ON (n.sub_session_id, n.workspace)" + ) + and fully_established + ) # NOTE: relationship created_by is intentionally NOT indexed in v1 — no consumer # yet; querying edges by contributor is unsupported until indexed in a future phase. @@ -1416,9 +1570,8 @@ async def get_node(self, node_id: str) -> dict[str, Any] | None: # Neo4j fallback try: result = await self._driver.execute_query( - "MATCH (n) WHERE n.node_id = $id AND n.workspace = $workspace " - "RETURN properties(n) AS props, labels(n) AS lbls", - {"id": node_id, "workspace": self.workspace}, + cast(LiteralString, _NODE_GET_BY_ID_CYPHER), + {"node_id": node_id, "workspace": self.workspace}, database_=self._database, ) records = result.records @@ -1464,8 +1617,7 @@ async def find_delegation_by_sub_session( # Neo4j fallback try: result = await self._driver.execute_query( - "MATCH (d:Delegation {sub_session_id: $sid, workspace: $workspace}) " - "RETURN properties(d) AS props", + cast(LiteralString, _DELEGATION_BY_SUB_SESSION_CYPHER), {"sid": sub_session_id, "workspace": workspace}, database_=self._database, ) @@ -1495,10 +1647,7 @@ async def get_edge(self, src_id: str, dst_id: str) -> dict[str, Any] | None: # Neo4j fallback try: result = await self._driver.execute_query( - "MATCH ()-[r]->() " - "WHERE r.src_id = $src_id AND r.dst_id = $dst_id " - "AND r.workspace = $workspace " - "RETURN properties(r) AS props", + cast(LiteralString, _EDGE_GET_BY_ENDPOINTS_CYPHER), {"src_id": src_id, "dst_id": dst_id, "workspace": self.workspace}, database_=self._database, ) @@ -1719,15 +1868,42 @@ async def _ensure_schema(self) -> None: any duplicates. This closes the "constraint created once, never retried" data-integrity gap, where a missing uniqueness constraint would otherwise let concurrent MERGE accrue duplicate Session/Event nodes until process restart. + + Two latches, checked in order, both O(1): + + - ``self._schema_initialized`` -- THIS store already established it. + - ``schema_ready()`` -- ANY caller in this PROCESS already established + it, including the FastAPI lifespan's own fail-closed cold-start pass. + In the server that is the branch which always fires: cold start + already ran ``ensure_neo4j_schema(..., fail_on_data_conflict=True)`` + and refused to boot otherwise, so a per-flush pass here can only + re-confirm what is already true -- at a cost of ~11 catalog + statements per flush, per session. See the ``_SCHEMA_READY`` block + above for why that mattered so much under pool starvation. + + When neither latch is set, the attempt is additionally rate-limited by + ``_schema_retry_allowed()``: a pass that could not complete (e.g. Neo4j + unreachable, connectivity errors swallowed so real events are not + dead-lettered) no longer retries on the very NEXT flush. Skipping a + retry is safe precisely because the pass is a no-op in the healthy case + and a storm in the degraded one -- the writes that follow it have never + depended on it having run on this particular flush. """ - if self._schema_initialized: + if self._schema_initialized or schema_ready(): + return + + if not _schema_retry_allowed(): return fully_established = await ensure_neo4j_schema(self._driver, self._database) if fully_established: self._schema_initialized = True - # else: leave the flag False so the NEXT flush retries schema init (self-heals - # once Neo4j is reachable / duplicates are cleared by the dedup pass). + # "Fully established" is a PROCESS-wide fact, not a per-store one -- + # latch it so sibling stores (one per session) never repeat the pass. + mark_schema_ready() + # else: leave both flags False so a LATER flush retries schema init, subject + # to the backoff above (self-heals once Neo4j is reachable / duplicates are + # cleared by the dedup pass). async def close(self) -> None: """Flush pending writes and close the driver, if this store owns it. diff --git a/pyproject.toml b/pyproject.toml index e813d97..8dd6420 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "context-intelligence-server" -version = "6.7.1" +version = "6.7.2" description = "Context Intelligence Server for Amplifier" requires-python = ">=3.11" dependencies = [ diff --git a/tests/conftest.py b/tests/conftest.py index 2c3b7c9..230c40a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -20,6 +20,36 @@ from context_intelligence_server.main import app, registry # noqa: E402 from context_intelligence_server.services import HookStateService # noqa: E402 +# Guarded import so the suite still COLLECTS against unfixed source (where the +# process-wide schema latch does not exist yet). Without this, reverting the +# fix would break collection for every test in the repo, and the new guards +# below would "fail" at import time rather than for the reason they exist. +try: # pragma: no cover - import resolution differs pre/post fix + from context_intelligence_server.neo4j_store import ( # noqa: E402 + reset_schema_state, + ) +except ImportError: # pragma: no cover - exercised only against unfixed code + + def reset_schema_state() -> None: # type: ignore[misc] + """No-op stand-in: unfixed code has no process-wide latch to reset.""" + + +@pytest.fixture(autouse=True) +def _reset_neo4j_schema_state() -> Generator[None, None, None]: + """Clear the process-wide Neo4j schema latch around EVERY test. + + ``neo4j_store._SCHEMA_READY`` is deliberately process-global: it is what + stops every per-session store from re-running the same ~11-statement + catalog pass on its first flush. Process-global state is also test-order + poison -- without this reset, the first test that fully establishes the + schema would silently short-circuit the schema path in every test that ran + after it, and those tests would pass or fail depending on collection order. + Resetting on both sides of the yield keeps each test's view independent. + """ + reset_schema_state() + yield + reset_schema_state() + # --------------------------------------------------------------------------- # Shared Neo4j mock helpers (used by POST /cypher tests) diff --git a/tests/neo4j/test_node_index_seek.py b/tests/neo4j/test_node_index_seek.py index 534f2da..ffe4d34 100644 --- a/tests/neo4j/test_node_index_seek.py +++ b/tests/neo4j/test_node_index_seek.py @@ -78,6 +78,24 @@ def _edge_merge_cypher(edge_type: str) -> str: ) +# The EXACT single-node read the production get_node() Neo4j fallback issues. +# Deliberately a SEPARATE try/except from the block above: _NODE_MATCH_BY_ID +# already exists on unfixed code, so folding this import in with it would make +# the whole block fall back and break the other tests for the wrong reason. +# The fallback is byte-identical to the *unfixed* get_node() query (label-free, +# $id rather than $node_id) so a revert still fails RED on AllNodesScan. +try: # pragma: no cover - import resolution differs pre/post fix + from context_intelligence_server.neo4j_store import _NODE_GET_BY_ID_CYPHER + + _GET_NODE_PARAM = "node_id" +except ImportError: # pragma: no cover - exercised only against unfixed code + _NODE_GET_BY_ID_CYPHER = ( + "MATCH (n) WHERE n.node_id = $id AND n.workspace = $workspace " + "RETURN properties(n) AS props, labels(n) AS lbls" + ) + _GET_NODE_PARAM = "id" + + def _collect_operators(plan: dict[str, Any]) -> list[str]: """Recursively collect every operatorType in a Neo4j EXPLAIN/PROFILE plan.""" ops: list[str] = [] @@ -358,3 +376,160 @@ async def test_label_patch_match_uses_index_seek_not_allnodesscan( "Label-write MATCH is not index-backed — expected a NodeIndexSeek " f"against idx_node_universal. Plan operators: {ops}" ) + + +async def test_get_node_fallback_uses_index_seek_not_allnodesscan( + neo4j_container: dict[str, Any], +) -> None: + """get_node()'s Neo4j fallback must seek the :Node index, never scan. + + This is the hot per-event read (touch_session -> get_node -> this query), + so a full scan here is not a slow edge case -- it saturates the shared + bolt pool and stalls ingest for every session. + + Observed in production on the 12.09M-node graph: the label-free form + planned as ``AllNodesScan`` (estimated 12,091,187 rows), each call ran + ~60s holding its pooled connection, 51 ran concurrently against a + 50-connection pool, and every further acquire died on the 30s + ``connection_acquisition_timeout`` -- dead-lettering live events. + + RED (unfixed): MATCH (n) WHERE n.node_id = $id ... is label-free, so no + label-scoped index is usable -> AllNodesScan. + GREEN (fixed): MATCH (n:Node {node_id, workspace}) -> NodeUniqueIndexSeek + against the :Node(node_id, workspace) composite index. + """ + await _flush_one_non_session_node(neo4j_container, "get-node-plan-1", {"v": 1}) + + ops = _explain_ops( + neo4j_container, + _NODE_GET_BY_ID_CYPHER, + **{_GET_NODE_PARAM: "get-node-plan-1", "workspace": "test"}, + ) + + assert ops, "EXPLAIN returned no plan operators for the get_node query" + assert not any("AllNodesScan" in op for op in ops), ( + "get_node()'s Neo4j fallback still does a full-graph AllNodesScan " + f"(the 12.09M-node pool-exhaustion stall). Plan operators: {ops}" + ) + assert any("IndexSeek" in op for op in ops), ( + "get_node()'s Neo4j fallback is not index-backed -- expected a " + f"NodeIndexSeek / NodeUniqueIndexSeek. Plan operators: {ops}" + ) + + +# The EXACT queries the production get_edge() / find_delegation_by_sub_session() +# fallbacks issue. Separate try/except blocks (not folded into the ones above) +# for the same reason documented there: those names already exist on unfixed +# code, so a shared block would fall back wholesale and break sibling tests for +# the wrong reason. Each fallback is byte-identical to its *unfixed* source, so +# a revert still fails RED on the real pre-fix plan. +try: # pragma: no cover - import resolution differs pre/post fix + from context_intelligence_server.neo4j_store import _EDGE_GET_BY_ENDPOINTS_CYPHER +except ImportError: # pragma: no cover - exercised only against unfixed code + _EDGE_GET_BY_ENDPOINTS_CYPHER = ( + "MATCH ()-[r]->() " + "WHERE r.src_id = $src_id AND r.dst_id = $dst_id " + "AND r.workspace = $workspace " + "RETURN properties(r) AS props" + ) + +try: # pragma: no cover - import resolution differs pre/post fix + from context_intelligence_server.neo4j_store import ( + _DELEGATION_BY_SUB_SESSION_CYPHER, + ) +except ImportError: # pragma: no cover - exercised only against unfixed code + _DELEGATION_BY_SUB_SESSION_CYPHER = ( + "MATCH (d:Delegation {sub_session_id: $sid, workspace: $workspace}) " + "RETURN properties(d) AS props" + ) + + +async def _flush_one_edge(container: dict[str, Any], src_id: str, dst_id: str) -> None: + """Drive one edge (and its two endpoint nodes) through the real flush path.""" + store = Neo4jGraphStore( + uri=container["bolt_url"], + auth=(container["user"], container["password"]), + workspace="test", + ) + try: + await store.upsert_node(src_id, {"labels": ["Event"]}) + await store.upsert_node(dst_id, {"labels": ["Event"]}) + await store.upsert_edge(src_id, dst_id, {"type": "CONTAINS"}) + await store.flush() + finally: + await store.close() + + +async def test_get_edge_fallback_does_not_scan_all_relationships( + neo4j_container: dict[str, Any], +) -> None: + """get_edge()'s Neo4j fallback must not plan as an AllRelationshipsScan. + + Relationship property indexes in Neo4j are TYPE-scoped, so the unfixed + ``MATCH ()-[r]->() WHERE r.src_id = ...`` -- which names no type -- could + not use any index and scanned every relationship in the graph. Same defect + class as the get_node() AllNodesScan, over a set that is larger still. + + RED (unfixed): MATCH ()-[r]->() WHERE r.src_id ... -> AllRelationshipsScan + GREEN (fixed): both endpoints anchored on :Node -> the planner seeks the + (node_id, workspace) unique index and expands only that + node's own relationships. + """ + await _flush_one_edge(neo4j_container, "edge-plan-src", "edge-plan-dst") + + ops = _explain_ops( + neo4j_container, + _EDGE_GET_BY_ENDPOINTS_CYPHER, + src_id="edge-plan-src", + dst_id="edge-plan-dst", + workspace="test", + ) + + assert ops, "EXPLAIN returned no plan operators for the get_edge query" + assert not any("AllRelationshipsScan" in op for op in ops), ( + f"get_edge() still scans every relationship in the graph. Plan operators: {ops}" + ) + assert not any("AllNodesScan" in op for op in ops), ( + f"get_edge() must not full-scan nodes either. Plan operators: {ops}" + ) + assert any("IndexSeek" in op for op in ops), ( + "get_edge() is not index-backed -- expected a NodeIndexSeek / " + f"NodeUniqueIndexSeek on the anchored endpoints. Plan operators: {ops}" + ) + + +async def test_delegation_lookup_uses_index_seek_not_label_scan( + neo4j_container: dict[str, Any], +) -> None: + """find_delegation_by_sub_session must seek its composite index, never label-scan. + + The query shape was always correct -- it just had no index to seek, so it + planned as a NodeByLabelScan across EVERY :Delegation node. This is a live + ingest path (the self-delegation resolver, whenever the parent Delegation + has already flushed out of the in-memory buffer), and it grows without + bound as delegation volume grows. + + RED (unfixed): no idx_delegation_sub_session -> NodeByLabelScan + GREEN (fixed): CREATE INDEX ... FOR (n:Delegation) + ON (n.sub_session_id, n.workspace) -> NodeIndexSeek + """ + # Any flush drives ensure_neo4j_schema, which creates the index under test. + await _flush_one_non_session_node(neo4j_container, "delegation-plan-seed", {"v": 1}) + + ops = _explain_ops( + neo4j_container, + _DELEGATION_BY_SUB_SESSION_CYPHER, + sid="sub-session-1", + workspace="test", + ) + + assert ops, "EXPLAIN returned no plan operators for the Delegation lookup" + assert not any("NodeByLabelScan" in op for op in ops), ( + "find_delegation_by_sub_session still scans every :Delegation node -- " + "the idx_delegation_sub_session composite index is missing or unusable. " + f"Plan operators: {ops}" + ) + assert any("IndexSeek" in op for op in ops), ( + "find_delegation_by_sub_session is not index-backed -- expected a " + f"NodeIndexSeek on (sub_session_id, workspace). Plan operators: {ops}" + ) diff --git a/tests/test_neo4j_store.py b/tests/test_neo4j_store.py index 2f919e4..dfe6233 100644 --- a/tests/test_neo4j_store.py +++ b/tests/test_neo4j_store.py @@ -23,15 +23,53 @@ import pytest from neo4j.exceptions import Neo4jError +from context_intelligence_server import neo4j_store as neo4j_store_module from context_intelligence_server.graph_store import GraphStore, QueryableStore from context_intelligence_server.neo4j_store import ( Neo4jGraphStore, _convert_temporal_props, _normalize_temporal, + _UNIVERSAL_NODE_LABEL, _validate_identifier, ensure_neo4j_schema, ) +# Guarded imports for names introduced by this change, so the suite still +# COLLECTS against unfixed source and the guards below fail for the reason they +# exist rather than at import time. Each fallback is byte-identical to the +# unfixed production value, so a revert still fails RED on the genuine pre-fix +# query / behaviour. +try: # pragma: no cover - import resolution differs pre/post fix + from context_intelligence_server.neo4j_store import ( + _DELEGATION_BY_SUB_SESSION_CYPHER, + _EDGE_GET_BY_ENDPOINTS_CYPHER, + ) +except ImportError: # pragma: no cover - exercised only against unfixed code + _EDGE_GET_BY_ENDPOINTS_CYPHER = ( + "MATCH ()-[r]->() " + "WHERE r.src_id = $src_id AND r.dst_id = $dst_id " + "AND r.workspace = $workspace " + "RETURN properties(r) AS props" + ) + _DELEGATION_BY_SUB_SESSION_CYPHER = ( + "MATCH (d:Delegation {sub_session_id: $sid, workspace: $workspace}) " + "RETURN properties(d) AS props" + ) + +try: # pragma: no cover - import resolution differs pre/post fix + from context_intelligence_server.neo4j_store import ( + mark_schema_ready, + schema_ready, + ) +except ImportError: # pragma: no cover - exercised only against unfixed code + + def mark_schema_ready() -> None: # type: ignore[misc] + """No-op stand-in: unfixed code has no process-wide schema latch.""" + + def schema_ready() -> bool: # type: ignore[misc] + """Always False: unfixed code has no process-wide schema latch.""" + return False + # --------------------------------------------------------------------------- # Helpers @@ -1015,9 +1053,57 @@ async def test_get_node_fallback_queries_by_node_id_property(): "Expected driver.execute_query to be called for Neo4j fallback" ) query: str = mock_execute.call_args[0][0] - assert "n.node_id" in query, ( - f"Fallback query must filter on 'n.node_id' (the property flush stores on nodes), " - f"but the issued query was: {query!r}" + # Syntax-agnostic: the key may appear as ``n.node_id = $node_id`` (WHERE + # form) or as ``{node_id: $node_id}`` (map-pattern form). What matters is + # that identity is keyed on the ``node_id`` property and never on ``id``. + assert "node_id" in query, ( + f"Fallback query must key on the 'node_id' property (the one flush stores " + f"on nodes), but the issued query was: {query!r}" + ) + assert "n.id" not in query, ( + f"Fallback query must NOT filter on 'n.id' -- no node carries that " + f"property, so it silently returns None. Issued query: {query!r}" + ) + params = mock_execute.call_args[0][1] + assert params.get("node_id") == "session-123", ( + f"Fallback must bind the node id to the $node_id parameter; got {params!r}" + ) + + +async def test_get_node_fallback_is_label_scoped_for_index_seek(): + """Post-flush fallback must scope the MATCH to the universal :Node label. + + Neo4j property indexes are label-scoped, so a label-free ``MATCH (n)`` + cannot use ANY index and plans as an AllNodesScan over the whole graph. + On the 12.09M-node production graph that made every get_node() fallback + take ~60s while holding its pooled bolt connection; 51 concurrent calls + exhausted the 50-connection shared pool and ingest dead-lettered live + events on the 30s connection_acquisition_timeout. + + The live query-plan proof lives in + tests/neo4j/test_node_index_seek.py, but that file is marked ``neo4j`` and + is deselected from the default suite -- so this unit-level guard is the one + that runs in ordinary CI. + + FAILS before fix -> query is 'MATCH (n) WHERE ...' (no label). + PASSES after fix -> query scopes to ':Node'. + """ + store = _make_store() + store._node_buffer = {} + + mock_result = MagicMock() + mock_result.records = [] + mock_execute = AsyncMock(return_value=mock_result) + store._driver.execute_query = mock_execute # type: ignore[attr-defined] + + await store.get_node("session-123") + + query: str = mock_execute.call_args[0][0] + assert f":{_UNIVERSAL_NODE_LABEL}" in query, ( + f"Fallback query must scope the MATCH to the universal " + f"':{_UNIVERSAL_NODE_LABEL}' label so the composite " + f"(node_id, workspace) index backs it -- a label-free MATCH plans as a " + f"full-graph AllNodesScan. Issued query: {query!r}" ) @@ -2483,10 +2569,21 @@ async def test_ensure_schema_does_not_latch_when_constraint_uncreated_on_connect "never retried and duplicate Session/Event nodes accrue." ) - async def test_ensure_schema_retries_until_established_then_latches(self) -> None: + async def test_ensure_schema_retries_until_established_then_latches( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: """Degraded first run retries on the next flush, then latches and stops re-running.""" from neo4j.exceptions import ServiceUnavailable + # An un-latched retry is now rate-limited (_SCHEMA_RETRY_BACKOFF_SECONDS) + # so a connectivity blip cannot become a per-flush DDL storm. This test + # is about the RETRY semantics, not the rate limit -- the rate limit has + # its own guard in TestSchemaRetryBackoff -- so collapse the window to + # zero and let the two flushes run back to back as they always did. + monkeypatch.setattr( + neo4j_store_module, "_SCHEMA_RETRY_BACKOFF_SECONDS", 0.0, raising=False + ) + store = _make_store() store._schema_initialized = False @@ -3464,3 +3561,408 @@ async def consume(self) -> None: session_calls = [c for c in captured if "n:Session" in str(c["statement"])] rows = session_calls[0]["kwargs"]["rows"] # type: ignore[index] assert "working_dir" not in rows[0] + + +# --------------------------------------------------------------------------- +# Wide-query containment: every persisted read must be able to SEEK an index. +# +# PR #98 fixed get_node()'s label-free MATCH, which planned as an AllNodesScan +# over the 12.09M-node production graph, held its pooled bolt connection for +# ~60s per call, and (51 concurrent, 50-connection pool) exhausted the shared +# pool until live activity records were dead-lettered. +# +# These guards cover the REST of that bug class rather than the one instance: +# the two remaining unindexed reads, and a source-level tripwire so the next +# hand-rolled query cannot re-introduce the pattern silently. +# --------------------------------------------------------------------------- + + +async def test_get_edge_fallback_is_node_anchored_not_all_relationships_scan(): + """get_edge()'s Neo4j fallback must anchor on :Node, never scan all relationships. + + ``MATCH ()-[r]->() WHERE r.src_id = ...`` names no relationship type, and + Neo4j relationship property indexes are TYPE-scoped, so no index could back + it: it plans as an AllRelationshipsScan. That is the same defect class as + the get_node() AllNodesScan over a set that is larger still -- a graph has + more relationships than nodes. + + FAILS before fix -> query is 'MATCH ()-[r]->() WHERE ...' (unanchored). + PASSES after fix -> both endpoints are :Node identity patterns, so the + planner seeks the (node_id, workspace) unique index + and walks only that node's own relationships. + """ + store = _make_store(workspace="ws-1") + store._edge_buffer = {} + + mock_result = MagicMock() + mock_result.records = [] + mock_execute = AsyncMock(return_value=mock_result) + store._driver.execute_query = mock_execute # type: ignore[attr-defined] + + await store.get_edge("src-1", "dst-1") + + query: str = mock_execute.call_args[0][0] + assert "MATCH ()-[r]->()" not in query.replace(" ", "").replace( + "MATCH()-[r]->()", "MATCH ()-[r]->()" + ), f"get_edge must not scan every relationship in the graph. Query: {query!r}" + assert query.count(f":{_UNIVERSAL_NODE_LABEL}") == 2, ( + f"get_edge must anchor BOTH endpoints on ':{_UNIVERSAL_NODE_LABEL}' so each " + f"is an index seek on (node_id, workspace). Query: {query!r}" + ) + assert query == _EDGE_GET_BY_ENDPOINTS_CYPHER, ( + "get_edge must issue the shared _EDGE_GET_BY_ENDPOINTS_CYPHER constant " + "rather than a hand-rolled string that can drift." + ) + + +async def test_get_edge_fallback_still_binds_endpoints_and_workspace(): + """The anchored rewrite must not change get_edge's parameter contract.""" + store = _make_store(workspace="ws-alpha") + store._edge_buffer = {} + + mock_result = MagicMock() + mock_result.records = [] + mock_execute = AsyncMock(return_value=mock_result) + store._driver.execute_query = mock_execute # type: ignore[attr-defined] + + await store.get_edge("a", "b") + + params = mock_execute.call_args[0][1] + assert params == {"src_id": "a", "dst_id": "b", "workspace": "ws-alpha"}, ( + f"get_edge must still bind src_id/dst_id/workspace unchanged; got {params!r}" + ) + + +async def test_delegation_lookup_uses_shared_indexed_constant(): + """find_delegation_by_sub_session must issue the constant the index backs.""" + store = _make_store(workspace="ws-1") + + mock_result = MagicMock() + mock_result.records = [] + mock_execute = AsyncMock(return_value=mock_result) + store._driver.execute_query = mock_execute # type: ignore[attr-defined] + + await store.find_delegation_by_sub_session("sub-1", "ws-1") + + query: str = mock_execute.call_args[0][0] + assert query == _DELEGATION_BY_SUB_SESSION_CYPHER, ( + "find_delegation_by_sub_session must issue the shared " + "_DELEGATION_BY_SUB_SESSION_CYPHER constant, which sits beside the " + f"idx_delegation_sub_session index that backs it. Query: {query!r}" + ) + assert "sub_session_id" in query and "workspace" in query, ( + "The lookup must key on BOTH indexed properties (sub_session_id, " + f"workspace) or the composite index cannot be seeked. Query: {query!r}" + ) + + +async def test_schema_creates_delegation_sub_session_index(): + """ensure_neo4j_schema must create the index backing the Delegation lookup. + + Without it the lookup names a label but no indexed property, so it plans as + a NodeByLabelScan over EVERY :Delegation node. That path is live ingest -- + the self-delegation resolver hits it whenever the parent Delegation has + already flushed out of the in-memory buffer -- and it grows without bound + as delegation volume grows. + """ + session = AsyncMock() + session.__aenter__ = AsyncMock(return_value=session) + session.run = AsyncMock(return_value=AsyncMock()) + driver = MagicMock() + driver.session = MagicMock(return_value=session) + + await ensure_neo4j_schema(driver) + + statements = [str(call.args[0]) for call in session.run.await_args_list] + delegation_idx = [ + s for s in statements if "Delegation" in s and "CREATE INDEX" in s + ] + assert delegation_idx, ( + "ensure_neo4j_schema must CREATE INDEX on :Delegation for the " + f"sub_session_id lookup. Statements issued: {statements}" + ) + assert any("sub_session_id" in s and "workspace" in s for s in delegation_idx), ( + "The Delegation index must be composite on (sub_session_id, workspace) " + f"-- the exact keys the lookup filters on. Got: {delegation_idx}" + ) + + +def test_no_unindexed_scan_patterns_in_neo4j_store_source(): + """Source-level tripwire: no NEW label-free / all-relationship query may ship. + + get_node() broke because it hand-rolled its own query string while three + sibling call sites composed theirs from the shared ``_NODE_MATCH_BY_ID`` + prefix -- so it inherited none of their index-scoped-ness, and none of the + existing plan guards covered it. Guarding one call site per outage does not + scale; this guards the whole module, including queries nobody has written + yet. + + The allow-list below is the complete, reviewed inventory of deliberately + label-free statements. Every one is either answered from Neo4j's O(1) + counts store or is an explicitly O(graph-size) repair path that only runs + under ``doctor --fix`` -- never on the hot ingest path. Adding an entry is + the review checkpoint: if a new query needs to be here, that needs to be an + argued decision, not an accident. + """ + import pathlib # noqa: PLC0415 + import re # noqa: PLC0415 + + source_path = ( + pathlib.Path(neo4j_store_module.__file__).resolve() # type: ignore[arg-type] + ) + lines = source_path.read_text(encoding="utf-8").splitlines() + + # MATCH (n) / MATCH (foo) -- a node pattern carrying no label. Neo4j property + # indexes are label-scoped, so such a pattern can never use one. + unlabelled_node = re.compile(r"MATCH\s*\(\s*[A-Za-z_][A-Za-z0-9_]*\s*\)") + # MATCH ()-[r]-> -- a relationship pattern carrying no type. Relationship + # property indexes are type-scoped, so likewise no index can back it. + untyped_relationship = re.compile(r"MATCH\s*\(\s*\)\s*-\s*\[") + + allowed = { + # Docstring prose explaining why the edge writer uses MERGE, not MATCH. + "Why MERGE and not MATCH: the old ``MATCH (src) MATCH (dst)`` was an inner join", + # _DUPLICATE_DETECT_CYPHER -- doctor --fix only, explicitly O(graph-size). + '"MATCH (n) "', + # _UNTAGGED_COUNT_CYPHER -- doctor --fix pre/post accounting only. + 'f"MATCH (n) WHERE NOT n:{_UNIVERSAL_NODE_LABEL} RETURN count(n) AS c"', + # _TOTAL_NODE_COUNT_CYPHER -- answered from the O(1) counts store. + '_TOTAL_NODE_COUNT_CYPHER = "MATCH (n) RETURN count(n) AS c"', + # backfill_node_labels -- doctor --fix only, batched IN TRANSACTIONS. + 'f"MATCH (n) WHERE NOT n:{_UNIVERSAL_NODE_LABEL} "', + } + + offenders: list[str] = [] + for lineno, line in enumerate(lines, start=1): + stripped = line.strip() + if stripped.startswith("#"): + continue # a comment cannot execute + if stripped in allowed: + continue + if unlabelled_node.search(line) or untyped_relationship.search(line): + offenders.append(f"{source_path.name}:{lineno}: {stripped}") + + assert not offenders, ( + "New unindexed scan pattern(s) in neo4j_store.py. A label-free MATCH " + "(n) or an untyped MATCH ()-[r]-> cannot use ANY Neo4j index and plans " + "as a full AllNodesScan / AllRelationshipsScan -- the 12.09M-node, " + "~60s-per-call, pool-exhausting stall this module has now hit twice. " + "Scope the pattern to a label/type it has an index for, or -- if the " + "scan is genuinely intended and off the hot path -- add the exact line " + "to this test's allow-list with a comment saying why.\n " + + "\n ".join(offenders) + ) + + +# --------------------------------------------------------------------------- +# Schema pass: process-wide latch and retry backoff. +# +# ensure_neo4j_schema ran once per STORE, and a store is created per SESSION, +# so every new session paid a full ~11-statement DDL pass on its first flush -- +# and, whenever that pass could not complete, paid it again on EVERY subsequent +# flush with no backoff. Each statement must acquire a pooled bolt connection, +# and _flush_body awaits _ensure_schema BEFORE writing any data, so under pool +# starvation this became a self-reinforcing storm competing with the very +# writes it exists to enable. +# --------------------------------------------------------------------------- + + +class TestSchemaProcessLatch: + """The schema latch is process-wide, and the lifespan can seed it.""" + + async def test_ensure_schema_skips_entirely_when_process_latch_set(self) -> None: + """A store must not re-run the DDL pass once ANY caller established it. + + This is the branch that always fires in the server: the FastAPI + lifespan runs ensure_neo4j_schema(fail_on_data_conflict=True) before + accepting a single request and refuses to boot unless it succeeds, so + a per-flush pass can only re-confirm what is already true -- at ~11 + catalog statements per flush, per session. + + FAILS before fix -> the per-store flag is the only latch, so the store + opens a session and re-runs the whole pass. + PASSES after fix -> schema_ready() short-circuits before any I/O. + """ + store = _make_store() + store._schema_initialized = False + store._driver.session = MagicMock( + side_effect=AssertionError( + "schema pass ran despite the process-wide latch being set" + ) + ) + + mark_schema_ready() + assert schema_ready() is True + + await store._ensure_schema() # must not touch the driver at all + + store._driver.session.assert_not_called() # type: ignore[attr-defined] + + async def test_successful_store_pass_latches_for_the_whole_process(self) -> None: + """One store fully establishing the schema spares its siblings the pass. + + Stores are per-session; without a process-wide latch, N concurrent + sessions each pay the same DDL pass on their first flush. + """ + assert schema_ready() is False, "autouse fixture must reset the latch" + + first = _make_store() + first._schema_initialized = False + ok_session = AsyncMock() + ok_session.__aenter__ = AsyncMock(return_value=ok_session) + ok_session.run = AsyncMock(return_value=AsyncMock()) + first._driver.session = MagicMock(return_value=ok_session) + + await first._ensure_schema() + + assert first._schema_initialized is True + assert schema_ready() is True, ( + "A fully established schema is a PROCESS-wide fact -- it must latch " + "for sibling stores, not just for the store that established it." + ) + + sibling = _make_store() + sibling._schema_initialized = False + sibling._driver.session = MagicMock( + side_effect=AssertionError("sibling store re-ran the schema pass") + ) + await sibling._ensure_schema() + sibling._driver.session.assert_not_called() # type: ignore[attr-defined] + + +class TestSchemaRetryBackoff: + """A failed schema pass must not retry on the very next flush.""" + + @staticmethod + def _degraded_session() -> AsyncMock: + from neo4j.exceptions import ServiceUnavailable # noqa: PLC0415 + + session = AsyncMock() + session.__aenter__ = AsyncMock(return_value=session) + + async def _run(statement: str, *args: object, **kwargs: object) -> AsyncMock: + if "CREATE CONSTRAINT" in str(statement): + raise ServiceUnavailable("connection refused") + return AsyncMock() + + session.run = AsyncMock(side_effect=_run) + return session + + async def test_failed_pass_does_not_retry_on_the_next_flush(self) -> None: + """Back-to-back flushes must issue ONE degraded pass, not two. + + Without backoff, a connectivity blip turns into a per-flush DDL storm: + every flush of every session re-attempts ~11 catalog statements, each + one waiting out the full connection-acquisition timeout against the + pool it is already starving. + + FAILS before fix -> both flushes run the pass (2 sessions opened). + PASSES after fix -> the second is suppressed by the backoff window. + """ + store = _make_store() + store._schema_initialized = False + sessions = [self._degraded_session(), self._degraded_session()] + store._driver.session = MagicMock(side_effect=sessions) + + await store._ensure_schema() + await store._ensure_schema() + + assert store._driver.session.call_count == 1, ( # type: ignore[attr-defined] + "A failed schema pass must not be retried on the immediately " + "following flush -- that is the per-flush DDL storm. Sessions " + f"opened: {store._driver.session.call_count}" # type: ignore[attr-defined] + ) + assert store._schema_initialized is False, ( + "Suppressing a retry must not latch a half-built schema." + ) + + async def test_retry_resumes_once_the_backoff_window_elapses( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Backoff delays the retry; it must never cancel the self-heal.""" + store = _make_store() + store._schema_initialized = False + degraded = self._degraded_session() + ok_session = AsyncMock() + ok_session.__aenter__ = AsyncMock(return_value=ok_session) + ok_session.run = AsyncMock(return_value=AsyncMock()) + store._driver.session = MagicMock(side_effect=[degraded, ok_session]) + + await store._ensure_schema() + assert store._schema_initialized is False + + # Simulate the window elapsing rather than sleeping through it. + monkeypatch.setattr( + neo4j_store_module, "_SCHEMA_RETRY_BACKOFF_SECONDS", 0.0, raising=False + ) + + await store._ensure_schema() + assert store._schema_initialized is True, ( + "Once the window elapses the retry must run and latch -- backoff " + "delays the self-heal, it does not disable it." + ) + + +class TestLifespanSchemaLatchSeeding: + """The lifespan may only seed the process latch on a FULLY established pass. + + ``ensure_neo4j_schema(fail_on_data_conflict=True)`` fails closed on a + :Node constraint DATA conflict, but a CONNECTIVITY failure on any single + index/constraint is swallowed and reported through the RETURN VALUE. A + lifespan that latched unconditionally would mark a half-built schema ready + and permanently disable the per-flush self-heal -- reopening the + "constraint created once, never retried" data-integrity gap. + """ + + async def test_incomplete_schema_pass_must_not_latch(self) -> None: + """A False return from the cold-start pass must leave the latch unset.""" + assert schema_ready() is False, "autouse fixture must reset the latch" + + # Mirror the lifespan's own decision rule against a degraded pass. + from neo4j.exceptions import ServiceUnavailable # noqa: PLC0415 + + session = AsyncMock() + session.__aenter__ = AsyncMock(return_value=session) + + async def _run(statement: str, *args: object, **kwargs: object) -> AsyncMock: + if "CREATE CONSTRAINT" in str(statement): + raise ServiceUnavailable("connection refused") + return AsyncMock() + + session.run = AsyncMock(side_effect=_run) + driver = MagicMock() + driver.session = MagicMock(return_value=session) + + fully_established = await ensure_neo4j_schema( + driver, fail_on_data_conflict=True + ) + assert fully_established is False, ( + "A swallowed connectivity failure must be reported via the return " + "value -- this is the signal the lifespan gates its seed on." + ) + + if fully_established: # pragma: no cover - guarding the wrong branch + mark_schema_ready() + + assert schema_ready() is False, ( + "Latching on an incomplete schema pass would permanently disable " + "the per-flush self-heal for the whole process." + ) + + async def test_complete_schema_pass_latches(self) -> None: + """A True return must seed the latch, so stores skip the redundant pass.""" + session = AsyncMock() + session.__aenter__ = AsyncMock(return_value=session) + session.run = AsyncMock(return_value=AsyncMock()) + driver = MagicMock() + driver.session = MagicMock(return_value=session) + + fully_established = await ensure_neo4j_schema( + driver, fail_on_data_conflict=True + ) + assert fully_established is True + + if fully_established: + mark_schema_ready() + + assert schema_ready() is True