fix(desktop): preserve active correction on warm resume (#69725)

Update the gateway's live turn projection after an accepted correction so a warm session resume does not add the stale original prompt beside the persisted correction.

Cover the inference-time correction path end to end and assert both redirect entry points refresh the live user text.

apps/desktop/e2e/correction-session-switch.spec.ts

@@ -17,6 +17,9 @@ const ORIGINAL_PROMPT = `${CORRECTION_SWITCH_TRIGGER}: original prompt must rema
 const CORRECTION = 'E2E correction must stay after the original prompt.'
 const TOOL_STARTED = 'Checking the long-running task before I continue.'
 const CORRECTED_REPLY = 'The corrected task finished.'
+const INFERENCE_SWITCH_TRIGGER = 'E2E_INFERENCE_SWITCH_TRIGGER'
+const INFERENCE_PROMPT = `${INFERENCE_SWITCH_TRIGGER}: original inference prompt must remain singular.`
+const INFERENCE_CORRECTION = `${INFERENCE_SWITCH_TRIGGER}: correction sent while inference is live.`
 
 async function send(page: Page, text: string): Promise<void> {
   const composer = page.locator('[contenteditable="true"]').first()
@@ -83,6 +86,13 @@ async function reopenOriginalSession(page: Page): Promise<void> {
   await openSidebarSession(page, ORIGINAL_PROMPT, ORIGINAL_PROMPT)
 }
 
+async function reopenInferenceSession(page: Page): Promise<void> {
+  const row = page.locator('[data-slot="sidebar"] button').filter({ hasText: INFERENCE_PROMPT }).first()
+  await row.waitFor({ state: 'visible', timeout: 30_000 })
+  await row.click()
+  await waitForTranscriptText(page, INFERENCE_PROMPT)
+}
+
 function relevantOrder(messages: string[]): string[] {
   return messages.filter(message => message.includes(ORIGINAL_PROMPT) || message.includes(CORRECTION))
 }
@@ -91,7 +101,9 @@ test.describe('correction session switch', () => {
   let fixture: MockBackendFixture | null = null
 
   test.beforeEach(async () => {
-    fixture = await setupMockBackend()
+    fixture = await setupMockBackend({
+      mockServer: { holdFirstStreamForPrompt: INFERENCE_SWITCH_TRIGGER },
+    })
     await waitForAppReady(fixture, 120_000)
   })
 
@@ -137,4 +149,29 @@ test.describe('correction session switch', () => {
 
     await waitForTranscriptText(page, CORRECTED_REPLY)
   })
+
+  test('keeps an inference-time correction visible through a warm session switch', async ({}, testInfo: TestInfo) => {
+    const { mock, page } = fixture!
+
+    await send(page, OTHER_SESSION_PROMPT)
+    await waitForTranscriptText(page, MOCK_REPLY)
+    await openFreshDraft(page, OTHER_SESSION_PROMPT)
+
+    await send(page, INFERENCE_PROMPT)
+    await mock.waitForHeldStream()
+    await waitForTranscriptText(page, INFERENCE_PROMPT)
+
+    await send(page, INFERENCE_CORRECTION)
+    await waitForTranscriptText(page, INFERENCE_CORRECTION)
+
+    await openSidebarSession(page, MOCK_REPLY, OTHER_SESSION_PROMPT)
+    await reopenInferenceSession(page)
+
+    expect(await textNodeOccurrences(page, INFERENCE_PROMPT)).toBe(1)
+    expect(await textNodeOccurrences(page, INFERENCE_CORRECTION)).toBe(1)
+    await page.screenshot({ path: testInfo.outputPath('inference-correction-after-warm-resume.png') })
+
+    mock.releaseHeldStream()
+    await waitForTranscriptText(page, MOCK_REPLY)
+  })
 })

tests/test_tui_gateway_queue_on_busy.py

@@ -60,11 +60,13 @@ def test_busy_interrupt_mode_redirects_active_turn(monkeypatch):
         ),
     )
     session = _session(agent=agent, running=True)
+    session["inflight_turn"] = {"user": "original request", "assistant": "partial reply"}
 
     resp = server._handle_busy_submit("r1", "sid", session, "redirect", "ws-1")
 
     assert resp["result"]["status"] == "redirected"
     assert seen == ["redirect"]
+    assert session["inflight_turn"]["user"] == "redirect"
     assert session.get("queued_prompt") is None
 
 

tests/test_tui_gateway_server.py

@@ -7051,6 +7051,7 @@ def test_session_redirect_calls_capable_core_agent(monkeypatch):
         redirect=lambda text: calls.append(text) or True,
     )
     session = _session(agent=agent)
+    session["inflight_turn"] = {"user": "original request", "assistant": "partial reply"}
     server._sessions["sid"] = session
     try:
         before = session.get("last_active")
@@ -7069,6 +7070,7 @@ def test_session_redirect_calls_capable_core_agent(monkeypatch):
         "text": "use Postgres",
     }
     assert calls == ["use Postgres"]
+    assert session["inflight_turn"]["user"] == "use Postgres"
     assert session.get("last_active") is not None
     assert before is None or session["last_active"] >= before
 

tui_gateway/server.py

@@ -5773,6 +5773,20 @@ 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:
+        return
+    turn = session.get("inflight_turn")
+    if not isinstance(turn, dict):
+        return
+    turn = dict(turn)
+    turn["user"] = user
+    turn["updated_at"] = time.time()
+    session["inflight_turn"] = turn
+
+
 def _clear_inflight_turn(session: dict) -> None:
     session["inflight_turn"] = None
 
@@ -5878,6 +5892,7 @@ def _handle_busy_submit(
         try:
             if agent.redirect(plain_text):
                 with session["history_lock"]:
+                    _replace_inflight_user(session, plain_text)
                     session["last_active"] = time.time()
                 return _ok(rid, {"status": "redirected"})
         except Exception:
@@ -9761,7 +9776,9 @@ def _(rid, params: dict) -> dict:
     except Exception as exc:
         return _err(rid, 5000, f"redirect failed: {exc}")
     if accepted:
-        session["last_active"] = time.time()
+        with session["history_lock"]:
+            _replace_inflight_user(session, text)
+            session["last_active"] = time.time()
     return _ok(
         rid,
         {"status": "redirected" if accepted else "rejected", "text": text},