gib/spoon
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(worker): make cancellation teardown failure-safe

Browse Source
This commit is contained in:
Gabriel Brown
2026-07-10 18:04:40 -04:00
parent 470af043fa
commit ac68484059
8 changed files with 604 additions and 162 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/agent-worker/src/codex-process.ts
+88
View File
@@ -0,0 +1,88 @@
type CommandResult = { exitCode: number; output: string };
const launchScript = `
import os
import sys
pid_file, job_id, *command = sys.argv[1:]
if os.getpgrp() != os.getpid():
os.setsid()
os.environ["SPOON_CODEX_JOB_ID"] = job_id
with open(pid_file, "w", encoding="utf-8") as handle:
handle.write(str(os.getpid()))
handle.flush()
os.fsync(handle.fileno())
os.execvpe(command[0], command, os.environ)
`;
const abortScript = `
import os
import signal
import sys
import time
pid_file, job_id = sys.argv[1:]
pid = None
for _ in range(20):
try:
with open(pid_file, encoding="utf-8") as handle:
pid = int(handle.read().strip())
break
except FileNotFoundError:
time.sleep(0.1)
if pid is None:
sys.exit(0)
try:
with open(f"/proc/{pid}/environ", "rb") as handle:
environment = handle.read().split(b"\\0")
marker = f"SPOON_CODEX_JOB_ID={job_id}".encode()
if marker not in environment or os.getpgid(pid) != pid:
raise RuntimeError("Codex process identity did not match the requested job")
os.killpg(pid, signal.SIGTERM)
for _ in range(20):
time.sleep(0.1)
try:
os.killpg(pid, 0)
except ProcessLookupError:
break
else:
os.killpg(pid, signal.SIGKILL)
except ProcessLookupError:
pass
finally:
try:
os.unlink(pid_file)
except FileNotFoundError:
pass
`;
export const managedCodexCommand = (args: {
jobId: string;
pidFile: string;
command: string[];
}) => [
'python3',
'-c',
launchScript,
args.pidFile,
args.jobId,
...args.command,
];
export const abortManagedCodexProcess = async (args: {
jobId: string;
pidFile: string;
execute: (command: string[]) => Promise<CommandResult>;
}) => {
const result = await args.execute([
'python3',
'-c',
abortScript,
args.pidFile,
args.jobId,
]);
if (result.exitCode !== 0) {
throw new Error(`Failed to abort Codex process: ${result.output}`);
}
};
apps/agent-worker/src/heartbeat-controller.ts
+42
View File
@@ -0,0 +1,42 @@
export const createHeartbeatController = (args: {
intervalMs: number;
heartbeat: (jobId: string) => Promise<{ cancelRequested: boolean }>;
onCancel: (jobId: string) => Promise<void>;
onError: (error: unknown) => void;
}) => {
const timers = new Map<string, ReturnType<typeof setInterval>>();
const inFlight = new Set<string>();
const stop = (jobId: string) => {
const timer = timers.get(jobId);
if (timer) clearInterval(timer);
timers.delete(jobId);
};
const tick = async (jobId: string) => {
if (inFlight.has(jobId) || !timers.has(jobId)) return;
inFlight.add(jobId);
try {
const result = await args.heartbeat(jobId);
if (result.cancelRequested) {
stop(jobId);
await args.onCancel(jobId);
}
} catch (error) {
args.onError(error);
} finally {
inFlight.delete(jobId);
}
};
return {
start: (jobId: string) => {
stop(jobId);
const timer = setInterval(() => void tick(jobId), args.intervalMs);
timer.unref();
timers.set(jobId, timer);
},
stop,
has: (jobId: string) => timers.has(jobId),
};
};
apps/agent-worker/src/worker.ts
+209 -154
View File
@@ -18,6 +18,7 @@ import type { NormalizedAgentEvent } from './agent-events';
import type { OpenCodeSession } from './opencode-session';
import type { BoxHandle } from './user-container';
import { normalizeCodexJsonLine } from './agent-events';
import { abortManagedCodexProcess, managedCodexCommand } from './codex-process';
import { prepareCodexWorkspaceFiles } from './codex-runtime';
import { env } from './env';
import {
@@ -28,6 +29,7 @@ import {
run,
} from './git';
import { getInstallationToken, openDraftPullRequest } from './github';
import { createHeartbeatController } from './heartbeat-controller';
import {
abortOpenCodeSession,
createOpenCodeSession,
@@ -45,6 +47,7 @@ import {
} from './runtime/docker';
import { acquireUserBox, runningBoxUsernames } from './user-container';
import { fetchUserEnvironment, materializeUserHome } from './user-environment';
import { teardownWorkspace } from './workspace-teardown';
type Claim = {
job: {
@@ -118,7 +121,9 @@ type ActiveWorkspace = {
opencodePassword?: string;
opencodeSession?: OpenCodeSession;
codexSessionId?: string;
codexPidFile?: string;
agentTurnActive?: boolean;
cancelRequested?: boolean;
resolveTurn?: () => void;
lastRecordedDiffSignature?: string;
// Captures the most recent Codex `error`/`turn.failed` event for the active
@@ -269,43 +274,17 @@ const heartbeatWorkspaceMutation = async (jobId: Id<'agentJobs'>) =>
jobId,
})) as { success: true; cancelRequested: boolean };
const heartbeatTimers = new Map<string, ReturnType<typeof setInterval>>();
const heartbeatController = createHeartbeatController({
intervalMs: 30_000,
heartbeat: async (jobId) =>
await heartbeatWorkspaceMutation(jobId as Id<'agentJobs'>),
onCancel: async (jobId) => await cancelWorkspace(jobId),
onError: (error) => console.error(error),
});
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 stopHeartbeat = (jobId: string) => heartbeatController.stop(jobId);
const startHeartbeat = (jobId: Id<'agentJobs'>) =>
heartbeatController.start(jobId);
const appendMessage = async (args: {
jobId: Id<'agentJobs'>;
@@ -789,7 +768,10 @@ const runCodexTurn = async (args: {
'.codex',
outputFileName,
);
const command = workspace.codexSessionId
if (workspace.cancelRequested) {
throw new Error('Workspace cancellation requested.');
}
const codexCommand = workspace.codexSessionId
? [
'codex',
'exec',
@@ -814,46 +796,65 @@ const runCodexTurn = async (args: {
workspace.containerRepo,
prompt,
];
const pidFileName = `codex-turn-${workspace.claim.job._id}.pid`;
const pidFileHostPath = path.join(workspace.workdir, '.codex', pidFileName);
const pidFileContainerPath = path.posix.join(
workspace.containerHome,
'.codex',
pidFileName,
);
workspace.codexPidFile = pidFileContainerPath;
const command = managedCodexCommand({
jobId: workspace.claim.job._id,
pidFile: pidFileContainerPath,
command: codexCommand,
});
const aiEnv = providerEnvironment(workspace.claim, workspace.containerHome);
const secretEnv = Object.fromEntries(
workspace.claim.secrets.map((secret) => [secret.name, secret.value]),
);
const result = await streamExecInContainer({
containerName: workspace.boxName,
containerCwd: workspace.containerRepo,
command,
environment: {
...aiEnv,
...secretEnv,
},
redact: workspace.redact,
timeoutMs: env.jobTimeoutMs,
onStdoutLine: async (line) => {
for (const event of normalizeCodexJsonLine(line)) {
await handleAgentEvent({
workspace,
event,
assistantMessageId,
assistantContent,
});
}
},
onStderrLine: async (line) => {
const trimmed = line.trim();
if (
trimmed &&
trimmed !== 'Reading additional input from stdin...' &&
!trimmed.includes('`[features].codex_hooks` is deprecated')
) {
await appendEvent(
workspace.claim.job._id,
'info',
'plan',
truncate(trimmed, 10_000),
);
}
},
});
let result;
try {
result = await streamExecInContainer({
containerName: workspace.boxName,
containerCwd: workspace.containerRepo,
command,
environment: {
...aiEnv,
...secretEnv,
},
redact: workspace.redact,
timeoutMs: env.jobTimeoutMs,
onStdoutLine: async (line) => {
for (const event of normalizeCodexJsonLine(line)) {
await handleAgentEvent({
workspace,
event,
assistantMessageId,
assistantContent,
});
}
},
onStderrLine: async (line) => {
const trimmed = line.trim();
if (
trimmed &&
trimmed !== 'Reading additional input from stdin...' &&
!trimmed.includes('`[features].codex_hooks` is deprecated')
) {
await appendEvent(
workspace.claim.job._id,
'info',
'plan',
truncate(trimmed, 10_000),
);
}
},
});
} finally {
workspace.codexPidFile = undefined;
await rm(pidFileHostPath, { force: true });
}
await appendEvent(
workspace.claim.job._id,
'info',
@@ -1485,32 +1486,46 @@ const runClaim = async (claim: Claim) => {
});
} catch (error) {
const message = error instanceof Error ? error.message : String(error);
await appendEvent(
jobId,
'error',
'cleanup',
truncate(redact(message), 20_000),
);
await addArtifact({
jobId,
kind: 'error',
title: 'Failure',
content: truncate(redact(message), 50_000),
contentType: 'text/plain',
});
await updateStatus(
jobId,
message.toLowerCase().includes('timed out') ? 'timed_out' : 'failed',
{ error: truncate(redact(message), 10_000) },
);
await markWorkspaceStopped(
jobId,
message.toLowerCase().includes('timed out') ? 'expired' : 'failed',
).catch((stopError: unknown) => {
console.error(stopError);
});
stopHeartbeat(jobId);
acquiredBoxHandle?.release();
const workspaceStatus = message.toLowerCase().includes('timed out')
? 'expired'
: 'failed';
try {
await appendEvent(
jobId,
'error',
'cleanup',
truncate(redact(message), 20_000),
);
await addArtifact({
jobId,
kind: 'error',
title: 'Failure',
content: truncate(redact(message), 50_000),
contentType: 'text/plain',
});
await updateStatus(
jobId,
workspaceStatus === 'expired' ? 'timed_out' : 'failed',
{ error: truncate(redact(message), 10_000) },
);
} finally {
const workspace = activeWorkspaces.get(jobId);
if (workspace) {
const teardown = await teardownActiveWorkspace({
jobId,
workspace,
abortAgent: true,
workspaceStatus,
});
reportTeardownErrors(teardown.errors);
} else {
stopHeartbeat(jobId);
await markWorkspaceStopped(jobId, workspaceStatus).catch(
(stopError: unknown) => console.error(stopError),
);
acquiredBoxHandle?.release();
}
}
}
};
@@ -1638,27 +1653,39 @@ export const getWorkspaceAgentStatus = (jobId: string) => {
};
};
export const abortWorkspaceAgent = async (jobId: string) => {
const workspace = resolveWorkspace(jobId);
const abortWorkspaceAgentForWorkspace = async (workspace: ActiveWorkspace) => {
if (workspace.opencodeSession) {
await abortOpenCodeSession(workspace.opencodeSession);
workspace.agentTurnActive = false;
workspace.resolveTurn?.();
workspace.resolveTurn = undefined;
await appendEvent(
workspace.claim.job._id,
'warn',
'cleanup',
'Agent turn aborted.',
);
return { success: true };
}
if (workspace.runtimeMode === 'codex_exec') {
throw new Error('Codex agent turns cannot be aborted from Spoon yet.');
} else if (workspace.runtimeMode === 'codex_exec' && workspace.codexPidFile) {
await abortManagedCodexProcess({
jobId: workspace.claim.job._id,
pidFile: workspace.codexPidFile,
execute: async (command) =>
await runExecInContainer({
containerName: workspace.boxName,
containerCwd: workspace.containerRepo,
command,
environment: {},
redact: workspace.redact,
timeoutMs: env.jobTimeoutMs,
}),
});
}
workspace.agentTurnActive = false;
workspace.resolveTurn?.();
workspace.resolveTurn = undefined;
await appendEvent(
workspace.claim.job._id,
'warn',
'cleanup',
'Agent turn aborted.',
);
return { success: true };
};
export const abortWorkspaceAgent = async (jobId: string) =>
await abortWorkspaceAgentForWorkspace(resolveWorkspace(jobId));
export const replyToInteraction = async (
jobId: string,
args: {
@@ -1869,6 +1896,60 @@ export const sendWorkspaceMessage = async (
}
};
const teardownActiveWorkspace = async (args: {
jobId: string;
workspace: ActiveWorkspace;
abortAgent: boolean;
appendCancellationEvent?: boolean;
workspaceStatus?: 'stopped' | 'expired' | 'failed';
}) => {
const { jobId, workspace } = args;
const containerName = workspace.containerName;
if (args.abortAgent) workspace.cancelRequested = true;
return await teardownWorkspace({
stopHeartbeat: () => stopHeartbeat(jobId),
removeActive: () => activeWorkspaces.delete(jobId),
appendEvent: args.appendCancellationEvent
? async () =>
await appendEvent(
workspace.claim.job._id,
'warn',
'cleanup',
'Cancellation requested; stopping workspace.',
)
: undefined,
abortAgent: args.abortAgent
? async () => await abortWorkspaceAgentForWorkspace(workspace)
: undefined,
markStopped: async () =>
await markWorkspaceStopped(workspace.claim.job._id, args.workspaceStatus),
closeAgent: () => workspace.opencodeSession?.close(),
stopContainer: containerName
? async () => await stopWorkspaceContainer(containerName)
: undefined,
release: () => workspace.boxHandle.release(),
});
};
const reportTeardownErrors = (errors: unknown[]) => {
for (const error of errors) console.error(error);
};
const cancelWorkspace = async (jobId: string) => {
const workspace = activeWorkspaces.get(jobId);
if (!workspace) {
stopHeartbeat(jobId);
return;
}
const result = await teardownActiveWorkspace({
jobId,
workspace,
abortAgent: true,
appendCancellationEvent: true,
});
reportTeardownErrors(result.errors);
};
export const openWorkspacePullRequest = async (jobId: string) => {
const workspace = resolveWorkspace(jobId);
const { claim, repoDir, redact } = workspace;
@@ -1931,16 +2012,14 @@ export const openWorkspacePullRequest = async (jobId: string) => {
pullRequestNumber: pullRequest.number,
summary: 'Draft PR opened from interactive workspace.',
});
await markWorkspaceStopped(claim.job._id);
workspace.opencodeSession?.close();
if (workspace.containerName) {
await stopWorkspaceContainer(workspace.containerName);
const teardown = await teardownActiveWorkspace({
jobId,
workspace,
abortAgent: false,
});
if (teardown.errors.length > 0) {
throw new AggregateError(teardown.errors, 'Workspace teardown failed.');
}
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).
workspace.boxHandle.release();
return {
pullRequestUrl: pullRequest.html_url,
pullRequestNumber: pullRequest.number,
@@ -1949,37 +2028,13 @@ 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);
} catch (error) {
errors.push(error);
}
try {
workspace.opencodeSession?.close();
} catch (error) {
errors.push(error);
}
try {
if (workspace.containerName) {
await stopWorkspaceContainer(workspace.containerName);
}
} catch (error) {
errors.push(error);
} finally {
activeWorkspaces.delete(jobId);
try {
workspace.boxHandle.release();
} catch (error) {
errors.push(error);
}
}
if (errors.length === 1) throw errors[0];
if (errors.length > 1) {
throw new AggregateError(errors, 'Workspace teardown failed.');
const teardown = await teardownActiveWorkspace({
jobId,
workspace,
abortAgent: true,
});
if (teardown.errors.length > 0) {
throw new AggregateError(teardown.errors, 'Workspace teardown failed.');
}
return { success: true };
};
apps/agent-worker/src/workspace-teardown.ts
+35
View File
@@ -0,0 +1,35 @@
type TeardownStep = () => unknown;
export const teardownWorkspace = async (steps: {
stopHeartbeat: TeardownStep;
removeActive: TeardownStep;
appendEvent?: TeardownStep;
abortAgent?: TeardownStep;
markStopped?: TeardownStep;
closeAgent?: TeardownStep;
stopContainer?: TeardownStep;
release: TeardownStep;
}) => {
const errors: unknown[] = [];
const call = (step: TeardownStep | undefined) => {
if (!step) return Promise.resolve();
try {
return Promise.resolve(step()).catch((error: unknown) => {
errors.push(error);
});
} catch (error) {
errors.push(error);
return Promise.resolve();
}
};
const immediate = [call(steps.stopHeartbeat), call(steps.removeActive)];
await Promise.all(immediate);
await call(steps.appendEvent);
await call(steps.abortAgent);
await call(steps.markStopped);
await call(steps.closeAgent);
await call(steps.stopContainer);
await call(steps.release);
return { errors };
};