596 lines
30 KiB
TypeScript
596 lines
30 KiB
TypeScript
import { createHash, randomUUID } from "node:crypto";
|
|
import { existsSync, mkdirSync, renameSync, rmSync, writeFileSync } from "node:fs";
|
|
import { dirname, join, resolve } from "node:path";
|
|
|
|
import Database from "better-sqlite3";
|
|
import type BetterSqlite3 from "better-sqlite3";
|
|
|
|
import type { GenerationAdapter, GenerationAdapterResult, NormalizedGenerationOutput } from "./ai-adapter-contract.js";
|
|
import { GatewayBalanceRuntime } from "./gateway-balance-runtime.js";
|
|
import { generationErrorRegistry, type GenerationErrorCategory } from "./generation-error-registry.js";
|
|
import { configureWorkerDatabase } from "./sqlite-connection.js";
|
|
|
|
const hardLimitBytes = 5_368_709_120;
|
|
const warningLimitBytes = 4_294_967_296;
|
|
const criticalLimitBytes = 4_831_838_208;
|
|
const maximumOutputBytes = 20 * 1_024 * 1_024;
|
|
const safeSourceCategory = /^[a-z][a-z0-9_]{0,79}$/;
|
|
|
|
interface JobRow {
|
|
config_snapshot_json: string;
|
|
error_category: GenerationErrorCategory | null;
|
|
final_credit_state: "committed" | "released" | null;
|
|
generation_id: string;
|
|
model_id: string;
|
|
owner_id: string;
|
|
project_id: string;
|
|
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 {
|
|
amount: number;
|
|
model_id: string;
|
|
status: "reserved" | "committed" | "released";
|
|
user_id: string;
|
|
}
|
|
|
|
export interface GenerationProcessingResult {
|
|
category: GenerationErrorCategory | null;
|
|
generationId: string;
|
|
outputAssetId: string | null;
|
|
projectId: string;
|
|
status: "running" | "succeeded" | "failed" | "rejected";
|
|
}
|
|
|
|
export interface LeaseRecoveryAction {
|
|
action: "poll" | "unknown_retryable";
|
|
generationId: string;
|
|
}
|
|
|
|
function iso(timestamp: number) {
|
|
return new Date(timestamp).toISOString();
|
|
}
|
|
|
|
function extensionFor(mimeType: NormalizedGenerationOutput["mimeType"]) {
|
|
return mimeType === "image/png" ? ".png" : mimeType === "image/jpeg" ? ".jpg" : ".webp";
|
|
}
|
|
|
|
function hasExpectedMagic(output: NormalizedGenerationOutput) {
|
|
const bytes = output.bytes;
|
|
if (output.mimeType === "image/png") return bytes.length >= 8 && bytes.subarray(0, 8).equals(Buffer.from([137, 80, 78, 71, 13, 10, 26, 10]));
|
|
if (output.mimeType === "image/jpeg") return bytes.length >= 3 && bytes[0] === 0xff && bytes[1] === 0xd8 && bytes[2] === 0xff;
|
|
return bytes.length >= 12 && bytes.subarray(0, 4).toString("ascii") === "RIFF" && bytes.subarray(8, 12).toString("ascii") === "WEBP";
|
|
}
|
|
|
|
function capacityClass(bytes: number) {
|
|
return {
|
|
capacity: bytes < warningLimitBytes ? "normal" : bytes < criticalLimitBytes ? "warning" : "critical",
|
|
status: bytes >= hardLimitBytes ? "full" : "active",
|
|
} as const;
|
|
}
|
|
|
|
export class GenerationProcessor {
|
|
readonly database: BetterSqlite3.Database;
|
|
private readonly adapter: GenerationAdapter;
|
|
private readonly clock: () => number;
|
|
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");
|
|
this.adapter = input.adapter;
|
|
this.clock = input.clock ?? Date.now;
|
|
this.dataRoot = resolve(input.dataRoot);
|
|
this.workerId = input.workerId;
|
|
const nativeBinding = process.env.DADA_SQLITE_NATIVE_BINDING;
|
|
this.database = new Database(input.databasePath, nativeBinding ? { nativeBinding } : undefined);
|
|
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 submission_ready = 1 AND (status = 'queued' OR (status = 'running' AND lease_owner = ?))
|
|
ORDER BY created_at, generation_id LIMIT 1
|
|
`).get(this.workerId) as { generation_id: string } | undefined;
|
|
return row ? this.processGeneration(row.generation_id) : undefined;
|
|
}
|
|
|
|
async processGeneration(generationId: string): Promise<GenerationProcessingResult> {
|
|
const replay = this.readReceipt(generationId);
|
|
if (replay) return replay;
|
|
const job = this.claim(generationId);
|
|
if (["succeeded", "failed", "rejected"].includes(job.status)) return this.terminalResult(job);
|
|
|
|
const references = this.database.prepare(`
|
|
SELECT managed_file_id FROM generation_reference_snapshots WHERE generation_id = ? ORDER BY position
|
|
`).all(generationId) as Array<{ managed_file_id: string }>;
|
|
let adapterResult: GenerationAdapterResult;
|
|
try {
|
|
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 {
|
|
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);
|
|
}
|
|
try {
|
|
return this.completeSuccess(job, output);
|
|
} catch {
|
|
return this.completeFailure(job, "unknown_retryable", "output_persist_failed", undefined, true);
|
|
}
|
|
}
|
|
|
|
private claim(generationId: string) {
|
|
return this.immediate(() => {
|
|
const row = this.readJob(generationId);
|
|
if (row.status === "queued") {
|
|
const now = this.clock();
|
|
const changed = this.database.prepare(`
|
|
UPDATE generation_jobs
|
|
SET status = 'running', lease_owner = ?, lease_expires_at = ?, heartbeat_at = ?, started_at = COALESCE(started_at, ?),
|
|
attempt_no = attempt_no + 1, updated_at = ?
|
|
WHERE generation_id = ? AND status = 'queued' AND submission_ready = 1
|
|
`).run(this.workerId, now + 30_000, now, now, now, generationId);
|
|
if (changed.changes !== 1) throw new Error("generation_claim_conflict");
|
|
return this.readJob(generationId);
|
|
}
|
|
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);
|
|
if (replay) return replay;
|
|
const latest = this.readJob(job.generation_id);
|
|
if (["succeeded", "failed", "rejected"].includes(latest.status)) return this.terminalResult(latest);
|
|
const status = category === "safety_rejected" ? "rejected" as const : "failed" as const;
|
|
this.finalizeCredit(job.generation_id, status);
|
|
if (balanceSignal) {
|
|
this.gatewayBalance.seedModels([{ gatewayAccountRef: balanceSignal.gatewayAccountRef, modelId: job.model_id }]);
|
|
this.gatewayBalance.recordInsufficient({
|
|
eventId: `generation-${job.generation_id}`,
|
|
gatewayAccountRef: balanceSignal.gatewayAccountRef,
|
|
impactScope: balanceSignal.impactScope,
|
|
modelId: job.model_id,
|
|
sourceCategory: "adapter_balance_signal",
|
|
});
|
|
}
|
|
const now = this.clock();
|
|
this.database.prepare(`
|
|
UPDATE generation_jobs
|
|
SET status = ?, error_category = ?, diagnostic_source_category = ?, final_credit_state = 'released',
|
|
upstream_outcome_known = ?, upstream_cost_reconciliation = ?, finished_at = ?, updated_at = ?,
|
|
lease_owner = NULL, lease_expires_at = NULL, heartbeat_at = NULL
|
|
WHERE generation_id = ? AND status = 'running'
|
|
`).run(
|
|
status,
|
|
category,
|
|
safeSourceCategory.test(sourceCategory) ? sourceCategory : "adapter_error",
|
|
upstreamOutcomeKnown ? 1 : 0,
|
|
reconciliation,
|
|
now,
|
|
now,
|
|
job.generation_id,
|
|
);
|
|
const result = this.terminalResult(this.readJob(job.generation_id));
|
|
this.writeReceipt(result, now);
|
|
return result;
|
|
});
|
|
}
|
|
|
|
private completeSuccess(job: JobRow, output: NormalizedGenerationOutput) {
|
|
const fileId = randomUUID();
|
|
const operationId = randomUUID();
|
|
const extension = extensionFor(output.mimeType);
|
|
const relativePath = `content/generated/${job.owner_id}/${fileId}${extension}`;
|
|
const destinationPath = resolve(this.dataRoot, relativePath);
|
|
const stagingDirectory = resolve(this.dataRoot, "staging", operationId);
|
|
const stagingPath = join(stagingDirectory, "payload.tmp");
|
|
if (!destinationPath.startsWith(`${this.dataRoot}\\`) && !destinationPath.startsWith(`${this.dataRoot}/`)) throw new Error("generation_output_path_invalid");
|
|
|
|
this.reserveStorage(operationId, output.bytes.byteLength);
|
|
let renamed = false;
|
|
try {
|
|
mkdirSync(stagingDirectory, { recursive: true });
|
|
writeFileSync(stagingPath, output.bytes, { flag: "wx" });
|
|
mkdirSync(dirname(destinationPath), { recursive: true });
|
|
renameSync(stagingPath, destinationPath);
|
|
renamed = true;
|
|
rmSync(stagingDirectory, { force: true, recursive: true });
|
|
return this.immediate(() => {
|
|
const replay = this.readReceipt(job.generation_id);
|
|
if (replay) return replay;
|
|
const now = this.clock();
|
|
const project = this.database.prepare("SELECT name, state_version, current_image_id FROM projects WHERE project_id = ? AND status = 'active'")
|
|
.get(job.project_id) as { current_image_id: string | null; name: string; state_version: number } | undefined;
|
|
if (!project) throw new Error("generation_project_unavailable");
|
|
const state = this.database.prepare("SELECT canvas_json FROM project_states WHERE project_id = ? ORDER BY state_version DESC LIMIT 1")
|
|
.get(job.project_id) as { canvas_json: string } | undefined;
|
|
if (!state) throw new Error("generation_project_state_unavailable");
|
|
const canvas = JSON.parse(state.canvas_json) as { background: { asset_id: string | null }; [key: string]: unknown };
|
|
if (project.current_image_id === null) canvas.background.asset_id = fileId;
|
|
this.database.prepare(`
|
|
INSERT INTO managed_files (file_id, file_kind, owner_ref, relative_path, byte_size, mime_type, sha256, status, created_at)
|
|
VALUES (?, 'generated', ?, ?, ?, ?, ?, 'committed', ?)
|
|
`).run(fileId, job.owner_id, relativePath, output.bytes.byteLength, output.mimeType, createHash("sha256").update(output.bytes).digest("hex"), iso(now));
|
|
this.database.prepare("INSERT INTO project_resource_files (project_id, managed_file_id, resource_kind, created_at) VALUES (?, ?, 'generated', ?)")
|
|
.run(job.project_id, fileId, now);
|
|
this.database.prepare("INSERT INTO project_asset_refs (reference_id, managed_file_id, reference_type, created_at) VALUES (?, ?, 'project', ?)")
|
|
.run(`project:${job.project_id}:${fileId}`, fileId, iso(now));
|
|
this.database.prepare(`
|
|
INSERT INTO generation_output_assets (generation_id, managed_file_id, pixel_width, pixel_height, created_at)
|
|
VALUES (?, ?, ?, ?, ?)
|
|
`).run(job.generation_id, fileId, output.pixelWidth, output.pixelHeight, now);
|
|
this.database.prepare("INSERT INTO project_images (image_id, project_id, generation_id, created_at) VALUES (?, ?, ?, ?)")
|
|
.run(fileId, job.project_id, job.generation_id, now);
|
|
this.database.prepare(`
|
|
UPDATE projects SET current_image_id = COALESCE(current_image_id, ?), updated_at = ?, state_version = state_version + 1
|
|
WHERE project_id = ?
|
|
`).run(fileId, now, job.project_id);
|
|
this.database.prepare("INSERT INTO project_states (project_id, state_version, name, canvas_json, created_at) VALUES (?, ?, ?, ?, ?)")
|
|
.run(job.project_id, project.state_version + 1, project.name, JSON.stringify(canvas), now);
|
|
this.consumeStorage(operationId, output.bytes.byteLength, now);
|
|
this.finalizeCredit(job.generation_id, "succeeded");
|
|
this.database.prepare(`
|
|
UPDATE generation_jobs
|
|
SET status = 'succeeded', error_category = NULL, final_credit_state = 'committed', output_asset_id = ?,
|
|
upstream_outcome_known = 1, upstream_cost_reconciliation = 'not_required', finished_at = ?, updated_at = ?,
|
|
lease_owner = NULL, lease_expires_at = NULL, heartbeat_at = NULL
|
|
WHERE generation_id = ? AND status = 'running'
|
|
`).run(fileId, now, now, job.generation_id);
|
|
const result = this.terminalResult(this.readJob(job.generation_id));
|
|
this.writeReceipt(result, now);
|
|
return result;
|
|
});
|
|
} catch (error) {
|
|
if (renamed && existsSync(destinationPath)) this.queueCompensation(relativePath, output.bytes.byteLength);
|
|
else rmSync(stagingDirectory, { force: true, recursive: true });
|
|
this.releaseStorage(operationId);
|
|
throw error;
|
|
}
|
|
}
|
|
|
|
private isSafeOutput(job: JobRow, output: NormalizedGenerationOutput) {
|
|
const pixels = { "1:1": [1080, 1080], "3:4": [1080, 1440], "4:3": [1440, 1080], "9:16": [1080, 1920] }[job.ratio];
|
|
return output.bytes.byteLength > 0 && output.bytes.byteLength <= maximumOutputBytes && hasExpectedMagic(output)
|
|
&& output.pixelWidth === pixels[0] && output.pixelHeight === pixels[1];
|
|
}
|
|
|
|
private finalizeCredit(generationId: string, outcome: "succeeded" | "failed" | "rejected") {
|
|
const reservation = this.database.prepare("SELECT * FROM credit_reservations WHERE generation_id = ?")
|
|
.get(generationId) as ReservationRow | undefined;
|
|
if (!reservation) throw new Error("generation_credit_reservation_missing");
|
|
if (reservation.status !== "reserved") return;
|
|
const account = this.database.prepare("SELECT available_balance, reserved_balance FROM credit_accounts WHERE user_id = ?")
|
|
.get(reservation.user_id) as { available_balance: number; reserved_balance: number } | undefined;
|
|
if (!account || account.reserved_balance < reservation.amount) throw new Error("generation_credit_invariant_failed");
|
|
const committed = outcome === "succeeded";
|
|
const availableAfter = committed ? account.available_balance : account.available_balance + reservation.amount;
|
|
const reservedAfter = account.reserved_balance - reservation.amount;
|
|
const now = this.clock();
|
|
const operationKey = `generation:${generationId}:${committed ? "commit" : "release"}`;
|
|
this.database.prepare("UPDATE credit_accounts SET available_balance = ?, reserved_balance = ?, updated_at = ? WHERE user_id = ?")
|
|
.run(availableAfter, reservedAfter, now, reservation.user_id);
|
|
this.database.prepare("UPDATE credit_reservations SET status = ?, finalized_at = ? WHERE generation_id = ? AND status = 'reserved'")
|
|
.run(committed ? "committed" : "released", now, generationId);
|
|
this.database.prepare(`
|
|
INSERT INTO credit_ledger (
|
|
ledger_id, user_id, operation_key, entry_type, amount, available_before, available_after,
|
|
reserved_before, reserved_after, created_at, reference_type, reference_id, model_id, reason, entry_status
|
|
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 'generation', ?, ?, NULL, ?)
|
|
`).run(
|
|
randomUUID(), reservation.user_id, operationKey, committed ? "generation_commit" : "generation_release",
|
|
committed ? -reservation.amount : reservation.amount, account.available_balance, availableAfter,
|
|
account.reserved_balance, reservedAfter, now, generationId, reservation.model_id, committed ? "committed" : "released",
|
|
);
|
|
this.database.prepare(`
|
|
INSERT INTO outbox_events (event_id, operation_key, topic, aggregate_type, aggregate_id, payload_json, status, created_at, published_at)
|
|
VALUES (?, ?, ?, 'generation', ?, ?, 'pending', ?, NULL)
|
|
`).run(randomUUID(), operationKey, committed ? "generation_credit_committed" : "generation_credit_released", generationId, JSON.stringify({ outcome }), now);
|
|
}
|
|
|
|
private reserveStorage(operationId: string, bytes: number) {
|
|
this.immediate(() => {
|
|
const state = this.database.prepare("SELECT managed_content_bytes, storage_status FROM local_backend_storage_state WHERE singleton = 1")
|
|
.get() as { managed_content_bytes: number; storage_status: string } | undefined;
|
|
const active = this.database.prepare("SELECT COALESCE(SUM(projected_bytes), 0) AS bytes FROM storage_reservations WHERE status = 'active'")
|
|
.get() as { bytes: number };
|
|
if (!state || state.storage_status !== "active" || state.managed_content_bytes + active.bytes + bytes > hardLimitBytes) {
|
|
throw new Error("generation_storage_unavailable");
|
|
}
|
|
this.database.prepare("INSERT INTO storage_reservations (reservation_id, operation_id, projected_bytes, status, created_at) VALUES (?, ?, ?, 'active', ?)")
|
|
.run(randomUUID(), operationId, bytes, iso(this.clock()));
|
|
});
|
|
}
|
|
|
|
private consumeStorage(operationId: string, bytes: number, now: number) {
|
|
this.database.prepare("UPDATE storage_reservations SET status = 'consumed', resolved_at = ? WHERE operation_id = ? AND status = 'active'")
|
|
.run(iso(now), operationId);
|
|
const state = this.database.prepare("SELECT managed_content_bytes FROM local_backend_storage_state WHERE singleton = 1")
|
|
.get() as { managed_content_bytes: number };
|
|
const next = state.managed_content_bytes + bytes;
|
|
const classification = capacityClass(next);
|
|
this.database.prepare(`
|
|
UPDATE local_backend_storage_state
|
|
SET managed_content_bytes = ?, capacity_notice_level = ?, storage_status = ?, measured_at = ?, version = version + 1
|
|
WHERE singleton = 1
|
|
`).run(next, classification.capacity, classification.status, iso(now));
|
|
}
|
|
|
|
private releaseStorage(operationId: string) {
|
|
try {
|
|
this.database.prepare("UPDATE storage_reservations SET status = 'released', resolved_at = ? WHERE operation_id = ? AND status = 'active'")
|
|
.run(iso(this.clock()), operationId);
|
|
} catch {
|
|
// Startup storage reconciliation remains the fallback if SQLite is unavailable.
|
|
}
|
|
}
|
|
|
|
private queueCompensation(relativePath: string, bytes: number) {
|
|
try {
|
|
this.database.prepare(`
|
|
INSERT OR IGNORE INTO file_cleanup_queue (
|
|
cleanup_id, managed_file_id, relative_path, byte_size, counts_toward_managed, reason, status, created_at
|
|
) VALUES (?, NULL, ?, ?, 0, 'compensation', 'pending', ?)
|
|
`).run(randomUUID(), relativePath, bytes, iso(this.clock()));
|
|
} catch {
|
|
// Startup reconciliation detects unindexed generated files as a final fallback.
|
|
}
|
|
}
|
|
|
|
private readJob(generationId: string) {
|
|
const row = this.database.prepare("SELECT * FROM generation_jobs WHERE generation_id = ? AND submission_ready = 1").get(generationId) as JobRow | undefined;
|
|
if (!row) throw new Error("generation_not_found");
|
|
return row;
|
|
}
|
|
|
|
private terminalResult(row: JobRow): GenerationProcessingResult {
|
|
if (!["succeeded", "failed", "rejected"].includes(row.status)) throw new Error("generation_not_terminal");
|
|
const output = this.database.prepare("SELECT managed_file_id FROM generation_output_assets WHERE generation_id = ?")
|
|
.get(row.generation_id) as { managed_file_id: string } | undefined;
|
|
return {
|
|
category: row.error_category,
|
|
generationId: row.generation_id,
|
|
outputAssetId: output?.managed_file_id ?? null,
|
|
projectId: row.project_id,
|
|
status: row.status as GenerationProcessingResult["status"],
|
|
};
|
|
}
|
|
|
|
private readReceipt(generationId: string) {
|
|
const row = this.database.prepare("SELECT response_json FROM generation_processor_receipts WHERE generation_id = ?")
|
|
.get(generationId) as { response_json: string } | undefined;
|
|
return row ? JSON.parse(row.response_json) as GenerationProcessingResult : undefined;
|
|
}
|
|
|
|
private writeReceipt(result: GenerationProcessingResult, now: number) {
|
|
this.database.prepare("INSERT INTO generation_processor_receipts (generation_id, response_json, created_at) VALUES (?, ?, ?)")
|
|
.run(result.generationId, JSON.stringify(result), now);
|
|
}
|
|
|
|
private immediate<T>(action: () => T) {
|
|
if (this.database.inTransaction) return action();
|
|
this.database.exec("BEGIN IMMEDIATE");
|
|
try {
|
|
const result = action();
|
|
this.database.exec("COMMIT");
|
|
return result;
|
|
} catch (error) {
|
|
if (this.database.inTransaction) this.database.exec("ROLLBACK");
|
|
throw error;
|
|
}
|
|
}
|
|
|
|
private ensureColumn(table: string, column: string, definition: string) {
|
|
const columns = this.database.prepare(`PRAGMA table_info(${table})`).all() as Array<{ name: string }>;
|
|
if (!columns.some((entry) => entry.name === column)) this.database.exec(`ALTER TABLE ${table} ADD COLUMN ${column} ${definition}`);
|
|
}
|
|
|
|
private migrate() {
|
|
this.ensureColumn("generation_jobs", "lease_owner", "TEXT");
|
|
this.ensureColumn("generation_jobs", "lease_expires_at", "INTEGER");
|
|
this.ensureColumn("generation_jobs", "heartbeat_at", "INTEGER");
|
|
this.ensureColumn("generation_jobs", "started_at", "INTEGER");
|
|
this.ensureColumn("generation_jobs", "attempt_no", "INTEGER NOT NULL DEFAULT 0");
|
|
this.ensureColumn("generation_jobs", "output_asset_id", "TEXT");
|
|
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),
|
|
pixel_width INTEGER NOT NULL CHECK (pixel_width > 0),
|
|
pixel_height INTEGER NOT NULL CHECK (pixel_height > 0),
|
|
created_at INTEGER NOT NULL
|
|
);
|
|
CREATE TABLE IF NOT EXISTS generation_processor_receipts (
|
|
generation_id TEXT PRIMARY KEY REFERENCES generation_jobs(generation_id),
|
|
response_json TEXT NOT NULL CHECK (json_valid(response_json)),
|
|
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());
|
|
}
|
|
}
|