From d07e40ccdaa2f2c82086e0d00b20fbffa350b15d Mon Sep 17 00:00:00 2001 From: Neil Date: Mon, 9 Nov 2026 06:06:57 -0800 Subject: [PATCH 1/4] fix(remote-runtime): requests wait for an in-flight VM wake; resumed remote agents continue a turn their host cut off; restart offers tolerate host clock skew --- .../hosted-agent-auto-resume.ts | 8 +++- .../hosted-session-lease.test.ts | 9 ++-- ...runtime-environment-recipe-vm-wake.test.ts | 38 ++++++++++++++++ .../runtime-environment-recipe-vm-wake.ts | 19 ++++++++ .../sleep/recipe-vm-in-flight-wake.test.ts | 30 +++++++++++++ .../workspaces/sleep/recipe-vm-operations.ts | 23 ++++++++++ .../sleep/recipe-vm-waker-registry.ts | 8 ++++ .../runtime-workspace-sleep-operations.ts | 10 +++++ .../sleep/workspace-sleep-runtime.ts | 15 +++++++ .../startup/main-process-runtime-service.ts | 2 + .../startup/cold-restore-continue.test.ts | 13 ++++++ .../startup/cold-restore-resume-startup.ts | 24 ++++++++++- .../startup/tui-agent-resume-continue.test.ts | 43 +++++++++++++++++++ .../startup/tui-agent-resume-startup.ts | 16 ++++++- 14 files changed, 250 insertions(+), 8 deletions(-) create mode 100644 src/main/runtime/workspaces/sleep/recipe-vm-in-flight-wake.test.ts create mode 100644 src/renderer/src/components/terminal-pane/pty-connection/startup/cold-restore-continue.test.ts create mode 100644 src/shared/agent-launch/startup/tui-agent-resume-continue.test.ts diff --git a/src/main/hosted-sandbox/hosted-agent-auto-resume.ts b/src/main/hosted-sandbox/hosted-agent-auto-resume.ts index 2c1a1fd8e47..6637066ec92 100644 --- a/src/main/hosted-sandbox/hosted-agent-auto-resume.ts +++ b/src/main/hosted-sandbox/hosted-agent-auto-resume.ts @@ -24,6 +24,10 @@ const ContinueAnswer = z.looseObject({ skipped: z.array(z.unknown()).optional() }) +/** `recordedAt` is the host's clock and `drainStartedAt` this machine's: tolerate skew between + * them so an offer the drain just recorded is not mistaken for an older one. */ +export const CLOCK_SKEW_TOLERANCE_MS = 2 * 60_000 + function hasCode(error: unknown, code: string): boolean { return typeof error === 'object' && error !== null && 'code' in error && error.code === code } @@ -47,7 +51,9 @@ export async function autoResumeAgentsAfterRollover( } const ids = list.data.sessions.flatMap((row) => { const offer = Offer.safeParse(row) - return offer.success && offer.data.recordedAt >= drainStartedAt ? [offer.data.sessionId] : [] + return offer.success && offer.data.recordedAt >= drainStartedAt - CLOCK_SKEW_TOLERANCE_MS + ? [offer.data.sessionId] + : [] }) const outcome = { offered: ids.length, continued: 0, skipped: 0, unsupported: false } for (let offset = 0; offset < ids.length; offset += MAX_RESTART_RESUME_SESSIONS) { diff --git a/src/main/hosted-sandbox/hosted-session-lease.test.ts b/src/main/hosted-sandbox/hosted-session-lease.test.ts index ef018530cd6..3400ef8afee 100644 --- a/src/main/hosted-sandbox/hosted-session-lease.test.ts +++ b/src/main/hosted-sandbox/hosted-session-lease.test.ts @@ -112,14 +112,15 @@ describe('agent auto-resume after rollover', () => { if (method === 'agentSession.restartResumable') { return { sessions: [ - { sessionId: 'old', recordedAt: 900 }, - { sessionId: 'a', recordedAt: 1000 }, - { sessionId: 'b', recordedAt: 1500 } + { sessionId: 'old', recordedAt: T0 - 10 * MIN }, + // The host's clock runs 30 s behind this machine's: still this drain's offer. + { sessionId: 'a', recordedAt: T0 - 30_000 }, + { sessionId: 'b', recordedAt: T0 + 500 } ] } } return { continued: [{ sessionId: 'a' }], skipped: ['b'] } - }, 1000) + }, T0) expect(calls).toEqual([ { method: 'agentSession.restartResumable', params: undefined }, { method: 'agentSession.restartContinue', params: { sessionIds: ['a', 'b'] } } diff --git a/src/main/ipc/runtime-environments/runtime-environment-recipe-vm-wake.test.ts b/src/main/ipc/runtime-environments/runtime-environment-recipe-vm-wake.test.ts index 6e3c5b7f1bb..009fd9fbe21 100644 --- a/src/main/ipc/runtime-environments/runtime-environment-recipe-vm-wake.test.ts +++ b/src/main/ipc/runtime-environments/runtime-environment-recipe-vm-wake.test.ts @@ -131,6 +131,44 @@ describe('a request to a recipe VM environment', () => { } ) + it('holds a request that does not open the workspace while an in-flight wake finishes', async () => { + recordVm('running') + ownerState = 'waking' + const awaitWake = vi.fn(async () => 'up' as const) + registerRecipeVmWaker({ + state, + noteOpened: vi.fn(), + holdOpen, + awaitWake, + openEnvironment: async () => 'up' + }) + noteRecipeVmWakePairing(VM_ENVIRONMENT.id, 7, true) + await expect(prepareRecipeVm(userDataPath, VM_ENVIRONMENT, 'none')).resolves.toEqual({ + wokenPairingRevision: null + }) + expect(awaitWake).toHaveBeenCalledWith('workspace-of-vm-1', false, 60_000, undefined) + // It never starts a wake of its own. + expect(holdOpen).not.toHaveBeenCalled() + }) + + it.each([ + ['waking', 'runtime_environment_waking'], + ['wake-failed', 'runtime_environment_wake_failed'], + ['slept', 'runtime_environment_asleep'], + ['asleep', 'runtime_environment_asleep'] + ] as const)('refuses with %s when the in-flight wake ends that way', async (outcome, code) => { + recordVm('running') + ownerState = 'waking' + registerRecipeVmWaker({ + state, + noteOpened: vi.fn(), + holdOpen, + awaitWake: async () => outcome, + openEnvironment: async () => 'up' + }) + await expect(codeOf('none')).resolves.toBe(code) + }) + // A phone gives up after 30 seconds, so its open is held for less. // A Disconnect made before the open holds: no client's open wakes the VM behind that server. diff --git a/src/main/ipc/runtime-environments/runtime-environment-recipe-vm-wake.ts b/src/main/ipc/runtime-environments/runtime-environment-recipe-vm-wake.ts index f906a8612f9..e1dd906d203 100644 --- a/src/main/ipc/runtime-environments/runtime-environment-recipe-vm-wake.ts +++ b/src/main/ipc/runtime-environments/runtime-environment-recipe-vm-wake.ts @@ -111,12 +111,31 @@ export async function prepareRecipeVm( return { wokenPairingRevision: woken !== wakeBefore ? (woken?.pairingRevision ?? null) : null } } const state = waker.state(vm.workspaceId, vm.recordSuspended) + if (state === 'waking' && waker.awaitWake) { + // The wake is already underway (often just reconnecting): a terminal or panel opened in that + // window waits for it instead of failing and never retrying. Nothing here starts a wake. + const wakeBefore = latestRecipeVmWake(environment.id) + const outcome = await waker.awaitWake( + vm.workspaceId, + vm.recordSuspended, + Math.min(holdMs, RECIPE_VM_IN_FLIGHT_WAKE_HOLD_MS), + signal + ) + if (outcome !== 'up') { + throw refusal(outcome) + } + const woken = latestRecipeVmWake(environment.id) + return { wokenPairingRevision: woken !== wakeBefore ? (woken?.pairingRevision ?? null) : null } + } if (state !== 'up') { throw refusal(state) } return UNCHANGED } +/** How long a request that does not open the workspace waits for a wake already in flight. */ +export const RECIPE_VM_IN_FLIGHT_WAKE_HOLD_MS = 60_000 + /** The pairing revision a recipe VM wake this request waited for left its server on, if any. */ export type RecipeVmPreparation = { wokenPairingRevision: number | null } const UNCHANGED: RecipeVmPreparation = { wokenPairingRevision: null } diff --git a/src/main/runtime/workspaces/sleep/recipe-vm-in-flight-wake.test.ts b/src/main/runtime/workspaces/sleep/recipe-vm-in-flight-wake.test.ts new file mode 100644 index 00000000000..bcc6b8075ee --- /dev/null +++ b/src/main/runtime/workspaces/sleep/recipe-vm-in-flight-wake.test.ts @@ -0,0 +1,30 @@ +import { describe, expect, it } from 'vitest' +import { awaitInFlightWake, WorkspaceSleepRequests, type VmOperation } from './recipe-vm-operations' + +describe('awaitInFlightWake', () => { + it('answers at once for a VM that is up or asleep, without starting anything', async () => { + const sleeps = new WorkspaceSleepRequests() + await expect(awaitInFlightWake([], false, sleeps, 'w', 1000)).resolves.toBe('up') + await expect(awaitInFlightWake([], true, sleeps, 'w', 1000)).resolves.toBe('asleep') + }) + + it('waits for the wake already in flight', async () => { + let finish = (): void => {} + const settled = new Promise((resolve) => { + finish = resolve + }) + const pending: VmOperation[] = [{ kind: 'wake', settled }] + const answer = awaitInFlightWake(pending, false, new WorkspaceSleepRequests(), 'w', 5000) + finish() + await expect(answer).resolves.toBe('up') + }) + + it('ends slept when a Sleep is asked for while it waits, and waking when the hold runs out', async () => { + const sleeps = new WorkspaceSleepRequests() + const never: VmOperation[] = [{ kind: 'wake', settled: new Promise(() => {}) }] + const slept = awaitInFlightWake(never, false, sleeps, 'w', 5000) + sleeps.noteRequested('w') + await expect(slept).resolves.toBe('slept') + await expect(awaitInFlightWake(never, false, sleeps, 'w', 5)).resolves.toBe('waking') + }) +}) diff --git a/src/main/runtime/workspaces/sleep/recipe-vm-operations.ts b/src/main/runtime/workspaces/sleep/recipe-vm-operations.ts index 00305978ffd..1778f20ef98 100644 --- a/src/main/runtime/workspaces/sleep/recipe-vm-operations.ts +++ b/src/main/runtime/workspaces/sleep/recipe-vm-operations.ts @@ -124,3 +124,26 @@ export async function holdOpenUntilUp( stopAbortWatch() } } + +/** Holds a request that does not open the workspace while the wake already in flight settles; a VM + * that is not waking answers at once with where it stands. Never starts a wake. */ +export async function awaitInFlightWake( + pending: readonly VmOperation[], + recordSuspended: boolean, + sleepRequests: WorkspaceSleepRequests, + worktreeId: string, + timeoutMs: number, + signal?: AbortSignal +): Promise { + signal?.throwIfAborted() + const state = vmStateOf(pending, recordSuspended) + if (state !== 'waking') { + return state === 'up' ? 'up' : 'asleep' + } + const sleep = sleepRequests.next(worktreeId) + try { + return await holdOpenUntilUp(pendingVmWake(pending), sleep.requested, timeoutMs, signal) + } finally { + sleep.stop() + } +} diff --git a/src/main/runtime/workspaces/sleep/recipe-vm-waker-registry.ts b/src/main/runtime/workspaces/sleep/recipe-vm-waker-registry.ts index b2697ae0ecd..cead8b03c5e 100644 --- a/src/main/runtime/workspaces/sleep/recipe-vm-waker-registry.ts +++ b/src/main/runtime/workspaces/sleep/recipe-vm-waker-registry.ts @@ -15,6 +15,14 @@ export type RecipeVmWaker = { timeoutMs: number, signal?: AbortSignal ) => Promise + /** Holds a request that does not open the workspace while a wake already in flight finishes; + * never starts one. */ + awaitWake?: ( + workspaceId: string, + recordSuspended: boolean, + timeoutMs: number, + signal?: AbortSignal + ) => Promise /** Holds an open of the workspace behind a saved environment (id or name) for the CLI. */ openEnvironment: (selector: string) => Promise } diff --git a/src/main/runtime/workspaces/sleep/runtime-workspace-sleep-operations.ts b/src/main/runtime/workspaces/sleep/runtime-workspace-sleep-operations.ts index cb1a161ffd6..ee3b05ffc87 100644 --- a/src/main/runtime/workspaces/sleep/runtime-workspace-sleep-operations.ts +++ b/src/main/runtime/workspaces/sleep/runtime-workspace-sleep-operations.ts @@ -161,6 +161,16 @@ export abstract class RuntimeWorkspaceSleepOperations extends OrcaRuntimeWithRes return this.workspaceSleep.holdOpen(worktreeId, recordSuspended, timeoutMs, signal) } + /** Holds a request that does not open this workspace while its in-flight wake finishes. */ + awaitWorkspaceVmWake( + worktreeId: string, + recordSuspended: boolean, + timeoutMs: number, + signal?: AbortSignal + ): Promise { + return this.workspaceSleep.awaitWake(worktreeId, recordSuspended, timeoutMs, signal) + } + /** The window opened a slept pane; main resumes it and the window attaches. */ resumeSleptPane(pane: RestoredPane, automatic?: boolean): Promise { return this.workspaceSleep.resumePane(pane, automatic) diff --git a/src/main/runtime/workspaces/sleep/workspace-sleep-runtime.ts b/src/main/runtime/workspaces/sleep/workspace-sleep-runtime.ts index 8bd911b0526..cb30955f5d3 100644 --- a/src/main/runtime/workspaces/sleep/workspace-sleep-runtime.ts +++ b/src/main/runtime/workspaces/sleep/workspace-sleep-runtime.ts @@ -1,4 +1,5 @@ import { + awaitInFlightWake, holdOpenUntilUp, pendingVmWake, trackPendingVmOperation, @@ -209,6 +210,20 @@ export class WorkspaceSleepCoordinator { } } + /** Holds a request that does not open the workspace while a wake already in flight finishes; + * never starts a wake or counts as a reopening, so background traffic leaves a VM asleep. */ + awaitWake(worktreeId: string, recordSuspended: boolean, timeoutMs: number, signal?: AbortSignal) { + const pending = this.pendingVmOperations.get(worktreeId) ?? [] + return awaitInFlightWake( + pending, + recordSuspended, + this.sleepRequests, + worktreeId, + timeoutMs, + signal + ) + } + private trackVmOperation(worktreeId: string, operation: VmOperation, run: Promise) { return trackPendingVmOperation(this.pendingVmOperations, worktreeId, operation, run) } diff --git a/src/main/startup/main-process-runtime-service.ts b/src/main/startup/main-process-runtime-service.ts index 7b0405af786..2fc9e8f28f4 100644 --- a/src/main/startup/main-process-runtime-service.ts +++ b/src/main/startup/main-process-runtime-service.ts @@ -166,6 +166,8 @@ export function initializeMainProcessRuntime(): OrcaRuntimeService { noteOpened: (workspaceId) => runtime.noteWorkspaceReopened(workspaceId), holdOpen: (workspaceId, suspended, timeoutMs, signal) => runtime.holdWorkspaceVmOpen(workspaceId, suspended, timeoutMs, signal), + awaitWake: (workspaceId, suspended, timeoutMs, signal) => + runtime.awaitWorkspaceVmWake(workspaceId, suspended, timeoutMs, signal), openEnvironment: (selector) => { const root = getRuntimeEnvironmentRegistryRoot() return openRecipeVmEnvironment(root, resolveEnvironment(root, selector)) diff --git a/src/renderer/src/components/terminal-pane/pty-connection/startup/cold-restore-continue.test.ts b/src/renderer/src/components/terminal-pane/pty-connection/startup/cold-restore-continue.test.ts new file mode 100644 index 00000000000..a0487b3e0a2 --- /dev/null +++ b/src/renderer/src/components/terminal-pane/pty-connection/startup/cold-restore-continue.test.ts @@ -0,0 +1,13 @@ +import { describe, expect, it, vi } from 'vitest' + +vi.mock('@/store/index', () => ({ useAppStore: { getState: () => ({}) } })) + +describe('shouldContinueInterruptedTurn', () => { + it('continues only a remote agent whose last status was working', async () => { + const { shouldContinueInterruptedTurn } = await import('./cold-restore-resume-startup') + expect(shouldContinueInterruptedTurn({ state: 'working' }, 'env-1')).toBe(true) + expect(shouldContinueInterruptedTurn({ state: 'waiting' }, 'env-1')).toBe(false) + expect(shouldContinueInterruptedTurn({ state: 'working' }, null)).toBe(false) + expect(shouldContinueInterruptedTurn(null, 'env-1')).toBe(false) + }) +}) diff --git a/src/renderer/src/components/terminal-pane/pty-connection/startup/cold-restore-resume-startup.ts b/src/renderer/src/components/terminal-pane/pty-connection/startup/cold-restore-resume-startup.ts index 983d4b9c96f..7abba5460db 100644 --- a/src/renderer/src/components/terminal-pane/pty-connection/startup/cold-restore-resume-startup.ts +++ b/src/renderer/src/components/terminal-pane/pty-connection/startup/cold-restore-resume-startup.ts @@ -6,7 +6,10 @@ import { normalizeExecutionHostId } from '../../../../../../shared/execution-hos import { parseWorkspaceKey } from '../../../../../../shared/workspace-scope' import { useAppStore } from '@/store/index' import { createNonSecureContextUuid as createBrowserUuid } from '../../../../../../shared/non-secure-context-uuid' -import { buildAgentResumeStartupPlan } from '../../../../../../shared/agent-launch/startup/tui-agent-resume-startup' +import { + INTERRUPTED_TURN_CONTINUE_PROMPT, + buildAgentResumeStartupPlan +} from '../../../../../../shared/agent-launch/startup/tui-agent-resume-startup' import { resolveRuntimeHostReportedPlatform, resolveRuntimeHostTerminalWindowsShell, @@ -31,6 +34,15 @@ import type { ConnectPanePtySession } from '../connect-pane-pty-session' export const NOT_REOPENED: HeldPaneReopenOutcome = { reopened: false, agentResumed: false } +/** A cold-restored agent on a remote host whose last status was working lost its turn when the + * host restarted; local panes are left alone (their PTYs outlive app restarts). */ +export function shouldContinueInterruptedTurn( + entry: { state: string } | null | undefined, + runtimeEnvironmentId: string | null +): boolean { + return Boolean(runtimeEnvironmentId) && entry?.state === 'working' +} + type AppState = ReturnType /** The agent session a cold restore of this pane would resume; read-only, so pre-spawn checks can share it. */ @@ -117,7 +129,15 @@ export function bindBuildColdRestoreAgentResumeStartup(session: ConnectPanePtySe platform: resumeTarget.platform, agentHostPlatform: resumeTarget.agentHostPlatform, shell: resumeTarget.shell, - launchCwd: session.worktree?.path + launchCwd: session.worktree?.path, + // A remote host that restarted under a working agent (reboot, crash, VM rollover) killed it + // mid-turn; its last status still says working. Carry the turn on instead of waiting. + continuePrompt: shouldContinueInterruptedTurn( + entry, + session.transport.getRuntimeEnvironmentId?.() ?? null + ) + ? INTERRUPTED_TURN_CONTINUE_PROMPT + : null }) if (!startupPlan) { return null diff --git a/src/shared/agent-launch/startup/tui-agent-resume-continue.test.ts b/src/shared/agent-launch/startup/tui-agent-resume-continue.test.ts new file mode 100644 index 00000000000..c8bb1747775 --- /dev/null +++ b/src/shared/agent-launch/startup/tui-agent-resume-continue.test.ts @@ -0,0 +1,43 @@ +import { describe, expect, it } from 'vitest' +import { buildAgentResumeStartupPlan } from './tui-agent-resume-startup' + +const providerSession = { key: 'session_id', id: 'claude-session-1' } as const + +describe('resume with a continue prompt', () => { + const base = { cmdOverrides: {}, platform: 'linux' as const, shell: 'posix' as const } + + it('submits the prompt on resume for Claude and Codex', () => { + const claude = buildAgentResumeStartupPlan({ + ...base, + agent: 'claude', + providerSession, + continuePrompt: 'continue' + }) + expect(claude?.launchCommand).toMatch(/'?--resume'? '?claude-session-1'? '?continue'?$/) + const codex = buildAgentResumeStartupPlan({ + ...base, + agent: 'codex', + providerSession: { key: 'session_id', id: 'codex-1' }, + continuePrompt: 'continue' + }) + expect(codex?.launchCommand).toMatch(/'?resume'? '?codex-1'? '?continue'?$/) + }) + + it('leaves other agents and empty prompts unchanged', () => { + const plain = buildAgentResumeStartupPlan({ ...base, agent: 'claude', providerSession }) + const blank = buildAgentResumeStartupPlan({ + ...base, + agent: 'claude', + providerSession, + continuePrompt: ' ' + }) + expect(blank?.launchCommand).toBe(plain?.launchCommand) + const cursor = buildAgentResumeStartupPlan({ + ...base, + agent: 'cursor', + providerSession: { key: 'conversation_id', id: 'c-1' }, + continuePrompt: 'continue' + }) + expect(cursor?.launchCommand).not.toMatch(/continue'?$/) + }) +}) diff --git a/src/shared/agent-launch/startup/tui-agent-resume-startup.ts b/src/shared/agent-launch/startup/tui-agent-resume-startup.ts index 3f981d79664..a1f61d88365 100644 --- a/src/shared/agent-launch/startup/tui-agent-resume-startup.ts +++ b/src/shared/agent-launch/startup/tui-agent-resume-startup.ts @@ -37,6 +37,9 @@ export function buildAgentResumeStartupPlan(args: { fork?: boolean /** OS that runs the agent when `platform` only picks quoting; null when unreported. */ agentHostPlatform?: NodeJS.Platform | null + /** First message the resumed session submits, e.g. to carry on a turn its host cut off. Only + * agents whose resume takes a prompt argument (Claude, Codex) send it; others ignore it. */ + continuePrompt?: string | null }): AgentStartupPlan | null { const resumeArgv = args.fork ? getAgentForkArgv(args.agent, args.providerSession) @@ -45,10 +48,15 @@ export function buildAgentResumeStartupPlan(args: { return null } // Codex otherwise asks for its old folder, even if deleted (#17745). Keep before resume so argv removal preserves the folder choice. - const argv = + const withCwd = args.agent === 'codex' && args.resumeInLaunchCwd ? [resumeArgv[0], '-c', 'tui.resume_cwd=current', ...resumeArgv.slice(1)] : resumeArgv + const continuePrompt = args.continuePrompt?.trim() + const argv = + continuePrompt && !args.fork && RESUME_PROMPT_AGENTS.has(args.agent) + ? [...withCwd, continuePrompt] + : withCwd const shell = resolveStartupShell(args.platform, args.shell) // Why: Codex restores a resumed thread's saved permissions unless the launch passes one, so // replaying Orca's launch-time flag would undo a change made inside the session (#10886). @@ -112,6 +120,12 @@ export function buildAgentResumeStartupPlan(args: { } } +/** Agents whose resume command accepts a positional prompt it submits on start. */ +const RESUME_PROMPT_AGENTS: ReadonlySet = new Set(['claude', 'codex']) + +/** What a resumed agent is told when its host went away mid-turn (restart, reboot, crash). */ +export const INTERRUPTED_TURN_CONTINUE_PROMPT = 'continue' + /** WSL panes start at a UNC path while the agent reports its Linux path. */ function isSameFolder(agentCwd: string, launchCwd: string): boolean { const launchPath = parseWslUncPath(launchCwd)?.linuxPath ?? launchCwd From eea4d2184ed5df78da52139bd442fd50302dafb3 Mon Sep 17 00:00:00 2001 From: Neil Date: Mon, 9 Nov 2026 06:49:54 -0800 Subject: [PATCH 2/4] fix(cloud): wake progress reports ready only once the woken server answers terminal requests --- .../hosted-workspace-runtime.ts | 22 +++++++++++++++++++ .../hosted-workspace-service.ts | 7 ++++++ 2 files changed, 29 insertions(+) diff --git a/src/main/hosted-sandbox/hosted-workspace-runtime.ts b/src/main/hosted-sandbox/hosted-workspace-runtime.ts index 1d54e278900..69b25b17502 100644 --- a/src/main/hosted-sandbox/hosted-workspace-runtime.ts +++ b/src/main/hosted-sandbox/hosted-workspace-runtime.ts @@ -17,6 +17,7 @@ import { } from '../environment-providers/environment-provider-driver' import { environmentHandleFromRecipeResult } from '../environment-providers/environment-provider-recipe-adapter' import { isHostedEntitlementDenied } from './hosted-environment-provider' +import { callRuntimeEnvironment } from '../ipc/runtime-environments/runtime-environment-transport-routing' import { HOSTED_SANDBOX_DRIVER_ID } from './hosted-sandbox-driver' import type { HostedWorkspaceService } from './hosted-workspace-service' @@ -163,3 +164,24 @@ export function setHostedWorkspaceService(next: HostedWorkspaceService | null): export function getHostedWorkspaceService(): HostedWorkspaceService | null { return current } + +/** Polls a cheap terminal listing until the woken server answers it (or the deadline passes). */ +export async function waitForTerminalsReachable( + userDataPath: string, + environmentId: string, + deadline: number +): Promise { + while (Date.now() < deadline) { + const response = await callRuntimeEnvironment( + userDataPath, + environmentId, + 'terminal.list', + { limit: 1 }, + 5_000 + ).catch(() => null) + if (response?.ok) { + return + } + await new Promise((resolve) => setTimeout(resolve, 250)) + } +} diff --git a/src/main/hosted-sandbox/hosted-workspace-service.ts b/src/main/hosted-sandbox/hosted-workspace-service.ts index 6cf52ac7689..9297e972799 100644 --- a/src/main/hosted-sandbox/hosted-workspace-service.ts +++ b/src/main/hosted-sandbox/hosted-workspace-service.ts @@ -27,6 +27,7 @@ import { failure, findCloudRuntime, readCachedCheckouts, + waitForTerminalsReachable, type CloudRuntime } from './hosted-workspace-runtime' @@ -99,6 +100,12 @@ export function installHostedWorkspaceService(deps: { while (Date.now() < deadline && wakes.has(workspaceId)) { const runtime = findCloudRuntime(deps.userDataPath, workspaceId)?.runtime if (runtime?.status === 'running') { + // Running is not yet usable: panes rebind to the restarted server's new terminals a moment + // later, and keys typed before that are dropped by design. Ready = the server answers a + // terminal listing through the window's own transport. + if (runtime.runtimeEnvironmentId) { + await waitForTerminalsReachable(deps.userDataPath, runtime.runtimeEnvironmentId, deadline) + } service.emitWakeProgress({ workspaceId, phase: 'ready', source: 'open', startedAt }) return } From 60223bef63948cdc3962c505bb1cad4459330f77 Mon Sep 17 00:00:00 2001 From: Neil Date: Mon, 9 Nov 2026 07:02:34 -0800 Subject: [PATCH 3/4] fix(cloud): review fixes for wake readiness, in-flight wake hold outside call slots, drain-scoped agent auto-resume, gated interrupted-turn continue --- .../hosted-agent-auto-resume.ts | 43 ++++++++++-- .../hosted-session-lease-scheduler.ts | 10 ++- .../hosted-session-lease.test.ts | 34 +++++++++- .../hosted-sandbox/hosted-session-registry.ts | 6 +- .../hosted-workspace-runtime.ts | 65 +++++++++++++++++-- .../hosted-workspace-service.ts | 41 +++++------- .../hosted-workspace-wake-ready.test.ts | 49 ++++++++++++++ ...runtime-environment-recipe-vm-wake.test.ts | 35 +++++----- .../runtime-environment-recipe-vm-wake.ts | 64 +++++++++++++----- .../runtime-environment-transport-routing.ts | 2 + .../startup/main-process-runtime-service.ts | 6 +- .../startup/cold-restore-continue.test.ts | 55 ++++++++++++++-- .../startup/cold-restore-resume-startup.ts | 33 ++++++++-- 13 files changed, 359 insertions(+), 84 deletions(-) create mode 100644 src/main/hosted-sandbox/hosted-workspace-wake-ready.test.ts diff --git a/src/main/hosted-sandbox/hosted-agent-auto-resume.ts b/src/main/hosted-sandbox/hosted-agent-auto-resume.ts index 6637066ec92..4eaae15a42d 100644 --- a/src/main/hosted-sandbox/hosted-agent-auto-resume.ts +++ b/src/main/hosted-sandbox/hosted-agent-auto-resume.ts @@ -27,14 +27,45 @@ const ContinueAnswer = z.looseObject({ /** `recordedAt` is the host's clock and `drainStartedAt` this machine's: tolerate skew between * them so an offer the drain just recorded is not mistaken for an older one. */ export const CLOCK_SKEW_TOLERANCE_MS = 2 * 60_000 +/** A host too slow to list its offers before a restart is treated as unable to (falls back to time). */ +const OFFER_SNAPSHOT_TIMEOUT_MS = 5_000 function hasCode(error: unknown, code: string): boolean { return typeof error === 'object' && error !== null && 'code' in error && error.code === code } +/** Offer ids the host already holds before a drain; offers in it were left for the user. Null when + * the host cannot answer (an unresponsive host is often why it is being restarted). */ +export async function listRestartOfferIds(call: RuntimeCall): Promise | null> { + try { + const listed = await Promise.race([ + call('agentSession.restartResumable'), + new Promise((resolve) => setTimeout(() => resolve(null), OFFER_SNAPSHOT_TIMEOUT_MS)) + ]) + const list = OfferList.safeParse(listed) + if (!list.success) { + return null + } + return new Set( + list.data.sessions.flatMap((row) => { + const offer = Offer.safeParse(row) + return offer.success ? [offer.data.sessionId] : [] + }) + ) + } catch { + return null + } +} + +/** + * Continues the offers this drain recorded. With a pre-drain snapshot of the host's offers, "recorded + * by this drain" is set membership (no clock comparison); without one it falls back to the host + * timestamp against this machine's drain start, allowing for clock skew. + */ export async function autoResumeAgentsAfterRollover( call: RuntimeCall, - drainStartedAt: number + drainStartedAt: number, + offersBeforeDrain: ReadonlySet | null = null ): Promise { let listed: unknown try { @@ -51,9 +82,13 @@ export async function autoResumeAgentsAfterRollover( } const ids = list.data.sessions.flatMap((row) => { const offer = Offer.safeParse(row) - return offer.success && offer.data.recordedAt >= drainStartedAt - CLOCK_SKEW_TOLERANCE_MS - ? [offer.data.sessionId] - : [] + if (!offer.success) { + return [] + } + const recordedByThisDrain = offersBeforeDrain + ? !offersBeforeDrain.has(offer.data.sessionId) + : offer.data.recordedAt >= drainStartedAt - CLOCK_SKEW_TOLERANCE_MS + return recordedByThisDrain ? [offer.data.sessionId] : [] }) const outcome = { offered: ids.length, continued: 0, skipped: 0, unsupported: false } for (let offset = 0; offset < ids.length; offset += MAX_RESTART_RESUME_SESSIONS) { diff --git a/src/main/hosted-sandbox/hosted-session-lease-scheduler.ts b/src/main/hosted-sandbox/hosted-session-lease-scheduler.ts index 26bde666efb..1e01085bb6c 100644 --- a/src/main/hosted-sandbox/hosted-session-lease-scheduler.ts +++ b/src/main/hosted-sandbox/hosted-session-lease-scheduler.ts @@ -23,7 +23,12 @@ export type HostedSessionLeasePorts = { warn: (minutesLeft: HostedLeaseWarning) => void /** Re-pairs with the rebooted orcad; resolves once the runtime answers again. */ reconnect: (ready: HostedEnvironmentReady) => Promise - resumeAgents: (drainStartedAt: number) => Promise + resumeAgents: ( + drainStartedAt: number, + offersBeforeDrain: ReadonlySet | null + ) => Promise + /** The host's restart offers before the drain, so only the drain's own are continued. */ + snapshotOffers?: () => Promise | null> onSuspended: () => void /** Replaces the driver suspend, e.g. to run the window's ordinary Sleep for the workspace. */ suspend?: () => Promise @@ -117,6 +122,7 @@ export class HostedSessionLeaseScheduler { private async rollover(reason: 'opportunistic' | 'forced'): Promise { const { driver, handle } = this.ports + const offersBeforeDrain = (await this.ports.snapshotOffers?.().catch(() => null)) ?? null const drainStartedAt = this.now() this.log(`[orca-cloud] ${reason} rollover`) const ready = await driver.rollover(handle) @@ -124,7 +130,7 @@ export class HostedSessionLeaseScheduler { this.warned.clear() await this.ports.reconnect(ready) try { - const outcome = await this.ports.resumeAgents(drainStartedAt) + const outcome = await this.ports.resumeAgents(drainStartedAt, offersBeforeDrain) this.log( `[orca-cloud] continued ${outcome.continued}/${outcome.offered} agents after rollover` ) diff --git a/src/main/hosted-sandbox/hosted-session-lease.test.ts b/src/main/hosted-sandbox/hosted-session-lease.test.ts index 3400ef8afee..6c508f641c5 100644 --- a/src/main/hosted-sandbox/hosted-session-lease.test.ts +++ b/src/main/hosted-sandbox/hosted-session-lease.test.ts @@ -1,6 +1,6 @@ import { describe, expect, it, vi } from 'vitest' import type { HostedEnvironmentReady } from './hosted-environment-provider' -import { autoResumeAgentsAfterRollover } from './hosted-agent-auto-resume' +import { autoResumeAgentsAfterRollover, listRestartOfferIds } from './hosted-agent-auto-resume' import { createHostedSandboxControlPlaneClient } from './hosted-sandbox-control-plane-client' import { createHostedSandboxDriver, HOSTED_SANDBOX_CAPABILITIES } from './hosted-sandbox-driver' import { decideHostedLeaseAction, type HostedActivitySample } from './hosted-session-lease-policy' @@ -98,7 +98,7 @@ describe('hosted session lease scheduler', () => { expect(await scheduler.tick()).toEqual({ kind: 'rollover', reason: 'forced' }) expect(reconnect).toHaveBeenCalledTimes(1) expect(reconnect.mock.calls[0]?.[0].session.startedAt).toBe(now) - expect(resumeAgents).toHaveBeenCalledWith(created.session.deadlineAt) + expect(resumeAgents).toHaveBeenCalledWith(created.session.deadlineAt, null) // The new session is far from its cap, so the next tick does nothing drastic. expect((await scheduler.tick()).kind).not.toBe('rollover') }) @@ -141,3 +141,33 @@ describe('agent auto-resume after rollover', () => { expect(call).toHaveBeenCalledTimes(1) }) }) + +describe('agent auto-resume with a pre-drain offer snapshot', () => { + it('continues only offers the host did not already hold, whatever the clocks say', async () => { + const offers = [ + { sessionId: 'left-for-user', recordedAt: T0 + 5_000 }, + { sessionId: 'this-drain', recordedAt: T0 - 9 * MIN } + ] + const calls: unknown[] = [] + const call = async (method: string, params?: Record) => { + calls.push({ method, params }) + return method === 'agentSession.restartResumable' ? { sessions: offers } : { continued: [{}] } + } + const before = await listRestartOfferIds(async () => ({ sessions: [offers[0]] })) + expect(before).toEqual(new Set(['left-for-user'])) + const outcome = await autoResumeAgentsAfterRollover(call, T0, before) + expect(calls).toContainEqual({ + method: 'agentSession.restartContinue', + params: { sessionIds: ['this-drain'] } + }) + expect(outcome.offered).toBe(1) + }) + + it('answers null when the host cannot list its offers', async () => { + await expect( + listRestartOfferIds(async () => { + throw new Error('down') + }) + ).resolves.toBeNull() + }) +}) diff --git a/src/main/hosted-sandbox/hosted-session-registry.ts b/src/main/hosted-sandbox/hosted-session-registry.ts index bc7094e7d19..82e39e59efc 100644 --- a/src/main/hosted-sandbox/hosted-session-registry.ts +++ b/src/main/hosted-sandbox/hosted-session-registry.ts @@ -18,7 +18,7 @@ import { noteRecipeVmWakePairing } from '../ephemeral-vm/recipe-vm-wake-pairings import { getEphemeralVmRecipeResultPairingCode } from '../../shared/ephemeral-vm-recipes' import { environmentHandleFromRecipeResult } from '../environment-providers/environment-provider-recipe-adapter' import type { HostedEnvironmentProviderDriver } from './hosted-environment-provider' -import { autoResumeAgentsAfterRollover } from './hosted-agent-auto-resume' +import { autoResumeAgentsAfterRollover, listRestartOfferIds } from './hosted-agent-auto-resume' import { HOSTED_SANDBOX_DRIVER_ID } from './hosted-sandbox-driver' import { createHostedRuntimeEnvironmentCall } from './hosted-runtime-environment-call' import type { HostedActivitySample } from './hosted-session-lease-policy' @@ -112,7 +112,9 @@ export function createHostedSessionRegistry(deps: { true ) }, - resumeAgents: (drainStartedAt) => autoResumeAgentsAfterRollover(call, drainStartedAt), + snapshotOffers: () => listRestartOfferIds(call), + resumeAgents: (drainStartedAt, offersBeforeDrain) => + autoResumeAgentsAfterRollover(call, drainStartedAt, offersBeforeDrain), // The window runs its ordinary Sleep, so the sidebar, chats and VM suspend stay one path. onSuspended: () => log(`[orca-cloud] idle sleep requested for ${runtime.workspaceName ?? runtime.id}`), diff --git a/src/main/hosted-sandbox/hosted-workspace-runtime.ts b/src/main/hosted-sandbox/hosted-workspace-runtime.ts index 69b25b17502..ffd4ce8feab 100644 --- a/src/main/hosted-sandbox/hosted-workspace-runtime.ts +++ b/src/main/hosted-sandbox/hosted-workspace-runtime.ts @@ -4,6 +4,7 @@ import { listEphemeralVmRuntimes } from '../../shared/ephemeral-vm-runtime-store import type { EphemeralVmRuntimeRecord } from '../../shared/ephemeral-vm-runtimes' import type { CloudOomEvent, + CloudWakeProgressEvent, CloudWorkspaceErrorCode, CloudWorkspaceResult } from '../../shared/hosted-sandbox/cloud-workspace-api' @@ -169,17 +170,73 @@ export function getHostedWorkspaceService(): HostedWorkspaceService | null { export async function waitForTerminalsReachable( userDataPath: string, environmentId: string, - deadline: number -): Promise { - while (Date.now() < deadline) { + deadline: number, + stillWanted: () => boolean = () => true +): Promise { + while (Date.now() < deadline && stillWanted()) { + const remaining = deadline - Date.now() + // Bounded by the deadline: the probe neither waits past it on a wake in flight nor on a slow server. + const signal = AbortSignal.timeout(Math.max(1, remaining)) const response = await callRuntimeEnvironment( userDataPath, environmentId, 'terminal.list', { limit: 1 }, - 5_000 + Math.min(5_000, Math.max(1, remaining)), + undefined, + undefined, + { signal, preparation: { recipeVmHoldMs: Math.max(0, remaining) } } ).catch(() => null) if (response?.ok) { + return stillWanted() + } + await new Promise((resolve) => setTimeout(resolve, 250)) + } + return false +} + +/** + * After a wake's step answers: `ready` once the runtime record says running AND the woken server + * answers a terminal listing (panes rebind to the restarted server's terminals a moment after + * running; keys typed before that are dropped by design). `failed` when it never answers. Nothing + * when this wake is no longer the current one (a newer wake, a failure or a Sleep took over). + */ +export async function announceWakeUsable( + userDataPath: string, + workspaceId: string, + startedAt: number, + hooks: { isCurrent: () => boolean; emit: (event: Omit) => void }, + timeoutMs = 60_000 +): Promise { + const deadline = Date.now() + timeoutMs + const running = (): EphemeralVmRuntimeRecord | null => { + const runtime = findCloudRuntime(userDataPath, workspaceId)?.runtime + return runtime?.status === 'running' ? runtime : null + } + while (Date.now() < deadline && hooks.isCurrent()) { + const runtime = running() + if (runtime) { + const reachable = runtime.runtimeEnvironmentId + ? await waitForTerminalsReachable( + userDataPath, + runtime.runtimeEnvironmentId, + deadline, + hooks.isCurrent + ) + : true + if (hooks.isCurrent() && running()) { + hooks.emit( + reachable + ? { workspaceId, phase: 'ready', source: 'open', startedAt } + : { + workspaceId, + phase: 'failed', + source: 'open', + startedAt, + message: 'The workspace did not answer after waking.' + } + ) + } return } await new Promise((resolve) => setTimeout(resolve, 250)) diff --git a/src/main/hosted-sandbox/hosted-workspace-service.ts b/src/main/hosted-sandbox/hosted-workspace-service.ts index 9297e972799..bd42a7ea2d1 100644 --- a/src/main/hosted-sandbox/hosted-workspace-service.ts +++ b/src/main/hosted-sandbox/hosted-workspace-service.ts @@ -17,7 +17,7 @@ import type { import { applyResumedRuntimePairing } from '../ephemeral-vm/runtime-resume-pairing' import { noteRecipeVmWakePairing } from '../ephemeral-vm/recipe-vm-wake-pairings' import { invalidateRuntimeEnvironmentTransport } from '../ipc/runtime-environments/runtime-environments' -import { autoResumeAgentsAfterRollover } from './hosted-agent-auto-resume' +import { autoResumeAgentsAfterRollover, listRestartOfferIds } from './hosted-agent-auto-resume' import type { HostedSandboxControlPlaneClient } from './hosted-sandbox-control-plane-client' import { toReady } from './hosted-sandbox-driver' import { createHostedRuntimeEnvironmentCall } from './hosted-runtime-environment-call' @@ -27,7 +27,7 @@ import { failure, findCloudRuntime, readCachedCheckouts, - waitForTerminalsReachable, + announceWakeUsable, type CloudRuntime } from './hosted-workspace-runtime' @@ -95,23 +95,11 @@ export function installHostedWorkspaceService(deps: { } /** The step answers before the window re-pairs; `ready` fires once the runtime record says running. */ - async function announceReadyWhenRunning(workspaceId: string, startedAt: number): Promise { - const deadline = Date.now() + 60_000 - while (Date.now() < deadline && wakes.has(workspaceId)) { - const runtime = findCloudRuntime(deps.userDataPath, workspaceId)?.runtime - if (runtime?.status === 'running') { - // Running is not yet usable: panes rebind to the restarted server's new terminals a moment - // later, and keys typed before that are dropped by design. Ready = the server answers a - // terminal listing through the window's own transport. - if (runtime.runtimeEnvironmentId) { - await waitForTerminalsReachable(deps.userDataPath, runtime.runtimeEnvironmentId, deadline) - } - service.emitWakeProgress({ workspaceId, phase: 'ready', source: 'open', startedAt }) - return - } - await new Promise((resolve) => setTimeout(resolve, 250)) - } - } + const announceReadyWhenRunning = (workspaceId: string, startedAt: number): Promise => + announceWakeUsable(deps.userDataPath, workspaceId, startedAt, { + isCurrent: () => wakes.get(workspaceId)?.startedAt === startedAt, + emit: (event) => service.emitWakeProgress(event) + }) const service: HostedWorkspaceService = { emitWakeProgress(event) { @@ -219,6 +207,11 @@ export function installHostedWorkspaceService(deps: { source: 'restart', startedAt }) + const call = found.runtime.runtimeEnvironmentId + ? createHostedRuntimeEnvironmentCall(deps.userDataPath, found.runtime.runtimeEnvironmentId) + : null + // Offers the host already holds were left for the user; only this restart's are continued. + const offersBeforeDrain = call ? await listRestartOfferIds(call) : null // The scheduler holds the old session's deadline and would warn/roll over the new one; a fresh one starts on sync. deps.stopSession(found.sandboxId) try { @@ -237,14 +230,10 @@ export function installHostedWorkspaceService(deps: { source: 'restart', startedAt }) - if (found.runtime.runtimeEnvironmentId) { + if (call) { // Same continuation as a rollover: chats the restart interrupted carry on by themselves. - const call = createHostedRuntimeEnvironmentCall( - deps.userDataPath, - found.runtime.runtimeEnvironmentId - ) - void autoResumeAgentsAfterRollover(call, startedAt).catch((error: unknown) => - log(`[orca-cloud] auto-resume after restart failed: ${String(error)}`) + void autoResumeAgentsAfterRollover(call, startedAt, offersBeforeDrain).catch( + (error: unknown) => log(`[orca-cloud] auto-resume after restart failed: ${String(error)}`) ) } return { ok: true, value: null } diff --git a/src/main/hosted-sandbox/hosted-workspace-wake-ready.test.ts b/src/main/hosted-sandbox/hosted-workspace-wake-ready.test.ts new file mode 100644 index 00000000000..73a810d260f --- /dev/null +++ b/src/main/hosted-sandbox/hosted-workspace-wake-ready.test.ts @@ -0,0 +1,49 @@ +import { describe, expect, it, vi } from 'vitest' + +const runtimes: { status: string; runtimeEnvironmentId?: string }[] = [] +let listOk = true +vi.mock('../ipc/runtime-environments/runtime-environment-transport-routing', () => ({ + callRuntimeEnvironment: vi.fn(async () => + listOk ? { ok: true, result: {} } : { ok: false, error: { code: 'x', message: 'x' } } + ) +})) +vi.mock('../../shared/ephemeral-vm-runtime-store', () => ({ + listEphemeralVmRuntimes: () => + runtimes.map((r) => ({ + ...r, + workspaceId: 'ws', + recipe: { provider: 'orca-cloud' }, + recipeResult: { + schemaVersion: 1, + pairingCode: 'x', + projectRoot: '/p', + orcaCloud: { sandboxId: 'sbx_1' } + } + })) +})) +vi.mock('../environment-providers/environment-provider-recipe-adapter', () => ({ + environmentHandleFromRecipeResult: () => ({ driverState: { sandboxId: 'sbx_1' } }) +})) + +describe('announceWakeUsable', () => { + it('emits ready once running and reachable', async () => { + const { announceWakeUsable } = await import('./hosted-workspace-runtime') + runtimes.splice(0, runtimes.length, { status: 'running', runtimeEnvironmentId: 'env' }) + listOk = true + const emit = vi.fn() + await announceWakeUsable('/u', 'ws', 1, { isCurrent: () => true, emit }, 2_000) + expect(emit).toHaveBeenCalledWith(expect.objectContaining({ phase: 'ready', startedAt: 1 })) + }) + + it('emits failed when the server never answers, and nothing for a superseded wake', async () => { + const { announceWakeUsable } = await import('./hosted-workspace-runtime') + runtimes.splice(0, runtimes.length, { status: 'running', runtimeEnvironmentId: 'env' }) + listOk = false + const emit = vi.fn() + await announceWakeUsable('/u', 'ws', 2, { isCurrent: () => true, emit }, 600) + expect(emit).toHaveBeenCalledWith(expect.objectContaining({ phase: 'failed' })) + const stale = vi.fn() + await announceWakeUsable('/u', 'ws', 3, { isCurrent: () => false, emit: stale }, 600) + expect(stale).not.toHaveBeenCalled() + }) +}) diff --git a/src/main/ipc/runtime-environments/runtime-environment-recipe-vm-wake.test.ts b/src/main/ipc/runtime-environments/runtime-environment-recipe-vm-wake.test.ts index 009fd9fbe21..8dc5621200a 100644 --- a/src/main/ipc/runtime-environments/runtime-environment-recipe-vm-wake.test.ts +++ b/src/main/ipc/runtime-environments/runtime-environment-recipe-vm-wake.test.ts @@ -13,7 +13,7 @@ import type { RecipeVmOpening, RecipeVmState } from '../../runtime/workspaces/sleep/recipe-vm-operations' -import { prepareRecipeVm } from './runtime-environment-recipe-vm-wake' +import { holdForInFlightRecipeVmWake, prepareRecipeVm } from './runtime-environment-recipe-vm-wake' import { RemoteRuntimeClientError } from '../../../shared/remote-runtime/remote-runtime-client-error' import { forgetRecipeVmWakePairingsForTests, @@ -131,7 +131,7 @@ describe('a request to a recipe VM environment', () => { } ) - it('holds a request that does not open the workspace while an in-flight wake finishes', async () => { + it('holds a request that does not open the workspace for an in-flight wake, outside its call slot', async () => { recordVm('running') ownerState = 'waking' const awaitWake = vi.fn(async () => 'up' as const) @@ -142,31 +142,34 @@ describe('a request to a recipe VM environment', () => { awaitWake, openEnvironment: async () => 'up' }) - noteRecipeVmWakePairing(VM_ENVIRONMENT.id, 7, true) - await expect(prepareRecipeVm(userDataPath, VM_ENVIRONMENT, 'none')).resolves.toEqual({ - wokenPairingRevision: null + await holdForInFlightRecipeVmWake(userDataPath, VM_ENVIRONMENT, 'terminal.list', 5_000) + expect(awaitWake).toHaveBeenCalledWith('workspace-of-vm-1', false, 5_000, undefined) + // Bounded by the caller's hold (0 = never wait) and never for an open (that wakes in its slot). + await holdForInFlightRecipeVmWake(userDataPath, VM_ENVIRONMENT, 'terminal.list', 5_000, { + preparation: { recipeVmHoldMs: 0 } }) - expect(awaitWake).toHaveBeenCalledWith('workspace-of-vm-1', false, 60_000, undefined) - // It never starts a wake of its own. + await holdForInFlightRecipeVmWake(userDataPath, VM_ENVIRONMENT, 'worktree.activate', 5_000) + expect(awaitWake).toHaveBeenCalledTimes(1) expect(holdOpen).not.toHaveBeenCalled() + // Inside the slot the request is still answered at once from the owner's state. + await expect(codeOf('none')).resolves.toBe('runtime_environment_waking') }) - it.each([ - ['waking', 'runtime_environment_waking'], - ['wake-failed', 'runtime_environment_wake_failed'], - ['slept', 'runtime_environment_asleep'], - ['asleep', 'runtime_environment_asleep'] - ] as const)('refuses with %s when the in-flight wake ends that way', async (outcome, code) => { + it('caps the hold at the in-flight maximum and does not wait on a VM that is not waking', async () => { recordVm('running') - ownerState = 'waking' + const awaitWake = vi.fn(async () => 'up' as const) registerRecipeVmWaker({ state, noteOpened: vi.fn(), holdOpen, - awaitWake: async () => outcome, + awaitWake, openEnvironment: async () => 'up' }) - await expect(codeOf('none')).resolves.toBe(code) + await holdForInFlightRecipeVmWake(userDataPath, VM_ENVIRONMENT, 'terminal.list', 600_000) + expect(awaitWake).not.toHaveBeenCalled() + ownerState = 'waking' + await holdForInFlightRecipeVmWake(userDataPath, VM_ENVIRONMENT, 'terminal.list', 600_000) + expect(awaitWake).toHaveBeenCalledWith('workspace-of-vm-1', false, 60_000, undefined) }) // A phone gives up after 30 seconds, so its open is held for less. diff --git a/src/main/ipc/runtime-environments/runtime-environment-recipe-vm-wake.ts b/src/main/ipc/runtime-environments/runtime-environment-recipe-vm-wake.ts index e1dd906d203..336a08849f9 100644 --- a/src/main/ipc/runtime-environments/runtime-environment-recipe-vm-wake.ts +++ b/src/main/ipc/runtime-environments/runtime-environment-recipe-vm-wake.ts @@ -14,6 +14,7 @@ import { RUNTIME_ENVIRONMENT_WAKE_FAILED_MESSAGE, RUNTIME_ENVIRONMENT_WAKING_ERROR_CODE, RUNTIME_ENVIRONMENT_WAKING_MESSAGE, + recipeVmWakeIntentFor, type RecipeVmWakeIntent } from '../../../shared/remote-runtime/recipe-vm-wake' import { getRecipeVmWaker } from '../../runtime/workspaces/sleep/recipe-vm-waker-registry' @@ -111,22 +112,6 @@ export async function prepareRecipeVm( return { wokenPairingRevision: woken !== wakeBefore ? (woken?.pairingRevision ?? null) : null } } const state = waker.state(vm.workspaceId, vm.recordSuspended) - if (state === 'waking' && waker.awaitWake) { - // The wake is already underway (often just reconnecting): a terminal or panel opened in that - // window waits for it instead of failing and never retrying. Nothing here starts a wake. - const wakeBefore = latestRecipeVmWake(environment.id) - const outcome = await waker.awaitWake( - vm.workspaceId, - vm.recordSuspended, - Math.min(holdMs, RECIPE_VM_IN_FLIGHT_WAKE_HOLD_MS), - signal - ) - if (outcome !== 'up') { - throw refusal(outcome) - } - const woken = latestRecipeVmWake(environment.id) - return { wokenPairingRevision: woken !== wakeBefore ? (woken?.pairingRevision ?? null) : null } - } if (state !== 'up') { throw refusal(state) } @@ -136,6 +121,29 @@ export async function prepareRecipeVm( /** How long a request that does not open the workspace waits for a wake already in flight. */ export const RECIPE_VM_IN_FLIGHT_WAKE_HOLD_MS = 60_000 +/** + * Before a request that does not open the workspace takes its call slot: while the VM's wake is + * already underway (often just reconnecting), wait for it, up to `holdMs`, instead of being refused + * and never retried. Never starts a wake and never refuses: the request's own preparation decides + * what happens next (it is refused there if the wake failed or is still running). + */ +export async function awaitRecipeVmInFlightWake( + userDataPath: string, + environment: RecipeVmEnvironment, + holdMs: number, + signal?: AbortSignal +): Promise { + const waker = getRecipeVmWaker() + if (holdMs <= 0 || !waker?.awaitWake) { + return + } + const vm = recipeVmOf(userDataPath, environment) + if (!vm || waker.state(vm.workspaceId, vm.recordSuspended) !== 'waking') { + return + } + await waker.awaitWake(vm.workspaceId, vm.recordSuspended, holdMs, signal).catch(() => undefined) +} + /** The pairing revision a recipe VM wake this request waited for left its server on, if any. */ export type RecipeVmPreparation = { wokenPairingRevision: number | null } const UNCHANGED: RecipeVmPreparation = { wokenPairingRevision: null } @@ -174,3 +182,27 @@ function refusal(outcome: Exclude | 'asleep'): RemoteRunt ) } } + +/** A request that does not open the workspace waits (bounded by its own timeout and hold) for a wake + * already in flight. Called before the request takes a per-environment call slot, so held + * background requests never block an open. */ +export function holdForInFlightRecipeVmWake( + userDataPath: string, + environment: RecipeVmEnvironment, + method: string, + timeoutMs: number | undefined, + options?: { preparation?: { recipeVmHoldMs?: number }; signal?: AbortSignal } +): Promise { + if (recipeVmWakeIntentFor(method) !== 'none') { + return Promise.resolve() + } + const holdMs = Math.min( + RECIPE_VM_IN_FLIGHT_WAKE_HOLD_MS, + options?.preparation?.recipeVmHoldMs ?? Infinity, + timeoutMs ?? DEFAULT_CALL_TIMEOUT_MS + ) + return awaitRecipeVmInFlightWake(userDataPath, environment, holdMs, options?.signal) +} + +/** The transport's default call timeout, which also bounds the hold for calls that set none. */ +const DEFAULT_CALL_TIMEOUT_MS = 15_000 diff --git a/src/main/ipc/runtime-environments/runtime-environment-transport-routing.ts b/src/main/ipc/runtime-environments/runtime-environment-transport-routing.ts index 0a4fcce7984..2ca31d1e8f7 100644 --- a/src/main/ipc/runtime-environments/runtime-environment-transport-routing.ts +++ b/src/main/ipc/runtime-environments/runtime-environment-transport-routing.ts @@ -23,6 +23,7 @@ import type { RemoteRuntimeTransportSubscription } from '../../../shared/remote- import { withRuntimeUnreachableHint } from './runtime-environment-unreachable-hint' import { decodeRuntimeEnvironmentSubscriptionEvents } from './runtime-environment-subscription-events' import { enqueueRuntimeCall } from './runtime-environment-call-queue' +import { holdForInFlightRecipeVmWake } from './runtime-environment-recipe-vm-wake' import { sendRemoteRuntimeConnectionRequestAbortable, sendRemoteRuntimeRequestAbortable @@ -80,6 +81,7 @@ export async function callRuntimeEnvironment( ) } const environment = resolveEnvironment(userDataPath, selector) + await holdForInFlightRecipeVmWake(userDataPath, environment, method, timeoutMs, options) // Connect failures reject; retain the dispatched endpoint for hints because queued re-pairing can change it. let endpoint = getPreferredPairedRuntimeConnection(environment).endpoint try { diff --git a/src/main/startup/main-process-runtime-service.ts b/src/main/startup/main-process-runtime-service.ts index 2fc9e8f28f4..38a5ae32e8a 100644 --- a/src/main/startup/main-process-runtime-service.ts +++ b/src/main/startup/main-process-runtime-service.ts @@ -128,13 +128,17 @@ export function initializeMainProcessRuntime(): OrcaRuntimeService { onRuntimeEnvironmentStatusPublished(trackRuntimeEnvironmentVerifiedTransitions(listener)) }, workspaceSleepPlatform: { + // Sleep's own calls and observations never wait on a wake in flight. callEnvironment: (environmentId, method, params, timeoutMs) => callRuntimeEnvironment( getRuntimeEnvironmentRegistryRoot(), environmentId, method, params, - timeoutMs + timeoutMs, + undefined, + undefined, + { preparation: { recipeVmHoldMs: 0 } } ), ownsVm: (workspaceId) => findWorkspaceEphemeralVmRuntime(app.getPath('userData'), workspaceId) !== undefined, diff --git a/src/renderer/src/components/terminal-pane/pty-connection/startup/cold-restore-continue.test.ts b/src/renderer/src/components/terminal-pane/pty-connection/startup/cold-restore-continue.test.ts index a0487b3e0a2..25f2b850fb7 100644 --- a/src/renderer/src/components/terminal-pane/pty-connection/startup/cold-restore-continue.test.ts +++ b/src/renderer/src/components/terminal-pane/pty-connection/startup/cold-restore-continue.test.ts @@ -3,11 +3,56 @@ import { describe, expect, it, vi } from 'vitest' vi.mock('@/store/index', () => ({ useAppStore: { getState: () => ({}) } })) describe('shouldContinueInterruptedTurn', () => { - it('continues only a remote agent whose last status was working', async () => { + const now = 1_800_000_000_000 + const working = (stateStartedAt: number) => ({ + state: 'working', + updatedAt: now - 1000, + stateStartedAt + }) + + it('continues a remote agent with a fresh working status, once per interrupted turn', async () => { const { shouldContinueInterruptedTurn } = await import('./cold-restore-resume-startup') - expect(shouldContinueInterruptedTurn({ state: 'working' }, 'env-1')).toBe(true) - expect(shouldContinueInterruptedTurn({ state: 'waiting' }, 'env-1')).toBe(false) - expect(shouldContinueInterruptedTurn({ state: 'working' }, null)).toBe(false) - expect(shouldContinueInterruptedTurn(null, 'env-1')).toBe(false) + const ctx = { paneKey: 'p1', hasSleepingRecord: false, now } + expect(shouldContinueInterruptedTurn(working(1), 'env-1', ctx)).toBe(true) + // A retried cold restore of the same turn does not send it again. + expect(shouldContinueInterruptedTurn(working(1), 'env-1', ctx)).toBe(false) + // A later interrupted turn does. + expect(shouldContinueInterruptedTurn(working(2), 'env-1', ctx)).toBe(true) + }) + + it('never continues after a deliberate sleep, for stale or non-working status, or locally', async () => { + const { shouldContinueInterruptedTurn } = await import('./cold-restore-resume-startup') + expect( + shouldContinueInterruptedTurn(working(10), 'env-1', { + paneKey: 'p2', + hasSleepingRecord: true, + now + }) + ).toBe(false) + expect( + shouldContinueInterruptedTurn( + { state: 'working', updatedAt: now - 31 * 60_000, stateStartedAt: 11 }, + 'env-1', + { + paneKey: 'p3', + hasSleepingRecord: false, + now + } + ) + ).toBe(false) + expect( + shouldContinueInterruptedTurn( + { state: 'waiting', updatedAt: now, stateStartedAt: 12 }, + 'env-1', + { paneKey: 'p4', hasSleepingRecord: false, now } + ) + ).toBe(false) + expect( + shouldContinueInterruptedTurn(working(13), null, { + paneKey: 'p5', + hasSleepingRecord: false, + now + }) + ).toBe(false) }) }) diff --git a/src/renderer/src/components/terminal-pane/pty-connection/startup/cold-restore-resume-startup.ts b/src/renderer/src/components/terminal-pane/pty-connection/startup/cold-restore-resume-startup.ts index 7abba5460db..7ab48a380e1 100644 --- a/src/renderer/src/components/terminal-pane/pty-connection/startup/cold-restore-resume-startup.ts +++ b/src/renderer/src/components/terminal-pane/pty-connection/startup/cold-restore-resume-startup.ts @@ -27,6 +27,7 @@ import { normalizeAgentProviderSession } from '../../../../../../shared/agent-session/lifecycle/agent-session-resume' +import { AGENT_STATUS_STALE_AFTER_MS } from '../../../../../../shared/agent-status/contracts/agent-status-freshness' import type { HeldPaneReopenOutcome } from '../../session-restore/newer-daemon/newer-daemon-hold-state' import type { ColdRestoreAgentResumeStartup } from './fresh-spawn-types' @@ -34,13 +35,32 @@ import type { ConnectPanePtySession } from '../connect-pane-pty-session' export const NOT_REOPENED: HeldPaneReopenOutcome = { reopened: false, agentResumed: false } -/** A cold-restored agent on a remote host whose last status was working lost its turn when the - * host restarted; local panes are left alone (their PTYs outlive app restarts). */ +/** Interrupted turns already continued, so a retried cold restore never sends `continue` twice. */ +const continuedInterruptedTurns = new Set() + +/** + * A cold-restored agent on a remote host whose last status was working lost its turn when the host + * restarted (reboot, crash, rollover): carry it on. Only with fresh evidence (a recent working + * status), never after a deliberate sleep (that leaves a sleeping record), once per interrupted turn. + * Local panes are left alone: their PTYs outlive app restarts. + */ export function shouldContinueInterruptedTurn( - entry: { state: string } | null | undefined, - runtimeEnvironmentId: string | null + entry: { state: string; updatedAt: number; stateStartedAt?: number } | null | undefined, + runtimeEnvironmentId: string | null, + context: { paneKey: string; hasSleepingRecord: boolean; now?: number } ): boolean { - return Boolean(runtimeEnvironmentId) && entry?.state === 'working' + if (!runtimeEnvironmentId || context.hasSleepingRecord || entry?.state !== 'working') { + return false + } + if ((context.now ?? Date.now()) - entry.updatedAt > AGENT_STATUS_STALE_AFTER_MS) { + return false + } + const turn = `${context.paneKey}\u0000${entry.stateStartedAt ?? entry.updatedAt}` + if (continuedInterruptedTurns.has(turn)) { + return false + } + continuedInterruptedTurns.add(turn) + return true } type AppState = ReturnType @@ -134,7 +154,8 @@ export function bindBuildColdRestoreAgentResumeStartup(session: ConnectPanePtySe // mid-turn; its last status still says working. Carry the turn on instead of waiting. continuePrompt: shouldContinueInterruptedTurn( entry, - session.transport.getRuntimeEnvironmentId?.() ?? null + session.transport.getRuntimeEnvironmentId?.() ?? null, + { paneKey: session.cacheKey, hasSleepingRecord: Boolean(sleepingRecord) } ) ? INTERRUPTED_TURN_CONTINUE_PROMPT : null From b0041f6de57aa22cedb901d7abc59a1d0fc3fde9 Mon Sep 17 00:00:00 2001 From: Neil Date: Mon, 9 Nov 2026 07:09:53 -0800 Subject: [PATCH 4/4] fix(cloud): held in-flight wake keeps the caller fence (re-pairing wake), interrupted-turn continue spent only on launch, offer snapshot timer cleared --- .../hosted-agent-auto-resume.ts | 7 ++- ...runtime-environment-recipe-vm-wake.test.ts | 18 ++++++++ .../runtime-environment-recipe-vm-wake.ts | 32 ++++++++++---- ...me-environment-request-preparation.test.ts | 43 +++++++++++++++++++ ...runtime-environment-request-preparation.ts | 14 ++++-- .../runtime-environment-transport-routing.ts | 25 +++++------ .../runtime-environment-usage-marking.ts | 15 +++++++ .../deferred-cold-restore-and-snapshot.ts | 2 + .../deferred-session-reattach-connect.ts | 5 +++ .../startup/cold-restore-continue.test.ts | 8 +++- .../startup/cold-restore-resume-startup.ts | 40 +++++++++++------ .../startup/fresh-spawn-types.ts | 2 + 12 files changed, 166 insertions(+), 45 deletions(-) create mode 100644 src/main/ipc/runtime-environments/runtime-environment-usage-marking.ts diff --git a/src/main/hosted-sandbox/hosted-agent-auto-resume.ts b/src/main/hosted-sandbox/hosted-agent-auto-resume.ts index 4eaae15a42d..fced9e50cd1 100644 --- a/src/main/hosted-sandbox/hosted-agent-auto-resume.ts +++ b/src/main/hosted-sandbox/hosted-agent-auto-resume.ts @@ -38,10 +38,13 @@ function hasCode(error: unknown, code: string): boolean { * the host cannot answer (an unresponsive host is often why it is being restarted). */ export async function listRestartOfferIds(call: RuntimeCall): Promise | null> { try { + let timer: ReturnType | undefined const listed = await Promise.race([ call('agentSession.restartResumable'), - new Promise((resolve) => setTimeout(() => resolve(null), OFFER_SNAPSHOT_TIMEOUT_MS)) - ]) + new Promise((resolve) => { + timer = setTimeout(() => resolve(null), OFFER_SNAPSHOT_TIMEOUT_MS) + }) + ]).finally(() => clearTimeout(timer)) const list = OfferList.safeParse(listed) if (!list.success) { return null diff --git a/src/main/ipc/runtime-environments/runtime-environment-recipe-vm-wake.test.ts b/src/main/ipc/runtime-environments/runtime-environment-recipe-vm-wake.test.ts index 8dc5621200a..13f7642f191 100644 --- a/src/main/ipc/runtime-environments/runtime-environment-recipe-vm-wake.test.ts +++ b/src/main/ipc/runtime-environments/runtime-environment-recipe-vm-wake.test.ts @@ -172,6 +172,24 @@ describe('a request to a recipe VM environment', () => { expect(awaitWake).toHaveBeenCalledWith('workspace-of-vm-1', false, 60_000, undefined) }) + it('reports the pairing a wake it waited through left, for the caller fence', async () => { + recordVm('running') + ownerState = 'waking' + registerRecipeVmWaker({ + state, + noteOpened: vi.fn(), + holdOpen, + awaitWake: async () => { + noteRecipeVmWakePairing(VM_ENVIRONMENT.id, 9, true) + return 'up' as const + }, + openEnvironment: async () => 'up' + }) + await expect( + holdForInFlightRecipeVmWake(userDataPath, VM_ENVIRONMENT, 'terminal.list', 5_000) + ).resolves.toEqual({ before: VM_ENVIRONMENT, wokenPairingRevision: 9 }) + }) + // A phone gives up after 30 seconds, so its open is held for less. // A Disconnect made before the open holds: no client's open wakes the VM behind that server. diff --git a/src/main/ipc/runtime-environments/runtime-environment-recipe-vm-wake.ts b/src/main/ipc/runtime-environments/runtime-environment-recipe-vm-wake.ts index 336a08849f9..dda73c5db0c 100644 --- a/src/main/ipc/runtime-environments/runtime-environment-recipe-vm-wake.ts +++ b/src/main/ipc/runtime-environments/runtime-environment-recipe-vm-wake.ts @@ -127,21 +127,35 @@ export const RECIPE_VM_IN_FLIGHT_WAKE_HOLD_MS = 60_000 * and never retried. Never starts a wake and never refuses: the request's own preparation decides * what happens next (it is refused there if the wake failed or is still running). */ -export async function awaitRecipeVmInFlightWake( +export async function awaitRecipeVmInFlightWake( userDataPath: string, - environment: RecipeVmEnvironment, + environment: E, holdMs: number, signal?: AbortSignal -): Promise { +): Promise | null> { const waker = getRecipeVmWaker() if (holdMs <= 0 || !waker?.awaitWake) { - return + return null } const vm = recipeVmOf(userDataPath, environment) if (!vm || waker.state(vm.workspaceId, vm.recordSuspended) !== 'waking') { - return + return null } + const before = latestRecipeVmWake(environment.id) await waker.awaitWake(vm.workspaceId, vm.recordSuspended, holdMs, signal).catch(() => undefined) + const woken = latestRecipeVmWake(environment.id) + // The wake this request waited through re-paired the server: the caller's fence names the pairing + // before it, exactly as for a request that woke the VM itself. + return woken !== before + ? { before: environment, wokenPairingRevision: woken?.pairingRevision ?? null } + : null +} + +/** A wake a request waited through before taking its call slot: the record from before it and + * the pairing it left. */ +export type HeldRecipeVmWake = { + before: E + wokenPairingRevision: number | null } /** The pairing revision a recipe VM wake this request waited for left its server on, if any. */ @@ -186,15 +200,15 @@ function refusal(outcome: Exclude | 'asleep'): RemoteRunt /** A request that does not open the workspace waits (bounded by its own timeout and hold) for a wake * already in flight. Called before the request takes a per-environment call slot, so held * background requests never block an open. */ -export function holdForInFlightRecipeVmWake( +export function holdForInFlightRecipeVmWake( userDataPath: string, - environment: RecipeVmEnvironment, + environment: E, method: string, timeoutMs: number | undefined, options?: { preparation?: { recipeVmHoldMs?: number }; signal?: AbortSignal } -): Promise { +): Promise | null> { if (recipeVmWakeIntentFor(method) !== 'none') { - return Promise.resolve() + return Promise.resolve(null) } const holdMs = Math.min( RECIPE_VM_IN_FLIGHT_WAKE_HOLD_MS, diff --git a/src/main/ipc/runtime-environments/runtime-environment-request-preparation.test.ts b/src/main/ipc/runtime-environments/runtime-environment-request-preparation.test.ts index 18e6360355c..38bc56ae4a8 100644 --- a/src/main/ipc/runtime-environments/runtime-environment-request-preparation.test.ts +++ b/src/main/ipc/runtime-environments/runtime-environment-request-preparation.test.ts @@ -56,3 +56,46 @@ describe('a request whose own recipe VM wake re-paired its server', () => { expect((await fence(1, 'before')).failure).toMatchObject({ ok: false }) }) }) + +describe('a fenced request that waited through a wake already in flight', () => { + beforeEach(() => { + prepare.mockClear() + // By the time the request takes its slot the record already holds the new pairing; its own + // preparation woke nothing. + prepared.current = { environment: AFTER_WAKE, before: AFTER_WAKE, wokenPairingRevision: null } + }) + + const heldFence = (pairingRevision: number, wokenPairingRevision: number | null) => + prepareFencedRuntimeRequest( + '/user-data', + 'env-vm', + 'terminal.list', + undefined, + { pairingRevision, runtimeId: 'before' }, + undefined, + { before: BEFORE as never, wokenPairingRevision } + ) + + it('passes the pre-wake fence and then expects the pairing that wake made', async () => { + await expect(heldFence(1, 2)).resolves.toMatchObject({ + failure: null, + expected: { pairingRevision: 2, runtimeId: undefined } + }) + }) + + it('still refuses a caller fenced to an older pairing', async () => { + expect((await heldFence(0, 2)).failure).toMatchObject({ ok: false }) + expect(prepare).not.toHaveBeenCalled() + }) + + it('without the held wake, the same call is refused as changed', async () => { + const result = await prepareFencedRuntimeRequest( + '/user-data', + 'env-vm', + 'terminal.list', + undefined, + { pairingRevision: 1, runtimeId: 'before' } + ) + expect(result.failure).toMatchObject({ ok: false }) + }) +}) diff --git a/src/main/ipc/runtime-environments/runtime-environment-request-preparation.ts b/src/main/ipc/runtime-environments/runtime-environment-request-preparation.ts index 5c5cc54739a..e204f046db3 100644 --- a/src/main/ipc/runtime-environments/runtime-environment-request-preparation.ts +++ b/src/main/ipc/runtime-environments/runtime-environment-request-preparation.ts @@ -22,14 +22,17 @@ export async function prepareFencedRuntimeRequest( method: string, preparation: ManagedRuntimeCallPreparation | undefined, expected: Expected, - signal?: AbortSignal + signal?: AbortSignal, + /** A recipe VM wake this request waited through before its call slot: its fence is checked + * against the record from before that wake, as for a wake the request made itself. */ + heldWake?: { before: KnownRuntimeEnvironment; wokenPairingRevision: number | null } ): Promise< | { failure: Rpc.RuntimeRpcResponse } | { failure: null; environment: KnownRuntimeEnvironment; expected: Expected } > { // A stale caller is refused before its request can wake (and re-pair) a VM. const stale = runtimeEnvironmentRevisionFailure( - resolveEnvironment(userDataPath, environmentId), + heldWake?.before ?? resolveEnvironment(userDataPath, environmentId), expected.pairingRevision, method, expected.runtimeId @@ -44,7 +47,10 @@ export async function prepareFencedRuntimeRequest( method, signal ) - const { environment, wokenPairingRevision: woken } = prepared + const { environment } = prepared + const woken = prepared.wokenPairingRevision ?? heldWake?.wokenPairingRevision ?? null + const before = + prepared.wokenPairingRevision !== null ? prepared.before : (heldWake?.before ?? prepared.before) const failure = woken === null ? runtimeEnvironmentRevisionFailure( @@ -54,7 +60,7 @@ export async function prepareFencedRuntimeRequest( expected.runtimeId ) : (runtimeEnvironmentRevisionFailure( - prepared.before, + before, expected.pairingRevision, method, expected.runtimeId diff --git a/src/main/ipc/runtime-environments/runtime-environment-transport-routing.ts b/src/main/ipc/runtime-environments/runtime-environment-transport-routing.ts index 2ca31d1e8f7..b3b6bf26455 100644 --- a/src/main/ipc/runtime-environments/runtime-environment-transport-routing.ts +++ b/src/main/ipc/runtime-environments/runtime-environment-transport-routing.ts @@ -15,7 +15,7 @@ import { resolveManagedRuntimeEnvironment, type ManagedRuntimeCallPreparation } from './runtime-environment-managed-tunnel' -import { recordRuntimeEnvironmentUsage } from '../../../shared/remote-runtime/host-store/runtime-environment-usage-record' +import { markEnvironmentUsedFromResponse } from './runtime-environment-usage-marking' import type * as Rpc from '../../../shared/runtime-rpc-envelope' import '../../../shared/remote-runtime/requests/remote-runtime-request-socket' import { subscribeRemoteRuntimeTransport } from '../../../shared/remote-runtime/requests/remote-runtime-subscription-transport' @@ -81,7 +81,13 @@ export async function callRuntimeEnvironment( ) } const environment = resolveEnvironment(userDataPath, selector) - await holdForInFlightRecipeVmWake(userDataPath, environment, method, timeoutMs, options) + const held = await holdForInFlightRecipeVmWake( + userDataPath, + environment, + method, + timeoutMs, + options + ) // Connect failures reject; retain the dispatched endpoint for hints because queued re-pairing can change it. let endpoint = getPreferredPairedRuntimeConnection(environment).endpoint try { @@ -98,7 +104,8 @@ export async function callRuntimeEnvironment( pairingRevision: expectedEnvironmentPairingRevision, runtimeId: options?.expectedEnvironmentRuntimeId }, - options?.signal + options?.signal, + held ?? undefined ) if (prepared.failure) { return prepared.failure @@ -292,18 +299,6 @@ export async function subscribeRuntimeEnvironment( } } -function markEnvironmentUsedFromResponse( - userDataPath: string, - environmentId: string, - response: Rpc.RuntimeRpcResponse -): void { - if (response.ok === true) { - recordRuntimeEnvironmentUsage(userDataPath, environmentId, { - runtimeId: response._meta.runtimeId - }) - } -} - function shouldUseCachedRequestConnection(method: string): boolean { return method === 'terminal.send' || method === 'terminal.updateViewport' } diff --git a/src/main/ipc/runtime-environments/runtime-environment-usage-marking.ts b/src/main/ipc/runtime-environments/runtime-environment-usage-marking.ts new file mode 100644 index 00000000000..83472942edb --- /dev/null +++ b/src/main/ipc/runtime-environments/runtime-environment-usage-marking.ts @@ -0,0 +1,15 @@ +import { recordRuntimeEnvironmentUsage } from '../../../shared/remote-runtime/host-store/runtime-environment-usage-record' +import type * as Rpc from '../../../shared/runtime-rpc-envelope' + +/** A successful answer records which server process served the environment. */ +export function markEnvironmentUsedFromResponse( + userDataPath: string, + environmentId: string, + response: Rpc.RuntimeRpcResponse +): void { + if (response.ok === true) { + recordRuntimeEnvironmentUsage(userDataPath, environmentId, { + runtimeId: response._meta.runtimeId + }) + } +} diff --git a/src/renderer/src/components/terminal-pane/pty-connection/deferred-cold-restore-and-snapshot.ts b/src/renderer/src/components/terminal-pane/pty-connection/deferred-cold-restore-and-snapshot.ts index 0c22d59bde5..d45980ceb60 100644 --- a/src/renderer/src/components/terminal-pane/pty-connection/deferred-cold-restore-and-snapshot.ts +++ b/src/renderer/src/components/terminal-pane/pty-connection/deferred-cold-restore-and-snapshot.ts @@ -1,3 +1,4 @@ +import { markInterruptedTurnContinued } from './startup/cold-restore-resume-startup' import { paneResumeRecord } from './startup/pane-resume-record' import { useAppStore } from '@/store/index' import type { PtyReplayDataMeta } from '../pty-transport/pty-transport-types' @@ -34,6 +35,7 @@ export function bindDeferredColdRestoreAndSnapshot(session: ConnectPanePtySessio if (!startup || startup.resumeSource?.cancelled) { return false } + markInterruptedTurnContinued(startup.interruptedTurnKey) const state = useAppStore.getState() if (!session.pendingStartupCommand && session.shouldDeliverStartupViaTerminalPaste) { session.pendingStartupCommand = { diff --git a/src/renderer/src/components/terminal-pane/pty-connection/reattach/deferred-session-reattach-connect.ts b/src/renderer/src/components/terminal-pane/pty-connection/reattach/deferred-session-reattach-connect.ts index 46f8bb740ad..052a8b77862 100644 --- a/src/renderer/src/components/terminal-pane/pty-connection/reattach/deferred-session-reattach-connect.ts +++ b/src/renderer/src/components/terminal-pane/pty-connection/reattach/deferred-session-reattach-connect.ts @@ -1,5 +1,6 @@ import { NOT_REOPENED, + markInterruptedTurnContinued, reopenPaneAfterMissingReattach } from '../startup/cold-restore-resume-startup' import { writeStartupCwdFallbackNotice } from '../startup/startup-cwd-fallback-notice' @@ -134,6 +135,10 @@ function connectDeferredSessionReattach( ) session.beginReattachLiveDataDeferral(outputCallbacks.generation) session.transportConnectInFlightSince = Date.now() + // The command goes to the server now: a `continue` it carries is spent. + if (coldRestoreStartup?.command) { + markInterruptedTurnContinued(coldRestoreStartup.interruptedTurnKey) + } const reattachPromise = session.transport.connect({ url: '', cols: session.cols, diff --git a/src/renderer/src/components/terminal-pane/pty-connection/startup/cold-restore-continue.test.ts b/src/renderer/src/components/terminal-pane/pty-connection/startup/cold-restore-continue.test.ts index 25f2b850fb7..16356718f15 100644 --- a/src/renderer/src/components/terminal-pane/pty-connection/startup/cold-restore-continue.test.ts +++ b/src/renderer/src/components/terminal-pane/pty-connection/startup/cold-restore-continue.test.ts @@ -11,10 +11,14 @@ describe('shouldContinueInterruptedTurn', () => { }) it('continues a remote agent with a fresh working status, once per interrupted turn', async () => { - const { shouldContinueInterruptedTurn } = await import('./cold-restore-resume-startup') + const { shouldContinueInterruptedTurn, interruptedTurnKey, markInterruptedTurnContinued } = + await import('./cold-restore-resume-startup') const ctx = { paneKey: 'p1', hasSleepingRecord: false, now } expect(shouldContinueInterruptedTurn(working(1), 'env-1', ctx)).toBe(true) - // A retried cold restore of the same turn does not send it again. + // Building (and discarding) the startup again does not spend the turn. + expect(shouldContinueInterruptedTurn(working(1), 'env-1', ctx)).toBe(true) + // Once the resume launched, a retried cold restore of the same turn does not send it again. + markInterruptedTurnContinued(interruptedTurnKey(working(1), 'p1')) expect(shouldContinueInterruptedTurn(working(1), 'env-1', ctx)).toBe(false) // A later interrupted turn does. expect(shouldContinueInterruptedTurn(working(2), 'env-1', ctx)).toBe(true) diff --git a/src/renderer/src/components/terminal-pane/pty-connection/startup/cold-restore-resume-startup.ts b/src/renderer/src/components/terminal-pane/pty-connection/startup/cold-restore-resume-startup.ts index 7ab48a380e1..216ecb72b10 100644 --- a/src/renderer/src/components/terminal-pane/pty-connection/startup/cold-restore-resume-startup.ts +++ b/src/renderer/src/components/terminal-pane/pty-connection/startup/cold-restore-resume-startup.ts @@ -38,6 +38,22 @@ export const NOT_REOPENED: HeldPaneReopenOutcome = { reopened: false, agentResum /** Interrupted turns already continued, so a retried cold restore never sends `continue` twice. */ const continuedInterruptedTurns = new Set() +/** Identity of the turn an agent was in when its host went away. */ +export function interruptedTurnKey( + entry: { updatedAt: number; stateStartedAt?: number }, + paneKey: string +): string { + return `${paneKey}\u0000${entry.stateStartedAt ?? entry.updatedAt}` +} + +/** Called when a resume that carries `continue` actually launches (startups are often built and + * discarded, e.g. on every reattach check). */ +export function markInterruptedTurnContinued(key: string | undefined): void { + if (key) { + continuedInterruptedTurns.add(key) + } +} + /** * A cold-restored agent on a remote host whose last status was working lost its turn when the host * restarted (reboot, crash, rollover): carry it on. Only with fresh evidence (a recent working @@ -55,12 +71,7 @@ export function shouldContinueInterruptedTurn( if ((context.now ?? Date.now()) - entry.updatedAt > AGENT_STATUS_STALE_AFTER_MS) { return false } - const turn = `${context.paneKey}\u0000${entry.stateStartedAt ?? entry.updatedAt}` - if (continuedInterruptedTurns.has(turn)) { - return false - } - continuedInterruptedTurns.add(turn) - return true + return !continuedInterruptedTurns.has(interruptedTurnKey(entry, context.paneKey)) } type AppState = ReturnType @@ -141,6 +152,12 @@ export function bindBuildColdRestoreAgentResumeStartup(session: ConnectPanePtySe session.executionHostId ) }) + const continueTurn = + entry !== null && + shouldContinueInterruptedTurn(entry, session.transport.getRuntimeEnvironmentId?.() ?? null, { + paneKey: session.cacheKey, + hasSleepingRecord: Boolean(sleepingRecord) + }) const startupPlan = buildAgentResumeStartupPlan({ agent, providerSession, @@ -152,13 +169,7 @@ export function bindBuildColdRestoreAgentResumeStartup(session: ConnectPanePtySe launchCwd: session.worktree?.path, // A remote host that restarted under a working agent (reboot, crash, VM rollover) killed it // mid-turn; its last status still says working. Carry the turn on instead of waiting. - continuePrompt: shouldContinueInterruptedTurn( - entry, - session.transport.getRuntimeEnvironmentId?.() ?? null, - { paneKey: session.cacheKey, hasSleepingRecord: Boolean(sleepingRecord) } - ) - ? INTERRUPTED_TURN_CONTINUE_PROMPT - : null + continuePrompt: continueTurn ? INTERRUPTED_TURN_CONTINUE_PROMPT : null }) if (!startupPlan) { return null @@ -192,6 +203,9 @@ export function bindBuildColdRestoreAgentResumeStartup(session: ConnectPanePtySe : {}) }, launchToken: coldRestoreLaunchToken, + ...(continueTurn && entry + ? { interruptedTurnKey: interruptedTurnKey(entry, session.cacheKey) } + : {}), useLiveEntry: Boolean(useLiveEntry), hasSleepingRecord: Boolean(sleepingRecord), sleepingRecordEntry diff --git a/src/renderer/src/components/terminal-pane/pty-connection/startup/fresh-spawn-types.ts b/src/renderer/src/components/terminal-pane/pty-connection/startup/fresh-spawn-types.ts index d3071dcc0a5..31926f92f3d 100644 --- a/src/renderer/src/components/terminal-pane/pty-connection/startup/fresh-spawn-types.ts +++ b/src/renderer/src/components/terminal-pane/pty-connection/startup/fresh-spawn-types.ts @@ -20,6 +20,8 @@ export type FreshSpawnOptions = { } export type ColdRestoreAgentResumeStartup = PendingStartupCommand & { + /** Set when the command carries `continue` for an interrupted turn; marked used on launch. */ + interruptedTurnKey?: string agent: ResumableTuiAgent resumeProviderSession: AgentProviderSessionMetadata connectionId?: string | null