AI Image Generator — TypeScript source

// Pure logic for the AI Image Generator (FEAT-068). Zero I/O — the island
// and the /api/image-gen route both depend on this module, so validation,
// gate math, response parsing, and URL state are all testable without a
// browser or a Worker. Mirrored in Go at cli/ai-image-generator (lock-step
// test vectors).

export const MODEL = '@cf/black-forest-labs/flux-1-schnell';
/** flux-1-schnell's design point. Model max is 8, but cost scales at
 *  9.6 neurons/step, so v1 pins 4 → a flat 57.6 neurons per image. */
export const STEPS = 4;
/** Our UX clamp — deliberately below the model's 2048-char limit. */
export const PROMPT_MAX_CHARS = 600;
export const IP_HOUR_CAP = 10;
/** Empirical (2026-09-10, live usage.neurons): one 1024² image = 172.8
 *  neurons — 3× the docs-table estimate. 40/day × 172.8 = 6,912, inside
 *  the 10,000 free allocation with ~30% headroom. */
export const DAILY_CAP = 40;

export type PromptInvalidReason = 'empty' | 'tooLong' | 'controlChars';

export function validatePrompt(
  raw: string
): { ok: true; prompt: string } | { ok: false; reason: PromptInvalidReason } {
  // Normalize line endings first so pasted CRLF/CR text keeps its structure.
  const prompt = raw.replace(/\r\n?/g, '\n').trim();
  if (prompt.length === 0) return { ok: false, reason: 'empty' };
  if (prompt.length > PROMPT_MAX_CHARS) return { ok: false, reason: 'tooLong' };
  // Control characters except \n (multi-line prompts stay legal): C0 minus
  // U+000A, DEL, and C1. A stray \r was already normalized above.
  if (/[\u0000-\u0009\u000b-\u001f\u007f-\u009f]/.test(prompt)) {
    return { ok: false, reason: 'controlChars' };
  }
  return { ok: true, prompt };
}

export function buildModelParams(prompt: string): { prompt: string; steps: number } {
  return { prompt, steps: STEPS };
}

function minutesToUtcMidnight(now: number): number {
  const msLeft = Date.UTC(
    new Date(now).getUTCFullYear(),
    new Date(now).getUTCMonth(),
    new Date(now).getUTCDate() + 1
  ) - now;
  return Math.max(1, Math.ceil(msLeft / 60_000));
}

export type GateVerdict =
  | { ok: true }
  | { ok: false; error: 'RATE_LIMITED_IP' | 'RATE_LIMITED_DAILY'; retryAfterMinutes: number };

export function checkGate(ipCount: number, dailyCount: number, now: number = Date.now()): GateVerdict {
  if (dailyCount >= DAILY_CAP) {
    return { ok: false, error: 'RATE_LIMITED_DAILY', retryAfterMinutes: minutesToUtcMidnight(now) };
  }
  if (ipCount >= IP_HOUR_CAP) {
    return { ok: false, error: 'RATE_LIMITED_IP', retryAfterMinutes: 60 };
  }
  return { ok: true };
}

export type GenError =
  | 'BAD_REQUEST'
  | 'RATE_LIMITED_IP'
  | 'RATE_LIMITED_DAILY'
  | 'MODEL_ERROR'
  | 'MODEL_UNAVAILABLE'
  | 'NETWORK';

export type GenerateResult =
  | { ok: true; dataUrl: string }
  | { ok: false; error: GenError; retryAfterMinutes?: number };

export function parseGenerateResponse(status: number, body: unknown): GenerateResult {
  if (status === 400) return { ok: false, error: 'BAD_REQUEST' };
  if (status === 429) {
    const b = (body ?? {}) as { error?: unknown; retryAfterMinutes?: unknown };
    const error =
      b.error === 'RATE_LIMITED_DAILY' ? 'RATE_LIMITED_DAILY' : 'RATE_LIMITED_IP';
    const retryAfterMinutes =
      typeof b.retryAfterMinutes === 'number' ? b.retryAfterMinutes : 60;
    return { ok: false, error, retryAfterMinutes };
  }
  if (status === 503) return { ok: false, error: 'MODEL_UNAVAILABLE' };
  if (status === 502) return { ok: false, error: 'MODEL_ERROR' };
  if (status === 200) {
    const url = (body as { dataUrl?: unknown } | null)?.dataUrl;
    if (typeof url === 'string' && url.startsWith('data:image/')) {
      return { ok: true, dataUrl: url };
    }
    return { ok: false, error: 'MODEL_ERROR' };
  }
  return { ok: false, error: 'NETWORK' };
}

export function encodeUrlState(prompt: string): string {
  return prompt ? `?prompt=${encodeURIComponent(prompt)}` : '';
}

export function decodeUrlState(search: string): string {
  return new URLSearchParams(search).get('prompt') ?? '';
}