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

This commit is contained in:
Gabriel Brown
2026-07-10 18:03:45 -04:00
parent cbbb168e7c
commit 470af043fa
5 changed files with 266 additions and 1 deletions
+47 -1
View File
@@ -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' };
},
});
+6
View File
@@ -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
@@ -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');
});
});
@@ -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 });
});
});