From ac684840595c89cfb78df7ee9f3ed3cac8838d56 Mon Sep 17 00:00:00 2001 From: Gabriel Brown Date: Fri, 10 Jul 2026 17:52:20 -0400 Subject: [PATCH] fix(worker): make cancellation teardown failure-safe --- apps/agent-worker/src/codex-process.ts | 88 +++++ apps/agent-worker/src/heartbeat-controller.ts | 42 ++ apps/agent-worker/src/worker.ts | 363 ++++++++++-------- apps/agent-worker/src/workspace-teardown.ts | 35 ++ .../tests/unit/codex-process.test.ts | 37 ++ .../tests/unit/heartbeat-controller.test.ts | 60 +++ .../tests/unit/workspace-teardown.test.ts | 62 +++ .../tests/unit/recover-stale-jobs.test.ts | 79 +++- 8 files changed, 604 insertions(+), 162 deletions(-) create mode 100644 apps/agent-worker/src/codex-process.ts create mode 100644 apps/agent-worker/src/heartbeat-controller.ts create mode 100644 apps/agent-worker/src/workspace-teardown.ts create mode 100644 apps/agent-worker/tests/unit/codex-process.test.ts create mode 100644 apps/agent-worker/tests/unit/heartbeat-controller.test.ts create mode 100644 apps/agent-worker/tests/unit/workspace-teardown.test.ts diff --git a/apps/agent-worker/src/codex-process.ts b/apps/agent-worker/src/codex-process.ts new file mode 100644 index 0000000..c6223a2 --- /dev/null +++ b/apps/agent-worker/src/codex-process.ts @@ -0,0 +1,88 @@ +type CommandResult = { exitCode: number; output: string }; + +const launchScript = ` +import os +import sys + +pid_file, job_id, *command = sys.argv[1:] +if os.getpgrp() != os.getpid(): + os.setsid() +os.environ["SPOON_CODEX_JOB_ID"] = job_id +with open(pid_file, "w", encoding="utf-8") as handle: + handle.write(str(os.getpid())) + handle.flush() + os.fsync(handle.fileno()) +os.execvpe(command[0], command, os.environ) +`; + +const abortScript = ` +import os +import signal +import sys +import time + +pid_file, job_id = sys.argv[1:] +pid = None +for _ in range(20): + try: + with open(pid_file, encoding="utf-8") as handle: + pid = int(handle.read().strip()) + break + except FileNotFoundError: + time.sleep(0.1) +if pid is None: + sys.exit(0) + +try: + with open(f"/proc/{pid}/environ", "rb") as handle: + environment = handle.read().split(b"\\0") + marker = f"SPOON_CODEX_JOB_ID={job_id}".encode() + if marker not in environment or os.getpgid(pid) != pid: + raise RuntimeError("Codex process identity did not match the requested job") + os.killpg(pid, signal.SIGTERM) + for _ in range(20): + time.sleep(0.1) + try: + os.killpg(pid, 0) + except ProcessLookupError: + break + else: + os.killpg(pid, signal.SIGKILL) +except ProcessLookupError: + pass +finally: + try: + os.unlink(pid_file) + except FileNotFoundError: + pass +`; + +export const managedCodexCommand = (args: { + jobId: string; + pidFile: string; + command: string[]; +}) => [ + 'python3', + '-c', + launchScript, + args.pidFile, + args.jobId, + ...args.command, +]; + +export const abortManagedCodexProcess = async (args: { + jobId: string; + pidFile: string; + execute: (command: string[]) => Promise; +}) => { + const result = await args.execute([ + 'python3', + '-c', + abortScript, + args.pidFile, + args.jobId, + ]); + if (result.exitCode !== 0) { + throw new Error(`Failed to abort Codex process: ${result.output}`); + } +}; diff --git a/apps/agent-worker/src/heartbeat-controller.ts b/apps/agent-worker/src/heartbeat-controller.ts new file mode 100644 index 0000000..c6e8380 --- /dev/null +++ b/apps/agent-worker/src/heartbeat-controller.ts @@ -0,0 +1,42 @@ +export const createHeartbeatController = (args: { + intervalMs: number; + heartbeat: (jobId: string) => Promise<{ cancelRequested: boolean }>; + onCancel: (jobId: string) => Promise; + onError: (error: unknown) => void; +}) => { + const timers = new Map>(); + const inFlight = new Set(); + + const stop = (jobId: string) => { + const timer = timers.get(jobId); + if (timer) clearInterval(timer); + timers.delete(jobId); + }; + + const tick = async (jobId: string) => { + if (inFlight.has(jobId) || !timers.has(jobId)) return; + inFlight.add(jobId); + try { + const result = await args.heartbeat(jobId); + if (result.cancelRequested) { + stop(jobId); + await args.onCancel(jobId); + } + } catch (error) { + args.onError(error); + } finally { + inFlight.delete(jobId); + } + }; + + return { + start: (jobId: string) => { + stop(jobId); + const timer = setInterval(() => void tick(jobId), args.intervalMs); + timer.unref(); + timers.set(jobId, timer); + }, + stop, + has: (jobId: string) => timers.has(jobId), + }; +}; diff --git a/apps/agent-worker/src/worker.ts b/apps/agent-worker/src/worker.ts index 9715b5a..4f3b572 100644 --- a/apps/agent-worker/src/worker.ts +++ b/apps/agent-worker/src/worker.ts @@ -18,6 +18,7 @@ import type { NormalizedAgentEvent } from './agent-events'; import type { OpenCodeSession } from './opencode-session'; import type { BoxHandle } from './user-container'; import { normalizeCodexJsonLine } from './agent-events'; +import { abortManagedCodexProcess, managedCodexCommand } from './codex-process'; import { prepareCodexWorkspaceFiles } from './codex-runtime'; import { env } from './env'; import { @@ -28,6 +29,7 @@ import { run, } from './git'; import { getInstallationToken, openDraftPullRequest } from './github'; +import { createHeartbeatController } from './heartbeat-controller'; import { abortOpenCodeSession, createOpenCodeSession, @@ -45,6 +47,7 @@ import { } from './runtime/docker'; import { acquireUserBox, runningBoxUsernames } from './user-container'; import { fetchUserEnvironment, materializeUserHome } from './user-environment'; +import { teardownWorkspace } from './workspace-teardown'; type Claim = { job: { @@ -118,7 +121,9 @@ type ActiveWorkspace = { opencodePassword?: string; opencodeSession?: OpenCodeSession; codexSessionId?: string; + codexPidFile?: string; agentTurnActive?: boolean; + cancelRequested?: boolean; resolveTurn?: () => void; lastRecordedDiffSignature?: string; // Captures the most recent Codex `error`/`turn.failed` event for the active @@ -269,43 +274,17 @@ const heartbeatWorkspaceMutation = async (jobId: Id<'agentJobs'>) => jobId, })) as { success: true; cancelRequested: boolean }; -const heartbeatTimers = new Map>(); +const heartbeatController = createHeartbeatController({ + intervalMs: 30_000, + heartbeat: async (jobId) => + await heartbeatWorkspaceMutation(jobId as Id<'agentJobs'>), + onCancel: async (jobId) => await cancelWorkspace(jobId), + onError: (error) => console.error(error), +}); -const stopHeartbeat = (jobId: string) => { - const timer = heartbeatTimers.get(jobId); - if (timer) clearInterval(timer); - heartbeatTimers.delete(jobId); -}; - -const startHeartbeat = (jobId: Id<'agentJobs'>) => { - stopHeartbeat(jobId); - const timer = setInterval(() => { - void (async () => { - try { - const result = await heartbeatWorkspaceMutation(jobId); - if (result.cancelRequested && activeWorkspaces.has(jobId)) { - await appendEvent( - jobId, - 'warn', - 'cleanup', - 'Cancellation requested; stopping workspace.', - ); - await abortWorkspaceAgent(jobId).catch((error: unknown) => { - console.error(error); - }); - await stopWorkspace(jobId).catch((error: unknown) => { - console.error(error); - }); - stopHeartbeat(jobId); - } - } catch (error) { - console.error(error); - } - })(); - }, 30_000); - timer.unref(); - heartbeatTimers.set(jobId, timer); -}; +const stopHeartbeat = (jobId: string) => heartbeatController.stop(jobId); +const startHeartbeat = (jobId: Id<'agentJobs'>) => + heartbeatController.start(jobId); const appendMessage = async (args: { jobId: Id<'agentJobs'>; @@ -789,7 +768,10 @@ const runCodexTurn = async (args: { '.codex', outputFileName, ); - const command = workspace.codexSessionId + if (workspace.cancelRequested) { + throw new Error('Workspace cancellation requested.'); + } + const codexCommand = workspace.codexSessionId ? [ 'codex', 'exec', @@ -814,46 +796,65 @@ const runCodexTurn = async (args: { workspace.containerRepo, prompt, ]; + const pidFileName = `codex-turn-${workspace.claim.job._id}.pid`; + const pidFileHostPath = path.join(workspace.workdir, '.codex', pidFileName); + const pidFileContainerPath = path.posix.join( + workspace.containerHome, + '.codex', + pidFileName, + ); + workspace.codexPidFile = pidFileContainerPath; + const command = managedCodexCommand({ + jobId: workspace.claim.job._id, + pidFile: pidFileContainerPath, + command: codexCommand, + }); const aiEnv = providerEnvironment(workspace.claim, workspace.containerHome); const secretEnv = Object.fromEntries( workspace.claim.secrets.map((secret) => [secret.name, secret.value]), ); - const result = await streamExecInContainer({ - containerName: workspace.boxName, - containerCwd: workspace.containerRepo, - command, - environment: { - ...aiEnv, - ...secretEnv, - }, - redact: workspace.redact, - timeoutMs: env.jobTimeoutMs, - onStdoutLine: async (line) => { - for (const event of normalizeCodexJsonLine(line)) { - await handleAgentEvent({ - workspace, - event, - assistantMessageId, - assistantContent, - }); - } - }, - onStderrLine: async (line) => { - const trimmed = line.trim(); - if ( - trimmed && - trimmed !== 'Reading additional input from stdin...' && - !trimmed.includes('`[features].codex_hooks` is deprecated') - ) { - await appendEvent( - workspace.claim.job._id, - 'info', - 'plan', - truncate(trimmed, 10_000), - ); - } - }, - }); + let result; + try { + result = await streamExecInContainer({ + containerName: workspace.boxName, + containerCwd: workspace.containerRepo, + command, + environment: { + ...aiEnv, + ...secretEnv, + }, + redact: workspace.redact, + timeoutMs: env.jobTimeoutMs, + onStdoutLine: async (line) => { + for (const event of normalizeCodexJsonLine(line)) { + await handleAgentEvent({ + workspace, + event, + assistantMessageId, + assistantContent, + }); + } + }, + onStderrLine: async (line) => { + const trimmed = line.trim(); + if ( + trimmed && + trimmed !== 'Reading additional input from stdin...' && + !trimmed.includes('`[features].codex_hooks` is deprecated') + ) { + await appendEvent( + workspace.claim.job._id, + 'info', + 'plan', + truncate(trimmed, 10_000), + ); + } + }, + }); + } finally { + workspace.codexPidFile = undefined; + await rm(pidFileHostPath, { force: true }); + } await appendEvent( workspace.claim.job._id, 'info', @@ -1485,32 +1486,46 @@ const runClaim = async (claim: Claim) => { }); } catch (error) { const message = error instanceof Error ? error.message : String(error); - await appendEvent( - jobId, - 'error', - 'cleanup', - truncate(redact(message), 20_000), - ); - await addArtifact({ - jobId, - kind: 'error', - title: 'Failure', - content: truncate(redact(message), 50_000), - contentType: 'text/plain', - }); - await updateStatus( - jobId, - message.toLowerCase().includes('timed out') ? 'timed_out' : 'failed', - { error: truncate(redact(message), 10_000) }, - ); - await markWorkspaceStopped( - jobId, - message.toLowerCase().includes('timed out') ? 'expired' : 'failed', - ).catch((stopError: unknown) => { - console.error(stopError); - }); - stopHeartbeat(jobId); - acquiredBoxHandle?.release(); + const workspaceStatus = message.toLowerCase().includes('timed out') + ? 'expired' + : 'failed'; + try { + await appendEvent( + jobId, + 'error', + 'cleanup', + truncate(redact(message), 20_000), + ); + await addArtifact({ + jobId, + kind: 'error', + title: 'Failure', + content: truncate(redact(message), 50_000), + contentType: 'text/plain', + }); + await updateStatus( + jobId, + workspaceStatus === 'expired' ? 'timed_out' : 'failed', + { error: truncate(redact(message), 10_000) }, + ); + } finally { + const workspace = activeWorkspaces.get(jobId); + if (workspace) { + const teardown = await teardownActiveWorkspace({ + jobId, + workspace, + abortAgent: true, + workspaceStatus, + }); + reportTeardownErrors(teardown.errors); + } else { + stopHeartbeat(jobId); + await markWorkspaceStopped(jobId, workspaceStatus).catch( + (stopError: unknown) => console.error(stopError), + ); + acquiredBoxHandle?.release(); + } + } } }; @@ -1638,27 +1653,39 @@ export const getWorkspaceAgentStatus = (jobId: string) => { }; }; -export const abortWorkspaceAgent = async (jobId: string) => { - const workspace = resolveWorkspace(jobId); +const abortWorkspaceAgentForWorkspace = async (workspace: ActiveWorkspace) => { if (workspace.opencodeSession) { await abortOpenCodeSession(workspace.opencodeSession); - workspace.agentTurnActive = false; - workspace.resolveTurn?.(); - workspace.resolveTurn = undefined; - await appendEvent( - workspace.claim.job._id, - 'warn', - 'cleanup', - 'Agent turn aborted.', - ); - return { success: true }; - } - if (workspace.runtimeMode === 'codex_exec') { - throw new Error('Codex agent turns cannot be aborted from Spoon yet.'); + } else if (workspace.runtimeMode === 'codex_exec' && workspace.codexPidFile) { + await abortManagedCodexProcess({ + jobId: workspace.claim.job._id, + pidFile: workspace.codexPidFile, + execute: async (command) => + await runExecInContainer({ + containerName: workspace.boxName, + containerCwd: workspace.containerRepo, + command, + environment: {}, + redact: workspace.redact, + timeoutMs: env.jobTimeoutMs, + }), + }); } + workspace.agentTurnActive = false; + workspace.resolveTurn?.(); + workspace.resolveTurn = undefined; + await appendEvent( + workspace.claim.job._id, + 'warn', + 'cleanup', + 'Agent turn aborted.', + ); return { success: true }; }; +export const abortWorkspaceAgent = async (jobId: string) => + await abortWorkspaceAgentForWorkspace(resolveWorkspace(jobId)); + export const replyToInteraction = async ( jobId: string, args: { @@ -1869,6 +1896,60 @@ export const sendWorkspaceMessage = async ( } }; +const teardownActiveWorkspace = async (args: { + jobId: string; + workspace: ActiveWorkspace; + abortAgent: boolean; + appendCancellationEvent?: boolean; + workspaceStatus?: 'stopped' | 'expired' | 'failed'; +}) => { + const { jobId, workspace } = args; + const containerName = workspace.containerName; + if (args.abortAgent) workspace.cancelRequested = true; + return await teardownWorkspace({ + stopHeartbeat: () => stopHeartbeat(jobId), + removeActive: () => activeWorkspaces.delete(jobId), + appendEvent: args.appendCancellationEvent + ? async () => + await appendEvent( + workspace.claim.job._id, + 'warn', + 'cleanup', + 'Cancellation requested; stopping workspace.', + ) + : undefined, + abortAgent: args.abortAgent + ? async () => await abortWorkspaceAgentForWorkspace(workspace) + : undefined, + markStopped: async () => + await markWorkspaceStopped(workspace.claim.job._id, args.workspaceStatus), + closeAgent: () => workspace.opencodeSession?.close(), + stopContainer: containerName + ? async () => await stopWorkspaceContainer(containerName) + : undefined, + release: () => workspace.boxHandle.release(), + }); +}; + +const reportTeardownErrors = (errors: unknown[]) => { + for (const error of errors) console.error(error); +}; + +const cancelWorkspace = async (jobId: string) => { + const workspace = activeWorkspaces.get(jobId); + if (!workspace) { + stopHeartbeat(jobId); + return; + } + const result = await teardownActiveWorkspace({ + jobId, + workspace, + abortAgent: true, + appendCancellationEvent: true, + }); + reportTeardownErrors(result.errors); +}; + export const openWorkspacePullRequest = async (jobId: string) => { const workspace = resolveWorkspace(jobId); const { claim, repoDir, redact } = workspace; @@ -1931,16 +2012,14 @@ export const openWorkspacePullRequest = async (jobId: string) => { pullRequestNumber: pullRequest.number, summary: 'Draft PR opened from interactive workspace.', }); - await markWorkspaceStopped(claim.job._id); - workspace.opencodeSession?.close(); - if (workspace.containerName) { - await stopWorkspaceContainer(workspace.containerName); + const teardown = await teardownActiveWorkspace({ + jobId, + workspace, + abortAgent: false, + }); + if (teardown.errors.length > 0) { + throw new AggregateError(teardown.errors, 'Workspace teardown failed.'); } - stopHeartbeat(jobId); - activeWorkspaces.delete(jobId); - // The persistent per-user home + ~/Code checkouts survive across sessions; - // release the box (reaped once no other thread/terminal holds it). - workspace.boxHandle.release(); return { pullRequestUrl: pullRequest.html_url, pullRequestNumber: pullRequest.number, @@ -1949,37 +2028,13 @@ export const openWorkspacePullRequest = async (jobId: string) => { export const stopWorkspace = async (jobId: string) => { const workspace = resolveWorkspace(jobId); - // Stop the heartbeat timer before teardown so it never outlives the - // workspace even if a teardown step throws. - stopHeartbeat(jobId); - const errors: unknown[] = []; - try { - await markWorkspaceStopped(workspace.claim.job._id); - } catch (error) { - errors.push(error); - } - try { - workspace.opencodeSession?.close(); - } catch (error) { - errors.push(error); - } - try { - if (workspace.containerName) { - await stopWorkspaceContainer(workspace.containerName); - } - } catch (error) { - errors.push(error); - } finally { - activeWorkspaces.delete(jobId); - try { - workspace.boxHandle.release(); - } catch (error) { - errors.push(error); - } - } - if (errors.length === 1) throw errors[0]; - if (errors.length > 1) { - throw new AggregateError(errors, 'Workspace teardown failed.'); + const teardown = await teardownActiveWorkspace({ + jobId, + workspace, + abortAgent: true, + }); + if (teardown.errors.length > 0) { + throw new AggregateError(teardown.errors, 'Workspace teardown failed.'); } return { success: true }; }; diff --git a/apps/agent-worker/src/workspace-teardown.ts b/apps/agent-worker/src/workspace-teardown.ts new file mode 100644 index 0000000..9fac3ab --- /dev/null +++ b/apps/agent-worker/src/workspace-teardown.ts @@ -0,0 +1,35 @@ +type TeardownStep = () => unknown; + +export const teardownWorkspace = async (steps: { + stopHeartbeat: TeardownStep; + removeActive: TeardownStep; + appendEvent?: TeardownStep; + abortAgent?: TeardownStep; + markStopped?: TeardownStep; + closeAgent?: TeardownStep; + stopContainer?: TeardownStep; + release: TeardownStep; +}) => { + const errors: unknown[] = []; + const call = (step: TeardownStep | undefined) => { + if (!step) return Promise.resolve(); + try { + return Promise.resolve(step()).catch((error: unknown) => { + errors.push(error); + }); + } catch (error) { + errors.push(error); + return Promise.resolve(); + } + }; + + const immediate = [call(steps.stopHeartbeat), call(steps.removeActive)]; + await Promise.all(immediate); + await call(steps.appendEvent); + await call(steps.abortAgent); + await call(steps.markStopped); + await call(steps.closeAgent); + await call(steps.stopContainer); + await call(steps.release); + return { errors }; +}; diff --git a/apps/agent-worker/tests/unit/codex-process.test.ts b/apps/agent-worker/tests/unit/codex-process.test.ts new file mode 100644 index 0000000..27cb15f --- /dev/null +++ b/apps/agent-worker/tests/unit/codex-process.test.ts @@ -0,0 +1,37 @@ +import { describe, expect, test, vi } from 'vitest'; + +import { + abortManagedCodexProcess, + managedCodexCommand, +} from '../../src/codex-process'; + +describe('managed Codex process', () => { + test('launches Codex in a job-specific process session', () => { + const command = managedCodexCommand({ + jobId: 'job-1', + pidFile: '/home/alice/.codex/job-1.pid', + command: ['codex', 'exec', '--json', 'fix it'], + }); + + expect(command.slice(0, 2)).toEqual(['python3', '-c']); + expect(command).toContain('/home/alice/.codex/job-1.pid'); + expect(command).toContain('job-1'); + expect(command.slice(-4)).toEqual(['codex', 'exec', '--json', 'fix it']); + }); + + test('invokes an isolated in-box kill for the matching job', async () => { + const execute = vi.fn().mockResolvedValue({ exitCode: 0, output: '' }); + + await abortManagedCodexProcess({ + jobId: 'job-1', + pidFile: '/home/alice/.codex/job-1.pid', + execute, + }); + + expect(execute).toHaveBeenCalledOnce(); + const [command] = execute.mock.calls[0] as [string[]]; + expect(command.slice(0, 2)).toEqual(['python3', '-c']); + expect(command).toContain('/home/alice/.codex/job-1.pid'); + expect(command).toContain('job-1'); + }); +}); diff --git a/apps/agent-worker/tests/unit/heartbeat-controller.test.ts b/apps/agent-worker/tests/unit/heartbeat-controller.test.ts new file mode 100644 index 0000000..d933bee --- /dev/null +++ b/apps/agent-worker/tests/unit/heartbeat-controller.test.ts @@ -0,0 +1,60 @@ +import { afterEach, beforeEach, describe, expect, test, vi } from 'vitest'; + +import { createHeartbeatController } from '../../src/heartbeat-controller'; + +describe('heartbeat controller', () => { + beforeEach(() => vi.useFakeTimers()); + afterEach(() => vi.useRealTimers()); + + test('does not overlap ticks while a heartbeat is in flight', async () => { + let resolveHeartbeat!: (value: { cancelRequested: boolean }) => void; + const heartbeat = vi.fn( + () => + new Promise<{ cancelRequested: boolean }>((resolve) => { + resolveHeartbeat = resolve; + }), + ); + const controller = createHeartbeatController({ + intervalMs: 30_000, + heartbeat, + onCancel: vi.fn(), + onError: vi.fn(), + }); + + controller.start('job-1'); + await vi.advanceTimersByTimeAsync(90_000); + expect(heartbeat).toHaveBeenCalledOnce(); + + resolveHeartbeat({ cancelRequested: false }); + await Promise.resolve(); + await vi.advanceTimersByTimeAsync(30_000); + expect(heartbeat).toHaveBeenCalledTimes(2); + controller.stop('job-1'); + }); + + test('removes the timer before awaiting cancellation teardown', async () => { + let finishCancel!: () => void; + const onCancel = vi.fn( + () => + new Promise((resolve) => { + finishCancel = resolve; + }), + ); + const heartbeat = vi.fn().mockResolvedValue({ cancelRequested: true }); + const controller = createHeartbeatController({ + intervalMs: 30_000, + heartbeat, + onCancel, + onError: vi.fn(), + }); + + controller.start('job-1'); + await vi.advanceTimersByTimeAsync(30_000); + expect(onCancel).toHaveBeenCalledOnce(); + expect(controller.has('job-1')).toBe(false); + + await vi.advanceTimersByTimeAsync(90_000); + expect(heartbeat).toHaveBeenCalledOnce(); + finishCancel(); + }); +}); diff --git a/apps/agent-worker/tests/unit/workspace-teardown.test.ts b/apps/agent-worker/tests/unit/workspace-teardown.test.ts new file mode 100644 index 0000000..3bc1fc2 --- /dev/null +++ b/apps/agent-worker/tests/unit/workspace-teardown.test.ts @@ -0,0 +1,62 @@ +import { describe, expect, test, vi } from 'vitest'; + +import { abortManagedCodexProcess } from '../../src/codex-process'; +import { teardownWorkspace } from '../../src/workspace-teardown'; + +describe('workspace teardown', () => { + test('event failure does not skip abort, stop, removal, or release', async () => { + const calls: string[] = []; + const record = (name: string) => vi.fn(() => calls.push(name)); + + const result = await teardownWorkspace({ + stopHeartbeat: record('heartbeat'), + removeActive: record('remove'), + appendEvent: vi.fn().mockRejectedValue(new Error('event failed')), + abortAgent: record('abort'), + markStopped: vi.fn().mockRejectedValue(new Error('mutation failed')), + closeAgent: record('close'), + stopContainer: vi.fn().mockRejectedValue(new Error('container failed')), + release: record('release'), + }); + + expect(calls).toEqual(['heartbeat', 'remove', 'abort', 'close', 'release']); + expect(result.errors).toHaveLength(3); + }); + + test('removes active state before asynchronous cancellation work', async () => { + const active = new Set(['job-1']); + let activeDuringAbort = true; + + await teardownWorkspace({ + stopHeartbeat: vi.fn(), + removeActive: () => active.delete('job-1'), + abortAgent: () => { + activeDuringAbort = active.has('job-1'); + }, + release: vi.fn(), + }); + + expect(activeDuringAbort).toBe(false); + expect(active.has('job-1')).toBe(false); + }); + + test('Codex cancellation invokes isolated kill and leaves no active workspace', async () => { + const active = new Set(['job-1']); + const execute = vi.fn().mockResolvedValue({ exitCode: 0, output: '' }); + + await teardownWorkspace({ + stopHeartbeat: vi.fn(), + removeActive: () => active.delete('job-1'), + abortAgent: async () => + await abortManagedCodexProcess({ + jobId: 'job-1', + pidFile: '/home/alice/.codex/job-1.pid', + execute, + }), + release: vi.fn(), + }); + + expect(execute).toHaveBeenCalledOnce(); + expect(active.has('job-1')).toBe(false); + }); +}); diff --git a/packages/backend/tests/unit/recover-stale-jobs.test.ts b/packages/backend/tests/unit/recover-stale-jobs.test.ts index 62285f9..33a4180 100644 --- a/packages/backend/tests/unit/recover-stale-jobs.test.ts +++ b/packages/backend/tests/unit/recover-stale-jobs.test.ts @@ -1,5 +1,5 @@ import { convexTest } from 'convex-test'; -import { describe, expect, test } from 'vitest'; +import { describe, expect, test, vi } from 'vitest'; import type { Id } from '../../convex/_generated/dataModel.js'; import { internal } from '../../convex/_generated/api.js'; @@ -7,9 +7,15 @@ import schema from '../../convex/schema'; const modules = import.meta.glob('../../convex/**/*.*s'); -const seedRunningJob = async ( +type ActiveStatus = 'claimed' | 'preparing' | 'running' | 'checks_running'; + +const seedJob = async ( t: ReturnType, - lastHeartbeatAt: number, + args: { + status?: ActiveStatus; + claimedAt?: number; + lastHeartbeatAt?: number; + }, ) => await t.mutation(async (ctx) => { const now = Date.now(); @@ -60,7 +66,7 @@ const seedRunningJob = async ( ownerId, agentRequestId: requestId, threadId, - status: 'running', + status: args.status ?? 'running', prompt: 'Recover the worker job', runtime: 'opencode', workspaceStatus: 'active', @@ -75,8 +81,10 @@ const seedRunningJob = async ( model: 'gpt-5.1-codex', reasoningEffort: 'medium', claimedBy: 'w1', - claimedAt: now, - lastHeartbeatAt, + claimedAt: args.claimedAt ?? now, + ...(args.lastHeartbeatAt === undefined + ? {} + : { lastHeartbeatAt: args.lastHeartbeatAt }), createdAt: now, updatedAt: now, }); @@ -102,7 +110,9 @@ const readState = async ( describe('recoverStaleJobs', () => { test('times out a running job with a stale heartbeat', async () => { const t = convexTest(schema, modules); - const ids = await seedRunningJob(t, Date.now() - 10 * 60 * 1000); + const ids = await seedJob(t, { + lastHeartbeatAt: Date.now() - 10 * 60 * 1000, + }); const result = await t.mutation(internal.agentJobs.recoverStaleJobs, {}); const state = await readState(t, ids); @@ -117,7 +127,7 @@ describe('recoverStaleJobs', () => { test('leaves a running job with a fresh heartbeat untouched', async () => { const t = convexTest(schema, modules); - const ids = await seedRunningJob(t, Date.now()); + const ids = await seedJob(t, { lastHeartbeatAt: Date.now() }); const result = await t.mutation(internal.agentJobs.recoverStaleJobs, {}); const state = await readState(t, ids); @@ -127,4 +137,57 @@ describe('recoverStaleJobs', () => { expect(state.request?.status).toBe('running'); expect(state.thread?.status).toBe('running'); }); + + test('uses claimedAt when a claimed job has no heartbeat', async () => { + const t = convexTest(schema, modules); + const ids = await seedJob(t, { + status: 'claimed', + claimedAt: Date.now() - 10 * 60 * 1000, + }); + + const result = await t.mutation(internal.agentJobs.recoverStaleJobs, {}); + const state = await readState(t, ids); + + expect(result).toEqual({ recovered: 1 }); + expect(state.job?.status).toBe('timed_out'); + }); + + test('times out a heartbeat exactly at the default threshold', async () => { + vi.useFakeTimers(); + try { + const now = new Date('2026-07-10T12:00:00.000Z'); + vi.setSystemTime(now); + const t = convexTest(schema, modules); + const ids = await seedJob(t, { + lastHeartbeatAt: now.getTime() - 3 * 60 * 1000, + }); + + const result = await t.mutation(internal.agentJobs.recoverStaleJobs, {}); + const state = await readState(t, ids); + + expect(result).toEqual({ recovered: 1 }); + expect(state.job?.status).toBe('timed_out'); + } finally { + vi.useRealTimers(); + } + }); + + test.each([ + 'claimed', + 'preparing', + 'running', + 'checks_running', + ])('recovers stale %s jobs', async (status) => { + const t = convexTest(schema, modules); + const ids = await seedJob(t, { + status, + lastHeartbeatAt: Date.now() - 10 * 60 * 1000, + }); + + const result = await t.mutation(internal.agentJobs.recoverStaleJobs, {}); + const state = await readState(t, ids); + + expect(result).toEqual({ recovered: 1 }); + expect(state.job?.status).toBe('timed_out'); + }); });