diff --git a/apps/desktop/e2e/correction-session-switch.spec.ts b/apps/desktop/e2e/correction-session-switch.spec.ts index 3ff5ec99d58e..9dc71dfa7912 100644 --- a/apps/desktop/e2e/correction-session-switch.spec.ts +++ b/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 { const composer = page.locator('[contenteditable="true"]').first() @@ -83,6 +86,13 @@ async function reopenOriginalSession(page: Page): Promise { await openSidebarSession(page, ORIGINAL_PROMPT, ORIGINAL_PROMPT) } +async function reopenInferenceSession(page: Page): Promise { + 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) + }) }) \ No newline at end of file diff --git a/tests/test_tui_gateway_queue_on_busy.py b/tests/test_tui_gateway_queue_on_busy.py index 0543d7e8b6b2..ea13596ef691 100644 --- a/tests/test_tui_gateway_queue_on_busy.py +++ b/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 diff --git a/tests/test_tui_gateway_server.py b/tests/test_tui_gateway_server.py index baeb4101c9c4..4387eab7d5cf 100644 --- a/tests/test_tui_gateway_server.py +++ b/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 diff --git a/tui_gateway/server.py b/tui_gateway/server.py index 84535d8c0471..e2057ff5db3b 100644 --- a/tui_gateway/server.py +++ b/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},