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
+51
View File
@@ -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);