@pentatonic-ai/ai-agent-sdk 0.9.5 → 0.10.0

package/packages/memory-engine-v2/tests/e2e_smoke.py

@@ -0,0 +1,356 @@
+"""End-to-end smoke tests for pentatonic-memory-engine v2.
+
+Run on the engine host against compat at http://localhost:8199. Verifies:
+
+  1. /health + /health/deep
+  2. Single /store works end-to-end (compat → extractor-sync → org-model +
+     vector-index)
+  3. /search finds the stored record, late-materialisation hydrates content
+  4. Idempotency: same arena+content → same content-hash ID
+  5. /store-batch works with the exact wire shape TES's
+     engineStoreBatch helper posts (this is the TES-cutover gate)
+  6. Cross-arena isolation: search in arena A doesn't return arena B records
+  7. /forget by id deletes from both stores
+  8. /forget by metadata_contains (with arena filter) deletes matching set
+  9. FORGET cascade: extracted entities for an event vanish when the event
+     is forgotten
+
+Exits 0 on success, 1 on any failure. Prints a clear pass/fail line per
+test so the SSM output is greppable.
+"""
+
+from __future__ import annotations
+
+import json
+import sys
+import time
+import urllib.request
+import urllib.error
+
+BASE = "http://localhost:8199"
+
+
+def _http(method: str, path: str, body: dict | None = None) -> tuple[int, dict]:
+    data = json.dumps(body).encode() if body is not None else None
+    req = urllib.request.Request(
+        f"{BASE}{path}",
+        method=method,
+        data=data,
+        headers={"Content-Type": "application/json"} if body else {},
+    )
+    try:
+        with urllib.request.urlopen(req, timeout=30) as r:
+            return r.status, json.loads(r.read())
+    except urllib.error.HTTPError as e:
+        return e.code, {"_error": e.read().decode(errors="replace")[:500]}
+
+
+_passes = 0
+_fails: list[str] = []
+
+
+def _assert(cond: bool, msg: str) -> None:
+    global _passes
+    if cond:
+        _passes += 1
+        print(f"  PASS  {msg}")
+    else:
+        _fails.append(msg)
+        print(f"  FAIL  {msg}")
+
+
+def section(name: str) -> None:
+    print(f"\n=== {name} ===")
+
+
+# ----------------------------------------------------------------------
+# 1. Health
+# ----------------------------------------------------------------------
+
+section("health")
+code, body = _http("GET", "/health")
+_assert(code == 200, f"GET /health returns 200 (got {code})")
+_assert(body.get("status") == "healthy", f"health.status == healthy (got {body.get('status')})")
+_assert(body.get("service") == "pme2-compat", f"health.service identifies pme2-compat")
+
+code, body = _http("GET", "/health/deep")
+_assert(code == 200, f"GET /health/deep returns 200 (got {code})")
+stores = body.get("stores", {})
+_assert(stores.get("org_model", {}).get("status") == "ok", "health/deep: org_model ok")
+_assert(stores.get("vector_index", {}).get("status") == "ok", "health/deep: vector_index ok")
+_assert(stores.get("embed_gateway", {}).get("status") == "ok", "health/deep: embed_gateway ok")
+_assert(stores.get("embed_gateway", {}).get("dim") == 4096, "health/deep: embed dim is 4096")
+
+
+# ----------------------------------------------------------------------
+# 2-4. /store + /search + idempotency
+# ----------------------------------------------------------------------
+
+section("/store + /search + idempotency")
+
+# Use a unique arena per test run so we don't collide with anything else
+# accumulating in the engine. Time-based suffix.
+ARENA_A = f"e2e-test-a-{int(time.time())}"
+ARENA_B = f"e2e-test-b-{int(time.time())}"
+
+# Single store.
+store_body = {
+    "content": "The quarterly review confirmed that the integration project is on track for Q3 delivery",
+    "metadata": {"arena": ARENA_A, "clientId": ARENA_A, "kind": "doc"},
+}
+code, resp = _http("POST", "/store", store_body)
+_assert(code == 200, f"POST /store returns 200 (got {code})")
+event_id_1 = resp.get("id")
+_assert(isinstance(event_id_1, str) and len(event_id_1) == 32, f"store returns 32-char hex id (got {event_id_1!r})")
+_assert(resp.get("layerId", "").startswith("ml_"), "store returns layerId in v1-compat shape")
+
+# Idempotency: same arena + content → same id.
+code2, resp2 = _http("POST", "/store", store_body)
+_assert(code2 == 200, "second /store with same content returns 200")
+_assert(resp2.get("id") == event_id_1, f"idempotent: same content → same id (got {resp2.get('id')} vs {event_id_1})")
+
+# Different content → different id.
+store_body2 = {**store_body, "content": "Different content same arena should yield a different id deterministically"}
+code3, resp3 = _http("POST", "/store", store_body2)
+_assert(code3 == 200, "third /store (different content) returns 200")
+_assert(resp3.get("id") != event_id_1, "different content → different id")
+event_id_2 = resp3.get("id")
+
+# Search should find both.
+time.sleep(1)  # let Qdrant settle
+search_body = {"query": "quarterly review integration project", "arena": ARENA_A, "limit": 5}
+code, sresp = _http("POST", "/search", search_body)
+_assert(code == 200, f"POST /search returns 200 (got {code})")
+results = sresp.get("results", [])
+_assert(len(results) >= 1, f"/search returns at least 1 result (got {len(results)})")
+ids_found = [r["id"] for r in results]
+_assert(event_id_1 in ids_found, f"/search finds first stored event ({event_id_1[:8]}...)")
+_assert(all(r.get("similarity", 0) > 0 for r in results), "every result has similarity > 0")
+# Late materialisation: content is hydrated from org-model
+_assert(any("integration project" in r.get("content", "") for r in results),
+        "late materialisation: full content (not just preview) in response")
+
+
+# ----------------------------------------------------------------------
+# 5. TES wire shape — exact engineStoreBatch body
+# ----------------------------------------------------------------------
+
+section("TES engineStoreBatch wire shape")
+
+# This is the exact body TES's SDK posts (see ai-agent-sdk packages/memory/src/engine.js).
+# Gate: v2 must accept this without modification.
+tes_arena = f"e2e-tes-{int(time.time())}"
+tes_body = {
+    "records": [
+        {
+            "id": "stable-dedup-key-1",
+            "content": "First record from the simulated TES consumer drain",
+            "metadata": {
+                "arena": tes_arena,
+                "layer_type": "episodic",
+                "actor_user_id": "u-philh",
+                "kind": "chat",
+                "channel": "slack",
+            },
+        },
+        {
+            "id": "stable-dedup-key-2",
+            "content": "Second record, same batch, different content payload",
+            "metadata": {
+                "arena": tes_arena,
+                "layer_type": "episodic",
+                "actor_user_id": "u-philh",
+                "kind": "chat",
+                "channel": "slack",
+            },
+        },
+        {
+            # No caller-supplied id — v2 derives content-hash id.
+            "content": "Third record without explicit id, server derives the content-hash",
+            "metadata": {
+                "arena": tes_arena,
+                "layer_type": "episodic",
+                "actor_user_id": "u-philh",
+                "kind": "chat",
+            },
+        },
+    ],
+}
+code, resp = _http("POST", "/store-batch", tes_body)
+_assert(code == 200, f"POST /store-batch (TES shape) returns 200 (got {code})")
+_assert(resp.get("status") == "ok", f"store-batch.status == ok (got {resp.get('status')})")
+_assert(resp.get("inserted") == 3, f"store-batch.inserted == 3 (got {resp.get('inserted')})")
+ids = resp.get("ids", [])
+_assert(len(ids) == 3, f"store-batch.ids has 3 entries (got {len(ids)})")
+_assert(all(isinstance(i, str) and len(i) == 32 for i in ids), "all ids are 32-char hex content-hashes")
+
+# Search for one of them.
+time.sleep(1)
+code, sresp = _http("POST", "/search", {"query": "simulated TES consumer", "arena": tes_arena, "limit": 5})
+_assert(code == 200, "/search (TES arena) returns 200")
+_assert(len(sresp.get("results", [])) >= 1, "TES batch records are findable via /search")
+
+
+# ----------------------------------------------------------------------
+# 6. Cross-arena isolation
+# ----------------------------------------------------------------------
+
+section("cross-arena isolation")
+
+# Write something to arena B.
+code, resp = _http("POST", "/store", {
+    "content": "Secrets that must not leak into the other tenant under any condition",
+    "metadata": {"arena": ARENA_B, "clientId": ARENA_B, "kind": "doc"},
+})
+_assert(code == 200, "store in arena B succeeds")
+
+time.sleep(1)
+# Search arena A — should NOT find the arena B record. Match is
+# case-insensitive — stored content uses capital "Secrets".
+code, sresp = _http("POST", "/search", {"query": "secrets must not leak", "arena": ARENA_A, "limit": 5})
+_assert(code == 200, "search in arena A returns 200")
+contents_a = [r.get("content", "").lower() for r in sresp.get("results", [])]
+_assert(not any("secrets" in c for c in contents_a),
+        "isolation: arena B's content does NOT appear in arena A search")
+
+# Search arena B — should find its own record.
+code, sresp = _http("POST", "/search", {"query": "secrets must not leak", "arena": ARENA_B, "limit": 5})
+contents_b = [r.get("content", "").lower() for r in sresp.get("results", [])]
+_assert(any("secrets" in c for c in contents_b),
+        "arena B can find its own records")
+
+
+# ----------------------------------------------------------------------
+# 7. /forget by id
+# ----------------------------------------------------------------------
+
+section("/forget by id")
+
+# Forget event_id_2 (the second arena-A record).
+code, fresp = _http("POST", "/forget", {"id": event_id_2})
+_assert(code == 200, "/forget by id returns 200")
+_assert(fresp.get("deleted") == 1, f"/forget deletes exactly 1 (got {fresp.get('deleted')})")
+
+time.sleep(1)
+# Now search arena A — event_id_1 should still be there, event_id_2 should be gone.
+code, sresp = _http("POST", "/search", {"query": "Different content same arena", "arena": ARENA_A, "limit": 5})
+ids_after = [r["id"] for r in sresp.get("results", [])]
+_assert(event_id_2 not in ids_after, f"forgotten event no longer in search results")
+
+
+# ----------------------------------------------------------------------
+# 8. /forget by metadata_contains
+# ----------------------------------------------------------------------
+
+section("/forget by metadata_contains")
+
+# Forget the entire TES batch by arena filter.
+code, fresp = _http("POST", "/forget", {"metadata_contains": {"arena": tes_arena}})
+_assert(code == 200, "/forget by metadata_contains returns 200")
+_assert(fresp.get("deleted", 0) >= 3, f"forget-by-arena deletes >= 3 (got {fresp.get('deleted')})")
+
+time.sleep(1)
+# Search should return nothing.
+code, sresp = _http("POST", "/search", {"query": "simulated TES consumer", "arena": tes_arena, "limit": 5})
+_assert(len(sresp.get("results", [])) == 0,
+        f"TES arena search after forget returns empty (got {len(sresp.get('results', []))})")
+
+
+# ----------------------------------------------------------------------
+# 9. /search requires arena (no implicit "general" wildcard)
+# ----------------------------------------------------------------------
+
+section("/search rejects unscoped query")
+
+code, sresp = _http("POST", "/search", {"query": "no arena specified", "limit": 5})
+_assert(code == 400, f"/search without arena returns 400 (got {code}) — v2 explicit-scope contract")
+
+
+# ----------------------------------------------------------------------
+# 10. TES engineSearch wire shape — arenas[] list + arena single-field
+# ----------------------------------------------------------------------
+
+section("TES engineSearch wire shape (arenas list)")
+
+# This is what the SDK's engineSearch helper actually posts: arenas[]
+# list (authoritative) + arena single-field (back-compat). v2 reads
+# arenas[] when present; ignoring arena is the right behaviour.
+tes_search_arena = f"e2e-tes-search-{int(time.time())}"
+tes_user_arena = f"{tes_search_arena}:u-philh"
+_http("POST", "/store", {
+    "content": "User-scoped memory only visible in the user-arena lane",
+    "metadata": {"arena": tes_user_arena, "clientId": tes_search_arena, "kind": "chat"},
+})
+_http("POST", "/store", {
+    "content": "Tenant-wide memory visible across the whole client",
+    "metadata": {"arena": tes_search_arena, "clientId": tes_search_arena, "kind": "doc"},
+})
+time.sleep(1)
+
+# engineSearch with userId composes arenas=[clientId, clientId:userId] —
+# caller sees their own AND tenant-wide records, never another user's.
+body = {
+    "query": "memory tenant user-scoped visible",
+    "arenas": [tes_search_arena, tes_user_arena],
+    "arena": tes_search_arena,
+    "limit": 10,
+    "min_score": 0.001,
+}
+code, sresp = _http("POST", "/search", body)
+_assert(code == 200, f"engineSearch wire shape returns 200 (got {code})")
+ids = [r["id"] for r in sresp.get("results", [])]
+arenas_seen = {r.get("metadata", {}).get("arena") for r in sresp.get("results", [])}
+_assert(len(sresp.get("results", [])) >= 2, f"multi-arena search returns both records (got {len(sresp.get('results', []))})")
+_assert(tes_search_arena in arenas_seen and tes_user_arena in arenas_seen,
+        f"both arena and user-arena records present (got {arenas_seen})")
+
+
+# ----------------------------------------------------------------------
+# 11. TES engineForget wire shape — top-level arena + id
+# ----------------------------------------------------------------------
+
+section("TES engineForget wire shape (top-level arena + id)")
+
+# Plant a record to forget.
+code, resp = _http("POST", "/store", {
+    "content": "Record specifically to be forgotten via the engineForget(id=...) wire shape",
+    "metadata": {"arena": ARENA_A, "clientId": ARENA_A, "kind": "doc"},
+})
+forget_target_id = resp["id"]
+
+# SDK engineForget(id=X) posts {arena: clientId, id: X}. v2 should use
+# id and ignore the redundant top-level arena.
+code, fresp = _http("POST", "/forget", {"arena": ARENA_A, "id": forget_target_id})
+_assert(code == 200, f"engineForget(id) with top-level arena returns 200 (got {code})")
+_assert(fresp.get("deleted") == 1, f"deletes exactly 1 by id (got {fresp.get('deleted')})")
+
+
+# ----------------------------------------------------------------------
+# 12. Endpoints NOT yet implemented (gracefully missing)
+# ----------------------------------------------------------------------
+
+section("Endpoints v2 doesn't implement yet")
+
+# personFacets calls /aggregate; peopleList calls /people-list-internal.
+# v2 doesn't implement these (waiting on keystone spec + recursive-CTE
+# spike). They should 404, not crash — TES resolvers have fallback
+# paths that handle missing aggregate endpoints gracefully.
+code, _ = _http("POST", "/aggregate", {"arena": ARENA_A, "contact_email": "a@b.com"})
+_assert(code == 404, f"/aggregate returns 404 (v2 doesn't implement yet) (got {code})")
+code, _ = _http("POST", "/people-list-internal", {"arenas": [ARENA_A]})
+_assert(code == 404, f"/people-list-internal returns 404 (v2 doesn't implement yet) (got {code})")
+
+
+# ----------------------------------------------------------------------
+# Done
+# ----------------------------------------------------------------------
+
+print(f"\n{'=' * 60}")
+if _fails:
+    print(f"FAIL — {len(_fails)} failure(s):")
+    for f in _fails:
+        print(f"  - {f}")
+    print(f"passes: {_passes}")
+    sys.exit(1)
+print(f"OK — all {_passes} assertions passed")
+sys.exit(0)