fix(agent-jobs): add heartbeat cancel and stale recovery

apps/agent-worker/src/worker.ts

@@ -262,6 +262,51 @@ const markWorkspaceStopped = async (
     workspaceStatus,
   });
 
+const heartbeatWorkspaceMutation = async (jobId: Id<'agentJobs'>) =>
+  (await client.mutation(api.agentJobs.heartbeatWorkspace, {
+    workerToken: env.workerToken,
+    workerId: env.workerId,
+    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);
+          });
+          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 appendMessage = async (args: {
   jobId: Id<'agentJobs'>;
   role: 'user' | 'assistant' | 'system' | 'tool';
@@ -1416,6 +1461,7 @@ const runClaim = async (claim: Claim) => {
     });
     activeWorkspaces.set(jobId, workspace);
     await markWorkspaceActive({ jobId });
+    startHeartbeat(jobId);
     await updateStatus(jobId, 'running', {
       summary: 'Workspace is active.',
     });
@@ -1463,6 +1509,7 @@ const runClaim = async (claim: Claim) => {
     ).catch((stopError: unknown) => {
       console.error(stopError);
     });
+    stopHeartbeat(jobId);
     acquiredBoxHandle?.release();
   }
 };
@@ -1889,6 +1936,7 @@ export const openWorkspacePullRequest = async (jobId: string) => {
   if (workspace.containerName) {
     await stopWorkspaceContainer(workspace.containerName);
   }
+  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).
@@ -1901,6 +1949,9 @@ 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);

packages/backend/convex/agentJobs.ts

@@ -1201,6 +1201,52 @@ export const failClaimInternal = internalMutation({
   },
 });
 
+export const recoverStaleJobs = internalMutation({
+  args: { staleMs: v.optional(v.number()) },
+  handler: async (ctx, { staleMs }) => {
+    const threshold = Date.now() - (staleMs ?? 3 * 60 * 1000);
+    const activeStatuses = [
+      'claimed',
+      'preparing',
+      'running',
+      'checks_running',
+    ] as const;
+    let recovered = 0;
+    for (const status of activeStatuses) {
+      const jobs = await ctx.db
+        .query('agentJobs')
+        .withIndex('by_status', (q) => q.eq('status', status))
+        .collect();
+      for (const job of jobs) {
+        const last = job.lastHeartbeatAt ?? job.claimedAt ?? job.createdAt;
+        if (last > threshold) continue;
+        const now = Date.now();
+        await ctx.db.patch(job._id, {
+          status: 'timed_out',
+          error:
+            job.error ??
+            'The worker stopped reporting progress; the job timed out.',
+          completedAt: now,
+          updatedAt: now,
+        });
+        await ctx.db.patch(job.agentRequestId, {
+          status: 'failed',
+          updatedAt: now,
+        });
+        if (job.threadId) {
+          await ctx.db.patch(job.threadId, {
+            status: 'failed',
+            updatedAt: now,
+            resolvedAt: now,
+          });
+        }
+        recovered += 1;
+      }
+    }
+    return { recovered };
+  },
+});
+
 export const updateStatus = mutation({
   args: {
     workerToken: v.string(),
@@ -1471,7 +1517,7 @@ export const heartbeatWorkspace = mutation({
       lastHeartbeatAt: Date.now(),
       updatedAt: Date.now(),
     });
-    return { success: true };
+    return { success: true, cancelRequested: job.status === 'cancelled' };
   },
 });
 

packages/backend/convex/crons.ts

@@ -12,6 +12,12 @@ crons.interval(
     limit: 10,
   },
 );
+crons.interval(
+  'Recover stale agent jobs',
+  { minutes: 1 },
+  internal.agentJobs.recoverStaleJobs,
+  {},
+);
 /* Example cron jobs
 crons.cron(
   // Run at 7:00 AM CST / 8:00 AM CDT

packages/backend/tests/unit/recover-stale-jobs.test.ts

@@ -0,0 +1,130 @@
+import { convexTest } from 'convex-test';
+import { describe, expect, test } from 'vitest';
+
+import type { Id } from '../../convex/_generated/dataModel.js';
+import { internal } from '../../convex/_generated/api.js';
+import schema from '../../convex/schema';
+
+const modules = import.meta.glob('../../convex/**/*.*s');
+
+const seedRunningJob = async (
+  t: ReturnType<typeof convexTest>,
+  lastHeartbeatAt: number,
+) =>
+  await t.mutation(async (ctx) => {
+    const now = Date.now();
+    const ownerId = await ctx.db.insert('users', {
+      email: 'owner@example.com',
+      name: 'Owner',
+    });
+    const spoonId = await ctx.db.insert('spoons', {
+      ownerId,
+      name: 'Recovery Spoon',
+      provider: 'github',
+      upstreamOwner: 'upstream',
+      upstreamRepo: 'recovery',
+      upstreamDefaultBranch: 'main',
+      upstreamUrl: 'https://github.com/upstream/recovery',
+      forkOwner: 'team',
+      forkRepo: 'recovery-spoon',
+      forkUrl: 'https://github.com/team/recovery-spoon',
+      visibility: 'private',
+      maintenanceMode: 'watch',
+      syncCadence: 'daily',
+      productionRefStrategy: 'default_branch',
+      status: 'active',
+      createdAt: now,
+      updatedAt: now,
+    });
+    const requestId = await ctx.db.insert('agentRequests', {
+      spoonId,
+      ownerId,
+      prompt: 'Recover the worker job',
+      status: 'running',
+      createdAt: now,
+      updatedAt: now,
+    });
+    const threadId = await ctx.db.insert('threads', {
+      ownerId,
+      spoonId,
+      title: 'Recover the worker job',
+      source: 'user_request',
+      status: 'running',
+      priority: 'normal',
+      relatedAgentRequestId: requestId,
+      createdAt: now,
+      updatedAt: now,
+    });
+    const jobId = await ctx.db.insert('agentJobs', {
+      spoonId,
+      ownerId,
+      agentRequestId: requestId,
+      threadId,
+      status: 'running',
+      prompt: 'Recover the worker job',
+      runtime: 'opencode',
+      workspaceStatus: 'active',
+      baseBranch: 'main',
+      workBranch: 'spoon/recover-worker-job',
+      forkOwner: 'team',
+      forkRepo: 'recovery-spoon',
+      forkUrl: 'https://github.com/team/recovery-spoon',
+      upstreamOwner: 'upstream',
+      upstreamRepo: 'recovery',
+      selectedSecretIds: [],
+      model: 'gpt-5.1-codex',
+      reasoningEffort: 'medium',
+      claimedBy: 'w1',
+      claimedAt: now,
+      lastHeartbeatAt,
+      createdAt: now,
+      updatedAt: now,
+    });
+    await ctx.db.patch(requestId, { agentJobId: jobId });
+    await ctx.db.patch(threadId, { latestAgentJobId: jobId });
+    return { jobId, requestId, threadId };
+  });
+
+const readState = async (
+  t: ReturnType<typeof convexTest>,
+  ids: {
+    jobId: Id<'agentJobs'>;
+    requestId: Id<'agentRequests'>;
+    threadId: Id<'threads'>;
+  },
+) =>
+  await t.run(async (ctx) => ({
+    job: await ctx.db.get(ids.jobId),
+    request: await ctx.db.get(ids.requestId),
+    thread: await ctx.db.get(ids.threadId),
+  }));
+
+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 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');
+    expect(state.job?.completedAt).toBeTypeOf('number');
+    expect(state.request?.status).toBe('failed');
+    expect(state.thread?.status).toBe('failed');
+    expect(state.thread?.resolvedAt).toBeTypeOf('number');
+  });
+
+  test('leaves a running job with a fresh heartbeat untouched', async () => {
+    const t = convexTest(schema, modules);
+    const ids = await seedRunningJob(t, Date.now());
+
+    const result = await t.mutation(internal.agentJobs.recoverStaleJobs, {});
+    const state = await readState(t, ids);
+
+    expect(result).toEqual({ recovered: 0 });
+    expect(state.job?.status).toBe('running');
+    expect(state.request?.status).toBe('running');
+    expect(state.thread?.status).toBe('running');
+  });
+});

packages/backend/tests/unit/update-status.test.ts

@@ -107,3 +107,35 @@ describe('updateStatus terminal guard', () => {
     expect(job?.status).toBe('checks_running');
   });
 });
+
+describe('heartbeatWorkspace cancellation propagation', () => {
+  test('reports cancellation while refreshing the workspace heartbeat', async () => {
+    const t = convexTest(schema, modules);
+    const jobId = await seedJob(t, 'cancelled');
+
+    const result = await t.mutation(api.agentJobs.heartbeatWorkspace, {
+      workerToken: 'test-worker-token',
+      workerId: 'worker-1',
+      jobId,
+    });
+
+    expect(result).toEqual({ success: true, cancelRequested: true });
+    const job = await t.run(async (ctx) => await ctx.db.get(jobId));
+    expect(job?.status).toBe('cancelled');
+    expect(job?.workspaceStatus).toBe('active');
+    expect(job?.lastHeartbeatAt).toBeTypeOf('number');
+  });
+
+  test('reports no cancellation for an active job', async () => {
+    const t = convexTest(schema, modules);
+    const jobId = await seedJob(t, 'running');
+
+    const result = await t.mutation(api.agentJobs.heartbeatWorkspace, {
+      workerToken: 'test-worker-token',
+      workerId: 'worker-1',
+      jobId,
+    });
+
+    expect(result).toEqual({ success: true, cancelRequested: false });
+  });
+});