@pentatonic-ai/ai-agent-sdk 0.10.38 → 0.10.39

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

@@ -116,6 +116,11 @@ def _envflag(name: str, default: str = "false") -> bool:
 # event_distillations ledger (migration 010) so escalation can be monitored and
 # student rows can be re-distilled later ("mop up").
 CASCADE_ENABLED = _envflag("CASCADE_ENABLED")
+# Async escalation (RFC-async-escalation.md). When true, the student-primary
+# pass does NOT call the teacher inline; escalated events become
+# 'escalated_pending' for the teacher-mop-up consumer. Default false = today's
+# synchronous inline cascade (instant revert).
+ASYNC_ESCALATION_ENABLED = os.environ.get("ASYNC_ESCALATION_ENABLED", "") not in ("", "0", "false")
 STUDENT_ENDPOINT = os.environ.get("STUDENT_ENDPOINT", "")
 STUDENT_API_KEY = os.environ.get("STUDENT_API_KEY", "")
 STUDENT_MODEL = os.environ.get("STUDENT_MODEL", "nuextract-2.0-4b-ft")
@@ -2622,7 +2627,7 @@ def claim_next_batch(conn: psycopg.Connection) -> list[dict[str, Any]]:
                     SELECT DISTINCT client FROM distillation_queue
                     WHERE (
                       status = 'pending'
-                      OR (status = 'claimed' AND claim_expires_at < NOW())
+                      OR (status = 'claimed' AND claim_expires_at < NOW() AND escalation_reason IS NULL)
                     ) AND attempts < %s
                   ) c
                   CROSS JOIN LATERAL (
@@ -2631,7 +2636,7 @@ def claim_next_batch(conn: psycopg.Connection) -> list[dict[str, Any]]:
                     WHERE client IS NOT DISTINCT FROM c.client
                       AND (
                         status = 'pending'
-                        OR (status = 'claimed' AND claim_expires_at < NOW())
+                        OR (status = 'claimed' AND claim_expires_at < NOW() AND escalation_reason IS NULL)
                       ) AND attempts < %s
                     ORDER BY COALESCE(priority, 100), emit_ts, id
                     LIMIT %s
@@ -2642,7 +2647,7 @@ def claim_next_batch(conn: psycopg.Connection) -> list[dict[str, Any]]:
               ))
               AND (
                 status = 'pending'
-                OR (status = 'claimed' AND claim_expires_at < NOW())
+                OR (status = 'claimed' AND claim_expires_at < NOW() AND escalation_reason IS NULL)
               ) AND attempts < %s
               FOR UPDATE SKIP LOCKED
             )
@@ -2654,6 +2659,47 @@ def claim_next_batch(conn: psycopg.Connection) -> list[dict[str, Any]]:
         return cur.fetchall()
 
 
+def claim_escalated_batch(
+    conn: psycopg.Connection,
+    worker_id: str,
+    batch_size: int,
+    claim_ttl_sec: int,
+    max_attempts: int,
+) -> list[dict[str, Any]]:
+    """Teacher-mop-up claim: lease `escalated_pending` rows (priority, emit_ts,
+    id order — no student fair-share needed, this is the teacher tier), leased
+    the same way as `claim_next_batch` (claimed/claimed_by/claimed_at/
+    claim_expires_at, FOR UPDATE SKIP LOCKED) so it is crash-safe. Also
+    re-claims expired-claimed rows whose prior state was escalated (a mop-up
+    worker that dies mid-batch leaves them `claimed`; once the lease expires
+    they must come back through THIS claim, not the student's `claim_next_batch`
+    — scoped via `escalation_reason IS NOT NULL`, which student-claimed rows
+    never carry). Served by idx_distillation_escalated."""
+    with conn.cursor(row_factory=psycopg.rows.dict_row) as cur:
+        cur.execute(
+            """
+            UPDATE distillation_queue SET
+              status = 'claimed',
+              claimed_by = %s,
+              claimed_at = NOW(),
+              claim_expires_at = NOW() + (%s || ' seconds')::interval
+            WHERE id IN (
+              SELECT id FROM distillation_queue
+              WHERE (
+                (status = 'escalated_pending')
+                OR (status = 'claimed' AND claim_expires_at < NOW() AND escalation_reason IS NOT NULL)
+              ) AND attempts < %s
+              ORDER BY COALESCE(priority, 100), emit_ts, id
+              LIMIT %s
+              FOR UPDATE SKIP LOCKED
+            )
+            RETURNING id, event_id, attempts, transient_attempts, escalation_reason
+            """,
+            (worker_id, claim_ttl_sec, max_attempts, batch_size),
+        )
+        return cur.fetchall()
+
+
 def fetch_event(conn: psycopg.Connection, event_id: str) -> dict[str, Any] | None:
     with conn.cursor(row_factory=psycopg.rows.dict_row) as cur:
         cur.execute(
@@ -2688,11 +2734,23 @@ def release_claim(conn: psycopg.Connection, queue_id: int, error: str) -> None:
     # Recoverable genuine-failure path (will retry) → count the attempt. This is
     # where the retry budget is spent — NOT at claim time — so a deploy-induced
     # reclaim never consumes it.
+    #
+    # status: CASE-guarded rather than a bare 'pending'. This function is
+    # reused by the teacher-mop-up path (_run_teacher, on non-exhausted
+    # failures) — a row claimed out of `escalated_pending` still carries
+    # `escalation_reason`. Sending it to bare 'pending' would let the
+    # STUDENT's claim_next_batch (unconditional status='pending' branch)
+    # claim and write it, silently defeating the escalation. Routing it back
+    # to 'escalated_pending' instead means only claim_escalated_batch can
+    # reclaim it. In SYNC mode escalation_reason is never set (see
+    # mark_escalated_pending / ASYNC_ESCALATION_ENABLED gate), so the CASE
+    # always yields 'pending' there — byte-for-byte identical to before.
     with conn.cursor() as cur:
         cur.execute(
             """
             UPDATE distillation_queue SET
-              status = 'pending',
+              status = CASE WHEN escalation_reason IS NOT NULL
+                       THEN 'escalated_pending' ELSE 'pending' END,
               claimed_by = NULL,
               claimed_at = NULL,
               claim_expires_at = NULL,
@@ -2704,6 +2762,44 @@ def release_claim(conn: psycopg.Connection, queue_id: int, error: str) -> None:
         )
 
 
+def mark_escalated_pending(
+    conn: psycopg.Connection, queue_id: int, reason: str
+) -> None:
+    """Student ran; a teacher-only pass is owed. Move the row to
+    'escalated_pending' WITHOUT touching `attempts` (the student didn't fail —
+    it deferred) and clear the student's claim so the teacher-mop-up consumer
+    can pick it up. The student's output is NOT written (deferred-write → the
+    teacher's write is the first and only write; no supersede)."""
+    with conn.cursor() as cur:
+        cur.execute(
+            """
+            UPDATE distillation_queue SET
+              status = 'escalated_pending',
+              escalation_reason = %s,
+              claimed_by = NULL,
+              claimed_at = NULL,
+              claim_expires_at = NULL
+            WHERE id = %s
+            """,
+            (reason, queue_id),
+        )
+
+
+def _route_escalations(
+    conn: psycopg.Connection,
+    escalate_items: list[dict[str, Any]],
+    reason_by_qid: dict[int, str],
+) -> list[dict[str, Any]]:
+    """Async: mark each escalated item 'escalated_pending' and return an empty
+    inline-teacher batch. Sync (flag off): return the items unchanged so the
+    caller runs the teacher inline exactly as before."""
+    if not ASYNC_ESCALATION_ENABLED:
+        return escalate_items
+    for item in escalate_items:
+        mark_escalated_pending(conn, item["id"], reason_by_qid.get(item["id"], "escalated"))
+    return []
+
+
 def _attempts_exhausted(attempts: int) -> bool:
     """Whether THIS processing failure should be terminal (mark_failed) rather
     than retried (release_claim). `attempts` is the row's value at claim time =
@@ -2759,11 +2855,20 @@ def requeue_transient(
     if transient_attempts + 1 >= MAX_TRANSIENT_ATTEMPTS:
         mark_failed(conn, queue_id, f"transient-exhausted: {error}")
         return
+    # status: CASE-guarded rather than a bare 'pending' — same reasoning as
+    # release_claim above. _run_teacher calls requeue_transient on ReadTimeout/
+    # 502/503/504 for BOTH student-primary and teacher-mop-up items; a row
+    # claimed out of `escalated_pending` still carries `escalation_reason` and
+    # must return there (only claim_escalated_batch may reclaim it), not to
+    # bare 'pending' where the student's claim_next_batch would grab and write
+    # it, silently defeating the escalation. In SYNC mode escalation_reason is
+    # never set, so the CASE always yields 'pending' — identical to before.
     with conn.cursor() as cur:
         cur.execute(
             """
             UPDATE distillation_queue SET
-              status = 'pending',
+              status = CASE WHEN escalation_reason IS NOT NULL
+                       THEN 'escalated_pending' ELSE 'pending' END,
               claimed_by = NULL,
               claimed_at = NULL,
               claim_expires_at = NULL,
@@ -2931,6 +3036,11 @@ async def _process_cascade(
         f"cascade: {len(callable_items) - len(escalate_items)} student-handled, "
         f"{len(escalate_items)} escalated to teacher"
     )
+    # RFC-async-escalation.md: flag-off ⇒ passthrough (identity-preserved), so
+    # the inline teacher call below runs exactly as before. Flag-on ⇒ each
+    # escalated item is marked 'escalated_pending' here and this returns [],
+    # so the inline teacher below is never called for this claim.
+    escalate_items = _route_escalations(conn, escalate_items, reason_by_qid)
     if escalate_items:
         await _run_teacher(http, conn, escalate_items, events_by_qid, False, reason_by_qid)
 
@@ -3206,18 +3316,68 @@ async def amain():
                 await process_batch(http, conn, items, stub_mode)
                 return len(items)
 
+        async def _run_one_mopup_batch() -> int:
+            """Teacher-mop-up: claim + process one batch of `escalated_pending`
+            rows on a DEDICATED connection (mirrors `_run_one_batch`). Only
+            scheduled when ASYNC_ESCALATION_ENABLED. Reuses `_run_teacher` —
+            the SAME teacher call → `_apply_extraction(producer="teacher")` →
+            trace → `mark_done` path the inline cascade escalation uses — so
+            there is exactly one teacher-write code path, invoked either
+            inline (sync escalation) or here (async drain). `reason_by_qid`
+            is populated from each row's `escalation_reason` (set by
+            `mark_escalated_pending` when the student deferred it) so the
+            cascade ledger still tags these as escalated."""
+            with psycopg.connect(PG_DSN, autocommit=True) as conn:
+                esc = claim_escalated_batch(
+                    conn, WORKER_ID, BATCH_SIZE, CLAIM_TTL_SEC, MAX_ATTEMPTS
+                )
+                if not esc:
+                    return 0
+                log.info(f"mop-up claimed {len(esc)} escalated item(s)")
+                events_by_qid: dict[int, dict[str, Any] | None] = {
+                    item["id"]: fetch_event(conn, item["event_id"]) for item in esc
+                }
+                reason_by_qid: dict[int, str] = {
+                    item["id"]: item.get("escalation_reason") or "escalated"
+                    for item in esc
+                }
+                # Mirror process_batch's missing-event guard — never hand
+                # _run_teacher an id whose event vanished.
+                callable_items = []
+                for item in esc:
+                    if events_by_qid[item["id"]] is None:
+                        log.warning(
+                            f"event {item['event_id']} missing — "
+                            f"marking queue {item['id']} done"
+                        )
+                        mark_done(conn, item["id"])
+                    else:
+                        callable_items.append(item)
+                if callable_items:
+                    await _run_teacher(
+                        http, conn, callable_items, events_by_qid, stub_mode, reason_by_qid
+                    )
+                return len(esc)
+
         # Keep up to MAX_INFLIGHT_BATCHES independent claim→apply pipelines
         # running at once. Their student/teacher calls all draw from the shared
         # pools above, so while one batch drains its slow tail the others keep
         # the GPUs fed. Total GPU concurrency is capped by the sems, NOT by this
-        # count.
+        # count. The teacher-mop-up drain (escalated_pending) is scheduled
+        # OUTSIDE this budget — one at a time, tracked separately — since it is
+        # the teacher tier, not the student fair-share pool; it only runs when
+        # ASYNC_ESCALATION_ENABLED.
         inflight: set[asyncio.Task] = set()
+        mopup_task: asyncio.Task | None = None
         while True:
             try:
                 while len(inflight) < MAX_INFLIGHT_BATCHES:
                     inflight.add(asyncio.create_task(_run_one_batch()))
-                done, inflight = await asyncio.wait(
-                    inflight, return_when=asyncio.FIRST_COMPLETED
+                if ASYNC_ESCALATION_ENABLED and mopup_task is None:
+                    mopup_task = asyncio.create_task(_run_one_mopup_batch())
+                wait_set = inflight | ({mopup_task} if mopup_task is not None else set())
+                done, pending = await asyncio.wait(
+                    wait_set, return_when=asyncio.FIRST_COMPLETED
                 )
                 claimed_any = False
                 for t in done:
@@ -3225,7 +3385,11 @@ async def amain():
                         if t.result() > 0:
                             claimed_any = True
                     except Exception as exc:
-                        log.error(f"batch task error: {exc}")
+                        label = "mop-up batch" if t is mopup_task else "batch"
+                        log.error(f"{label} task error: {exc}")
+                inflight = pending & inflight
+                if mopup_task is not None and mopup_task not in pending:
+                    mopup_task = None
                 # All finished batches were empty claims ⇒ queue drained; back
                 # off before re-claiming so idle workers don't hot-spin. Under a
                 # real backlog every claim returns items and this never sleeps.

package/packages/memory-engine-v2/org-model/migrations/014_escalated_pending.sql

@@ -0,0 +1,22 @@
+-- 014 — escalated_pending state for async escalation (RFC-async-escalation.md).
+--
+-- The student-primary consumer marks escalated events 'escalated_pending' (no
+-- write) instead of calling the teacher inline; the teacher-mop-up consumer
+-- claims and writes them once. Additive + safe online; apply BEFORE the code.
+
+ALTER TABLE distillation_queue
+  ADD COLUMN IF NOT EXISTS escalation_reason text;
+
+-- Recreate the status CHECK to admit the new value. NOT VALID + VALIDATE avoids
+-- the long ACCESS EXCLUSIVE validation lock on the large prod table.
+ALTER TABLE distillation_queue DROP CONSTRAINT IF EXISTS valid_status;
+ALTER TABLE distillation_queue ADD CONSTRAINT valid_status
+  CHECK (status = ANY (ARRAY['pending','claimed','done','failed','escalated_pending']))
+  NOT VALID;
+ALTER TABLE distillation_queue VALIDATE CONSTRAINT valid_status;
+
+-- Mop-up claim + autoscaler count. Mirrors idx_distillation_fair_claim shape.
+-- NOTE: CREATE INDEX CONCURRENTLY cannot run inside a transaction block.
+-- The schema manager applies this as a standalone statement.
+CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_distillation_escalated
+  ON distillation_queue (priority, emit_ts, id) WHERE status = 'escalated_pending';

package/packages/memory-engine-v2/scripts/redistill.py

@@ -49,6 +49,21 @@ from datetime import datetime, timezone
 
 CLEAN_PROMPT_HASH = "6ccfe70f1286a131"  # SDK 0.10.19 (#126 + #129); verify vs worker.SYSTEM_PROMPT_HASH
 
+# ASYNC_ESCALATION_ENABLED env-string check identical to extractor-async/worker.py
+# (RFC-async-escalation.md Task 6). When on, redistill's re-enqueue must go
+# straight to the teacher (escalated_pending) — the student-primary consumer
+# claims plain 'pending' rows and would otherwise hit these first, which is
+# the opposite of what a stale/dirty-hash re-distill wants.
+ASYNC_ESCALATION_ENABLED = os.environ.get("ASYNC_ESCALATION_ENABLED", "") not in ("", "0", "false")
+
+ENQUEUE_SQL = (
+    "INSERT INTO distillation_queue (event_id, status, attempts, priority, escalation_reason) "
+    "VALUES (%s, 'escalated_pending', 0, 200, 'redistill')"
+    if ASYNC_ESCALATION_ENABLED
+    else "INSERT INTO distillation_queue (event_id, status, attempts, priority) "
+         "VALUES (%s, 'pending', 0, 200)"
+)
+
 
 def _connect(dsn: str):
     import psycopg
@@ -97,7 +112,7 @@ def triage(cur, arena: str, clean_hash: str, dirty_hashes: list[str] | None, lim
               WHERE t.event_id = e.id AND t.system_prompt_hash = %s)
             AND NOT EXISTS (                          -- not already queued (so loops/batches advance)
               SELECT 1 FROM distillation_queue q
-              WHERE q.event_id = e.id AND q.status IN ('pending','claimed'))
+              WHERE q.event_id = e.id AND q.status IN ('pending','claimed','escalated_pending'))
             AND (
               EXISTS (SELECT 1 FROM distillation_traces t2     -- dirty teacher trace
                       WHERE t2.event_id = e.id {dirty_trace})
@@ -120,7 +135,7 @@ def triage(cur, arena: str, clean_hash: str, dirty_hashes: list[str] | None, lim
                     AND cardinality(f.provenance_event_ids) = 1) AS solo_fact_n,
                EXISTS (SELECT 1 FROM distillation_queue q
                        WHERE q.event_id = s.event_id
-                         AND q.status IN ('pending','claimed')) AS in_flight
+                         AND q.status IN ('pending','claimed','escalated_pending')) AS in_flight
         FROM stale s
         ORDER BY fact_n DESC
         """,
@@ -211,13 +226,14 @@ def main() -> int:
                     eid = c["event_id"]
                     if args.supersede_facts:
                         deleted += dump_and_delete_solo_facts(cur, eid, rb)
-                    cur.execute(
-                        # priority=200 (low): re-distills are maintenance — they
-                        # must never compete with live client traffic. Explicit
-                        # value wins over the enrichment trigger's mapping
-                        # (migration 012).
-                        "INSERT INTO distillation_queue (event_id, status, attempts, priority) "
-                        "VALUES (%s, 'pending', 0, 200)", (eid,))
+                    # priority=200 (low) in both modes: re-distills are
+                    # maintenance — they must never compete with live client
+                    # traffic. Explicit value wins over the enrichment
+                    # trigger's mapping (migration 012). Under
+                    # ASYNC_ESCALATION_ENABLED, ENQUEUE_SQL routes straight to
+                    # 'escalated_pending' (teacher-only, bypasses the
+                    # student) — see module docstring + Task 6.
+                    cur.execute(ENQUEUE_SQL, (eid,))
                     enq += 1
                     led.write(json.dumps({
                         "event_id": eid, "arena": args.arena,

package/packages/memory-engine-v2/scripts/test_redistill.py

@@ -0,0 +1,81 @@
+"""Guards redistill's enqueue against the latent "hits student first" gap
+(RFC-async-escalation.md Task 6): under async escalation, re-enqueued stale
+events must go straight to the teacher (`escalated_pending`), not `pending`
+(which the student-primary consumer would claim first). Pure-logic tests —
+no DB — importing `redistill` must not require psycopg (it's imported lazily
+inside `_connect`)."""
+from __future__ import annotations
+
+import importlib
+import os
+import sys
+
+HERE = os.path.dirname(__file__)
+
+
+def _load():
+    sys.path.insert(0, HERE)
+    sys.modules.pop("redistill", None)
+    return importlib.import_module("redistill")
+
+
+def test_redistill_enqueues_teacher_only(monkeypatch) -> None:
+    monkeypatch.setenv("ASYNC_ESCALATION_ENABLED", "1")
+    redistill = _load()
+    assert "escalated_pending" in redistill.ENQUEUE_SQL
+    assert "'redistill'" in redistill.ENQUEUE_SQL
+
+
+def test_redistill_enqueues_pending_when_flag_off(monkeypatch) -> None:
+    monkeypatch.delenv("ASYNC_ESCALATION_ENABLED", raising=False)
+    redistill = _load()
+    assert "'pending'" in redistill.ENQUEUE_SQL
+    assert "escalated_pending" not in redistill.ENQUEUE_SQL
+    assert "escalation_reason" not in redistill.ENQUEUE_SQL
+
+
+def test_redistill_flag_off_explicitly(monkeypatch) -> None:
+    for off_value in ("0", "false"):
+        monkeypatch.setenv("ASYNC_ESCALATION_ENABLED", off_value)
+        redistill = _load()
+        assert "'pending'" in redistill.ENQUEUE_SQL
+        assert "escalated_pending" not in redistill.ENQUEUE_SQL
+
+
+class _FakeCursor:
+    """Captures the SQL text passed to execute() without touching a DB."""
+
+    def __init__(self) -> None:
+        self.query: str | None = None
+        self.params: list | None = None
+
+    def execute(self, query, params=None) -> None:
+        self.query = query
+        self.params = params
+
+    def fetchall(self):
+        return []
+
+
+def test_triage_not_already_queued_filter_includes_escalated_pending(monkeypatch) -> None:
+    """The CTE's `NOT EXISTS (... status IN (...))` filter must also treat
+    `escalated_pending` as in-flight, or a redistill re-run before the teacher
+    mop-up drains the queue would re-select and duplicate-enqueue events
+    already sitting in escalated_pending (violates exactly-once)."""
+    monkeypatch.delenv("ASYNC_ESCALATION_ENABLED", raising=False)
+    redistill = _load()
+    cur = _FakeCursor()
+    redistill.triage(cur, "arena%", redistill.CLEAN_PROMPT_HASH, None, 50)
+    assert "AND q.status IN ('pending','claimed','escalated_pending')" in cur.query
+
+
+def test_triage_in_flight_column_includes_escalated_pending(monkeypatch) -> None:
+    """The `in_flight` column subquery must widen the same way, so an event
+    already in escalated_pending is correctly reported as in-flight (not
+    re-enqueued) rather than silently duplicated."""
+    monkeypatch.delenv("ASYNC_ESCALATION_ENABLED", raising=False)
+    redistill = _load()
+    cur = _FakeCursor()
+    redistill.triage(cur, "arena%", redistill.CLEAN_PROMPT_HASH, None, 50)
+    in_flight_clause = "AND q.status IN ('pending','claimed','escalated_pending')) AS in_flight"
+    assert in_flight_clause in cur.query