From 2b27c171cab8592f648813269a72fb46b5ba8a69 Mon Sep 17 00:00:00 2001 From: Brooklyn Nicholson Date: Wed, 22 Jul 2026 12:48:14 -0500 Subject: [PATCH] fix(redirect): cover the build-window and post-reconnect correction races MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two narrow timing windows (reported by null-runner) silently downgraded a mid-turn correction to a plain next-turn message on the desktop client: - Turn-build window: a fresh turn flips running=True and builds the agent asynchronously, so session["agent"] is briefly None. session.redirect answered 4010 "unsupported", which the renderer's catch swallowed into a lost follow-up. Queue the correction server-side instead and return status="queued" — lossless, and honest about what happened. - Stale runtime id after reconnect: session.redirect 404s on a sid the gateway no longer maps. redirectPrompt now resumes the stored session and retries once, mirroring stopPrompt, so a correction fired right after a reconnect isn't dropped. The desktop treats "queued" like "redirected": the correction reaches the model either way, so it's recorded once as a real user message. --- .../hooks/use-prompt-actions/index.test.tsx | 73 +++++++++++++++++++ .../session/hooks/use-prompt-actions/index.ts | 50 ++++++++++--- apps/desktop/src/app/types.ts | 2 +- tests/test_tui_gateway_server.py | 43 +++++++++++ tui_gateway/server.py | 9 +++ 5 files changed, 164 insertions(+), 13 deletions(-) diff --git a/apps/desktop/src/app/session/hooks/use-prompt-actions/index.test.tsx b/apps/desktop/src/app/session/hooks/use-prompt-actions/index.test.tsx index 6fb71e1b6c3..f83aa123572 100644 --- a/apps/desktop/src/app/session/hooks/use-prompt-actions/index.test.tsx +++ b/apps/desktop/src/app/session/hooks/use-prompt-actions/index.test.tsx @@ -777,6 +777,79 @@ describe('usePromptActions redirectPrompt', () => { expect(await handle!.redirectPrompt(' ')).toBe(false) expect(requestGateway).not.toHaveBeenCalled() }) + + it('accepts a queued redirect during the agent-build window and records the correction', async () => { + // running=True but the agent is still building: the gateway queues the + // correction instead of rejecting, so the composer must NOT re-queue it. + const requestGateway = vi.fn(async () => ({ status: 'queued' }) as never) + + let handle: HarnessHandle | null = null + const capturedStates: Record[] = [] + await actRender( + (handle = h)} + onSeedState={state => capturedStates.push(state)} + refreshSessions={async () => undefined} + requestGateway={requestGateway} + /> + ) + + expect(await handle!.redirectPrompt('build-window nudge')).toBe(true) + expect(requestGateway).toHaveBeenCalledWith('session.redirect', { + session_id: RUNTIME_SESSION_ID, + text: 'build-window nudge' + }) + expect(requestGateway).not.toHaveBeenCalledWith('prompt.submit', expect.anything()) + expect((capturedStates.at(-1)?.messages as unknown[]).at(-1)).toMatchObject({ + role: 'user', + parts: [{ type: 'text', text: 'build-window nudge' }] + }) + }) + + it('resumes the stored session and retries once when session.redirect reports "session not found"', async () => { + const STORED_SESSION_ID = 'stored-db-xyz789' + const RECOVERED_SESSION_ID = 'rt-recovered-456' + const calls: { method: string; params?: Record }[] = [] + let redirectAttempts = 0 + + const requestGateway = vi.fn(async (method: string, params?: Record) => { + calls.push({ method, params }) + + if (method === 'session.redirect') { + redirectAttempts += 1 + + if (redirectAttempts === 1) { + throw new Error('session not found') + } + + return { status: 'redirected' } as never + } + + if (method === 'session.resume') { + return { session_id: RECOVERED_SESSION_ID } as never + } + + return {} as never + }) + + let handle: HarnessHandle | null = null + await actRender( + (handle = h)} + refreshSessions={async () => undefined} + requestGateway={requestGateway} + storedSessionId={STORED_SESSION_ID} + /> + ) + await waitFor(() => expect(handle).not.toBeNull()) + + expect(await handle!.redirectPrompt('reconnect nudge')).toBe(true) + expect(calls.map(c => c.method)).toEqual(['session.redirect', 'session.resume', 'session.redirect']) + expect(calls[0]?.params).toEqual({ session_id: RUNTIME_SESSION_ID, text: 'reconnect nudge' }) + expect(calls[1]?.params).toEqual({ session_id: STORED_SESSION_ID, source: 'desktop' }) + expect(calls[2]?.params).toEqual({ session_id: RECOVERED_SESSION_ID, text: 'reconnect nudge' }) + expect(handle!.activeSessionIdRef.current).toBe(RECOVERED_SESSION_ID) + }) }) describe('usePromptActions restoreToMessage', () => { 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 852e822b189..e7d2c3c36f2 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 @@ -612,28 +612,54 @@ export function usePromptActions({ return false } - try { - const result = await requestGateway('session.redirect', { - session_id: sessionId, - text - }) + // Accepted whether the live turn was redirected in place or queued for + // the next turn (the build window, before the agent is wired) — either + // way the correction reaches the model, so record it once as a real user + // message after the interrupted checkpoint, matching the durable core + // transcript rather than a system note that changes role after reload. + const send = async (id: string): Promise => { + const result = await requestGateway('session.redirect', { session_id: id, text }) - if (result?.status === 'redirected') { + if (result?.status === 'redirected' || result?.status === 'queued') { triggerHaptic('submit') - // Match the durable core transcript: the correction is a real user - // message after the interrupted assistant checkpoint, not a system - // note that changes role after reload. - appendSessionTextMessage(sessionId, 'user', text) + appendSessionTextMessage(id, 'user', text) return true } - } catch { + + return false + } + + try { + return await send(sessionId) + } catch (err) { + // A stale runtime id after reconnect 404s ("session not found"): resume + // the stored session and retry once, mirroring stopPrompt so a + // correction right after a reconnect isn't lost to the race. + if (isSessionNotFoundError(err) && selectedStoredSessionIdRef.current) { + try { + const resumed = await requestGateway<{ session_id: string }>('session.resume', { + session_id: selectedStoredSessionIdRef.current, + source: 'desktop' + }) + + const recoveredId = resumed?.session_id + + if (recoveredId) { + activeSessionIdRef.current = recoveredId + + return await send(recoveredId) + } + } catch { + // fall through — caller queues so nothing is lost + } + } // Swallow — caller queues the text so nothing is lost. } return false }, - [activeSessionId, activeSessionIdRef, appendSessionTextMessage, requestGateway] + [activeSessionId, activeSessionIdRef, appendSessionTextMessage, requestGateway, selectedStoredSessionIdRef] ) const reloadFromMessage = useCallback( diff --git a/apps/desktop/src/app/types.ts b/apps/desktop/src/app/types.ts index a9d09165c76..e12e92286d0 100644 --- a/apps/desktop/src/app/types.ts +++ b/apps/desktop/src/app/types.ts @@ -61,7 +61,7 @@ export interface SessionSteerResponse { } export interface SessionRedirectResponse { - status?: 'redirected' | 'rejected' + status?: 'redirected' | 'queued' | 'rejected' text?: string } diff --git a/tests/test_tui_gateway_server.py b/tests/test_tui_gateway_server.py index 48fa63c7b82..9a197f13acd 100644 --- a/tests/test_tui_gateway_server.py +++ b/tests/test_tui_gateway_server.py @@ -6864,6 +6864,49 @@ def test_session_redirect_calls_capable_core_agent(monkeypatch): assert before is None or session["last_active"] >= before +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 + # (lossless, reaches the model next turn), not hard-reject as unsupported. + session = _session(running=True) + session["agent"] = None + server._sessions["sid"] = session + try: + resp = server.handle_request( + { + "id": "1", + "method": "session.redirect", + "params": {"session_id": "sid", "text": "wait, use SQLite"}, + } + ) + finally: + server._sessions.pop("sid", None) + + assert resp["result"] == {"status": "queued", "text": "wait, use SQLite"} + assert session["queued_prompt"]["text"] == "wait, use SQLite" + + +def test_session_redirect_rejects_when_idle_without_agent(monkeypatch): + # No live turn and no agent: nothing to redirect, and we must not queue a + # phantom turn — keep the explicit unsupported rejection. + session = _session(running=False) + session["agent"] = None + server._sessions["sid"] = session + try: + resp = server.handle_request( + { + "id": "1", + "method": "session.redirect", + "params": {"session_id": "sid", "text": "hi"}, + } + ) + finally: + server._sessions.pop("sid", None) + + assert resp["error"]["code"] == 4010 + assert session.get("queued_prompt") is None + + def test_session_info_includes_mcp_servers(monkeypatch): fake_status = [ {"name": "github", "transport": "http", "tools": 12, "connected": True}, diff --git a/tui_gateway/server.py b/tui_gateway/server.py index 15ed6b4a7ad..c364e8fdc23 100644 --- a/tui_gateway/server.py +++ b/tui_gateway/server.py @@ -9645,6 +9645,15 @@ def _(rid, params: dict) -> dict: if err: return err agent = session.get("agent") + # Turn-build window: a fresh turn flips running=True and kicks off an async + # agent build, so session["agent"] is briefly None. That is not an + # unsupported runtime — queue the correction server-side so it reaches the + # model as the next turn, instead of a misleading 4010 the client silently + # swallows into a lost follow-up. + if agent is None and session.get("running"): + _enqueue_prompt(session, text, current_transport() or _stdio_transport) + session["last_active"] = time.time() + return _ok(rid, {"status": "queued", "text": text}) if ( agent is None or getattr(agent, "_supports_active_turn_redirect", False) is not True