feat(harness): validate image attachments

2026-09-07 16:01:16 +08:00
parent f9ea4ce429
commit 1eaa85c09a
5 changed files with 458 additions and 0 deletions
@@ -0,0 +1,121 @@
+import { mkdir, mkdtemp, rm } from "node:fs/promises";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { Readable } from "node:stream";
+import sharp from "sharp";
+import { afterEach, beforeEach, describe, expect, it } from "vitest";
+import * as harness from "../src/index.js";
+
+type MediaType = "image/jpeg" | "image/png" | "image/webp";
+type Upload = { name: string; mediaType: MediaType; stream: Readable };
+type Attachment = { attachmentId: string; sessionId: string; name: string; mediaType: MediaType; byteSize: number; width: number; height: number; status: "ready" };
+type Store = {
+  add(sessionId: string, upload: Upload): Promise<Attachment>;
+  get(sessionId: string, attachmentId: string): Attachment | undefined;
+  openNormalized(sessionId: string, attachmentId: string): Promise<Buffer>;
+  remove(sessionId: string, attachmentId: string): Promise<boolean>;
+  removeAll(sessionId: string): Promise<void>;
+  toPiImages(sessionId: string, attachmentIds: string[]): Promise<Array<{ type: "image"; data: string; mimeType: string }>>;
+};
+
+let root = "";
+
+beforeEach(async () => {
+  root = await mkdtemp(join(tmpdir(), "agent-studio-attachments-"));
+  await mkdir(join(root, "ses_a"), { recursive: true });
+  await mkdir(join(root, "ses_b"), { recursive: true });
+});
+
+afterEach(async () => {
+  await rm(root, { recursive: true, force: true });
+});
+
+function createStore(options: Record<string, unknown> = {}): Store {
+  const Constructor = Reflect.get(harness, "FileAttachmentStore") as unknown as
+    | (new (options: Record<string, unknown>) => Store)
+    | undefined;
+  expect(Constructor, "FileAttachmentStore must be exported").toBeDefined();
+  return new Constructor!({ sessionsRoot: root, ...options });
+}
+
+function upload(name: string, mediaType: MediaType, bytes: Buffer): Upload {
+  return { name, mediaType, stream: Readable.from(bytes) };
+}
+
+async function image(format: "jpeg" | "png" | "webp", width = 4, height = 3): Promise<Buffer> {
+  return sharp({ create: { width, height, channels: 3, background: "#7c3aed" } })[format]().toBuffer();
+}
+
+describe("FileAttachmentStore", () => {
+  it.each([
+    ["photo.jpg", "image/jpeg", "jpeg"],
+    ["chart.png", "image/png", "png"],
+    ["creative.webp", "image/webp", "webp"],
+  ] as const)("accepts and normalizes %s", async (name, mediaType, format) => {
+    const store = createStore();
+    const item = await store.add("ses_a", upload(name, mediaType, await image(format)));
+
+    expect(item).toMatchObject({ sessionId: "ses_a", name, mediaType, width: 4, height: 3, status: "ready" });
+    expect(item.attachmentId).toMatch(/^att_/);
+    expect((await store.openNormalized("ses_a", item.attachmentId)).byteLength).toBe(item.byteSize);
+  });
+
+  it("rejects declared PNG with JPEG bytes", async () => {
+    await expect(createStore().add("ses_a", upload("fake.png", "image/png", await image("jpeg"))))
+      .rejects.toMatchObject({ code: "IMAGE_SIGNATURE_MISMATCH" });
+  });
+
+  it("rejects unsupported SVG before decoding", async () => {
+    const svg = Buffer.from('<svg xmlns="http://www.w3.org/2000/svg"><rect width="1" height="1"/></svg>');
+    await expect(
+      createStore().add("ses_a", { name: "vector.svg", mediaType: "image/png", stream: Readable.from(svg) }),
+    ).rejects.toMatchObject({ code: "IMAGE_TYPE_NOT_ALLOWED" });
+  });
+
+  it("enforces stream byte and decoded dimension limits", async () => {
+    const png = await image("png", 10, 10);
+    await expect(createStore({ maxBytes: 12 }).add("ses_a", upload("large.png", "image/png", png)))
+      .rejects.toMatchObject({ code: "IMAGE_TOO_LARGE" });
+    await expect(createStore({ maxWidth: 5 }).add("ses_a", upload("wide.png", "image/png", png)))
+      .rejects.toMatchObject({ code: "IMAGE_DIMENSIONS_EXCEEDED" });
+    await expect(createStore({ maxPixels: 50 }).add("ses_a", upload("pixels.png", "image/png", png)))
+      .rejects.toMatchObject({ code: "IMAGE_DIMENSIONS_EXCEEDED" });
+  });
+
+  it("does not expose cross-session attachment metadata or bytes", async () => {
+    const store = createStore();
+    const item = await store.add("ses_a", upload("chart.png", "image/png", await image("png")));
+
+    expect(store.get("ses_b", item.attachmentId)).toBeUndefined();
+    await expect(store.openNormalized("ses_b", item.attachmentId)).rejects.toMatchObject({ code: "ATTACHMENT_NOT_FOUND" });
+  });
+
+  it("maps ready attachments to Pi images, enforces count, and removes idempotently", async () => {
+    const store = createStore({ maxAttachmentsPerMessage: 1 });
+    const first = await store.add("ses_a", upload("one.png", "image/png", await image("png")));
+    const second = await store.add("ses_a", upload("two.png", "image/png", await image("png")));
+
+    const images = await store.toPiImages("ses_a", [first.attachmentId]);
+    expect(images[0]).toMatchObject({ type: "image", mimeType: "image/png" });
+    expect(Buffer.from(images[0]!.data, "base64").byteLength).toBe(first.byteSize);
+    await expect(store.toPiImages("ses_a", [first.attachmentId, second.attachmentId])).rejects.toMatchObject({ code: "TOO_MANY_ATTACHMENTS" });
+    await expect(store.remove("ses_a", first.attachmentId)).resolves.toBe(true);
+    await expect(store.remove("ses_a", first.attachmentId)).resolves.toBe(false);
+    await store.removeAll("ses_a");
+    expect(store.get("ses_a", second.attachmentId)).toBeUndefined();
+  });
+
+  it("emits safe metadata events without image data", async () => {
+    const Store = Reflect.get(harness, "InMemoryEventStore") as new () => {
+      listAfter(sessionId: string): { events: unknown[] };
+    };
+    const events = new Store();
+    const store = createStore({ eventStore: events });
+    await store.add("ses_a", upload("chart.png", "image/png", await image("png")));
+
+    const serialized = JSON.stringify(events.listAfter("ses_a"));
+    expect(serialized).toContain("attachment.uploaded");
+    expect(serialized).not.toContain("base64");
+    expect(serialized).not.toContain("iVBOR");
+  });
+});