diff --git a/apps/desktop/src/app/session/hooks/use-session-actions/utils.test.ts b/apps/desktop/src/app/session/hooks/use-session-actions/utils.test.ts index eb2633e54001..a83f8a384026 100644 --- a/apps/desktop/src/app/session/hooks/use-session-actions/utils.test.ts +++ b/apps/desktop/src/app/session/hooks/use-session-actions/utils.test.ts @@ -435,6 +435,39 @@ describe('preserveLocalPendingTurnMessages', () => { expect(preserveLocalPendingTurnMessages(compressedAuthority, pollutedWarmCache)).toBe(compressedAuthority) }) + // A mid-turn redirect inserts its correction as a SECOND optimistic user row + // for the same turn. Keeping only the newest dropped the prompt that started + // it, so a resume repainted the thread with the user's message missing. + it('keeps every optimistic user row in the live run after a mid-turn redirect', () => { + const previous = [ + msg('user-1000', 'user', 'remove the session counts'), + msg('user-2000', 'user', 'hurry up'), + msg('assistant-stream-1', 'assistant', 'Moving.', { pending: true }) + ] + + expect(preserveLocalPendingTurnMessages([], previous).map(message => message.id)).toEqual([ + 'user-1000', + 'user-2000', + 'assistant-stream-1' + ]) + }) + + it('still drops optimistic rows separated from the live run by an assistant reply', () => { + const previous = [ + msg('user-stale', 'user', 'compressed-away prompt'), + msg('assistant-stale', 'assistant', 'compressed-away reply'), + msg('user-1000', 'user', 'the live prompt'), + msg('user-2000', 'user', 'the correction'), + msg('assistant-stream-1', 'assistant', 'Moving.', { pending: true }) + ] + + expect(preserveLocalPendingTurnMessages([], previous).map(message => message.id)).toEqual([ + 'user-1000', + 'user-2000', + 'assistant-stream-1' + ]) + }) + // #67603: the gateway persists model-switch / personality notices as role=user // ([System: …], tui_gateway/server.py). A single trailing marker is already // handled by the latestAuthoritativeUser guard above, but TWO switches around @@ -538,6 +571,46 @@ describe('preserveLocalPendingTurnMessages', () => { }) describe('appendLiveSessionProjection', () => { + // Corrections typed while a turn ran are their own user bubbles on the same + // turn. Resume must rebuild the prompt AND every correction, in order. + it('projects mid-turn redirect corrections after the prompt that started the turn', () => { + const restored = appendLiveSessionProjection([], { + session_id: 'runtime-1', + inflight: { + user: 'remove the session counts', + corrections: ['hurry up', 'and the worktree ones'], + assistant: 'Moving.', + streaming: true + } + }) + + expect(restored.map(message => message.parts.map(part => ('text' in part ? part.text : '')).join(''))).toEqual([ + 'remove the session counts', + 'hurry up', + 'and the worktree ones', + 'Moving.' + ]) + }) + + it('does not re-project a correction the transcript already persisted', () => { + const stored = [msg('stored-user', 'user', 'remove the session counts'), msg('stored-fix', 'user', 'hurry up')] + + const restored = appendLiveSessionProjection(stored, { + session_id: 'runtime-1', + inflight: { + user: 'remove the session counts', + corrections: ['hurry up'], + assistant: 'Moving.', + streaming: true + } + }) + + expect(restored.filter(message => message.role === 'user').map(message => message.id)).toEqual([ + 'stored-user', + 'stored-fix' + ]) + }) + it('does not duplicate the inflight user when the persisted turn carries @image refs', () => { // By the time a stored transcript reaches appendLiveSessionProjection it // has already been run through toChatMessages, so the @image directive has diff --git a/apps/desktop/src/app/session/hooks/use-session-actions/utils.ts b/apps/desktop/src/app/session/hooks/use-session-actions/utils.ts index f06069ca75fd..f6e19919f061 100644 --- a/apps/desktop/src/app/session/hooks/use-session-actions/utils.ts +++ b/apps/desktop/src/app/session/hooks/use-session-actions/utils.ts @@ -326,6 +326,26 @@ export function preserveLocalPendingTurnMessages( .reverse() .find(message => message.role === 'user' && message.id.startsWith('user-')) + // A mid-turn redirect inserts its correction as a second optimistic user row + // directly before the live reply, so one turn can own a contiguous RUN of + // them. Preserving only the newest keeps the correction and drops the prompt + // that started the turn. Widen to the run — but only the contiguous one: any + // `user-*` row separated by an assistant reply is stale post-compression + // history, which is what the newest-only rule exists to discard. + const liveOptimisticUsers = new Set() + + if (newestOptimisticUser) { + for (let index = previousMessages.indexOf(newestOptimisticUser); index >= 0; index -= 1) { + const candidate = previousMessages[index] + + if (candidate.role !== 'user' || !candidate.id.startsWith('user-')) { + break + } + + liveOptimisticUsers.add(candidate) + } + } + const latestAuthoritativeUser = [...nextMessages].reverse().find(message => message.role === 'user') const preserved: ChatMessage[] = [] @@ -346,7 +366,7 @@ export function preserveLocalPendingTurnMessages( continue } - if (isOptimisticUser && message !== newestOptimisticUser) { + if (isOptimisticUser && !liveOptimisticUsers.has(message)) { continue } @@ -392,13 +412,27 @@ export function appendLiveSessionProjection( const inflightUser = projection.inflight?.user?.trim() ?? '' const inflightAssistant = projection.inflight?.assistant ?? '' const inflightStreaming = Boolean(projection.inflight?.streaming) + + // Mid-turn redirect corrections. They are additional user bubbles belonging + // to this same turn, ordered after the prompt that started it. + const inflightCorrections = (projection.inflight?.corrections ?? []) + .map(correction => correction?.trim() ?? '') + .filter(Boolean) + // A retained failed turn (the gateway keeps error snapshots replayable when // the terminal frame may have been lost to a disconnect) — surface the // failure on the projected row instead of rendering the partial as healthy. const inflightError = projection.inflight?.error?.trim() ?? '' const queuedUser = projection.queued?.user?.trim() ?? '' - if (!inflightUser && !inflightAssistant && !inflightStreaming && !inflightError && !queuedUser) { + if ( + !inflightUser && + !inflightAssistant && + !inflightStreaming && + !inflightError && + !queuedUser && + !inflightCorrections.length + ) { return messages } @@ -409,10 +443,20 @@ export function appendLiveSessionProjection( // both makes a backgrounded prompt appear twice when its session is reopened. // Only suppress the projection when the latest authoritative user row is the // same turn — older identical prompts must not hide a newly accepted repeat. - const latestUser = [...messages].reverse().find(message => message.role === 'user') + // A mid-turn redirect gives that turn a RUN of user rows (prompt + + // corrections), so match the contiguous run ending at the latest user row + // rather than the single last one. + const latestUserIndex = messages.map(message => message.role).lastIndexOf('user') + const latestUserRun: ChatMessage[] = [] - const inflightUserAlreadyPersisted = - latestUser && textWithoutImageRefs(chatMessageText(latestUser)) === textWithoutImageRefs(inflightUser) + for (let index = latestUserIndex; index >= 0 && messages[index].role === 'user'; index -= 1) { + latestUserRun.unshift(messages[index]) + } + + const persistedInLatestRun = (text: string): boolean => + latestUserRun.some(message => textWithoutImageRefs(chatMessageText(message)) === textWithoutImageRefs(text)) + + const inflightUserAlreadyPersisted = Boolean(inflightUser) && persistedInLatestRun(inflightUser) if (inflightUser && !inflightUserAlreadyPersisted) { projected.push({ @@ -422,6 +466,22 @@ export function appendLiveSessionProjection( }) } + // Corrections typed while the turn ran. Each is its own bubble, placed after + // the original prompt and before the reply they redirected — the same order + // the live transcript showed. Skip any the transcript already holds so a + // resume doesn't double them. + for (const [index, correction] of inflightCorrections.entries()) { + if (persistedInLatestRun(correction)) { + continue + } + + projected.push({ + id: `user-inflight-correction-${index}-${sessionId}`, + role: 'user', + parts: [textPart(correction)] + }) + } + // Keep a pending assistant boundary even before the first delta when a // queued user turn follows it. This preserves the two distinct turns. if (inflightAssistant || inflightStreaming || inflightError || (inflightUser && queuedUser)) { diff --git a/apps/desktop/src/lib/inflight-turn-journal.test.ts b/apps/desktop/src/lib/inflight-turn-journal.test.ts index ffa1c4300651..41d7a362d469 100644 --- a/apps/desktop/src/lib/inflight-turn-journal.test.ts +++ b/apps/desktop/src/lib/inflight-turn-journal.test.ts @@ -238,3 +238,63 @@ describe('mergeInFlightMessages', () => { expect(result.caughtUp).toBe(false) }) }) + +describe('mid-turn redirect corrections', () => { + beforeEach(() => { + window.localStorage.clear() + vi.useFakeTimers() + }) + + afterEach(() => { + vi.useRealTimers() + }) + + // A redirect inserts its correction as a second user row directly before the + // live reply, so the turn opens with a RUN of user rows. Journaling only back + // to the nearest one lost the prompt that actually started the turn — the + // vanishing user bubble. + it('journals the whole user run, not just the correction', () => { + persistInFlightTurnState({ + awaitingResponse: false, + busy: true, + messages: [ + user('user-1', 'remove the session counts'), + user('user-2', 'hurry up'), + assistant('assistant-stream-1', 'Moving.', { pending: true }) + ], + storedSessionId: 'stored-redirect', + streamId: 'assistant-stream-1', + turnStartedAt: Date.now() + }) + vi.advanceTimersByTime(400) + + const journaled = readInFlightTurnJournal('stored-redirect')?.messages ?? [] + + expect(journaled.map(message => message.parts.map(part => (part as { text: string }).text).join(''))).toEqual([ + 'remove the session counts', + 'hurry up', + 'Moving.' + ]) + }) + + it('still stops at an assistant boundary so prior turns are not journaled', () => { + persistInFlightTurnState({ + awaitingResponse: false, + busy: true, + messages: [ + user('user-old', 'an earlier turn'), + assistant('assistant-old', 'an earlier answer'), + user('user-1', 'the live prompt'), + assistant('assistant-stream-1', 'Moving.', { pending: true }) + ], + storedSessionId: 'stored-boundary', + streamId: 'assistant-stream-1', + turnStartedAt: Date.now() + }) + vi.advanceTimersByTime(400) + + const journaled = readInFlightTurnJournal('stored-boundary')?.messages ?? [] + + expect(journaled.map(message => message.id)).toEqual(['user-1', 'assistant-stream-1']) + }) +}) diff --git a/apps/desktop/src/lib/inflight-turn-journal.ts b/apps/desktop/src/lib/inflight-turn-journal.ts index 8056fa7af102..e93a0f097540 100644 --- a/apps/desktop/src/lib/inflight-turn-journal.ts +++ b/apps/desktop/src/lib/inflight-turn-journal.ts @@ -216,6 +216,14 @@ function recoverableTail(messages: ChatMessage[], streamId: null | string): Chat if (visible[index].role === 'user') { start = index + // A mid-turn redirect inserts its correction as another user row right + // before the live reply, so the turn can open with a RUN of user rows. + // Keep walking back over them: stopping at the nearest one journals the + // correction alone and loses the prompt that actually started the turn. + while (start > 0 && visible[start - 1].role === 'user') { + start -= 1 + } + break } } diff --git a/apps/desktop/src/types/hermes.ts b/apps/desktop/src/types/hermes.ts index 669813b0ab25..8857bca43277 100644 --- a/apps/desktop/src/types/hermes.ts +++ b/apps/desktop/src/types/hermes.ts @@ -523,6 +523,9 @@ export interface SessionResumeResponse { } inflight?: null | { assistant?: string + /** Mid-turn redirect corrections, oldest first. The turn's original prompt + * stays in `user`; these are the follow-ups typed while it ran. */ + corrections?: string[] /** Retained failed turn: the error the terminal frame carried (the frame * itself may have been lost to a disconnect). */ error?: string diff --git a/tests/test_tui_gateway_queue_on_busy.py b/tests/test_tui_gateway_queue_on_busy.py index ea13596ef691..1fffb7db3f68 100644 --- a/tests/test_tui_gateway_queue_on_busy.py +++ b/tests/test_tui_gateway_queue_on_busy.py @@ -66,7 +66,9 @@ def test_busy_interrupt_mode_redirects_active_turn(monkeypatch): assert resp["result"]["status"] == "redirected" assert seen == ["redirect"] - assert session["inflight_turn"]["user"] == "redirect" + # Appended, not overwritten: the original prompt must stay recoverable. + assert session["inflight_turn"]["user"] == "original request" + assert session["inflight_turn"]["corrections"] == ["redirect"] assert session.get("queued_prompt") is None diff --git a/tests/test_tui_gateway_server.py b/tests/test_tui_gateway_server.py index 64ec7f75b2f7..6d3bae7b612b 100644 --- a/tests/test_tui_gateway_server.py +++ b/tests/test_tui_gateway_server.py @@ -7531,11 +7531,56 @@ def test_session_redirect_calls_capable_core_agent(monkeypatch): "text": "use Postgres", } assert calls == ["use Postgres"] - assert session["inflight_turn"]["user"] == "use Postgres" + # The correction is recorded alongside the prompt that started the turn, + # never over it — resume must be able to rebuild both bubbles. + assert session["inflight_turn"]["user"] == "original request" + assert session["inflight_turn"]["corrections"] == ["use Postgres"] assert session.get("last_active") is not None assert before is None or session["last_active"] >= before +def test_session_redirect_records_correction_without_erasing_prompt(): + """A redirect must not overwrite the turn's original user text. + + The inflight snapshot is the only thing session.resume can replay, so + overwriting ``user`` erased the prompt that started the turn and the + client repainted the thread with the user's message missing. + """ + session = {} + server._start_inflight_turn(session, "remove the session counts") + server._append_inflight_delta(session, "Moving.") + server._record_inflight_correction(session, "hurry up") + server._record_inflight_correction(session, "and the worktree ones") + + snapshot = server._inflight_snapshot(session) + assert snapshot is not None + + assert snapshot["user"] == "remove the session counts" + assert snapshot["corrections"] == ["hurry up", "and the worktree ones"] + + +def test_inflight_snapshot_omits_corrections_when_none_recorded(): + session = {} + server._start_inflight_turn(session, "just the prompt") + + snapshot = server._inflight_snapshot(session) + assert snapshot is not None + assert "corrections" not in snapshot + + +def test_new_turn_does_not_inherit_prior_turn_corrections(): + session = {} + server._start_inflight_turn(session, "first prompt") + server._record_inflight_correction(session, "first correction") + server._start_inflight_turn(session, "second prompt") + + snapshot = server._inflight_snapshot(session) + assert snapshot is not None + + assert snapshot["user"] == "second prompt" + assert "corrections" not in snapshot + + def test_session_redirect_queues_during_agent_build_window(monkeypatch): # A fresh turn flips running=True and builds the agent asynchronously, so # session["agent"] is briefly None. A correction landing here must queue diff --git a/tui_gateway/server.py b/tui_gateway/server.py index 39253e14561b..d32cbb4fac18 100644 --- a/tui_gateway/server.py +++ b/tui_gateway/server.py @@ -6387,16 +6387,25 @@ def _append_inflight_delta(session: dict, delta: Any) -> None: session["inflight_turn"] = turn -def _replace_inflight_user(session: dict, text: Any) -> None: - """Reflect an accepted correction as the live turn's current user text.""" - user = _inflight_text(text) - if not user: +def _record_inflight_correction(session: dict, text: Any) -> None: + """Record an accepted mid-turn correction on the live turn. + + The correction is appended, never written over ``user``: a resuming client + must be able to rebuild BOTH bubbles. Overwriting the slot erased the + prompt that started the turn from the only snapshot resume can read, so a + reconnect (or a dev hot-reload that wipes the renderer cache) repainted the + thread with the user's original message missing. + """ + correction = _inflight_text(text) + if not correction: return turn = session.get("inflight_turn") if not isinstance(turn, dict): return turn = dict(turn) - turn["user"] = user + corrections = list(turn.get("corrections") or []) + corrections.append(correction) + turn["corrections"] = corrections turn["updated_at"] = time.time() session["inflight_turn"] = turn @@ -6689,7 +6698,7 @@ def _handle_busy_submit( try: if agent.redirect(plain_text): with session["history_lock"]: - _replace_inflight_user(session, plain_text) + _record_inflight_correction(session, plain_text) session["last_active"] = time.time() return _ok(rid, {"status": "redirected"}) except Exception: @@ -6760,6 +6769,11 @@ def _inflight_snapshot(session: dict) -> dict | None: "streaming": streaming, "user": user, } + corrections = [c for c in (turn.get("corrections") or []) if str(c).strip()] + if corrections: + # Mid-turn redirects. Carried alongside the original prompt (not over + # it) so resume can rebuild every user bubble the turn produced. + snapshot["corrections"] = [str(c) for c in corrections] if error: # Retained failed turn (see _fail_inflight_turn): carry the error # semantics so a resuming client can rebuild the failed-turn bubble @@ -10753,7 +10767,7 @@ def _(rid, params: dict) -> dict: return _err(rid, 5000, f"redirect failed: {exc}") if accepted: with session["history_lock"]: - _replace_inflight_user(session, text) + _record_inflight_correction(session, text) session["last_active"] = time.time() return _ok( rid,