diff --git a/tests/test_judge_storage.py b/tests/test_judge_storage.py index ff6f58b5..a333e392 100644 --- a/tests/test_judge_storage.py +++ b/tests/test_judge_storage.py @@ -287,6 +287,40 @@ class TestIntentVerdictBulkInsert: assert v1["risk_level"] == "low" and v1["tier"] == "heuristic" assert v2["risk_level"] == "high" and v2["tier"] == "llm" + def test_bulk_insert_pk_collision_skips_only_colliding_row(self, db): + """Regression: the async judge daemon can UPSERT a fallback row — + reusing a heuristic verdict_id from the incoming batch — BEFORE + ``approve_tools`` runs the bulk write. The bulk insert must skip + just that row (keeping the daemon's tier upgrade) instead of + aborting the whole statement and silently losing every sibling + row in the batch.""" + # Daemon won the race: fallback row already sits on b2's PK. + db.upsert_intent_verdict( + **_make_verdict_kwargs( + verdict_id="b2", + call_id="c2", + tier="llm_fallback", + judge_model="judge-model", + ) + ) + db.create_intent_verdicts_bulk( + [ + _make_verdict_kwargs(verdict_id="b1", call_id="c1"), + _make_verdict_kwargs(verdict_id="b2", call_id="c2"), # collides + _make_verdict_kwargs(verdict_id="b3", call_id="c3"), + ] + ) + # Siblings landed despite the mid-batch collision. + for vid in ("b1", "b3"): + v = db.get_intent_verdict(vid) + assert v is not None, f"sibling row {vid} lost to the collision" + assert v["tier"] == "heuristic" + # The colliding row kept the daemon's upgrade, not the bulk stamp. + v2 = db.get_intent_verdict("b2") + assert v2 is not None + assert v2["tier"] == "llm_fallback" + assert v2["judge_model"] == "judge-model" + # --------------------------------------------------------------------------- # List queries diff --git a/turnstone/core/session_ui_base.py b/turnstone/core/session_ui_base.py index 14f305ba..72ae2c91 100644 --- a/turnstone/core/session_ui_base.py +++ b/turnstone/core/session_ui_base.py @@ -1480,25 +1480,22 @@ class SessionUIBase: } for v in verdicts ] - # Plain INSERT (not UPSERT) at the bulk site. The race - # where a daemon-judge verdict lands BEFORE this bulk - # write IS reachable today: ``_evaluate_intent`` - # (session.py) spawns the daemon thread before - # ``approve_tools`` is called, and the daemon's first - # emission (heuristic-only short batch, fast LLM response, - # or cancel-event ``_deliver_fallbacks`` from judge.py) - # can fire ``_persist_intent_verdict`` before this bulk - # INSERT runs. Outcome of that race is unchanged by the - # per-row UPSERT switch: the bulk INSERT statement aborts - # on PK collision regardless of whether the colliding row - # was planted by INSERT or UPSERT, and the wrapping - # ``try/except`` swallows it. Race A (daemon fires AFTER - # bulk) IS improved by the fix: heuristic→llm_fallback - # upgrade-in-place now lands. Future bulk-side hardening - # (``ON CONFLICT DO NOTHING``) would preserve the OTHER - # rows in the batch when one collides, but would keep the - # daemon's ``tier`` ("llm"/"llm_fallback") for the - # colliding row instead of the bulk's heuristic stamp. + # The daemon-judge race where a verdict lands BEFORE this + # bulk write IS reachable: ``_evaluate_intent`` (session.py) + # spawns the daemon thread before ``approve_tools`` is + # called, and the daemon's first emission (heuristic-only + # short batch, fast LLM response, or cancel-event + # ``_deliver_fallbacks`` from judge.py) can fire + # ``_persist_intent_verdict`` first — a fallback UPSERT + # plants the very ``verdict_id`` this batch is about to + # INSERT. The bulk site inserts ``ON CONFLICT DO NOTHING`` + # so that one collision skips only its own row: the rest of + # the batch still lands, and the colliding row keeps the + # daemon's ``llm_fallback`` tier upgrade instead of being + # regressed to the heuristic stamp. (Plain INSERT here used + # to abort the entire statement — and the ``try/except`` + # below swallowed it — discarding the whole batch's + # heuristic rows.) storage.create_intent_verdicts_bulk(rows) except Exception: log.debug("Failed to bulk-persist intent verdicts", exc_info=True) diff --git a/turnstone/core/storage/_postgresql.py b/turnstone/core/storage/_postgresql.py index d766239f..8a6f46fe 100644 --- a/turnstone/core/storage/_postgresql.py +++ b/turnstone/core/storage/_postgresql.py @@ -3642,6 +3642,16 @@ class PostgreSQLBackend: conn.commit() def create_intent_verdicts_bulk(self, verdicts: list[dict[str, Any]]) -> None: + # ON CONFLICT DO NOTHING per row: the async judge daemon can land + # an UPSERT (which INSERTs the heuristic verdict_id it reuses) + # before this bulk write runs. A plain INSERT would abort the + # whole statement on that one collision — silently discarding + # every OTHER row in the batch (the caller swallows storage + # errors). Skipping just the colliding row also keeps the + # daemon's tier upgrade ("llm_fallback") instead of regressing + # it to the bulk's heuristic stamp. + from sqlalchemy.dialects.postgresql import insert as pg_insert + if not verdicts: return now = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S") @@ -3667,7 +3677,10 @@ class PostgreSQLBackend: for v in verdicts ] with self._conn() as conn: - conn.execute(sa.insert(intent_verdicts), rows) + conn.execute( + pg_insert(intent_verdicts).on_conflict_do_nothing(index_elements=["verdict_id"]), + rows, + ) conn.commit() def get_intent_verdict(self, verdict_id: str) -> dict[str, Any] | None: diff --git a/turnstone/core/storage/_protocol.py b/turnstone/core/storage/_protocol.py index 86edda74..7b054dd8 100644 --- a/turnstone/core/storage/_protocol.py +++ b/turnstone/core/storage/_protocol.py @@ -1736,8 +1736,10 @@ class StorageBackend(Protocol): Used by :meth:`SessionUIBase._persist_intent_verdict` for every async LLM-tier delivery; the synchronous heuristic-bulk path - (:meth:`create_intent_verdicts_bulk`) stays as plain INSERT - since each heuristic UUID is freshly generated per turn. + (:meth:`create_intent_verdicts_bulk`) inserts with per-row + ``ON CONFLICT DO NOTHING`` instead — its UUIDs are freshly + generated per turn, but the daemon can race a fallback UPSERT + of one of those same IDs in ahead of the bulk write. """ ... @@ -1754,6 +1756,15 @@ class StorageBackend(Protocol): vocabulary. Used by the synchronous heuristic-verdict persistence loop in ``approve_tools`` so a tool-heavy turn doesn't pay N×commit latency before the approval prompt renders. + + Inserts ``ON CONFLICT (verdict_id) DO NOTHING``: the async judge + daemon's first delivery can UPSERT a fallback row — which reuses + a heuristic ``verdict_id`` from this very batch — before the + bulk write runs. Aborting the whole statement on that collision + (plain-INSERT behavior) silently discarded every other row in + the batch; skipping just the colliding row keeps the rest AND + preserves the daemon's ``llm_fallback`` tier upgrade rather + than regressing it to the heuristic stamp. """ ... diff --git a/turnstone/core/storage/_sqlite.py b/turnstone/core/storage/_sqlite.py index e2b6472a..335ea449 100644 --- a/turnstone/core/storage/_sqlite.py +++ b/turnstone/core/storage/_sqlite.py @@ -3815,6 +3815,16 @@ class SQLiteBackend: conn.commit() def create_intent_verdicts_bulk(self, verdicts: list[dict[str, Any]]) -> None: + # ON CONFLICT DO NOTHING per row: the async judge daemon can land + # an UPSERT (which INSERTs the heuristic verdict_id it reuses) + # before this bulk write runs. A plain INSERT would abort the + # whole statement on that one collision — silently discarding + # every OTHER row in the batch (the caller swallows storage + # errors). Skipping just the colliding row also keeps the + # daemon's tier upgrade ("llm_fallback") instead of regressing + # it to the bulk's heuristic stamp. + from sqlalchemy.dialects.sqlite import insert as sqlite_insert + if not verdicts: return now = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S") @@ -3840,7 +3850,12 @@ class SQLiteBackend: for v in verdicts ] with self._conn() as conn: - conn.execute(sa.insert(intent_verdicts), rows) + conn.execute( + sqlite_insert(intent_verdicts).on_conflict_do_nothing( + index_elements=["verdict_id"] + ), + rows, + ) conn.commit() def get_intent_verdict(self, verdict_id: str) -> dict[str, Any] | None: