Harden Live worker recovery

AI-assisted: Codex
2026-09-14 07:06:45 +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() {