atelier/sidecar/generations/queue.ts

import { join } from 'node:path';
import { claimableGenerations, failInterrupted, markFailed, markRunning, markSucceeded } from '../db';
import type { Generation } from '../db';
import { IMAGES_DIR, INPUTS_DIR, MAX_CONCURRENT_GENERATIONS } from '../env';
import { adapterByName } from '../providers';
import { keyFor } from '../providers/credentials';
import { inspectImage } from './inspect';
import type { ImageEditRequest, ImageRequest, UploadedImage } from '../providers';

/**
 * The worker pool.
 *
 * Postgres holds the queue, not memory. A row is the job, so a restart loses nothing that was waiting
 * and `failInterrupted()` can tell the caller the truth about what was mid-flight. The only in-memory
 * state is how many workers are awake, which is exactly the state that should not survive a restart.
 */
let running = 0;
let sweeping = false;

/** Wake the pool. Safe to call on every enqueue and after every completion; it settles on its own. */
export function pump(): void {
  if (sweeping) return;
  sweeping = true;
  void sweep().finally(() => {
    sweeping = false;
  });
}

async function sweep(): Promise<void> {
  while (running < MAX_CONCURRENT_GENERATIONS) {
    const free = MAX_CONCURRENT_GENERATIONS - running;
    const candidates = await claimableGenerations(free);
    if (candidates.length === 0) return;

    let started = 0;
    for (const candidate of candidates) {
      // markRunning only transitions from 'queued', so two sweeps racing for one row cannot both win
      // and the image is never paid for twice.
      const claimed = await markRunning(candidate.id);
      if (!claimed) continue;

      started += 1;
      running += 1;
      void run(claimed).finally(() => {
        running -= 1;
        pump();
      });
    }

    // Everything on offer was taken by someone else; nothing to do until the next enqueue.
    if (started === 0) return;
  }
}

async function run(row: Generation): Promise<void> {
  const started = Date.now();

  try {
    const adapter = adapterByName(row.provider);
    if (!adapter) throw new Error(`Unknown adapter "${row.provider}"`);

    // The adapter and model come from the ROW's own columns, and the key is read at RUN time: a key
    // rotated between queueing and running should be the new one, and a row queued before a key existed
    // at all runs the moment there is one.
    const apiKey = keyFor(row.provider);
    if (!apiKey) throw new Error(`No API key for provider "${row.provider}"`);

    const result =
      row.operation === 'edit'
        ? await runEdit(row, adapter, apiKey)
        : await adapter.generate(
            { ...(row.options ?? {}), prompt: row.prompt, model: row.model, n: 1 } as ImageRequest,
            apiKey,
          );

    const image = result.images[0];
    if (!image) throw new Error('Provider returned no image');

    // The row id is the filename: one lookup gets you from a served URL back to the prompt that made it.
    const filename = `${row.id}.${image.format}`;
    const path = join(IMAGES_DIR, filename);
    await Bun.write(path, image.bytes);

    // Measured after the write, from the file itself.
    const facts = await inspectImage(path);

    await markSucceeded(row.id, {
      imagePath: filename,
      ...facts,
      costUsd: adapter.costUsd(row.model, result.usage),
      usage: result.usage ?? null,
      result: { ...result.echoed, revised_prompt: image.revisedPrompt },
      durationMs: Date.now() - started,
    });
  } catch (error) {
    await markFailed(row.id, error instanceof Error ? error.message : String(error), Date.now() - started);
  }
}

const MIME: Record<string, string> = {
  '.png': 'image/png',
  '.webp': 'image/webp',
  '.jpg': 'image/jpeg',
  '.jpeg': 'image/jpeg',
};

/** Read the uploads back off disk and hand them to the provider. */
async function runEdit(row: Generation, adapter: NonNullable<ReturnType<typeof adapterByName>>, apiKey: string) {
  if (!adapter.edit) throw new Error(`Adapter "${adapter.name}" cannot edit images`);

  const load = async (name: string): Promise<UploadedImage> => {
    const file = Bun.file(join(INPUTS_DIR, name));
    if (!(await file.exists())) throw new Error(`Input image ${name} is missing`);
    const extension = name.slice(name.lastIndexOf('.')).toLowerCase();
    return { filename: name, bytes: new Uint8Array(await file.arrayBuffer()), type: MIME[extension] ?? 'image/png' };
  };

  const images = await Promise.all((row.inputImages ?? []).map(load));
  const mask = row.maskImage ? await load(row.maskImage) : undefined;

  const request = { ...(row.options ?? {}), prompt: row.prompt, model: row.model } as ImageEditRequest;
  return adapter.edit(request, images, mask, apiKey);
}

/**
 * Called once at boot, before the pool starts.
 *
 * Rows left `running` have no worker behind them — whatever was going to finish them died with the
 * last process. Left alone they stay running forever and a caller polling one waits on an answer that
 * is never coming.
 */
export async function startQueue(): Promise<void> {
  // Once per process, matching the standalone repo's own guard against `bun --hot` double-starting the
  // sweep loop. This process is not hot-reloaded under PM2, but the guard is free and correct either way.
  const globalForQueue = globalThis as unknown as { __atelierQueueStarted?: boolean };
  if (globalForQueue.__atelierQueueStarted) return;
  globalForQueue.__atelierQueueStarted = true;

  const interrupted = await failInterrupted();
  if (interrupted > 0) console.log(`[atelier:queue] failed ${interrupted} generation(s) interrupted by a restart`);
  pump();
}