Merge pull request #72339 from NousResearch/bb/redirect-user-row

Preserve the original prompt when a mid-turn redirect corrects a turn

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

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<ChatMessage>()
+
+  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)) {

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'])
+  })
+})

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
     }
   }

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

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
 
 

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

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,