atelier/sidecar/generations/response.ts

import { join } from 'node:path';
import type { Generation } from '../db';
import { IMAGES_DIR } from '../env';
import { serializeGenerationImage } from '../serialize';

export type JobStatus = 'queued' | 'running' | 'succeeded' | 'failed' | 'partial';

/**
 * The status of a request, from the rows of its images.
 *
 * `partial` exists because `n` images can end differently: three succeed and one trips the content
 * filter. Reporting that as `failed` throws away three images that exist and were paid for, and
 * reporting it as `succeeded` hides a failure the caller should see.
 */
export function jobStatus(rows: Generation[]): JobStatus {
  if (rows.some((row) => row.status === 'running')) return 'running';
  if (rows.some((row) => row.status === 'queued')) return 'queued';

  const failed = rows.filter((row) => row.status === 'failed').length;
  if (failed === 0) return 'succeeded';
  if (failed === rows.length) return 'failed';
  return 'partial';
}

export const isSettled = (status: JobStatus) => status === 'succeeded' || status === 'failed' || status === 'partial';

/** `/files/<name>`, relative from this plugin's own mount. See serialize.ts for why. */
export const imageUrl = (row: Generation) => `/files/${row.imagePath}`;

/**
 * OpenAI's `ImagesResponse`, built from our rows.
 *
 * Deliberate divergence from OpenAI's own shape, and the reason Atelier exists: `url` is always
 * populated, including for the GPT image models, where OpenAI returns base64 only. `b64_json` is
 * included only when the caller asked for it, so the common response stays small.
 */
export async function imagesResponse(rows: Generation[], includeBase64: boolean) {
  const succeeded = rows.filter((row) => row.status === 'succeeded' && row.imagePath);

  const data = await Promise.all(
    succeeded.map(async (row) => {
      const item: Record<string, unknown> = { url: imageUrl(row) };

      const revised = (row.result as { revised_prompt?: string } | null)?.revised_prompt;
      if (revised) item.revised_prompt = revised;

      if (includeBase64) {
        const bytes = await Bun.file(join(IMAGES_DIR, row.imagePath!)).arrayBuffer();
        item.b64_json = Buffer.from(bytes).toString('base64');
      }

      return item;
    }),
  );

  const first = rows[0];
  const echoed = (first?.result ?? {}) as Record<string, unknown>;

  const totalCost = rows.reduce((sum, row) => sum + Number(row.costUsd ?? 0), 0);

  return {
    created: Math.floor((first?.createdAt?.getTime() ?? Date.now()) / 1000),
    data,
    ...(echoed.background ? { background: echoed.background } : {}),
    ...(echoed.output_format ? { output_format: echoed.output_format } : {}),
    ...(echoed.size ? { size: echoed.size } : {}),
    ...(echoed.quality ? { quality: echoed.quality } : {}),
    usage: mergeUsage(rows),
    // Not OpenAI's, and namespaced so it cannot collide with a field they add later.
    atelier: {
      id: first?.requestId,
      status: jobStatus(rows),
      model: first?.model,
      provider: first?.provider,
      cost_usd: totalCost > 0 ? totalCost.toFixed(6) : null,
      duration_ms: rows.reduce((max, row) => Math.max(max, row.durationMs ?? 0), 0),
      // Through the shared serializer, so a field here and the same field in /images cannot drift
      // into two spellings again.
      images: rows.map(serializeGenerationImage),
    },
  };
}

/** Sum the per-image usage back into one object shaped like OpenAI's. */
function mergeUsage(rows: Generation[]) {
  const totals = { total_tokens: 0, input_tokens: 0, output_tokens: 0, text_tokens: 0, image_tokens: 0 };

  for (const row of rows) {
    const usage = row.usage as {
      total_tokens?: number;
      input_tokens?: number;
      output_tokens?: number;
      input_tokens_details?: Record<string, number>;
    } | null;
    if (!usage) continue;
    totals.total_tokens += usage.total_tokens ?? 0;
    totals.input_tokens += usage.input_tokens ?? 0;
    totals.output_tokens += usage.output_tokens ?? 0;
    totals.text_tokens += usage.input_tokens_details?.text_tokens ?? 0;
    totals.image_tokens += usage.input_tokens_details?.image_tokens ?? 0;
  }

  return {
    total_tokens: totals.total_tokens,
    input_tokens: totals.input_tokens,
    output_tokens: totals.output_tokens,
    input_tokens_details: { text_tokens: totals.text_tokens, image_tokens: totals.image_tokens },
  };
}