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

2026-07-10 18:03:45 -04:00
parent cbbb168e7c
commit 470af043fa
5 changed files with 266 additions and 1 deletions
@@ -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');
+  });
+});