From 470af043fa16b8cb45a6c7cc1fbfadc5edd58da9 Mon Sep 17 00:00:00 2001 From: Gabriel Brown Date: Fri, 10 Jul 2026 17:38:31 -0400 Subject: [PATCH] fix(agent-jobs): add heartbeat cancel and stale recovery --- apps/agent-worker/src/worker.ts | 51 +++++++ packages/backend/convex/agentJobs.ts | 48 ++++++- packages/backend/convex/crons.ts | 6 + .../tests/unit/recover-stale-jobs.test.ts | 130 ++++++++++++++++++ .../backend/tests/unit/update-status.test.ts | 32 +++++ 5 files changed, 266 insertions(+), 1 deletion(-) create mode 100644 packages/backend/tests/unit/recover-stale-jobs.test.ts diff --git a/apps/agent-worker/src/worker.ts b/apps/agent-worker/src/worker.ts index 3c6bc3b..9715b5a 100644 --- a/apps/agent-worker/src/worker.ts +++ b/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>(); + +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); diff --git a/packages/backend/convex/agentJobs.ts b/packages/backend/convex/agentJobs.ts index f820f02..019578c 100644 --- a/packages/backend/convex/agentJobs.ts +++ b/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' }; }, }); diff --git a/packages/backend/convex/crons.ts b/packages/backend/convex/crons.ts index 0b65e24..070130e 100644 --- a/packages/backend/convex/crons.ts +++ b/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 diff --git a/packages/backend/tests/unit/recover-stale-jobs.test.ts b/packages/backend/tests/unit/recover-stale-jobs.test.ts new file mode 100644 index 0000000..62285f9 --- /dev/null +++ b/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, + 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, + 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'); + }); +}); diff --git a/packages/backend/tests/unit/update-status.test.ts b/packages/backend/tests/unit/update-status.test.ts index 0758db1..691e5d9 100644 --- a/packages/backend/tests/unit/update-status.test.ts +++ b/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 }); + }); +});