feat: complete TASK-WP3-04 worker recovery

apps/api/src/generation-submission.ts

@@ -260,6 +260,8 @@ export class GenerationSubmissionService {
     const prompt = normalizePrompt(input.prompt);
     const model = this.models.readModel(input.modelId);
     if (!model) throw new GenerationSubmissionError("generation_blocked", { errorCategory: "model_disabled" });
+    const workerState = this.database.prepare("SELECT status FROM worker_runtime_state WHERE singleton = 1").get() as { status: "ready" | "degraded" } | undefined;
+    if (workerState?.status === "degraded") throw new GenerationSubmissionError("generation_blocked", { errorCategory: "gateway_contract_invalid" });
     if (model.configVersion !== input.modelConfigVersion || model.creditCost !== input.confirmedCreditCost) {
       throw new GenerationSubmissionError("model_config_stale", {
         latest: { configVersion: model.configVersion, creditCost: model.creditCost, modelId: model.modelId },
@@ -501,6 +503,14 @@ export class GenerationSubmissionService {
     this.ensureColumn("generation_jobs", "submission_ready", "INTEGER NOT NULL DEFAULT 0");
     this.ensureColumn("generation_jobs", "config_snapshot_json", "TEXT");
     this.database.exec(`
+      CREATE TABLE IF NOT EXISTS worker_runtime_state (
+        singleton INTEGER PRIMARY KEY CHECK (singleton = 1),
+        status TEXT NOT NULL CHECK (status IN ('ready', 'degraded')),
+        reason TEXT,
+        owner_id TEXT,
+        lock_expires_at INTEGER,
+        updated_at INTEGER NOT NULL
+      );
       CREATE UNIQUE INDEX IF NOT EXISTS generation_jobs_active_owner
       ON generation_jobs(owner_id) WHERE status IN ('queued', 'running');
       CREATE UNIQUE INDEX IF NOT EXISTS generation_jobs_client_submission
@@ -530,5 +540,8 @@ export class GenerationSubmissionService {
       WHEN dada_allow_privacy_purge() <> 1
       BEGIN SELECT RAISE(ABORT, 'generation_reference_snapshot_immutable'); END;
     `);
+    this.ensureColumn("worker_runtime_state", "owner_id", "TEXT");
+    this.ensureColumn("worker_runtime_state", "lock_expires_at", "INTEGER");
+    this.database.prepare("INSERT OR IGNORE INTO worker_runtime_state (singleton, status, reason, owner_id, lock_expires_at, updated_at) VALUES (1, 'ready', NULL, NULL, NULL, ?)").run(this.clock());
   }
 }

apps/worker/src/ai-adapter-contract.ts

@@ -18,6 +18,7 @@ export interface NormalizedGenerationOutput {
 
 export type GenerationAdapterResult =
   | { outputs: readonly NormalizedGenerationOutput[]; status: "completed" }
+  | { status: "pending"; upstreamJobReference: string }
   | {
       balanceSignal?: { gatewayAccountRef: string; impactScope: "model" | "account" | "unknown" };
       category: GenerationErrorCategory;
@@ -27,6 +28,7 @@ export type GenerationAdapterResult =
 
 export interface GenerationAdapter {
   start(request: GenerationAdapterRequest): Promise<GenerationAdapterResult>;
+  poll?(upstreamJobReference: string): Promise<GenerationAdapterResult>;
 }
 
 type MockResult = GenerationAdapterResult & { unsafeRaw?: string };
@@ -41,6 +43,7 @@ export class MockGenerationAdapter implements GenerationAdapter {
 
   async start(request: GenerationAdapterRequest): Promise<GenerationAdapterResult> {
     this.calls.push({ generationId: request.generationId, modelId: request.modelId });
+    if (this.result.status === "pending") return { status: "pending", upstreamJobReference: this.result.upstreamJobReference };
     if (this.result.status === "completed") {
       return { outputs: this.result.outputs.map((output) => ({ ...output, bytes: Buffer.from(output.bytes) })), status: "completed" };
     }

apps/worker/src/generation-processor.ts

@@ -27,6 +27,10 @@ interface JobRow {
   prompt: string;
   ratio: "3:4" | "1:1" | "4:3" | "9:16";
   status: "queued" | "running" | "succeeded" | "failed" | "rejected";
+  heartbeat_at: number | null;
+  lease_expires_at: number | null;
+  lease_owner: string | null;
+  upstream_job_reference: string | null;
 }
 
 interface ReservationRow {
@@ -41,7 +45,12 @@ export interface GenerationProcessingResult {
   generationId: string;
   outputAssetId: string | null;
   projectId: string;
-  status: "succeeded" | "failed" | "rejected";
+  status: "running" | "succeeded" | "failed" | "rejected";
+}
+
+export interface LeaseRecoveryAction {
+  action: "poll" | "unknown_retryable";
+  generationId: string;
 }
 
 function iso(timestamp: number) {
@@ -73,6 +82,8 @@ export class GenerationProcessor {
   private readonly dataRoot: string;
   private readonly gatewayBalance: GatewayBalanceRuntime;
   private readonly workerId: string;
+  private closed = false;
+  private readonly lockTimer: ReturnType<typeof setInterval>;
 
   constructor(input: { adapter: GenerationAdapter; clock?: () => number; dataRoot: string; databasePath: string; workerId: string }) {
     if (!input.workerId) throw new Error("generation_worker_id_required");
@@ -84,18 +95,41 @@ export class GenerationProcessor {
     configureWorkerDatabase(this.database);
     this.migrate();
     this.gatewayBalance = new GatewayBalanceRuntime({ clock: this.clock, database: this.database });
+    try {
+      this.setReady();
+      try {
+        this.recoverExpiredLeases();
+      } catch {
+        this.setDegraded("lease_recovery_failed");
+      }
+      this.lockTimer = setInterval(() => this.renewWorkerLock(), 10_000);
+    } catch (error) {
+      this.database.close();
+      throw error;
+    }
   }
 
   close() {
+    if (this.closed) return;
+    this.closed = true;
+    clearInterval(this.lockTimer);
+    this.immediate(() => {
+      const now = this.clock();
+      this.database.prepare("UPDATE worker_runtime_state SET status = 'degraded', reason = ?, owner_id = NULL, lock_expires_at = NULL, updated_at = ? WHERE singleton = 1 AND owner_id = ?")
+        .run("worker_stopped", now, this.workerId);
+    });
+    this.gatewayBalance.close();
     this.database.close();
   }
 
   async processNext() {
+    if (this.readWorkerState().status === "degraded") return undefined;
+    this.recoverExpiredLeases();
     const row = this.database.prepare(`
       SELECT generation_id FROM generation_jobs
-      WHERE status = 'queued' AND submission_ready = 1
+      WHERE submission_ready = 1 AND (status = 'queued' OR (status = 'running' AND lease_owner = ?))
       ORDER BY created_at, generation_id LIMIT 1
-    `).get() as { generation_id: string } | undefined;
+    `).get(this.workerId) as { generation_id: string } | undefined;
     return row ? this.processGeneration(row.generation_id) : undefined;
   }
 
@@ -110,19 +144,25 @@ export class GenerationProcessor {
     `).all(generationId) as Array<{ managed_file_id: string }>;
     let adapterResult: GenerationAdapterResult;
     try {
-      adapterResult = await this.adapter.start({
-        configSnapshot: JSON.parse(job.config_snapshot_json) as Record<string, unknown>,
-        generationId,
-        modelId: job.model_id,
-        prompt: job.prompt,
-        ratio: job.ratio,
-        referenceAssetIds: references.map((row) => row.managed_file_id),
-      });
+      if (job.upstream_job_reference) {
+        if (!this.adapter.poll) return this.completeFailure(job, "unknown_retryable", "poll_unsupported", undefined, false, "pending_manual_review");
+        adapterResult = await this.withHeartbeat(generationId, () => this.adapter.poll!(job.upstream_job_reference!));
+      } else {
+        adapterResult = await this.withHeartbeat(generationId, () => this.adapter.start({
+            configSnapshot: JSON.parse(job.config_snapshot_json) as Record<string, unknown>,
+            generationId,
+            modelId: job.model_id,
+            prompt: job.prompt,
+            ratio: job.ratio,
+            referenceAssetIds: references.map((row) => row.managed_file_id),
+          }));
+      }
     } catch {
-      adapterResult = { category: "unknown_retryable", sourceCategory: "adapter_exception", status: "failed" };
+      return this.completeFailure(job, "unknown_retryable", "adapter_exception", undefined, false, "pending_manual_review");
     }
 
     if (adapterResult.status === "failed") return this.completeFailure(job, adapterResult.category, adapterResult.sourceCategory, adapterResult.balanceSignal);
+    if (adapterResult.status === "pending") return this.persistPending(job.generation_id, adapterResult.upstreamJobReference);
     const output = adapterResult.outputs.length === 1 ? adapterResult.outputs[0] : undefined;
     if (!output || !this.isSafeOutput(job, output)) {
       return this.completeFailure(job, adapterResult.outputs.length === 1 ? "unknown_retryable" : "gateway_contract_invalid", "output_validation_failed", undefined, true);
@@ -148,17 +188,123 @@ export class GenerationProcessor {
         if (changed.changes !== 1) throw new Error("generation_claim_conflict");
         return this.readJob(generationId);
       }
-      if (row.status === "running" && row.final_credit_state === null) return row;
+      if (row.status === "running" && row.final_credit_state === null && row.lease_owner === this.workerId && (row.lease_expires_at ?? 0) >= this.clock()) return row;
+      if (row.status === "running" && row.final_credit_state === null) throw new Error("generation_claim_conflict");
       return row;
     });
   }
 
+  private async withHeartbeat<T>(generationId: string, action: () => Promise<T>) {
+    const timer = setInterval(() => {
+      try {
+        this.heartbeat(generationId);
+      } catch {
+        this.setDegraded("heartbeat_failed");
+      }
+    }, 10_000);
+    try {
+      return await action();
+    } finally {
+      clearInterval(timer);
+    }
+  }
+
+  private renewWorkerLock() {
+    try {
+      const now = this.clock();
+      this.database.prepare("UPDATE worker_runtime_state SET lock_expires_at = ?, updated_at = ? WHERE singleton = 1 AND owner_id = ?")
+        .run(now + 30_000, now, this.workerId);
+    } catch {
+      if (!this.closed) this.setDegraded("worker_lock_renewal_failed");
+    }
+  }
+
+  heartbeat(generationId: string) {
+    return this.immediate(() => {
+      const now = this.clock();
+      const changed = this.database.prepare(`
+        UPDATE generation_jobs SET heartbeat_at = ?, lease_expires_at = ?, updated_at = ?
+        WHERE generation_id = ? AND status = 'running' AND lease_owner = ?
+      `).run(now, now + 30_000, now, generationId, this.workerId);
+      if (changed.changes !== 1) throw new Error("generation_heartbeat_conflict");
+      this.database.prepare("UPDATE worker_runtime_state SET lock_expires_at = ?, updated_at = ? WHERE singleton = 1 AND owner_id = ?")
+        .run(now + 30_000, now, this.workerId);
+      return { generationId, heartbeat_at: now, lease_expires_at: now + 30_000 };
+    });
+  }
+
+  recoverExpiredLeases(): LeaseRecoveryAction[] {
+    return this.immediate(() => {
+      const now = this.clock();
+      const rows = this.database.prepare(`
+        SELECT * FROM generation_jobs
+        WHERE status = 'running' AND submission_ready = 1 AND lease_expires_at IS NOT NULL AND lease_expires_at < ?
+        ORDER BY created_at, generation_id
+      `).all(now) as JobRow[];
+      const actions: LeaseRecoveryAction[] = [];
+      for (const row of rows) {
+        if (row.upstream_job_reference) {
+          this.database.prepare(`
+            UPDATE generation_jobs SET lease_owner = ?, lease_expires_at = ?, heartbeat_at = ?, updated_at = ?
+            WHERE generation_id = ? AND status = 'running' AND lease_expires_at < ?
+          `).run(this.workerId, now + 30_000, now, now, row.generation_id, now);
+          actions.push({ action: "poll", generationId: row.generation_id });
+        } else {
+          this.completeFailure(row, "unknown_retryable", "lease_expired_without_reference", undefined, false, "pending_manual_review");
+          actions.push({ action: "unknown_retryable", generationId: row.generation_id });
+        }
+      }
+      return actions;
+    });
+  }
+
+  setDegraded(reason: string) {
+    return this.immediate(() => {
+      const now = this.clock();
+      this.database.prepare("UPDATE worker_runtime_state SET status = 'degraded', reason = ?, updated_at = ? WHERE singleton = 1").run(reason, now);
+      return this.readWorkerState();
+    });
+  }
+
+  setReady() {
+    return this.immediate(() => {
+      const now = this.clock();
+      const current = this.database.prepare("SELECT owner_id, lock_expires_at FROM worker_runtime_state WHERE singleton = 1")
+        .get() as { lock_expires_at: number | null; owner_id: string | null } | undefined;
+      if (current?.owner_id && current.owner_id !== this.workerId && (current.lock_expires_at ?? 0) >= now) throw new Error("worker_instance_lock_conflict");
+      this.database.prepare("UPDATE worker_runtime_state SET status = 'ready', reason = NULL, owner_id = ?, lock_expires_at = ?, updated_at = ? WHERE singleton = 1")
+        .run(this.workerId, now + 30_000, now);
+      return this.readWorkerState();
+    });
+  }
+
+  private readWorkerState() {
+    return this.database.prepare("SELECT status, reason, updated_at FROM worker_runtime_state WHERE singleton = 1").get() as { reason: string | null; status: "ready" | "degraded"; updated_at: number };
+  }
+
+  private persistPending(generationId: string, upstreamJobReference: string): GenerationProcessingResult {
+    return this.immediate(() => {
+      if (!/^[A-Za-z0-9][A-Za-z0-9:_.-]{0,159}$/.test(upstreamJobReference)) {
+        const job = this.readJob(generationId);
+        return this.completeFailure(job, "gateway_contract_invalid", "upstream_reference_invalid");
+      }
+      const now = this.clock();
+      this.database.prepare(`
+        UPDATE generation_jobs SET upstream_job_reference = ?, heartbeat_at = ?, lease_expires_at = ?, updated_at = ?
+        WHERE generation_id = ? AND status = 'running' AND lease_owner = ?
+      `).run(upstreamJobReference, now, now + 30_000, now, generationId, this.workerId);
+      const job = this.readJob(generationId);
+      return { category: null, generationId, outputAssetId: null, projectId: job.project_id, status: "running" };
+    });
+  }
+
   private completeFailure(
     job: JobRow,
     category: GenerationErrorCategory,
     sourceCategory: string,
     balanceSignal?: { gatewayAccountRef: string; impactScope: "model" | "account" | "unknown" },
     upstreamOutcomeKnown = false,
+    reconciliation: "not_required" | "pending_manual_review" = upstreamOutcomeKnown ? "pending_manual_review" : "not_required",
   ) {
     return this.immediate(() => {
       const replay = this.readReceipt(job.generation_id);
@@ -189,7 +335,7 @@ export class GenerationProcessor {
         category,
         safeSourceCategory.test(sourceCategory) ? sourceCategory : "adapter_error",
         upstreamOutcomeKnown ? 1 : 0,
-        upstreamOutcomeKnown ? "pending_manual_review" : "not_required",
+        reconciliation,
         now,
         now,
         job.generation_id,
@@ -418,7 +564,16 @@ export class GenerationProcessor {
     this.ensureColumn("generation_jobs", "diagnostic_source_category", "TEXT");
     this.ensureColumn("generation_jobs", "upstream_outcome_known", "INTEGER NOT NULL DEFAULT 0");
     this.ensureColumn("generation_jobs", "upstream_cost_reconciliation", "TEXT NOT NULL DEFAULT 'not_applicable'");
+    this.ensureColumn("generation_jobs", "upstream_job_reference", "TEXT");
     this.database.exec(`
+      CREATE TABLE IF NOT EXISTS worker_runtime_state (
+        singleton INTEGER PRIMARY KEY CHECK (singleton = 1),
+        status TEXT NOT NULL CHECK (status IN ('ready', 'degraded')),
+        reason TEXT,
+        owner_id TEXT,
+        lock_expires_at INTEGER,
+        updated_at INTEGER NOT NULL
+      );
       CREATE TABLE IF NOT EXISTS generation_output_assets (
         generation_id TEXT PRIMARY KEY REFERENCES generation_jobs(generation_id),
         managed_file_id TEXT NOT NULL UNIQUE REFERENCES managed_files(file_id),
@@ -432,5 +587,8 @@ export class GenerationProcessor {
         created_at INTEGER NOT NULL
       );
     `);
+    this.ensureColumn("worker_runtime_state", "owner_id", "TEXT");
+    this.ensureColumn("worker_runtime_state", "lock_expires_at", "INTEGER");
+    this.database.prepare("INSERT OR IGNORE INTO worker_runtime_state (singleton, status, reason, owner_id, lock_expires_at, updated_at) VALUES (1, 'ready', NULL, NULL, NULL, ?)").run(this.clock());
   }
 }