feat: better tool call followup on compat mode

tests/api/answer/test_base_routes.py

@@ -1181,12 +1181,18 @@ class TestCompleteStreamWalAcceptance:
 
 @pytest.mark.unit
 class TestStreamingHeartbeatSeed:
-    """Regression guard: when the row flips to ``streaming`` we must seed
-    ``last_heartbeat_at`` so the watchdog doesn't fall back to
-    ``timestamp`` (creation time) on slow LLM cold-starts (>idle_secs).
+    """Regression guard: the reserved row must carry a fresh
+    ``last_heartbeat_at`` from generation start so the reconciler doesn't fall
+    back to ``timestamp`` (creation time) on slow LLM cold-starts or while a
+    reasoning model streams only ``thought`` chunks. The heartbeat is seeded
+    once before the first chunk and re-stamped when the row flips to
+    ``streaming``; the ``pending → streaming`` status transition itself still
+    fires exactly once on the first ``answer``/``sources`` chunk.
     """
 
-    def test_heartbeat_seeded_on_first_chunk(self, mock_mongo_db, flask_app):
+    def test_heartbeat_seeded_at_generation_start_and_on_first_chunk(
+        self, mock_mongo_db, flask_app,
+    ):
         from application.api.answer.routes.base import BaseAnswerResource
 
         with flask_app.app_context():
@@ -1221,7 +1227,7 @@ class TestStreamingHeartbeatSeed:
             )
 
             # update_message_status flips the row to ``streaming`` exactly
-            # once (idempotent via streaming_marked).
+            # once (idempotent via streaming_marked) on the first answer chunk.
             status_calls = [
                 c for c in resource.conversation_service
                 .update_message_status.call_args_list
@@ -1230,14 +1236,18 @@ class TestStreamingHeartbeatSeed:
             assert len(status_calls) == 1
             assert status_calls[0].args[0] == "msg1"
 
-            # heartbeat seed runs exactly once at the same flip — multiple
-            # chunks don't re-stamp inside this window (the throttled
-            # _heartbeat_streaming path is gated by STREAM_HEARTBEAT_INTERVAL).
+            # Heartbeat is stamped twice for this 2-answer-chunk stream: once at
+            # the generation-start seed (before the first chunk, while still
+            # ``pending``) and once when the first answer chunk flips the row to
+            # ``streaming``. The second answer chunk does NOT re-stamp — the
+            # per-chunk ``_heartbeat_streaming`` pump is throttled by
+            # STREAM_HEARTBEAT_INTERVAL and the two chunks fall inside one
+            # interval under real ``time.monotonic``.
             hb_calls = (
                 resource.conversation_service.heartbeat_message.call_args_list
             )
-            assert len(hb_calls) == 1
-            assert hb_calls[0].args[0] == "msg1"
+            assert len(hb_calls) == 2
+            assert all(c.args[0] == "msg1" for c in hb_calls)
 
     def test_heartbeat_seed_skipped_without_reserved_message_id(
         self, mock_mongo_db, flask_app,

tests/api/v1/test_reasoning_heartbeat_liveness.py

@@ -0,0 +1,599 @@
+"""Streaming-heartbeat liveness for reasoning models that "think" before
+answering (``complete_stream`` in ``application/api/answer/routes/base.py``).
+
+Background
+----------
+The reconciler (``ReconciliationRepository.find_and_lock_stuck_messages``)
+fails a ``pending``/``streaming`` row whose effective freshness —
+``GREATEST(timestamp, message_metadata.last_heartbeat_at)`` — is older than
+5 minutes. ``complete_stream`` keeps that freshness up via a streaming
+heartbeat.
+
+The bug these tests pin: the heartbeat pump used to early-return until the row
+had been flipped to ``streaming``, and the row is only flipped on the first
+``answer``/``sources`` chunk. A reasoning model (e.g. ``reasoning_effort:
+high``) that streams only ``thought`` chunks for minutes before its first
+answer token therefore left the row ``pending`` with a *frozen* heartbeat, so
+the reconciler would falsely fail a live request.
+
+The fix makes the heartbeat a true "agent is producing output / is alive"
+signal: it pumps whenever a ``reserved_message_id`` exists (regardless of the
+``streaming`` status flip), and is seeded once at generation start. The
+``pending → streaming`` status transition itself is unchanged (still driven by
+the first ``answer``/``sources`` chunk).
+
+Residual (documented, intentionally not covered here): a model that emits *no*
+chunks at all — not even ``thought`` — for >5 min would still go stale, since
+the pump only ticks when a chunk flows. Covering a fully-silent stream would
+need a background-thread heartbeat or a higher threshold; both are out of
+scope for this surgical fix.
+
+These tests drive the real ``complete_stream`` against an ephemeral Postgres
+database (``pg_engine``) and reuse the wiring/fakes style of
+``tests/api/v1/test_v1_tool_pause_finalization.py``.
+"""
+
+from __future__ import annotations
+
+import uuid
+from contextlib import contextmanager
+from typing import Any, Dict, List, Optional
+
+import pytest
+from sqlalchemy import text
+
+from application.api.answer.routes.base import BaseAnswerResource
+from application.api.answer.services.conversation_service import ConversationService
+from application.storage.db.repositories.conversations import ConversationsRepository
+from application.storage.db.repositories.reconciliation import (
+    ReconciliationRepository,
+)
+
+
+# ---------------------------------------------------------------------------
+# Fakes
+# ---------------------------------------------------------------------------
+
+
+class _FakeLLM:
+    """Minimal LLM stand-in so ``complete_stream`` can stamp ``_request_id``."""
+
+    def __init__(self) -> None:
+        self._request_id: Optional[str] = None
+        self.model_id = "gpt-5"
+
+
+class _ThoughtOnlyAgent:
+    """Agent that streams only ``thought`` chunks and never answers.
+
+    Models a reasoning model mid-"thinking" phase: the stream is alive and
+    producing output, but no ``answer``/``sources`` chunk has arrived yet, so
+    the row is still ``pending`` (``streaming`` not marked).
+    """
+
+    def __init__(self, n_thoughts: int = 4) -> None:
+        self.llm = _FakeLLM()
+        self._n_thoughts = n_thoughts
+
+    def gen(self, query: str = ""):
+        for i in range(self._n_thoughts):
+            yield {"thought": f"reasoning step {i}..."}
+
+
+class _ThoughtThenAnswerAgent:
+    """Agent that thinks (``thought`` chunks) then emits a final ``answer``.
+
+    Used for the regression case: the row must still flip to ``streaming`` on
+    the first ``answer`` chunk and finalize ``complete``.
+    """
+
+    ANSWER_TEXT = "The answer is 42."
+
+    def __init__(self, n_thoughts: int = 3) -> None:
+        self.llm = _FakeLLM()
+        self._n_thoughts = n_thoughts
+
+    def gen(self, query: str = ""):
+        for i in range(self._n_thoughts):
+            yield {"thought": f"reasoning step {i}..."}
+        yield {"answer": self.ANSWER_TEXT}
+
+
+class _NoopJournalWriter:
+    """No-op journal writer so the test exercises DB row state, not the
+    message-events journal / Redis broadcast (orthogonal to heartbeats)."""
+
+    def __init__(self, *args, **kwargs) -> None:
+        pass
+
+    def record(self, *args, **kwargs) -> None:
+        pass
+
+    def flush(self) -> None:
+        pass
+
+    def close(self) -> None:
+        pass
+
+
+# ---------------------------------------------------------------------------
+# Test harness
+# ---------------------------------------------------------------------------
+
+
+def _seed_user(conn, user_id: str) -> None:
+    conn.execute(
+        text(
+            "INSERT INTO users (user_id) VALUES (:u) "
+            "ON CONFLICT (user_id) DO NOTHING"
+        ),
+        {"u": user_id},
+    )
+
+
+def _row(conn, message_id: str) -> Dict[str, Any]:
+    """Fetch ``(status, message_metadata)`` for a reserved row."""
+    r = conn.execute(
+        text(
+            "SELECT status, message_metadata FROM conversation_messages "
+            "WHERE id = CAST(:m AS uuid)"
+        ),
+        {"m": message_id},
+    ).fetchone()
+    return {"status": r[0], "metadata": r[1]} if r is not None else {}
+
+
+@contextmanager
+def _wire_db(engine, monkeypatch):
+    """Point conversation_service / continuation_service / base at ``engine``.
+
+    Each helper opens its own short-lived connection (matching production),
+    so we hand out fresh connections from the same ephemeral engine and swap
+    the journal writer for a no-op. Mirrors the helper in
+    ``test_v1_tool_pause_finalization.py``.
+    """
+    from application.api.answer.services import conversation_service as conv_mod
+    from application.api.answer.services import continuation_service as cont_mod
+    from application.api.answer.routes import base as base_mod
+
+    @contextmanager
+    def _session():
+        conn = engine.connect()
+        txn = conn.begin()
+        try:
+            yield conn
+            txn.commit()
+        except Exception:
+            txn.rollback()
+            raise
+        finally:
+            conn.close()
+
+    @contextmanager
+    def _readonly():
+        conn = engine.connect()
+        try:
+            yield conn
+        finally:
+            conn.close()
+
+    monkeypatch.setattr(conv_mod, "db_session", _session)
+    monkeypatch.setattr(conv_mod, "db_readonly", _readonly)
+    monkeypatch.setattr(cont_mod, "db_session", _session)
+    monkeypatch.setattr(cont_mod, "db_readonly", _readonly)
+    monkeypatch.setattr(base_mod, "db_session", _session)
+    monkeypatch.setattr(base_mod, "db_readonly", _readonly)
+    monkeypatch.setattr(base_mod, "BatchedJournalWriter", _NoopJournalWriter)
+    monkeypatch.setattr(base_mod, "record_event", lambda *a, **kw: None)
+    monkeypatch.setattr(base_mod, "publish_user_event", lambda *a, **kw: None)
+    yield
+
+
+class _FakeMonotonic:
+    """Deterministic ``time.monotonic`` that jumps forward on every call.
+
+    ``complete_stream`` reads ``time.monotonic()`` once per loop iteration
+    (via ``_heartbeat_streaming``) plus at seed/mark points. Advancing by more
+    than ``STREAM_HEARTBEAT_INTERVAL`` (60s) on each read guarantees the
+    interval gate fires every iteration, so a thought-only stream attempts a
+    heartbeat on every ``thought`` chunk.
+    """
+
+    def __init__(self, step: float = 120.0) -> None:
+        self._t = 0.0
+        self._step = step
+
+    def __call__(self) -> float:
+        self._t += self._step
+        return self._t
+
+
+def _make_base() -> BaseAnswerResource:
+    base = BaseAnswerResource.__new__(BaseAnswerResource)
+    base.default_model_id = "gpt-5"
+    base.conversation_service = ConversationService()
+    return base
+
+
+def _drain(gen) -> List[str]:
+    return list(gen)
+
+
+# ---------------------------------------------------------------------------
+# 1. Pump-during-reasoning: a thought-only stream heartbeats while pending
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.integration
+class TestHeartbeatPumpsDuringReasoning:
+    """A reasoning model that only streams ``thought`` chunks (no answer yet)
+    must still keep its ``last_heartbeat_at`` fresh while the row is ``pending``
+    — otherwise the reconciler falsely fails a live request.
+
+    Pre-fix the pump early-returns while ``streaming`` is unmarked, so a
+    thought-only phase produces NO heartbeat and these assertions fail.
+    """
+
+    def test_thought_only_stream_heartbeats_while_pending(
+        self, pg_engine, monkeypatch
+    ):
+        from application.api.answer.routes import base as base_mod
+
+        user_id = f"user-{uuid.uuid4().hex[:8]}"
+        with pg_engine.begin() as conn:
+            _seed_user(conn, user_id)
+
+        # Count heartbeat invocations directly on the service so the test is
+        # explicit that the pump fired during the thought-only phase.
+        heartbeat_calls: List[str] = []
+        real_heartbeat = ConversationService.heartbeat_message
+
+        def _counting_heartbeat(self, message_id: str) -> bool:
+            heartbeat_calls.append(message_id)
+            return real_heartbeat(self, message_id)
+
+        with _wire_db(pg_engine, monkeypatch):
+            monkeypatch.setattr(
+                ConversationService, "heartbeat_message", _counting_heartbeat
+            )
+            # Make every loop iteration cross the heartbeat interval.
+            monkeypatch.setattr(base_mod.time, "monotonic", _FakeMonotonic())
+
+            base = _make_base()
+            agent = _ThoughtOnlyAgent(n_thoughts=4)
+            frames = _drain(
+                base.complete_stream(
+                    question="hard reasoning question",
+                    agent=agent,
+                    conversation_id=None,
+                    user_api_key=None,
+                    decoded_token={"sub": user_id},
+                    should_persist=True,
+                    model_id="gpt-5",
+                )
+            )
+
+        # Sanity: the stream only ever emitted thoughts (plus framing/end),
+        # never an answer — so ``streaming`` was never marked during the loop.
+        joined = "\n".join(frames)
+        assert '"type": "thought"' in joined
+        assert '"type": "answer"' not in joined
+
+        # Resolve the reserved row.
+        with pg_engine.connect() as conn:
+            msg_id = conn.execute(
+                text(
+                    "SELECT cm.id FROM conversation_messages cm "
+                    "JOIN conversations c ON c.id = cm.conversation_id "
+                    "WHERE c.user_id = :u ORDER BY cm.timestamp DESC LIMIT 1"
+                ),
+                {"u": user_id},
+            ).scalar()
+            assert msg_id is not None
+            row = _row(conn, str(msg_id))
+
+        # Core assertion: the heartbeat fired for this row during a thought-only
+        # phase. Pre-fix the pump is gated behind ``streaming_marked`` and never
+        # runs, so ``heartbeat_calls`` is empty.
+        assert str(msg_id) in heartbeat_calls, (
+            "heartbeat_message was never called for the pending reasoning row; "
+            "the heartbeat pump is still gated behind the streaming flip"
+        )
+        # And it is observable on the row: a non-null last_heartbeat_at.
+        assert row["metadata"] is not None
+        assert row["metadata"].get("last_heartbeat_at") is not None
+
+    def test_heartbeat_advances_last_heartbeat_at_while_row_stays_pending(
+        self, pg_engine, monkeypatch
+    ):
+        """End-of-stream the answer arrives and the row finalizes, so to observe
+        the *pending* heartbeat we snapshot the row mid-stream: after the first
+        ``thought`` chunk the row must be ``pending`` with a fresh
+        ``last_heartbeat_at`` already stamped (seed-at-start), and a subsequent
+        ``thought`` must bump it further — all before any ``answer``.
+        """
+        from application.api.answer.routes import base as base_mod
+
+        user_id = f"user-{uuid.uuid4().hex[:8]}"
+        with pg_engine.begin() as conn:
+            _seed_user(conn, user_id)
+
+        observed: List[Dict[str, Any]] = []
+
+        class _ObservingAgent:
+            """Yields thoughts; after each, the harness inspects the DB row."""
+
+            def __init__(self) -> None:
+                self.llm = _FakeLLM()
+
+            def gen(self, query: str = ""):
+                # message_id is surfaced before this generator runs, so by the
+                # time we yield, the reserved row already exists. We can't read
+                # it from inside gen() (no msg id handle), so just emit; the
+                # snapshots happen in the consuming loop below via a wrapper.
+                for i in range(3):
+                    yield {"thought": f"step {i}"}
+
+        with _wire_db(pg_engine, monkeypatch):
+            monkeypatch.setattr(base_mod.time, "monotonic", _FakeMonotonic())
+            base = _make_base()
+            agent = _ObservingAgent()
+
+            gen = base.complete_stream(
+                question="reason about it",
+                agent=agent,
+                conversation_id=None,
+                user_api_key=None,
+                decoded_token={"sub": user_id},
+                should_persist=True,
+                model_id="gpt-5",
+            )
+
+            # Pull the first frame (the ``message_id`` event) to learn the row.
+            first = next(gen)
+            assert "message_id" in first
+            msg_id = None
+            with pg_engine.connect() as conn:
+                msg_id = conn.execute(
+                    text(
+                        "SELECT cm.id FROM conversation_messages cm "
+                        "JOIN conversations c ON c.id = cm.conversation_id "
+                        "WHERE c.user_id = :u ORDER BY cm.timestamp DESC LIMIT 1"
+                    ),
+                    {"u": user_id},
+                ).scalar()
+            assert msg_id is not None
+
+            # Drain the rest, snapshotting the row's heartbeat after each frame.
+            for _frame in gen:
+                with pg_engine.connect() as conn:
+                    observed.append(_row(conn, str(msg_id)))
+
+        # While the stream was thought-only the row stayed ``pending`` and the
+        # heartbeat advanced (was non-null on each snapshot).
+        pending_snaps = [s for s in observed if s.get("status") == "pending"]
+        assert pending_snaps, "expected at least one pending snapshot mid-stream"
+        for snap in pending_snaps:
+            assert snap["metadata"] is not None
+            assert snap["metadata"].get("last_heartbeat_at") is not None
+
+
+# ---------------------------------------------------------------------------
+# 2. Seed-at-start: the reserved row has a heartbeat from generation start
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.integration
+class TestHeartbeatSeededAtGenerationStart:
+    """Right after ``complete_stream`` reserves the row and begins consuming the
+    generator (before any ``answer``), the row must already carry a non-null
+    ``last_heartbeat_at`` — so even a model that takes a while to emit its first
+    chunk is covered from t=0, not only from the first interval tick.
+    """
+
+    def test_reserved_row_has_heartbeat_before_first_answer(
+        self, pg_engine, monkeypatch
+    ):
+        # This case keeps real ``time.monotonic`` so it asserts the *seed*
+        # (which runs before the loop, with no time advance), not an
+        # interval-driven pump — so no ``base.time`` monkeypatch is needed.
+        user_id = f"user-{uuid.uuid4().hex[:8]}"
+        with pg_engine.begin() as conn:
+            _seed_user(conn, user_id)
+
+        captured: Dict[str, Any] = {}
+
+        class _SlowFirstChunkAgent:
+            """Yields a single thought then stops — emulates a model that has
+            produced no answer yet by the time we inspect the row."""
+
+            def __init__(self) -> None:
+                self.llm = _FakeLLM()
+
+            def gen(self, query: str = ""):
+                # Inspect the reserved row at the very first generator step,
+                # before any answer/sources chunk has been processed.
+                with pg_engine.connect() as conn:
+                    mid = conn.execute(
+                        text(
+                            "SELECT cm.id FROM conversation_messages cm "
+                            "JOIN conversations c ON c.id = cm.conversation_id "
+                            "WHERE c.user_id = :u "
+                            "ORDER BY cm.timestamp DESC LIMIT 1"
+                        ),
+                        {"u": user_id},
+                    ).scalar()
+                    captured["row"] = _row(conn, str(mid)) if mid else {}
+                yield {"thought": "starting to reason"}
+
+        with _wire_db(pg_engine, monkeypatch):
+            # Keep real monotonic time so this asserts the *seed*, not an
+            # interval-driven pump (no time advance happens before gen runs).
+            base = _make_base()
+            agent = _SlowFirstChunkAgent()
+            _drain(
+                base.complete_stream(
+                    question="seed check",
+                    agent=agent,
+                    conversation_id=None,
+                    user_api_key=None,
+                    decoded_token={"sub": user_id},
+                    should_persist=True,
+                    model_id="gpt-5",
+                )
+            )
+
+        row = captured.get("row") or {}
+        # The seed runs before the loop, so the row is still ``pending`` and
+        # already has a heartbeat stamped.
+        assert row.get("status") == "pending"
+        assert row.get("metadata") is not None
+        assert row["metadata"].get("last_heartbeat_at") is not None
+
+
+# ---------------------------------------------------------------------------
+# 3. End-to-end reconciler guarantee: a fresh-heartbeat pending row is safe
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.integration
+class TestReconcilerProtectsFreshlyHeartbeatedPendingRow:
+    """The whole point of the fix: a ``pending`` row older than 5 min by
+    ``timestamp`` but with a fresh ``last_heartbeat_at`` (the reasoning pump)
+    must NOT be swept by ``find_and_lock_stuck_messages``; a stale-heartbeat
+    ``pending`` row still IS swept.
+
+    This mirrors the ``streaming`` cases in
+    ``tests/storage/db/repositories/test_reconciliation.py`` but for the
+    ``pending`` (still-reasoning) status the fix is about.
+    """
+
+    def _seed_pending_with_heartbeat(
+        self, conn, *, hb_minutes_ago: float, ts_minutes_ago: int = 20,
+    ) -> str:
+        conv = ConversationsRepository(conn).create("u-recon", "reasoning recon")
+        row = conn.execute(
+            text(
+                """
+                INSERT INTO conversation_messages (
+                    conversation_id, position, prompt, response, status,
+                    user_id, timestamp, message_metadata
+                )
+                VALUES (
+                    CAST(:cid AS uuid), 0, 'p', '', 'pending', 'u-recon',
+                    clock_timestamp() - make_interval(mins => :ts),
+                    jsonb_build_object(
+                        'last_heartbeat_at',
+                        to_jsonb(
+                            clock_timestamp()
+                            - make_interval(secs => :hb_secs)
+                        )
+                    )
+                )
+                RETURNING id
+                """
+            ),
+            {
+                "cid": conv["id"],
+                "ts": ts_minutes_ago,
+                "hb_secs": int(hb_minutes_ago * 60),
+            },
+        ).fetchone()
+        return str(row[0])
+
+    def test_fresh_heartbeat_pending_row_is_not_swept(self, pg_conn):
+        # timestamp 20 min old (well past the 5-min age gate) but the reasoning
+        # pump heartbeat is only 30s old → the row is live, must be skipped.
+        msg_id = self._seed_pending_with_heartbeat(
+            pg_conn, hb_minutes_ago=0.5, ts_minutes_ago=20
+        )
+        rows = ReconciliationRepository(pg_conn).find_and_lock_stuck_messages()
+        assert all(str(r["id"]) != msg_id for r in rows), (
+            "a still-reasoning pending row with a fresh heartbeat was "
+            "incorrectly selected for reconciliation"
+        )
+
+    def test_stale_heartbeat_pending_row_is_swept(self, pg_conn):
+        # Both timestamp and heartbeat are older than 5 min → genuinely stuck.
+        msg_id = self._seed_pending_with_heartbeat(
+            pg_conn, hb_minutes_ago=10, ts_minutes_ago=20
+        )
+        rows = ReconciliationRepository(pg_conn).find_and_lock_stuck_messages()
+        assert any(str(r["id"]) == msg_id for r in rows), (
+            "a pending row stale by both timestamp and heartbeat should be "
+            "selected for reconciliation"
+        )
+
+
+# ---------------------------------------------------------------------------
+# 4. Regression: a normal answer turn still marks streaming + heartbeats
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.integration
+class TestNormalAnswerTurnUnchanged:
+    """A normal turn that streams an ``answer`` must still flip the row to
+    ``streaming`` on the first answer chunk (status semantics unchanged) and
+    finalize ``complete``. The heartbeat liveness change must not regress this.
+    """
+
+    def test_answer_turn_marks_streaming_and_finalizes_complete(
+        self, pg_engine, monkeypatch
+    ):
+        from application.api.answer.routes import base as base_mod
+
+        user_id = f"user-{uuid.uuid4().hex[:8]}"
+        with pg_engine.begin() as conn:
+            _seed_user(conn, user_id)
+
+        # Observe the status transition: snapshot after the first answer frame.
+        statuses_seen: List[str] = []
+        real_update = ConversationService.update_message_status
+
+        def _recording_update(self, message_id: str, status: str) -> bool:
+            statuses_seen.append(status)
+            return real_update(self, message_id, status)
+
+        with _wire_db(pg_engine, monkeypatch):
+            monkeypatch.setattr(
+                ConversationService, "update_message_status", _recording_update
+            )
+            monkeypatch.setattr(base_mod.time, "monotonic", _FakeMonotonic())
+            base = _make_base()
+            agent = _ThoughtThenAnswerAgent(n_thoughts=2)
+            frames = _drain(
+                base.complete_stream(
+                    question="normal question",
+                    agent=agent,
+                    conversation_id=None,
+                    user_api_key=None,
+                    decoded_token={"sub": user_id},
+                    should_persist=True,
+                    model_id="gpt-5",
+                )
+            )
+
+        joined = "\n".join(frames)
+        assert '"type": "answer"' in joined
+        assert '"type": "end"' in joined
+
+        # The first answer chunk marked the row ``streaming`` exactly once.
+        assert statuses_seen == ["streaming"]
+
+        with pg_engine.connect() as conn:
+            row = conn.execute(
+                text(
+                    "SELECT status, response, message_metadata "
+                    "FROM conversation_messages cm "
+                    "JOIN conversations c ON c.id = cm.conversation_id "
+                    "WHERE c.user_id = :u ORDER BY cm.timestamp DESC LIMIT 1"
+                ),
+                {"u": user_id},
+            ).fetchone()
+
+        # Final state: terminal ``complete`` with the answer, and a heartbeat
+        # was stamped along the way (liveness intact).
+        assert row[0] == "complete"
+        assert row[1] == _ThoughtThenAnswerAgent.ANSWER_TEXT
+        assert row[2] is not None
+        assert row[2].get("last_heartbeat_at") is not None

tests/api/v1/test_v1_continuation_ownership.py

@@ -0,0 +1,235 @@
+"""Foreign-``conversation_id`` ownership gate on the ``/v1`` continuation path.
+
+A ``/v1`` tool continuation persists its final answer into a client-supplied
+``conversation_id``. Ownership IS enforced deep in the stack:
+
+    route → ``build_continuation_from_messages`` → ``build_agent("")``
+          → ``initialize`` → ``_load_conversation_history``
+          → ``get_conversation(conversation_id, owner)``  ← owner-scoped read
+          → ``None`` for a foreign conversation
+          → ``ValueError("Conversation not found or unauthorized")``
+          → route's ``except ValueError`` → HTTP 400
+
+…but the existing E2E suite mocks ``build_agent``, so this security gate is
+never actually exercised. This test seeds a conversation owned by user A and an
+agent owned by user B, then POSTs a continuation with B's api_key targeting A's
+conversation. The ownership read runs against the **real** DB (not mocked): we
+only short-circuit ``StreamProcessor.initialize`` to the
+``_load_conversation_history`` step so the test doesn't have to stand up the
+full agent/model/source/retriever machinery to reach the gate.
+
+Asserted: the route returns 400 AND no assistant turn is appended to A's
+conversation (the foreign write is rejected before persistence).
+"""
+
+from __future__ import annotations
+
+import json
+import uuid
+from typing import Any, Dict, List
+
+import pytest
+from flask import Flask
+from sqlalchemy import text
+
+from application.api.answer.services import stream_processor as sp_mod
+from application.api.v1.routes import v1_bp
+from application.storage.db.repositories.agents import AgentsRepository
+from application.storage.db.repositories.conversations import ConversationsRepository
+
+# Reuse the route-level DB wiring + seed helpers from the tool-pause suite.
+from tests.api.v1.test_v1_tool_pause_finalization import (
+    _seed_user,
+    _wire_v1_route_db,
+)
+
+
+CLIENT_TOOL = {
+    "type": "function",
+    "function": {
+        "name": "get_weather",
+        "description": "Get the weather for a city.",
+        "parameters": {
+            "type": "object",
+            "properties": {"city": {"type": "string"}},
+            "required": ["city"],
+        },
+    },
+}
+
+
+def _seed_conversation(conn, user_id: str, name: str) -> str:
+    """Create a conversation owned by ``user_id`` with one finalized turn."""
+    repo = ConversationsRepository(conn)
+    conv = repo.create(user_id, name)
+    conv_id = str(conv["id"])
+    repo.reserve_message(
+        conv_id,
+        prompt="hello",
+        placeholder_response="hi there",
+        status="complete",
+    )
+    return conv_id
+
+
+def _seed_agent(conn, user_id: str, key: str) -> None:
+    AgentsRepository(conn).create(user_id, "Weather Agent", "published", key=key)
+
+
+def _build_app() -> Flask:
+    app = Flask(__name__)
+    app.register_blueprint(v1_bp)
+    return app
+
+
+def _row_count(conn, conv_id: str) -> int:
+    return conn.execute(
+        text(
+            "SELECT count(*) FROM conversation_messages "
+            "WHERE conversation_id = CAST(:c AS uuid)"
+        ),
+        {"c": conv_id},
+    ).scalar()
+
+
+def _only_run_history_load(self) -> None:
+    """Stand-in for ``StreamProcessor.initialize`` that runs ONLY the
+    ownership-checked history load.
+
+    The real ownership gate (``_load_conversation_history`` →
+    ``get_conversation(conversation_id, owner)``) is NOT mocked — it runs
+    against the real DB. We skip the surrounding agent/model/source/retriever
+    configuration only so the test can reach the gate without provisioning a
+    full agent runtime; the ``ValueError`` raised by the gate for a foreign
+    conversation propagates exactly as in production.
+    """
+    self._load_conversation_history()
+
+
+@pytest.mark.integration
+class TestForeignConversationOwnership:
+    """A continuation targeting another user's conversation is rejected at the
+    ownership gate (400) and never appends a turn to that conversation."""
+
+    PENDING_MESSAGES: List[Dict[str, Any]] = [
+        {"role": "user", "content": "weather in SF?"},
+        {
+            "role": "assistant",
+            "content": None,
+            "tool_calls": [
+                {
+                    "id": "call_abc",
+                    "type": "function",
+                    "function": {
+                        "name": "get_weather",
+                        "arguments": json.dumps({"city": "SF"}),
+                    },
+                }
+            ],
+        },
+        {"role": "tool", "tool_call_id": "call_abc", "content": "72F sunny"},
+    ]
+
+    def test_continuation_into_foreign_conversation_is_rejected(
+        self, pg_engine, monkeypatch
+    ):
+        user_a = f"userA-{uuid.uuid4().hex[:8]}"  # owns the conversation
+        user_b = f"userB-{uuid.uuid4().hex[:8]}"  # owns the agent (the caller)
+        api_key = f"key-{uuid.uuid4().hex[:8]}"
+
+        with pg_engine.begin() as conn:
+            _seed_user(conn, user_a)
+            _seed_user(conn, user_b)
+            # Agent owned by B → the v1 route resolves decoded_token={"sub": B}.
+            _seed_agent(conn, user_b, api_key)
+            # Conversation owned by A, with one existing complete turn.
+            conv_id = _seed_conversation(conn, user_a, "A's private chat")
+
+        with pg_engine.connect() as conn:
+            rows_before = _row_count(conn, conv_id)
+        assert rows_before == 1  # the seeded turn
+
+        app = _build_app()
+
+        # Drive the REAL ownership read; only short-circuit the surrounding
+        # initialize steps so we reach the gate without a full agent runtime.
+        with _wire_v1_route_db(pg_engine, monkeypatch):
+            monkeypatch.setattr(
+                sp_mod.StreamProcessor, "initialize", _only_run_history_load
+            )
+            with app.test_client() as c:
+                resp = c.post(
+                    "/v1/chat/completions",
+                    headers={"Authorization": f"Bearer {api_key}"},
+                    json={
+                        # B (the agent owner) tries to write into A's conversation.
+                        "conversation_id": conv_id,
+                        "messages": self.PENDING_MESSAGES,
+                        "tools": [CLIENT_TOOL],
+                        "docsgpt": {"save_conversation": True},
+                    },
+                )
+
+        # The ownership gate raised ValueError → route returns 400.
+        assert resp.status_code == 400, resp.get_data(as_text=True)
+        body = resp.get_json()
+        assert body["error"]["type"] == "invalid_request"
+
+        # No assistant turn was appended to A's conversation — the foreign
+        # write was rejected before any persistence.
+        with pg_engine.connect() as conn:
+            rows_after = _row_count(conn, conv_id)
+        assert rows_after == rows_before == 1
+
+    def test_owner_can_continue_into_own_conversation(self, pg_engine, monkeypatch):
+        """Control: the SAME continuation into the caller's OWN conversation
+        passes the ownership gate (no 400 from the gate). Proves the 400 above
+        is the ownership check firing, not an unrelated error in the path.
+        """
+        owner = f"owner-{uuid.uuid4().hex[:8]}"
+        api_key = f"key-{uuid.uuid4().hex[:8]}"
+
+        with pg_engine.begin() as conn:
+            _seed_user(conn, owner)
+            _seed_agent(conn, owner, api_key)
+            # Conversation owned by the agent owner (the caller).
+            conv_id = _seed_conversation(conn, owner, "owner's chat")
+
+        app = _build_app()
+
+        # Capture whether the gate passed: a flag the stubbed initialize sets
+        # only if ``_load_conversation_history`` returned without raising.
+        gate = {"passed": False}
+
+        def _history_then_flag(self):
+            self._load_conversation_history()
+            gate["passed"] = True
+            # Stop here: returning lets build_continuation_from_messages proceed
+            # to create_agent, which needs a full runtime we deliberately skip.
+            # Raise a sentinel the test recognises, distinct from the ownership
+            # ValueError, so we can assert the gate itself did not reject us.
+            raise RuntimeError("__gate_passed_sentinel__")
+
+        with _wire_v1_route_db(pg_engine, monkeypatch):
+            monkeypatch.setattr(
+                sp_mod.StreamProcessor, "initialize", _history_then_flag
+            )
+            with app.test_client() as c:
+                resp = c.post(
+                    "/v1/chat/completions",
+                    headers={"Authorization": f"Bearer {api_key}"},
+                    json={
+                        "conversation_id": conv_id,
+                        "messages": (
+                            TestForeignConversationOwnership.PENDING_MESSAGES
+                        ),
+                        "tools": [CLIENT_TOOL],
+                        "docsgpt": {"save_conversation": True},
+                    },
+                )
+
+        # The ownership read accepted the owner (gate passed); the request then
+        # fails on the deliberately-skipped runtime (500), NOT on a 400 from the
+        # gate. This isolates the ownership decision from the rest of the path.
+        assert gate["passed"] is True
+        assert resp.status_code == 500, resp.get_data(as_text=True)

tests/api/v1/test_v1_idempotency.py

@@ -0,0 +1,439 @@
+"""Layer-1 idempotency for ``/v1/chat/completions`` (non-streaming, OpenAI-style).
+
+Dropping the native ``resume_from_tool_actions`` path removed its
+``mark_resuming`` guard, so a duplicated/retried non-streaming POST could
+re-run the agent → a duplicate answer row + double token billing. The fix
+honors a client-supplied ``Idempotency-Key`` header: a retry returns the
+*stored first response* instead of re-running.
+
+These tests pin both layers against an ephemeral Postgres (``pg_engine``):
+
+- the ``application.api.v1.idempotency`` helper contract directly
+  (completed → cached, fresh-pending → 409, stale-pending → re-claim,
+  non-2xx → not cached); and
+- the real ``/v1/chat/completions`` route end-to-end: the same key twice
+  returns the cached body and does NOT re-run the agent or append a second
+  conversation row; no key / different keys run independently; a 5xx first
+  response is not cached so a retry re-runs.
+"""
+
+from __future__ import annotations
+
+import uuid
+from contextlib import contextmanager
+from typing import Any, Dict, List
+from unittest.mock import patch
+
+import pytest
+from flask import Flask
+from sqlalchemy import text
+
+from application.api.v1 import idempotency as v1_idem
+from application.api.v1.idempotency import (
+    STALE_PENDING_SECONDS,
+    TASK_NAME,
+    claim_or_replay,
+    finalize,
+    release,
+    scoped_key,
+)
+from application.api.v1.routes import v1_bp
+
+# Reuse the route-level DB wiring + fake answering agent from the tool-pause
+# suite so a real two-POST round-trip runs against the ephemeral Postgres.
+from tests.api.v1.test_v1_tool_pause_finalization import (
+    _seed_agent,
+    _seed_user,
+    _wire_v1_route_db,
+)
+
+
+# ---------------------------------------------------------------------------
+# Helper-contract tests (drive the real task_dedup row directly)
+# ---------------------------------------------------------------------------
+
+
+def _make_response(app: Flask, body: Dict[str, Any], status: int):
+    """Build a real Flask Response inside an app context (for ``finalize``)."""
+    from flask import jsonify, make_response
+
+    with app.test_request_context():
+        return make_response(jsonify(body), status)
+
+
+def _wire_idem_db(engine, monkeypatch) -> None:
+    """Point the v1 idempotency helper's ``db_session``/``db_readonly`` at
+    the ephemeral engine (each call opens its own short-lived connection)."""
+
+    @contextmanager
+    def _session():
+        conn = engine.connect()
+        txn = conn.begin()
+        try:
+            yield conn
+            txn.commit()
+        except Exception:
+            txn.rollback()
+            raise
+        finally:
+            conn.close()
+
+    @contextmanager
+    def _readonly():
+        conn = engine.connect()
+        try:
+            yield conn
+        finally:
+            conn.close()
+
+    monkeypatch.setattr(v1_idem, "db_session", _session)
+    monkeypatch.setattr(v1_idem, "db_readonly", _readonly)
+
+
+def _backdate_pending(engine, key: str, secs_ago: int) -> None:
+    with engine.begin() as conn:
+        conn.execute(
+            text(
+                "UPDATE task_dedup "
+                "SET created_at = clock_timestamp() - make_interval(secs => :s) "
+                "WHERE idempotency_key = :k"
+            ),
+            {"s": secs_ago, "k": key},
+        )
+
+
+@pytest.mark.integration
+class TestIdempotencyHelperContract:
+    """The three-outcome claim-before-process contract, against a real row."""
+
+    def test_scoped_key_namespaces_by_agent(self):
+        assert scoped_key("abc", "agent-1") == "agent-1:abc"
+        # Two agents with the same key value never collide.
+        assert scoped_key("abc", "agent-1") != scoped_key("abc", "agent-2")
+        # Opt-in: missing either component disables idempotency.
+        assert scoped_key(None, "agent-1") is None
+        assert scoped_key("abc", None) is None
+
+    def test_first_claim_then_completed_replays_cached(self, pg_engine, monkeypatch):
+        # ``claim_or_replay`` builds Flask responses (``jsonify``) for the
+        # replay/409 paths, so it runs inside an app context — exactly as it
+        # does in the real route.
+        app = Flask(__name__)
+        _wire_idem_db(pg_engine, monkeypatch)
+        key = f"agent:{uuid.uuid4().hex}"
+
+        with app.app_context():
+            claimed, replay = claim_or_replay(key)
+            assert claimed is True
+            assert replay is None
+
+            # Finalize a 200 with a body, then a retry replays it byte-for-byte
+            # without re-claiming.
+            body = {"id": "chatcmpl-1", "choices": [{"x": 1}]}
+            finalize(key, _make_response(app, body, 200))
+
+            claimed2, replay2 = claim_or_replay(key)
+            assert claimed2 is False
+            assert replay2 is not None
+            assert replay2.status_code == 200
+            assert replay2.get_json() == body
+
+    def test_fresh_pending_returns_409(self, pg_engine, monkeypatch):
+        app = Flask(__name__)
+        _wire_idem_db(pg_engine, monkeypatch)
+        key = f"agent:{uuid.uuid4().hex}"
+
+        with app.app_context():
+            claimed, _ = claim_or_replay(key)
+            assert claimed is True  # first request in flight, never finalized
+
+            # A concurrent same-key request sees a fresh pending row → 409.
+            claimed2, replay2 = claim_or_replay(key)
+            assert claimed2 is False
+            assert replay2 is not None
+            assert replay2.status_code == 409
+            assert replay2.get_json()["error"]["type"] == "idempotency_conflict"
+
+    def test_stale_pending_is_reclaimed(self, pg_engine, monkeypatch):
+        app = Flask(__name__)
+        _wire_idem_db(pg_engine, monkeypatch)
+        key = f"agent:{uuid.uuid4().hex}"
+
+        with app.app_context():
+            claimed, _ = claim_or_replay(key)
+            assert claimed is True
+            # The original request "died" before finalize; age the claim past
+            # the safety window so a retry may re-claim instead of 409-ing
+            # forever.
+            _backdate_pending(pg_engine, key, secs_ago=STALE_PENDING_SECONDS + 5)
+
+            claimed2, replay2 = claim_or_replay(key)
+            assert claimed2 is True
+            assert replay2 is None
+
+    def test_non_2xx_response_is_not_cached(self, pg_engine, monkeypatch):
+        app = Flask(__name__)
+        _wire_idem_db(pg_engine, monkeypatch)
+        key = f"agent:{uuid.uuid4().hex}"
+
+        with app.app_context():
+            claimed, _ = claim_or_replay(key)
+            assert claimed is True
+            # A 5xx must release the claim (not cache the error).
+            finalize(key, _make_response(app, {"error": {"message": "boom"}}, 500))
+
+        with pg_engine.connect() as conn:
+            row = conn.execute(
+                text("SELECT * FROM task_dedup WHERE idempotency_key = :k"),
+                {"k": key},
+            ).fetchone()
+        assert row is None  # released → a genuine retry can re-claim
+
+    def test_release_drops_pending_claim(self, pg_engine, monkeypatch):
+        app = Flask(__name__)
+        _wire_idem_db(pg_engine, monkeypatch)
+        key = f"agent:{uuid.uuid4().hex}"
+        with app.app_context():
+            claim_or_replay(key)
+            release(key)
+        with pg_engine.connect() as conn:
+            row = conn.execute(
+                text("SELECT * FROM task_dedup WHERE idempotency_key = :k"),
+                {"k": key},
+            ).fetchone()
+        assert row is None
+
+
+# ---------------------------------------------------------------------------
+# Route-level tests — same key twice does NOT re-run the agent
+# ---------------------------------------------------------------------------
+
+
+class _CountingAnswerAgent:
+    """Agent that answers once per ``gen`` and records how often it ran.
+
+    A class-level counter survives across the fresh instances handed back by
+    the patched ``build_agent``, so the route-level test can assert the agent
+    ran exactly once across two same-key POSTs.
+    """
+
+    ANSWER_TEXT = "The answer is 42."
+    gen_calls = 0
+
+    def __init__(self) -> None:
+        from tests.api.v1.test_v1_tool_pause_finalization import (
+            _FakeClientToolExecutor,
+            _FakeLLM,
+        )
+
+        self.llm = _FakeLLM()
+        self.tool_executor = _FakeClientToolExecutor()
+        self.conversation_id = None
+        self.initial_user_id = None
+        self.tools: List[Dict[str, Any]] = []
+
+    def gen(self, query: str = ""):
+        # Write the base-class counter explicitly so subclasses (e.g. a
+        # failing variant) share one run-count rather than shadowing it.
+        _CountingAnswerAgent.gen_calls += 1
+        yield {"answer": self.ANSWER_TEXT}
+
+
+def _build_app() -> Flask:
+    app = Flask(__name__)
+    app.register_blueprint(v1_bp)
+    return app
+
+
+def _post_chat(client, body, api_key, idem_key=None):
+    headers = {"Authorization": f"Bearer {api_key}"}
+    if idem_key is not None:
+        headers["Idempotency-Key"] = idem_key
+    return client.post("/v1/chat/completions", headers=headers, json=body)
+
+
+def _conv_count(conn, user_id: str) -> int:
+    return conn.execute(
+        text("SELECT count(*) FROM conversations WHERE user_id = :u"),
+        {"u": user_id},
+    ).scalar()
+
+
+def _row_count(conn, user_id: str) -> int:
+    return conn.execute(
+        text(
+            "SELECT count(*) FROM conversation_messages cm "
+            "JOIN conversations c ON c.id = cm.conversation_id "
+            "WHERE c.user_id = :u"
+        ),
+        {"u": user_id},
+    ).scalar()
+
+
+@pytest.mark.integration
+class TestV1IdempotencyRoute:
+    """Drive the real ``/v1/chat/completions`` route with/without a key."""
+
+    QUESTION = {"messages": [{"role": "user", "content": "what is the answer?"}],
+                "docsgpt": {"save_conversation": True}}
+
+    def _fake_build_agent(self, question):  # noqa: ARG002
+        return _CountingAnswerAgent()
+
+    def test_same_key_twice_replays_without_rerunning_agent(
+        self, pg_engine, monkeypatch
+    ):
+        user_id = f"user-{uuid.uuid4().hex[:8]}"
+        api_key = f"key-{uuid.uuid4().hex[:8]}"
+        with pg_engine.begin() as conn:
+            _seed_user(conn, user_id)
+            _seed_agent(conn, user_id, api_key)
+
+        app = _build_app()
+        idem_key = f"idem-{uuid.uuid4().hex}"
+        _CountingAnswerAgent.gen_calls = 0
+
+        with _wire_v1_route_db(pg_engine, monkeypatch), patch(
+            "application.api.answer.services.stream_processor.StreamProcessor"
+            ".build_agent",
+            self._fake_build_agent,
+        ):
+            _wire_idem_db(pg_engine, monkeypatch)
+            with app.test_client() as c:
+                resp1 = _post_chat(c, self.QUESTION, api_key, idem_key=idem_key)
+                resp2 = _post_chat(c, self.QUESTION, api_key, idem_key=idem_key)
+
+        assert resp1.status_code == 200, resp1.get_data(as_text=True)
+        assert resp2.status_code == 200, resp2.get_data(as_text=True)
+        body1, body2 = resp1.get_json(), resp2.get_json()
+        # The retry returns the cached first response byte-for-byte.
+        assert body2 == body1
+        assert body2["choices"][0]["message"]["content"] == (
+            _CountingAnswerAgent.ANSWER_TEXT
+        )
+
+        # Real dedup: the agent ran exactly once across the two POSTs...
+        assert _CountingAnswerAgent.gen_calls == 1
+        # ...and no duplicate conversation / answer row was appended.
+        with pg_engine.connect() as conn:
+            assert _conv_count(conn, user_id) == 1
+            # One conversation, one terminal answer row (not two).
+            assert _row_count(conn, user_id) == 1
+
+    def test_no_key_runs_each_time(self, pg_engine, monkeypatch):
+        """Opt-in: without a key, behavior is today's — each POST runs."""
+        user_id = f"user-{uuid.uuid4().hex[:8]}"
+        api_key = f"key-{uuid.uuid4().hex[:8]}"
+        with pg_engine.begin() as conn:
+            _seed_user(conn, user_id)
+            _seed_agent(conn, user_id, api_key)
+
+        app = _build_app()
+        _CountingAnswerAgent.gen_calls = 0
+
+        with _wire_v1_route_db(pg_engine, monkeypatch), patch(
+            "application.api.answer.services.stream_processor.StreamProcessor"
+            ".build_agent",
+            self._fake_build_agent,
+        ):
+            _wire_idem_db(pg_engine, monkeypatch)
+            with app.test_client() as c:
+                resp1 = _post_chat(c, self.QUESTION, api_key)
+                resp2 = _post_chat(c, self.QUESTION, api_key)
+
+        assert resp1.status_code == 200
+        assert resp2.status_code == 200
+        # Two independent runs → two conversations.
+        assert _CountingAnswerAgent.gen_calls == 2
+        with pg_engine.connect() as conn:
+            assert _conv_count(conn, user_id) == 2
+
+    def test_different_keys_run_independently(self, pg_engine, monkeypatch):
+        user_id = f"user-{uuid.uuid4().hex[:8]}"
+        api_key = f"key-{uuid.uuid4().hex[:8]}"
+        with pg_engine.begin() as conn:
+            _seed_user(conn, user_id)
+            _seed_agent(conn, user_id, api_key)
+
+        app = _build_app()
+        _CountingAnswerAgent.gen_calls = 0
+
+        with _wire_v1_route_db(pg_engine, monkeypatch), patch(
+            "application.api.answer.services.stream_processor.StreamProcessor"
+            ".build_agent",
+            self._fake_build_agent,
+        ):
+            _wire_idem_db(pg_engine, monkeypatch)
+            with app.test_client() as c:
+                _post_chat(c, self.QUESTION, api_key, idem_key="k-A")
+                _post_chat(c, self.QUESTION, api_key, idem_key="k-B")
+
+        assert _CountingAnswerAgent.gen_calls == 2
+        with pg_engine.connect() as conn:
+            assert _conv_count(conn, user_id) == 2
+
+    def test_5xx_first_response_is_not_cached_and_retry_reruns(
+        self, pg_engine, monkeypatch
+    ):
+        """A failed first response must NOT be cached: a retry re-runs."""
+        user_id = f"user-{uuid.uuid4().hex[:8]}"
+        api_key = f"key-{uuid.uuid4().hex[:8]}"
+        with pg_engine.begin() as conn:
+            _seed_user(conn, user_id)
+            _seed_agent(conn, user_id, api_key)
+
+        app = _build_app()
+        idem_key = f"idem-{uuid.uuid4().hex}"
+
+        class _BoomAgent(_CountingAnswerAgent):
+            def gen(self, query: str = ""):
+                _CountingAnswerAgent.gen_calls += 1
+                raise RuntimeError("upstream exploded")
+                yield  # pragma: no cover - keeps gen a generator
+
+        boom_state = {"fail": True}
+
+        def _build(self, question):  # noqa: ARG001
+            return _BoomAgent() if boom_state["fail"] else _CountingAnswerAgent()
+
+        _CountingAnswerAgent.gen_calls = 0
+
+        with _wire_v1_route_db(pg_engine, monkeypatch), patch(
+            "application.api.answer.services.stream_processor.StreamProcessor"
+            ".build_agent",
+            _build,
+        ):
+            _wire_idem_db(pg_engine, monkeypatch)
+            with app.test_client() as c:
+                resp1 = _post_chat(c, self.QUESTION, api_key, idem_key=idem_key)
+                # First attempt blew up (500) and released the claim. The key
+                # is agent-id scoped at the route, so confirm release via a
+                # table-wide count for this task_name (no live claim left).
+                assert resp1.status_code == 500, resp1.get_data(as_text=True)
+                with pg_engine.connect() as conn:
+                    rows = conn.execute(
+                        text(
+                            "SELECT count(*) FROM task_dedup "
+                            "WHERE task_name = :n"
+                        ),
+                        {"n": TASK_NAME},
+                    ).scalar()
+                assert rows == 0  # error was not cached; claim released
+
+                # Retry: the agent now answers; idempotency did NOT cache the error.
+                boom_state["fail"] = False
+                resp2 = _post_chat(c, self.QUESTION, api_key, idem_key=idem_key)
+
+        assert resp2.status_code == 200, resp2.get_data(as_text=True)
+        # Agent ran twice total: the failing first + the successful retry.
+        assert _CountingAnswerAgent.gen_calls == 2
+        # The successful retry IS now cached.
+        with pg_engine.connect() as conn:
+            completed = conn.execute(
+                text(
+                    "SELECT count(*) FROM task_dedup "
+                    "WHERE task_name = :n AND status = 'completed'"
+                ),
+                {"n": TASK_NAME},
+            ).scalar()
+        assert completed == 1