diff --git a/apps/worker/src/run-task/__tests__/run-task.test.ts b/apps/worker/src/run-task/__tests__/run-task.test.ts index 32d7ccec2..7ba74876d 100644 --- a/apps/worker/src/run-task/__tests__/run-task.test.ts +++ b/apps/worker/src/run-task/__tests__/run-task.test.ts @@ -1560,19 +1560,33 @@ describe('runTask', () => { ); }); - it('skips external sleep handoff when the task resolves as failed', async () => { + it.each([ + [ + 'uses the external sleep handoff with a retention deadline', + Date.now(), + 1, + ], + ['skips the external sleep handoff without a retention deadline', null, 0], + ])('%s when the task resolves as failed', async (_name, sleepAt, calls) => { resolveStatusMock.mockReturnValueOnce({ status: RunStatus.Failed }); - waitForShutdownMock.mockResolvedValueOnce({ - sessionId: 'failed-session', - cancelTriggeredAt: undefined, - lastMessageAt: Date.now(), - taskFinishedAt: undefined, - taskAbortedAt: undefined, - lastErrorMessage: - 'The provider returned an error: Input exceeds context window.', + waitForExternalSleepActionMock.mockResolvedValueOnce({ + claimed: true, + completed: true, + }); + waitForShutdownMock.mockImplementationOnce(async () => { + harnessManagerInstances.at(-1)!.currentSleepAt = sleepAt; + return { + sessionId: 'failed-session', + cancelTriggeredAt: undefined, + lastMessageAt: Date.now(), + taskFinishedAt: undefined, + taskAbortedAt: undefined, + lastErrorMessage: + 'The provider returned an error: Input exceeds context window.', + }; }); - await runTask({ + const result = await runTask({ taskRun: { id: 112, taskId: 'task-112', @@ -1607,7 +1621,10 @@ describe('runTask', () => { } as never, }); - expect(waitForExternalSleepActionMock).not.toHaveBeenCalled(); + expect(waitForExternalSleepActionMock).toHaveBeenCalledTimes(calls); + expect(result.status).toBe( + calls === 1 ? RunStatus.Completed : RunStatus.Failed, + ); }); it('checks Slack drain during sleep fallback when the worker started before thread ts was persisted', async () => { diff --git a/apps/worker/src/run-task/run-task.ts b/apps/worker/src/run-task/run-task.ts index 743f633d1..1e0e240e1 100644 --- a/apps/worker/src/run-task/run-task.ts +++ b/apps/worker/src/run-task/run-task.ts @@ -2085,6 +2085,8 @@ export const runTask = async ({ // // Terminal cancel must skip this handoff: publish no due sleepAt and // finish as Canceled instead of becoming a snapshot/standby candidate. + // Failed turns still use the ordinary idle retention path so a model or + // provider error cannot discard an otherwise healthy workspace. // // NOTE: Snapshot creation ultimately tears down the provider runtime. Vercel // does this as part of snapshot creation, while Modal explicitly terminates @@ -2097,25 +2099,27 @@ export const runTask = async ({ // messages). The drain check below is kept as a fallback for edge cases where // the snapshot fails or times out but the worker survives. const skipSleepAfterTerminalCancel = Boolean(finalState.cancelTriggeredAt); - const skipSleepAfterTerminalFailure = - resolvedResult.status === RunStatus.Failed; + const skipSleepWithoutRetentionDeadline = + resolvedResult.status === RunStatus.Failed && + harnessManager.getSleepAt() == null; if (skipSleepAfterTerminalCancel) { logger.info( `[runTask] Skipping external sleep handoff after terminal cancel for task run ${taskRun.id}`, ); } - if (skipSleepAfterTerminalFailure) { + if (skipSleepWithoutRetentionDeadline && !skipSleepAfterTerminalCancel) { logger.info( - `[runTask] Skipping external sleep handoff after terminal failure for task run ${taskRun.id}`, + `[runTask] Skipping external sleep handoff without a retention deadline for task run ${taskRun.id}`, ); } let sleepActionTriggered = false; + let sleepActionCompleted = false; if ( !skipExternalSleepAction && !skipSleepAfterTerminalCancel && - !skipSleepAfterTerminalFailure + !skipSleepWithoutRetentionDeadline ) { // BullMQ may claim the sleep action and snapshot the filesystem while // the handoff helper polls below. The harness has already shut down, so @@ -2123,10 +2127,11 @@ export const runTask = async ({ // dequeue response. await scrubSandboxSecretsBeforeSnapshot(logger, { homeDir, runtimeEnv }); - ({ claimed: sleepActionTriggered } = await waitForExternalSleepAction({ - taskRun, - logger, - })); + ({ claimed: sleepActionTriggered, completed: sleepActionCompleted } = + await waitForExternalSleepAction({ + taskRun, + logger, + })); } // Fallback: check for pending Linear messages that arrived during the snapshot @@ -2209,7 +2214,12 @@ export const runTask = async ({ } else { taskCancellation.abortController.abort(); } - return resolvedResult; + // BullMQ owns terminal state after a completed snapshot/standby handoff. + // Do not let a stale pre-handoff failure overwrite that completion if the + // provider leaves this worker alive long enough to return normally. + return sleepActionCompleted + ? { status: RunStatus.Completed } + : resolvedResult; } finally { activeWorkerCrashContext = null; } diff --git a/apps/worker/src/sandbox-server/lib/__tests__/harness-manager.test.ts b/apps/worker/src/sandbox-server/lib/__tests__/harness-manager.test.ts index b73c985fc..113af6d97 100644 --- a/apps/worker/src/sandbox-server/lib/__tests__/harness-manager.test.ts +++ b/apps/worker/src/sandbox-server/lib/__tests__/harness-manager.test.ts @@ -2805,8 +2805,9 @@ describe('HarnessManager error status', () => { } }); - it('treats a provider-error abort as a failed shutdown', async () => { - const { harness, manager } = createManager(); + it('keeps a terminal provider-error abort open for follow-up until normal idle shutdown', async () => { + vi.useFakeTimers(); + const { harness, manager } = createManager({ keepaliveMs: 60_000 }); try { manager.initializeWithoutPrompt(); @@ -2835,14 +2836,27 @@ describe('HarnessManager error status', () => { payload: ['task-provider-error'], } as TaskEvent); + expect(manager.getStatus()).toMatchObject({ + phase: 'waiting_for_prompt', + lastErrorMessage: 'The provider returned an error: API key is invalid.', + }); + expect(manager.getState().taskAbortedAt).toBeUndefined(); + expect(manager.getSleepAt()).not.toBeNull(); + + vi.advanceTimersByTime(59_000); + expect(manager.getStatus().phase).toBe('waiting_for_prompt'); + + vi.advanceTimersByTime(1_000); await expect(manager.waitForShutdown()).resolves.toMatchObject({ lastErrorMessage: 'The provider returned an error: API key is invalid.', taskAbortedAt: undefined, }); - expect(manager.getSleepAt()).toBeNull(); + expect(manager.getStatus().phase).toBe('shutting_down'); + expect(manager.getSleepAt()).not.toBeNull(); } finally { manager.dispose(); harness.dispose(); + vi.useRealTimers(); } }); @@ -2887,7 +2901,7 @@ describe('HarnessManager error status', () => { } }); - it('clears lastErrorMessage when a follow-up prompt is sent successfully', () => { + it('accepts a follow-up after a terminal provider error', () => { const { harness, manager } = createManager(); try { @@ -2908,15 +2922,21 @@ describe('HarnessManager error status', () => { message: { ts: Date.now(), type: 'say', - say: 'error', - text: 'temporary glitch', + say: 'terminal_provider_error', + text: 'The provider returned an error: API key is invalid.', }, }, ], } as TaskEvent); + harness.emitTaskEvent({ + eventName: TaskEventName.TaskAborted, + payload: ['task-error-clear'], + } as TaskEvent); + expect(manager.getStatus().phase).toBe('waiting_for_prompt'); expect(manager.getStatus().lastErrorMessage).toBeDefined(); - manager.sendFollowUpPrompt({ prompt: 'retry' }); + expect(manager.sendFollowUpPrompt({ prompt: 'retry' })).toBe(true); + expect(manager.getStatus().phase).toBe('running'); expect(manager.getStatus().lastErrorMessage).toBeUndefined(); } finally { manager.dispose(); @@ -2953,6 +2973,7 @@ describe('HarnessManager error status', () => { lastErrorMessage: 'OpenCode session creation did not respond within 90s', }); + expect(manager.getSleepAt()).toBeNull(); } finally { manager.dispose(); harness.dispose(); diff --git a/apps/worker/src/sandbox-server/lib/harness-manager.ts b/apps/worker/src/sandbox-server/lib/harness-manager.ts index 2612f0a51..c12567076 100644 --- a/apps/worker/src/sandbox-server/lib/harness-manager.ts +++ b/apps/worker/src/sandbox-server/lib/harness-manager.ts @@ -722,12 +722,14 @@ export class HarnessManager extends EventEmitter { return null; } - // Failed shutdowns must reach finishRun so channel integrations can report - // the error. Snapshotting would otherwise finalize the run as completed. + // Failures without a usable runtime session still need terminal + // finalization. A provider error from an existing session is a failed turn + // instead, so it remains eligible for the ordinary idle retention path. if ( this.phase === 'shutting_down' && this.state.lastErrorMessage && - !this.state.taskFinishedAt + !this.state.taskFinishedAt && + !this.terminalProviderErrorPending ) { return null; } @@ -1362,16 +1364,13 @@ export class HarnessManager extends EventEmitter { if (payload[0] === this.state.sessionId) { this.logger.info(`[HarnessManager] Task aborted: ${payload[0]}`); - // A provider error is terminal, unlike a user-initiated abort. Preserve - // the error and shut down without setting the cancellation stamp so the - // worker resolves this run as Failed rather than Canceled. - if (this.terminalProviderErrorPending && !this.state.cancelTriggeredAt) { - this.triggerShutdown(); - return; + // A terminal provider error ends the current model turn, not the live + // task session. Keep its error visible and leave the session available + // for follow-ups until the ordinary idle keepalive expires. + if (!this.terminalProviderErrorPending || this.state.cancelTriggeredAt) { + this.state.taskAbortedAt = Date.now(); } - this.state.taskAbortedAt = Date.now(); - if (this.runtimeQueuedMessagesCount > 0) { // Abort takes priority: never downgrade a deferred abort to completion. this.deferredTurnSettlement = HarnessEvent.TaskAborted;