From a0c8333a0d85bcd3b8b1642e0bead23be8f4a9f8 Mon Sep 17 00:00:00 2001 From: colombod Date: Wed, 26 Aug 2026 12:22:12 +0000 Subject: [PATCH] feat(session): working_dir end-to-end, agent persistence, and IncompleteSession heal-forward MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Session-node data-quality changes, stacked on the durable-ingest cursor work. No cursor/queue code -- that lives in the base branch underneath. working_dir is now a complete feature: EventRequest accepts an optional top-level working_dir envelope field (blank/whitespace-only rejected), the events endpoint lifts it into the event data, and ensure_session_node writes it onto the Session node. Populate-if-missing -- a node created before working_dir was known is backfilled by the first subsequent event carrying a non-empty value; an already-set value is never overwritten (the Neo4j write uses coalesce(n.working_dir, row.working_dir) rather than last-write-wins, and working_dir is excluded from the blind SET n += row.props). agent is backfilled on the existing-node branch with the same populate-if-missing rule. The agent name for a spawned sub-session arrives only on the parent's delegate:agent_spawned event, but the child's own session:start can create the node first, so the parent event routinely lands on the existing-node branch where agent was previously dropped -- leaving :Session.agent empty and undercounting WHERE s.agent = ... queries. IncompleteSession heal-forward: a stale IncompleteSession marker (stamped when a session:end drained before its session:start/fork) is stripped the moment the real start/fork is processed, leaving only the correct terminal label. classify() routes every start/fork transition through _heal_forward. Tests: heal-forward unit matrix + neo4j race; working_dir non-overwrite unit + neo4j race; agent non-overwrite unit + neo4j race. Version 7.1.0. 🤖 Generated with [Amplifier](https://github.com/microsoft/amplifier) Co-Authored-By: Amplifier <240397093+microsoft-amplifier@users.noreply.github.com> --- CHANGELOG.md | 27 +++ .../handlers/data_layer_2/session.py | 91 ++++++--- context_intelligence_server/main.py | 5 + context_intelligence_server/models.py | 23 +++ context_intelligence_server/neo4j_store.py | 25 ++- context_intelligence_server/services.py | 30 ++- pyproject.toml | 2 +- tests/handlers/data_layer_2/test_session.py | 172 ++++++++++++++-- tests/neo4j/test_agent_field_ordering_race.py | 142 +++++++++++++ .../test_incomplete_session_heal_forward.py | 191 ++++++++++++++++++ tests/neo4j/test_working_dir_non_overwrite.py | 158 +++++++++++++++ tests/test_services.py | 127 ++++++++++++ uv.lock | 2 +- 13 files changed, 933 insertions(+), 62 deletions(-) create mode 100644 tests/neo4j/test_agent_field_ordering_race.py create mode 100644 tests/neo4j/test_incomplete_session_heal_forward.py create mode 100644 tests/neo4j/test_working_dir_non_overwrite.py diff --git a/CHANGELOG.md b/CHANGELOG.md index d2e6da36..4fe75839 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,33 @@ All notable changes to the Context Intelligence Server are recorded here. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [7.1.0] + +### Added + +- **`working_dir` end-to-end on the Session node.** `EventRequest` accepts an + optional top-level `working_dir` envelope field (rejected only when + blank/whitespace-only); the events endpoint lifts it into the event data, and + `ensure_session_node` writes it onto the Session node. Populate-if-missing: a + node created before `working_dir` was known is backfilled by the first + subsequent event that carries a non-empty value. An already-set value is + never overwritten — the Neo4j write uses + `coalesce(n.working_dir, row.working_dir)` rather than last-write-wins. + +### Fixed + +- **`agent` persists across the delivery-order race.** The agent name for a + spawned sub-session arrives only on the parent's `delegate:agent_spawned` + event, but the child's own `session:start` can create the Session node first. + `ensure_session_node` now backfills `agent` on the existing-node branch with + the same populate-if-missing rule as `working_dir`, so it is no longer + silently dropped (which left `:Session.agent` empty and undercounted + `WHERE s.agent = ...` queries). +- **IncompleteSession heal-forward.** A stale `IncompleteSession` marker — + stamped when a session's `session:end` drained before its `session:start`/ + `session:fork` — is now stripped the moment the real start/fork is processed, + leaving only the correct terminal label. + ## [7.0.0] ### Changed (breaking) diff --git a/context_intelligence_server/handlers/data_layer_2/session.py b/context_intelligence_server/handlers/data_layer_2/session.py index e2f5771d..b0ccb7f3 100644 --- a/context_intelligence_server/handlers/data_layer_2/session.py +++ b/context_intelligence_server/handlers/data_layer_2/session.py @@ -68,46 +68,36 @@ class SessionLabelStateMachine: ForkedSession > SubSession > RootSession in specificity (terminal ordering). """ + @staticmethod + def _heal_forward(transition: LabelTransition) -> LabelTransition: + """Strip a stale IncompleteSession marker on a start/fork transition. + + IncompleteSession is only correct at session:end on a bare session whose + start/fork was lost. Its co-occurrence with a start/fork is always the + out-of-order case -- a forked sub-session's session:end drained before + its session:start/fork, stamping the marker first. Removing an absent + label is a no-op, so applying this to every start/fork result is safe and + makes the invariant impossible to miss when a branch is added later. + """ + if "IncompleteSession" in transition.remove: + return transition + return LabelTransition( + add=transition.add, remove=[*transition.remove, "IncompleteSession"] + ) + def classify( self, current_type: str | None, event: str, has_parent: bool ) -> LabelTransition: # StubSession is a plain observability marker, not part of the # terminal lattice; every branch assigning a real terminal label - # also clears it (removing an absent label is a no-op). + # also clears it (removing an absent label is a no-op). Start/fork + # results are passed through _heal_forward so a stale IncompleteSession + # marker is stripped on the transition (it is only correct at end). if event == "start": - if current_type in ("ForkedSession", "SubSession"): - return LabelTransition() - if current_type == "RootSession": - if has_parent: - return LabelTransition( - add=["SubSession", "SST_EVENT"], - remove=["RootSession", "StubSession"], - ) - return LabelTransition() - # bare session (current_type is None) - if has_parent: - return LabelTransition( - add=["Session", "SubSession", "SST_EVENT"], - remove=["StubSession"], - ) - return LabelTransition( - add=["RootSession", "Session", "SST_EVENT"], - remove=["StubSession"], - ) + return self._heal_forward(self._classify_start(current_type, has_parent)) if event == "fork": - if current_type == "ForkedSession": - return LabelTransition() - if current_type in ("RootSession", "SubSession"): - return LabelTransition( - add=["ForkedSession", "SST_EVENT"], - remove=[current_type, "StubSession"], - ) - # bare session (current_type is None) - return LabelTransition( - add=["Session", "ForkedSession", "SST_EVENT"], - remove=["StubSession"], - ) + return self._heal_forward(self._classify_fork(current_type, has_parent)) if event == "end": if current_type is not None: @@ -120,6 +110,43 @@ def classify( raise ValueError(f"classify() received unknown event: {event!r}") + @staticmethod + def _classify_start(current_type: str | None, has_parent: bool) -> LabelTransition: + if current_type in ("ForkedSession", "SubSession"): + return LabelTransition() + if current_type == "RootSession": + if has_parent: + return LabelTransition( + add=["SubSession", "SST_EVENT"], + remove=["RootSession", "StubSession"], + ) + return LabelTransition() + # bare session (current_type is None) + if has_parent: + return LabelTransition( + add=["Session", "SubSession", "SST_EVENT"], + remove=["StubSession"], + ) + return LabelTransition( + add=["RootSession", "Session", "SST_EVENT"], + remove=["StubSession"], + ) + + @staticmethod + def _classify_fork(current_type: str | None, has_parent: bool) -> LabelTransition: + if current_type == "ForkedSession": + return LabelTransition() + if current_type in ("RootSession", "SubSession"): + return LabelTransition( + add=["ForkedSession", "SST_EVENT"], + remove=[current_type, "StubSession"], + ) + # bare session (current_type is None) + return LabelTransition( + add=["Session", "ForkedSession", "SST_EVENT"], + remove=["StubSession"], + ) + class SessionHandler: """Handles session lifecycle events. diff --git a/context_intelligence_server/main.py b/context_intelligence_server/main.py index c1ffb22e..dfb9cf31 100644 --- a/context_intelligence_server/main.py +++ b/context_intelligence_server/main.py @@ -1140,6 +1140,11 @@ async def post_events( body = await http_request.body() body_obj = json.loads(body) body_obj["created_by"] = contributor_id # overwrite, never setdefault + # Lift the optional top-level working_dir envelope field into data so the + # Session-node write sees it; absent/empty leaves Session.working_dir null + # for a later event to populate. + if request.working_dir and isinstance(body_obj.get("data"), dict): + body_obj["data"]["working_dir"] = request.working_dir body = json.dumps(body_obj, separators=(",", ":")).encode() await registry.queue_manager.append(worker_key, body) # Bytes are on disk: the key may be burned now (a failed append diff --git a/context_intelligence_server/models.py b/context_intelligence_server/models.py index 03cf0793..6a8bc52b 100644 --- a/context_intelligence_server/models.py +++ b/context_intelligence_server/models.py @@ -14,10 +14,19 @@ class EventRequest(BaseModel): The Amplifier client must always supply workspace on every event. Events without workspace (e.g. an incorrectly configured hook) are rejected at the endpoint with HTTP 422. + + working_dir is OPTIONAL — the bundle hook emits it + as a top-level envelope field alongside workspace, but older clients/events + won't have it. Absent/empty is fine and leaves the Session node's + working_dir property null. Populate-if-missing: the Session node's + working_dir is filled in by the first subsequent event (including a + re-import via the upload CLI) that carries a non-empty value, but an + already-populated value is never overwritten. """ event: str workspace: str + working_dir: str | None = None idempotency_key: str | None = None data: dict[str, Any] @@ -29,6 +38,20 @@ def workspace_must_not_be_empty(cls, v: str) -> str: raise ValueError("workspace must not be empty") return v + @field_validator("working_dir") + @classmethod + def working_dir_must_not_be_blank(cls, v: str | None) -> str | None: + """Allow ``None`` (working_dir is optional, unlike workspace) but reject + blank/whitespace-only strings. + + Mirrors ``workspace_must_not_be_empty``'s normalize-or-reject stance: a + whitespace-only value (e.g. ``" "``) is never a legitimate path and must + not write through to the Session node verbatim. + """ + if v is not None and not v.strip(): + raise ValueError("working_dir must not be blank") + return v + class EventResponse(BaseModel): """Response returned after an event is accepted.""" diff --git a/context_intelligence_server/neo4j_store.py b/context_intelligence_server/neo4j_store.py index 08b7cc3e..6dad68a7 100644 --- a/context_intelligence_server/neo4j_store.py +++ b/context_intelligence_server/neo4j_store.py @@ -1009,12 +1009,19 @@ def _chunk_list( def _build_node_props(data: dict[str, Any], workspace: str) -> dict[str, Any]: """Assemble the sanitized props dict for a single node row. - ``labels`` and ``created_by`` are excluded from the returned dict: + ``labels``, ``created_by``, and ``working_dir`` are excluded from the returned dict: - ``labels`` are applied separately via ``SET n:Label`` statements (not stored as props). - ``created_by`` travels ONLY as the ``$created_by`` query param (never a node property) so it cannot affect node identity or clobber the ``ON CREATE SET n.created_by`` stamp. + - ``working_dir`` is excluded from the blind ``SET n += row.props`` so the + Session-node write can apply ``coalesce(n.working_dir, row.working_dir)`` + instead of last-write-wins -- an already-set working_dir is never clobbered. """ - raw = {k: v for k, v in data.items() if k not in ("labels", "created_by")} + raw = { + k: v + for k, v in data.items() + if k not in ("labels", "created_by", "working_dir") + } _convert_temporal_props(raw) # ISO str -> datetime, in place props = Neo4jGraphStore._sanitize_properties(raw) props["workspace"] = workspace @@ -1056,6 +1063,13 @@ async def _write_batch( row: dict[str, Any] = {"node_id": node_id, "props": props} if "Session" in labels: + # working_dir is a Session-only property carried as a separate + # top-level row key (not in props) so the MERGE below can coalesce + # it rather than blindly overwrite. Omitted when empty/absent, so + # coalesce(n.working_dir, null) is a no-op for such rows. + working_dir_value = data.get("working_dir") + if working_dir_value: + row["working_dir"] = working_dir_value session_rows.append(row) else: other_rows.append(row) @@ -1085,7 +1099,12 @@ async def _write_batch( f"MERGE (n:{_UNIVERSAL_NODE_LABEL} " "{node_id: row.node_id, workspace: row.props.workspace}) " "ON CREATE SET n.created_by = $created_by " - "SET n += row.props, n:Session", + "SET n += row.props, n:Session " + # Populate-if-missing: fill working_dir only when it is still null, + # never clobber an already-set value (a concurrent/replica writer or + # an earlier event may have populated it). Rows without a working_dir + # carry a null row key, so this is a no-op for them. + "SET n.working_dir = coalesce(n.working_dir, row.working_dir)", rows=session_rows, created_by=created_by, ) diff --git a/context_intelligence_server/services.py b/context_intelligence_server/services.py index 7182af82..68a2ca0f 100644 --- a/context_intelligence_server/services.py +++ b/context_intelligence_server/services.py @@ -343,10 +343,28 @@ async def ensure_session_node(self, session_id: str, data: dict[str, Any]) -> No if existing is not None: # Upsert a stub so this worker's own flush uses MERGE (idempotent) # instead of racing a second worker into creating a duplicate node. - await self.graph.upsert_node( - session_id, - {"labels": ["Session"], "status": "running", "session_id": session_id}, - ) + stub_data: dict[str, Any] = { + "labels": ["Session"], + "status": "running", + "session_id": session_id, + } + # Populate-if-missing: backfill working_dir on a node created before + # working_dir was known (or by a bare reference). Only when the event + # supplies one AND the node still lacks it; the DB-level coalesce + # guarantees an already-set value is never clobbered. + if data.get("working_dir") and not existing.get("working_dir"): + stub_data["working_dir"] = data["working_dir"] + # Same populate-if-missing rule for `agent`. The agent name for a + # spawned sub-session arrives ONLY on the parent's + # delegate:agent_spawned event; the child's own session:start (no + # top-level agent) can create the node first, so that later parent + # event routinely lands HERE. Without this, the parent's `agent` was + # silently dropped, leaving :Session.agent empty and breaking any + # `WHERE s.agent = ...` query. Only writes when the event supplies it + # AND the node still lacks it, so an already-set value is preserved. + if data.get("agent") and not existing.get("agent"): + stub_data["agent"] = data["agent"] + await self.graph.upsert_node(session_id, stub_data) self._seen_sessions.add(session_id) return @@ -366,6 +384,10 @@ async def ensure_session_node(self, session_id: str, data: dict[str, Any]) -> No node_data["started_at"] = _ts if "agent" in data: node_data["agent"] = data["agent"] + # Lift working_dir onto the Session node when the event carries one; + # absent/empty leaves it null for a later event to populate. + if data.get("working_dir"): + node_data["working_dir"] = data["working_dir"] await self.graph.upsert_node(session_id, node_data) self._seen_sessions.add(session_id) # only cache after successful write diff --git a/pyproject.toml b/pyproject.toml index 7e771dfc..cc9105dd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "context-intelligence-server" -version = "7.0.0" +version = "7.1.0" description = "Context Intelligence Server for Amplifier" requires-python = ">=3.11" dependencies = [ diff --git a/tests/handlers/data_layer_2/test_session.py b/tests/handlers/data_layer_2/test_session.py index 19983caf..767120c7 100644 --- a/tests/handlers/data_layer_2/test_session.py +++ b/tests/handlers/data_layer_2/test_session.py @@ -15,12 +15,11 @@ from typing import Any import pytest - from context_intelligence_server.handlers.data_layer_2.session import ( + _TYPE_LABELS, LabelTransition, SessionHandler, SessionLabelStateMachine, - _TYPE_LABELS, _current_type, _parent_of, _warn_if_dual_terminal, @@ -28,7 +27,6 @@ from context_intelligence_server.services import GraphState, HookStateService from context_intelligence_server.utils import make_node_id - # --------------------------------------------------------------------------- # _RecordingGraphState — spy helper for TestSessionHandlerUsesSessionLabeledMerge # --------------------------------------------------------------------------- @@ -2078,6 +2076,14 @@ class TestClassifyMatrix: remove, because by the time a node reaches one of those states its StubSession marker has already been cleared by the transition that got it there. + + "IncompleteSession" heal-forward: EVERY start/fork cell (including the + former pure no-ops) now also removes "IncompleteSession" — see + _heal_forward() in session.py. A node may carry a stale IncompleteSession + marker from an earlier out-of-order session:end regardless of which + start/fork branch it reaches, so every start/fork transition strips it. + The end/None cells are unchanged — session:end still stamps + IncompleteSession when genuinely no type is known at end. """ CASES: list[tuple[str, str | None, bool, list[str], list[str]]] = [ @@ -2085,78 +2091,78 @@ class TestClassifyMatrix: # ------------------------------------------------------------------ # event=start # ------------------------------------------------------------------ - ("start", "ForkedSession", True, [], []), - ("start", "ForkedSession", False, [], []), - ("start", "SubSession", True, [], []), - ("start", "SubSession", False, [], []), - ("start", "RootSession", False, [], []), + ("start", "ForkedSession", True, [], ["IncompleteSession"]), + ("start", "ForkedSession", False, [], ["IncompleteSession"]), + ("start", "SubSession", True, [], ["IncompleteSession"]), + ("start", "SubSession", False, [], ["IncompleteSession"]), + ("start", "RootSession", False, [], ["IncompleteSession"]), ( "start", "RootSession", True, ["SubSession", "SST_EVENT"], - ["RootSession", "StubSession"], + ["RootSession", "StubSession", "IncompleteSession"], ), ( "start", None, True, ["Session", "SubSession", "SST_EVENT"], - ["StubSession"], + ["StubSession", "IncompleteSession"], ), ( "start", None, False, ["RootSession", "Session", "SST_EVENT"], - ["StubSession"], + ["StubSession", "IncompleteSession"], ), # ------------------------------------------------------------------ # event=fork # ------------------------------------------------------------------ - ("fork", "ForkedSession", True, [], []), - ("fork", "ForkedSession", False, [], []), + ("fork", "ForkedSession", True, [], ["IncompleteSession"]), + ("fork", "ForkedSession", False, [], ["IncompleteSession"]), ( "fork", "RootSession", True, ["ForkedSession", "SST_EVENT"], - ["RootSession", "StubSession"], + ["RootSession", "StubSession", "IncompleteSession"], ), ( "fork", "RootSession", False, ["ForkedSession", "SST_EVENT"], - ["RootSession", "StubSession"], + ["RootSession", "StubSession", "IncompleteSession"], ), ( "fork", "SubSession", True, ["ForkedSession", "SST_EVENT"], - ["SubSession", "StubSession"], + ["SubSession", "StubSession", "IncompleteSession"], ), ( "fork", "SubSession", False, ["ForkedSession", "SST_EVENT"], - ["SubSession", "StubSession"], + ["SubSession", "StubSession", "IncompleteSession"], ), ( "fork", None, True, ["Session", "ForkedSession", "SST_EVENT"], - ["StubSession"], + ["StubSession", "IncompleteSession"], ), ( "fork", None, False, ["Session", "ForkedSession", "SST_EVENT"], - ["StubSession"], + ["StubSession", "IncompleteSession"], ), # ------------------------------------------------------------------ # event=end @@ -2165,7 +2171,8 @@ class TestClassifyMatrix: # not a fabricated Root/Sub terminal. has_parent is irrelevant here — # the server never guesses; it marks and surfaces the health signal. # IncompleteSession is a confirmed (if incomplete) terminal, so - # StubSession is cleared here too. + # StubSession is cleared here too. UNCHANGED by heal-forward: end is + # not a start/fork transition. ("end", None, True, ["IncompleteSession", "SST_EVENT"], ["StubSession"]), ("end", None, False, ["IncompleteSession", "SST_EVENT"], ["StubSession"]), ("end", "RootSession", True, [], []), @@ -2241,7 +2248,7 @@ async def test_ensure_session_node_existing_branch_does_not_add_stub_session( "already-enriched", {"labels": ["Session", "RootSession"], "status": "running"}, ) - # _seen_sessions cache is empty, so this call hits Tier 2 (graph query) + # _seen_sessions cache is empty, so this call hits the graph query # and takes the "existing is not None" branch. await services.ensure_session_node("already-enriched", {}) @@ -2882,3 +2889,126 @@ async def test_current_type_returns_none_for_incomplete_session_node( "_current_type must ignore IncompleteSession and return None, " "so a late start/fork can still classify the session normally" ) + + # ----------------------------------------------------------------------- + # Heal-forward: out-of-order end -> fork/start race + # + # A forked sub-session's session:end can drain, in an independent queue, + # BEFORE its session:fork/session:start. classify() sees current_type=None + # at end -> stamps IncompleteSession. When the real fork/start is then + # processed, it must strip the stale marker (heal-forward), leaving the + # node with ONLY the real terminal label + # (out-of-order end-before-start/fork heal-forward). + # ----------------------------------------------------------------------- + + async def test_out_of_order_end_then_fork_heals_incomplete_session( + self, services: HookStateService + ) -> None: + """end processed BEFORE fork: stale IncompleteSession must be stripped + the moment the real session:fork arrives, leaving ForkedSession only. + """ + handler = SessionHandler(services) + + # session:end drains first (simulating the cross-queue race) — stamps + # IncompleteSession on the bare node. + await handler( + "session:end", + {"session_id": "s-race-fork", "timestamp": "2026-01-01T01:00:00Z"}, + ) + node = await services.graph.get_node("s-race-fork") + assert node is not None + assert "IncompleteSession" in node["labels"], ( + "Precondition: out-of-order end must stamp the stale marker" + ) + + # The real session:fork arrives late. + await handler( + "session:fork", + { + "session_id": "s-race-fork", + "parent_id": "p-race", + "timestamp": "2026-01-01T00:00:00Z", + }, + ) + + node = await services.graph.get_node("s-race-fork") + assert node is not None + labels = node["labels"] + assert "ForkedSession" in labels, "Real terminal must be assigned by fork" + assert "IncompleteSession" not in labels, ( + "Heal-forward: session:fork must strip the stale IncompleteSession " + "marker left by the out-of-order session:end" + ) + + async def test_out_of_order_end_then_start_heals_incomplete_session( + self, services: HookStateService + ) -> None: + """end processed BEFORE start (no parent): stale IncompleteSession must + be stripped the moment the real session:start arrives, leaving + RootSession only. + """ + handler = SessionHandler(services) + + await handler( + "session:end", + {"session_id": "s-race-root", "timestamp": "2026-01-01T01:00:00Z"}, + ) + node = await services.graph.get_node("s-race-root") + assert node is not None + assert "IncompleteSession" in node["labels"], ( + "Precondition: out-of-order end must stamp the stale marker" + ) + + await handler( + "session:start", + {"session_id": "s-race-root", "timestamp": "2026-01-01T00:00:00Z"}, + ) + + node = await services.graph.get_node("s-race-root") + assert node is not None + labels = node["labels"] + assert "RootSession" in labels, "Real terminal must be assigned by start" + assert "IncompleteSession" not in labels, ( + "Heal-forward: session:start must strip the stale IncompleteSession " + "marker left by the out-of-order session:end" + ) + + async def test_out_of_order_end_then_start_with_parent_heals_to_subsession( + self, services: HookStateService + ) -> None: + """end processed BEFORE start (with parent): stale IncompleteSession + must be stripped, leaving SubSession only. + """ + handler = SessionHandler(services) + + await handler( + "session:end", + { + "session_id": "s-race-sub", + "parent_id": "p-race-sub", + "timestamp": "2026-01-01T01:00:00Z", + }, + ) + node = await services.graph.get_node("s-race-sub") + assert node is not None + assert "IncompleteSession" in node["labels"], ( + "Precondition: out-of-order end must stamp the stale marker" + ) + + await handler( + "session:start", + { + "session_id": "s-race-sub", + "parent_id": "p-race-sub", + "timestamp": "2026-01-01T00:00:00Z", + }, + ) + + node = await services.graph.get_node("s-race-sub") + assert node is not None + labels = node["labels"] + assert "SubSession" in labels, "Real terminal must be assigned by start" + assert "IncompleteSession" not in labels, ( + "Heal-forward: session:start must strip the stale IncompleteSession " + "marker left by the out-of-order session:end" + ) diff --git a/tests/neo4j/test_agent_field_ordering_race.py b/tests/neo4j/test_agent_field_ordering_race.py new file mode 100644 index 00000000..16ae73c9 --- /dev/null +++ b/tests/neo4j/test_agent_field_ordering_race.py @@ -0,0 +1,142 @@ +"""E2E reproduction for issue #484: `agent` dropped on the child-first ordering. + +Runs against an ISOLATED, throwaway Neo4j container (the ``neo4j_container`` +fixture in ``tests/neo4j/conftest.py`` — random ports, ``remove=True``, torn +down after the session). NEVER touches the production/shared store. + +The race +-------- +A spawned sub-session's ``agent`` name arrives on the PARENT's +``delegate:agent_spawned`` event (top-level ``agent``). The CHILD's own +``session:start`` carries no top-level ``agent``. Both are processed through +``ensure_session_node`` (pipeline step 2). ``ensure_session_node``'s "node +already exists" (Tier-2) branch historically upserted only +``{labels, status, session_id}`` — so when the CHILD's ``session:start`` created +the node first, the PARENT's later ``delegate:agent_spawned`` (which DOES carry +``{"agent": ...}``) hit the Tier-2 branch and its ``agent`` was silently dropped, +leaving ``:Session.agent`` permanently empty. + +This test drives the real handlers across two independent ``Neo4jGraphStore`` +instances sharing one Neo4j (exactly the two-drainer condition), forcing the +CHILD-first ordering deterministically, then asserts ``agent`` is persisted. + +RED before the services.py fix, GREEN after. + +Run: uv run pytest tests/neo4j/test_agent_field_ordering_race.py -v -m neo4j +""" + +from __future__ import annotations + +import uuid +from typing import Any + +import pytest + +from context_intelligence_server.handlers.data_layer_2.session import SessionHandler +from context_intelligence_server.handlers.data_layer_3.delegation import ( + DelegationHandler, +) +from context_intelligence_server.neo4j_store import ( + Neo4jGraphStore, + ensure_neo4j_schema, +) +from context_intelligence_server.services import HookStateService + +pytestmark = pytest.mark.neo4j + + +async def _neo4j_agent(store: Neo4jGraphStore, node_id: str) -> str | None: + """Read the `agent` property of a node straight from Neo4j (not the buffer).""" + rows = await store.execute_query( + "MATCH (n) WHERE n.node_id = $id AND n.workspace = $workspace " + "RETURN n.agent AS agent", + {"id": node_id, "workspace": store.workspace}, + workspace="*", + ) + return rows[0]["agent"] if rows else None + + +def _ts(n: int = 0) -> str: + return f"2026-01-01T00:{n:02d}:00Z" + + +@pytest.mark.neo4j +class TestAgentFieldChildFirstOrdering: + """#484: agent must survive when the child's session:start lands first.""" + + async def test_child_start_before_parent_spawn_persists_agent( + self, neo4j_container: dict[str, Any] + ) -> None: + """CHILD session:start creates the node first (no agent); the PARENT's + later delegate:agent_spawned must still persist `agent` onto it.""" + auth = (neo4j_container["user"], neo4j_container["password"]) + bolt = neo4j_container["bolt_url"] + ws = f"test-agent-484-{uuid.uuid4().hex[:8]}" + + from neo4j import AsyncGraphDatabase + + driver = AsyncGraphDatabase.driver(bolt, auth=auth) + await ensure_neo4j_schema(driver) + await driver.close() + + parent_id = f"parent-{uuid.uuid4().hex[:8]}" + child_id = f"child-{uuid.uuid4().hex[:8]}" + tool_call_id = f"tc-{uuid.uuid4().hex[:8]}" + expected_agent = "foundation:git-ops" + + # --- CHILD's drainer resources --- + child_store = Neo4jGraphStore(uri=bolt, auth=auth, workspace=ws) + child_services = HookStateService(workspace=ws, graph_store=child_store) + session_handler_child = SessionHandler(child_services) + + # --- PARENT's drainer resources --- + parent_store = Neo4jGraphStore(uri=bolt, auth=auth, workspace=ws) + parent_services = HookStateService(workspace=ws, graph_store=parent_store) + await parent_services.ensure_session_node(parent_id, {}) + delegation_handler = DelegationHandler(parent_services) + + # Step 1 — CHILD's own session:start creates the sub-session node FIRST, + # carrying NO top-level agent (agent name is only nested in metadata). + await session_handler_child( + "session:start", + { + "session_id": child_id, + "parent_id": parent_id, + "timestamp": _ts(1), + "metadata": {"agent_name": expected_agent}, + }, + ) + await child_store.flush() + + # Precondition: node exists in Neo4j but has no agent yet. + assert await _neo4j_agent(child_store, child_id) is None + + # Step 2 — PARENT's delegate:agent_spawned arrives LATER, carrying the + # top-level agent. This is the only event that supplies `agent` to + # ensure_session_node, and it now hits the Tier-2 existing-node branch. + await delegation_handler( + "delegate:agent_spawned", + { + "session_id": parent_id, + "parent_session_id": parent_id, + "sub_session_id": child_id, + "agent": expected_agent, + "tool_call_id": tool_call_id, + "timestamp": _ts(0), + }, + ) + await parent_store.flush() + + # Assert — the sub-session node carries the agent end-to-end in Neo4j. + verify_store = Neo4jGraphStore(uri=bolt, auth=auth, workspace=ws) + try: + actual = await _neo4j_agent(verify_store, child_id) + assert actual == expected_agent, ( + f"#484 REPRODUCED: sub-session {child_id} has agent={actual!r}, " + f"expected {expected_agent!r}. The parent's delegate:agent_spawned " + f"agent value was dropped by ensure_session_node's Tier-2 branch." + ) + finally: + await verify_store.close() + await child_store.close() + await parent_store.close() diff --git a/tests/neo4j/test_incomplete_session_heal_forward.py b/tests/neo4j/test_incomplete_session_heal_forward.py new file mode 100644 index 00000000..065a63f7 --- /dev/null +++ b/tests/neo4j/test_incomplete_session_heal_forward.py @@ -0,0 +1,191 @@ +"""Neo4j integration proof for IncompleteSession heal-forward. + +Encodes the out-of-order race that produces the ~99% false-positive +IncompleteSession population: a forked sub-session's session:end drains +(independent per-session queue) BEFORE its session:fork/session:start. +SessionLabelStateMachine.classify() must strip the stale IncompleteSession +marker the moment the real start/fork is processed +(out-of-order end-before-start/fork). + +This closes a real coverage gap: no prior real-Neo4j test exercised +set_labels() with a non-empty remove_labels list. It runs the real +SessionHandler (not just the pure classify() unit) against a live Neo4j +container, flushes, and reads back with a fresh Cypher query (bypassing the +in-memory buffer) to prove the label is PHYSICALLY removed from the node. + +It also proves non-interaction with the terminal-label lattice +normalization (_LATTICE_NORMALIZATION in neo4j_store.py): removing +IncompleteSession must not disturb the RootSession/SubSession/ForkedSession +convergence guarantee, since IncompleteSession is not a member of +_TERMINAL_LABELS. + +Run: uv run pytest tests/neo4j/test_incomplete_session_heal_forward.py -v -m neo4j +""" + +from __future__ import annotations + +from typing import Any + +import pytest + +from context_intelligence_server.handlers.data_layer_2.session import SessionHandler + +pytestmark = pytest.mark.neo4j + + +async def _neo4j_labels(services: Any, node_id: str) -> list[str]: + """Read labels directly from Neo4j (bypasses the in-memory buffer).""" + rows = await services.graph.execute_query( + "MATCH (n) WHERE n.node_id = $id AND n.workspace = $workspace " + "RETURN labels(n) AS lbls", + {"id": node_id, "workspace": services.graph.workspace}, + workspace="*", + ) + return list(rows[0]["lbls"]) if rows else [] + + +@pytest.mark.neo4j +class TestIncompleteSessionHealForward: + """Out-of-order end -> fork/start must physically strip IncompleteSession.""" + + async def test_out_of_order_end_then_fork_strips_incomplete_session_in_neo4j( + self, neo4j_services: Any + ) -> None: + """end (bare, stamps IncompleteSession) -> flush -> fork -> flush: + the real Neo4j node must end with ForkedSession only, no + IncompleteSession, proving REMOVE n:IncompleteSession actually ran. + """ + handler = SessionHandler(neo4j_services) + session_id = "child-heal-fork-001" + + # session:end drains first (out-of-order race) -- stamps IncompleteSession + await handler( + "session:end", + {"session_id": session_id, "timestamp": "2026-01-01T10:00:00Z"}, + ) + await neo4j_services.graph.flush() + + mid_labels = await _neo4j_labels(neo4j_services, session_id) + assert "IncompleteSession" in mid_labels, ( + f"precondition: out-of-order end must stamp IncompleteSession in " + f"Neo4j; got {mid_labels}" + ) + + # The real session:fork arrives late, in its own flush cycle. + await handler( + "session:fork", + { + "session_id": session_id, + "parent_id": "parent-heal-fork-001", + "timestamp": "2026-01-01T09:59:59Z", + }, + ) + await neo4j_services.graph.flush() + + final_labels = await _neo4j_labels(neo4j_services, session_id) + assert "ForkedSession" in final_labels, ( + f"expected ForkedSession in {final_labels}" + ) + assert "IncompleteSession" not in final_labels, ( + f"heal-forward must physically REMOVE n:IncompleteSession at " + f"flush; still present in Neo4j: {final_labels}" + ) + + async def test_out_of_order_end_then_start_strips_incomplete_session_in_neo4j( + self, neo4j_services: Any + ) -> None: + """end (bare, stamps IncompleteSession) -> flush -> start (no parent) + -> flush: the real Neo4j node must end with RootSession only. + """ + handler = SessionHandler(neo4j_services) + session_id = "root-heal-start-001" + + await handler( + "session:end", + {"session_id": session_id, "timestamp": "2026-01-01T10:00:00Z"}, + ) + await neo4j_services.graph.flush() + + mid_labels = await _neo4j_labels(neo4j_services, session_id) + assert "IncompleteSession" in mid_labels, ( + f"precondition: out-of-order end must stamp IncompleteSession in " + f"Neo4j; got {mid_labels}" + ) + + await handler( + "session:start", + {"session_id": session_id, "timestamp": "2026-01-01T09:59:59Z"}, + ) + await neo4j_services.graph.flush() + + final_labels = await _neo4j_labels(neo4j_services, session_id) + assert "RootSession" in final_labels, f"expected RootSession in {final_labels}" + assert "IncompleteSession" not in final_labels, ( + f"heal-forward must physically REMOVE n:IncompleteSession at " + f"flush; still present in Neo4j: {final_labels}" + ) + + async def test_heal_forward_does_not_disturb_terminal_lattice_normalization( + self, neo4j_services: Any + ) -> None: + """Healing IncompleteSession must not interact with, or break, the + RootSession/SubSession/ForkedSession lattice-normalization guarantee. + + Drives: end (bare, stamps IncompleteSession) -> flush -> start WITH a + parent (assigns SubSession) -> flush -> fork (reclassifies to + ForkedSession, the lattice's specificity ordering) -> flush. The node + must converge to exactly ONE terminal label (ForkedSession) with + IncompleteSession and the stale SubSession both absent -- proving the + IncompleteSession REMOVE and the terminal-lattice REMOVE/SET both ran + correctly and did not clobber each other. + """ + handler = SessionHandler(neo4j_services) + session_id = "lattice-heal-001" + parent_id = "lattice-heal-parent-001" + + await handler( + "session:end", + {"session_id": session_id, "timestamp": "2026-01-01T10:00:00Z"}, + ) + await neo4j_services.graph.flush() + + await handler( + "session:start", + { + "session_id": session_id, + "parent_id": parent_id, + "timestamp": "2026-01-01T09:59:58Z", + }, + ) + await neo4j_services.graph.flush() + + mid_labels = await _neo4j_labels(neo4j_services, session_id) + assert "SubSession" in mid_labels, f"expected SubSession in {mid_labels}" + assert "IncompleteSession" not in mid_labels, ( + f"IncompleteSession must already be healed after start: {mid_labels}" + ) + + await handler( + "session:fork", + { + "session_id": session_id, + "parent_id": parent_id, + "timestamp": "2026-01-01T09:59:59Z", + }, + ) + await neo4j_services.graph.flush() + + final_labels = await _neo4j_labels(neo4j_services, session_id) + terminals = [ + lbl + for lbl in final_labels + if lbl in ("RootSession", "SubSession", "ForkedSession") + ] + assert terminals == ["ForkedSession"], ( + f"lattice must converge to exactly one terminal (ForkedSession); " + f"got {terminals} in {final_labels}" + ) + assert "IncompleteSession" not in final_labels, ( + f"IncompleteSession must remain healed through the lattice " + f"reclassification: {final_labels}" + ) diff --git a/tests/neo4j/test_working_dir_non_overwrite.py b/tests/neo4j/test_working_dir_non_overwrite.py new file mode 100644 index 00000000..1b1549e8 --- /dev/null +++ b/tests/neo4j/test_working_dir_non_overwrite.py @@ -0,0 +1,158 @@ +"""Live E2E tests: working_dir is never clobbered at the DB level. + +Root cause being guarded here: prior to this fix, the ONLY guarantee that an +already-populated ``working_dir`` is never overwritten lived in the Python layer +(``services.py``'s ``if data.get("working_dir") and not existing.get("working_dir")`` +populate-if-missing check). That check reads the buffered/graph node BEFORE the +write is issued, so it cannot protect against a cross-writer or replica race: a +second concurrent flush that read the node before the first write committed would +still see no working_dir, and its ``SET n += row.props`` would blindly overwrite +whatever the first writer just set. + +The fix adds a genuine DB-level guarantee: the Session-node MERGE in +``_write_batch`` excludes working_dir from the blind ``+=`` merge and instead +applies ``SET n.working_dir = coalesce(n.working_dir, row.working_dir)`` -- a +non-overwrite rule enforced by Neo4j itself, at the same MERGE lock hold, immune +to read-then-write races between writers. + +Requires Docker and the docker Python package. Skip-if-absent via the +``neo4j_container`` fixture in tests/neo4j/conftest.py. + +Run explicitly: + cd amplifier-context-intelligence + uv run pytest tests/neo4j/test_working_dir_non_overwrite.py -v -m neo4j +""" + +from __future__ import annotations + +from typing import Any + +import pytest +from context_intelligence_server.neo4j_store import Neo4jGraphStore +from neo4j import GraphDatabase + +pytestmark = pytest.mark.neo4j + + +async def _flush_session_node( + container: dict[str, Any], node_id: str, data: dict[str, Any] +) -> None: + """Drive a single Session node through the real flush path via a FRESH store. + + A fresh ``Neo4jGraphStore`` per call simulates independent writers (e.g. two + drainer workers, or a writer racing a replica) rather than two writes queued + through the same in-process buffer -- the scenario the Python-layer + populate-if-missing check cannot see. + """ + store = Neo4jGraphStore( + uri=container["bolt_url"], + auth=(container["user"], container["password"]), + workspace="test", + ) + try: + await store.upsert_node(node_id, {"labels": ["Session"], **data}) + await store.flush() + finally: + await store.close() + + +def _read_working_dir(container: dict[str, Any], node_id: str) -> Any: + driver = GraphDatabase.driver( + container["bolt_url"], auth=(container["user"], container["password"]) + ) + try: + with driver.session() as session: + rec = session.run( + "MATCH (n:Session {node_id: $nid, workspace: $ws}) " + "RETURN n.working_dir AS wd", + nid=node_id, + ws="test", + ).single() + assert rec is not None, f"Session node {node_id!r} was not written" + return rec["wd"] + finally: + driver.close() + + +async def test_existing_working_dir_survives_conflicting_later_write( + neo4j_container: dict[str, Any], +) -> None: + """DB-level guarantee: an already-set working_dir is never overwritten, + even by a second independent writer (simulating a cross-writer/replica race). + + RED (unfixed): ``SET n += row.props`` blindly overwrites -> working_dir + becomes "/y" -> this assertion fails. + GREEN (fixed): ``coalesce(n.working_dir, row.working_dir)`` keeps the + existing value -> working_dir stays "/x". + """ + node_id = "sess-wd-no-clobber-live" + + # Writer 1: establishes working_dir="/x". + await _flush_session_node( + neo4j_container, + node_id, + {"status": "running", "working_dir": "/x"}, + ) + assert _read_working_dir(neo4j_container, node_id) == "/x" + + # Writer 2 (independent store instance): tries to write a DIFFERENT + # working_dir for the SAME node -- must be rejected at the DB level. + await _flush_session_node( + neo4j_container, + node_id, + {"status": "running", "working_dir": "/y"}, + ) + + assert _read_working_dir(neo4j_container, node_id) == "/x", ( + "An already-populated working_dir must never be overwritten by a " + "later/concurrent writer -- DB-level coalesce guarantee failed" + ) + + +async def test_working_dir_fills_gap_at_db_level_when_previously_absent( + neo4j_container: dict[str, Any], +) -> None: + """DB-level populate-if-missing: a node with no working_dir gets filled in + by a later write that supplies one (coalesce(null, value) -> value). + + Mirrors the Python-layer guarantee (services.py) but proves it also holds + purely at the DB level, independent of the in-process buffer. + """ + node_id = "sess-wd-fill-gap-live" + + # Writer 1: no working_dir supplied. + await _flush_session_node(neo4j_container, node_id, {"status": "running"}) + assert _read_working_dir(neo4j_container, node_id) is None + + # Writer 2: supplies working_dir for the first time. + await _flush_session_node( + neo4j_container, + node_id, + {"status": "running", "working_dir": "/first-value"}, + ) + + assert _read_working_dir(neo4j_container, node_id) == "/first-value", ( + "working_dir must be filled in at the DB level once a writer supplies " + "a value for a node that previously had none" + ) + + +async def test_working_dir_absent_write_does_not_clear_existing_value( + neo4j_container: dict[str, Any], +) -> None: + """A later write that omits working_dir entirely must not null out an + already-set value (coalesce(n.working_dir, null) -> unchanged). + """ + node_id = "sess-wd-absent-no-clear-live" + + await _flush_session_node( + neo4j_container, node_id, {"status": "running", "working_dir": "/keep-me"} + ) + assert _read_working_dir(neo4j_container, node_id) == "/keep-me" + + # Second write carries no working_dir key at all (e.g. a touch/heartbeat event). + await _flush_session_node(neo4j_container, node_id, {"status": "still-running"}) + + assert _read_working_dir(neo4j_container, node_id) == "/keep-me", ( + "A write that omits working_dir must never null out an already-set value" + ) diff --git a/tests/test_services.py b/tests/test_services.py index 1a78d01a..89d36199 100644 --- a/tests/test_services.py +++ b/tests/test_services.py @@ -791,3 +791,130 @@ def _fail_on_dl3(obj, **kw): assert svc.data_layer_2.iteration_count == 3 assert svc.data_layer_2.orch_run_seq == 5 assert svc.data_layer_2.active_orch_run_id == "OLD" + + +# --------------------------------------------------------------------------- +# ensure_session_node populate-if-missing: working_dir + agent +# +# The `working_dir` (bundle-hook envelope field) and `agent` (parent's +# delegate:agent_spawned) values can each arrive on a LATER event than the one +# that first creates the Session node. On the existing-node branch they must be +# backfilled when the node still lacks them, and NEVER overwrite an already-set +# value. +# +# Two HookStateService instances share ONE GraphState to model the real +# two-writer condition (each worker has its own cold _seen_sessions cache, so +# the second writer genuinely reaches the graph-query existing-node branch +# rather than short-circuiting on the warm-cache fast path). +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +class TestEnsureSessionNodeWorkingDir: + """working_dir survives a later-event delivery and is never clobbered.""" + + async def test_existing_node_populates_working_dir_from_later_event(self) -> None: + """A node created WITHOUT working_dir gets it backfilled by a later + event carrying one.""" + graph = GraphState() + svc_first = HookStateService(graph_store=graph) + svc_later = HookStateService(graph_store=graph) + + await svc_first.ensure_session_node( + "sess-wd", {"timestamp": "2026-01-01T00:00:00Z"} + ) + node = await graph.get_node("sess-wd") + assert node is not None + assert node.get("working_dir") is None # precondition + + await svc_later.ensure_session_node( + "sess-wd", {"working_dir": "/home/u/project"} + ) + node = await graph.get_node("sess-wd") + assert node is not None + assert node.get("working_dir") == "/home/u/project" + + async def test_existing_node_does_not_clobber_working_dir(self) -> None: + """A later working_dir-less (or different) event must NOT overwrite an + already-set working_dir.""" + graph = GraphState() + svc_first = HookStateService(graph_store=graph) + svc_later = HookStateService(graph_store=graph) + + await svc_first.ensure_session_node( + "sess-wd2", {"working_dir": "/home/u/original"} + ) + # A later event carrying a DIFFERENT working_dir must not win. + await svc_later.ensure_session_node("sess-wd2", {"working_dir": "/tmp/other"}) + node = await graph.get_node("sess-wd2") + assert node is not None + assert node.get("working_dir") == "/home/u/original" + + async def test_new_node_lifts_working_dir(self) -> None: + """A first event carrying working_dir writes it straight onto the node.""" + graph = GraphState() + svc = HookStateService(graph_store=graph) + await svc.ensure_session_node( + "sess-wd3", + {"timestamp": "2026-01-01T00:00:00Z", "working_dir": "/home/u/fresh"}, + ) + node = await graph.get_node("sess-wd3") + assert node is not None + assert node.get("working_dir") == "/home/u/fresh" + + +# --------------------------------------------------------------------------- +# Issue #484: ensure_session_node must not drop `agent` on the existing-node +# branch. The `agent` value for a spawned sub-session arrives on the parent's +# delegate:agent_spawned event ({"agent": ...}), while the child's own +# session:start (no top-level agent) can create the node first. When the parent +# event then hits the existing-node branch, `agent` must still be persisted +# (same populate-if-missing rule already applied to working_dir). +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +class TestEnsureSessionNodeAgentField: + """#484 regression: the agent name survives the child-first ordering race.""" + + async def test_existing_node_persists_agent_from_later_event(self) -> None: + """Child creates the node WITHOUT agent; a later writer carrying + {"agent": ...} (the parent's delegate:agent_spawned) must persist it.""" + graph = GraphState() + svc_child = HookStateService(graph_store=graph) + svc_parent = HookStateService(graph_store=graph) + + # Child's session:start reaches ensure_session_node first — no top-level agent. + await svc_child.ensure_session_node( + "child-484", {"timestamp": "2026-01-01T00:00:00Z"} + ) + node = await graph.get_node("child-484") + assert node is not None + assert node.get("agent") is None # precondition + + # Parent's delegate:agent_spawned arrives later, carrying the agent name. + # Second writer's cache is cold, so this reaches the existing-node branch. + await svc_parent.ensure_session_node( + "child-484", {"agent": "foundation:git-ops"} + ) + node = await graph.get_node("child-484") + assert node is not None + assert node.get("agent") == "foundation:git-ops" + + async def test_existing_node_does_not_clobber_agent(self) -> None: + """A later agent-less writer must NOT wipe an already-set agent.""" + graph = GraphState() + svc_parent = HookStateService(graph_store=graph) + svc_child = HookStateService(graph_store=graph) + + # Parent creates the node first, with the agent set. + await svc_parent.ensure_session_node( + "child-484b", {"agent": "foundation:git-ops"} + ) + # Child's later agent-less call hits the existing branch and must not clobber. + await svc_child.ensure_session_node( + "child-484b", {"timestamp": "2026-01-01T00:00:00Z"} + ) + node = await graph.get_node("child-484b") + assert node is not None + assert node.get("agent") == "foundation:git-ops" diff --git a/uv.lock b/uv.lock index be93615c..439c0b41 100644 --- a/uv.lock +++ b/uv.lock @@ -233,7 +233,7 @@ wheels = [ [[package]] name = "context-intelligence-server" -version = "7.0.0" +version = "7.1.0" source = { editable = "." } dependencies = [ { name = "aiofiles" },