evOo-local/packages/adapters/src/media/image-operations.postgres.test.ts

import { randomUUID } from "node:crypto";
import { mkdtemp, rm } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import type { JobPublisher } from "@rakazo/adapter-kit";
import type { MediaJobBlock, MessageBlock } from "@rakazo/contracts";
import { ComfyUiSettingsSchema } from "@rakazo/contracts";
import type { PrismaClient } from "@rakazo/db";
import { createDb } from "@rakazo/db";
import { afterAll, beforeAll, describe, expect, it } from "vitest";
import { LocalArtifactStore } from "../artifacts.js";
import { createMediaRuntime, MediaJobRejectedError } from "../media-runtime.js";
import { ComfyUiEmulator } from "../testing/comfyui-emulator.js";
import { createImageOperations } from "./image-operations.js";

const describePostgres =
  process.env.VERIFY_DATABASE && process.env.DATABASE_URL ? describe.sequential : describe.skip;

describePostgres("ComfyUI operations in the media runtime (PostgreSQL + ComfyUI emulator)", () => {
  let db: ReturnType<typeof createDb>;
  let prisma: PrismaClient;
  let root: string;
  const users: string[] = [];

  beforeAll(async () => {
    db = createDb(process.env.DATABASE_URL!);
    prisma = db.prisma;
    root = await mkdtemp(join(tmpdir(), "image-operations-pg-"));
  });

  afterAll(async () => {
    if (!db) return;
    await prisma.mediaJob.deleteMany({ where: { userId: { in: users } } });
    await prisma.artifact.deleteMany({ where: { userId: { in: users } } });
    await prisma.organization.deleteMany({ where: { id: { in: users } } });
    await prisma.user.deleteMany({ where: { id: { in: users } } });
    await prisma.$disconnect();
    await db.pool.end();
    await rm(root, { recursive: true, force: true });
  });

  async function setup() {
    const id = randomUUID();
    users.push(id);
    await prisma.user.create({
      data: { id, name: "Media test", email: `${id}@example.test`, emailVerified: true },
    });
    await prisma.organization.create({
      data: { id, name: "Media test", slug: id, createdAt: new Date() },
    });
    await prisma.space.create({ data: { id, organizationId: id, name: "Media", isDefault: true } });
    await prisma.member.create({
      data: { id, organizationId: id, userId: id, role: "owner", createdAt: new Date() },
    });
    const bot = await prisma.bot.create({
      data: { spaceId: id, userId: id, name: "Studio", color: "test-color" },
    });
    const thread = await prisma.thread.create({ data: { spaceId: id, userId: id, botId: bot.id } });
    const emulator = new ComfyUiEmulator();
    emulator.pendingPolls = 0;
    const jobs: JobPublisher = {
      enqueue: async () => undefined,
      cancel: async () => undefined,
      close: async () => undefined,
    };
    const runtime = createMediaRuntime(
      {
        prisma,
        events: { notify: async () => undefined },
        jobs,
        artifacts: new LocalArtifactStore(join(root, id, "artifacts")),
        registrations: createImageOperations({
          loadSettings: async () => ({
            comfyui: ComfyUiSettingsSchema.parse({ baseUrl: "http://127.0.0.1:8188" }),
            transcription: null,
          }),
          fetch: emulator.fetch,
        }),
      },
      { tempRoot: join(root, id, "jobs") },
    );
    const scope = {
      spaceId: id,
      userId: id,
      botId: bot.id,
      groupId: null,
      threadId: thread.id,
      runId: null,
    };
    return { emulator, runtime, scope };
  }

  async function card(messageId: string, jobId: string): Promise<MediaJobBlock | undefined> {
    const message = await prisma.message.findUniqueOrThrow({ where: { id: messageId } });
    return (message.blocks as MessageBlock[]).find(
      (block): block is MediaJobBlock => block.kind === "media" && block.jobId === jobId,
    );
  }

  it("generates images onto the card and edits one into a linked version", async () => {
    const { runtime, scope } = await setup();
    const generated = await runtime.start({
      ...scope,
      operation: "generate_image",
      input: { prompt: "a lighthouse at dusk", aspect: "landscape", count: 2, seed: 3 },
    });
    await runtime.runJob(generated.jobId);
    const job = await runtime.get(generated.jobId);
    expect(job).toMatchObject({ status: "done", title: "a lighthouse at dusk" });
    expect(job?.outputs.map((output) => [output.mimeType, output.width, output.height])).toEqual([
      ["image/png", 1344, 768],
      ["image/png", 1344, 768],
    ]);
    expect(await card(generated.messageId, generated.jobId)).toMatchObject({
      status: "done",
      outputs: job?.outputs,
    });

    const source = job!.outputs[0]!.artifactId;
    const edited = await runtime.start({
      ...scope,
      operation: "edit_image",
      input: { artifact_id: source, resize: { aspect: "vertical", fit: "blur" }, format: "jpeg" },
    });
    await runtime.runJob(edited.jobId);
    const edit = await runtime.get(edited.jobId);
    expect(edit).toMatchObject({ status: "done", title: "9:16 · JPEG" });
    const [output] = edit!.outputs;
    expect(output).toMatchObject({ mimeType: "image/jpeg", width: 1080, height: 1920 });
    const artifact = await prisma.artifact.findUniqueOrThrow({ where: { id: output!.artifactId } });
    expect(artifact).toMatchObject({ parentArtifactId: source, width: 1080, height: 1920 });
  });

  it("refuses source artifacts the user cannot read", async () => {
    const { runtime, scope } = await setup();
    const other = await setup();
    const generated = await other.runtime.start({
      ...other.scope,
      operation: "generate_image",
      input: { prompt: "private" },
    });
    await other.runtime.runJob(generated.jobId);
    const foreign = (await other.runtime.get(generated.jobId))!.outputs[0]!.artifactId;
    await expect(
      runtime.start({
        ...scope,
        operation: "generate_video",
        input: { prompt: "animate it", source_artifact_id: foreign },
      }),
    ).rejects.toThrow(MediaJobRejectedError);
    await expect(
      runtime.start({
        ...scope,
        operation: "edit_image",
        input: { artifact_id: foreign, rotate: 90 },
      }),
    ).rejects.toThrow(MediaJobRejectedError);
  });

  it("waits while the GPU lacks free memory", async () => {
    const { emulator, runtime, scope } = await setup();
    emulator.gpu = { totalMb: 24_576, freeMb: 4_000, torchReservedMb: 0, torchFreeMb: 0 };
    const video = await runtime.start({
      ...scope,
      operation: "generate_video",
      input: { prompt: "rain on a window" },
    });
    await runtime.runJob(video.jobId);
    expect(await runtime.get(video.jobId)).toMatchObject({ status: "waiting", outputs: [] });
    expect(emulator.prompts).toHaveLength(0);
  });
});