From 2e75f9dc4841ec33aa4459966e4a89965233d490 Mon Sep 17 00:00:00 2001 From: Flownium <157689911+itsflownium@users.noreply.github.com> Date: Tue, 25 Aug 2026 07:56:29 +1000 Subject: [PATCH] fix(desktop): preserve terminal state during reconnect hydration --- .../hooks/use-session-actions.test.tsx | 157 +++++++++++++++++- .../hooks/use-session-actions/index.ts | 73 ++++---- 2 files changed, 196 insertions(+), 34 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 3414003748..781dc8f275 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 @@ -2066,21 +2066,43 @@ describe('resumeSession warm-cache mapping integrity', () => { } const runtimeIdByStoredSessionIdRef: MutableRefObject> = { - current: new Map([['stored-warm', 'runtime-warm']]) + current: new Map([ + ['stored-warm', 'runtime-warm'], + ['stored-legacy', 'runtime-legacy'] + ]) } const sessionStateByRuntimeIdRef: MutableRefObject> = { - current: new Map([['runtime-warm', clientState('stored-warm')]]) + current: new Map([ + ['runtime-warm', clientState('stored-warm')], + ['runtime-legacy', clientState('stored-legacy')] + ]) } // Same-name rows without a source tag are not authoritative for an // explicit owner. Metadata must be re-read from the captured connection. - setSessions([storedSession({ id: 'stored-warm', profile: 'default' })]) + setSessions([ + storedSession({ id: 'stored-warm', profile: 'default' }), + storedSession({ id: 'stored-legacy', profile: 'default' }) + ]) vi.mocked(getSession).mockImplementation(async id => storedSession({ id, profile: 'default' })) vi.mocked(getLatestSessionMessages).mockImplementation(async id => ({ messages: [], session_id: id }) as never) vi.mocked(requestGatewayForAgent).mockImplementation(async (_connectionId, _profile, method, params) => { if (method === 'session.activate') { - throw new Error('Method not found') + if (params?.session_id === 'runtime-legacy') { + throw new Error('Method not found') + } + + return { + info: {}, + message_count: 0, + messages: [], + messages_omitted: true, + resumed: 'stored-warm', + running: false, + session_id: 'runtime-warm', + session_key: 'stored-warm' + } as never } if (method === 'session.usage') { @@ -2116,6 +2138,7 @@ describe('resumeSession warm-cache mapping integrity', () => { await waitFor(() => expect(resume).not.toBeNull()) await resume!('stored-warm', true, ownerRoute) + await resume!('stored-legacy', true, ownerRoute) await resume!('stored-cold', true, ownerRoute) await resume!('stored-cold', true, { connectionId: 'source-b', @@ -2144,7 +2167,7 @@ describe('resumeSession warm-cache mapping integrity', () => { expect.objectContaining({ session_id: 'runtime-warm' }) ) expect(requestGatewayForAgent).toHaveBeenCalledWith('source-a', 'default', 'session.usage', { - session_id: 'runtime-warm' + session_id: 'runtime-legacy' }) expect(requestGatewayForAgent).toHaveBeenCalledWith( 'source-a', @@ -2706,6 +2729,130 @@ describe('resumeSession warm-cache mapping integrity', () => { expect(JSON.stringify(resumedState?.messages)).not.toContain('partial before disconnect') }) + it('keeps a terminal live state when a running reconnect finishes during transcript hydration', async () => { + const runtimeIdByStoredSessionIdRef: MutableRefObject> = { + current: new Map([['stored-A', 'rt-A']]) + } + + const state = clientState('stored-A') + state.busy = true + state.awaitingResponse = true + state.turnLive = true + state.turnStartedAt = 1_700_000_123_000 + state.messages = [ + { + id: 'cached-user', + role: 'user', + parts: [{ type: 'text', text: 'long running prompt' }] + }, + { + id: 'assistant-stream-rt-A', + role: 'assistant', + parts: [{ type: 'text', text: 'partial before terminal event' }], + pending: true + } + ] + state.streamId = 'assistant-stream-rt-A' + + const sessionStateByRuntimeIdRef: MutableRefObject> = { + current: new Map([['rt-A', state]]) + } + + const persisted = deferred>>() + + vi.mocked(getLatestSessionMessages).mockReturnValue(persisted.promise) + + 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: 2, + messages: [], + messages_omitted: true, + running: true, + turn_started_at: 1_700_000_123, + inflight: { + user: 'long running prompt', + assistant: 'partial before terminal event', + streaming: true + }, + 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()) + + const resumePromise = resume!('stored-A', true) + + await waitFor(() => expect(getLatestSessionMessages).toHaveBeenCalledTimes(1)) + expect(sessionStateByRuntimeIdRef.current.get('rt-A')).toMatchObject({ + awaitingResponse: true, + busy: true, + turnLive: true + }) + + // The rebound transport delivers the terminal state while REST is still + // pending. Settle the same stream row the real terminal event owns; the + // later durable hydration may reconcile messages but cannot revive it. + const liveTerminalState = sessionStateByRuntimeIdRef.current.get('rt-A')! + sessionStateByRuntimeIdRef.current.set('rt-A', { + ...liveTerminalState, + adoptedRunningTurn: false, + awaitingResponse: false, + busy: false, + messages: liveTerminalState.messages.map(message => + message.id === 'assistant-stream-rt-A' + ? { + ...message, + parts: [{ type: 'text' as const, text: 'complete durable answer' }], + pending: false + } + : message + ), + streamId: null, + turnLive: false, + turnStartedAt: null + }) + + await act(async () => { + persisted.resolve({ + messages: [ + { content: 'long running prompt', role: 'user', timestamp: 1 }, + { content: 'complete durable answer', role: 'assistant', timestamp: 2 } + ], + session_id: 'stored-A' + } as never) + await resumePromise + }) + + expect(JSON.stringify(resumedState?.messages)).toContain('complete durable answer') + expect(JSON.stringify(resumedState?.messages)).not.toContain('partial before terminal event') + expect(resumedState).toMatchObject({ + adoptedRunningTurn: false, + awaitingResponse: false, + busy: false, + turnLive: false, + turnStartedAt: null + }) + }) + it('preserves cached image attachments through an idle persisted transcript refresh', async () => { const runtimeIdByStoredSessionIdRef: MutableRefObject> = { current: new Map([['stored-A', 'rt-A']]) 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 5186749970..84a7e463dd 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 @@ -973,16 +973,6 @@ export function useSessionActions({ sessionStateByRuntimeIdRef.current.delete(cachedRuntimeId) dropSessionState(cachedRuntimeId) } else { - // session.activate is the ordering barrier for reconnect recovery: - // it atomically rebinds a running turn before returning. If the - // turn is already terminal, this post-barrier REST read sees its - // durable final row; if it is still running, later deltas/finish - // events arrive on the newly attached transport and the - // concurrent-overlay step below preserves them. - const persistedTranscriptPromise = shouldRefreshPersistedTranscript - ? getLatestSessionMessages(storedSessionId, sessionRestScope).catch(() => null) - : null - const pendingApproval = restorePendingApproval(activated, cachedRuntimeId) const pendingClarifyState = restorePendingClarifyFromSnapshot( @@ -1040,6 +1030,49 @@ export function useSessionActions({ ? activated.turn_started_at * 1000 : null + // Settle the activation snapshot before transcript hydration. + // Once the attached transport reports a later terminal event, + // that live state is authoritative and must not be overwritten + // by the older `running` value after the REST request resolves. + const activatedLivenessState = updateSessionState( + cachedRuntimeId, + state => ({ + ...state, + ...(runtimeInfo ?? {}), + busy: running, + awaitingResponse: running && !pendingClarify, + // Resumed onto an already-running turn — that IS backend + // proof the turn is live (no message.start will replay). + turnLive: state.turnLive || running, + needsInput: + pendingApproval || + Boolean(pendingClarify) || + (clarifyAuthoritativelyAbsent ? false : state.needsInput), + // Adopting someone else's turn: we'll stream its reply + // without ever having received its prompt, so the settle + // path must not take the "I saw it all" shortcut. + adoptedRunningTurn: state.adoptedRunningTurn || running, + turnStartedAt: running ? (activatedTurnStartedAt ?? state.turnStartedAt ?? Date.now()) : null + }), + storedSessionId + ) + + busyRef.current = running + setBusy(running) + setAwaitingResponse(running && !pendingClarify) + syncSessionStateToView(cachedRuntimeId, activatedLivenessState) + + // session.activate is the ordering barrier for reconnect recovery: + // it atomically rebinds a running turn before returning. If the + // turn is already terminal, this post-barrier REST read sees its + // durable final row; if it is still running, later deltas/finish + // events arrive on the newly attached transport. Hydration below + // reconciles only messages, so those events also retain liveness + // authority while the request is pending. + const persistedTranscriptPromise = shouldRefreshPersistedTranscript + ? getLatestSessionMessages(storedSessionId, sessionRestScope).catch(() => null) + : null + // The persisted REST transcript is the display authority: a live // runtime may carry only the agent's compressed context projection, // which is intentionally smaller than the user-visible conversation. @@ -1124,22 +1157,7 @@ export function useSessionActions({ cachedRuntimeId, state => ({ ...state, - ...(runtimeInfo ?? {}), messages: visibleActivatedMessages, - busy: running, - awaitingResponse: running, - // Resumed onto an already-running turn — that IS backend - // proof the turn is live (no message.start will replay). - turnLive: state.turnLive || running, - needsInput: - pendingApproval || - Boolean(pendingClarify) || - (clarifyAuthoritativelyAbsent ? false : state.needsInput), - // Adopting someone else's turn: we'll stream its reply - // without ever having received its prompt, so the settle - // path must not take the "I saw it all" shortcut. - adoptedRunningTurn: state.adoptedRunningTurn || running, - turnStartedAt: running ? (activatedTurnStartedAt ?? state.turnStartedAt ?? Date.now()) : null, ...(pendingClarifyProjection ? { awaitingResponse: false, @@ -1149,16 +1167,13 @@ export function useSessionActions({ : {}), ...(clearedClarifyProjection ? { - streamId: running ? (clearedClarifyProjection.streamId ?? state.streamId) : null + streamId: state.busy ? (clearedClarifyProjection.streamId ?? state.streamId) : null } : {}) }), storedSessionId ) - busyRef.current = running - setBusy(running) - setAwaitingResponse(running && !pendingClarify) syncSessionStateToView(cachedRuntimeId, activatedState) // Cache backend transcript truth only. The pending/running bit and // any synthetic clarify row are a live resume projection and must