diff --git a/apps/desktop/src/app/chat/session-tile-actions.ts b/apps/desktop/src/app/chat/session-tile-actions.ts index 719f9648b275d..b7c038a1563a6 100644 --- a/apps/desktop/src/app/chat/session-tile-actions.ts +++ b/apps/desktop/src/app/chat/session-tile-actions.ts @@ -41,7 +41,10 @@ import { planEdit, planReload, planRestore, + rebindSurvivorRowIds, runRewindSubmit, + type SurvivorUserRowIds, + survivorRowIdsFrom, truncateSubmitParams } from '../session/hooks/use-prompt-actions/rewind' import { useSubmitPrompt } from '../session/hooks/use-prompt-actions/submit' @@ -391,6 +394,22 @@ export function useSessionTileActions({ runtimeId, scope, storedSessionId }: Ses [requestGateway] ) + // After a durable rewind the surviving bubbles' cached rowIds are stale (the + // gateway re-inserted the kept prefix as new SQLite rows). Rebind them to the + // authoritative post-rewrite ids so the NEXT rewind/edit/regenerate doesn't + // send a dead id and get refused with 4018 (consecutive-rewind staleness, + // #83202 review). + const applySurvivorRowIds = useCallback( + (survivorRowIds: SurvivorUserRowIds | undefined) => { + if (!survivorRowIds) { + return + } + + update(state => ({ ...state, messages: rebindSurvivorRowIds(state.messages, survivorRowIds) })) + }, + [update] + ) + const reloadFromMessage = useCallback( async (parentId: string | null) => { const state = readState() @@ -408,7 +427,7 @@ export function useSessionTileActions({ runtimeId, scope, storedSessionId }: Ses update(current => applyReloadOptimistic(current, plan)) try { - await requestGateway( + const result = await requestGateway<{ survivor_user_row_ids?: unknown }>( 'prompt.submit', { session_id: runtimeIdRef.current, @@ -417,12 +436,14 @@ export function useSessionTileActions({ runtimeId, scope, storedSessionId }: Ses }, PROMPT_SUBMIT_REQUEST_TIMEOUT_MS ) + + applySurvivorRowIds(survivorRowIdsFrom(result)) } catch (err) { update(current => ({ ...current, busy: false, awaitingResponse: false })) notifyError(err, copy.regenerateFailed) } }, - [copy.regenerateFailed, readState, requestGateway, update] + [applySurvivorRowIds, copy.regenerateFailed, readState, requestGateway, update] ) const restoreToMessage = useCallback( @@ -440,13 +461,15 @@ export function useSessionTileActions({ runtimeId, scope, storedSessionId }: Ses update(state => applyRewindOptimistic(state, plan.sourceIndex)) try { - await submitRewind(plan.text, plan.truncateOrdinal, wasBusy, plan.truncateMessageId, plan.truncateRowId) + applySurvivorRowIds( + await submitRewind(plan.text, plan.truncateOrdinal, wasBusy, plan.truncateMessageId, plan.truncateRowId) + ) } catch (err) { update(state => ({ ...state, busy: false, awaitingResponse: false, messages })) throw err } }, - [readMessages, readState, submitRewind, update] + [applySurvivorRowIds, readMessages, readState, submitRewind, update] ) const editMessage = useCallback( @@ -469,13 +492,15 @@ export function useSessionTileActions({ runtimeId, scope, storedSessionId }: Ses update(state => applyRewindOptimistic(state, plan.sourceIndex, plan.editedMessage)) try { - await submitRewind(plan.text, plan.truncateOrdinal, wasBusy, plan.truncateMessageId, plan.truncateRowId) + applySurvivorRowIds( + await submitRewind(plan.text, plan.truncateOrdinal, wasBusy, plan.truncateMessageId, plan.truncateRowId) + ) } catch (err) { update(state => ({ ...state, busy: false, awaitingResponse: false, messages })) notifyError(err, copy.editFailed) } }, - [copy.editFailed, readMessages, readState, submitRewind, update] + [applySurvivorRowIds, copy.editFailed, readMessages, readState, submitRewind, update] ) // Branch-visibility sync (assistant-ui hides non-active branches). diff --git a/apps/desktop/src/app/session/hooks/use-prompt-actions/index.ts b/apps/desktop/src/app/session/hooks/use-prompt-actions/index.ts index 9da23ed455ac3..1229d33a66b8d 100644 --- a/apps/desktop/src/app/session/hooks/use-prompt-actions/index.ts +++ b/apps/desktop/src/app/session/hooks/use-prompt-actions/index.ts @@ -56,7 +56,10 @@ import { planEdit, planReload, planRestore, + rebindSurvivorRowIds, runRewindSubmit, + type SurvivorUserRowIds, + survivorRowIdsFrom, truncateSubmitParams } from './rewind' import { useSlashCommand } from './slash' @@ -784,6 +787,25 @@ export function usePromptActions({ [activeSessionIdRef, appendSessionTextMessage, requestGateway, selectedStoredSessionIdRef, updateSessionState] ) + // After a durable rewind the surviving bubbles' cached rowIds are stale (the + // gateway re-inserted the kept prefix as new SQLite rows). Rebind them to the + // authoritative post-rewrite ids so the NEXT rewind/edit/regenerate doesn't + // send a dead id and get refused with 4018 (consecutive-rewind staleness, + // #83202 review). + const applySurvivorRowIds = useCallback( + (sessionId: string, survivorRowIds: SurvivorUserRowIds | undefined) => { + if (!survivorRowIds) { + return + } + + updateSessionState(sessionId, state => ({ + ...state, + messages: rebindSurvivorRowIds(state.messages, survivorRowIds) + })) + }, + [updateSessionState] + ) + const reloadFromMessage = useCallback( async (parentId: string | null) => { // Ref, not the closure-captured prop — a truncating resubmit aimed at a @@ -804,7 +826,7 @@ export function usePromptActions({ updateSessionState(sessionId, state => applyReloadOptimistic(state, plan)) try { - await requestGateway( + const result = await requestGateway<{ survivor_user_row_ids?: unknown }>( 'prompt.submit', { session_id: sessionId, @@ -813,6 +835,8 @@ export function usePromptActions({ }, PROMPT_SUBMIT_REQUEST_TIMEOUT_MS ) + + applySurvivorRowIds(sessionId, survivorRowIdsFrom(result)) } catch (err) { updateSessionState(sessionId, state => ({ ...state, @@ -822,7 +846,7 @@ export function usePromptActions({ notifyError(err, copy.regenerateFailed) } }, - [activeSessionIdRef, copy.regenerateFailed, requestGateway, updateSessionState] + [activeSessionIdRef, applySurvivorRowIds, copy.regenerateFailed, requestGateway, updateSessionState] ) // Cursor-style "restore checkpoint": rewind the conversation to a past user @@ -889,7 +913,7 @@ export function usePromptActions({ updateSessionState(sessionId, state => applyRewindOptimistic(state, plan.sourceIndex)) try { - await submitRewindPrompt( + const survivorRowIds = await submitRewindPrompt( sessionId, plan.text, plan.truncateOrdinal, @@ -897,6 +921,8 @@ export function usePromptActions({ busyRef.current || $busy.get(), plan.truncateRowId ) + + applySurvivorRowIds(sessionId, survivorRowIds) } catch (err) { // The rewind never landed (e.g. the gateway stayed busy past the retry // deadline). Roll the optimistic truncation back to the full original @@ -914,7 +940,7 @@ export function usePromptActions({ throw err } }, - [activeSessionIdRef, busyRef, submitRewindPrompt, updateSessionState] + [activeSessionIdRef, applySurvivorRowIds, busyRef, submitRewindPrompt, updateSessionState] ) const editMessage = useCallback( @@ -944,7 +970,7 @@ export function usePromptActions({ updateSessionState(sessionId, state => applyRewindOptimistic(state, plan.sourceIndex, plan.editedMessage)) try { - await submitRewindPrompt( + const survivorRowIds = await submitRewindPrompt( sessionId, plan.text, plan.truncateOrdinal, @@ -952,6 +978,8 @@ export function usePromptActions({ busyRef.current || $busy.get(), plan.truncateRowId ) + + applySurvivorRowIds(sessionId, survivorRowIds) } catch (err) { // Roll the optimistic edit/truncation back to the original history so the // UI stays in sync with what's persisted instead of stranding a partial @@ -963,7 +991,7 @@ export function usePromptActions({ notifyError(err, copy.editFailed) } }, - [activeSessionIdRef, busyRef, copy.editFailed, submitRewindPrompt, updateSessionState] + [activeSessionIdRef, applySurvivorRowIds, busyRef, copy.editFailed, submitRewindPrompt, updateSessionState] ) const handleThreadMessagesChange = useCallback( diff --git a/apps/desktop/src/app/session/hooks/use-prompt-actions/rewind.test.ts b/apps/desktop/src/app/session/hooks/use-prompt-actions/rewind.test.ts index 4f32e720a5e9f..2ba2c438b58da 100644 --- a/apps/desktop/src/app/session/hooks/use-prompt-actions/rewind.test.ts +++ b/apps/desktop/src/app/session/hooks/use-prompt-actions/rewind.test.ts @@ -1,6 +1,8 @@ import { describe, expect, it } from 'vitest' -import { truncateSubmitParams } from './rewind' +import { type ChatMessage, textPart } from '@/lib/chat-messages' + +import { rebindSurvivorRowIds, survivorRowIdsFrom, truncateSubmitParams } from './rewind' describe('truncateSubmitParams', () => { it('omits truncation fields when no ordinal is set', () => { @@ -58,3 +60,68 @@ describe('truncateSubmitParams', () => { }) }) }) + +describe('survivorRowIdsFrom', () => { + it('returns undefined when the field is absent or not an array (older gateway)', () => { + expect(survivorRowIdsFrom(undefined)).toBeUndefined() + expect(survivorRowIdsFrom({ status: 'streaming' })).toBeUndefined() + expect(survivorRowIdsFrom({ survivor_user_row_ids: 'nope' })).toBeUndefined() + }) + + it('keeps integer ids and nulls anything else', () => { + expect(survivorRowIdsFrom({ survivor_user_row_ids: [7, null, 9.5, '11', 12] })).toEqual([7, null, null, null, 12]) + }) +}) + +describe('rebindSurvivorRowIds', () => { + const user = (id: string, rowId?: number, hidden?: boolean): ChatMessage => ({ + id, + role: 'user', + parts: [textPart(`text ${id}`)], + ...(rowId !== undefined ? { rowId } : {}), + ...(hidden ? { hidden } : {}) + }) + const assistant = (id: string, rowId?: number): ChatMessage => ({ + id, + role: 'assistant', + parts: [textPart(`reply ${id}`)], + ...(rowId !== undefined ? { rowId } : {}) + }) + + it('rebinds surviving visible user turns positionally and clears the resubmitted turn', () => { + // Post-rewind state: two survivors + the resubmitted turn (stale rowId 5). + const messages = [user('u0', 1), assistant('a0', 2), user('u1', 3), assistant('a1', 4), user('u2', 5)] + const rebound = rebindSurvivorRowIds(messages, [7, 9]) + + expect(rebound[0].rowId).toBe(7) + expect(rebound[2].rowId).toBe(9) + // Resubmitted turn is past the survivor list — its durable id doesn't + // exist yet, and keeping the stale one would 4018 the next rewind. + expect(rebound[4].rowId).toBeUndefined() + // Assistant rows are untouched (only user turns are rewind targets). + expect(rebound[1].rowId).toBe(2) + }) + + it('clears the cached id for null entries instead of keeping a stale one', () => { + const messages = [user('u0', 1), user('u1', 3)] + const rebound = rebindSurvivorRowIds(messages, [null, 9]) + + expect(rebound[0].rowId).toBeUndefined() + expect(rebound[1].rowId).toBe(9) + }) + + it('skips hidden user turns — same visible-user filter as the ordinal math', () => { + const messages = [user('u0', 1), user('hidden', 2, true), user('u1', 3)] + const rebound = rebindSurvivorRowIds(messages, [7, 9]) + + expect(rebound[0].rowId).toBe(7) + expect(rebound[1].rowId).toBe(2) // hidden: untouched + expect(rebound[2].rowId).toBe(9) + }) + + it('preserves object identity when nothing changes', () => { + const messages = [user('u0', 7)] + + expect(rebindSurvivorRowIds(messages, [7])[0]).toBe(messages[0]) + }) +}) diff --git a/apps/desktop/src/app/session/hooks/use-prompt-actions/rewind.ts b/apps/desktop/src/app/session/hooks/use-prompt-actions/rewind.ts index 23d9d3809580b..d4b02639dab30 100644 --- a/apps/desktop/src/app/session/hooks/use-prompt-actions/rewind.ts +++ b/apps/desktop/src/app/session/hooks/use-prompt-actions/rewind.ts @@ -26,6 +26,61 @@ import { type RequestGateway = (method: string, params?: Record, timeoutMs?: number) => Promise +/** + * Post-rewrite durable ids of the surviving visible user turns, in visible-user + * ordinal order — the gateway's `survivor_user_row_ids` on a truncating + * `prompt.submit`. A rewind's `replace_messages` re-inserts the kept prefix as + * NEW SQLite rows, so every pre-rewind `ChatMessage.rowId` on a surviving + * bubble is stale the moment the rewind lands; targeting one on the next + * rewind/edit/regenerate gets a fail-closed 4018 from the gateway. `null` + * means that turn has no durable id (drop the cached one, don't keep a stale + * one). Absent entirely = the submit didn't truncate a durable session (or an + * older gateway) — leave state untouched. + */ +export type SurvivorUserRowIds = readonly (null | number)[] + +interface PromptSubmitResult { + status?: string + survivor_user_row_ids?: unknown +} + +export function survivorRowIdsFrom(result: PromptSubmitResult | undefined): SurvivorUserRowIds | undefined { + const raw = result?.survivor_user_row_ids + + if (!Array.isArray(raw)) { + return undefined + } + + return raw.map(entry => (typeof entry === 'number' && Number.isInteger(entry) ? entry : null)) +} + +/** + * Rebind the surviving visible user turns to their authoritative post-rewind + * row ids (positional, same visible-user filter `visibleUserOrdinal` uses — + * the exact parity truncate ordinals already rely on). Turns past the end of + * the survivor list — the resubmitted turn itself, whose durable id doesn't + * exist yet — and `null` entries get their cached rowId cleared instead: a + * stale id now addresses an archived row and would be refused with 4018. + */ +export function rebindSurvivorRowIds(messages: ChatMessage[], survivorRowIds: SurvivorUserRowIds): ChatMessage[] { + let ordinal = 0 + + return messages.map(message => { + if (message.role !== 'user' || message.hidden) { + return message + } + + const next = ordinal < survivorRowIds.length ? survivorRowIds[ordinal] : null + ordinal += 1 + + if (typeof next === 'number') { + return message.rowId === next ? message : { ...message, rowId: next } + } + + return message.rowId === undefined ? message : { ...message, rowId: undefined } + }) +} + /** * Build `prompt.submit` truncation params. `confirm_truncate` states that this * submit really is a rewind/edit/regenerate: the gateway drops history only for @@ -70,6 +125,10 @@ export function truncateSubmitParams( * / `truncate_before_message_id` / `truncate_before_row_id` (drops that user turn + everything after). * Idle rewinds submit directly; live/stuck turns interrupt first, and a raced * "session busy" response interrupts + retries through the shared busy gate. + * + * Resolves with the gateway's post-rewrite survivor row ids (see + * `SurvivorUserRowIds`) so the caller can rebind surviving bubbles, or + * undefined when the submit didn't truncate a durable transcript. */ export async function runRewindSubmit( requestGateway: RequestGateway, @@ -80,7 +139,7 @@ export async function runRewindSubmit( interruptFirst: boolean, recovery?: { storedSessionId?: null | string; onSessionRecovered?: (sessionId: string) => void }, truncateRowId?: number -): Promise { +): Promise { // Recovery may rebind the live id mid-flight; interrupt/submit must both // follow it rather than pinning the dead one. let liveSessionId = sessionId @@ -94,7 +153,7 @@ export async function runRewindSubmit( } const submitFor = (targetId: string) => - requestGateway( + requestGateway( 'prompt.submit', { session_id: targetId, @@ -105,15 +164,22 @@ export async function runRewindSubmit( ) const submit = async () => { - const { sessionId: usedId } = await withSessionNotFoundResume(liveSessionId, recovery?.storedSessionId, submitFor, { - requestGateway, - onRecovered: recoveredId => { - liveSessionId = recoveredId - recovery?.onSessionRecovered?.(recoveredId) + const { result, sessionId: usedId } = await withSessionNotFoundResume( + liveSessionId, + recovery?.storedSessionId, + submitFor, + { + requestGateway, + onRecovered: recoveredId => { + liveSessionId = recoveredId + recovery?.onSessionRecovered?.(recoveredId) + } } - }) + ) liveSessionId = usedId + + return survivorRowIdsFrom(result) } if (interruptFirst) { @@ -121,14 +187,15 @@ export async function runRewindSubmit( } try { - await submit() + return await submit() } catch (err) { if (!isSessionBusyError(err)) { throw err } await interrupt() - await withSessionBusyRetry(submit) + + return await withSessionBusyRetry(submit) } } diff --git a/tests/test_tui_gateway_server.py b/tests/test_tui_gateway_server.py index 561d0b67d278f..9f624eda25e76 100644 --- a/tests/test_tui_gateway_server.py +++ b/tests/test_tui_gateway_server.py @@ -17820,3 +17820,105 @@ def test_prompt_submit_row_id_db_fallback_ordinal_mapping_verifies_content( assert len(sess["history"]) == 5 finally: server._sessions.pop(sid, None) + + +def test_prompt_submit_consecutive_rewinds_with_returned_survivor_row_ids( + monkeypatch, tmp_path +): + """#83202 review (consecutive-rewind staleness): replace_messages re-inserts + the surviving prefix as NEW rows, so the pre-rewind client row ids die on + the first rewind. The submit response must return the fresh survivor ids, + and a second rewind using them must succeed where the stale id fail-closes. + """ + from hermes_state import SessionDB + + db = SessionDB(db_path=tmp_path / "rowid-consec.db") + session_key = "real-db-consec-rewind" + db.create_session(session_key, "cli") + msgs = [ + {"role": "user", "content": "first"}, + {"role": "assistant", "content": "reply 1"}, + {"role": "user", "content": "second"}, + {"role": "assistant", "content": "reply 2"}, + {"role": "user", "content": "third"}, + {"role": "assistant", "content": "reply 3"}, + ] + with db._lock: + db._insert_message_rows(db._conn, session_key, msgs) + db._conn.commit() + original_row_ids = [m["_row_id"] for m in msgs] + + sess = _session(history=[dict(m) for m in msgs], session_key=session_key) + sid = "real-db-consec-rewind-sid" + server._sessions[sid] = sess + monkeypatch.setattr(server, "_get_db", lambda: db) + monkeypatch.setattr(server, "_start_agent_build", lambda *a, **k: None) + monkeypatch.setattr(server, "_start_inflight_turn", lambda *a, **k: None) + + try: + # Rewind 1: cut before "third" (last user turn). Survivors: turns + # "first" + "second" (+ assistant replies) — re-inserted as NEW rows. + resp1 = server.handle_request( + { + "id": "1", + "method": "prompt.submit", + "params": { + "session_id": sid, + "text": "rewound third", + "truncate_before_row_id": original_row_ids[4], + "truncate_before_user_ordinal": 2, + "confirm_truncate": True, + }, + } + ) + assert resp1.get("error") is None, resp1 + survivors = resp1["result"].get("survivor_user_row_ids") + # Fresh ids for the two surviving user turns, in visible-user order. + assert isinstance(survivors, list) and len(survivors) == 2 + assert all(isinstance(r, int) for r in survivors) + # They must be NEW rows — the old ids are archived (active=0) now. + assert set(survivors).isdisjoint(set(original_row_ids)) + sess["running"] = False + + # Rewind 2a: the STALE pre-rewind id for "second" must fail closed. + stale_resp = server.handle_request( + { + "id": "2", + "method": "prompt.submit", + "params": { + "session_id": sid, + "text": "rewound second (stale id)", + "truncate_before_row_id": original_row_ids[2], + "truncate_before_user_ordinal": 1, + "confirm_truncate": True, + }, + } + ) + assert stale_resp.get("error") is not None + assert stale_resp["error"]["code"] == 4018 + assert len(sess["history"]) == 4 # nothing cut + + # Rewind 2b: the RETURNED survivor id for "second" must succeed. + resp2 = server.handle_request( + { + "id": "3", + "method": "prompt.submit", + "params": { + "session_id": sid, + "text": "rewound second (fresh id)", + "truncate_before_row_id": survivors[1], + "truncate_before_user_ordinal": 1, + "confirm_truncate": True, + }, + } + ) + assert resp2.get("error") is None, resp2 + assert len(sess["history"]) == 2 + assert sess["history"][0]["content"] == "first" + active = db.get_messages_as_conversation(session_key) + assert [m["content"] for m in active] == ["first", "reply 1"] + # And the second response rebinds again: one surviving user turn. + survivors2 = resp2["result"].get("survivor_user_row_ids") + assert isinstance(survivors2, list) and len(survivors2) == 1 + finally: + server._sessions.pop(sid, None) diff --git a/tui_gateway/methods_prompt.py b/tui_gateway/methods_prompt.py index 3547c5b150bd6..d4fdb5c3dccec 100644 --- a/tui_gateway/methods_prompt.py +++ b/tui_gateway/methods_prompt.py @@ -349,6 +349,10 @@ def _(rid, params: dict) -> dict: # claim so this prompt starts normally instead of being stranded in a # queue whose drain already ran. + # Filled when this submit performed a truncation against a durable session: + # the fresh post-rewrite row ids of the surviving user turns, for client + # rowId rebinding (see comment at the assignment site). + survivor_user_row_ids = None with session["history_lock"]: # A watch session's run lives in the PARENT turn, so its own running # flag is False — without this, typing mid-run builds a second agent @@ -586,6 +590,22 @@ def _(rid, params: dict) -> dict: ) session["history"] = truncated session["history_version"] = int(session.get("history_version", 0)) + 1 + if db is not None: + # replace_messages re-inserted the surviving prefix as NEW rows + # and stamped fresh _row_id values onto these same dicts. + # Surface the surviving user-turn ids (in visible-user-ordinal + # order) so the client can rebind its cached rowId stamps — + # otherwise a second rewind targeting an older surviving turn + # sends the pre-rewind id and the fail-closed resolver refuses + # it with 4018 (#83202 review: consecutive-rewind staleness). + # Ordinal order matches the client's visible-user filter the + # same way truncate ordinals already do. Entries are None when + # a row somehow has no stamp — the client must drop its cached + # id for that turn rather than keep a stale one. + survivor_user_row_ids = [ + _message_row_id(truncated[i]) + for i in _history_user_indices(truncated) + ] session["running"] = True session["_turn_cancel_requested"] = False session["last_active"] = time.time() @@ -594,6 +614,13 @@ def _(rid, params: dict) -> dict: if turn_isolation: isolated_response = _submit_prompt_to_compute_host(rid, sid, session, text) if not isolated_response.get("error"): + if survivor_user_row_ids is not None: + # The truncation already happened inline above (memory + DB), + # before compute-host dispatch — the rebind payload applies to + # this path exactly as it does to the inline one. + isolated_response["result"][ + "survivor_user_row_ids" + ] = survivor_user_row_ids return isolated_response logger.warning( "compute-host dispatch failed for session %s; falling back inline: %s", @@ -678,7 +705,17 @@ def _(rid, params: dict) -> dict: # `running` flag (a turn that died without clearing it) and recover the latter. session["_run_thread"] = run_thread run_thread.start() - return _ok(rid, {"status": "streaming"}) + return _ok( + rid, + { + "status": "streaming", + **( + {"survivor_user_row_ids": survivor_user_row_ids} + if survivor_user_row_ids is not None + else {} + ), + }, + ) @method("clipboard.paste")