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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -116,6 +116,11 @@ def _envflag(name: str, default: str = "false") -> bool:
116
116
  # event_distillations ledger (migration 010) so escalation can be monitored and
117
117
  # student rows can be re-distilled later ("mop up").
118
118
  CASCADE_ENABLED = _envflag("CASCADE_ENABLED")
119
+ # Async escalation (RFC-async-escalation.md). When true, the student-primary
120
+ # pass does NOT call the teacher inline; escalated events become
121
+ # 'escalated_pending' for the teacher-mop-up consumer. Default false = today's
122
+ # synchronous inline cascade (instant revert).
123
+ ASYNC_ESCALATION_ENABLED = os.environ.get("ASYNC_ESCALATION_ENABLED", "") not in ("", "0", "false")
119
124
  STUDENT_ENDPOINT = os.environ.get("STUDENT_ENDPOINT", "")
120
125
  STUDENT_API_KEY = os.environ.get("STUDENT_API_KEY", "")
121
126
  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]]:
2622
2627
  SELECT DISTINCT client FROM distillation_queue
2623
2628
  WHERE (
2624
2629
  status = 'pending'
2625
- OR (status = 'claimed' AND claim_expires_at < NOW())
2630
+ OR (status = 'claimed' AND claim_expires_at < NOW() AND escalation_reason IS NULL)
2626
2631
  ) AND attempts < %s
2627
2632
  ) c
2628
2633
  CROSS JOIN LATERAL (
@@ -2631,7 +2636,7 @@ def claim_next_batch(conn: psycopg.Connection) -> list[dict[str, Any]]:
2631
2636
  WHERE client IS NOT DISTINCT FROM c.client
2632
2637
  AND (
2633
2638
  status = 'pending'
2634
- OR (status = 'claimed' AND claim_expires_at < NOW())
2639
+ OR (status = 'claimed' AND claim_expires_at < NOW() AND escalation_reason IS NULL)
2635
2640
  ) AND attempts < %s
2636
2641
  ORDER BY COALESCE(priority, 100), emit_ts, id
2637
2642
  LIMIT %s
@@ -2642,11 +2647,11 @@ def claim_next_batch(conn: psycopg.Connection) -> list[dict[str, Any]]:
2642
2647
  ))
2643
2648
  AND (
2644
2649
  status = 'pending'
2645
- OR (status = 'claimed' AND claim_expires_at < NOW())
2650
+ OR (status = 'claimed' AND claim_expires_at < NOW() AND escalation_reason IS NULL)
2646
2651
  ) AND attempts < %s
2647
2652
  FOR UPDATE SKIP LOCKED
2648
2653
  )
2649
- RETURNING id, event_id, attempts
2654
+ RETURNING id, event_id, attempts, transient_attempts
2650
2655
  """,
2651
2656
  (WORKER_ID, CLAIM_TTL_SEC, MAX_ATTEMPTS, MAX_ATTEMPTS,
2652
2657
  BATCH_SIZE, BATCH_SIZE, MAX_ATTEMPTS),
@@ -2654,6 +2659,47 @@ def claim_next_batch(conn: psycopg.Connection) -> list[dict[str, Any]]:
2654
2659
  return cur.fetchall()
2655
2660
 
2656
2661
 
2662
+ def claim_escalated_batch(
2663
+ conn: psycopg.Connection,
2664
+ worker_id: str,
2665
+ batch_size: int,
2666
+ claim_ttl_sec: int,
2667
+ max_attempts: int,
2668
+ ) -> list[dict[str, Any]]:
2669
+ """Teacher-mop-up claim: lease `escalated_pending` rows (priority, emit_ts,
2670
+ id order — no student fair-share needed, this is the teacher tier), leased
2671
+ the same way as `claim_next_batch` (claimed/claimed_by/claimed_at/
2672
+ claim_expires_at, FOR UPDATE SKIP LOCKED) so it is crash-safe. Also
2673
+ re-claims expired-claimed rows whose prior state was escalated (a mop-up
2674
+ worker that dies mid-batch leaves them `claimed`; once the lease expires
2675
+ they must come back through THIS claim, not the student's `claim_next_batch`
2676
+ — scoped via `escalation_reason IS NOT NULL`, which student-claimed rows
2677
+ never carry). Served by idx_distillation_escalated."""
2678
+ with conn.cursor(row_factory=psycopg.rows.dict_row) as cur:
2679
+ cur.execute(
2680
+ """
2681
+ UPDATE distillation_queue SET
2682
+ status = 'claimed',
2683
+ claimed_by = %s,
2684
+ claimed_at = NOW(),
2685
+ claim_expires_at = NOW() + (%s || ' seconds')::interval
2686
+ WHERE id IN (
2687
+ SELECT id FROM distillation_queue
2688
+ WHERE (
2689
+ (status = 'escalated_pending')
2690
+ OR (status = 'claimed' AND claim_expires_at < NOW() AND escalation_reason IS NOT NULL)
2691
+ ) AND attempts < %s
2692
+ ORDER BY COALESCE(priority, 100), emit_ts, id
2693
+ LIMIT %s
2694
+ FOR UPDATE SKIP LOCKED
2695
+ )
2696
+ RETURNING id, event_id, attempts, transient_attempts, escalation_reason
2697
+ """,
2698
+ (worker_id, claim_ttl_sec, max_attempts, batch_size),
2699
+ )
2700
+ return cur.fetchall()
2701
+
2702
+
2657
2703
  def fetch_event(conn: psycopg.Connection, event_id: str) -> dict[str, Any] | None:
2658
2704
  with conn.cursor(row_factory=psycopg.rows.dict_row) as cur:
2659
2705
  cur.execute(
@@ -2688,11 +2734,23 @@ def release_claim(conn: psycopg.Connection, queue_id: int, error: str) -> None:
2688
2734
  # Recoverable genuine-failure path (will retry) → count the attempt. This is
2689
2735
  # where the retry budget is spent — NOT at claim time — so a deploy-induced
2690
2736
  # reclaim never consumes it.
2737
+ #
2738
+ # status: CASE-guarded rather than a bare 'pending'. This function is
2739
+ # reused by the teacher-mop-up path (_run_teacher, on non-exhausted
2740
+ # failures) — a row claimed out of `escalated_pending` still carries
2741
+ # `escalation_reason`. Sending it to bare 'pending' would let the
2742
+ # STUDENT's claim_next_batch (unconditional status='pending' branch)
2743
+ # claim and write it, silently defeating the escalation. Routing it back
2744
+ # to 'escalated_pending' instead means only claim_escalated_batch can
2745
+ # reclaim it. In SYNC mode escalation_reason is never set (see
2746
+ # mark_escalated_pending / ASYNC_ESCALATION_ENABLED gate), so the CASE
2747
+ # always yields 'pending' there — byte-for-byte identical to before.
2691
2748
  with conn.cursor() as cur:
2692
2749
  cur.execute(
2693
2750
  """
2694
2751
  UPDATE distillation_queue SET
2695
- status = 'pending',
2752
+ status = CASE WHEN escalation_reason IS NOT NULL
2753
+ THEN 'escalated_pending' ELSE 'pending' END,
2696
2754
  claimed_by = NULL,
2697
2755
  claimed_at = NULL,
2698
2756
  claim_expires_at = NULL,
@@ -2704,6 +2762,44 @@ def release_claim(conn: psycopg.Connection, queue_id: int, error: str) -> None:
2704
2762
  )
2705
2763
 
2706
2764
 
2765
+ def mark_escalated_pending(
2766
+ conn: psycopg.Connection, queue_id: int, reason: str
2767
+ ) -> None:
2768
+ """Student ran; a teacher-only pass is owed. Move the row to
2769
+ 'escalated_pending' WITHOUT touching `attempts` (the student didn't fail —
2770
+ it deferred) and clear the student's claim so the teacher-mop-up consumer
2771
+ can pick it up. The student's output is NOT written (deferred-write → the
2772
+ teacher's write is the first and only write; no supersede)."""
2773
+ with conn.cursor() as cur:
2774
+ cur.execute(
2775
+ """
2776
+ UPDATE distillation_queue SET
2777
+ status = 'escalated_pending',
2778
+ escalation_reason = %s,
2779
+ claimed_by = NULL,
2780
+ claimed_at = NULL,
2781
+ claim_expires_at = NULL
2782
+ WHERE id = %s
2783
+ """,
2784
+ (reason, queue_id),
2785
+ )
2786
+
2787
+
2788
+ def _route_escalations(
2789
+ conn: psycopg.Connection,
2790
+ escalate_items: list[dict[str, Any]],
2791
+ reason_by_qid: dict[int, str],
2792
+ ) -> list[dict[str, Any]]:
2793
+ """Async: mark each escalated item 'escalated_pending' and return an empty
2794
+ inline-teacher batch. Sync (flag off): return the items unchanged so the
2795
+ caller runs the teacher inline exactly as before."""
2796
+ if not ASYNC_ESCALATION_ENABLED:
2797
+ return escalate_items
2798
+ for item in escalate_items:
2799
+ mark_escalated_pending(conn, item["id"], reason_by_qid.get(item["id"], "escalated"))
2800
+ return []
2801
+
2802
+
2707
2803
  def _attempts_exhausted(attempts: int) -> bool:
2708
2804
  """Whether THIS processing failure should be terminal (mark_failed) rather
2709
2805
  than retried (release_claim). `attempts` is the row's value at claim time =
@@ -2713,6 +2809,77 @@ def _attempts_exhausted(attempts: int) -> bool:
2713
2809
  return attempts + 1 >= MAX_ATTEMPTS
2714
2810
 
2715
2811
 
2812
+ # ── transient vs genuine failures ─────────────────────────────────────
2813
+ # Transient = infra hiccups that resolve on their own: the teacher timing out
2814
+ # under load (ReadTimeout) or the teacher-lb returning 502/503/504 while the
2815
+ # fleet is cold-starting from zero. These must NOT burn the genuine `attempts`
2816
+ # budget (3) — 3 quick ReadTimeouts used to march a row to `failed`, and failed
2817
+ # rows are never requeued, so the extraction was silently lost. Instead they get
2818
+ # their own larger budget (MAX_TRANSIENT_ATTEMPTS) on the `transient_attempts`
2819
+ # counter (migration 013), leaving `attempts` untouched so the row stays
2820
+ # claimable (`attempts < MAX_ATTEMPTS` in claim_next_batch). ReadTimeouts are
2821
+ # self-spaced by LLM_TIMEOUT_SEC and the fair-share round-robin spreads re-claims
2822
+ # across the backlog, so this does not hot-loop; MAX_TRANSIENT_ATTEMPTS bounds a
2823
+ # genuine poison chunk (one that times out even on a healthy multi-box fleet).
2824
+ MAX_TRANSIENT_ATTEMPTS = int(os.environ.get("MAX_TRANSIENT_ATTEMPTS", "12"))
2825
+ _TRANSIENT_MARKERS = (
2826
+ "ReadTimeout",
2827
+ "TimeoutException",
2828
+ "ConnectTimeout",
2829
+ "ConnectError",
2830
+ "PoolTimeout",
2831
+ "502",
2832
+ "503",
2833
+ "504",
2834
+ "Bad Gateway",
2835
+ "Gateway Time-out",
2836
+ "Service Unavailable",
2837
+ "Temporarily Unavailable",
2838
+ )
2839
+
2840
+
2841
+ def _is_transient(err: str) -> bool:
2842
+ """True if `err` looks like a recoverable infra hiccup (teacher timeout /
2843
+ LB gateway error) rather than a genuine extraction/DB failure."""
2844
+ return any(m in err for m in _TRANSIENT_MARKERS)
2845
+
2846
+
2847
+ def requeue_transient(
2848
+ conn: psycopg.Connection, queue_id: int, transient_attempts: int, error: str
2849
+ ) -> None:
2850
+ """Recoverable INFRA failure. Release the claim back to pending WITHOUT
2851
+ touching `attempts` (preserves the genuine-failure budget), counting instead
2852
+ against `transient_attempts`. Once that budget is exhausted the row IS marked
2853
+ failed — a chunk that times out this many times is genuinely unprocessable,
2854
+ not a passing hiccup."""
2855
+ if transient_attempts + 1 >= MAX_TRANSIENT_ATTEMPTS:
2856
+ mark_failed(conn, queue_id, f"transient-exhausted: {error}")
2857
+ return
2858
+ # status: CASE-guarded rather than a bare 'pending' — same reasoning as
2859
+ # release_claim above. _run_teacher calls requeue_transient on ReadTimeout/
2860
+ # 502/503/504 for BOTH student-primary and teacher-mop-up items; a row
2861
+ # claimed out of `escalated_pending` still carries `escalation_reason` and
2862
+ # must return there (only claim_escalated_batch may reclaim it), not to
2863
+ # bare 'pending' where the student's claim_next_batch would grab and write
2864
+ # it, silently defeating the escalation. In SYNC mode escalation_reason is
2865
+ # never set, so the CASE always yields 'pending' — identical to before.
2866
+ with conn.cursor() as cur:
2867
+ cur.execute(
2868
+ """
2869
+ UPDATE distillation_queue SET
2870
+ status = CASE WHEN escalation_reason IS NOT NULL
2871
+ THEN 'escalated_pending' ELSE 'pending' END,
2872
+ claimed_by = NULL,
2873
+ claimed_at = NULL,
2874
+ claim_expires_at = NULL,
2875
+ transient_attempts = transient_attempts + 1,
2876
+ last_error = %s
2877
+ WHERE id = %s
2878
+ """,
2879
+ (error[:1024], queue_id),
2880
+ )
2881
+
2882
+
2716
2883
  # --------------------------------------------------------------------
2717
2884
  # Main loop
2718
2885
  # --------------------------------------------------------------------
@@ -2869,6 +3036,11 @@ async def _process_cascade(
2869
3036
  f"cascade: {len(callable_items) - len(escalate_items)} student-handled, "
2870
3037
  f"{len(escalate_items)} escalated to teacher"
2871
3038
  )
3039
+ # RFC-async-escalation.md: flag-off ⇒ passthrough (identity-preserved), so
3040
+ # the inline teacher call below runs exactly as before. Flag-on ⇒ each
3041
+ # escalated item is marked 'escalated_pending' here and this returns [],
3042
+ # so the inline teacher below is never called for this claim.
3043
+ escalate_items = _route_escalations(conn, escalate_items, reason_by_qid)
2872
3044
  if escalate_items:
2873
3045
  await _run_teacher(http, conn, escalate_items, events_by_qid, False, reason_by_qid)
2874
3046
 
@@ -2924,7 +3096,11 @@ async def _run_teacher(
2924
3096
  log.warning(
2925
3097
  f"extraction failed queue_id={queue_id} attempts={attempts}: {err}"
2926
3098
  )
2927
- if _attempts_exhausted(attempts):
3099
+ if _is_transient(err):
3100
+ requeue_transient(
3101
+ conn, queue_id, item.get("transient_attempts", 0), err
3102
+ )
3103
+ elif _attempts_exhausted(attempts):
2928
3104
  mark_failed(conn, queue_id, err)
2929
3105
  else:
2930
3106
  release_claim(conn, queue_id, err)
@@ -3022,8 +3198,13 @@ def _apply_extraction(
3022
3198
  f"db upsert failed queue_id={queue_id} attempts={attempts}: {err}"
3023
3199
  )
3024
3200
  # The transaction rolled back; the claim bookkeeping runs on the
3025
- # now-autocommit connection so it persists independently.
3026
- if _attempts_exhausted(attempts):
3201
+ # now-autocommit connection so it persists independently. A transient DB
3202
+ # blip (connection/pool timeout) requeues without burning the genuine
3203
+ # budget; real errors (UniqueViolation, schema mismatch) don't match
3204
+ # _is_transient and take the normal attempts path.
3205
+ if _is_transient(err):
3206
+ requeue_transient(conn, queue_id, item.get("transient_attempts", 0), err)
3207
+ elif _attempts_exhausted(attempts):
3027
3208
  mark_failed(conn, queue_id, err)
3028
3209
  else:
3029
3210
  release_claim(conn, queue_id, err)
@@ -3135,18 +3316,68 @@ async def amain():
3135
3316
  await process_batch(http, conn, items, stub_mode)
3136
3317
  return len(items)
3137
3318
 
3319
+ async def _run_one_mopup_batch() -> int:
3320
+ """Teacher-mop-up: claim + process one batch of `escalated_pending`
3321
+ rows on a DEDICATED connection (mirrors `_run_one_batch`). Only
3322
+ scheduled when ASYNC_ESCALATION_ENABLED. Reuses `_run_teacher` —
3323
+ the SAME teacher call → `_apply_extraction(producer="teacher")` →
3324
+ trace → `mark_done` path the inline cascade escalation uses — so
3325
+ there is exactly one teacher-write code path, invoked either
3326
+ inline (sync escalation) or here (async drain). `reason_by_qid`
3327
+ is populated from each row's `escalation_reason` (set by
3328
+ `mark_escalated_pending` when the student deferred it) so the
3329
+ cascade ledger still tags these as escalated."""
3330
+ with psycopg.connect(PG_DSN, autocommit=True) as conn:
3331
+ esc = claim_escalated_batch(
3332
+ conn, WORKER_ID, BATCH_SIZE, CLAIM_TTL_SEC, MAX_ATTEMPTS
3333
+ )
3334
+ if not esc:
3335
+ return 0
3336
+ log.info(f"mop-up claimed {len(esc)} escalated item(s)")
3337
+ events_by_qid: dict[int, dict[str, Any] | None] = {
3338
+ item["id"]: fetch_event(conn, item["event_id"]) for item in esc
3339
+ }
3340
+ reason_by_qid: dict[int, str] = {
3341
+ item["id"]: item.get("escalation_reason") or "escalated"
3342
+ for item in esc
3343
+ }
3344
+ # Mirror process_batch's missing-event guard — never hand
3345
+ # _run_teacher an id whose event vanished.
3346
+ callable_items = []
3347
+ for item in esc:
3348
+ if events_by_qid[item["id"]] is None:
3349
+ log.warning(
3350
+ f"event {item['event_id']} missing — "
3351
+ f"marking queue {item['id']} done"
3352
+ )
3353
+ mark_done(conn, item["id"])
3354
+ else:
3355
+ callable_items.append(item)
3356
+ if callable_items:
3357
+ await _run_teacher(
3358
+ http, conn, callable_items, events_by_qid, stub_mode, reason_by_qid
3359
+ )
3360
+ return len(esc)
3361
+
3138
3362
  # Keep up to MAX_INFLIGHT_BATCHES independent claim→apply pipelines
3139
3363
  # running at once. Their student/teacher calls all draw from the shared
3140
3364
  # pools above, so while one batch drains its slow tail the others keep
3141
3365
  # the GPUs fed. Total GPU concurrency is capped by the sems, NOT by this
3142
- # count.
3366
+ # count. The teacher-mop-up drain (escalated_pending) is scheduled
3367
+ # OUTSIDE this budget — one at a time, tracked separately — since it is
3368
+ # the teacher tier, not the student fair-share pool; it only runs when
3369
+ # ASYNC_ESCALATION_ENABLED.
3143
3370
  inflight: set[asyncio.Task] = set()
3371
+ mopup_task: asyncio.Task | None = None
3144
3372
  while True:
3145
3373
  try:
3146
3374
  while len(inflight) < MAX_INFLIGHT_BATCHES:
3147
3375
  inflight.add(asyncio.create_task(_run_one_batch()))
3148
- done, inflight = await asyncio.wait(
3149
- inflight, return_when=asyncio.FIRST_COMPLETED
3376
+ if ASYNC_ESCALATION_ENABLED and mopup_task is None:
3377
+ mopup_task = asyncio.create_task(_run_one_mopup_batch())
3378
+ wait_set = inflight | ({mopup_task} if mopup_task is not None else set())
3379
+ done, pending = await asyncio.wait(
3380
+ wait_set, return_when=asyncio.FIRST_COMPLETED
3150
3381
  )
3151
3382
  claimed_any = False
3152
3383
  for t in done:
@@ -3154,7 +3385,11 @@ async def amain():
3154
3385
  if t.result() > 0:
3155
3386
  claimed_any = True
3156
3387
  except Exception as exc:
3157
- log.error(f"batch task error: {exc}")
3388
+ label = "mop-up batch" if t is mopup_task else "batch"
3389
+ log.error(f"{label} task error: {exc}")
3390
+ inflight = pending & inflight
3391
+ if mopup_task is not None and mopup_task not in pending:
3392
+ mopup_task = None
3158
3393
  # All finished batches were empty claims ⇒ queue drained; back
3159
3394
  # off before re-claiming so idle workers don't hot-spin. Under a
3160
3395
  # real backlog every claim returns items and this never sleeps.
@@ -0,0 +1,27 @@
1
+ -- 013 — separate the TRANSIENT retry budget from the genuine-failure budget.
2
+ --
3
+ -- The distiller shares one `attempts` counter for every failure, capped at
4
+ -- MAX_ATTEMPTS (3) and enforced by the claim query's `attempts < MAX_ATTEMPTS`
5
+ -- filter. That conflates two very different failures:
6
+ -- * genuine — parse error, NULL-subject fact, schema mismatch: real, should
7
+ -- give up after a few tries.
8
+ -- * transient — teacher ReadTimeout under load, teacher-lb 502/503/504 during
9
+ -- a scale-from-zero window: infra hiccups that resolve on their
10
+ -- own once the fleet catches up / a box finishes booting.
11
+ --
12
+ -- Because they shared a counter, a transient hiccup burned the genuine budget:
13
+ -- 3 quick ReadTimeouts → `failed`, and since failed rows are never requeued the
14
+ -- extraction was silently LOST (measured 154 ReadTimeout losses in a single day
15
+ -- during a bulk-import whale, 2026-07-14).
16
+ --
17
+ -- This adds a dedicated `transient_attempts` counter. worker.py routes transient
18
+ -- errors here (incrementing THIS, leaving `attempts` untouched so the row stays
19
+ -- claimable — `attempts` remains < MAX_ATTEMPTS) up to MAX_TRANSIENT_ATTEMPTS,
20
+ -- after which the row is marked failed like any genuinely stuck work. Genuine
21
+ -- errors keep using `attempts` exactly as before.
22
+ --
23
+ -- Metadata-only on PostgreSQL (constant DEFAULT, no table rewrite) — safe to run
24
+ -- online against the ~250k-row prod table. Idempotent.
25
+
26
+ ALTER TABLE distillation_queue
27
+ ADD COLUMN IF NOT EXISTS transient_attempts integer NOT NULL DEFAULT 0;
@@ -0,0 +1,22 @@
1
+ -- 014 — escalated_pending state for async escalation (RFC-async-escalation.md).
2
+ --
3
+ -- The student-primary consumer marks escalated events 'escalated_pending' (no
4
+ -- write) instead of calling the teacher inline; the teacher-mop-up consumer
5
+ -- claims and writes them once. Additive + safe online; apply BEFORE the code.
6
+
7
+ ALTER TABLE distillation_queue
8
+ ADD COLUMN IF NOT EXISTS escalation_reason text;
9
+
10
+ -- Recreate the status CHECK to admit the new value. NOT VALID + VALIDATE avoids
11
+ -- the long ACCESS EXCLUSIVE validation lock on the large prod table.
12
+ ALTER TABLE distillation_queue DROP CONSTRAINT IF EXISTS valid_status;
13
+ ALTER TABLE distillation_queue ADD CONSTRAINT valid_status
14
+ CHECK (status = ANY (ARRAY['pending','claimed','done','failed','escalated_pending']))
15
+ NOT VALID;
16
+ ALTER TABLE distillation_queue VALIDATE CONSTRAINT valid_status;
17
+
18
+ -- Mop-up claim + autoscaler count. Mirrors idx_distillation_fair_claim shape.
19
+ -- NOTE: CREATE INDEX CONCURRENTLY cannot run inside a transaction block.
20
+ -- The schema manager applies this as a standalone statement.
21
+ CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_distillation_escalated
22
+ ON distillation_queue (priority, emit_ts, id) WHERE status = 'escalated_pending';
@@ -49,6 +49,21 @@ from datetime import datetime, timezone
49
49
 
50
50
  CLEAN_PROMPT_HASH = "6ccfe70f1286a131" # SDK 0.10.19 (#126 + #129); verify vs worker.SYSTEM_PROMPT_HASH
51
51
 
52
+ # ASYNC_ESCALATION_ENABLED env-string check identical to extractor-async/worker.py
53
+ # (RFC-async-escalation.md Task 6). When on, redistill's re-enqueue must go
54
+ # straight to the teacher (escalated_pending) — the student-primary consumer
55
+ # claims plain 'pending' rows and would otherwise hit these first, which is
56
+ # the opposite of what a stale/dirty-hash re-distill wants.
57
+ ASYNC_ESCALATION_ENABLED = os.environ.get("ASYNC_ESCALATION_ENABLED", "") not in ("", "0", "false")
58
+
59
+ ENQUEUE_SQL = (
60
+ "INSERT INTO distillation_queue (event_id, status, attempts, priority, escalation_reason) "
61
+ "VALUES (%s, 'escalated_pending', 0, 200, 'redistill')"
62
+ if ASYNC_ESCALATION_ENABLED
63
+ else "INSERT INTO distillation_queue (event_id, status, attempts, priority) "
64
+ "VALUES (%s, 'pending', 0, 200)"
65
+ )
66
+
52
67
 
53
68
  def _connect(dsn: str):
54
69
  import psycopg
@@ -97,7 +112,7 @@ def triage(cur, arena: str, clean_hash: str, dirty_hashes: list[str] | None, lim
97
112
  WHERE t.event_id = e.id AND t.system_prompt_hash = %s)
98
113
  AND NOT EXISTS ( -- not already queued (so loops/batches advance)
99
114
  SELECT 1 FROM distillation_queue q
100
- WHERE q.event_id = e.id AND q.status IN ('pending','claimed'))
115
+ WHERE q.event_id = e.id AND q.status IN ('pending','claimed','escalated_pending'))
101
116
  AND (
102
117
  EXISTS (SELECT 1 FROM distillation_traces t2 -- dirty teacher trace
103
118
  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
120
135
  AND cardinality(f.provenance_event_ids) = 1) AS solo_fact_n,
121
136
  EXISTS (SELECT 1 FROM distillation_queue q
122
137
  WHERE q.event_id = s.event_id
123
- AND q.status IN ('pending','claimed')) AS in_flight
138
+ AND q.status IN ('pending','claimed','escalated_pending')) AS in_flight
124
139
  FROM stale s
125
140
  ORDER BY fact_n DESC
126
141
  """,
@@ -211,13 +226,14 @@ def main() -> int:
211
226
  eid = c["event_id"]
212
227
  if args.supersede_facts:
213
228
  deleted += dump_and_delete_solo_facts(cur, eid, rb)
214
- cur.execute(
215
- # priority=200 (low): re-distills are maintenance they
216
- # must never compete with live client traffic. Explicit
217
- # value wins over the enrichment trigger's mapping
218
- # (migration 012).
219
- "INSERT INTO distillation_queue (event_id, status, attempts, priority) "
220
- "VALUES (%s, 'pending', 0, 200)", (eid,))
229
+ # priority=200 (low) in both modes: re-distills are
230
+ # maintenance they must never compete with live client
231
+ # traffic. Explicit value wins over the enrichment
232
+ # trigger's mapping (migration 012). Under
233
+ # ASYNC_ESCALATION_ENABLED, ENQUEUE_SQL routes straight to
234
+ # 'escalated_pending' (teacher-only, bypasses the
235
+ # student) see module docstring + Task 6.
236
+ cur.execute(ENQUEUE_SQL, (eid,))
221
237
  enq += 1
222
238
  led.write(json.dumps({
223
239
  "event_id": eid, "arena": args.arena,
@@ -0,0 +1,81 @@
1
+ """Guards redistill's enqueue against the latent "hits student first" gap
2
+ (RFC-async-escalation.md Task 6): under async escalation, re-enqueued stale
3
+ events must go straight to the teacher (`escalated_pending`), not `pending`
4
+ (which the student-primary consumer would claim first). Pure-logic tests —
5
+ no DB — importing `redistill` must not require psycopg (it's imported lazily
6
+ inside `_connect`)."""
7
+ from __future__ import annotations
8
+
9
+ import importlib
10
+ import os
11
+ import sys
12
+
13
+ HERE = os.path.dirname(__file__)
14
+
15
+
16
+ def _load():
17
+ sys.path.insert(0, HERE)
18
+ sys.modules.pop("redistill", None)
19
+ return importlib.import_module("redistill")
20
+
21
+
22
+ def test_redistill_enqueues_teacher_only(monkeypatch) -> None:
23
+ monkeypatch.setenv("ASYNC_ESCALATION_ENABLED", "1")
24
+ redistill = _load()
25
+ assert "escalated_pending" in redistill.ENQUEUE_SQL
26
+ assert "'redistill'" in redistill.ENQUEUE_SQL
27
+
28
+
29
+ def test_redistill_enqueues_pending_when_flag_off(monkeypatch) -> None:
30
+ monkeypatch.delenv("ASYNC_ESCALATION_ENABLED", raising=False)
31
+ redistill = _load()
32
+ assert "'pending'" in redistill.ENQUEUE_SQL
33
+ assert "escalated_pending" not in redistill.ENQUEUE_SQL
34
+ assert "escalation_reason" not in redistill.ENQUEUE_SQL
35
+
36
+
37
+ def test_redistill_flag_off_explicitly(monkeypatch) -> None:
38
+ for off_value in ("0", "false"):
39
+ monkeypatch.setenv("ASYNC_ESCALATION_ENABLED", off_value)
40
+ redistill = _load()
41
+ assert "'pending'" in redistill.ENQUEUE_SQL
42
+ assert "escalated_pending" not in redistill.ENQUEUE_SQL
43
+
44
+
45
+ class _FakeCursor:
46
+ """Captures the SQL text passed to execute() without touching a DB."""
47
+
48
+ def __init__(self) -> None:
49
+ self.query: str | None = None
50
+ self.params: list | None = None
51
+
52
+ def execute(self, query, params=None) -> None:
53
+ self.query = query
54
+ self.params = params
55
+
56
+ def fetchall(self):
57
+ return []
58
+
59
+
60
+ def test_triage_not_already_queued_filter_includes_escalated_pending(monkeypatch) -> None:
61
+ """The CTE's `NOT EXISTS (... status IN (...))` filter must also treat
62
+ `escalated_pending` as in-flight, or a redistill re-run before the teacher
63
+ mop-up drains the queue would re-select and duplicate-enqueue events
64
+ already sitting in escalated_pending (violates exactly-once)."""
65
+ monkeypatch.delenv("ASYNC_ESCALATION_ENABLED", raising=False)
66
+ redistill = _load()
67
+ cur = _FakeCursor()
68
+ redistill.triage(cur, "arena%", redistill.CLEAN_PROMPT_HASH, None, 50)
69
+ assert "AND q.status IN ('pending','claimed','escalated_pending')" in cur.query
70
+
71
+
72
+ def test_triage_in_flight_column_includes_escalated_pending(monkeypatch) -> None:
73
+ """The `in_flight` column subquery must widen the same way, so an event
74
+ already in escalated_pending is correctly reported as in-flight (not
75
+ re-enqueued) rather than silently duplicated."""
76
+ monkeypatch.delenv("ASYNC_ESCALATION_ENABLED", raising=False)
77
+ redistill = _load()
78
+ cur = _FakeCursor()
79
+ redistill.triage(cur, "arena%", redistill.CLEAN_PROMPT_HASH, None, 50)
80
+ in_flight_clause = "AND q.status IN ('pending','claimed','escalated_pending')) AS in_flight"
81
+ assert in_flight_clause in cur.query