Harden Live worker recovery

AI-assisted: Codex
2026-09-12 06:06:37 +03:00 · 2026-07-13 10:44:37 -07:00
parent a274f93c4e
commit c6dfd22329
7 changed files with 296 additions and 21 deletions
@@ -16,6 +16,7 @@
 import fs from 'node:fs';
 import path from 'node:path';
 import { isGeneratedFile } from './lib/is-generated.mjs';
+import { getLiveDir } from './lib/impeccable-paths.mjs';
 import { readBuffer as readManualEditsBuffer, writeBuffer as writeManualEditsBuffer } from './live/manual-edits-buffer.mjs';
 import { withSourceLockSync } from './live/source-lock.mjs';
 import {
@@ -72,6 +73,32 @@ Output (JSON):
  if (!id) { console.error('Missing --id'); process.exit(1); }
  if (!isDiscard && !variantNum) { console.error('Need --discard or --variant N'); process.exit(1); }

+  const requestedOperation = isDiscard ? 'discard' : 'accept';
+  const priorReceipt = readAcceptReceipt(process.cwd(), id);
+  if (priorReceipt) {
+    const sameOperation = priorReceipt.operation === requestedOperation
+      && (isDiscard || String(priorReceipt.variantId) === String(variantNum));
+    console.log(JSON.stringify(sameOperation
+      ? { ...priorReceipt.result, handled: true, alreadyApplied: true }
+      : {
+          handled: false,
+          error: 'accept_receipt_conflict',
+          priorOperation: priorReceipt.operation,
+          priorVariantId: priorReceipt.variantId ?? null,
+        }));
+    return;
+  }
+  const emitResult = (result) => {
+    if (result?.handled !== false) {
+      writeAcceptReceipt(process.cwd(), id, {
+        operation: requestedOperation,
+        variantId: isDiscard ? null : String(variantNum),
+        result,
+      });
+    }
+    console.log(JSON.stringify(result));
+  };
+
  let paramValues = null;
  if (paramValuesRaw) {
    try { paramValues = JSON.parse(paramValuesRaw); }
@@ -104,13 +131,13 @@ Output (JSON):
      } catch (err) {
        result = { handled: false, error: err.message };
      }
-      console.log(JSON.stringify({
+      emitResult({
        ...result,
        file: vueComponentManifest.sourceFile,
        carbonize: false,
        previewMode: 'vue-component',
        componentDir: vueComponentManifest.componentDir,
-      }));
+      });
      return;
    }

@@ -133,7 +160,7 @@ Output (JSON):
        carbonize: false,
      };
    }
-    console.log(JSON.stringify(result));
+    emitResult(result);
    return;
  }

@@ -153,13 +180,13 @@ Output (JSON):
      } catch (err) {
        result = { handled: false, error: err.message };
      }
-      console.log(JSON.stringify({
+      emitResult({
        ...result,
        file: svelteComponentManifest.sourceFile,
        carbonize: false,
        previewMode: 'svelte-component',
        componentDir: svelteComponentManifest.componentDir,
-      }));
+      });
      return;
    }

@@ -189,7 +216,7 @@ Output (JSON):
    if (result.carbonize) {
      result.todo = 'REQUIRED before next poll: carbonize cleanup in ' + result.file + '. See reference/live.md "Required after accept".';
    }
-    console.log(JSON.stringify({ handled: result.handled !== false, ...result }));
+    emitResult({ handled: result.handled !== false, ...result });
    return;
  }

@@ -221,7 +248,7 @@ Output (JSON):

  if (isDiscard) {
    const result = handleDiscard(id, lines, targetFile);
-    console.log(JSON.stringify({ handled: true, file: relFile, carbonize: false, ...result }));
+    emitResult({ handled: true, file: relFile, carbonize: false, ...result });
  } else {
    const result = handleAccept(id, variantNum, lines, targetFile, paramValues);
    const acceptedOriginalText = result.acceptedOriginalText || '';
@@ -242,7 +269,7 @@ Output (JSON):
        // Non-fatal; the buffer stays as-is and the user can discard later.
      }
    }
-    console.log(JSON.stringify({ handled: true, file: relFile, ...result }));
+    emitResult({ handled: true, file: relFile, ...result });
  }
 }

@@ -887,6 +914,28 @@ function searchDir(dir, query, seen, depth) {
 // Utilities
 // ---------------------------------------------------------------------------

+function acceptReceiptPath(cwd, id) {
+  return path.join(getLiveDir(cwd), 'accept-receipts', `${id}.json`);
+}
+
+function readAcceptReceipt(cwd, id) {
+  try { return JSON.parse(fs.readFileSync(acceptReceiptPath(cwd, id), 'utf-8')); } catch { return null; }
+}
+
+function writeAcceptReceipt(cwd, id, receipt) {
+  const file = acceptReceiptPath(cwd, id);
+  fs.mkdirSync(path.dirname(file), { recursive: true });
+  const value = {
+    id,
+    ...receipt,
+    completedAt: new Date().toISOString(),
+  };
+  const temporary = `${file}.${process.pid}.${Date.now()}.tmp`;
+  fs.writeFileSync(temporary, JSON.stringify(value, null, 2) + '\n', 'utf-8');
+  fs.renameSync(temporary, file);
+  return value;
+}
+
 function argVal(args, flag) {
  const idx = args.indexOf(flag);
  return idx !== -1 && idx + 1 < args.length ? args[idx + 1] : null;
@@ -160,6 +160,7 @@ export async function fetchNextEvent(base, token, {
  resolveTypes,
  perRequestTimeoutMs = PER_REQUEST_TIMEOUT_MS,
  leaseMs = DEFAULT_EVENT_LEASE_MS,
+  signal,
 } = {}) {
  while (true) {
    if (totalDeadline && Date.now() >= totalDeadline) {
@@ -177,7 +178,7 @@ export async function fetchNextEvent(base, token, {
    });
    const normalizedTypes = normalizePollTypes(resolveTypes ? await resolveTypes() : types);
    if (normalizedTypes.length > 0) query.set('types', normalizedTypes.join(','));
-    const res = await fetch(`${base}/poll?${query}`);
+    const res = await fetch(`${base}/poll?${query}`, { signal });

    if (res.status === 401) {
      const err = new Error('Authentication failed. The server token may have changed.');
@@ -199,7 +200,7 @@ export async function fetchNextEvent(base, token, {
  }
 }

-export async function augmentEventWithAcceptHandling(event, base, token) {
+export async function augmentEventWithAcceptHandling(event, base, token, { deferReply = false } = {}) {
  if (event.type !== 'accept' && event.type !== 'discard') return event;

  const __dirname = path.dirname(fileURLToPath(import.meta.url));
@@ -217,6 +218,15 @@ export async function augmentEventWithAcceptHandling(event, base, token) {
    event._acceptResult = { handled: false, mode: 'error', error: err.message };
  }

+  if (deferReply) {
+    event._completionAck = { ok: false, deferred: true };
+    return event;
+  }
+  await completeAcceptHandling(event, base, token);
+  return event;
+}
+
+export async function completeAcceptHandling(event, base, token) {
  const completionType = completionTypeForAcceptResult(event.type, event._acceptResult);
  try {
    await postReply(base, token, {
@@ -233,7 +243,6 @@ export async function augmentEventWithAcceptHandling(event, base, token) {
  if (!event._completionAck) {
    event._completionAck = completionAckForAcceptResult(event.id, completionType, event._acceptResult);
  }
-
  return event;
 }

@@ -290,6 +290,17 @@ function acknowledgePendingEvent(id, sourceEventType) {
  return acknowledged;
 }

+function releasePendingEvent(id, sourceEventType) {
+  const entry = state.pendingEvents.find((item) => (
+    item.event?.id === id
+    && (!sourceEventType || item.event?.type === sourceEventType)
+  ));
+  if (!entry) return null;
+  entry.leaseUntil = 0;
+  scheduleLeaseFlush();
+  return entry.event;
+}
+
 function retirePendingGeneration(id) {
  if (!id) return 0;
  let retired = 0;
@@ -1042,6 +1053,21 @@ function handlePollPost(req, res) {
      return;
    }
    const sourceEventType = msg.sourceEventType || inferSourceEventType(msg);
+    if (msg.type === 'retry') {
+      const releasedEvent = releasePendingEvent(msg.id, sourceEventType);
+      if (!releasedEvent) {
+        res.writeHead(msg.id ? 404 : 400, { 'Content-Type': 'application/json' });
+        res.end(JSON.stringify({
+          error: msg.id ? 'unknown_poll_retry_id' : 'missing_poll_retry_id',
+          id: msg.id,
+        }));
+        return;
+      }
+      flushPendingPolls();
+      res.writeHead(200, { 'Content-Type': 'application/json' });
+      res.end(JSON.stringify({ ok: true, released: true }));
+      return;
+    }
    const pendingEventBeforeAck = findPendingEventById(msg.id, sourceEventType);
    if (pendingEventBeforeAck?.type === 'steer' && msg.type === 'steer_done'
        && !msg.file && !(typeof msg.message === 'string' && msg.message.trim())) {
@@ -26,6 +26,7 @@ import {
 } from './codex-worker.mjs';
 import {
  augmentEventWithAcceptHandling,
+  completeAcceptHandling,
  fetchNextEvent,
  postReply,
  requiresAgentReply,
@@ -46,6 +47,7 @@ export class CodexLiveWorkerSupervisor {
    scriptsDir,
    fetchEvent = fetchNextEvent,
    handleAccept = augmentEventWithAcceptHandling,
+    completeAccept = completeAcceptHandling,
    reply = postReply,
    publishCheckpoint = postVariantCheckpoint,
    publishPhase = postAgentPhase,
@@ -62,6 +64,7 @@ export class CodexLiveWorkerSupervisor {
    this.scriptsDir = scriptsDir;
    this.fetchEvent = fetchEvent;
    this.handleAccept = handleAccept;
+    this.completeAccept = completeAccept;
    this.reply = reply;
    this.publishCheckpoint = publishCheckpoint;
    this.publishPhase = publishPhase;
@@ -73,6 +76,9 @@ export class CodexLiveWorkerSupervisor {
    this.active = null;
    this.canceled = new Set();
    this.queuedGenerationIds = new Set();
+    this.pollAbortController = null;
+    this.activePoll = null;
+    this.failure = null;
    this.thread = null;
    this.threadReady = Promise.resolve(null);
    this.model = null;
@@ -115,11 +121,24 @@ export class CodexLiveWorkerSupervisor {
  async run() {
    if (!this.thread) await this.initialize();
    this.running = true;
+    this.pollAbortController = new AbortController();
    while (this.running) {
-      const event = await this.fetchEvent(this.base, this.token, {
-        types: CODEX_WORKER_EVENT_TYPES,
-        leaseMs: CODEX_WORKER_EVENT_LEASE_MS,
-      });
+      let event;
+      try {
+        const poll = this.fetchEvent(this.base, this.token, {
+          types: CODEX_WORKER_EVENT_TYPES,
+          leaseMs: CODEX_WORKER_EVENT_LEASE_MS,
+          signal: this.pollAbortController.signal,
+        });
+        this.activePoll = poll;
+        event = await poll;
+      } catch (error) {
+        if (!this.running && (error?.name === 'AbortError' || this.pollAbortController.signal.aborted)) break;
+        throw error;
+      } finally {
+        this.activePoll = null;
+      }
+      if (!this.running) break;
      if (!event || event.type === 'timeout') continue;
      if (event.type === 'exit') {
        await this.cancelActive('live_exit');
@@ -134,15 +153,21 @@ export class CodexLiveWorkerSupervisor {
        // interrupt round trip before it can update source and reply.
        void this.cancelActive(event.type, event.id);
        if (replaceBusyThread) this.rotateWorkerThread(event.type);
-        const handled = await this.handleAccept(event, this.base, this.token);
+        const handled = await this.handleAccept(event, this.base, this.token, {
+          deferReply: event.type === 'accept',
+        });
        if (event.type === 'accept' && handled?._acceptResult?.carbonize === true) {
          await this.postCleanup(this.base, this.token, {
+            id: event.id,
            sessionId: event.id,
            file: handled._acceptResult.file,
            variantId: event.variantId,
            acceptResult: handled._acceptResult,
          });
        }
+        if (handled?._completionAck?.deferred === true) {
+          await this.completeAccept(handled, this.base, this.token);
+        }
        continue;
      }
      if (event.type === 'generate') {
@@ -165,7 +190,7 @@ export class CodexLiveWorkerSupervisor {
      }
    }
    await this.queue.catch(() => {});
-    await this.shutdown({ archive: true });
+    await this.shutdown({ archive: !this.failure });
  }

  async processGeneration(event) {
@@ -247,6 +272,18 @@ export class CodexLiveWorkerSupervisor {
  }

  async runGenerationPhase(event, phase, arrivedVariants) {
+    for (let attempt = 0; attempt < 2; attempt += 1) {
+      try {
+        return await this.runGenerationPhaseOnce(event, phase, arrivedVariants);
+      } catch (error) {
+        const sourceChangedDuringGeneration = error?.code === 'publish_source_hash_mismatch';
+        if (!sourceChangedDuringGeneration || attempt > 0 || this.isCanceled(event.id)) throw error;
+        this.log(`source changed during ${event.id} ${phase}; re-preparing once before publication`);
+      }
+    }
+  }
+
+  async runGenerationPhaseOnce(event, phase, arrivedVariants) {
    if (this.isCanceled(event.id)) return;
    const phaseStartedAt = Date.now();
    await this.publishPhase(this.base, this.token, {
@@ -391,12 +428,28 @@ export class CodexLiveWorkerSupervisor {
  async handleGenerationFailure(event, error) {
    if (this.isCanceled(event.id) || error.code === 'TURN_INTERRUPTED') return;
    this.log(`generation ${event.id} failed: ${error.stack || error.message}`);
+    this.failure = {
+      eventId: event.id,
+      error: error.message,
+      failedAt: new Date().toISOString(),
+    };
+    this.running = false;
+    this.pollAbortController?.abort();
+    if (this.activePoll) {
+      await Promise.race([
+        this.activePoll.catch(() => null),
+        new Promise((resolve) => {
+          const timer = setTimeout(resolve, 250);
+          timer.unref?.();
+        }),
+      ]);
+    }
    await this.reply(this.base, this.token, {
      id: event.id,
-      type: 'error',
+      type: 'retry',
      sourceEventType: event.type,
-      message: `Dedicated Codex worker failed: ${error.message}`,
    }).catch(() => {});
+    this.writeState('failed', this.failure);
  }

  isCanceled(eventId) {
@@ -428,7 +481,10 @@ export class CodexLiveWorkerSupervisor {
      }
    }
    await this.client.close().catch(() => {});
-    this.writeState(archived ? 'archived' : 'stopped', { archived });
+    this.writeState(
+      this.failure ? 'failed' : archived ? 'archived' : 'stopped',
+      { archived, ...(this.failure || {}) },
+    );
  }

  status() {
@@ -74,6 +74,33 @@ describe('live-accept — style-element edge cases', () => {
    assert.ok(!after.includes('original text'), 'original content dropped');
  });

+  it('replays a durable receipt when Accept is retried after source was already written', () => {
+    const html = `<body>
+  <!-- impeccable-variants-start RECEIPT1 -->
+  <div data-impeccable-variants="RECEIPT1" data-impeccable-variant-count="2" style="display: contents">
+    <div data-impeccable-variant="original"><p>original</p></div>
+    <style data-impeccable-css="RECEIPT1" />
+    <div data-impeccable-variant="1"><p>accepted once</p></div>
+    <div data-impeccable-variant="2" style="display: none"><p>other</p></div>
+  </div>
+  <!-- impeccable-variants-end RECEIPT1 -->
+</body>`;
+    writeFileSync(join(tmp, 'page.html'), html);
+
+    const first = runAccept(tmp, ['--id', 'RECEIPT1', '--variant', '1']);
+    const afterFirst = readFileSync(join(tmp, 'page.html'), 'utf-8');
+    const replay = runAccept(tmp, ['--id', 'RECEIPT1', '--variant', '1']);
+
+    assert.equal(first.handled, true);
+    assert.equal(replay.handled, true);
+    assert.equal(replay.alreadyApplied, true);
+    assert.equal(replay.file, 'page.html');
+    assert.equal(readFileSync(join(tmp, 'page.html'), 'utf-8'), afterFirst);
+    const conflict = runAccept(tmp, ['--id', 'RECEIPT1', '--variant', '2']);
+    assert.equal(conflict.handled, false);
+    assert.equal(conflict.error, 'accept_receipt_conflict');
+  });
+
  // Variant: same-line <style>…</style> block should also be treated as a
  // single skipped unit; the line has both open and close tags.
  it('finds the accepted variant after a single-line <style>…</style> block', () => {
@@ -207,6 +207,7 @@ describe('Codex Live worker supervisor ownership and lifecycle', () => {
  it('queues carbonize cleanup onto the foreground control lane', async () => {
    const cwd = mkdtempSync(path.join(tmpdir(), 'codex-supervisor-carbonize-'));
    const cleanups = [];
+    const order = [];
    const client = fakeClient();
    const supervisor = new CodexLiveWorkerSupervisor({
      cwd,
@@ -219,8 +220,10 @@ describe('Codex Live worker supervisor ownership and lifecycle', () => {
      handleAccept: async (event) => ({
        ...event,
        _acceptResult: { handled: true, carbonize: true, file: 'src/App.jsx' },
+        _completionAck: { ok: false, deferred: true },
      }),
-      postCleanup: async (_base, _token, event) => { cleanups.push(event); },
+      postCleanup: async (_base, _token, event) => { cleanups.push(event); order.push('cleanup'); },
+      completeAccept: async () => { order.push('accept_ack'); },
    });
    supervisor.running = true;
    supervisor.thread = { id: 'live-worker-thread' };
@@ -233,11 +236,13 @@ describe('Codex Live worker supervisor ownership and lifecycle', () => {
    };
    await supervisor.run();
    assert.deepEqual(cleanups, [{
+      id: 'abc12345',
      sessionId: 'abc12345',
      file: 'src/App.jsx',
      variantId: '1',
      acceptResult: { handled: true, carbonize: true, file: 'src/App.jsx' },
    }]);
+    assert.deepEqual(order, ['cleanup', 'accept_ack']);
  });

  it('renews but never queues the same long-running generation twice', async () => {
@@ -288,6 +293,63 @@ describe('Codex Live worker supervisor ownership and lifecycle', () => {
    assert.equal(client.calls.resumeDedicatedThread.length, 1);
  });

+  it('relinquishes Generate and advertises foreground fallback after permanent failure', async () => {
+    const cwd = mkdtempSync(path.join(tmpdir(), 'codex-supervisor-fallback-'));
+    const replies = [];
+    const statePath = path.join(cwd, 'state.json');
+    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,
+      scriptsDir: path.join(cwd, 'skill/scripts'),
+      reply: async (_base, _token, value) => { replies.push(value); },
+    });
+    supervisor.running = true;
+    supervisor.pollAbortController = new AbortController();
+
+    await supervisor.handleGenerationFailure(
+      { type: 'generate', id: 'recoverable-generation' },
+      new Error('app-server remained unavailable after reconnect'),
+    );
+
+    assert.equal(supervisor.running, false);
+    assert.equal(supervisor.pollAbortController.signal.aborted, true);
+    assert.deepEqual(replies, [{
+      id: 'recoverable-generation',
+      type: 'retry',
+      sourceEventType: 'generate',
+    }]);
+    const state = JSON.parse(readFileSync(statePath, 'utf-8'));
+    assert.equal(state.status, 'failed');
+    assert.equal(state.eventId, 'recoverable-generation');
+    assert.match(state.error, /app-server remained unavailable/);
+  });
+
+  it('re-prepares once when foreground cleanup changes source during generation', async () => {
+    const cwd = mkdtempSync(path.join(tmpdir(), 'codex-supervisor-source-race-'));
+    const supervisor = createSupervisor({
+      cwd,
+      statePath: path.join(cwd, 'state.json'),
+      client: fakeClient(),
+    });
+    let attempts = 0;
+    supervisor.runGenerationPhaseOnce = async () => {
+      attempts += 1;
+      if (attempts === 1) {
+        const error = new Error('publish_source_hash_mismatch');
+        error.code = 'publish_source_hash_mismatch';
+        throw error;
+      }
+    };
+
+    await supervisor.runGenerationPhase({ id: 'source-race' }, 'first', 1);
+
+    assert.equal(attempts, 2);
+  });
+
  it('resumes progressive delivery from durable variant checkpoints', async () => {
    const cwd = mkdtempSync(path.join(tmpdir(), 'codex-supervisor-checkpoint-resume-'));
    const phases = [];
@@ -2652,6 +2652,52 @@ colors: {}
    }
  });

+  it('releases a failed worker Generate lease without consuming or broadcasting it', async () => {
+    await drainPolls(server);
+    const id = 'fa11bac1';
+    const generated = await fetch(`http://localhost:${server.port}/events`, {
+      method: 'POST',
+      headers: { 'Content-Type': 'application/json' },
+      body: JSON.stringify({
+        token: server.token,
+        type: 'generate',
+        id,
+        action: 'bolder',
+        count: 3,
+        element: { outerHTML: '<article>fallback</article>', tagName: 'article' },
+      }),
+    });
+    assert.equal(generated.status, 200);
+    const leased = await fetch(`http://localhost:${server.port}/poll?token=${server.token}&types=generate&timeout=100&leaseMs=5000`).then((response) => response.json());
+    assert.equal(leased.id, id);
+
+    const retried = await fetch(`http://localhost:${server.port}/poll`, {
+      method: 'POST',
+      headers: { 'Content-Type': 'application/json' },
+      body: JSON.stringify({
+        token: server.token,
+        id,
+        type: 'retry',
+        sourceEventType: 'generate',
+      }),
+    });
+    assert.equal(retried.status, 200);
+    assert.equal((await retried.json()).released, true);
+
+    const fallback = await fetch(`http://localhost:${server.port}/poll?token=${server.token}&types=generate&timeout=100&leaseMs=100`).then((response) => response.json());
+    assert.equal(fallback.id, id);
+    assert.equal(fallback.type, 'generate');
+    const status = await fetch(`http://localhost:${server.port}/status?token=${server.token}`).then((response) => response.json());
+    assert.equal(status.pendingEvents.some((event) => event.id === id && event.type === 'generate'), true);
+
+    const done = await fetch(`http://localhost:${server.port}/poll`, {
+      method: 'POST',
+      headers: { 'Content-Type': 'application/json' },
+      body: JSON.stringify({ token: server.token, id, type: 'done', sourceEventType: 'generate' }),
+    });
+    assert.equal(done.status, 200);
+  });
+
  it('wakes a parked poll as soon as a missed-ack lease expires', async () => {
    await drainPolls(server);