image-studio/commands/generation.test.ts

import { describe, expect, test } from "bun:test";
import type { GenerationCandidate } from "@editor/state";
import { createInitialAppState } from "@editor/initial-state";
import {
  generationAddCandidateCommand,
  generationApplyCandidateAsLayerCommand,
  generationRemoveCandidateCommand,
  generationClearCandidatesCommand,
  generationReuseCandidateSettingsCommand,
  generationReplaceCandidatePixelsCommand,
  generationSetCompareModeCommand,
  generationFailJobCommand,
  generationStartJobCommand,
  generationSucceedJobCommand,
  generationCancelJobCommand,
} from "./generation";

describe("generation commands", () => {
  test("adds, selects, and removes candidates", () => {
    const state = createInitialAppState("Test");
    const candidate = generationCandidate("candidate-1");

    const added = generationAddCandidateCommand.execute({ state }, { candidate });
    const removed = generationRemoveCandidateCommand.execute({ state: added }, { candidateId: candidate.id });

    expect(added.editor.generation.candidates).toEqual([candidate]);
    expect(added.editor.generation.selectedCandidateId).toBe(candidate.id);
    expect(added.editor.generation.compareMode).toBe("result");
    expect(removed.editor.generation.candidates).toEqual([]);
    expect(removed.editor.generation.selectedCandidateId).toBeUndefined();
    expect(removed.editor.generation.compareMode).toBe("result");
  });

  test("sets generation compare mode without touching candidates", () => {
    const state = generationAddCandidateCommand.execute({ state: createInitialAppState("Test") }, { candidate: generationCandidate("candidate-1") });

    const next = generationSetCompareModeCommand.execute({ state }, { mode: "split" });

    expect(next.editor.generation.candidates).toEqual([generationCandidate("candidate-1")]);
    expect(next.editor.generation.selectedCandidateId).toBe("candidate-1");
    expect(next.editor.generation.compareMode).toBe("split");
  });

  test("bounds retained candidate source memory as well as candidate count", () => {
    const largeSource = "x".repeat(34 * 1024 * 1024);
    const first = { ...generationCandidate("candidate-1"), source: largeSource };
    const second = { ...generationCandidate("candidate-2"), source: largeSource };
    const withFirst = generationAddCandidateCommand.execute({ state: createInitialAppState("Test") }, { candidate: first });
    const withSecond = generationAddCandidateCommand.execute({ state: withFirst }, { candidate: second });

    expect(withSecond.editor.generation.candidates.map((candidate) => candidate.id)).toEqual(["candidate-2"]);
  });

  test("clears the candidate session and resets comparison", () => {
    const withCandidate = generationAddCandidateCommand.execute(
      { state: createInitialAppState("Test") },
      { candidate: generationCandidate("candidate-1") },
    );
    const comparing = generationSetCompareModeCommand.execute({ state: withCandidate }, { mode: "split" });

    const next = generationClearCandidatesCommand.execute({ state: comparing }, undefined);

    expect(next.editor.generation.candidates).toEqual([]);
    expect(next.editor.generation.selectedCandidateId).toBeUndefined();
    expect(next.editor.generation.compareMode).toBe("result");
  });

  test("reuses a candidate's full settings with its resolved seed", () => {
    const candidate = generationCandidate("candidate-1");
    candidate.settings = {
      ...candidate.settings,
      prompt: "Stored prompt",
      model: "stored-model.safetensors",
      seed: -1,
      inpaint: { ...candidate.settings.inpaint, maskBlur: 7 },
    };
    candidate.seed = 987654;
    const state = generationAddCandidateCommand.execute({ state: createInitialAppState("Test") }, { candidate });

    const next = generationReuseCandidateSettingsCommand.execute({ state }, { candidateId: candidate.id });

    expect(next.editor.tools.generate).toEqual({ ...candidate.settings, seed: 987654 });
    expect(next.editor.generation).toBe(state.editor.generation);
  });

  test("applies candidates as top-level layers", () => {
    const withRemainingCandidate = generationAddCandidateCommand.execute(
      { state: documentWithSourceLayer() },
      { candidate: generationCandidate("candidate-2") },
    );
    const withSelectedCandidate = generationAddCandidateCommand.execute(
      { state: withRemainingCandidate },
      { candidate: generationCandidate("candidate-1") },
    );
    const state = generationSetCompareModeCommand.execute({ state: withSelectedCandidate }, { mode: "split" });

    const next = generationApplyCandidateAsLayerCommand.execute(
      { state },
      { candidateId: "candidate-1", assetId: "generated-asset", layerId: "generated-layer" },
    );

    expect(next.document.assets.find((asset) => asset.id === "generated-asset")?.source).toBe("generated-source");
    expect(next.document.assets.find((asset) => asset.id === "generated-asset")?.provenance).toMatchObject({
      kind: "generated",
      candidateId: "candidate-1",
      acceptance: "layer",
      seed: 123,
    });
    expect(next.document.artboards[0]?.layers[0]?.id).toBe("generated-layer");
    expect(next.editor.selection).toEqual({ artboardId: "a1", layerIds: ["generated-layer"] });
    expect(next.editor.generation).toEqual({
      candidates: [generationCandidate("candidate-2")],
      selectedCandidateId: "candidate-2",
      compareMode: "split",
      jobs: [],
      resources: { status: "idle" },
    });
  });

  test("replaces source asset pixels for inpaint candidates", () => {
    const withRemainingCandidate = generationAddCandidateCommand.execute(
      { state: documentWithSourceLayer() },
      { candidate: generationCandidate("candidate-2", true) },
    );
    const withSelectedCandidate = generationAddCandidateCommand.execute(
      { state: withRemainingCandidate },
      { candidate: generationCandidate("candidate-1", true) },
    );
    const state = generationSetCompareModeCommand.execute({ state: withSelectedCandidate }, { mode: "before" });

    const next = generationReplaceCandidatePixelsCommand.execute(
      { state },
      { candidateId: "candidate-1", source: "composited-source", mimeType: "image/png" },
    );

    expect(next.document.assets.find((asset) => asset.id === "source-asset")?.source).toBe("composited-source");
    expect(next.document.assets.find((asset) => asset.id === "source-asset")?.mimeType).toBe("image/png");
    expect(next.document.assets.find((asset) => asset.id === "source-asset")?.provenance).toMatchObject({
      kind: "generated",
      candidateId: "candidate-1",
      acceptance: "replacement",
      inpaint: {
        sourceAssetId: "source-asset",
        maskAssetId: "mask-asset",
      },
    });
    expect(next.editor.selection).toEqual({ artboardId: "a1", layerIds: ["source-layer"] });
    expect(next.editor.generation).toEqual({
      candidates: [generationCandidate("candidate-2", true)],
      selectedCandidateId: "candidate-2",
      compareMode: "before",
      jobs: [],
      resources: { status: "idle" },
    });
  });

  test("tracks one durable generation job through completion", () => {
    const state = createInitialAppState("Test");
    const running = generationStartJobCommand.execute({ state }, { jobId: "job-1", kind: "generate", label: "Generating", startedAt: 100 });
    const duplicate = generationStartJobCommand.execute({ state: running }, { jobId: "job-2", kind: "regenerate", label: "Regenerate", startedAt: 101 });
    const completed = generationSucceedJobCommand.execute({ state: duplicate }, { jobId: "job-1", finishedAt: 150 });

    expect(duplicate).toBe(running);
    expect(completed.editor.generation.jobs[0]).toEqual({ id: "job-1", kind: "generate", label: "Generating", status: "succeeded", startedAt: 100, finishedAt: 150, error: undefined });
  });

  test("preserves generation errors in authoritative state", () => {
    const running = generationStartJobCommand.execute({ state: createInitialAppState("Test") }, { jobId: "job-1", kind: "replace", label: "Replacing pixels", startedAt: 100 });
    const failed = generationFailJobCommand.execute({ state: running }, { jobId: "job-1", finishedAt: 125, error: "Backend unavailable" });

    expect(failed.editor.generation.jobs[0]).toMatchObject({ id: "job-1", status: "failed", error: "Backend unavailable", finishedAt: 125 });
  });

  test("records explicit cancellation separately from failure", () => {
    const running = generationStartJobCommand.execute({ state: createInitialAppState("Test") }, { jobId: "job-1", kind: "generate", label: "Generating", startedAt: 100 });
    const cancelled = generationCancelJobCommand.execute({ state: running }, { jobId: "job-1", finishedAt: 110 });

    expect(cancelled.editor.generation.jobs[0]).toMatchObject({ id: "job-1", status: "cancelled", finishedAt: 110 });
  });
});

function documentWithSourceLayer() {
  return {
    ...createInitialAppState("Test"),
    document: {
      ...createInitialAppState("Test").document,
      assets: [
        { id: "source-asset", name: "Source", mimeType: "image/png", source: "source", intrinsicSize: { w: 100, h: 100 } },
        { id: "mask-asset", name: "Mask", mimeType: "image/png", source: "mask", intrinsicSize: { w: 100, h: 100 } },
      ],
      artboards: [
        {
          id: "a1",
          name: "Artboard 1",
          bounds: { x: 0, y: 0, w: 100, h: 100 },
          backgroundColor: "transparent",
          visible: true,
          locked: false,
          layers: [
            raster("mask-layer", "Mask", "mask-asset"),
            { ...raster("source-layer", "Source", "source-asset"), clippingMask: { maskLayerId: "mask-layer" } },
          ],
        },
      ],
    },
  };
}

function generationCandidate(id: string, inpaint = false): GenerationCandidate {
  const candidate: GenerationCandidate = {
    id,
    source: "generated-source",
    mimeType: "image/png",
    intrinsicSize: { w: 64, h: 64 },
    mode: inpaint ? "inpaint" : "text-to-image",
    settings: createInitialAppState("Test").editor.tools.generate,
    seed: 123,
    width: 64,
    height: 64,
    placement: {
      artboardId: "a1",
      layerName: "Generated",
      transform: { position: { x: 5, y: 6 }, scale: { x: 1, y: 1 }, rotation: 0 },
    },
  };

  return inpaint
    ? {
        ...candidate,
        inputImage: "input",
        maskImage: "mask",
        inpaint: {
          targetLayerId: "source-layer",
          regionId: "region",
          sourceAssetId: "source-asset",
          maskAssetId: "mask-asset",
          inputImage: "input",
          maskImage: "mask",
          editMaskImage: "edit-mask",
          blendMaskImage: "blend-mask",
          revision: { source: "source-revision", mask: "mask-revision" },
          crop: {
            assetBounds: { x: 0, y: 0, w: 64, h: 64 },
            documentBounds: { x: 0, y: 0, w: 64, h: 64 },
            padding: 12,
            maskedAreaOnly: true,
          },
          mask: {
            polarity: "hidden",
            activeBounds: { x: 10, y: 10, w: 20, h: 20 },
          },
          backend: {
            growMaskBy: 6,
            maskedContent: "neutral",
            maskBlur: 0,
            maskFeather: 0,
            maskExpand: 0,
            cropPadding: 12,
          },
        },
      }
    : candidate;
}

function raster(id: string, name: string, assetId: string) {
  return {
    id,
    type: "raster" as const,
    name,
    visible: true,
    locked: false,
    opacity: 1,
    assetId,
    transform: { position: { x: 0, y: 0 }, scale: { x: 1, y: 1 }, rotation: 0 },
  };
}