fix(worker): make cancellation teardown failure-safe

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<CommandResult>;
+}) => {
+  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}`);
+  }
+};

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<void>;
+  onError: (error: unknown) => void;
+}) => {
+  const timers = new Map<string, ReturnType<typeof setInterval>>();
+  const inFlight = new Set<string>();
+
+  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),
+  };
+};

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<string, ReturnType<typeof setInterval>>();
-
-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);
+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),
 });
-          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,11 +796,26 @@ 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({
+  let result;
+  try {
+    result = await streamExecInContainer({
       containerName: workspace.boxName,
       containerCwd: workspace.containerRepo,
       command,
@@ -854,6 +851,10 @@ const runCodexTurn = async (args: {
         }
       },
     });
+  } finally {
+    workspace.codexPidFile = undefined;
+    await rm(pidFileHostPath, { force: true });
+  }
   await appendEvent(
     workspace.claim.job._id,
     'info',
@@ -1485,6 +1486,10 @@ const runClaim = async (claim: Claim) => {
     });
   } catch (error) {
     const message = error instanceof Error ? error.message : String(error);
+    const workspaceStatus = message.toLowerCase().includes('timed out')
+      ? 'expired'
+      : 'failed';
+    try {
       await appendEvent(
         jobId,
         'error',
@@ -1500,18 +1505,28 @@ const runClaim = async (claim: Claim) => {
       });
       await updateStatus(
         jobId,
-      message.toLowerCase().includes('timed out') ? 'timed_out' : 'failed',
+        workspaceStatus === 'expired' ? 'timed_out' : 'failed',
         { error: truncate(redact(message), 10_000) },
       );
-    await markWorkspaceStopped(
+    } finally {
+      const workspace = activeWorkspaces.get(jobId);
+      if (workspace) {
+        const teardown = await teardownActiveWorkspace({
           jobId,
-      message.toLowerCase().includes('timed out') ? 'expired' : 'failed',
-    ).catch((stopError: unknown) => {
-      console.error(stopError);
+          workspace,
+          abortAgent: true,
+          workspaceStatus,
         });
+        reportTeardownErrors(teardown.errors);
+      } else {
         stopHeartbeat(jobId);
+        await markWorkspaceStopped(jobId, workspaceStatus).catch(
+          (stopError: unknown) => console.error(stopError),
+        );
         acquiredBoxHandle?.release();
       }
+    }
+  }
 };
 
 export const listWorkspaceTree = async (jobId: string) => {
@@ -1638,10 +1653,24 @@ 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);
+  } 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;
@@ -1652,13 +1681,11 @@ export const abortWorkspaceAgent = async (jobId: string) => {
     'Agent turn aborted.',
   );
   return { success: true };
-  }
-  if (workspace.runtimeMode === 'codex_exec') {
-    throw new Error('Codex agent turns cannot be aborted from Spoon yet.');
-  }
-  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 };
 };

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 };
+};

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');
+  });
+});

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<void>((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();
+  });
+});

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);
+  });
+});

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<typeof convexTest>,
-  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<ActiveStatus>([
+    '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');
+  });
 });