@pentatonic-ai/ai-agent-sdk 0.10.5 → 0.10.7

package/packages/memory-engine-v2/extractor-async/worker.py

@@ -41,6 +41,14 @@ import psycopg.rows
 
 from confidence import corroborated_confidence
 from entity_id import entity_id, normalize_surface_form
+from extraction_schema import (
+    ALLOWED_ENT_TYPES,
+    ALLOWED_FCT_CATEGORIES,
+    EXTRACTION_SCHEMA,
+    MAX_ENTITIES_PER_EVENT,
+    MAX_FACTS_PER_EVENT,
+    MAX_RELATIONSHIPS_PER_EVENT,
+)
 from noise_filter import is_noise_entity_name
 from sensitive_filter import SKIP_SENSITIVE_CONTENT, is_sensitive_event
 
@@ -90,15 +98,93 @@ DISTILL_TRACE_ENABLED = os.environ.get(
 # chunk via a JSONDecodeError. Pipe-delimited records, one per line,
 # recover at line granularity — a malformed line skips itself, the rest
 # of the chunk lands. See 2026-05-18 ops notes.
+#
+# 2026-06-11 update: guided JSON is back as an OPT-IN second mode
+# (DISTILL_OUTPUT_MODE=guided_json, default "kv" — a no-op until an
+# operator flips it). Both halves of the 2026-05-18 removal rationale
+# are answered this time:
+#   (a) the self-hosted Qwen2.5-7B vLLM box enforces structured output
+#       via logit masking (xgrammar/outlines) — the model CANNOT emit
+#       schema-invalid bytes, unlike the old VL gateway which
+#       half-ignored response_format;
+#   (b) blast radius is solved structurally — the schema is an array
+#       of per-event objects (see extraction_schema.py), so one
+#       event's content can't corrupt another's parse; the only
+#       residual failure is max_tokens truncation, and
+#       _parse_guided_json salvages every complete event object.
+# ALLOWED_ENT_TYPES / ALLOWED_FCT_CATEGORIES now live in
+# extraction_schema.py (imported above) so the schema enums and the
+# KV prompt pin to the same single source.
 EVENT_HEADER_RE = re.compile(r"^===?\s*event\s+(\d+)\s*===?\s*$", re.IGNORECASE)
-ALLOWED_ENT_TYPES = {
-    "person", "org", "product", "place", "project",
-    "concept", "topic", "date", "other",
-}
-ALLOWED_FCT_CATEGORIES = {
-    "decision", "commitment", "state", "mention",
-    "observation", "preference",
-}
+
+# Output mode flag. "kv" (default) keeps today's pipe-delimited path
+# byte-for-byte; "guided_json" switches the prompt, request params and
+# parser. Anything unrecognised falls back to "kv" — fail-safe.
+DISTILL_OUTPUT_MODE = os.environ.get("DISTILL_OUTPUT_MODE", "kv").strip().lower()
+if DISTILL_OUTPUT_MODE not in ("kv", "guided_json"):
+    log.warning(
+        f"DISTILL_OUTPUT_MODE={DISTILL_OUTPUT_MODE!r} unrecognised — using 'kv'"
+    )
+    DISTILL_OUTPUT_MODE = "kv"
+
+# How the structured-output schema is attached to the request in
+# guided_json mode. The repo carries no pin for the engine box's vLLM
+# version, so this is operator-selectable:
+#   - "response_format" (default): OpenAI-style
+#     response_format={"type":"json_schema","json_schema":{...}} —
+#     supported by vLLM >= 0.6.x OpenAI-compat server.
+#   - "guided_json": vLLM's legacy extension param (top-level
+#     `guided_json` in the request body; what openai-client users pass
+#     via extra_body). FALLBACK for older vLLM builds that predate
+#     json_schema response_format.
+# Exactly one is sent — some vLLM versions reject requests that carry
+# both guided-decoding params at once.
+DISTILL_GUIDED_PARAM_STYLE = os.environ.get(
+    "DISTILL_GUIDED_PARAM_STYLE", "response_format"
+).strip().lower()
+if DISTILL_GUIDED_PARAM_STYLE not in ("response_format", "guided_json"):
+    log.warning(
+        f"DISTILL_GUIDED_PARAM_STYLE={DISTILL_GUIDED_PARAM_STYLE!r} unrecognised "
+        f"— using 'response_format'"
+    )
+    DISTILL_GUIDED_PARAM_STYLE = "response_format"
+
+# Optional chat-template kwargs forwarded verbatim on every chat
+# completion (vLLM extension: top-level `chat_template_kwargs`).
+# Needed for thinking-capable teachers: Qwen3.x chat templates default
+# enable_thinking=true, which burns the max_tokens budget on reasoning
+# the distiller never reads. The 2026-06-11 teacher bake-off ran the
+# Qwen3.6 lanes with {"enable_thinking": false}, so the prod swap must
+# send the same switch for its traces to match the benchmarked
+# distribution. Unset (default) sends nothing — the request body stays
+# byte-identical for teachers without template switches (Qwen2.5).
+DISTILL_CHAT_TEMPLATE_KWARGS: dict[str, Any] | None = None
+_raw_ctk = os.environ.get("DISTILL_CHAT_TEMPLATE_KWARGS", "").strip()
+if _raw_ctk:
+    try:
+        _parsed_ctk = json.loads(_raw_ctk)
+        if not isinstance(_parsed_ctk, dict):
+            raise ValueError("must be a JSON object")
+        DISTILL_CHAT_TEMPLATE_KWARGS = _parsed_ctk
+    except ValueError as e:
+        log.warning(f"DISTILL_CHAT_TEMPLATE_KWARGS invalid ({e}) — ignoring")
+
+# JSON output carries structural overhead (braces, quotes, key names)
+# the KV format doesn't, so guided mode gets its own per-event token
+# budget. Truncation is guided mode's ONLY parse-failure mode (the
+# schema enforcer guarantees validity up to the cut), so this errs
+# higher than the KV 300.
+#
+# NOTE the budget is SHARED across the chunk (max_tokens = this × N
+# events per request). A fully-maxed event (8 ent / 6 fct with 140-char
+# statements / 6 rel + JSON overhead) is ~1.1k output tokens, so chunk
+# size and this value must be chosen together against the server's
+# max_model_len. Raised 400→900 after prod showed 15% of 5-event chunks
+# truncating on `length` (2026-06-12); prod now runs EVENTS_PER_LLM_CALL=3
+# so 3×900 output + ~2.1k prompt stays well inside the L40S 8192 ctx.
+LLM_MAX_TOKENS_PER_EVENT_JSON = int(
+    os.environ.get("LLM_MAX_TOKENS_PER_EVENT_JSON", "900")
+)
 
 
 # --------------------------------------------------------------------
@@ -165,10 +251,71 @@ A whole file is one entity, not twenty.
 - Output ONLY the formatted records. No header, no footer, no prose."""
 
 
+# Guided-JSON variant of BATCH_SYSTEM_PROMPT. Same CONTENT rules
+# (conservatism, per-event caps, code-content rule, subject-must-be-a-
+# declared-entity, email-alias pairing, statement <= 140 chars, never
+# skip an event) — only the output-format scaffolding changes. The
+# pipe-format anchoring ("COUNT THE PIPES", pipe/newline substitution)
+# is dropped: vLLM's guided decoding enforces the schema mechanically,
+# so the prompt no longer needs to beg for format compliance, and JSON
+# string escaping makes the pipe/newline substitution rules moot.
+GUIDED_JSON_SYSTEM_PROMPT = """You extract structured knowledge from N \
+events for a personal-memory graph.
+
+You will receive N events, each prefixed with `[event K]`. Respond \
+with a single JSON object: {"events": [...]} containing one object \
+per input event. Be conservative — only emit things explicitly stated.
+
+Each per-event object has:
+  "index": the zero-indexed event number, matching the input `[event K]`.
+  "entities": array of {"name", "type", "email"?}.
+  "facts": array of {"category", "subject", "predicate", "object", "statement"}.
+  "relationships": array of {"from", "to", "type"}.
+
+RULES:
+- NEVER skip an event — if an event has nothing to extract, emit its \
+object with "index" set and empty arrays.
+- entities: type ∈ {person, org, product, place, project, concept, \
+topic, date, other}.
+  email (OPTIONAL, person only): when the event body or attributes
+  show an email address that unambiguously identifies the person,
+  include it. This pairs the name+email forms so a later event seeing
+  only the email resolves to the same entity. Omit the key otherwise.
+- facts: category ∈ {decision, commitment, state, mention, \
+observation, preference}.
+  subject MUST be an entity name declared in THIS event's "entities".
+  predicate is a short verb phrase (e.g. "agreed to", "owns", "works at").
+  object MAY be an entity name OR a literal string OR null if absent.
+  statement ≤ 140 characters, a self-contained sentence.
+  WORKED EXAMPLE: {"category": "commitment", "subject": "Timothy \
+Bradley", "predicate": "agreed to", "object": "SAFE amendments", \
+"statement": "Timothy confirmed the SAFE amendments are set (14 May 2026)"}
+- relationships: "from" and "to" MUST be entity names declared in THIS \
+event's "entities". "type" is a short verb / preposition phrase.
+- HARD CAPS per event: 8 entities, 6 facts, 6 relationships. Pick the \
+most salient.
+- For code / technical content: extract only top-level services, \
+modules, or domain concepts. NOT variables, types, or method names. \
+A whole file is one entity, not twenty.
+- Output ONLY the JSON object. No markdown fences, no prose."""
+
+
+# The system prompt actually sent to the LLM under the current output
+# mode. Everything downstream (request body, trace fingerprint) hangs
+# off this so the two can never disagree.
+ACTIVE_SYSTEM_PROMPT = (
+    GUIDED_JSON_SYSTEM_PROMPT
+    if DISTILL_OUTPUT_MODE == "guided_json"
+    else BATCH_SYSTEM_PROMPT
+)
+
 # Teacher-prompt fingerprint for trace logging. If the prompt changes,
 # the hash changes — lets training-data exports filter by teacher
-# version so we never mix outputs from a retired prompt.
-SYSTEM_PROMPT_HASH = hashlib.sha256(BATCH_SYSTEM_PROMPT.encode()).hexdigest()[:16]
+# version so we never mix outputs from a retired prompt. Computed from
+# the ACTIVE prompt, so flipping DISTILL_OUTPUT_MODE auto-segments
+# distillation_traces into a new teacher version (KV-format traces and
+# guided-JSON traces never mix in a training export).
+SYSTEM_PROMPT_HASH = hashlib.sha256(ACTIVE_SYSTEM_PROMPT.encode()).hexdigest()[:16]
 
 
 # --------------------------------------------------------------------
@@ -353,14 +500,234 @@ def _split_event_blocks(text: str, expected_n: int) -> list[str]:
     return slices
 
 
+# --------------------------------------------------------------------
+# Guided-JSON parsing (DISTILL_OUTPUT_MODE=guided_json)
+# --------------------------------------------------------------------
+
+
+def _load_guided_payload(text: str) -> dict[str, Any] | None:
+    """Parse the guided-JSON chunk output into the {"events": [...]}
+    payload, salvaging what's complete if the output was truncated.
+
+    Under guided decoding the server's logit masking guarantees every
+    emitted byte is schema-consistent, so the ONLY way the payload can
+    fail to parse is max_tokens truncation mid-stream. Salvage is
+    therefore simple and structural: walk back to the last complete
+    `}` (the close of the last fully-emitted event object), close the
+    events array + root object, and re-parse. Each step back drops at
+    most one (incomplete) event — per-event degradation, never
+    chunk-level loss. Returns None if nothing parseable remains."""
+    raw = (text or "").strip()
+    if not raw:
+        return None
+    # Defensive fence strip — can't occur under guided decoding, but
+    # the bake-off script replays this parser over unguided output too.
+    if raw.startswith("```"):
+        raw = raw.strip("`").strip()
+        if raw.lower().startswith("json"):
+            raw = raw[4:].lstrip()
+    try:
+        payload = json.loads(raw)
+        return payload if isinstance(payload, dict) else None
+    except json.JSONDecodeError:
+        pass
+    # Truncated: trim to the last complete `}` of the events array and
+    # close the structure. Walk back through `}` occurrences until a
+    # candidate parses (bounded — each iteration discards at least one
+    # char, and 200 closing braces covers far more events than a chunk
+    # can hold).
+    end = len(raw)
+    for _ in range(200):
+        idx = raw.rfind("}", 0, end)
+        if idx < 0:
+            return None
+        candidate = raw[: idx + 1] + "]}"
+        try:
+            payload = json.loads(candidate)
+            return payload if isinstance(payload, dict) else None
+        except json.JSONDecodeError:
+            end = idx
+    return None
+
+
+def _resolve_event_index(ev: dict[str, Any], pos: int, expected_n: int) -> int | None:
+    """Map a parsed event object to its result slot. Trust the model's
+    "index" field when it's a valid in-range int (it mirrors the
+    `[event K]` input header); fall back to array position otherwise.
+    None = undeliverable (both out of range) — the object is dropped
+    without corrupting any other event's slot."""
+    idx = ev.get("index")
+    if isinstance(idx, int) and not isinstance(idx, bool) and 0 <= idx < expected_n:
+        return idx
+    if 0 <= pos < expected_n:
+        return pos
+    return None
+
+
+def _parse_guided_json(text: str, expected_n: int) -> list[dict[str, Any]]:
+    """Parse guided-JSON output into per-event extraction dicts —
+    sibling of _parse_kv_records, returning the IDENTICAL shape
+    ({"entities": [...], "facts": [...], "relationships": [...]}, with
+    entity emails promoted into "aliases") so the upsert path and trace
+    logging are untouched by the output-mode flip.
+
+    Defensive beyond what guided decoding guarantees: truncation is
+    salvaged per-event (see _load_guided_payload), per-item junk is
+    skipped, the per-event hard caps are re-enforced, and string fields
+    are normalised exactly as the KV parser normalises them (strip,
+    lowercase type/category, `-`/empty/null object → None, non-email
+    "email" values dropped). Always returns expected_n entries."""
+    results: list[dict[str, Any]] = [
+        {"entities": [], "facts": [], "relationships": []} for _ in range(expected_n)
+    ]
+    payload = _load_guided_payload(text)
+    if payload is None:
+        return results
+    events = payload.get("events")
+    if not isinstance(events, list):
+        return results
+    for pos, ev in enumerate(events):
+        if not isinstance(ev, dict):
+            continue
+        idx = _resolve_event_index(ev, pos, expected_n)
+        if idx is None:
+            continue
+        target = results[idx]
+        ents = ev.get("entities")
+        for e in (ents if isinstance(ents, list) else [])[:MAX_ENTITIES_PER_EVENT]:
+            if not isinstance(e, dict):
+                continue
+            name = str(e.get("name") or "").strip()
+            if not name:
+                continue
+            etype = str(e.get("type") or "").strip().lower()
+            ent: dict[str, Any] = {"type": etype, "name": name}
+            # Mirror the KV 4th-field rule: promote into aliases only
+            # when it actually looks like an email; drop junk silently.
+            email = e.get("email")
+            if isinstance(email, str):
+                email = email.strip()
+                if email and "@" in email and " " not in email:
+                    ent["aliases"] = [email]
+            target["entities"].append(ent)
+        facts = ev.get("facts")
+        for f in (facts if isinstance(facts, list) else [])[:MAX_FACTS_PER_EVENT]:
+            if not isinstance(f, dict):
+                continue
+            stmt = str(f.get("statement") or "").strip()
+            if not stmt:
+                continue
+            obj = f.get("object")
+            obj = obj.strip() if isinstance(obj, str) else None
+            target["facts"].append(
+                {
+                    "category": str(f.get("category") or "").strip().lower(),
+                    "subject": str(f.get("subject") or "").strip(),
+                    "predicate": str(f.get("predicate") or "").strip(),
+                    "object": None if obj in (None, "", "-", "null", "None") else obj,
+                    "statement": stmt,
+                }
+            )
+        rels = ev.get("relationships")
+        for r in (rels if isinstance(rels, list) else [])[:MAX_RELATIONSHIPS_PER_EVENT]:
+            if not isinstance(r, dict):
+                continue
+            frm = str(r.get("from") or "").strip()
+            to = str(r.get("to") or "").strip()
+            rtype = str(r.get("type") or "").strip()
+            if frm and to and rtype:
+                target["relationships"].append({"from": frm, "to": to, "type": rtype})
+    return results
+
+
+def _guided_event_slices(text: str, expected_n: int) -> list[str]:
+    """Per-event raw slices for trace logging in guided mode — the
+    JSON-mode sibling of _split_event_blocks, same shape contract
+    (expected_n entries, missing events as empty strings). Each slice
+    is the model's event object re-serialised verbatim-in-content
+    (key order preserved, non-ASCII kept) so distillation_traces stays
+    a faithful (input, output) training pair."""
+    slices: list[str] = [""] * expected_n
+    payload = _load_guided_payload(text)
+    if payload is None:
+        return slices
+    events = payload.get("events")
+    if not isinstance(events, list):
+        return slices
+    for pos, ev in enumerate(events):
+        if not isinstance(ev, dict):
+            continue
+        idx = _resolve_event_index(ev, pos, expected_n)
+        if idx is not None:
+            slices[idx] = json.dumps(ev, ensure_ascii=False)
+    return slices
+
+
+def _build_request_body(user_prompt: str, n: int) -> dict[str, Any]:
+    """Chat-completions request body for one N-event chunk. Pure —
+    everything mode-dependent (prompt, token budget, structured-output
+    params) keys off the module-level flags so this is unit-testable.
+
+    kv mode (default): byte-for-byte the pre-flag body — KV-text
+    output, no guided_json / response_format. The benefit of
+    structured-output enforcement was half-ignored by the old VL
+    upstream, and the KV parser recovers from per-line drift.
+
+    guided_json mode: attaches EXTRACTION_SCHEMA via ONE of the two
+    vLLM structured-output param styles (DISTILL_GUIDED_PARAM_STYLE;
+    some vLLM versions reject requests carrying both at once):
+      - response_format {"type": "json_schema", ...} — OpenAI-style,
+        current vLLM (default).
+      - top-level guided_json — vLLM's legacy extension param (what
+        openai-client callers pass via extra_body), fallback for older
+        server builds.
+    """
+    body: dict[str, Any] = {
+        "model": LLM_MODEL,
+        "messages": [
+            {"role": "system", "content": ACTIVE_SYSTEM_PROMPT},
+            {"role": "user", "content": user_prompt},
+        ],
+        "temperature": 0.0,
+        "max_tokens": (
+            LLM_MAX_TOKENS_PER_EVENT_JSON
+            if DISTILL_OUTPUT_MODE == "guided_json"
+            else LLM_MAX_TOKENS_PER_EVENT
+        ) * n,
+    }
+    if DISTILL_CHAT_TEMPLATE_KWARGS:
+        body["chat_template_kwargs"] = DISTILL_CHAT_TEMPLATE_KWARGS
+    if DISTILL_OUTPUT_MODE == "guided_json":
+        if DISTILL_GUIDED_PARAM_STYLE == "guided_json":
+            body["guided_json"] = EXTRACTION_SCHEMA
+        else:
+            body["response_format"] = {
+                "type": "json_schema",
+                "json_schema": {
+                    "name": "memory_extraction",
+                    "strict": True,
+                    "schema": EXTRACTION_SCHEMA,
+                },
+            }
+    return body
+
+
 async def call_llm_batch(
     client: httpx.AsyncClient, events: list[dict[str, Any]]
 ) -> list[dict[str, Any]]:
     """Send N events in a single chat-completion call, return the list
-    of per-event extraction dicts in input order. The model emits
-    pipe-delimited KV records (see BATCH_SYSTEM_PROMPT); the parser is
-    line-tolerant so a malformed record skips itself rather than
-    failing the chunk. Raises only on transport failure or completely
+    of per-event extraction dicts in input order.
+
+    kv mode (default): the model emits pipe-delimited KV records (see
+    BATCH_SYSTEM_PROMPT); the parser is line-tolerant so a malformed
+    record skips itself rather than failing the chunk.
+
+    guided_json mode: the model emits the EXTRACTION_SCHEMA-constrained
+    JSON envelope under server-side guided decoding; the parser
+    salvages complete event objects from a truncated stream so failure
+    degrades per-event, never per-chunk. Both parsers return the same
+    per-event dict shape, so everything downstream of this function is
+    mode-agnostic. Raises only on transport failure or completely
     empty output."""
     n = len(events)
     if n == 0:
@@ -378,20 +745,7 @@ async def call_llm_batch(
         build_event_block(i, ev) for i, ev in enumerate(events)
     )
 
-    body: dict[str, Any] = {
-        "model": LLM_MODEL,
-        "messages": [
-            {"role": "system", "content": BATCH_SYSTEM_PROMPT},
-            {"role": "user", "content": user_prompt},
-        ],
-        "temperature": 0.0,
-        "max_tokens": LLM_MAX_TOKENS_PER_EVENT * n,
-        # KV-text output — no guided_json / response_format. The
-        # benefit of structured-output enforcement was already
-        # half-ignored by VL upstream, and the parser now recovers
-        # from per-line drift so the schema enforcement isn't worth
-        # the JSON brittleness it brought.
-    }
+    body = _build_request_body(user_prompt, n)
     r = await client.post(LLM_ENDPOINT, json=body, headers=headers)
     r.raise_for_status()
     data = r.json()
@@ -400,12 +754,16 @@ async def call_llm_batch(
         text = data.get("message", {}).get("content", "")
     if not text:
         raise RuntimeError(f"llm returned no content: {json.dumps(data)[:300]}")
-    parsed = _parse_kv_records(text, n)
+    if DISTILL_OUTPUT_MODE == "guided_json":
+        parsed = _parse_guided_json(text, n)
+        slices = _guided_event_slices(text, n)
+    else:
+        parsed = _parse_kv_records(text, n)
+        slices = _split_event_blocks(text, n)
     # Attach the per-event raw slice so downstream trace logging gets
     # the model's verbatim output for THIS event without re-splitting
     # the chunk-level text. Parser semantics are unaffected — the
     # raw_slice key is ignored by upsert paths.
-    slices = _split_event_blocks(text, n)
     for record, slice_text in zip(parsed, slices):
         record["raw_slice"] = slice_text
     return parsed
@@ -1141,7 +1499,9 @@ async def amain():
         f"endpoint={LLM_ENDPOINT or '(stub)'}, model={LLM_MODEL}, "
         f"poll={POLL_INTERVAL_SEC}s, claim={BATCH_SIZE}, "
         f"events_per_call={EVENTS_PER_LLM_CALL}, "
-        f"concurrent_calls={CONCURRENT_LLM_CALLS})"
+        f"concurrent_calls={CONCURRENT_LLM_CALLS}, "
+        f"output_mode={DISTILL_OUTPUT_MODE}, "
+        f"prompt_hash={SYSTEM_PROMPT_HASH})"
     )
     stub_mode = not LLM_ENDPOINT
     if stub_mode:

package/packages/memory-engine-v2/extractor-sync/server.py

@@ -56,11 +56,15 @@ _pool: AsyncConnectionPool | None = None
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     global _pool
+    # Default (tuple) row factory — _upsert_entities and friends index
+    # fetchone() rows positionally, matching extractor-async's worker.
+    # A dict_row factory here turns row[0] into KeyError: 0 on the
+    # entity-merge path (2026-06-11 prod incident: every extract that
+    # re-saw a known entity 500'd; only never-seen-entity events stored).
     _pool = AsyncConnectionPool(
         conninfo=PG_DSN,
         min_size=8,
         max_size=50,
-        kwargs={"row_factory": psycopg.rows.dict_row},
         open=False,
     )
     await _pool.open()
@@ -89,7 +93,7 @@ class ExtractRequest(BaseModel):
     clientId: str
     userId: str | None = None
     event_type: str = "STORE_MEMORY"
-    source_kind: str  # 'chat' | 'note' | 'doc' | 'event' | 'ticket' | 'commit' | 'system' | 'agent'
+    source_kind: str  # 'chat' | 'note' | 'doc' | 'event' | 'ticket' | 'commit' | 'system' | 'agent' | 'code_reference'
     source_id: str | None = None
     content: str
     attributes: dict[str, Any] = {}

package/packages/memory-engine-v2/extractor-sync/test_paired_extraction.py

@@ -22,8 +22,14 @@ import pytest
 
 
 # Load extractor-sync's server.py as a module so we can call its
-# private helpers directly.
+# private helpers directly. server.py flat-imports its siblings
+# (entity_id) the way the container's WORKDIR layout resolves them, so
+# this directory must be on sys.path — otherwise exec_module raises
+# ImportError and the module-level skip below silently swallows the
+# whole suite whenever pytest runs from the repo root.
 _THIS = Path(__file__).resolve().parent
+if str(_THIS) not in sys.path:
+    sys.path.insert(0, str(_THIS))
 _SPEC = importlib.util.spec_from_file_location("extractor_sync_server",
                                                 _THIS / "server.py")
 assert _SPEC and _SPEC.loader
@@ -206,3 +212,78 @@ def test_extract_event_organizer_object_form() -> None:
     assert len(entities) == 1
     assert entities[0]["canonical_name"] == "X Person"
     assert "x@example.com" in entities[0]["aliases"]
+
+
+# ----------------------------------------------------------------------
+# _upsert_entities — merge path indexes rows positionally
+# ----------------------------------------------------------------------
+#
+# Regression for the 2026-06-11 prod incident: the pool was configured
+# with row_factory=dict_row while _upsert_entities did `row[0]`, so the
+# merge branch (entity already known) raised KeyError: 0 and every
+# extract that re-saw a known entity 500'd. Only never-seen-entity
+# events could store. Two guards:
+#   1. the pool must keep psycopg's default tuple row factory
+#      (matching extractor-async's worker, which also indexes
+#      positionally), and
+#   2. the merge branch must work against tuple rows end-to-end.
+
+import asyncio
+import inspect
+
+
+class _FakeCursor:
+    """Quacks like psycopg.AsyncCursor, returning TUPLE rows — the
+    shape the pool's default row factory produces. If the pool ever
+    grows a custom row_factory again, update this fake to match it or
+    test_pool_keeps_default_tuple_row_factory will flag the drift."""
+
+    def __init__(self, existing_id: str | None) -> None:
+        self.executed: list[tuple[str, object]] = []
+        self._existing_id = existing_id
+
+    async def execute(self, sql: str, params: object = None) -> None:
+        self.executed.append((" ".join(sql.split()), params))
+
+    async def fetchone(self):
+        return (self._existing_id,) if self._existing_id else None
+
+
+def _entity_stub() -> dict:
+    return {
+        "id": "e_new",
+        "arena": "arena1",
+        "entity_type": "person",
+        "canonical_name": "Alice One",
+        "aliases": ["Alice One", "alice@example.com"],
+        "provenance_event_ids": ["evt1"],
+        "participant_set": ["arena1"],
+        "disclosure_class": "private",
+    }
+
+
+def test_pool_keeps_default_tuple_row_factory() -> None:
+    src = inspect.getsource(sync_server.lifespan)
+    assert "row_factory" not in src, (
+        "extractor-sync's pool must use psycopg's default tuple rows: "
+        "_upsert_entities indexes fetchone() results positionally."
+    )
+
+
+def test_upsert_entities_merge_branch_with_tuple_rows() -> None:
+    """Entity already exists → UPDATE branch runs, id taken from row[0]."""
+    cur = _FakeCursor(existing_id="e_existing")
+    asyncio.run(sync_server._upsert_entities(cur, [_entity_stub()]))
+    updates = [(s, p) for s, p in cur.executed if s.startswith("UPDATE entities")]
+    assert len(updates) == 1
+    _, params = updates[0]
+    assert params[-1] == "e_existing"  # WHERE id = %s ← row[0]
+    assert not any(s.startswith("INSERT INTO entities") for s, _ in cur.executed)
+
+
+def test_upsert_entities_insert_branch_when_no_match() -> None:
+    cur = _FakeCursor(existing_id=None)
+    asyncio.run(sync_server._upsert_entities(cur, [_entity_stub()]))
+    inserts = [s for s, _ in cur.executed if s.startswith("INSERT INTO entities")]
+    assert len(inserts) == 1
+    assert not any(s.startswith("UPDATE entities") for s, _ in cur.executed)

package/packages/memory-engine-v2/org-model/migrations/004_source_kind_code_reference.sql

@@ -0,0 +1,12 @@
+-- 004: accept 'code_reference' source events (SDK corpus ingest).
+--
+-- The SDK corpus module (packages/memory/src/corpus/) emits events with
+-- source_kind='code_reference' (code-signature ingest, adapters.js).
+-- The enum predates that feature, so those events bounced with
+-- InvalidTextRepresentation and could never be stored — observed in
+-- prod 2026-06-11 as persistent /extract 500s + producer retry loops.
+--
+-- ALTER TYPE ... ADD VALUE cannot run inside a transaction block;
+-- apply with autocommit (psql's default per-statement behaviour).
+-- Applied manually to prod (pme2-org-model) on 2026-06-11.
+ALTER TYPE source_kind ADD VALUE IF NOT EXISTS 'code_reference';

package/packages/memory-engine-v2/org-model/migrations/005_fk_indexes.sql

@@ -0,0 +1,20 @@
+-- 005: index every column that references events(id).
+--
+-- events has four referencing constraints:
+--   distillation_queue.event_id   ON DELETE CASCADE
+--   vector_provenance.event_id    ON DELETE CASCADE
+--   distillation_traces.event_id  ON DELETE CASCADE
+--   events.forgets (self)         ON DELETE SET NULL
+--
+-- Postgres does NOT auto-index FK referencing columns. Without these,
+-- every DELETE on events seq-scans each referencing table per deleted
+-- row to enforce the constraint — the 2026-06-11 arena-scoped nuke of
+-- ~70k events ran for HOURS until the missing indexes were created
+-- on-box. (distillation_queue.event_id already had idx_distillation_
+-- event_id from 003; listed here for completeness via IF NOT EXISTS.)
+--
+-- All idempotent; applied manually to prod (pme2-org-model) 2026-06-12.
+CREATE INDEX IF NOT EXISTS idx_distillation_event_id ON distillation_queue(event_id);
+CREATE INDEX IF NOT EXISTS idx_traces_event_id ON distillation_traces(event_id);
+CREATE INDEX IF NOT EXISTS idx_vector_provenance_event_id ON vector_provenance(event_id);
+CREATE INDEX IF NOT EXISTS idx_events_forgets ON events(forgets);