From 1614257d375f5bb7f4e213ef5b32b797429194d1 Mon Sep 17 00:00:00 2001 From: Teknium <127238744+teknium1@users.noreply.github.com> Date: Fri, 16 Oct 2026 20:23:48 -0700 Subject: [PATCH] fix(gateway): /retry's rewrite commits in its re-send's own admission Under a session authority, messaging /retry rewrote the transcript and only then admitted the re-send. A drain, retirement or full FIFO that began after the pre-check deleted the last exchange and then refused the re-send, so the turn was lost (#139257, MSG-11's race). The handler now only plans the rewrite. It reads the active row ids before the transcript and binds the plan to the re-send (gateway/session_retry_admission.py). admit_native takes the binding for that event's own source object and runs the rewrite as a write guard after its admission gate and queue cap, inside the BEGIN IMMEDIATE that inserts the row. A refusal rolls both back. A transcript changed since the read (a newer turn, a compaction) refuses both, and /retry answers "failed; transcript unchanged". Both rewrite paths go through it: plain turns (replace_messages_in_transaction, pinned on the active ids) and compaction carriers (plan_user_turn_rewind + rewind_in_transaction with the plan's pins). The transcript and the admission ledger share one SessionDB handle per profile, which the path checks; without an authority the old rewrite-then-send path is unchanged. replace_messages and rewind_user_turn are split into their in-transaction and planning halves with no behaviour change for their existing callers. Review finding by @unsupportedpastels on #140185 Refs #139257 Refs #134889 (MSG-11) --- gateway/session_authority.py | 8 +- gateway/session_retry_admission.py | 78 +++++++++++ gateway/slash_commands_session.py | 81 +++++++++++ hermes_state_messages.py | 58 +++++--- hermes_state_rewind.py | 102 +++++++++----- tests/gateway/test_retry_authority_refusal.py | 127 +++++++++++++++++- 6 files changed, 392 insertions(+), 62 deletions(-) create mode 100644 gateway/session_retry_admission.py diff --git a/gateway/session_authority.py b/gateway/session_authority.py index d280068ce9..0c24a3b807 100644 --- a/gateway/session_authority.py +++ b/gateway/session_authority.py @@ -375,14 +375,20 @@ class SessionAuthority: source.thread_id, source.user_id], separators=(',', ':')) request_id = str(payload['native_text_v1']['event']['message_id'] or uuid.uuid4().hex) committed = [ref] # the row's target: the write may retarget a rotated-away reservation + # A /retry re-send carries its transcript rewrite: it commits in this admission's own write. + from gateway import session_retry_admission + rewrite = session_retry_admission.take(event, self.db) def admitted(row): committed[0] = SessionRef(self.profile_id, row.get('target_session_id') or ref.session_id) + if rewrite is not None and rewrite.applied: + rewrite.on_commit() return self._admitted(committed[0], event, row) native_write_started.set(True) # shielded from here: a cancelled caller may still commit return await tracked_write(self, partial( self._admit_native_write, principal_id='messaging:' + identity, session_id=ref.session_id, - request_id=request_id, payload=payload, authorize=self._admission_gate(), + request_id=request_id, payload=payload, + authorize=session_retry_admission.compose(rewrite, self._admission_gate()), raw=getattr(event, 'raw_message', None)), then=admitted, ordered=True, after=lambda _receipt: self._schedule_admitted(committed[0])) diff --git a/gateway/session_retry_admission.py b/gateway/session_retry_admission.py new file mode 100644 index 0000000000..652fc8d340 --- /dev/null +++ b/gateway/session_retry_admission.py @@ -0,0 +1,78 @@ +"""Messaging ``/retry`` under a session authority: the transcript rewrite commits in the re-sent +turn's own admission transaction (#139257 step 2, MSG-11). + +``/retry`` drops the last exchange and re-sends its user text. Rewriting first and admitting +afterwards loses the turn whenever admission refuses (a drain or retirement that begins after the +pre-check, a full FIFO). Here the handler only PLANS the rewrite: it reads the transcript, pins the +active row ids it read, and binds the plan to the re-send. ``admit_native`` takes the binding and +runs the rewrite as a write guard after its own admission gate and queue cap, in the same +``BEGIN IMMEDIATE`` that inserts the row. A refusal rolls both back; a transcript that changed since +the read refuses both (``RetryRewriteRefused``). +""" +from __future__ import annotations + +import sqlite3 +from contextvars import ContextVar +from dataclasses import dataclass, field +from typing import Any, Callable, Optional + +from hermes_state_runtime import RuntimeStoreError + +_bound: ContextVar[Optional[RetryRewrite]] = ContextVar('retry_rewrite', default=None) + + +class RetryRewriteRefused(RuntimeError): + """The planned rewrite no longer applies (the transcript changed since the read, a turn lease is + live, the session closed): the re-send was refused with it and nothing was written.""" + + +@dataclass +class RetryRewrite: + """One ``/retry`` rewrite, applied at most once inside the re-send's admission write.""" + db: Any + session_id: str + source: Any # the re-send's own SessionSource object: only that admission takes the binding + write: Callable[[Any], None] # (conn) -> None, inside the admission transaction + on_commit: Callable[[], None] # on the loop, once the admission that applied it committed + applied: bool = False + taken: bool = field(default=False, repr=False) + + def guard(self, conn): + self.applied = False # the writer may retry the whole transaction + try: + self.write(conn) + except (RuntimeStoreError, sqlite3.OperationalError): + raise # gate refusals and writer contention keep their own handling + except Exception as exc: + raise RetryRewriteRefused(str(exc)) from exc + self.applied = True + + +def bind(rewrite: RetryRewrite): + return _bound.set(rewrite) + + +def unbind(token) -> None: + _bound.reset(token) + + +def take(event, db) -> Optional[RetryRewrite]: + """The bound rewrite for this native admission, once, and only for the re-send itself (the same + source object on the same store). A drain task inheriting the context never takes it.""" + rewrite = _bound.get() + if (rewrite is None or rewrite.taken or getattr(event, 'source', None) is not rewrite.source + or rewrite.db is not db): + return None + rewrite.taken = True + return rewrite + + +def compose(rewrite: Optional[RetryRewrite], authorize): + """``authorize`` (the admission gate) first, then the rewrite, in one write transaction.""" + if rewrite is None: + return authorize + + def guard(conn): + authorize(conn) + rewrite.guard(conn) + return guard diff --git a/gateway/slash_commands_session.py b/gateway/slash_commands_session.py index 06d88d88cd..7969173441 100644 --- a/gateway/slash_commands_session.py +++ b/gateway/slash_commands_session.py @@ -392,10 +392,17 @@ class GatewaySessionCommandsMixin: source = event.source session_entry = await self.async_session_store.get_or_create_session(source) + # Under a session authority the rewrite commits in the re-send's own admission, fenced on + # the active rows read here (before the transcript, so any later write fails the fence). + owned = self._retry_owned_store(session_entry.session_id) try: + pinned = owned and await asyncio.to_thread(owned.get_active_message_ids, session_entry.session_id) history = await self.async_session_store.load_transcript(session_entry.session_id) except TranscriptReadError: return history_unreadable() + except Exception: + logger.debug("/retry could not read the active rows of %s", session_entry.session_id, exc_info=True) + return history_unreadable() last_user_idx = next((i for i in range(len(history) - 1, -1, -1) if user_originated_turn_view(history[i]) is not None), None) if last_user_idx is None: @@ -415,6 +422,9 @@ class GatewaySessionCommandsMixin: if readmission_reply is not None: return readmission_reply + if owned: + return await self._retry_in_admission( + event, session_entry, owned, pinned, truncated, last_user_msg, handoff is not None) if handoff is not None: # Composite carrier (one row = retained summary + live ask): the carrier-aware rewind # archives it and inserts the pure scaffold atomically, reselecting the latest carrier @@ -444,6 +454,77 @@ class GatewaySessionCommandsMixin: carry_callback(event, turn) return await self._handle_message(turn) + def _retry_owned_store(self, session_id: str): + """The SessionDB that holds both *session_id*'s transcript and its authority's admission + ledger, or None (no authority: the rewrite-then-send path). One handle means one write + transaction can carry the rewrite and the admission (#139257 step 2).""" + from gateway.session_authorities import active_authority + authority = active_authority(self) + if authority is None: + return None + db = self.session_store._db_for_session_id(session_id) + return db if db is not None and db is authority.db else None + + async def _retry_in_admission(self, event, session_entry, db, pinned, truncated, resend_text, + composite: bool) -> Optional[str]: + """Plan the rewrite, bind it to the re-send and let ``admit_native`` commit both in one write + (``gateway.session_retry_admission``). A refusal leaves the transcript as it was; a + transcript changed since the read refuses both.""" + from gateway import session_retry_admission as retry_admission + from gateway.session_ingress_context import carry_callback + from hermes_state_rewind import RewindTargetUnavailableError + from hermes_state_runtime import RuntimeStoreError + session_id = session_entry.session_id + if composite: + # The carrier-aware rewind, planned on the same snapshot: archive the carrier, keep its scaffold. + try: + plan = await asyncio.to_thread(db.plan_user_turn_rewind, session_id, -1, + require_retryable=True, require_composite=True) + except ValueError as exc: + if isinstance(exc, RewindTargetUnavailableError): + return t("gateway.retry.failed_unchanged") + return t("gateway.retry.unsafe", error=exc) + except RuntimeError: # the history changed under the plan, or the target has no row id + logger.debug("/retry carrier plan refused for %s", session_id, exc_info=True) + return t("gateway.retry.failed_unchanged") + if plan.expected_active_ids != pinned: + return t("gateway.retry.failed_unchanged") + resend_text = plan.live_text + + def write(conn): + from hermes_state_mutation_transcript import rewind_in_transaction + rewind_in_transaction(db, conn, session_id, plan.target_row_id, preserve_compaction_handoff=True, + expected_active_ids=plan.expected_active_ids, + expected_target_content=plan.expected_target_content) + else: + def write(conn): + db.replace_messages_in_transaction(conn, session_id, [dict(m) for m in truncated], active_only=True, + reject_active_turn_lease=True, expected_active_ids=pinned) + + def committed(): + self.session_store._clear_dirty_transcript(session_id) + session_entry.last_prompt_tokens = 0 # transcript was truncated + self._record_model_friction("retry", event.source, session_id) + rewrite = retry_admission.RetryRewrite(db=db, session_id=session_id, source=event.source, + write=write, on_commit=committed) + turn = MessageEvent( + text=resend_text, message_type=MessageType.TEXT, source=event.source, + raw_message=event.raw_message, channel_prompt=event.channel_prompt) + carry_callback(event, turn) + token = retry_admission.bind(rewrite) + try: + return await self._handle_message(turn) + except retry_admission.RetryRewriteRefused as exc: + logger.info("/retry refused with the transcript unchanged for %s: %s", session_id, exc) + return t("gateway.retry.failed_unchanged") + except RuntimeStoreError as exc: + if not rewrite.applied and exc.reason == 'runtime_draining': + # The drain began after the pre-check: the re-send was refused, nothing rewritten. + return t("gateway.busy.drain_rejected", action=self._status_action_gerund()) + raise + finally: + retry_admission.unbind(token) + def _retry_readmission_reply(self, source: SessionSource, resend_text: str) -> Optional[str]: """Refusal the /retry re-send would meet in ``_handle_message``, if any — checked before the transcript rewrite so a refused /retry leaves the transcript untouched. Mirrors the gates diff --git a/hermes_state_messages.py b/hermes_state_messages.py index b7e4db3813..397faeb3ea 100644 --- a/hermes_state_messages.py +++ b/hermes_state_messages.py @@ -804,31 +804,45 @@ class SessionMessagesMixin(SessionDetachedResultsMixin): the whole transcript on every rewind). The live view is identical either way; only the durability of the dropped turns differs. """ - from hermes_state_errors import CompressionSessionClosedError def _do(conn): - if reject_active_turn_lease: - self._check_transcript_write_guards( - conn, session_id, None, reject_active_turn_lease=True, reject_active_compression_lock=True) - elif _ended_by_compression(conn.execute(_ENDED_ROW_SQL, (session_id,)).fetchone()): - raise CompressionSessionClosedError(session_id) - kept = kept_tool_calls = 0 - if archive_dropped: - # Only the first len(messages)+1 live rows matter: the prefix to match plus the row whose - # id anchors the archive UPDATE (which itself covers every later row via `id >= ?`). - live = conn.execute(_LIVE_IDENTITY_SQL, (session_id, len(messages) + 1)).fetchall() - kept = self._stamp_kept_live_prefix(live, messages) - kept_tool_calls = sum(_tool_calls_len(row[4], scalar=1) for row in live[:kept]) - if kept < len(live): - # FTS triggers don't fire on `active`: replaced turns stay searchable (include_inactive=True). - conn.execute("UPDATE messages SET active = 0 WHERE session_id = ? AND active = 1 AND id >= ?", - (session_id, live[kept][0])) - else: - conn.execute(f"DELETE FROM messages WHERE session_id = ?{' AND active = 1' if active_only else ''}", (session_id,)) - inserted, inserted_tool_calls = self._insert_message_rows(conn, session_id, messages[kept:]) - conn.execute(f"{_SET_COUNTERS_SQL} WHERE id = ?", - (kept + inserted, kept_tool_calls + inserted_tool_calls, session_id)) + self.replace_messages_in_transaction( + conn, session_id, messages, active_only=active_only, archive_dropped=archive_dropped, + reject_active_turn_lease=reject_active_turn_lease) self._execute_transcript_write(_do, messages) + def replace_messages_in_transaction(self, conn, session_id: str, messages: list[dict[str, Any]], *, + active_only: bool = False, archive_dropped: bool = False, + reject_active_turn_lease: bool = False, + expected_active_ids: Optional[list[int]] = None) -> None: + """:meth:`replace_messages` inside the caller's write transaction (the gateway ``/retry`` + rewrite commits in its re-send's own admission). ``expected_active_ids`` pins the active set + the replacement was computed from: a row committed since refuses (``RuntimeError``).""" + from hermes_state_errors import CompressionSessionClosedError + if reject_active_turn_lease: + self._check_transcript_write_guards( + conn, session_id, None, reject_active_turn_lease=True, reject_active_compression_lock=True) + elif _ended_by_compression(conn.execute(_ENDED_ROW_SQL, (session_id,)).fetchone()): + raise CompressionSessionClosedError(session_id) + if expected_active_ids is not None and [ + int(row[0]) for row in conn.execute(_ACTIVE_IDS_SQL, (session_id,))] != expected_active_ids: + raise RuntimeError("active transcript changed before the rewrite could be persisted") + kept = kept_tool_calls = 0 + if archive_dropped: + # Only the first len(messages)+1 live rows matter: the prefix to match plus the row whose + # id anchors the archive UPDATE (which itself covers every later row via `id >= ?`). + live = conn.execute(_LIVE_IDENTITY_SQL, (session_id, len(messages) + 1)).fetchall() + kept = self._stamp_kept_live_prefix(live, messages) + kept_tool_calls = sum(_tool_calls_len(row[4], scalar=1) for row in live[:kept]) + if kept < len(live): + # FTS triggers don't fire on `active`: replaced turns stay searchable (include_inactive=True). + conn.execute("UPDATE messages SET active = 0 WHERE session_id = ? AND active = 1 AND id >= ?", + (session_id, live[kept][0])) + else: + conn.execute(f"DELETE FROM messages WHERE session_id = ?{' AND active = 1' if active_only else ''}", (session_id,)) + inserted, inserted_tool_calls = self._insert_message_rows(conn, session_id, messages[kept:]) + conn.execute(f"{_SET_COUNTERS_SQL} WHERE id = ?", + (kept + inserted, kept_tool_calls + inserted_tool_calls, session_id)) + @classmethod def _row_identity(cls, role: str, content: Any, tool_call_id: Any, tool_calls: Any) -> tuple: """The (role, content, tool_call_id, tool_calls) columns a message writes — the identity the diff --git a/hermes_state_rewind.py b/hermes_state_rewind.py index 1b83a34c80..3748b7158d 100644 --- a/hermes_state_rewind.py +++ b/hermes_state_rewind.py @@ -17,6 +17,20 @@ class RewindTargetUnavailableError(ValueError): required a compaction carrier. Surfaces map this to their own "nothing to undo" message.""" +@dataclass +class RewindPlan: + target_row_id: int + expected_active_ids: list[int] + expected_target_content: Any + scaffold: Optional[dict[str, Any]] + prefix: list[dict[str, Any]] + durable_prefix: list[dict[str, Any]] + live_view: dict[str, Any] + live_text: Optional[str] + durable_user: list[int] + user_ordinal: int + + @dataclass class RewindOutcome: prefix: list[dict[str, Any]] # history to install: the warm prefix when ``warm_history`` was given, else durable @@ -55,11 +69,57 @@ class SessionRewindMixin: carrier. ``adopt_row_ids`` (TUI): copy durable ``_row_id`` identities onto the installed warm prefix so clients can address follow-ups by row; the CLI leaves its history shape alone. Out-of-range / wrong-shape targets raise :class:`RewindTargetUnavailableError`.""" - from agent.context_compressor import ( - _DB_PERSISTED_MARKER, history_before_user_originated_turn, retryable_user_text, - split_user_originated_turn, user_originated_turn_view) + from agent.context_compressor import _DB_PERSISTED_MARKER from agent.message_content import flatten_message_text from agent.message_metadata import MESSAGE_UID, message_uid_or_none + + plan = self.plan_user_turn_rewind(session_id, user_ordinal, warm_history=warm_history, + require_retryable=require_retryable, require_composite=require_composite) + (prefix, durable_prefix, live_view, live_text, scaffold, durable_user, + user_ordinal) = (plan.prefix, plan.durable_prefix, plan.live_view, plan.live_text, plan.scaffold, + plan.durable_user, plan.user_ordinal) + try: + result = self.rewind_to_message( + session_id, plan.target_row_id, preserve_compaction_handoff=scaffold is not None, + expected_active_ids=plan.expected_active_ids, expected_target_content=plan.expected_target_content) + except ValueError as exc: # target vanished / changed role under us: same class of failure as out-of-range + raise RewindTargetUnavailableError(str(exc)) from exc + if scaffold is not None: + replacement_id = result.get("replacement_message_id") + if not isinstance(replacement_id, int) or not durable_prefix: + raise RuntimeError("rewind did not retain its compaction handoff") + # The installed scaffold IS the replacement row: carry its identity, not just its row id. + durable_prefix[-1].update({"_row_id": replacement_id, _DB_PERSISTED_MARKER: True}) + if replacement_uid := result.get("replacement_message_uid"): + durable_prefix[-1][MESSAGE_UID] = replacement_uid + prefix[-1] = durable_prefix[-1] + if adopt_row_ids and prefix is not durable_prefix and len(prefix) == len(durable_prefix) and all( + warm.get("role") == durable_message.get("role") + and bool(warm.get("display_kind")) == bool(durable_message.get("display_kind")) + and _comparison_content(warm) == _comparison_content(durable_message) + for warm, durable_message in zip(prefix, durable_prefix) + ): + # Clients address follow-ups by durable row id: keep the richer warm content, adopt the identities. + for warm, durable_message in zip(prefix, durable_prefix): + if isinstance(row_id := durable_message.get("_row_id"), int): + warm["_row_id"] = row_id + if uid := message_uid_or_none(durable_message): + warm[MESSAGE_UID] = uid + return RewindOutcome( + prefix=prefix, live_view=live_view, + live_text=live_text if live_text is not None else flatten_message_text(live_view.get("content")), + rewound_count=int(result.get("rewound_count", 0)), turns_undone=len(durable_user) - user_ordinal) + + def plan_user_turn_rewind( + self, session_id: str, user_ordinal: int, *, warm_history: Optional[list[dict[str, Any]]] = None, + require_retryable: bool = False, require_composite: bool = False, + ) -> RewindPlan: + """Every read and check :meth:`rewind_user_turn` makes before its write, with no write: the + gateway ``/retry`` commits the planned rewind inside its re-send's admission transaction + (``rewind_in_transaction`` with the plan's pins), so a refused re-send leaves it unapplied.""" + from agent.context_compressor import ( + history_before_user_originated_turn, retryable_user_text, + split_user_originated_turn, user_originated_turn_view) from agent.session_persistence import _is_ephemeral_scaffolding expected_active_ids = self.get_active_message_ids(session_id) @@ -102,34 +162,8 @@ class SessionRewindMixin: (user_originated_turn_view(m) for m in stored if m.get("_row_id") == target_row_id), None) if stored_view is None: raise RuntimeError(_HISTORY_CHANGED) - try: - result = self.rewind_to_message( - session_id, target_row_id, preserve_compaction_handoff=scaffold is not None, - expected_active_ids=expected_active_ids, expected_target_content=stored_view.get("content")) - except ValueError as exc: # target vanished / changed role under us: same class of failure as out-of-range - raise RewindTargetUnavailableError(str(exc)) from exc - if scaffold is not None: - replacement_id = result.get("replacement_message_id") - if not isinstance(replacement_id, int) or not durable_prefix: - raise RuntimeError("rewind did not retain its compaction handoff") - # The installed scaffold IS the replacement row: carry its identity, not just its row id. - durable_prefix[-1].update({"_row_id": replacement_id, _DB_PERSISTED_MARKER: True}) - if replacement_uid := result.get("replacement_message_uid"): - durable_prefix[-1][MESSAGE_UID] = replacement_uid - prefix[-1] = durable_prefix[-1] - if adopt_row_ids and prefix is not durable_prefix and len(prefix) == len(durable_prefix) and all( - warm.get("role") == durable_message.get("role") - and bool(warm.get("display_kind")) == bool(durable_message.get("display_kind")) - and _comparison_content(warm) == _comparison_content(durable_message) - for warm, durable_message in zip(prefix, durable_prefix) - ): - # Clients address follow-ups by durable row id: keep the richer warm content, adopt the identities. - for warm, durable_message in zip(prefix, durable_prefix): - if isinstance(row_id := durable_message.get("_row_id"), int): - warm["_row_id"] = row_id - if uid := message_uid_or_none(durable_message): - warm[MESSAGE_UID] = uid - return RewindOutcome( - prefix=prefix, live_view=live_view, - live_text=live_text if live_text is not None else flatten_message_text(live_view.get("content")), - rewound_count=int(result.get("rewound_count", 0)), turns_undone=len(durable_user) - user_ordinal) + return RewindPlan( + target_row_id=target_row_id, expected_active_ids=expected_active_ids, + expected_target_content=stored_view.get("content"), scaffold=scaffold, prefix=prefix, + durable_prefix=durable_prefix, live_view=live_view, live_text=live_text, durable_user=durable_user, + user_ordinal=user_ordinal) diff --git a/tests/gateway/test_retry_authority_refusal.py b/tests/gateway/test_retry_authority_refusal.py index 83b6edb9d7..d440cb6b9d 100644 --- a/tests/gateway/test_retry_authority_refusal.py +++ b/tests/gateway/test_retry_authority_refusal.py @@ -12,6 +12,7 @@ from pathlib import Path import pytest +from agent.i18n import t from gateway.config import GatewayConfig, Platform from gateway.platforms.event import MessageEvent, MessageType from gateway.run import GatewayRunner @@ -102,14 +103,130 @@ async def test_retry_into_a_full_fifo_keeps_the_exchange(served, monkeypatch): assert reply and 'waiting' in reply assert _contents(runner, sid) == ['retry me', 'old answer'] +async def _admitting_exchange(runner, root, rows=None): + """Like ``_exchange`` but the real ``_handle_message`` runs, so the re-send reaches the + authority's ``admit_native``; only the drain is held so the admitted row stays queued.""" + await initialize_gateway_runtime(runner) + register_transport_home(runner, None, root) + adapter = _FakeAdapter() + runner.adapters = {Platform.DISCORD: adapter} + runner._is_user_authorized_for_source = lambda *a, **k: True + source = SessionSource(platform=Platform.DISCORD, chat_id='c1', chat_type='dm', + user_id='u1', user_name='probe') + source._transport_adapter_ref = weakref.ref(adapter) + entry = await runner.async_session_store.get_or_create_session(source) + for row in rows or ({'role': 'user', 'content': 'retry me'}, {'role': 'assistant', 'content': 'old answer'}): + runner.session_store.append_to_transcript(entry.session_id, row) + from gateway.session_authorities import active_authority + authority = active_authority(runner) + authority._schedule = lambda ref: None + return source, entry.session_id, authority + + +def _queued_texts(authority, sid): + import hermes_state_runtime + rows = hermes_state_runtime.list_session_admissions(authority.db, session_id=authority.logical_owner(sid)) + return [row['payload']['text'] for row in rows if row['status'] == 'queued'] + @pytest.mark.asyncio -async def test_an_admissible_retry_still_replaces_the_exchange_once(served): +async def test_a_drain_that_begins_after_the_precheck_leaves_the_exchange(served, monkeypatch): + """MSG-11's race (#139257 step 2): a drain that engages after the pre-check refuses the re-send + inside its admission write, and the rewrite in that same write rolls back with it.""" runner, root = served - source, sid, sent = await _exchange(runner, root) + source, sid, authority = await _admitting_exchange(runner, root) + original = type(runner)._retry_readmission_reply + + def precheck_then_drain(self, *args, **kwargs): + reply = original(self, *args, **kwargs) + self._draining = True # the drain begins once the pre-check has passed + return reply + monkeypatch.setattr(type(runner), '_retry_readmission_reply', precheck_then_drain) reply = await _retry(runner, root, source) - assert reply == 'resent' - assert sent == ['retry me'] - assert _contents(runner, sid) == [] + assert reply and '⏳' in reply + assert _contents(runner, sid) == ['retry me', 'old answer'] + assert _queued_texts(authority, sid) == [] + + +@pytest.mark.asyncio +async def test_a_turn_committed_after_the_read_refuses_the_retry_unchanged(served, monkeypatch): + """The rewrite is fenced on the transcript it was computed from: a row committed between the + read and the admission refuses both, so a newer exchange is never erased.""" + runner, root = served + source, sid, authority = await _admitting_exchange(runner, root) + original = type(runner)._retry_readmission_reply + + def precheck_then_append(self, *args, **kwargs): + self.session_store.append_to_transcript(sid, {'role': 'user', 'content': 'newer ask'}) + return original(self, *args, **kwargs) + monkeypatch.setattr(type(runner), '_retry_readmission_reply', precheck_then_append) + + reply = await _retry(runner, root, source) + + assert reply == t('gateway.retry.failed_unchanged') + assert _contents(runner, sid) == ['retry me', 'old answer', 'newer ask'] + assert _queued_texts(authority, sid) == [] + + +@pytest.mark.asyncio +async def test_an_admitted_retry_rewrites_in_its_own_admission(served): + """A successful /retry still replaces the last exchange exactly once, in its admission.""" + runner, root = served + source, sid, authority = await _admitting_exchange(runner, root, rows=( + {'role': 'user', 'content': 'first'}, {'role': 'assistant', 'content': 'kept'}, + {'role': 'user', 'content': 'retry me'}, {'role': 'assistant', 'content': 'old answer'})) + import asyncio + task = asyncio.ensure_future(_retry(runner, root, source)) + for _ in range(200): + if _queued_texts(authority, sid): + break + await asyncio.sleep(0.01) + assert _queued_texts(authority, sid) == ['retry me'] + assert _contents(runner, sid) == ['first', 'kept'] + task.cancel() + + +def _carrier(ask='REAL ASK'): + from agent.context_compressor import HISTORICAL_TASK_HEADING, SUMMARY_PREFIX, _SUMMARY_END_MARKER + return {'role': 'user', + 'content': f'{SUMMARY_PREFIX}\n{HISTORICAL_TASK_HEADING}\nold task\n\n{_SUMMARY_END_MARKER}\n\n{ask}'} + + +@pytest.mark.asyncio +async def test_a_refused_carrier_retry_keeps_the_carrier(served, monkeypatch): + """The compaction-carrier rewind takes the same path: refused in the write, nothing archived.""" + runner, root = served + source, sid, authority = await _admitting_exchange( + runner, root, rows=(_carrier(), {'role': 'assistant', 'content': 'failed answer'})) + original = type(runner)._retry_readmission_reply + + def precheck_then_drain(self, *args, **kwargs): + reply = original(self, *args, **kwargs) + self._draining = True + return reply + monkeypatch.setattr(type(runner), '_retry_readmission_reply', precheck_then_drain) + + reply = await _retry(runner, root, source) + + assert reply and '⏳' in reply + assert _contents(runner, sid) == [_carrier()['content'], 'failed answer'] + assert _queued_texts(authority, sid) == [] + + +@pytest.mark.asyncio +async def test_an_admitted_carrier_retry_keeps_the_scaffold_and_resends_the_ask(served): + runner, root = served + source, sid, authority = await _admitting_exchange( + runner, root, rows=(_carrier(), {'role': 'assistant', 'content': 'failed answer'})) + import asyncio + task = asyncio.ensure_future(_retry(runner, root, source)) + for _ in range(200): + if _queued_texts(authority, sid): + break + await asyncio.sleep(0.01) + assert _queued_texts(authority, sid) == ['REAL ASK'] + active = runner.session_store.load_transcript(sid) + assert len(active) == 1 and active[0]['display_kind'] == 'hidden' and 'REAL ASK' not in active[0]['content'] + task.cancel()