NianAIGC/lib/server/public-api-jobs.ts

import { createHash } from "node:crypto";
import { assemblePrompt, type PromptAssemblyInput, type PromptMaterial } from "@/lib/prompt/assembler";
import { findGenerationJobByIdempotency } from "@/lib/server/data-store";
import { submitImageJob, type SubmitImageJobInput } from "@/lib/server/generation-service";
import { DEFAULT_OWNER_ID } from "@/lib/server/runtime";
import { submitVideoJob, type SubmitVideoJobInput } from "@/lib/server/video-generation-service";
import type { PublicApiClient } from "@/lib/server/public-api-auth";
import type { EnabledImageCapability, GenerationCapability, GenerationJob } from "@/lib/types";

export class PublicApiConflictError extends Error {
  status = 409;

  constructor(message: string) {
    super(message);
    this.name = "PublicApiConflictError";
  }
}

export type PublicJobCreateBody = {
  capability?: GenerationCapability;
  prompt?: string;
  inputUrls?: string[];
  imageUrls?: string[];
  inputAssetIds?: string[];
  materials?: PromptMaterial[];
  promptAssembly?: PromptAssemblyInput;
  settings?: SubmitVideoJobInput["settings"];
  priority?: number;
  webhookUrl?: string;
  idempotencyKey?: string;
  scale?: number;
  width?: number;
  height?: number;
  min_ratio?: number;
  max_ratio?: number;
  force_single?: boolean;
  resolution?: "4k" | "8k";
  quality?: string;
  seed?: number;
};

export async function createPublicGenerationJob(input: {
  client: PublicApiClient;
  body: PublicJobCreateBody;
  request: Request;
  origin: string;
}): Promise<{ job: GenerationJob; reused: boolean }> {
  const capability = input.body.capability || "image.generate";
  const idempotencyKey = input.request.headers.get("idempotency-key") || input.body.idempotencyKey;
  const fingerprint = idempotencyKey ? fingerprintBody(input.body) : undefined;
  if (idempotencyKey && fingerprint) {
    const existing = await findGenerationJobByIdempotency(input.client.id, idempotencyKey);
    if (existing) {
      if (existing.idempotencyFingerprint !== fingerprint) {
        throw new PublicApiConflictError("Idempotency key was already used with a different request body.");
      }
      return { job: existing, reused: true };
    }
  }

  const common = {
    ownerId: DEFAULT_OWNER_ID,
    externalClientId: input.client.id,
    idempotencyKey,
    idempotencyFingerprint: fingerprint,
    priority: normalizePriority(input.body.priority),
    webhookUrl: normalizeWebhookUrl(input.body.webhookUrl),
    maxAttempts: 3
  };

  if (capability === "video.generate") {
    const job = await submitVideoJob({
      ...input.body,
      ...common,
      mode: "video",
      materials: input.body.materials || input.body.promptAssembly?.materials || []
    } as SubmitVideoJobInput, input.origin);
    return { job, reused: false };
  }

  const imageCapability = normalizeImageCapability(capability);
  const assembled = input.body.promptAssembly
    ? assemblePrompt({
        ...input.body.promptAssembly,
        mode: "image",
        materials: input.body.materials || input.body.promptAssembly.materials || []
      })
    : undefined;
  const materialImages = (input.body.materials || assembled?.materials || [])
    .filter((material) => material.type === "image")
    .map((material) => material.url);
  const job = await submitImageJob({
    ...common,
    capability: imageCapability,
    prompt: input.body.prompt || assembled?.prompt,
    imageUrls: input.body.imageUrls || input.body.inputUrls || materialImages,
    inputAssetIds: input.body.inputAssetIds || (input.body.materials || []).map((material) => material.id).filter(Boolean) as string[],
    scale: asNumber(input.body.scale),
    width: asNumber(input.body.width),
    height: asNumber(input.body.height),
    min_ratio: asNumber(input.body.min_ratio),
    max_ratio: asNumber(input.body.max_ratio),
    force_single: Boolean(input.body.force_single),
    resolution: input.body.resolution,
    quality: normalizeQuality(input.body.quality),
    seed: asNumber(input.body.seed)
  } satisfies SubmitImageJobInput, input.origin);
  return { job, reused: false };
}

function normalizeImageCapability(capability: GenerationCapability): EnabledImageCapability {
  if (capability === "image.generate" || capability === "image.inpaint" || capability === "image.upscale") {
    return capability;
  }
  throw new Error(`Unsupported image capability: ${capability}`);
}

function normalizePriority(value: unknown): number {
  const parsed = Number(value);
  if (!Number.isFinite(parsed)) return 0;
  return Math.max(-100, Math.min(100, Math.trunc(parsed)));
}

function normalizeWebhookUrl(value: unknown): string | undefined {
  if (typeof value !== "string" || !value.trim()) return undefined;
  const url = new URL(value.trim());
  if (!["http:", "https:"].includes(url.protocol)) throw new Error("webhookUrl must be an HTTP or HTTPS URL.");
  return url.toString();
}

function asNumber(value: unknown): number | undefined {
  if (value === undefined || value === null || value === "") return undefined;
  const parsed = Number(value);
  return Number.isFinite(parsed) ? parsed : undefined;
}

function normalizeQuality(value: unknown): string | undefined {
  if (typeof value !== "string") return undefined;
  const normalized = value.trim().toLowerCase();
  return ["low", "medium", "high"].includes(normalized) ? normalized : undefined;
}

function fingerprintBody(body: PublicJobCreateBody): string {
  const { idempotencyKey: _idempotencyKey, ...fingerprintSource } = body;
  return createHash("sha256").update(stableStringify(fingerprintSource)).digest("hex");
}

function stableStringify(value: unknown): string {
  if (Array.isArray(value)) return `[${value.map(stableStringify).join(",")}]`;
  if (value && typeof value === "object") {
    return `{${Object.entries(value as Record<string, unknown>)
      .sort(([left], [right]) => left.localeCompare(right))
      .map(([key, item]) => `${JSON.stringify(key)}:${stableStringify(item)}`)
      .join(",")}}`;
  }
  return JSON.stringify(value);
}