@pentatonic-ai/ai-agent-sdk 0.7.13 → 0.8.1

package/packages/memory-engine/tests/test_channel_stat_reader.py

@@ -0,0 +1,437 @@
+"""Integration tests for the /aggregate-internal reader fast path.
+
+Sister file to ``test_channel_stat_rollups.py``: where that file
+covers the *writer* Cypher (ChannelStat nodes are maintained on every
+store), this file covers the *reader* — proves that
+``aggregate_internal`` actually reads from the denormalisation when
+the conditions are right, falls through to the edge walk when they
+aren't, and produces the same response shape either way.
+
+The PR review for #33 flagged that the writer was thoroughly tested
+but the reader's fast path only had indirect coverage via the
+endpoint integration test in ``test_aggregate.py``. This file closes
+that gap by importing the proxy module and calling ``aggregate_internal``
+directly, so a regression in the fast-path branch can't be masked by
+the silent fall-through to the edge walk.
+
+Gated on NEO4J_TEST_URI + NEO4J_TEST_PASSWORD; skip cleanly when
+those env vars are absent so unit-only test runs stay fast.
+
+Run:
+
+    cd packages/memory-engine
+    NEO4J_TEST_URI=bolt://localhost:17687 \\
+    NEO4J_TEST_PASSWORD=testpassword \\
+    .venv/bin/python -m pytest tests/test_channel_stat_reader.py -v
+"""
+from __future__ import annotations
+
+import asyncio
+import importlib.util
+import os
+import sys
+import uuid
+from pathlib import Path
+
+import pytest
+
+
+_NEO4J_URI = os.environ.get("NEO4J_TEST_URI")
+_NEO4J_USER = os.environ.get("NEO4J_TEST_USER", "neo4j")
+_NEO4J_PASSWORD = os.environ.get("NEO4J_TEST_PASSWORD")
+
+_skip_no_neo4j = pytest.mark.skipif(
+    not (_NEO4J_URI and _NEO4J_PASSWORD),
+    reason="set NEO4J_TEST_URI + NEO4J_TEST_PASSWORD to run integration tests",
+)
+
+
+ENGINE_ROOT = Path(__file__).resolve().parent.parent / "engine" / "services" / "l2"
+sys.path.insert(0, str(ENGINE_ROOT))
+
+
+@pytest.fixture(scope="module")
+def proxy_module():
+    """Load l2-hybridrag-proxy as a module and point it at the test
+    Neo4j. The module reads NEO4J_URI/NEO4J_AUTH as module-level
+    constants at import time; we override them after load so the
+    request handler dials the test instance instead of the default
+    localhost:7687.
+
+    Skip cleanly if fastapi/neo4j aren't importable (matches the
+    pattern in test_aggregate.py)."""
+    spec = importlib.util.spec_from_file_location(
+        "l2_proxy_module",
+        ENGINE_ROOT / "l2-hybridrag-proxy.py",
+    )
+    assert spec and spec.loader
+    try:
+        mod = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(mod)
+    except ImportError:
+        pytest.skip("l2 proxy deps unavailable in this venv (fine for unit-only runs)")
+    mod.NEO4J_URI = _NEO4J_URI
+    mod.NEO4J_AUTH = (_NEO4J_USER, _NEO4J_PASSWORD)
+    return mod
+
+
+@pytest.fixture
+def neo4j_driver():
+    """Per-test driver + cleanup. Two arenas so isolation tests can
+    run side by side without trampling each other."""
+    from neo4j import GraphDatabase
+
+    driver = GraphDatabase.driver(_NEO4J_URI, auth=(_NEO4J_USER, _NEO4J_PASSWORD))
+    arenas = [f"rdr_a_{uuid.uuid4().hex[:8]}", f"rdr_b_{uuid.uuid4().hex[:8]}"]
+    yield driver, arenas
+    with driver.session() as session:
+        for arena in arenas:
+            session.run(
+                "MATCH (n) WHERE n.arena = $arena DETACH DELETE n",
+                arena=arena,
+            )
+    driver.close()
+
+
+def _write_with_stats(
+    session,
+    arena: str,
+    cid: str,
+    email: str,
+    channel: str,
+    direction: str,
+    occurred_at: str,
+) -> None:
+    """Mirror the writer Cypher in /index-internal-batch's email path,
+    including the ChannelStat rollup. Keeping this inline rather than
+    importing from test_channel_stat_rollups.py so this file can be
+    read top-to-bottom without cross-file detective work."""
+    now_iso = "2026-05-11T00:00:00Z"
+    session.run(
+        """
+        MERGE (c:Chunk {id: $cid})
+        SET c.text = 't', c.path = 'p', c.arena = $arena,
+            c.created_at = $now
+        """,
+        cid=cid, arena=arena, now=now_iso,
+    )
+    session.run(
+        """
+        MERGE (p:Entity:Person {arena: $arena, email: $email})
+        ON CREATE SET p.created_at = $now
+        WITH p
+        MATCH (c:Chunk {arena: $arena, id: $cid})
+        MERGE (p)-[r:COMMUNICATED]->(c)
+        ON CREATE SET r.channel = $channel,
+                      r.direction = $direction,
+                      r.occurred_at = $occurred_at,
+                      r.weight = 1.0,
+                      r._counted = false
+        WITH p, r
+        FOREACH (_ IN CASE WHEN r._counted = false THEN [1] ELSE [] END |
+          MERGE (s:ChannelStat {arena: $arena, person_email: $email, channel: $channel})
+          ON CREATE SET s.count = 0,
+                        s.inbound = 0,
+                        s.outbound = 0,
+                        s.first_seen = $occurred_at,
+                        s.last_seen = $occurred_at,
+                        s.created_at = $now
+          SET s.count = s.count + 1,
+              s.inbound = s.inbound + (CASE WHEN $direction = 'inbound' THEN 1 ELSE 0 END),
+              s.outbound = s.outbound + (CASE WHEN $direction = 'outbound' THEN 1 ELSE 0 END),
+              s.first_seen = CASE
+                WHEN $occurred_at < coalesce(s.first_seen, $occurred_at)
+                  THEN $occurred_at
+                ELSE s.first_seen END,
+              s.last_seen = CASE
+                WHEN $occurred_at > coalesce(s.last_seen, '')
+                  THEN $occurred_at
+                ELSE s.last_seen END,
+              s.updated_at = $now
+          MERGE (p)-[:HAS_STAT]->(s)
+          SET r._counted = true
+        )
+        """,
+        arena=arena, email=email, cid=cid,
+        channel=channel, direction=direction,
+        occurred_at=occurred_at, now=now_iso,
+    )
+
+
+def _write_edges_only(
+    session,
+    arena: str,
+    cid: str,
+    email: str,
+    channel: str,
+    direction: str,
+    occurred_at: str,
+) -> None:
+    """Write Person + COMMUNICATED + Chunk *without* a ChannelStat.
+    Simulates pre-rollout data: older tenants whose ingest never went
+    through the new writer. Reader fast path must fall through to the
+    edge walk and still return correct results for these rows."""
+    now_iso = "2026-05-11T00:00:00Z"
+    session.run(
+        """
+        MERGE (c:Chunk {id: $cid})
+        SET c.text = 't', c.path = 'p', c.arena = $arena,
+            c.created_at = $now
+        MERGE (p:Entity:Person {arena: $arena, email: $email})
+        ON CREATE SET p.created_at = $now
+        MERGE (p)-[r:COMMUNICATED]->(c)
+        ON CREATE SET r.channel = $channel,
+                      r.direction = $direction,
+                      r.occurred_at = $occurred_at,
+                      r.weight = 1.0
+        """,
+        cid=cid, arena=arena, email=email,
+        channel=channel, direction=direction,
+        occurred_at=occurred_at, now=now_iso,
+    )
+
+
+def _ensure_indexes(session) -> None:
+    """Run the same index/constraint statements the writer runs at the
+    top of /index-internal-batch. Tests that exercise the constraint
+    behaviour need it present; the writer fixture in
+    test_channel_stat_rollups.py doesn't fire this path."""
+    session.run(
+        "CREATE INDEX channelstat_arena_email IF NOT EXISTS "
+        "FOR (s:ChannelStat) ON (s.arena, s.person_email)"
+    )
+    session.run(
+        "CREATE CONSTRAINT channelstat_unique IF NOT EXISTS "
+        "FOR (s:ChannelStat) REQUIRE (s.arena, s.person_email, s.channel) IS UNIQUE"
+    )
+
+
+def _call_aggregate(proxy_module, **kwargs):
+    """Invoke aggregate_internal as the FastAPI route would, but
+    without going through HTTP. The handler is async; this helper
+    wraps the asyncio.run boilerplate so individual tests stay terse."""
+    req = proxy_module.AggregateInternalRequest(**kwargs)
+    return asyncio.run(proxy_module.aggregate_internal(req))
+
+
+# ---------------------------------------------------------------------------
+# UNIQUE constraint.
+# ---------------------------------------------------------------------------
+
+
+@_skip_no_neo4j
+def test_unique_constraint_present_after_index_setup(neo4j_driver) -> None:
+    """Pin the constraint exists. Without it, concurrent writers can
+    create rival ChannelStat nodes for the same (arena, email,
+    channel) tuple — MERGE doesn't lock without a constraint."""
+    driver, _ = neo4j_driver
+    with driver.session() as session:
+        _ensure_indexes(session)
+        names = {
+            rec["name"]
+            for rec in session.run("SHOW CONSTRAINTS YIELD name")
+        }
+        assert "channelstat_unique" in names
+
+
+@_skip_no_neo4j
+def test_unique_constraint_rejects_duplicate_channelstat(neo4j_driver) -> None:
+    """End-to-end: with the constraint in place, an attempt to CREATE
+    a second node for the same key fails. Belt-and-braces against the
+    MERGE pattern ever falling back to CREATE under contention."""
+    from neo4j.exceptions import ConstraintError
+
+    driver, (arena, _) = neo4j_driver
+    with driver.session() as session:
+        _ensure_indexes(session)
+        session.run(
+            "CREATE (s:ChannelStat {arena: $arena, person_email: $email, "
+            "channel: $channel, count: 1})",
+            arena=arena, email="alex@x.io", channel="email",
+        )
+        with pytest.raises(ConstraintError):
+            session.run(
+                "CREATE (s:ChannelStat {arena: $arena, person_email: $email, "
+                "channel: $channel, count: 1})",
+                arena=arena, email="alex@x.io", channel="email",
+            )
+
+
+# ---------------------------------------------------------------------------
+# Fast-path reader behaviour.
+# ---------------------------------------------------------------------------
+
+
+@_skip_no_neo4j
+def test_fast_path_returns_per_channel_buckets_from_stats(
+    neo4j_driver, proxy_module
+) -> None:
+    """When ChannelStats exist for the contact, /aggregate-internal
+    should return one bucket per channel, ordered by count desc."""
+    driver, (arena, _) = neo4j_driver
+    email = "alex@x.io"
+    with driver.session() as session:
+        _ensure_indexes(session)
+        # 3 emails (2 in, 1 out), 1 slack (in).
+        _write_with_stats(session, arena, "c1", email, "email", "inbound", "2026-05-09T09:00:00Z")
+        _write_with_stats(session, arena, "c2", email, "email", "outbound", "2026-05-09T10:00:00Z")
+        _write_with_stats(session, arena, "c3", email, "email", "inbound", "2026-05-09T11:00:00Z")
+        _write_with_stats(session, arena, "c4", email, "slack", "inbound", "2026-05-09T08:00:00Z")
+
+    out = _call_aggregate(
+        proxy_module,
+        arena=arena, contact_email=email, group_by=["channel"],
+    )
+    assert out.total == 4
+    assert out.last_seen == "2026-05-09T11:00:00Z"
+    assert len(out.buckets) == 2
+    # Busiest first.
+    email_bucket = out.buckets[0]
+    slack_bucket = out.buckets[1]
+    assert email_bucket.keys == {"channel": "email"}
+    assert email_bucket.count == 3
+    assert email_bucket.inbound == 2
+    assert email_bucket.outbound == 1
+    assert slack_bucket.keys == {"channel": "slack"}
+    assert slack_bucket.count == 1
+    assert slack_bucket.inbound == 1
+
+
+@_skip_no_neo4j
+def test_fast_path_returns_single_bucket_when_group_by_empty(
+    neo4j_driver, proxy_module
+) -> None:
+    """Empty group_by collapses ChannelStat rows into one global
+    bucket — totals summed across channels, inbound/outbound likewise,
+    last_seen = max across all channels, first_seen = min."""
+    driver, (arena, _) = neo4j_driver
+    email = "alex@x.io"
+    with driver.session() as session:
+        _ensure_indexes(session)
+        _write_with_stats(session, arena, "c1", email, "email", "inbound", "2026-05-08T09:00:00Z")
+        _write_with_stats(session, arena, "c2", email, "email", "outbound", "2026-05-09T10:00:00Z")
+        _write_with_stats(session, arena, "c3", email, "slack", "inbound", "2026-05-07T15:00:00Z")
+
+    out = _call_aggregate(
+        proxy_module,
+        arena=arena, contact_email=email, group_by=[],
+    )
+    assert out.total == 3
+    assert len(out.buckets) == 1
+    bucket = out.buckets[0]
+    assert bucket.keys == {}
+    assert bucket.count == 3
+    assert bucket.inbound == 2
+    assert bucket.outbound == 1
+    # Time bounds span the full range across channels.
+    assert bucket.first_seen == "2026-05-07T15:00:00Z"
+    assert bucket.last_seen == "2026-05-09T10:00:00Z"
+
+
+@_skip_no_neo4j
+def test_fast_path_falls_through_to_edge_walk_when_stats_absent(
+    neo4j_driver, proxy_module
+) -> None:
+    """The forward-only optimisation: pre-rollout data has
+    COMMUNICATED edges but no ChannelStat nodes. The fast-path check
+    must fall through silently and the edge-walk path must still
+    return the correct buckets — this is the contract that means no
+    migration is needed."""
+    driver, (arena, _) = neo4j_driver
+    email = "legacy@example.com"
+    with driver.session() as session:
+        _ensure_indexes(session)
+        # Edges only — no ChannelStat written, simulating an older
+        # tenant whose data was ingested before this PR.
+        _write_edges_only(session, arena, "c1", email, "email", "inbound", "2026-05-09T09:00:00Z")
+        _write_edges_only(session, arena, "c2", email, "email", "outbound", "2026-05-09T10:00:00Z")
+        # Sanity: confirm no ChannelStat exists for this contact.
+        rows = list(session.run(
+            "MATCH (s:ChannelStat {arena: $arena, person_email: $email}) RETURN s",
+            arena=arena, email=email,
+        ))
+        assert rows == []
+
+    out = _call_aggregate(
+        proxy_module,
+        arena=arena, contact_email=email, group_by=["channel"],
+    )
+    # Edge-walk path took over and returned the same shape.
+    assert out.total == 2
+    assert len(out.buckets) == 1
+    assert out.buckets[0].keys == {"channel": "email"}
+    assert out.buckets[0].inbound == 1
+    assert out.buckets[0].outbound == 1
+
+
+@_skip_no_neo4j
+def test_fast_path_arena_isolated(neo4j_driver, proxy_module) -> None:
+    """A's aggregate must never reflect B's stats even when both
+    arenas have the same email — the channelstat_arena_email index is
+    keyed on arena first to enforce this. Companion to the writer-
+    side isolation test in test_channel_stat_rollups.py."""
+    driver, (arena_a, arena_b) = neo4j_driver
+    email = "shared@example.com"
+    with driver.session() as session:
+        _ensure_indexes(session)
+        _write_with_stats(session, arena_a, "ca1", email, "email", "inbound", "2026-05-09T09:00:00Z")
+        _write_with_stats(session, arena_b, "cb1", email, "email", "inbound", "2026-05-09T10:00:00Z")
+        _write_with_stats(session, arena_b, "cb2", email, "slack", "outbound", "2026-05-09T11:00:00Z")
+
+    out_a = _call_aggregate(
+        proxy_module, arena=arena_a, contact_email=email, group_by=["channel"],
+    )
+    out_b = _call_aggregate(
+        proxy_module, arena=arena_b, contact_email=email, group_by=["channel"],
+    )
+    assert out_a.total == 1
+    assert len(out_a.buckets) == 1
+    assert out_a.buckets[0].keys == {"channel": "email"}
+    assert out_b.total == 2
+    assert len(out_b.buckets) == 2
+
+
+@_skip_no_neo4j
+def test_fast_path_and_edge_walk_produce_equivalent_totals(
+    neo4j_driver, proxy_module
+) -> None:
+    """Equivalence: for data that has BOTH stats and edges (the
+    normal post-rollout case), the fast-path response must match what
+    the edge walk would have returned. Caught a class of bug where
+    the writer's rollup drifts from the edges — surface it via
+    response divergence rather than waiting for users to notice
+    relationships-UI inconsistency."""
+    driver, (arena, _) = neo4j_driver
+    email = "alex@x.io"
+    with driver.session() as session:
+        _ensure_indexes(session)
+        _write_with_stats(session, arena, "c1", email, "email", "inbound", "2026-05-08T09:00:00Z")
+        _write_with_stats(session, arena, "c2", email, "email", "outbound", "2026-05-09T10:00:00Z")
+        _write_with_stats(session, arena, "c3", email, "slack", "inbound", "2026-05-07T15:00:00Z")
+
+        # Compute the edge-walk answer by hand from the COMMUNICATED
+        # rows. This is the ground truth the rollup should reflect.
+        edge_rows = list(session.run(
+            "MATCH (p:Person {arena: $arena, email: $email})-[r:COMMUNICATED]->(:Chunk {arena: $arena})\n"
+            "WITH r.channel AS channel, r.direction AS direction\n"
+            "RETURN channel,\n"
+            "       count(*) AS count,\n"
+            "       sum(CASE WHEN direction = 'inbound' THEN 1 ELSE 0 END) AS inbound,\n"
+            "       sum(CASE WHEN direction = 'outbound' THEN 1 ELSE 0 END) AS outbound\n"
+            "ORDER BY count DESC",
+            arena=arena, email=email,
+        ))
+        ground_truth = {
+            rec["channel"]: (int(rec["count"]), int(rec["inbound"]), int(rec["outbound"]))
+            for rec in edge_rows
+        }
+
+    out = _call_aggregate(
+        proxy_module, arena=arena, contact_email=email, group_by=["channel"],
+    )
+    fast_path = {
+        b.keys["channel"]: (b.count, b.inbound, b.outbound)
+        for b in out.buckets
+    }
+    assert fast_path == ground_truth
+    assert out.total == sum(c for c, _, _ in ground_truth.values())