Improve Live worker recovery and Accept latency

AI-assisted: Codex
This commit is contained in:
Paul Bakaus
2026-07-13 09:56:48 -07:00
parent 8a6c4ab486
commit 2008381e91
2 changed files with 140 additions and 29 deletions
+49 -28
View File
@@ -30,6 +30,7 @@ import {
postReply,
requiresAgentReply,
} from '../live-poll.mjs';
import { createLiveSessionStore } from './session-store.mjs';
export const CODEX_WORKER_EVENT_TYPES = Object.freeze(['generate', 'accept', 'discard', 'prefetch']);
export const CODEX_WORKER_EVENT_LEASE_MS = 15_000;
@@ -49,6 +50,7 @@ export class CodexLiveWorkerSupervisor {
publishCheckpoint = postVariantCheckpoint,
publishPhase = postAgentPhase,
postCleanup = postCarbonizeCleanup,
sessionStore = null,
log = () => {},
}) {
this.cwd = path.resolve(cwd);
@@ -64,6 +66,7 @@ export class CodexLiveWorkerSupervisor {
this.publishCheckpoint = publishCheckpoint;
this.publishPhase = publishPhase;
this.postCleanup = postCleanup;
this.sessionStore = sessionStore || createLiveSessionStore({ cwd: this.cwd });
this.log = log;
this.running = false;
this.queue = Promise.resolve();
@@ -131,7 +134,10 @@ export class CodexLiveWorkerSupervisor {
}
if (event.type === 'accept' || event.type === 'discard') {
this.canceled.add(event.id);
await this.cancelActive(event.type, event.id);
// Cancellation fences publication synchronously. Do not make the
// deterministic Accept/Discard path wait on a slow app-server
// interrupt round trip before it can update source and reply.
void this.cancelActive(event.type, event.id);
const handled = await this.handleAccept(event, this.base, this.token);
if (event.type === 'accept' && handled?._acceptResult?.carbonize === true) {
await this.postCleanup(this.base, this.token, {
@@ -175,12 +181,18 @@ export class CodexLiveWorkerSupervisor {
this.active = { eventId: event.id, turnId: null };
this.writeState('working', { eventId: event.id });
try {
if (this.config.delivery === 'progressive' && Number(event.count || 0) > 1) {
await this.runGenerationPhase(event, 'first', 1);
const expectedVariants = Number(event.count || 1);
const snapshot = this.sessionStore.getSnapshot(event.id, { includeCompleted: true });
const sameEpoch = Number(snapshot?.generationEpoch || 1) === Number(event.generationEpoch || 1);
const arrivedVariants = sameEpoch ? Number(snapshot?.arrivedVariants || 0) : 0;
if (this.config.delivery === 'progressive' && expectedVariants > 1) {
if (arrivedVariants < 1) await this.runGenerationPhase(event, 'first', 1);
if (this.isCanceled(event.id)) return;
await this.runGenerationPhase(event, 'final', Number(event.count));
} else {
await this.runGenerationPhase(event, 'atomic', Number(event.count || 1));
if (arrivedVariants < expectedVariants) {
await this.runGenerationPhase(event, 'final', expectedVariants);
}
} else if (arrivedVariants < expectedVariants) {
await this.runGenerationPhase(event, 'atomic', expectedVariants);
}
if (this.isCanceled(event.id)) return;
await this.reply(this.base, this.token, {
@@ -226,34 +238,43 @@ export class CodexLiveWorkerSupervisor {
cwd: this.cwd,
});
let publishedFromMessage = false;
let publicationPromise = null;
let earlyCandidateError = null;
const publishCandidate = async (answer) => {
if (publishedFromMessage || this.isCanceled(event.id)) return;
if (!publicationPromise) {
publicationPromise = (async () => {
await this.publishPhase(this.base, this.token, {
eventId: event.id,
phase: phase === 'final' ? 'remaining_variants_validating' : 'first_variant_validating',
durationMs: Date.now() - phaseStartedAt,
});
applyCodexWorkerOutput({
output: answer,
prepared,
phase,
expectedVariants: Number(event.count || arrivedVariants),
cwd: this.cwd,
maxBytes: this.config.maxArtifactBytes,
});
if (this.isCanceled(event.id)) return;
const published = publishCodexWorkerPhase({ event, prepared, arrivedVariants, cwd: this.cwd });
await this.publishCheckpoint(this.base, this.token, {
event,
published,
scaffold: event.scaffold,
arrivedVariants,
});
publishedFromMessage = true;
})();
}
const pendingPublication = publicationPromise;
try {
await this.publishPhase(this.base, this.token, {
eventId: event.id,
phase: phase === 'final' ? 'remaining_variants_validating' : 'first_variant_validating',
durationMs: Date.now() - phaseStartedAt,
});
applyCodexWorkerOutput({
output: answer,
prepared,
phase,
expectedVariants: Number(event.count || arrivedVariants),
cwd: this.cwd,
maxBytes: this.config.maxArtifactBytes,
});
if (this.isCanceled(event.id)) return;
const published = publishCodexWorkerPhase({ event, prepared, arrivedVariants, cwd: this.cwd });
await this.publishCheckpoint(this.base, this.token, {
event,
published,
scaffold: event.scaffold,
arrivedVariants,
});
publishedFromMessage = true;
await pendingPublication;
} catch (error) {
earlyCandidateError = error;
} finally {
if (publicationPromise === pendingPublication) publicationPromise = null;
}
};
const result = await this.runTurnWithReconnect({
+91 -1
View File
@@ -107,6 +107,44 @@ describe('Codex Live worker supervisor ownership and lifecycle', () => {
assert.equal(supervisor.canceled.has('generation-1'), true);
});
it('does not block deterministic Accept on the app-server interrupt round trip', async () => {
const cwd = mkdtempSync(path.join(tmpdir(), 'codex-supervisor-fast-accept-'));
const client = fakeClient();
let releaseInterrupt;
const interruptReleased = new Promise((resolve) => { releaseInterrupt = resolve; });
client.interruptTurn = async (threadId, turnId) => {
client.calls.interruptTurn.push({ threadId, turnId });
await interruptReleased;
};
let acceptStarted = false;
const events = [
{ type: 'accept', id: 'generation-1', variantId: '1' },
{ type: 'exit' },
];
const supervisor = new CodexLiveWorkerSupervisor({
cwd,
base: 'http://localhost:1',
token: 'token',
client,
config: { model: null, effort: 'low', delivery: 'progressive', maxArtifactBytes: 2_000_000 },
statePath: path.join(cwd, 'state.json'),
scriptsDir: path.join(cwd, 'skill/scripts'),
fetchEvent: async () => events.shift(),
handleAccept: async () => {
acceptStarted = true;
releaseInterrupt();
return { _acceptResult: { handled: true, carbonize: false } };
},
});
supervisor.thread = { id: 'live-worker-thread' };
supervisor.active = { eventId: 'generation-1', turnId: 'turn-1' };
await supervisor.run();
assert.equal(acceptStarted, true);
assert.equal(client.calls.interruptTurn.length >= 1, true);
});
it('interrupts a canceled turn whose id arrives after Accept', async () => {
const cwd = mkdtempSync(path.join(tmpdir(), 'codex-supervisor-late-turn-'));
const client = fakeClient();
@@ -208,6 +246,55 @@ describe('Codex Live worker supervisor ownership and lifecycle', () => {
assert.equal(client.calls.resumeDedicatedThread.length, 1);
});
it('resumes progressive delivery from durable variant checkpoints', async () => {
const cwd = mkdtempSync(path.join(tmpdir(), 'codex-supervisor-checkpoint-resume-'));
const phases = [];
const replies = [];
const supervisor = new CodexLiveWorkerSupervisor({
cwd,
base: 'http://localhost:1',
token: 'token',
client: fakeClient(),
config: { model: null, effort: 'low', delivery: 'progressive', maxArtifactBytes: 2_000_000 },
statePath: path.join(cwd, 'state.json'),
scriptsDir: path.join(cwd, 'skill/scripts'),
reply: async (_base, _token, value) => { replies.push(value); },
});
supervisor.runGenerationPhase = async (_event, phase, arrivedVariants) => {
phases.push({ phase, arrivedVariants });
};
const partialId = 'resume-partial';
const store = createLiveSessionStore({ cwd, sessionId: partialId });
store.appendEvent({ type: 'generate', id: partialId, count: 3, generationEpoch: 1 });
store.appendEvent({ type: 'checkpoint', id: partialId, phase: 'cycling', revision: 1, arrivedVariants: 1 });
await supervisor.processGeneration({
type: 'generate',
id: partialId,
count: 3,
generationEpoch: 1,
scaffold: { file: 'src/App.jsx' },
});
assert.deepEqual(phases, [{ phase: 'final', arrivedVariants: 3 }]);
assert.equal(replies.at(-1).type, 'done');
phases.length = 0;
const completeId = 'resume-complete';
const completeStore = createLiveSessionStore({ cwd, sessionId: completeId });
completeStore.appendEvent({ type: 'generate', id: completeId, count: 3, generationEpoch: 1 });
completeStore.appendEvent({ type: 'checkpoint', id: completeId, phase: 'cycling', revision: 2, arrivedVariants: 3 });
await supervisor.processGeneration({
type: 'generate',
id: completeId,
count: 3,
generationEpoch: 1,
scaffold: { file: 'src/App.jsx' },
});
assert.deepEqual(phases, []);
assert.equal(replies.at(-1).id, completeId);
assert.equal(replies.at(-1).type, 'done');
});
it('archives its dedicated thread during clean Live shutdown', async () => {
const cwd = mkdtempSync(path.join(tmpdir(), 'codex-supervisor-close-'));
const client = fakeClient();
@@ -268,7 +355,10 @@ describe('Codex Live worker supervisor ownership and lifecycle', () => {
const prompt = input.find((item) => item.type === 'text').text;
const artifactPath = JSON.parse(prompt.match(/Return exactly one file whose path is ("[^"]+")/)[1]);
const message = JSON.stringify({ files: [{ path: artifactPath, content: turn === 1 ? first : final }] });
await onAgentMessage?.(message);
await Promise.all([
onAgentMessage?.(message),
onAgentMessage?.(message),
]);
return { message };
};
const replies = [];