atelier/sidecar/api/images.ts

import { randomUUID } from 'node:crypto';
import { extname, join } from 'node:path';
import {
  addToCollection,
  createRequest,
  findOrCreateCollection,
  getCollection,
  getRequest,
  listGenerations,
} from '../db';
import type { NewGeneration } from '../db';
import { createRouter } from '../create-router';
import { resolveForRequest } from '../providers/credentials';
import { validateImageEditRequest, validateImageRequest } from '../generations/validate';
import { INPUTS_DIR } from '../env';
import { pump } from '../generations/queue';
import { imagesResponse, isSettled, jobStatus } from '../generations/response';
import { serializeGeneration } from '../serialize';

export const imagesRouter = createRouter();

/**
 * The longest a caller may ask us to hold a connection.
 *
 * Three timeouts sit in front of the answer and the innermost has to be the one that fires, or the
 * caller gets a dropped connection instead of a status:
 *
 *   nginx  proxy_read_timeout  320s   (data/nginx/custom/server_proxy.conf)
 *   Bun    idleTimeout         255s   (server.ts — 255 is Bun's hard ceiling, 256 is rejected)
 *   here   MAX_WAIT_SECONDS    240s
 *
 * 240 rather than the 300 the documentation once promised: Bun cannot be told to wait longer than
 * 255, so a 300s wait would be cut by our own runtime at 255 with no body — the exact failure this
 * ordering exists to prevent. The ceiling is Bun's, not the proxy's, and raising it further means
 * changing runtimes rather than configuration.
 *
 * Was 55 until the proxy was raised from its 60s default, which forced every caller to poll two to
 * six times for a single 40s generation.
 */
const MAX_WAIT_SECONDS = 240;

/**
 * Resolve `?collection=` to a collection, creating one if the value is a name we have not seen.
 *
 * A query parameter and not a body field, deliberately. The body is documented as OpenAI's, byte for
 * byte, so that a client written against their API needs no changes — putting an Atelier-only key in
 * it would make that false for the sake of one string.
 *
 * A uuid must already exist: it can only have come from us, so a miss is a real mistake and answering
 * it by creating a collection named after a uuid helps nobody. A name is find-or-create, which is the
 * whole point — "put this afternoon's work in one place" should not need a setup step first.
 */
async function resolveCollection(value: string | undefined) {
  const name = value?.trim();
  if (!name) return { collection: undefined };

  if (UUID_PATTERN.test(name)) {
    const collection = await getCollection(name);
    return collection ? { collection } : { error: `No collection with id "${name}".` };
  }

  return { collection: await findOrCreateCollection(name) };
}

import { openaiError, UUID_PATTERN } from '../errors';

/**
 * `POST /api/images/generations` — OpenAI's request body, verbatim.
 *
 * Answers 202 with a job rather than the image. OpenAI's own endpoint blocks for the length of the
 * generation; holding that open for `n: 10` would tie a connection to work that takes minutes, and the
 * callers here are agents that can poll. The finished job returns their exact response shape.
 */
imagesRouter.post('/generations', async (ctx) => {
  const body = await ctx.req.json().catch(() => null);

  const parsed = validateImageRequest(body);
  if ('error' in parsed) {
    return ctx.json(openaiError(parsed.error.message, 'invalid_request_error', parsed.error.param), 400);
  }

  const request = parsed.request;

  // `?provider=` names one explicitly; otherwise the model decides, falling back to the default.
  const resolvedProvider = await resolveForRequest({ providerRef: ctx.req.query('provider'), model: request.model });
  if ('error' in resolvedProvider) {
    const type = resolvedProvider.status === 404 ? 'not_found' : 'configuration_error';
    return ctx.json(openaiError(resolvedProvider.error, type), resolvedProvider.status);
  }
  const { provider, adapter } = resolvedProvider;

  const resolved = await resolveCollection(ctx.req.query('collection'));
  if (resolved.error) return ctx.json(openaiError(resolved.error, 'not_found', 'collection'), 404);

  // The request's model wins over the provider's, so a provider can serve a model nobody configured.
  const model = request.model ?? provider.model;
  const requestId = randomUUID();
  const count = request.n ?? 1;

  // Everything except prompt/model/n — those have columns of their own, and duplicating them into
  // options would give two answers to the same question the first time one was edited.
  const { prompt, model: _model, n: _n, ...options } = request;

  const rows: NewGeneration[] = Array.from({ length: count }, (_, index) => ({
    requestId,
    index,
    status: 'queued' as const,
    provider: adapter.name,
    providerId: provider.id,
    model,
    prompt,
    options,
  }));

  const created = await createRequest(rows);

  // Filed at request time rather than on completion: the collection is a statement about what was
  // asked for, and a generation that fails is still part of the afternoon it belonged to.
  if (resolved.collection)
    await addToCollection(
      resolved.collection.id,
      created.map((row) => row.id),
    );

  pump();

  return ctx.json(
    {
      id: requestId,
      status: 'queued',
      created: Math.floor(Date.now() / 1000),
      n: count,
      ...(resolved.collection ? { collection: { id: resolved.collection.id, name: resolved.collection.name } } : {}),
    },
    202,
  );
});

/**
 * `GET /api/images/generations/:id` — poll, or wait.
 *
 * `?wait=<seconds>` holds the request open until the job settles, up to 300s. Agents are the callers
 * here and a long poll spares every one of them from writing the same retry loop; without the
 * parameter it answers immediately.
 */
imagesRouter.get('/generations/:id', async (ctx) => {
  const id = ctx.req.param('id');

  // Checked before the query, because Postgres rejects a malformed uuid at the cast and the failure
  // surfaces as a 500 — a client mistake reported as a server fault, outside the error envelope.
  // A caller cannot tell "never existed" from "does not exist" and should not have to.
  if (!UUID_PATTERN.test(id)) {
    return ctx.json(openaiError(`No generation request with id "${id}".`, 'not_found'), 404);
  }

  const waitParam = Number(ctx.req.query('wait') ?? 0);
  const waitMs = Math.min(Math.max(Number.isFinite(waitParam) ? waitParam : 0, 0), MAX_WAIT_SECONDS) * 1000;

  const deadline = Date.now() + waitMs;

  for (;;) {
    const rows = await getRequest(id);
    if (rows.length === 0) return ctx.json(openaiError(`No generation request with id "${id}".`, 'not_found'), 404);

    const status = jobStatus(rows);
    if (isSettled(status) || Date.now() >= deadline) {
      if (!isSettled(status)) {
        return ctx.json({ id, status, created: Math.floor(rows[0]!.createdAt.getTime() / 1000) }, 202);
      }
      const includeBase64 = (rows[0]!.options as { response_format?: string } | null)?.response_format === 'b64_json';
      return ctx.json(await imagesResponse(rows, includeBase64), status === 'failed' ? 500 : 200);
    }

    await new Promise((resolve) => setTimeout(resolve, 500));
  }
});

/** Everything generated, newest first. Not OpenAI's — the history is ours. */
imagesRouter.get('/', async (ctx) => {
  const limit = Math.min(Math.max(Number(ctx.req.query('limit') ?? 50), 1), 200);
  const rows = await listGenerations(limit);
  return ctx.json({ data: rows.map(serializeGeneration) });
});

/** png / webp / jpg, per what the endpoint accepts. Kept from the upload so the provider sees a real name. */
const SAFE_EXTENSION = /^\.(png|jpe?g|webp)$/i;

/**
 * `POST /api/images/edits` — OpenAI's multipart body, verbatim.
 *
 * Same job model as generations: 202 with an id, the finished job returns their response shape.
 *
 * The uploads are written to disk before the row is queued, and that ordering is the point. The work
 * happens after the request has returned, so bytes held only in the request's memory would be gone by
 * the time anything needed them — and an edit whose input has vanished cannot even be retried.
 */
imagesRouter.post('/edits', async (ctx) => {
  let form: FormData;
  try {
    form = await ctx.req.formData();
  } catch {
    return ctx.json(openaiError('Request must be multipart/form-data.', 'invalid_request_error'), 400);
  }

  // Both `image` and `image[]` — the endpoint takes either, and clients send both spellings.
  const uploads = [...form.getAll('image'), ...form.getAll('image[]')].filter(
    (value): value is File => value instanceof File && value.size > 0,
  );

  if (uploads.length === 0) {
    return ctx.json(openaiError("'image' is required.", 'invalid_request_error', 'image'), 400);
  }
  if (uploads.length > 16) {
    return ctx.json(openaiError('At most 16 images may be provided.', 'invalid_request_error', 'image'), 400);
  }

  const maskValue = form.get('mask');
  const mask = maskValue instanceof File && maskValue.size > 0 ? maskValue : undefined;

  const fields: Record<string, string> = {};
  for (const [key, value] of form.entries()) {
    if (typeof value === 'string') fields[key] = value;
  }

  const parsed = validateImageEditRequest(fields);
  if ('error' in parsed) {
    return ctx.json(openaiError(parsed.error.message, 'invalid_request_error', parsed.error.param), 400);
  }

  const request = parsed.request;
  const resolvedProvider = await resolveForRequest({ providerRef: ctx.req.query('provider'), model: request.model });
  if ('error' in resolvedProvider) {
    const type = resolvedProvider.status === 404 ? 'not_found' : 'configuration_error';
    return ctx.json(openaiError(resolvedProvider.error, type), resolvedProvider.status);
  }
  const { provider, adapter } = resolvedProvider;

  if (!adapter.edit) {
    return ctx.json(openaiError(`Adapter "${adapter.name}" cannot edit images.`, 'invalid_request_error'), 400);
  }

  const store = async (file: File) => {
    const extension = SAFE_EXTENSION.test(extname(file.name)) ? extname(file.name).toLowerCase() : '.png';
    const name = `${randomUUID()}${extension}`;
    await Bun.write(join(INPUTS_DIR, name), await file.arrayBuffer());
    return name;
  };

  const inputImages = await Promise.all(uploads.map(store));
  const maskImage = mask ? await store(mask) : null;

  const resolved = await resolveCollection(ctx.req.query('collection'));
  if (resolved.error) return ctx.json(openaiError(resolved.error, 'not_found', 'collection'), 404);

  const model = request.model ?? provider.model;
  const requestId = randomUUID();
  const count = request.n ?? 1;

  const { prompt, model: _model, n: _n, ...options } = request;

  const created = await createRequest(
    Array.from({ length: count }, (_, index) => ({
      requestId,
      index,
      operation: 'edit' as const,
      status: 'queued' as const,
      provider: adapter.name,
      providerId: provider.id,
      model,
      prompt,
      options,
      inputImages,
      maskImage,
    })),
  );

  if (resolved.collection)
    await addToCollection(
      resolved.collection.id,
      created.map((row) => row.id),
    );

  pump();

  return ctx.json(
    {
      id: requestId,
      status: 'queued',
      created: Math.floor(Date.now() / 1000),
      n: count,
      ...(resolved.collection ? { collection: { id: resolved.collection.id, name: resolved.collection.name } } : {}),
    },
    202,
  );
});