@pentatonic-ai/ai-agent-sdk 0.8.0 → 0.8.2

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

@@ -0,0 +1,412 @@
+"""Integration tests for L3 arena isolation + typed-entity writes.
+
+Two flavours:
+
+  - Neo4j-backed integration tests (run when ``NEO4J_TEST_URI`` and
+    ``NEO4J_TEST_PASSWORD`` env vars are set). These spin up the actual
+    writer logic and verify cross-arena isolation against a live
+    Neo4j. Skip cleanly when env is absent so the unit-test job stays
+    Neo4j-free.
+
+  - Pure-unit tests using a stub session that records every Cypher
+    call. Fast, hermetic, validate the structural invariants we care
+    about: each typed-entity write carries arena, COMMUNICATED edges
+    carry channel + direction, etc.
+
+Run:
+
+    cd packages/memory-engine
+    .venv/bin/python -m pytest tests/test_l3_arena_isolation.py -v
+
+Run with Neo4j:
+
+    NEO4J_TEST_URI=bolt://localhost:7687 \\
+    NEO4J_TEST_PASSWORD=test \\
+    .venv/bin/python -m pytest tests/test_l3_arena_isolation.py -v
+"""
+from __future__ import annotations
+
+import os
+import sys
+import uuid
+from pathlib import Path
+
+import pytest
+
+
+# Make the engine module importable without packaging it.
+ENGINE_ROOT = Path(__file__).resolve().parent.parent / "engine" / "services" / "l2"
+sys.path.insert(0, str(ENGINE_ROOT))
+
+
+# ---------------------------------------------------------------------------
+# Stub session for unit tests — records calls without hitting Neo4j.
+# ---------------------------------------------------------------------------
+
+
+class _Recorder:
+    """Stand-in for a Neo4j session.run that records every call.
+
+    Just enough surface for the writer block in the engine to think it's
+    talking to Neo4j: ``run(cypher, **params)`` returns an object whose
+    ``single()`` / iteration / ``data()`` calls all return empty.
+    """
+
+    def __init__(self) -> None:
+        self.calls: list[tuple[str, dict]] = []
+
+    def run(self, cypher: str, **params) -> "_Recorder":
+        self.calls.append((cypher, params))
+        return self
+
+    def single(self) -> dict:
+        return {}
+
+    def __iter__(self):
+        return iter([])
+
+
+def _has_arena_in_pattern(cypher: str, label: str, var: str = "") -> bool:
+    """True if every occurrence of (…label…) in the cypher names arena.
+
+    Parses each node pattern as ``(var:Label1:Label2:… {props})``,
+    finds the ones whose label list contains the target ``label``, and
+    asserts each one is arena-scoped (either via ``arena:`` in the
+    property bag or via ``<var>.arena`` somewhere in the same block).
+
+    Parameters
+    ----------
+    cypher : the Cypher block to inspect.
+    label  : target label (e.g. ``"Entity"``, ``"Chunk"``).
+    var    : optional variable filter — when set, only patterns that
+             bind this variable are checked. Useful when a block has
+             multiple patterns of the same label and we want to assert
+             one specific one.
+    """
+    import re
+
+    pattern = re.compile(
+        r"\("
+        r"\s*(?P<v>[A-Za-z_]\w*)?"
+        r"\s*(?P<labels>(?::[A-Za-z_]\w*)+)"  # one or more labels
+        r"\s*(?P<props>\{[^{}]*\})?"
+        r"\s*\)",
+        re.MULTILINE,
+    )
+    target = f":{label}"
+    found_any = False
+    for m in pattern.finditer(cypher):
+        labels = m.group("labels") or ""
+        # Require an exact-label match so :Entity matches `:Entity` and
+        # `:Entity:Concept` but not `:Entitlement` (re.search alone
+        # would treat the latter as a hit).
+        label_tokens = re.findall(r":([A-Za-z_]\w*)", labels)
+        if label not in label_tokens:
+            continue
+        v = m.group("v") or ""
+        if var and v != var:
+            continue
+        found_any = True
+        props = m.group("props") or ""
+        if "arena" in props:
+            continue
+        if v and re.search(rf"\b{re.escape(v)}\.arena\b", cypher):
+            continue
+        return False
+    return found_any
+
+
+# ---------------------------------------------------------------------------
+# Unit tests — exercise the writer block via the stub session by
+# calling its Cypher directly.
+# ---------------------------------------------------------------------------
+
+
+def test_writer_concept_path_carries_arena() -> None:
+    """The Concept-extraction Cypher always names arena on Entity + Chunk."""
+    rec = _Recorder()
+    arena = "acme"
+    rec.run(
+        """
+        MERGE (e:Entity:Concept {arena: $arena, name: $name})
+        ON CREATE SET e.type = 'Concept',
+                      e.created_at = $now,
+                      e.weight = 1.0
+        WITH e
+        MATCH (c:Chunk {arena: $arena, id: $cid})
+        MERGE (e)-[r:MENTIONS]->(c)
+        ON CREATE SET r.weight = 1.0, r.created_at = $now
+        ON MATCH SET r.weight = coalesce(r.weight, 1.0) + 0.1
+        """,
+        arena=arena, name="Pricing", cid="chunk-1", now="t",
+    )
+    cypher, params = rec.calls[-1]
+    assert _has_arena_in_pattern(cypher, "Entity")
+    assert _has_arena_in_pattern(cypher, "Chunk")
+    assert params["arena"] == arena
+
+
+def test_writer_person_email_carries_arena_email_and_communicated_edge() -> None:
+    """Metadata-driven Person email path tags arena + builds COMMUNICATED edge."""
+    rec = _Recorder()
+    arena = "acme"
+    rec.run(
+        """
+        MERGE (p:Entity:Person {arena: $arena, email: $email})
+        ON CREATE SET p.created_at = $now,
+                      p.first_seen = $occurred_at,
+                      p.last_seen = $occurred_at
+        ON MATCH SET p.last_seen = CASE
+              WHEN $occurred_at > coalesce(p.last_seen, '')
+                THEN $occurred_at
+              ELSE p.last_seen END
+        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
+        """,
+        arena=arena, email="alex@acme.com", cid="c-1",
+        channel="email", direction="inbound",
+        occurred_at="2026-05-08T00:00:00Z", now="t",
+    )
+    cypher, params = rec.calls[-1]
+    assert _has_arena_in_pattern(cypher, "Person")
+    assert _has_arena_in_pattern(cypher, "Chunk")
+    assert "COMMUNICATED" in cypher
+    assert params["channel"] == "email"
+    assert params["direction"] == "inbound"
+
+
+def test_co_occurs_edges_are_arena_scoped_on_both_endpoints() -> None:
+    """CO_OCCURS Cypher matches both endpoints in the same arena."""
+    rec = _Recorder()
+    rec.run(
+        """
+        MATCH (a:Entity:Concept {arena: $arena, name: $a})
+        MATCH (b:Entity:Concept {arena: $arena, name: $b})
+        MERGE (a)-[r:CO_OCCURS]->(b)
+        ON CREATE SET r.weight = 0.5, r.created_at = $now
+        ON MATCH SET r.weight = coalesce(r.weight, 0.5) + 0.05
+        """,
+        arena="acme", a="Pricing", b="Negotiation", now="t",
+    )
+    cypher, _ = rec.calls[-1]
+    # Both endpoints carry arena in the property bag.
+    assert cypher.count("arena: $arena") == 2
+
+
+def test_known_as_edge_links_email_and_name_within_arena() -> None:
+    """KNOWN_AS edge connects the (arena, name) and (arena, email) Person nodes."""
+    rec = _Recorder()
+    rec.run(
+        """
+        MATCH (n:Person {arena: $arena, name: $name})
+        MATCH (e:Person {arena: $arena, email: $email})
+        MERGE (n)-[:KNOWN_AS]->(e)
+        """,
+        arena="acme", name="Alex Tong", email="alex@acme.com",
+    )
+    cypher, _ = rec.calls[-1]
+    assert _has_arena_in_pattern(cypher, "Person")
+    assert "KNOWN_AS" in cypher
+
+
+# ---------------------------------------------------------------------------
+# Self-test: the helper above flags the bug the lint missed in v1.
+# ---------------------------------------------------------------------------
+
+
+def test_helper_flags_unscoped_person_pattern() -> None:
+    bad_cypher = """
+        MERGE (p:Entity:Person {email: $email})
+        ON CREATE SET p.created_at = $now
+        MATCH (c:Chunk {arena: $arena, id: $cid})
+        MERGE (p)-[:MENTIONS]->(c)
+    """
+    assert not _has_arena_in_pattern(bad_cypher, "Person"), \
+        "helper must flag the unscoped Person pattern"
+    assert _has_arena_in_pattern(bad_cypher, "Chunk"), \
+        "Chunk pattern with arena in property bag should pass"
+
+
+# ---------------------------------------------------------------------------
+# Neo4j-backed integration tests — only run when env is set.
+# ---------------------------------------------------------------------------
+
+
+_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",
+)
+
+
+@pytest.fixture
+def neo4j_driver():
+    """Open a Neo4j driver and clean test data on teardown.
+
+    Uses a randomised arena pair so concurrent test runs don't trample
+    each other; tears down by deleting nodes scoped to those arenas
+    (never a global wipe — this fixture must be safe against a
+    populated dev database).
+    """
+    from neo4j import GraphDatabase  # local import keeps unit-only runs neo4j-free
+
+    driver = GraphDatabase.driver(_NEO4J_URI, auth=(_NEO4J_USER, _NEO4J_PASSWORD))
+    arenas = [f"test_a_{uuid.uuid4().hex[:8]}", f"test_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()
+
+
+@_skip_no_neo4j
+def test_two_arenas_get_distinct_person_nodes_for_same_email(neo4j_driver) -> None:
+    """Same contact_email in two arenas → two :Person nodes, not one."""
+    driver, (arena_a, arena_b) = neo4j_driver
+    email = "shared@example.com"
+    chunk_id = lambda a: f"c_{a}"
+
+    with driver.session() as session:
+        for arena in (arena_a, arena_b):
+            # Materialise the chunk that the Person attaches to.
+            session.run(
+                """
+                MERGE (c:Chunk {id: $cid})
+                SET c.arena = $arena, c.text = 'test', c.path = 'test',
+                    c.created_at = '2026-05-08T00:00:00Z'
+                """,
+                cid=chunk_id(arena), arena=arena,
+            )
+            # Apply the same writer block as the engine would.
+            session.run(
+                """
+                MERGE (p:Entity:Person {arena: $arena, email: $email})
+                ON CREATE SET p.created_at = $now,
+                              p.first_seen = $occurred_at,
+                              p.last_seen = $occurred_at
+                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
+                """,
+                arena=arena, email=email, cid=chunk_id(arena),
+                channel="email", direction="inbound",
+                occurred_at="2026-05-08T00:00:00Z", now="2026-05-08T00:00:00Z",
+            )
+
+        # Assert two distinct nodes, one per arena, both with the same email.
+        result = session.run(
+            "MATCH (p:Person) WHERE p.email = $email RETURN p.arena AS arena",
+            email=email,
+        )
+        seen = sorted(rec["arena"] for rec in result if rec["arena"] in (arena_a, arena_b))
+        assert seen == sorted([arena_a, arena_b]), (
+            f"expected exactly one Person per arena for the same email, got {seen!r}"
+        )
+
+
+@_skip_no_neo4j
+def test_arena_scoped_search_does_not_traverse_other_arena(neo4j_driver) -> None:
+    """A graph search scoped to arena_a can't return arena_b nodes."""
+    driver, (arena_a, arena_b) = neo4j_driver
+    name = "Shared Concept"
+
+    with driver.session() as session:
+        for arena in (arena_a, arena_b):
+            session.run(
+                """
+                MERGE (c:Chunk {id: $cid})
+                SET c.arena = $arena, c.text = 't', c.path = 'p',
+                    c.created_at = '2026-05-08T00:00:00Z'
+                MERGE (e:Entity:Concept {arena: $arena, name: $name})
+                ON CREATE SET e.weight = 1.0, e.created_at = '2026-05-08T00:00:00Z'
+                MERGE (e)-[:MENTIONS]->(c)
+                """,
+                cid=f"c_{arena}", arena=arena, name=name,
+            )
+
+        # Arena-scoped lookup as the engine search now does it.
+        result = session.run(
+            """
+            MATCH (n:Entity {name: $name, arena: $arena})
+            RETURN n.arena AS arena, n.name AS name
+            """,
+            name=name, arena=arena_a,
+        )
+        rows = list(result)
+        assert len(rows) == 1
+        assert rows[0]["arena"] == arena_a, (
+            f"arena-scoped match returned wrong arena: {rows[0]['arena']!r}"
+        )
+
+        # And the unscoped query (intentionally cross-tenant) still
+        # finds two — proves the data is there, isolation is real.
+        result = session.run(
+            "MATCH (n:Entity {name: $name}) RETURN n.arena AS arena", name=name,
+        )
+        all_rows = sorted(
+            r["arena"] for r in result if r["arena"] in (arena_a, arena_b)
+        )
+        assert all_rows == sorted([arena_a, arena_b])
+
+
+@_skip_no_neo4j
+def test_forget_arena_b_leaves_arena_a_intact(neo4j_driver) -> None:
+    """Tenant-scoped delete on arena_b doesn't affect arena_a."""
+    driver, (arena_a, arena_b) = neo4j_driver
+
+    with driver.session() as session:
+        for arena in (arena_a, arena_b):
+            session.run(
+                """
+                MERGE (c:Chunk {id: $cid})
+                SET c.arena = $arena, c.text = 't', c.path = 'p',
+                    c.created_at = '2026-05-08T00:00:00Z'
+                MERGE (e:Entity:Concept {arena: $arena, name: 'thing'})
+                ON CREATE SET e.weight = 1.0, e.created_at = '2026-05-08T00:00:00Z'
+                MERGE (e)-[:MENTIONS]->(c)
+                """,
+                cid=f"c_{arena}", arena=arena,
+            )
+
+        # Tenant-scoped delete: same shape as forget-internal.
+        session.run(
+            "MATCH (c:Chunk {arena: $arena}) DETACH DELETE c", arena=arena_b,
+        )
+        session.run(
+            "MATCH (e:Entity {arena: $arena}) DETACH DELETE e", arena=arena_b,
+        )
+
+        # arena_a still intact.
+        a_chunks = session.run(
+            "MATCH (c:Chunk {arena: $arena}) RETURN count(c) AS n", arena=arena_a,
+        ).single()["n"]
+        a_entities = session.run(
+            "MATCH (e:Entity {arena: $arena}) RETURN count(e) AS n", arena=arena_a,
+        ).single()["n"]
+        assert a_chunks >= 1
+        assert a_entities >= 1
+
+        # arena_b gone.
+        b_chunks = session.run(
+            "MATCH (c:Chunk {arena: $arena}) RETURN count(c) AS n", arena=arena_b,
+        ).single()["n"]
+        b_entities = session.run(
+            "MATCH (e:Entity {arena: $arena}) RETURN count(e) AS n", arena=arena_b,
+        ).single()["n"]
+        assert b_chunks == 0
+        assert b_entities == 0