atelier/sidecar/db/generations/schema.ts

import { pgTable, uuid, text, jsonb, integer, numeric, boolean, timestamp, index, check } from 'drizzle-orm/pg-core';
import { sql } from 'drizzle-orm';

/**
 * The set of states a generation can be in. Closed, so it is a `check()` rather than free text —
 * otherwise a status the code stopped writing sits in old rows with nothing structural saying it
 * became illegal.
 */
export const GENERATION_STATUSES = ['queued', 'running', 'succeeded', 'failed'] as const;
export type GenerationStatus = (typeof GENERATION_STATUSES)[number];

/**
 * Which provider endpoint made the row. Closed, so a check() rather than free text.
 *
 * One table for both rather than a second one for edits: everything that makes a generation worth
 * recording — prompt, model, cost, duration, what landed on disk — is identical, and history that has
 * to be read from two places is history that gets read from one.
 */
export const GENERATION_OPERATIONS = ['generate', 'edit'] as const;
export type GenerationOperation = (typeof GENERATION_OPERATIONS)[number];

/**
 * One row per image we were asked to make, whatever became of it.
 *
 * The row is written when the request arrives, not when it succeeds: a generation that failed, or that
 * a restart interrupted, is a thing that happened and cost time, and the prompt that produced it is
 * worth as much as one that worked.
 *
 * ── Why the image is a path and not bytes ──
 *
 * The files live under ATELIER_DATA_DIR, outside the repo and outside Postgres. They are served
 * unauthenticated as ordinary images, so they need to be a file a static route can hand to
 * `new Response(Bun.file(...))` — not a column something has to decode first.
 */
export const generations = pgTable(
  'generations',
  {
    /**
     * Also the filename on disk. One id means one image, so a row can be turned into a served URL
     * without a second lookup.
     */
    id: uuid('id').primaryKey().defaultRandom(),

    /**
     * The request these images came from — `n: 4` is four rows sharing one requestId, and the job the
     * caller polls. Per-image rows rather than an array column because cost, duration and failure are
     * per image: one of four can fail, and a row that averaged them would say nothing true.
     */
    requestId: uuid('request_id').notNull(),

    /** Position within the request, so `data[]` comes back in the order it was asked for. */
    index: integer('index').notNull().default(0),

    operation: text('operation').notNull().default('generate').$type<GenerationOperation>(),

    status: text('status').notNull().default('queued').$type<GenerationStatus>(),

    /**
     * The adapter and model, as text, denormalised on purpose.
     *
     * `providerId` below points at the configured provider, but that row can be renamed, repointed at
     * another model, or deleted — and none of that may change what an existing image was made with.
     * The row is a pointer to configuration; these two are the historical fact.
     */
    provider: text('provider').notNull(),
    model: text('model').notNull(),

    /** The provider row this came from, when there was one. Null for rows made before providers existed. */
    providerId: uuid('provider_id'),

    prompt: text('prompt').notNull(),

    /** Whatever else the caller asked for — size, quality, background. Provider-shaped, so opaque here. */
    options: jsonb('options').$type<Record<string, unknown>>(),

    /**
     * For edits: the uploaded files this was derived from, as names under the inputs directory.
     *
     * Kept because an edit without its input is unreadable afterwards — the prompt alone does not say
     * what was edited, and "make the hat blue" means nothing without the hat. Names rather than bytes,
     * for the same reason the output is a file.
     */
    inputImages: jsonb('input_images').$type<string[]>(),
    maskImage: text('mask_image'),

    /** Filename under the images directory. Null until the bytes exist. */
    imagePath: text('image_path'),

    // ── What actually landed on disk ──
    //
    // Measured from the written file, never echoed from the request. What was asked for and what
    // arrived are exactly the two things that come apart, and the caller is an agent that cannot look
    // at the image to tell the difference.
    //
    // `hasAlpha` and `opaque` are deliberately both here and they are not the same question.
    // `hasAlpha` is whether the file carries an alpha channel; `opaque` is whether every pixel in it
    // is fully opaque. A PNG with an alpha channel that the model never used is hasAlpha: true,
    // opaque: true — which is a request for a transparent background that silently did not happen,
    // and the single most consequential failure for an asset destined to be an app icon.
    width: integer('width'),
    height: integer('height'),
    hasAlpha: boolean('has_alpha'),
    opaque: boolean('opaque'),

    /**
     * Cost in USD. `numeric` rather than a float: money that goes through binary floating point comes
     * back slightly wrong, and these get summed.
     */
    costUsd: numeric('cost_usd', { precision: 12, scale: 6 }),

    /** Raw usage the provider reported, so cost can be recomputed if the pricing we applied was wrong. */
    usage: jsonb('usage').$type<Record<string, unknown>>(),

    /** What the provider said back, minus the image bytes — `revised_prompt`, echoed size and quality. */
    result: jsonb('result').$type<Record<string, unknown>>(),

    durationMs: integer('duration_ms'),

    error: text('error'),

    createdAt: timestamp('created_at', { withTimezone: true }).notNull().defaultNow(),
    startedAt: timestamp('started_at', { withTimezone: true }),
    finishedAt: timestamp('finished_at', { withTimezone: true }),
  },
  (table) => [
    check('generations_status_check', sql`${table.status} in ('queued', 'running', 'succeeded', 'failed')`),
    check('generations_operation_check', sql`${table.operation} in ('generate', 'edit')`),
    // The queue reads "what is still to do, oldest first" on every tick and after every restart.
    index('generations_status_created_idx').on(table.status, table.createdAt),
    // Every poll of a job is this lookup.
    index('generations_request_idx').on(table.requestId, table.index),
  ],
);