diff --git a/src/main/ipc/renderer-shutdown-checkpoint-sleep-records.test.ts b/src/main/ipc/renderer-shutdown-checkpoint-sleep-records.test.ts new file mode 100644 index 00000000000..d5121cc5f83 --- /dev/null +++ b/src/main/ipc/renderer-shutdown-checkpoint-sleep-records.test.ts @@ -0,0 +1,127 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { mkdtempSync, rmSync } from 'node:fs' +import { join } from 'node:path' +import { tmpdir } from 'node:os' +import { closeTestStores, createStore, testState } from '../persistence-test-harness' +import { + addWorkspace, + emptySession, + GIT_KEY, + leaf +} from '../../shared/workspace-layout/workspace-layout-session.test-fixture' +import type { SleepingAgentSessionRecord } from '../../shared/agent-session-resume' +import type { WorkspaceSessionState } from '../../shared/workspace-session-state-types' + +const { handlers } = vi.hoisted(() => ({ + handlers: new Map unknown>() +})) + +vi.mock('../../shared/layout-owner-switch', () => ({ ORCA_LAYOUT_OWNER_IS_MAIN: true })) +vi.mock('electron', () => ({ + app: { getPath: () => testState.dir }, + safeStorage: { isEncryptionAvailable: () => false }, + ipcMain: { + handle: ( + channel: string, + handler: (event: { returnValue?: unknown }, args: unknown) => unknown + ) => { + handlers.set(channel, handler) + }, + on: ( + channel: string, + handler: (event: { returnValue?: unknown }, args: unknown) => unknown + ) => { + handlers.set(channel, handler) + } + } +})) +vi.mock('../telemetry/client', () => ({ track: vi.fn() })) +vi.mock('../telemetry/cohort-classifier', () => ({ getCohortAtEmit: vi.fn().mockReturnValue({}) })) +vi.mock('../ssh/ssh-config-parser', () => ({ + loadUserSshConfig: vi.fn(), + sshConfigHostsToTargets: vi.fn() +})) + +import { registerRendererShutdownCheckpointHandler } from './renderer-shutdown-checkpoint' + +const PANE_A = `tab-a:${leaf(1)}` +const PANE_B = `tab-b:${leaf(2)}` + +function session(): WorkspaceSessionState { + return addWorkspace(emptySession(), GIT_KEY, [ + { + id: 'group-1', + tabs: [ + { id: 'tab-a', leaves: [[leaf(1), 'pty-1']] }, + { id: 'tab-b', leaves: [[leaf(2), 'pty-2']] } + ] + } + ]) +} + +function record(paneKey: string, capturedAt: number): SleepingAgentSessionRecord { + return { + paneKey, + tabId: paneKey.split(':')[0], + worktreeId: GIT_KEY, + agent: 'claude', + providerSession: { key: 'session_id', id: `session-${capturedAt}` }, + prompt: 'work', + state: 'working', + capturedAt, + updatedAt: capturedAt + } +} + +function stage(state: WorkspaceSessionState, sleepRecords?: SleepingAgentSessionRecord[]) { + const event: { returnValue?: unknown } = {} + handlers.get('app:stage-before-unload-sync')!(event, { + sessions: [{ state, hostId: 'local' }], + ui: {}, + ...(sleepRecords ? { sleepRecords } : {}) + }) + return event.returnValue +} + +describe('the quit checkpoint in a layout-owner build', () => { + let store: ReturnType + beforeEach(() => { + testState.dir = mkdtempSync(join(tmpdir(), 'orca-test-')) + store = createStore() + registerRendererShutdownCheckpointHandler(store) + }) + afterEach(async () => { + await closeTestStores() + rmSync(testState.dir, { recursive: true, force: true }) + }) + + it('records the quit capture’s resume records in main before staging', () => { + store.setWorkspaceSession(session(), 'local') + const captured = record(PANE_A, 2) + + expect( + stage({ ...session(), sleepingAgentSessionsByPaneKey: { [PANE_A]: captured } }, [captured]) + ).toEqual({ ok: true }) + + expect(store.getWorkspaceSession('local').sleepingAgentSessionsByPaneKey?.[PANE_A]).toEqual( + captured + ) + }) + + it('keeps a wake main recorded that the window had not seen yet', () => { + // Main woke tab-b (a phone showed it) after the window last saw its record. + store.setWorkspaceSession(session(), 'local') + const captured = record(PANE_A, 2) + + stage( + { + ...session(), + sleepingAgentSessionsByPaneKey: { [PANE_A]: captured, [PANE_B]: record(PANE_B, 1) } + }, + [captured] + ) + + const records = store.getWorkspaceSession('local').sleepingAgentSessionsByPaneKey ?? {} + expect(Object.keys(records)).toEqual([PANE_A]) + }) +}) diff --git a/src/main/ipc/renderer-shutdown-checkpoint.ts b/src/main/ipc/renderer-shutdown-checkpoint.ts index ac84f45e10b..62009cb2597 100644 --- a/src/main/ipc/renderer-shutdown-checkpoint.ts +++ b/src/main/ipc/renderer-shutdown-checkpoint.ts @@ -1,5 +1,11 @@ import { ipcMain } from 'electron' -import type { ExecutionHostId } from '../../shared/execution-host' +import { + LOCAL_EXECUTION_HOST_ID, + normalizeExecutionHostId, + type ExecutionHostId +} from '../../shared/execution-host' +import type { SleepingAgentSessionRecord } from '../../shared/agent-session-resume' +import { applyCommandToSession } from '../persistence/terminal-topology/terminal-pane-admission' import type { PersistedUIState } from '../../shared/persisted-ui-state-types' import type { WorkspaceSessionState } from '../../shared/workspace-session-state-types' import type { Store } from '../persistence' @@ -11,6 +17,37 @@ import { ORCA_LAYOUT_OWNER_IS_MAIN } from '../../shared/layout-owner-switch' type StageBeforeUnloadSyncArgs = { sessions: { state: WorkspaceSessionState; hostId?: ExecutionHostId }[] ui: Partial + /** Layout-owner builds: the resume records the quit capture changed. */ + sleepRecords?: SleepingAgentSessionRecord[] +} + +/** + * The partition with the quit capture's resume records written, each workspace's through the + * record-only sleep command a live capture sends: only what the capture changed, so a wake main + * recorded after the window last looked stays. Records of panes the partition lacks change nothing. + */ +function withQuitSleepRecords( + hostId: string | undefined, + session: WorkspaceSessionState, + records: readonly SleepingAgentSessionRecord[] +): WorkspaceSessionState { + const byWorkspace = Map.groupBy(records, (record) => record.worktreeId) + let current = session + for (const [workspace, workspaceRecords] of byWorkspace) { + const change = applyCommandToSession( + normalizeExecutionHostId(hostId) ?? LOCAL_EXECUTION_HOST_ID, + current, + { + type: 'sleep', + workspace, + paneKeys: workspaceRecords.map((record) => record.paneKey), + records: workspaceRecords, + recordOnly: true + } + ) + current = change.session ?? current + } + return current } export type ShutdownCheckpointResult = { ok: boolean } @@ -56,7 +93,11 @@ export function registerRendererShutdownCheckpointHandler(store: Store): void { if (!isFrozenOrcadSourceSessionPartition(store, hostId)) { // Switch point: main keeps its own layout over the window's and takes its view and facts. if (ORCA_LAYOUT_OWNER_IS_MAIN) { - const stored = store.getWorkspaceSession(hostId) + const stored = withQuitSleepRecords( + hostId, + store.getWorkspaceSession(hostId), + args.sleepRecords ?? [] + ) reportWindowSaveDrift(stored, state, hostId, 'shutdown-checkpoint') store.stageWorkspaceSessionBeforeUnload( windowViewSaveForMain(stored, state, hostId), diff --git a/src/preload/api/app-api.ts b/src/preload/api/app-api.ts index 7e7decb8316..168e23df818 100644 --- a/src/preload/api/app-api.ts +++ b/src/preload/api/app-api.ts @@ -1,4 +1,5 @@ import type { AppIdentity } from '../../shared/app-identity' +import type { SleepingAgentSessionRecord } from '../../shared/agent-session-resume' import type { E2EConfig } from '../../shared/e2e-config' import type { ExecutionHostId } from '../../shared/execution-host' import type { @@ -38,6 +39,8 @@ export type AppApi = { stageBeforeUnloadSync: (args: { sessions: { state: WorkspaceSessionState; hostId?: ExecutionHostId }[] ui: Partial + /** Layout-owner builds: the resume records the quit capture changed. */ + sleepRecords?: SleepingAgentSessionRecord[] }) => void /** Resolves once the last staged checkpoint is durably written; rejects if that * write failed, so a reload/restart can abort instead of losing the snapshot. */ diff --git a/src/renderer/src/app-shell/shutdown-checkpoint-persist.test.ts b/src/renderer/src/app-shell/shutdown-checkpoint-persist.test.ts index 9bb9137c405..834d82b9883 100644 --- a/src/renderer/src/app-shell/shutdown-checkpoint-persist.test.ts +++ b/src/renderer/src/app-shell/shutdown-checkpoint-persist.test.ts @@ -21,6 +21,19 @@ function makeDeps( } describe('createShutdownCheckpointPersist', () => { + it('stages the resume records the quit capture changed, for main to record', () => { + const records = [{ paneKey: 'tab-1:leaf-1' }] as never + const deps = makeDeps({ captureSleepingAgentSessions: () => records }) + + createShutdownCheckpointPersist(deps).run() + + expect(deps.stageBeforeUnloadSync).toHaveBeenCalledWith({ + sessions: [{ state: { activeTabId: 't1' } }], + ui: { activeView: 'workspace' }, + sleepRecords: records + }) + }) + it('does not fail the checkpoint when the sleeping-agent quit capture throws (STA-5505)', () => { vi.spyOn(console, 'error').mockImplementation(() => {}) const deps = makeDeps({ diff --git a/src/renderer/src/app-shell/shutdown-checkpoint-persist.ts b/src/renderer/src/app-shell/shutdown-checkpoint-persist.ts index b8a1be67eb2..23194e31bc8 100644 --- a/src/renderer/src/app-shell/shutdown-checkpoint-persist.ts +++ b/src/renderer/src/app-shell/shutdown-checkpoint-persist.ts @@ -1,4 +1,5 @@ import type { PersistedUIState } from '../../../shared/persisted-ui-state-types' +import type { SleepingAgentSessionRecord } from '../../../shared/agent-session-resume' import { formatShutdownCheckpointFailureReason } from '../../../shared/renderer-shutdown-events' import type { WorkspaceSessionHostSnapshot } from '../lib/workspace-session-host-persistence' import { recordRendererCrashBreadcrumb } from '../lib/crash-breadcrumb-recorder' @@ -6,13 +7,16 @@ import { recordRendererCrashBreadcrumb } from '../lib/crash-breadcrumb-recorder' export type ShutdownCheckpointStageArgs = { sessions: WorkspaceSessionHostSnapshot[] ui: Partial + /** Layout-owner builds: the resume records the quit capture changed, which main records. */ + sleepRecords?: SleepingAgentSessionRecord[] } export type ShutdownCheckpointPersistDeps = { shouldCaptureSession: () => boolean /** Per-pane failures are already swallowed inside the capture loop. */ captureTerminalBuffers: () => void - captureSleepingAgentSessions: () => void + /** Returns the records it changed when main records them (a layout-owner build). */ + captureSleepingAgentSessions: () => SleepingAgentSessionRecord[] | void buildSessionSnapshots: () => WorkspaceSessionHostSnapshot[] buildUiPatch: () => Partial hasDirtyOpenFiles: () => boolean @@ -41,10 +45,11 @@ export function createShutdownCheckpointPersist( let fullStagingFailedOnPriorAttempt = false const run = (): void => { const shouldCaptureSession = deps.shouldCaptureSession() + let sleepRecords: SleepingAgentSessionRecord[] | void = undefined if (shouldCaptureSession) { deps.captureTerminalBuffers() try { - deps.captureSleepingAgentSessions() + sleepRecords = deps.captureSleepingAgentSessions() } catch (error) { // Why: blocking the shutdown here is what stranded STA-5505. The cost of // continuing is real but bounded — done panes keep their weaker live-origin @@ -74,7 +79,8 @@ export function createShutdownCheckpointPersist( try { deps.stageBeforeUnloadSync({ sessions: degraded ? [] : sessionSnapshots, - ui: deps.buildUiPatch() + ui: deps.buildUiPatch(), + ...(sleepRecords && sleepRecords.length > 0 && !degraded ? { sleepRecords } : {}) }) } catch (error) { // A durable-only stage has nothing safer left to fall back to. diff --git a/src/renderer/src/app-shell/use-app-session-persistence.ts b/src/renderer/src/app-shell/use-app-session-persistence.ts index 5b2235a3944..5f453bc8145 100644 --- a/src/renderer/src/app-shell/use-app-session-persistence.ts +++ b/src/renderer/src/app-shell/use-app-session-persistence.ts @@ -1,5 +1,6 @@ import { useEffect } from 'react' import { useAppStore } from '../store' +import { ORCA_LAYOUT_OWNER_IS_MAIN } from '../../../shared/layout-owner-switch' import { createSessionWriteSubscriber } from '../lib/session-write-subscriber' import { buildActiveViewUnloadPatch } from '../lib/active-view-persist' import { @@ -19,6 +20,8 @@ import { createShutdownCheckpointGuard } from '../lib/shutdown-checkpoint-guard' import { createShutdownCheckpointPersist } from './shutdown-checkpoint-persist' +import { changedSleepingRecords } from '../store/terminals/terminal-sleep-layout-owner' +import type { SleepingAgentSessionRecord } from '../../../shared/agent-session-resume' import { shutdownBufferCaptures } from '../components/terminal-pane/shutdown-buffer-captures' import { dispatchWindowCloseRequest, @@ -81,8 +84,17 @@ export function useAppSessionPersistence(): void { // which is in-memory. Capture them into the persisted sleeping-session // map so a daemon/session death while the app is closed can still // cold-restore via the agent's resume command (#5232). - captureSleepingAgentSessions: () => - useAppStore.getState().captureAllSleepingAgentSessions('quit'), + captureSleepingAgentSessions: (): SleepingAgentSessionRecord[] | undefined => { + const previous = useAppStore.getState().sleepingAgentSessionsByPaneKey + useAppStore.getState().captureAllSleepingAgentSessions('quit') + // Switch point: the stage writes only the capture's changes into main's records. + if (ORCA_LAYOUT_OWNER_IS_MAIN) { + const state = useAppStore.getState() + return changedSleepingRecords(state, previous, state.sleepingAgentSessionsByPaneKey) + .recorded + } + return undefined + }, // Why: re-read state after capture() calls populated scrollback buffers // into the store via Zustand setters. The shouldCaptureSession read is // only for the gating flags and would miss those updates. diff --git a/src/renderer/src/store/terminals/terminal-sleep-layout-owner.ts b/src/renderer/src/store/terminals/terminal-sleep-layout-owner.ts index a4119022819..0f55d173388 100644 --- a/src/renderer/src/store/terminals/terminal-sleep-layout-owner.ts +++ b/src/renderer/src/store/terminals/terminal-sleep-layout-owner.ts @@ -40,35 +40,49 @@ export function recordPaneSleepThroughLayoutOwner( } } +type HeldRecordsState = Pick + +/** + * The resume records `next` sets or replaces since `previous`, and the ones it drops, of panes in + * tabs the window holds: a closed tab's or removed workspace's records went with it in main, and a + * capture from a stale agent status names a tab main no longer has. + */ +export function changedSleepingRecords( + state: HeldRecordsState, + previous: SleepingRecords, + next: SleepingRecords +): { recorded: SleepingAgentSessionRecord[]; dropped: [string, SleepingAgentSessionRecord][] } { + const held = ([paneKey, record]: [string, SleepingAgentSessionRecord]): boolean => + state.tabsByWorktree[record.worktreeId]?.some( + (tab) => tab.id === parsePaneKey(paneKey)?.tabId + ) ?? false + return { + recorded: Object.entries(next) + .filter((entry) => previous[entry[0]] !== entry[1] && held(entry)) + .map(([, record]) => record), + dropped: Object.entries(previous).filter((entry) => !(entry[0] in next) && held(entry)) + } +} + /** * Every change of the window's resume records, as record-only commands: live and quit captures of * running agents, records a resume or close dropped, and the rest no sleep or wake command sends. * Without these main learned the records only from the window's session save. */ export function recordSleepingRecordChangesThroughLayoutOwner( - state: WorktreeRuntimeOwnerState & Pick, + state: WorktreeRuntimeOwnerState & HeldRecordsState, previous: SleepingRecords, next: SleepingRecords ): void { - // Only panes of tabs the window holds: a closed tab's or removed workspace's records went with it - // in main, and a capture from a stale agent status names a tab main no longer has. - const held = ([paneKey, record]: [string, SleepingAgentSessionRecord]): boolean => - state.tabsByWorktree[record.worktreeId]?.some( - (tab) => tab.id === parsePaneKey(paneKey)?.tabId - ) ?? false - const changedByWorktree = new Map>() - for (const [paneKey, record] of Object.entries(next)) { - if (previous[paneKey] !== record && held([paneKey, record])) { - changedByWorktree.set(record.worktreeId, { - ...changedByWorktree.get(record.worktreeId), - [paneKey]: record - }) - } + const { recorded, dropped } = changedSleepingRecords(state, previous, next) + for (const [worktreeId, records] of Map.groupBy(recorded, (record) => record.worktreeId)) { + recordPaneSleepThroughLayoutOwner( + state, + worktreeId, + Object.fromEntries(records.map((record) => [record.paneKey, record])), + true + ) } - for (const [worktreeId, records] of changedByWorktree) { - recordPaneSleepThroughLayoutOwner(state, worktreeId, records, true) - } - const dropped = Object.entries(previous).filter((entry) => !(entry[0] in next) && held(entry)) for (const [tabId, entries] of byTab(dropped, ([paneKey]) => paneKey)) { recordLayoutApply(state, entries[0]![1].worktreeId, { type: 'wake',