From 2637aa607f2017ec5f638ce1843b946312ccbd48 Mon Sep 17 00:00:00 2001 From: Austin Pickett Date: Sat, 18 Jul 2026 16:12:59 -0400 Subject: [PATCH] fix(desktop): preserve in-flight turns across gateway reconnects (rebase of #66234) (#67114) * fix(desktop): preserve turns across gateway reconnect * chore(release): map UnathiCodex attribution * fix(desktop): prefer rotated resume projection * fix(desktop): refresh warm session transcripts * chore(contributors): use email mapping file * fix(desktop): restore live prompts after restart --------- Co-authored-by: UnathiCodex --- .../hooks/use-session-actions.test.tsx | 311 +++++++++++++++++- .../hooks/use-session-actions/index.ts | 193 +++++++++-- .../hooks/use-session-actions/utils.test.ts | 71 ++++ .../hooks/use-session-actions/utils.ts | 122 ++++++- apps/desktop/src/types/hermes.ts | 12 + contributors/emails/theunathi@gmail.com | 1 + tests/test_tui_gateway_server.py | 40 +++ tui_gateway/server.py | 24 +- 8 files changed, 741 insertions(+), 33 deletions(-) diff --git a/apps/desktop/src/app/session/hooks/use-session-actions.test.tsx b/apps/desktop/src/app/session/hooks/use-session-actions.test.tsx index 1b725f657..2b32437a0 100644 --- a/apps/desktop/src/app/session/hooks/use-session-actions.test.tsx +++ b/apps/desktop/src/app/session/hooks/use-session-actions.test.tsx @@ -212,14 +212,18 @@ describe('createBackendSessionForSend profile routing', () => { // (b) arm $resumeFailedSessionId so use-route-resume can retry. A resume that // succeeds must NOT leave the flag armed. function ResumeHarness({ + onStateUpdate, onReady, requestGateway, runtimeIdByStoredSessionIdRef, + selectedStoredSessionId = null, sessionStateByRuntimeIdRef }: { + onStateUpdate?: (sessionId: string, state: ClientSessionState) => void onReady: (resume: (storedSessionId: string, replaceRoute?: boolean) => Promise) => void requestGateway: (method: string, params?: Record) => Promise runtimeIdByStoredSessionIdRef?: MutableRefObject> + selectedStoredSessionId?: string | null sessionStateByRuntimeIdRef?: MutableRefObject> }) { const ref = (value: T): MutableRefObject => ({ current: value }) @@ -235,11 +239,16 @@ function ResumeHarness({ requestGateway, resetViewSync: vi.fn(), runtimeIdByStoredSessionIdRef: runtimeIdByStoredSessionIdRef ?? ref(new Map()), - selectedStoredSessionId: null, - selectedStoredSessionIdRef: ref(null), + selectedStoredSessionId, + selectedStoredSessionIdRef: ref(selectedStoredSessionId), sessionStateByRuntimeIdRef: sessionStateByRuntimeIdRef ?? ref(new Map()), syncSessionStateToView: vi.fn(), - updateSessionState: (_sessionId, updater) => updater({} as ClientSessionState) + updateSessionState: (sessionId, updater) => { + const next = updater({} as ClientSessionState) + onStateUpdate?.(sessionId, next) + + return next + } }) useEffect(() => { @@ -321,6 +330,155 @@ describe('resumeSession failure recovery', () => { expect($messages.get().length).toBeGreaterThan(0) }) + it('preserves an optimistic user message during a same-session reconnect', async () => { + setMessages([ + { + id: 'stored-user', + role: 'user', + parts: [{ type: 'text', text: 'earlier question' }] + }, + { + id: 'stored-assistant', + role: 'assistant', + parts: [{ type: 'text', text: 'earlier answer' }] + }, + { + id: 'user-optimistic', + role: 'user', + parts: [{ type: 'text', text: 'message sent during reconnect' }] + } + ]) + + const storedMessages = [ + { content: 'earlier question', role: 'user', timestamp: 1 }, + { content: 'earlier answer', role: 'assistant', timestamp: 2 } + ] + + vi.mocked(getSessionMessages).mockResolvedValue({ messages: storedMessages, session_id: 'stored-1' } as never) + + const requestGateway = vi.fn(async (method: string) => { + if (method === 'session.resume') { + return { + session_id: 'runtime-1', + session_key: 'stored-1', + resumed: 'stored-1', + message_count: 2, + messages: storedMessages, + info: {} + } as never + } + + return {} as never + }) + + let resume: ((storedSessionId: string, replaceRoute?: boolean) => Promise) | null = null + render( + (resume = r)} requestGateway={requestGateway} selectedStoredSessionId="stored-1" /> + ) + await waitFor(() => expect(resume).not.toBeNull()) + await resume!('stored-1', true) + + expect($messages.get().map(message => message.id)).toContain('user-optimistic') + }) + + it('restores the in-flight turn and queued user prompt after a full renderer restart', async () => { + const storedMessages = [ + { content: 'earlier question', role: 'user', timestamp: 1 }, + { content: 'earlier answer', role: 'assistant', timestamp: 2 } + ] + + vi.mocked(getSessionMessages).mockResolvedValue({ messages: storedMessages, session_id: 'stored-1' } as never) + + const requestGateway = vi.fn(async (method: string) => { + if (method === 'session.resume') { + return { + session_id: 'runtime-1', + session_key: 'stored-1', + resumed: 'stored-1', + message_count: storedMessages.length, + messages: storedMessages, + running: true, + inflight: { + user: 'current prompt', + assistant: 'partial answer', + streaming: true + }, + queued: { user: 'newest prompt' }, + info: {} + } as never + } + + return {} as never + }) + + let resumedState: ClientSessionState | undefined + let resume: ((storedSessionId: string, replaceRoute?: boolean) => Promise) | null = null + render( + (resume = ready)} + onStateUpdate={(_sessionId, state) => (resumedState = state)} + requestGateway={requestGateway} + /> + ) + await waitFor(() => expect(resume).not.toBeNull()) + await resume!('stored-1', true) + + const renderedMessages = JSON.stringify(resumedState?.messages) + expect(renderedMessages).toContain('current prompt') + expect(renderedMessages).toContain('partial answer') + expect(renderedMessages).toContain('newest prompt') + }) + + it('uses the continuation projection when resume rotates an equal-length stored transcript', async () => { + const parentMessages = [ + { content: 'question before compression', role: 'user', timestamp: 1 }, + { content: 'answer before compression', role: 'assistant', timestamp: 2 } + ] + + const continuationMessages = [ + { content: 'prompt after compression', role: 'user', timestamp: 3 }, + { content: 'answer after compression', role: 'assistant', timestamp: 4 } + ] + + vi.mocked(getSessionMessages).mockResolvedValue({ + messages: parentMessages, + session_id: 'stored-1' + } as never) + + const requestGateway = vi.fn(async (method: string) => { + if (method === 'session.resume') { + return { + session_id: 'runtime-continuation', + session_key: 'stored-continuation', + resumed: 'stored-continuation', + message_count: continuationMessages.length, + messages: continuationMessages, + info: {} + } as never + } + + return {} as never + }) + + let resumedState: ClientSessionState | undefined + let resume: ((storedSessionId: string, replaceRoute?: boolean) => Promise) | null = null + + render( + (resume = ready)} + onStateUpdate={(_sessionId, state) => (resumedState = state)} + requestGateway={requestGateway} + /> + ) + await waitFor(() => expect(resume).not.toBeNull()) + await resume!('stored-1', true) + + const renderedMessages = JSON.stringify(resumedState?.messages) + expect(renderedMessages).toContain('prompt after compression') + expect(renderedMessages).toContain('answer after compression') + expect(renderedMessages).not.toContain('answer before compression') + }) + it('does NOT throw out of the fallback when REST also fails (no unhandled rejection)', async () => { const requestGateway = vi.fn(async (method: string) => { if (method === 'session.resume') { @@ -601,9 +759,10 @@ describe('resumeSession warm-cache mapping integrity', () => { expect(sessionStateByRuntimeIdRef.current.has('rt-recycled')).toBe(false) }) - it('honours a warm cache entry whose stored id matches (no needless refetch)', async () => { + it('honours a warm cache entry whose stored id matches and refreshes its persisted transcript', async () => { // Correctly-wired mapping: 'rt-A' <-> 'stored-A'. The fast-path should trust - // it and never reach session.resume (only the lightweight usage probe). + // it and never reach session.resume. session.activate refreshes the live + // projection and, critically, rebinds its event transport after reconnect. const runtimeIdByStoredSessionIdRef: MutableRefObject> = { current: new Map([['stored-A', 'rt-A']]) } @@ -613,8 +772,141 @@ describe('resumeSession warm-cache mapping integrity', () => { } const requestGateway = vi.fn(async (method: string) => { - if (method === 'session.usage') { - return { input: 0, output: 0, total: 0 } as never + if (method === 'session.activate') { + return { + session_id: 'rt-A', + session_key: 'stored-A', + resumed: 'stored-A', + message_count: 0, + messages: [], + running: false, + info: {} + } as never + } + + return {} as never + }) + + vi.mocked(getSessionMessages).mockResolvedValue({ messages: [], session_id: 'stored-A' } as never) + + let resume: ((storedSessionId: string, replaceRoute?: boolean) => Promise) | null = null + render( + (resume = r)} + requestGateway={requestGateway} + runtimeIdByStoredSessionIdRef={runtimeIdByStoredSessionIdRef} + sessionStateByRuntimeIdRef={sessionStateByRuntimeIdRef} + /> + ) + await waitFor(() => expect(resume).not.toBeNull()) + await resume!('stored-A', true) + + // Fast-path served the session from cache: no full resume RPC, mapping intact. + // The persisted transcript still refreshes in parallel because the runtime + // projection can differ even when its row count matches. + const methods = requestGateway.mock.calls.map(([method]) => method) + expect(methods).toContain('session.activate') + expect(methods).not.toContain('session.resume') + expect(getSessionMessages).toHaveBeenCalledWith('stored-A', undefined) + expect(runtimeIdByStoredSessionIdRef.current.get('stored-A')).toBe('rt-A') + }) + + it('repairs an idle warm cache from a divergent equal-length persisted transcript', async () => { + const runtimeIdByStoredSessionIdRef: MutableRefObject> = { + current: new Map([['stored-A', 'rt-A']]) + } + + const state = clientState('stored-A') + state.messages = [ + { + id: 'cached-user', + role: 'user', + parts: [{ type: 'text', text: 'stale runtime prompt' }] + }, + { + id: 'cached-assistant', + role: 'assistant', + parts: [{ type: 'text', text: 'stale runtime answer' }] + } + ] + + const sessionStateByRuntimeIdRef: MutableRefObject> = { + current: new Map([['rt-A', state]]) + } + + const staleRuntimeMessages = [ + { content: 'stale runtime prompt', role: 'user', timestamp: 1 }, + { content: 'stale runtime answer', role: 'assistant', timestamp: 2 } + ] + + const persistedMessages = [ + { content: 'prompt saved after compression', role: 'user', timestamp: 3 }, + { content: 'answer saved after compression', role: 'assistant', timestamp: 4 } + ] + + vi.mocked(getSessionMessages).mockResolvedValue({ + messages: persistedMessages, + session_id: 'stored-A' + } as never) + + const requestGateway = vi.fn(async (method: string) => { + if (method === 'session.activate') { + return { + session_id: 'rt-A', + session_key: 'stored-A', + resumed: 'stored-A', + message_count: staleRuntimeMessages.length, + messages: staleRuntimeMessages, + running: false, + info: {} + } as never + } + + return {} as never + }) + + let resumedState: ClientSessionState | undefined + let resume: ((storedSessionId: string, replaceRoute?: boolean) => Promise) | null = null + + render( + (resume = ready)} + onStateUpdate={(_sessionId, next) => (resumedState = next)} + requestGateway={requestGateway} + runtimeIdByStoredSessionIdRef={runtimeIdByStoredSessionIdRef} + sessionStateByRuntimeIdRef={sessionStateByRuntimeIdRef} + /> + ) + await waitFor(() => expect(resume).not.toBeNull()) + await resume!('stored-A', true) + + const renderedMessages = JSON.stringify(resumedState?.messages) + expect(renderedMessages).toContain('prompt saved after compression') + expect(renderedMessages).toContain('answer saved after compression') + expect(renderedMessages).not.toContain('stale runtime answer') + }) + + it('keeps a warm runtime and optimistic turn on a transient activation timeout', async () => { + const runtimeIdByStoredSessionIdRef: MutableRefObject> = { + current: new Map([['stored-A', 'rt-A']]) + } + + const state = clientState('stored-A') + state.messages = [ + { + id: 'user-optimistic', + role: 'user', + parts: [{ type: 'text', text: 'do not lose me' }] + } + ] + + const sessionStateByRuntimeIdRef: MutableRefObject> = { + current: new Map([['rt-A', state]]) + } + + const requestGateway = vi.fn(async (method: string) => { + if (method === 'session.activate') { + throw new Error('request timed out: session.activate') } return {} as never @@ -632,10 +924,9 @@ describe('resumeSession warm-cache mapping integrity', () => { await waitFor(() => expect(resume).not.toBeNull()) await resume!('stored-A', true) - // Fast-path served the session from cache: no full resume RPC, mapping intact. - const methods = requestGateway.mock.calls.map(([method]) => method) - expect(methods).not.toContain('session.resume') + expect(requestGateway.mock.calls.map(([method]) => method)).not.toContain('session.resume') expect(runtimeIdByStoredSessionIdRef.current.get('stored-A')).toBe('rt-A') + expect(sessionStateByRuntimeIdRef.current.get('rt-A')?.messages[0]?.id).toBe('user-optimistic') }) }) diff --git a/apps/desktop/src/app/session/hooks/use-session-actions/index.ts b/apps/desktop/src/app/session/hooks/use-session-actions/index.ts index 1e1e37405..84acfda63 100644 --- a/apps/desktop/src/app/session/hooks/use-session-actions/index.ts +++ b/apps/desktop/src/app/session/hooks/use-session-actions/index.ts @@ -5,7 +5,8 @@ import type { NavigateFunction } from 'react-router-dom' import { revealTreePane } from '@/components/pane-shell/tree/store' import { deleteSession, getSessionMessages, setSessionArchived } from '@/hermes' import { useI18n } from '@/i18n' -import { preserveLocalAssistantErrors, toChatMessages } from '@/lib/chat-messages' +import { type ChatMessage, preserveLocalAssistantErrors, toChatMessages } from '@/lib/chat-messages' +import { isMissingRpcMethod } from '@/lib/gateway-rpc' import { setSessionYolo } from '@/lib/yolo-session' import { clearQueuedPrompts } from '@/store/composer-queue' import { $pinnedSessionIds } from '@/store/layout' @@ -62,12 +63,14 @@ import { NEW_CHAT_ROUTE, sessionRoute, SETTINGS_ROUTE } from '../../../routes' import type { ClientSessionState, SidebarNavItem } from '../../../types' import { + appendLiveSessionProjection, applyRuntimeInfo, applyStoredSessionPreviewRuntimeInfo, type BranchMessage, chatMessageArraysEquivalent, isSessionGoneError, patchSessionWorkspace, + preserveLocalPendingTurnMessages, reconcileResumeMessages, resolveStoredSession, sessionMatchesStoredId, @@ -117,6 +120,19 @@ function applyStoredUsage(stored: { input_tokens?: number | null; output_tokens? setCurrentUsage(current => ({ ...current, input, output, total: input + output })) } +function reconcileAuthoritativeMessages( + authoritativeMessages: SessionResumeResponse['messages'], + previousMessages: ChatMessage[], + liveProjection?: Pick +): ChatMessage[] { + const authoritative = toChatMessages(authoritativeMessages) + const withLiveProjection = liveProjection ? appendLiveSessionProjection(authoritative, liveProjection) : authoritative + const reconciled = reconcileResumeMessages(withLiveProjection, previousMessages) + const withPendingTurn = preserveLocalPendingTurnMessages(reconciled, previousMessages) + + return preserveLocalAssistantErrors(withPendingTurn, previousMessages) +} + // `session.create` params from the current profile + sticky-UI model/effort/fast, // ensuring the gateway is on that profile first. Shared by the primary send path // and the "open in split" tile path; `cwd` is the one thing that differs (the @@ -431,6 +447,8 @@ export function useSessionActions({ async (storedSessionId: string, replaceRoute = false) => { const requestId = resumeRequestRef.current + 1 resumeRequestRef.current = requestId + const resumedSameSelectedSession = selectedStoredSessionIdRef.current === storedSessionId + const resumeStartMessages = resumedSameSelectedSession ? $messages.get() : [] const isCurrentResume = () => resumeRequestRef.current === requestId && selectedStoredSessionIdRef.current === storedSessionId @@ -462,7 +480,7 @@ export function useSessionActions({ // session being resumed. A pooled profile backend that gets idle-reaped // and respawned (pruneSecondaryGateways) re-mints runtime ids, so a // recycled id can resolve to a live-but-DIFFERENT session's cache entry. - // The session.usage 404 guard below only catches a fully-DEAD id — a + // The session.activate 404 guard below only catches a fully-DEAD id — a // recycled-live id 200s, so an unchecked hit paints the wrong transcript // under the current route (the "open chat A, chat B loads" bug). On a // mismatch the mapping is cross-wired: purge both sides and report a miss @@ -489,7 +507,10 @@ export function useSessionActions({ if (!takeWarmCache()) { setActiveSessionId(null) activeSessionIdRef.current = null - setMessages([]) + + if (!resumedSameSelectedSession) { + setMessages([]) + } } // Swap the single live gateway to this session's profile before any @@ -517,7 +538,7 @@ export function useSessionActions({ const stored = $sessions.get().find(session => sessionMatchesStoredId(session, storedSessionId)) ?? storedForProfile - const cachedViewState = + let cachedViewState = !cachedState.model && stored?.model != null ? { ...cachedState, @@ -525,6 +546,14 @@ export function useSessionActions({ } : cachedState + if (resumedSameSelectedSession) { + const messages = preserveLocalPendingTurnMessages(cachedViewState.messages, resumeStartMessages) + + if (messages !== cachedViewState.messages) { + cachedViewState = { ...cachedViewState, messages } + } + } + if (cachedViewState !== cachedState) { sessionStateByRuntimeIdRef.current.set(cachedRuntimeId, cachedViewState) publishSessionState(cachedRuntimeId, cachedViewState) @@ -535,6 +564,17 @@ export function useSessionActions({ sessionStateByRuntimeIdRef.current.delete(cachedRuntimeId) dropSessionState(cachedRuntimeId) } else { + // Paint the warm cache immediately, but also refresh the persisted + // transcript in parallel. A resumed runtime carries the agent's + // compression projection, which can have the same row count as the + // stored conversation while containing different rows. Trusting that + // projection alone made completed prompts disappear after an app + // restart whenever this warm path short-circuited the cold REST + // prefetch. Watch mirrors stay live-only by design. + const persistedTranscriptPromise = isWatchWindow() + ? null + : getSessionMessages(storedSessionId, sessionProfile).catch(() => null) + setFreshDraftReady(false) clearNotifications() setSelectedStoredSessionId(storedSessionId) @@ -547,27 +587,111 @@ export function useSessionActions({ setSessionStartedAt(Date.now()) try { - const usage = await requestGateway('session.usage', { session_id: cachedRuntimeId }) + let activated: SessionResumeResponse | null = null + + try { + activated = await requestGateway('session.activate', { + session_id: cachedRuntimeId, + cols: 96 + }) + } catch (error) { + // Compatibility for older backends. Modern backends require + // session.activate here because it rebinds the live session's + // event transport to this newly-opened WebSocket. + if (!isMissingRpcMethod(error)) { + throw error + } + + const usage = await requestGateway('session.usage', { session_id: cachedRuntimeId }) + + if (!isCurrentResume()) { + return + } + + if (usage) { + setCurrentUsage(current => ({ ...current, ...usage })) + } + + return + } if (!isCurrentResume()) { return } - if (usage) { - setCurrentUsage(current => ({ ...current, ...usage })) - } + if (activated.session_key && activated.session_key !== storedSessionId) { + runtimeIdByStoredSessionIdRef.current.delete(storedSessionId) + sessionStateByRuntimeIdRef.current.delete(cachedRuntimeId) + dropSessionState(cachedRuntimeId) + } else { + const runtimeInfo = applyRuntimeInfo(activated.info) - return - } catch { + let activatedMessages = + activated.messages.length || activated.inflight || activated.queued + ? reconcileAuthoritativeMessages(activated.messages, cachedViewState.messages, activated) + : cachedViewState.messages + + const running = Boolean(activated.running ?? cachedViewState.busy) + + // While idle, the persisted REST transcript is the display + // authority: session.activate returns the runtime's compressed + // context projection, not necessarily the complete conversation. + // During a live turn, keep the runtime/cache projection so an + // accepted but not-yet-persisted prompt or stream is never lost. + if (!running && persistedTranscriptPromise) { + const persisted = await persistedTranscriptPromise + + if (!isCurrentResume()) { + return + } + + const activatedStoredSessionId = activated.session_key || activated.resumed + + const persistedMatchesActivatedSession = + !persisted?.session_id || + !activatedStoredSessionId || + persisted.session_id === activatedStoredSessionId + + if (persisted && persistedMatchesActivatedSession) { + activatedMessages = reconcileAuthoritativeMessages(persisted.messages, activatedMessages) + } + } + + const activatedState = updateSessionState( + cachedRuntimeId, + state => ({ + ...state, + ...(runtimeInfo ?? {}), + messages: activatedMessages, + busy: running, + awaitingResponse: running + }), + storedSessionId + ) + + busyRef.current = running + setBusy(running) + setAwaitingResponse(running) + syncSessionStateToView(cachedRuntimeId, activatedState) + + return + } + } catch (error) { // The cached runtime id was minted by a prior backend instance. A // pooled profile backend that gets idle-reaped (pruneSecondaryGateways) // and respawned across a profile swap mints fresh ids, so this mapping // now 404s ("session not found"). Drop it and fall through to a full - // resume that rebinds a live runtime id. + // resume that rebinds a live runtime id. A transient timeout or + // transport error is NOT proof that the session is dead: keep the + // cache and optimistic turn intact for the next reconnect attempt. if (!isCurrentResume()) { return } + if (!isSessionGoneError(error)) { + return + } + runtimeIdByStoredSessionIdRef.current.delete(storedSessionId) sessionStateByRuntimeIdRef.current.delete(cachedRuntimeId) dropSessionState(cachedRuntimeId) @@ -585,7 +709,7 @@ export function useSessionActions({ // session's transcript would leak into this cold resume ("switching // sessions shows the same messages"). Clear it so the loader/prefetch // paints fresh; guarded so the normal cold path (already cleared) no-ops. - if ($messages.get().length > 0) { + if (!resumedSameSelectedSession && $messages.get().length > 0) { setMessages([]) } @@ -610,7 +734,14 @@ export function useSessionActions({ try { const watchWindow = isWatchWindow() - let localSnapshot = $messages.get() + + let localSnapshot = resumedSameSelectedSession + ? preserveLocalPendingTurnMessages($messages.get(), resumeStartMessages) + : $messages.get() + + let prefetchApplied = false + let prefetchedMessageCount = 0 + let prefetchedStoredSessionId: string | null = null // REST transcript prefetch and the gateway resume RPC are independent // — run them concurrently so a big session's wall time is @@ -641,7 +772,14 @@ export function useSessionActions({ const storedMessages = await prefetchPromise if (isCurrentResume()) { - localSnapshot = preserveLocalAssistantErrors(toChatMessages(storedMessages.messages), $messages.get()) + const previousMessages = resumedSameSelectedSession + ? preserveLocalPendingTurnMessages($messages.get(), resumeStartMessages) + : $messages.get() + + localSnapshot = reconcileAuthoritativeMessages(storedMessages.messages, previousMessages) + prefetchApplied = true + prefetchedMessageCount = storedMessages.messages.length + prefetchedStoredSessionId = storedMessages.session_id || storedSessionId if (!chatMessageArraysEquivalent($messages.get(), localSnapshot)) { setMessages(localSnapshot) @@ -665,14 +803,25 @@ export function useSessionActions({ // skip converting/reconciling the resume payload entirely — on a // 1000+-message session that second conversion plus the deep // equivalence compare costs over a second of main-thread time. + const resumedStoredSessionId = resumed.session_key || resumed.resumed + + const prefetchMatchesResumedSession = + !prefetchedStoredSessionId || !resumedStoredSessionId || prefetchedStoredSessionId === resumedStoredSessionId + + const hasLiveProjection = Boolean(resumed.inflight || resumed.queued) + const preferredMessages = - localSnapshot.length > 0 + prefetchApplied && + prefetchMatchesResumedSession && + !hasLiveProjection && + resumed.messages.length <= prefetchedMessageCount ? localSnapshot : (() => { - const resumedMessages = preserveLocalAssistantErrors( - reconcileResumeMessages(toChatMessages(resumed.messages), currentMessages), - currentMessages - ) + const previousMessages = resumedSameSelectedSession + ? preserveLocalPendingTurnMessages(currentMessages, resumeStartMessages) + : currentMessages + + const resumedMessages = reconcileAuthoritativeMessages(resumed.messages, previousMessages, resumed) return chatMessageArraysEquivalent(currentMessages, resumedMessages) ? currentMessages : resumedMessages })() @@ -735,7 +884,11 @@ export function useSessionActions({ return } - setMessages(preserveLocalAssistantErrors(toChatMessages(fallback.messages), $messages.get())) + const previousMessages = resumedSameSelectedSession + ? preserveLocalPendingTurnMessages($messages.get(), resumeStartMessages) + : $messages.get() + + setMessages(reconcileAuthoritativeMessages(fallback.messages, previousMessages)) } catch (e) { // Fallback also failed: nothing to paint. Leave whatever messages are // already shown and fall through to arm the resume-failure latch so 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 1a5390210..84e0d232f 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 @@ -6,11 +6,13 @@ import { $activeGatewayProfile } from '@/store/profile' import type { SessionInfo } from '@/types/hermes' import { + appendLiveSessionProjection, applyRuntimeInfo, chatMessageArraysEquivalent, chatMessagesEquivalent, chatPartsEquivalent, isSessionGoneError, + preserveLocalPendingTurnMessages, reconcileResumeMessages, sessionMatchesStoredId, sessionShouldHaveTranscript, @@ -289,3 +291,72 @@ describe('reconcileResumeMessages', () => { expect(out.parts.some(p => p.type === 'reasoning')).toBe(true) }) }) + +describe('preserveLocalPendingTurnMessages', () => { + it('keeps an optimistic user turn and pending assistant when the server projection is behind', () => { + const next = [msg('1-user', 'user', 'first'), msg('2-assistant', 'assistant', 'first answer')] + + const previous = [ + ...next, + msg('user-optimistic', 'user', 'new question'), + msg('assistant-stream-1', 'assistant', 'partial answer', { pending: true }) + ] + + expect(preserveLocalPendingTurnMessages(next, previous).map(message => message.id)).toEqual([ + '1-user', + '2-assistant', + 'user-optimistic', + 'assistant-stream-1' + ]) + }) + + it('drops the local copies once the same role ordinals are authoritative', () => { + const previous = [ + msg('1-user', 'user', 'first'), + msg('2-assistant', 'assistant', 'first answer'), + msg('user-optimistic', 'user', 'new question'), + msg('assistant-stream-1', 'assistant', 'partial answer', { pending: true }) + ] + + const next = [ + msg('1-user-stored', 'user', 'first'), + msg('2-assistant-stored', 'assistant', 'first answer'), + msg('3-user-stored', 'user', 'new question'), + msg('4-assistant-stored', 'assistant', 'complete answer') + ] + + expect(preserveLocalPendingTurnMessages(next, previous)).toBe(next) + }) +}) + +describe('appendLiveSessionProjection', () => { + it('restores the running turn and accepted queued prompt after a renderer restart', () => { + const stored = [msg('stored-user', 'user', 'earlier'), msg('stored-assistant', 'assistant', 'earlier answer')] + + const restored = appendLiveSessionProjection(stored, { + session_id: 'runtime-1', + inflight: { + user: 'current prompt', + assistant: 'partial answer', + streaming: true + }, + queued: { user: 'newest prompt' } + }) + + expect(restored.map(message => message.role)).toEqual(['user', 'assistant', 'user', 'assistant', 'user']) + expect(restored.map(message => message.parts.map(part => ('text' in part ? part.text : '')).join(''))).toEqual([ + 'earlier', + 'earlier answer', + 'current prompt', + 'partial answer', + 'newest prompt' + ]) + expect(restored[3]).toMatchObject({ id: 'assistant-stream-runtime-1', pending: true }) + }) + + it('preserves the original array when no live projection exists', () => { + const stored = [msg('stored-user', 'user', 'earlier')] + + expect(appendLiveSessionProjection(stored, { session_id: 'runtime-1' })).toBe(stored) + }) +}) 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 e7a7ec1ad..927b72d6d 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 @@ -1,5 +1,5 @@ import { getSession } from '@/hermes' -import { type ChatMessage, chatMessageText } from '@/lib/chat-messages' +import { assistantTextPart, type ChatMessage, chatMessageText, textPart } from '@/lib/chat-messages' import { normalizePersonalityValue } from '@/lib/chat-runtime' import { embeddedImageUrls, textWithoutEmbeddedImages } from '@/lib/embedded-images' import { reconcileApprovalModeForProfile } from '@/store/approval-mode' @@ -26,7 +26,7 @@ import { // it from here; the canonical definition lives in @/store/session. export { sessionMatchesStoredId } import { reportBackendContract, reportInstallMethodWarning } from '@/store/updates' -import type { SessionCreateResponse, SessionInfo, SessionRuntimeInfo } from '@/types/hermes' +import type { SessionCreateResponse, SessionInfo, SessionResumeResponse, SessionRuntimeInfo } from '@/types/hermes' import type { ClientSessionState } from '../../../types' @@ -224,6 +224,124 @@ export function reconcileResumeMessages(nextMessages: ChatMessage[], previousMes }) } +/** + * Keep the local tail of a turn while a reconnect hydrates an older server + * projection. The user's optimistic row exists before prompt.submit persists + * it, and the pending assistant row exists before message.complete commits it; + * dropping either makes an accepted turn appear to vanish during transport + * churn. + * + * Authoritative rows use different ids, so match by role ordinal. A matching + * user row is considered committed only when its visible text also matches; + * any authoritative assistant at the same ordinal supersedes the local stream. + */ +export function preserveLocalPendingTurnMessages( + nextMessages: ChatMessage[], + previousMessages: ChatMessage[] +): ChatMessage[] { + if (!previousMessages.length) { + return nextMessages + } + + const nextByRoleOrdinal = new Map() + const nextRoleCounts = new Map() + + for (const message of nextMessages) { + const ordinal = nextRoleCounts.get(message.role) ?? 0 + nextRoleCounts.set(message.role, ordinal + 1) + nextByRoleOrdinal.set(`${message.role}:${ordinal}`, message) + } + + const nextIds = new Set(nextMessages.map(message => message.id)) + const previousRoleCounts = new Map() + const preserved: ChatMessage[] = [] + + for (const message of previousMessages) { + const ordinal = previousRoleCounts.get(message.role) ?? 0 + previousRoleCounts.set(message.role, ordinal + 1) + + const isOptimisticUser = message.role === 'user' && message.id.startsWith('user-') + + const isPendingAssistant = + message.role === 'assistant' && (message.pending === true || message.id.startsWith('assistant-stream-')) + + if ((!isOptimisticUser && !isPendingAssistant) || nextIds.has(message.id)) { + continue + } + + const authoritative = nextByRoleOrdinal.get(`${message.role}:${ordinal}`) + + if (authoritative) { + if (isPendingAssistant) { + continue + } + + if (chatMessageText(authoritative).trim() === chatMessageText(message).trim()) { + continue + } + } + + preserved.push(message) + } + + return preserved.length ? [...nextMessages, ...preserved] : nextMessages +} + +/** + * Append the backend-only tail of a live turn to a stored transcript. + * + * Session history is committed only when a turn finishes. During a reconnect, + * `inflight` is therefore the authority for the currently running user/assistant + * pair, while `queued` is an accepted next-turn prompt waiting in gateway + * memory. Stable ids let repeated activate/resume hydration reconcile instead + * of growing duplicate rows. + */ +export function appendLiveSessionProjection( + messages: ChatMessage[], + projection: Pick +): ChatMessage[] { + const inflightUser = projection.inflight?.user?.trim() ?? '' + const inflightAssistant = projection.inflight?.assistant ?? '' + const inflightStreaming = Boolean(projection.inflight?.streaming) + const queuedUser = projection.queued?.user?.trim() ?? '' + + if (!inflightUser && !inflightAssistant && !inflightStreaming && !queuedUser) { + return messages + } + + const sessionId = projection.session_id || 'session' + const projected: ChatMessage[] = [] + + if (inflightUser) { + projected.push({ + id: `user-inflight-${sessionId}`, + role: 'user', + parts: [textPart(inflightUser)] + }) + } + + // 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 || (inflightUser && queuedUser)) { + projected.push({ + id: `assistant-stream-${sessionId}`, + role: 'assistant', + parts: inflightAssistant ? [assistantTextPart(inflightAssistant)] : [], + pending: inflightStreaming + }) + } + + if (queuedUser) { + projected.push({ + id: `user-queued-${sessionId}`, + role: 'user', + parts: [textPart(queuedUser)] + }) + } + + return projected.length ? [...messages, ...projected] : messages +} + export interface BranchMessage { content: string role: ChatMessage['role'] diff --git a/apps/desktop/src/types/hermes.ts b/apps/desktop/src/types/hermes.ts index 0c4a4f8f7..f1914834b 100644 --- a/apps/desktop/src/types/hermes.ts +++ b/apps/desktop/src/types/hermes.ts @@ -386,11 +386,23 @@ export interface SessionMessagesResponse { } export interface SessionResumeResponse { + inflight?: null | { + assistant?: string + streaming?: boolean + user?: string + } + queued?: null | { + user?: string + } info?: SessionRuntimeInfo message_count: number messages: SessionMessage[] resumed: string + running?: boolean session_id: string + session_key?: string + started_at?: number + status?: string } export interface SessionRuntimeInfo { diff --git a/contributors/emails/theunathi@gmail.com b/contributors/emails/theunathi@gmail.com index da16c209d..ad2207957 100644 --- a/contributors/emails/theunathi@gmail.com +++ b/contributors/emails/theunathi@gmail.com @@ -1,2 +1,3 @@ UnathiCodex # PR contribution (tui_gateway: keep busy submits resume-safe) +# PR #66234 (Desktop reconnect and transcript reconciliation) diff --git a/tests/test_tui_gateway_server.py b/tests/test_tui_gateway_server.py index 30b6bcbfb..3c7386c8a 100644 --- a/tests/test_tui_gateway_server.py +++ b/tests/test_tui_gateway_server.py @@ -7976,6 +7976,46 @@ def test_session_activate_returns_inflight_stream_before_completion(monkeypatch) server._sessions.pop("sid-live", None) +def test_session_activate_returns_prompt_queued_during_busy_turn(monkeypatch): + """A full client restart must recover an accepted next-turn prompt. + + Busy prompts are intentionally not durable until they drain. Their only + authoritative copy is ``queued_prompt``, so the live projection must expose + that copy without leaking the transport object. + """ + monkeypatch.setattr(server, "_load_busy_input_mode", lambda: "queue") + monkeypatch.setattr(server, "_session_info", lambda agent: {"model": agent.model}) + agent = types.SimpleNamespace(model="model-live") + session = _session( + agent=agent, + running=True, + inflight_turn={ + "assistant": "partial answer", + "streaming": True, + "user": "current prompt", + }, + ) + server._sessions["sid-live"] = session + try: + queued = server._handle_busy_submit( + "submit", "sid-live", session, "newest prompt", object() + ) + assert queued["result"]["status"] == "queued" + + activated = server.handle_request( + { + "id": "activate", + "method": "session.activate", + "params": {"session_id": "sid-live"}, + } + ) + + assert activated["result"]["queued"] == {"user": "newest prompt"} + assert "transport" not in activated["result"]["queued"] + finally: + server._sessions.pop("sid-live", None) + + def test_session_activate_switches_live_session_without_closing_siblings(monkeypatch): monkeypatch.setattr(server, "_session_info", lambda agent: {"model": agent.model}) server._sessions["sid-a"] = _session( diff --git a/tui_gateway/server.py b/tui_gateway/server.py index a79c06947..523dada17 100644 --- a/tui_gateway/server.py +++ b/tui_gateway/server.py @@ -5602,6 +5602,21 @@ def _inflight_snapshot(session: dict) -> dict | None: } +def _queued_prompt_snapshot(session: dict) -> dict | None: + """Return the accepted next-turn prompt without its transport handle. + + A busy ``prompt.submit`` lives only in ``session["queued_prompt"]`` until + the current turn winds down. Desktop may reconnect or restart during that + window, so the live-session projection must carry the user-visible text; + otherwise the accepted prompt disappears until it finally drains. + """ + queued = session.get("queued_prompt") + if not isinstance(queued, dict): + return None + user = _inflight_text(queued.get("text")) + return {"user": user} if user else None + + # ── Methods: session ───────────────────────────────────────────────── @@ -6438,8 +6453,12 @@ def _session_live_item(sid: str, session: dict, current_sid: str = "") -> dict: history = list(session.get("history") or []) status = _session_live_status(sid, session) inflight = _inflight_snapshot(session) + queued = _queued_prompt_snapshot(session) preview = _message_preview(history) - if inflight: + if queued: + preview = queued.get("user") or preview + preview = " ".join(str(preview).split())[:160] + elif inflight: preview = inflight.get("assistant") or inflight.get("user") or preview preview = " ".join(str(preview).split())[:160] now = time.time() @@ -6508,6 +6527,7 @@ def _live_session_payload( session.get("history") or [] ) inflight = _inflight_snapshot(session) + queued = _queued_prompt_snapshot(session) running = bool(session.get("running")) payload = { "info": _fallback_session_info(session), @@ -6521,6 +6541,8 @@ def _live_session_payload( } if inflight: payload["inflight"] = inflight + if queued: + payload["queued"] = queued return payload