pi-image-paths/index.test.ts

// Drives the extension the way pi does: a stand-in for pi records the handlers it registers, a
// temp directory holds real image files, the editor rewrite runs off a fake input box, and the
// entry the extension appends is rendered back through the renderer it registered.
//
// The load-bearing test is "the sent image is drawn below the message": pi renders no image
// blocks on user messages, so without the appended entry the transcript carries only the marker.
// Run with `npm test`.

import assert from "node:assert/strict";
import { mkdtempSync, writeFileSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { test } from "node:test";

import { setCapabilities } from "@earendil-works/pi-tui";

import imagePaths, { extractImages, resolveImageMarkers } from "./index";
import type { ImageContent } from "./index";

/** 1x1 transparent PNG. */
const PNG =
  "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAF" +
  "AAH/q842iQAAAABJRU5ErkJggg==";

type Handler = (event: any, ctx: any) => any;
type Renderer = (entry: any, options: any, theme: any) => any;

const THEME = { fg: (_color: string, text: string) => text, bold: (text: string) => text };

const tick = (ms = 20) => new Promise((resolve) => setTimeout(resolve, ms));

/** A directory with a screenshot, a screenshot whose name has a space, and a text file. */
function fixture() {
  const dir = mkdtempSync(join(tmpdir(), "pi-image-paths-"));
  const shot = join(dir, "shot.png");
  const spaced = join(dir, "two words.png");
  writeFileSync(shot, Buffer.from(PNG, "base64"));
  writeFileSync(spaced, Buffer.from(PNG, "base64"));
  writeFileSync(join(dir, "notes.txt"), "not an image");
  return { dir, shot, spaced };
}

/** A stand-in for pi's extension API, plus the state a test needs to inspect. */
function load(options: { showImages?: boolean; imageWidthCells?: number } = {}) {
  const handlers = new Map<string, Handler>();
  const entries: { type: string; data: any }[] = [];
  const notices: string[] = [];
  let entryType: string | undefined;
  let renderer: Renderer | undefined;
  let editorText = "";
  let terminalInput: ((data: string) => unknown) | undefined;

  const pi = {
    on: (name: string, handler: Handler) => {
      handlers.set(name, handler);
      return () => handlers.delete(name);
    },
    registerEntryRenderer: (type: string, fn: Renderer) => {
      entryType = type;
      renderer = fn;
    },
    registerCommand: () => {},
    registerShortcut: () => {},
    getSettings: () => ({
      terminal: {
        showImages: options.showImages ?? true,
        imageWidthCells: options.imageWidthCells,
      },
    }),
    appendEntry: (type: string, data: unknown) => entries.push({ type, data }),
  };

  imagePaths(pi as unknown as Parameters<typeof imagePaths>[0]);

  return {
    handlers,
    entries,
    notices,
    entryType: () => entryType,
    renderer: () => renderer as Renderer,
    terminalInput: () => terminalInput,
    editorText: () => editorText,
    setEditorText: (text: string) => {
      editorText = text;
    },
    ctx: (cwd: string, mode = "tui") => ({
      mode,
      hasUI: true,
      cwd,
      ui: {
        onTerminalInput: (fn: (data: string) => unknown) => {
          terminalInput = fn;
          return () => {
            terminalInput = undefined;
          };
        },
        getEditorText: () => editorText,
        setEditorText: (text: string) => {
          editorText = text;
        },
        notify: (message: string) => notices.push(message),
      },
    }),
  };
}

/** A registered handler, asserted present so tests fail loudly instead of silently skipping. */
function handler(loaded: ReturnType<typeof load>, name: string): Handler {
  const found = loaded.handlers.get(name);
  assert.ok(found, `no ${name} handler registered`);
  return found;
}

/** The terminal-input listener the session_start handler subscribed. */
function terminalInput(loaded: ReturnType<typeof load>): (data: string) => unknown {
  const found = loaded.terminalInput();
  assert.ok(found, "no terminal input listener registered");
  return found;
}

function userMessage(images = true) {
  return {
    role: "user",
    content: images
      ? [{ type: "text", text: "[Image #1] why is the header clipped?" }, { type: "image", data: PNG, mimeType: "image/png" }]
      : [{ type: "text", text: "no pictures here" }],
  };
}

// ---------------------------------------------------------------- paths in the prompt

test("a bare absolute path becomes a marker and an attachment", () => {
  const { dir, shot } = fixture();
  const extracted = extractImages(`look at ${shot} please`, dir);
  assert.equal(extracted.text, "look at [Image #1] please");
  assert.deepEqual(extracted.files, [{ file: shot, ext: "png" }]);
  assert.deepEqual(extracted.missing, []);
});

test("several paths number in order of appearance", () => {
  const { dir, shot, spaced } = fixture();
  const extracted = extractImages(`${shot} and "${spaced}"`, dir);
  // Quoting is the prompt's, so the marker lands inside it.
  assert.equal(extracted.text, '[Image #1] and "[Image #2]"');
  assert.deepEqual(
    extracted.files.map((f) => f.file),
    [shot, spaced],
  );
});

test("a path with escaped spaces resolves as one file", () => {
  const { dir, spaced } = fixture();
  const extracted = extractImages(`see ${spaced.replaceAll(" ", "\\ ")}`, dir);
  assert.equal(extracted.text, "see [Image #1]");
  assert.deepEqual(extracted.files, [{ file: spaced, ext: "png" }]);
});

test("a path with unescaped spaces is joined and disambiguated against the filesystem", () => {
  const { dir, spaced } = fixture();
  const extracted = extractImages(`see ${spaced} now`, dir);
  assert.equal(extracted.text, "see [Image #1] now");
  assert.deepEqual(extracted.files, [{ file: spaced, ext: "png" }]);
});

test("startNumber offsets the first marker", () => {
  const { dir, shot } = fixture();
  assert.equal(extractImages(shot, dir, 4).text, "[Image #4]");
});

test("a missing image is reported, not swallowed", () => {
  const { dir } = fixture();
  const missing = join(dir, "gone.png");
  const extracted = extractImages(`look at ${missing}`, dir);
  assert.equal(extracted.text, `look at ${missing}`);
  assert.deepEqual(extracted.files, []);
  assert.deepEqual(extracted.missing, [missing]);
});

test("URLs and non-image extensions are left alone", () => {
  const { dir } = fixture();
  const text = `https://example.com/a.png ${join(dir, "notes.txt")}`;
  const extracted = extractImages(text, dir);
  assert.equal(extracted.text, text);
  assert.deepEqual(extracted.files, []);
  assert.deepEqual(extracted.missing, []);
});

test("markers resolve through the registry, unknown ones stay literal", () => {
  const registry = new Map([
    [1, { file: "/tmp/one.png", ext: "png" }],
    [2, { file: "/tmp/two.png", ext: "png" }],
  ]);
  const resolved = resolveImageMarkers("[Image #2] then [Image #1] and [Image #7]", registry);
  assert.equal(resolved.text, "[Image #1] then [Image #2] and [Image #7]");
  assert.deepEqual(
    resolved.files.map((f) => f.file),
    ["/tmp/two.png", "/tmp/one.png"],
  );
});

test("text without a marker short-circuits", () => {
  assert.deepEqual(resolveImageMarkers("nothing here", new Map()), {
    text: "nothing here",
    files: [],
    missing: [],
  });
});

// ---------------------------------------------------------------- the editor rewrite

test("the live rewrite collapses a dropped path to a marker", async () => {
  const { dir, shot } = fixture();
  const loaded = load();
  const ctx = loaded.ctx(dir);
  handler(loaded, "session_start")({ type: "session_start" }, ctx);

  loaded.setEditorText(`look at ${shot}`);
  terminalInput(loaded)("x");
  await tick(250);

  assert.equal(loaded.editorText(), "look at [Image #1]");
});

test("submitting with Enter never triggers a rewrite", async () => {
  const { dir, shot } = fixture();
  const loaded = load();
  const ctx = loaded.ctx(dir);
  handler(loaded, "session_start")({ type: "session_start" }, ctx);

  loaded.setEditorText(`look at ${shot}`);
  terminalInput(loaded)("\r");
  await tick(250);

  assert.equal(loaded.editorText(), `look at ${shot}`);
});

test("paste placeholders are left for pi to expand", async () => {
  const { dir, shot } = fixture();
  const loaded = load();
  const ctx = loaded.ctx(dir);
  handler(loaded, "session_start")({ type: "session_start" }, ctx);

  loaded.setEditorText(`[paste #1 +20 lines] ${shot}`);
  terminalInput(loaded)("x");
  await tick(250);

  assert.equal(loaded.editorText(), `[paste #1 +20 lines] ${shot}`);
});

test("a rewritten marker still attaches on send", async () => {
  const { dir, shot, spaced } = fixture();
  const loaded = load();
  const ctx = loaded.ctx(dir);
  handler(loaded, "session_start")({ type: "session_start" }, ctx);

  loaded.setEditorText(`${shot} and ${spaced}`);
  terminalInput(loaded)("x");
  await tick(250);

  const result = await handler(loaded, "input")(
    { type: "input", text: loaded.editorText(), source: "interactive" },
    ctx,
  );
  assert.equal(result.action, "transform");
  assert.equal(result.text, "[Image #1] and [Image #2]");
  assert.deepEqual(
    result.images.map((image: ImageContent) => image.mimeType),
    ["image/png", "image/png"],
  );
});

// ---------------------------------------------------------------- attaching on send

test("the input handler attaches the file behind a path left in the prompt", async () => {
  const { dir, shot } = fixture();
  const loaded = load();
  const result = await handler(loaded, "input")(
    { type: "input", text: `look at ${shot}`, source: "interactive" },
    loaded.ctx(dir),
  );
  assert.equal(result.action, "transform");
  assert.equal(result.text, "look at [Image #1]");
  assert.deepEqual(result.images, [{ type: "image", data: PNG, mimeType: "image/png" }]);
});

test("a path that does not exist warns and leaves the prompt alone", async () => {
  const { dir } = fixture();
  const loaded = load();
  const missing = join(dir, "gone.png");
  const result = await handler(loaded, "input")(
    { type: "input", text: `look at ${missing}`, source: "interactive" },
    loaded.ctx(dir),
  );
  assert.equal(result.action, "continue");
  assert.equal(loaded.notices.length, 1);
  assert.match(loaded.notices[0], /gone\.png/);
});

test("extension-injected input is never rewritten", async () => {
  const { dir, shot } = fixture();
  const loaded = load();
  const result = await handler(loaded, "input")({ type: "input", text: shot, source: "extension" }, loaded.ctx(dir));
  assert.deepEqual(result, { action: "continue" });
});

// ---------------------------------------------------------------- drawing the sent image

test("the sent image is drawn below the message", async () => {
  const { dir } = fixture();
  const loaded = load();
  const messageEnd = handler(loaded, "message_end");

  messageEnd({ type: "message_end", message: userMessage() }, loaded.ctx(dir));
  assert.equal(loaded.entries.length, 0, "appended before the session stored the message");
  await tick();

  assert.equal(loaded.entries.length, 1);
  assert.equal(loaded.entries[0].type, loaded.entryType());
  assert.deepEqual(loaded.entries[0].data.images, [{ data: PNG, mimeType: "image/png" }]);
});

test("showImages off appends nothing while the image still reaches the model", async () => {
  const { dir } = fixture();
  const loaded = load({ showImages: false });
  handler(loaded, "message_end")({ type: "message_end", message: userMessage() }, loaded.ctx(dir));
  await tick();
  assert.deepEqual(loaded.entries, []);
});

test("only user messages with images get an entry", async () => {
  const { dir } = fixture();
  const loaded = load();
  const messageEnd = handler(loaded, "message_end");

  messageEnd({ type: "message_end", message: userMessage(false) }, loaded.ctx(dir));
  messageEnd({ type: "message_end", message: { role: "assistant", content: [{ type: "text", text: "hi" }] } }, loaded.ctx(dir));
  messageEnd({ type: "message_end", message: userMessage() }, loaded.ctx(dir, "rpc"));
  await tick();

  assert.deepEqual(loaded.entries, []);
});

test("the renderer draws the image in the transcript", () => {
  setCapabilities({ images: "kitty", trueColor: true, hyperlinks: false });
  const loaded = load({ imageWidthCells: 40 });
  const entry = {
    type: "custom",
    customType: loaded.entryType(),
    data: { images: [{ data: PNG, mimeType: "image/png" }] },
  };

  const component = loaded.renderer()(entry, { expanded: false }, THEME);
  assert.ok(component, "no component for a drawable entry");
  const lines = component.render(40);
  assert.ok(lines.length > 0);
  assert.ok(
    lines.some((line: string) => line.includes("\x1b_G")),
    "no kitty image escape in the rendered lines",
  );
});

test("the renderer stays out of the way when drawing is off or impossible", () => {
  const entry = {
    type: "custom",
    customType: "image-paths:attached",
    data: { images: [{ data: PNG, mimeType: "image/png" }] },
  };

  setCapabilities({ images: "kitty", trueColor: true, hyperlinks: false });
  const off = load({ showImages: false });
  assert.equal(off.renderer()(entry, { expanded: false }, THEME), undefined);

  setCapabilities({ images: null, trueColor: true, hyperlinks: false });
  const terminal = load();
  assert.equal(terminal.renderer()(entry, { expanded: false }, THEME), undefined);

  setCapabilities({ images: "kitty", trueColor: true, hyperlinks: false });
  const empty = load();
  assert.equal(
    empty.renderer()({ ...entry, data: {} }, { expanded: false }, THEME),
    undefined,
  );
});