diff --git a/docs/image-attachments.md b/docs/image-attachments.md new file mode 100644 index 000000000..fdc25439f --- /dev/null +++ b/docs/image-attachments.md @@ -0,0 +1,65 @@ +# Temporary files for chat images + +Pasted, dropped, and selected chat images are still sent as inline image inputs +and shown in the composer and conversation. Pi Web also saves the received image +bytes to temporary files on the server and appends their absolute paths, numbered +in attachment order, to the agent's message. Local file and image tools can use +these paths without asking the user to locate or save the image again (#635). + +Files are created before the prompt is submitted. A failed write rejects the +submission and removes that batch's partial files. A rejected prompt preflight +also removes its batch; a provider failure after acceptance preserves the files +for follow-up work. Normal prompts, streaming prompts, steering and follow-ups +use the same storage path. + +## Storage and lifetime + +- Files live below the OS temporary directory (`os.tmpdir()`, normally `$TMPDIR` + on Unix) in a private `pi-web-clipboard-v1--/batch-` directory. +- File names are `pi-web-clipboard-.`. The extension comes from + the declared image MIME type; client-supplied names and paths are never used. +- Directories use mode `0700` and files use `0600` on POSIX systems. Windows + access remains governed by the user's temporary-directory ACLs. +- The existing limits of 10 images and 10 MiB per image remain. Malformed Base64 + and MIME types without a supported file extension are rejected. This is not a + full image decoder: existing image decoding remains the responsibility of the + browser, model provider and image tools. +- Files belong to the server process, so releasing an idle session wrapper, + reloading its tools, or forking a session does not immediately delete them. +- The process checks hourly and removes its temporary files after 24 hours of + inactivity. New prompt/steering/follow-up submissions refresh that period; + active agent work and queued messages defer expiry for all images in the + process. These files are not a durable storage API or a disk quota mechanism. +- Normal process exit removes its files. On the next image upload and hourly + thereafter, directories older than 24 hours whose owning process no longer + exists are collected. Live or unverifiable owners, other users' directories + on POSIX, unrelated names and symlinks are skipped. PID reuse conservatively + delays collection. Cleanup errors are logged by the server. + +Paths can expire or disappear on restart or OS temporary-directory cleanup. +Historical messages keep their inline image data, but their old paths are not +automatically recreated. Reattach the image when a fresh path is needed. No +cross-session or cross-process path permanence is promised. The file browser's +allowed roots are not broadened to include the whole temporary directory. + +## Image fidelity + +The temporary file contains exactly the bytes received by the server. The +existing browser compression still applies: images larger than 1 MiB, except +GIFs, may be resized to a maximum 1024-pixel side and encoded as JPEG at 0.85 +quality when that produces a smaller payload. This feature does not archive +the original upload before compression. + +The paths are local to the server. A model receives the inline image and the +path annotation; whether it calls a file tool depends on the task and enabled +tools. Providing a path does not give a text-only model vision capability. + +## Verification + +`lib/image-materialization.test.mjs` covers byte preservation, limits, private +permissions, isolated batches, partial-write rollback and garbage collection. +`lib/rpc-manager-images.test.mjs` checks all message entry points, rejection +versus acceptance cleanup, wrapper disposal and actual Pi `read` tool access. +The browser suite creates non-sensitive JPEG/PNG fixtures, exercises picker, +paste and drop on desktop/mobile, and checks the inputs at the SDK extension +boundary against readable server files. It does not contact a model provider. diff --git a/e2e/README.md b/e2e/README.md index 745eca914..ea2ada6c8 100644 --- a/e2e/README.md +++ b/e2e/README.md @@ -34,6 +34,8 @@ Coverage: - Unknown sessions and paths outside the fixture project are rejected. - A local extension checks dialog keyboard navigation, Esc cancellation, collapse/expand draft preservation, countdown display, and server-side expiry. +- JPEG/PNG selection, paste and drop preserve previews and inline image inputs; + a local input extension verifies server-side temporary paths and image bytes. Model prompts, live model streaming, and agent execution are outside this suite. Failures save a screenshot, Playwright trace, and server log under diff --git a/e2e/image-attachments.mjs b/e2e/image-attachments.mjs new file mode 100644 index 000000000..f339f1845 --- /dev/null +++ b/e2e/image-attachments.mjs @@ -0,0 +1,79 @@ +import assert from "node:assert/strict"; +import { readFile, stat } from "node:fs/promises"; +import { extname, join } from "node:path"; + +export const imageExtensionSource = ` +import { readFile, writeFile } from "node:fs/promises"; +import { join } from "node:path"; +export default function (pi) { + pi.on("input", async (event, ctx) => { + if (!event.text.startsWith("E2E image ")) return { action: "continue" }; + const paths = [...event.text.matchAll(/^\\d+\\. (".*")$/gm)].map(match => JSON.parse(match[1])); + const files = await Promise.all(paths.map(async path => ({ path, data: (await readFile(path)).toString("base64") }))); + await writeFile(join(ctx.cwd, "image-capture.json"), JSON.stringify({ text: event.text, images: event.images, files })); + ctx.ui.notify(event.text.split("\\n")[0].trim() + " captured"); + return { action: "handled" }; + }); +}`; + +export async function checkImageAttachments(page, project, artifacts, width) { + const fixtures = await page.evaluate(() => { + const canvas = document.createElement("canvas"); + canvas.width = 16; + canvas.height = 16; + const ctx = canvas.getContext("2d"); + ctx.fillStyle = "#d64a39"; + ctx.fillRect(0, 0, 8, 16); + ctx.fillStyle = "#269dc5"; + ctx.fillRect(8, 0, 8, 16); + return ["image/jpeg", "image/png"].map(mimeType => ({ mimeType, data: canvas.toDataURL(mimeType).split(",")[1] })); + }); + const allPaths = new Set(); + for (const entry of ["picker", "paste", "drop"]) { + const input = page.locator("textarea").last(); + const text = `E2E image ${width} ${entry}`; + await input.fill(text); + if (entry === "picker") { + await page.locator('input[type="file"][accept="image/*"]').setInputFiles(fixtures.map((image, index) => ({ + name: `fixture-${index}.${index === 0 ? "jpg" : "png"}`, + mimeType: image.mimeType, + buffer: Buffer.from(image.data, "base64"), + }))); + } else { + await input.evaluate((element, { entry, fixtures }) => { + const transfer = new DataTransfer(); + fixtures.forEach((image, index) => { + const bytes = Uint8Array.from(atob(image.data), char => char.charCodeAt(0)); + transfer.items.add(new File([bytes], `fixture-${index}`, { type: image.mimeType })); + }); + let event; + if (entry === "paste") event = new ClipboardEvent("paste", { clipboardData: transfer, bubbles: true, cancelable: true }); + else event = new DragEvent("drop", { dataTransfer: transfer, bubbles: true, cancelable: true }); + element.dispatchEvent(event); + }, { entry, fixtures }); + } + await page.waitForFunction(() => document.querySelectorAll('img[src^="blob:"]').length === 2); + await page.screenshot({ path: join(artifacts, `images-${entry}-${width}.png`) }); + const requestPromise = page.waitForRequest(request => request.method() === "POST" + && /\/api\/agent\/[^/]+$/.test(new URL(request.url()).pathname) + && request.postDataJSON()?.message === text); + await page.getByRole("button", { name: "Send", exact: true }).click(); + const sent = (await requestPromise).postDataJSON(); + await page.getByText(`${text} captured`, { exact: true }).waitFor(); + await page.getByRole("button", { name: "Stop agent", exact: true }).waitFor({ state: "hidden" }); + const captured = JSON.parse(await readFile(join(project, "image-capture.json"), "utf8")); + assert.deepEqual(sent.images, fixtures.map(image => ({ type: "image", ...image }))); + assert.deepEqual(captured.images, sent.images); + assert.equal(captured.files.length, 2); + for (let index = 0; index < captured.files.length; index++) { + const file = captured.files[index]; + assert.equal(file.data, fixtures[index].data); + assert.equal(extname(file.path), index === 0 ? ".jpg" : ".png"); + assert.ok(!allPaths.has(file.path)); + allPaths.add(file.path); + assert.deepEqual(await readFile(file.path), Buffer.from(fixtures[index].data, "base64")); + if (process.platform !== "win32") assert.equal((await stat(file.path)).mode & 0o777, 0o600); + } + } + console.log(`PASS: ${width}px JPEG/PNG picker, paste, drop, preview, inline inputs and readable server files`); +} diff --git a/e2e/run.mjs b/e2e/run.mjs index fd68dba51..ba86126aa 100644 --- a/e2e/run.mjs +++ b/e2e/run.mjs @@ -11,6 +11,7 @@ import { fileURLToPath } from "node:url"; import { chromium } from "playwright"; import { checkExtensionDialogs, extensionSource } from "./extension-dialog.mjs"; import { checkChatAppearance } from "./chat-appearance.mjs"; +import { checkImageAttachments, imageExtensionSource } from "./image-attachments.mjs"; const root = dirname(dirname(fileURLToPath(import.meta.url))); const mode = process.env.E2E_SERVER_MODE || "dev"; @@ -32,7 +33,15 @@ const text = (i) => `E2E message ${String(i).padStart(4, "0")}`; const ids = (start, end) => Array.from({ length: end - start }, (_, i) => `e${start + i}`); function message(id, parentId, role, content) { - return { type: "message", id, parentId, timestamp, message: { role, content } }; + const entry = { type: "message", id, parentId, timestamp, message: { role, content } }; + if (role === "assistant") { + // Real session instances compute their own usage, so fixtures must provide the fields the SDK expects. + entry.message.usage = { + input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }; + } + return entry; } function writeSession(id, entries) { @@ -60,6 +69,7 @@ try { // Seed before startup so the first catalogue scan sees every fixture. mkdirSync(join(agentDir, "extensions")); writeFileSync(join(agentDir, "extensions", "e2e-dialog.js"), extensionSource); + writeFileSync(join(agentDir, "extensions", "e2e-images.js"), imageExtensionSource); const longEntries = Array.from({ length: 5000 }, (_, i) => message(`e${i}`, i ? `e${i - 1}` : null, i % 2 ? "assistant" : "user", text(i))); longEntries.splice(1, 0, message("alternate", "e0", "user", "E2E alternate history branch")); @@ -352,6 +362,7 @@ try { await heading.waitFor({ state: "visible" }); } await checkExtensionDialogs(page, artifacts, viewport.width); + await checkImageAttachments(page, project, artifacts, viewport.width); if (viewport.width > 600) { await page.goto(`${base}/?session=${RICH}`, { waitUntil: "domcontentloaded" }); await page.locator(".markdown-code-block pre").waitFor(); diff --git a/lib/image-materialization.test.mjs b/lib/image-materialization.test.mjs new file mode 100644 index 000000000..f83c06cf9 --- /dev/null +++ b/lib/image-materialization.test.mjs @@ -0,0 +1,164 @@ +import assert from "node:assert/strict"; +import fs from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join, dirname, extname } from "node:path"; +import test from "node:test"; +import { execFile } from "node:child_process"; +import { promisify } from "node:util"; +import { createJiti } from "jiti"; + +const jiti = createJiti(import.meta.url, { tsconfigPaths: true }); +const { ImageAttachmentStore, IMAGE_IDLE_TTL_MS, appendImagePaths } = await jiti.import("./image-materialization.ts"); +const bytes = Buffer.from([0xff, 0xd8, 0xff, 0xe0, 0, 16, 74, 70, 73, 70, 0, 0x80, 0xff, 0xd9]); +const jpeg = { type: "image", mimeType: "image/jpeg", data: bytes.toString("base64") }; + +async function fixture(t, options = {}) { + const root = await fs.mkdtemp(join(tmpdir(), "pi-web-image-test-")); + const store = new ImageAttachmentStore({ root, ...options }); + t.after(async () => { store.disposeSync(); await fs.rm(root, { recursive: true, force: true }); }); + return { root, store }; +} + +test("materializes binary bytes with MIME extensions, private permissions and unique batches", async (t) => { + const { store } = await fixture(t); + const batches = await Promise.all([ + store.materialize([jpeg, { ...jpeg, mimeType: "image/png" }]), + store.materialize([jpeg]), + ]); + const paths = batches.flatMap(batch => batch.paths); + assert.equal(new Set(paths).size, 3); + assert.deepEqual(paths.map(extname), [".jpg", ".png", ".jpg"]); + assert.notEqual(dirname(paths[0]), dirname(paths[2])); + for (const path of paths) { + assert.deepEqual(await fs.readFile(path), bytes); + if (process.platform !== "win32") { + assert.equal((await fs.stat(path)).mode & 0o777, 0o600); + assert.equal((await fs.stat(dirname(path))).mode & 0o777, 0o700); + } + } + await batches[0].rollback(); + await assert.rejects(fs.stat(paths[0]), { code: "ENOENT" }); + assert.deepEqual(await fs.readFile(paths[2]), bytes); +}); + +test("rejects invalid inputs before creating files", async (t) => { + const { store, root } = await fixture(t); + for (const images of [ + [{ ...jpeg, data: "broken!" }], + [{ ...jpeg, mimeType: "image/../../bad" }], + [{ ...jpeg, mimeType: "image/unknown" }], + Array.from({ length: 11 }, () => jpeg), + [{ ...jpeg, data: Buffer.alloc(10 * 1024 * 1024 + 1).toString("base64") }], + ]) await assert.rejects(store.materialize(images)); + assert.deepEqual(await fs.readdir(root), []); +}); + +test("a failed batch removes partial files without deleting another batch", async (t) => { + const { store } = await fixture(t); + const kept = await store.materialize([jpeg]); + const write = fs.writeFile; + let calls = 0; + const mocked = t.mock.method(fs, "writeFile", async (...args) => { + await write(...args); + if (++calls === 2) throw new Error("fixture disk full"); + }); + await assert.rejects(store.materialize([jpeg, jpeg]), /fixture disk full/); + mocked.mock.restore(); + assert.deepEqual(await fs.readFile(kept.paths[0]), bytes); + assert.equal((await fs.readdir(dirname(dirname(kept.paths[0])))).length, 1); +}); + +test("idle cleanup protects active work and later prompts, then expires temporary paths", async (t) => { + let now = Date.now(); + let busy = false; + const { store } = await fixture(t, { now: () => now, isBusy: () => busy }); + const batch = await store.materialize([jpeg]); + now += IMAGE_IDLE_TTL_MS + 1; + busy = true; + await store.collect(); + assert.deepEqual(await fs.readFile(batch.paths[0]), bytes); + busy = false; + now += IMAGE_IDLE_TTL_MS - 1; + store.touch(); + await store.collect(); + assert.deepEqual(await fs.readFile(batch.paths[0]), bytes); + now += IMAGE_IDLE_TTL_MS + 1; + await store.collect(); + await assert.rejects(fs.stat(batch.paths[0]), { code: "ENOENT" }); + const next = await store.materialize([jpeg]); + assert.deepEqual(await fs.readFile(next.paths[0]), bytes); +}); + +test("orphan cleanup skips live owners, unrelated directories, and symlinks", async (t) => { + const now = Date.now(); + const { root, store } = await fixture(t, { now: () => now, isProcessAlive: pid => pid === process.pid }); + const dead = join(root, "pi-web-clipboard-v1-999999-ABCdef"); + const live = join(root, `pi-web-clipboard-v1-${process.pid}-ABCdef`); + const unrelated = join(root, "other-app"); + for (const directory of [dead, live, unrelated]) { + await fs.mkdir(directory, { mode: 0o700 }); + await fs.writeFile(join(directory, "keep"), "fixture"); + const old = new Date(now - IMAGE_IDLE_TTL_MS - 1); + await fs.utimes(directory, old, old); + } + const link = join(root, "pi-web-clipboard-v1-999998-ABCdef"); + await fs.symlink(unrelated, link); + await store.collect(); + await assert.rejects(fs.stat(dead), { code: "ENOENT" }); + assert.equal(await fs.readFile(join(live, "keep"), "utf8"), "fixture"); + assert.equal(await fs.readFile(join(unrelated, "keep"), "utf8"), "fixture"); + assert.equal((await fs.lstat(link)).isSymbolicLink(), true); +}); + +test("path annotation preserves slash command token and safely quotes paths", () => { + assert.equal(appendImagePaths("hello", []), "hello"); + const annotated = appendImagePaths("/command", ["/tmp/a b.jpg", "/tmp/a\nb.png"]); + assert.equal(annotated.split(" ")[0], "/command"); + assert.ok(annotated.includes(JSON.stringify("/tmp/a b.jpg"))); + assert.ok(annotated.includes(JSON.stringify("/tmp/a\nb.png"))); +}); + +test("cleanup waits for an in-flight batch and starts its lifetime after writing", async (t) => { + let now = Date.now(); + const { store } = await fixture(t, { now: () => now }); + let release; + let started; + const gate = new Promise(resolve => { release = resolve; }); + const entered = new Promise(resolve => { started = resolve; }); + const write = fs.writeFile; + t.mock.method(fs, "writeFile", async (...args) => { + started(); + await gate; + return write(...args); + }); + const pending = store.materialize([jpeg]); + await entered; + now += IMAGE_IDLE_TTL_MS + 1; + const collected = store.collect(); + release(); + const batch = await pending; + await collected; + assert.deepEqual(await fs.readFile(batch.paths[0]), bytes); +}); + +test("normal process exit cleans its files without being kept alive by the GC timer", async (t) => { + const { root } = await fixture(t); + const code = ` + const { createJiti } = require("jiti"); + const jiti = createJiti(process.cwd() + "/image-exit-test.mjs"); + (async () => { + const { getImageAttachmentStore } = await jiti.import("./lib/image-materialization.ts"); + const batch = await getImageAttachmentStore(() => false).materialize(${JSON.stringify([jpeg])}); + console.log(JSON.stringify(batch.paths)); + })().catch(error => { console.error(error); process.exitCode = 1; }); + `; + const { stdout } = await promisify(execFile)(process.execPath, ["-e", code], { + cwd: new URL("..", import.meta.url), + env: { ...process.env, TMPDIR: root, TMP: root, TEMP: root }, + timeout: 15_000, + }); + const paths = JSON.parse(stdout); + assert.equal(paths.length, 1); + await assert.rejects(fs.stat(paths[0]), { code: "ENOENT" }); + assert.deepEqual(await fs.readdir(root), []); +}); diff --git a/lib/image-materialization.ts b/lib/image-materialization.ts new file mode 100644 index 000000000..e5ce4acdc --- /dev/null +++ b/lib/image-materialization.ts @@ -0,0 +1,169 @@ +import fs from "node:fs/promises"; +import { rmSync } from "node:fs"; +import { randomUUID } from "node:crypto"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { validateAgentImages, type Base64ImageAttachment } from "./image-attachments"; + +export const IMAGE_IDLE_TTL_MS = 24 * 60 * 60 * 1000; +const GC_INTERVAL_MS = 60 * 60 * 1000; +const OWNER_DIRECTORY = /^pi-web-clipboard-v1-([1-9]\d*)-[a-zA-Z0-9]{6}$/; +const MIME_EXTENSIONS: Record = { + "image/jpeg": "jpg", "image/png": "png", "image/gif": "gif", "image/webp": "webp", + "image/avif": "avif", "image/svg+xml": "svg", "image/bmp": "bmp", + "image/tiff": "tiff", "image/heic": "heic", "image/heif": "heif", + "image/x-icon": "ico", "image/vnd.microsoft.icon": "ico", "image/x-ms-bmp": "bmp", +}; + +export interface MaterializedImages { + paths: string[]; + rollback: () => Promise; +} + +type AgentImageAttachment = Base64ImageAttachment & { type: "image" }; + +function isProcessAlive(pid: number): boolean { + try { + process.kill(pid, 0); + return true; + } catch (error) { + // Permission errors and other unexpected failures must not count as process exit. + return (error as NodeJS.ErrnoException).code !== "ESRCH"; + } +} + +export class ImageAttachmentStore { + private readonly root: string; + private readonly now: () => number; + private readonly isBusy: () => boolean; + private readonly isProcessAlive: (pid: number) => boolean; + private directory?: string; + private lastActivity: number; + private tail: Promise = Promise.resolve(); + + constructor(options: { + root?: string; + now?: () => number; + isBusy?: () => boolean; + isProcessAlive?: (pid: number) => boolean; + } = {}) { + this.root = options.root ?? tmpdir(); + this.now = options.now ?? Date.now; + this.isBusy = options.isBusy ?? (() => false); + this.isProcessAlive = options.isProcessAlive ?? isProcessAlive; + this.lastActivity = this.now(); + } + + touch(): void { + this.lastActivity = this.now(); + } + + private exclusive(operation: () => Promise): Promise { + // Cleanup, rollback and writes share one queue so async deletion cannot race a new batch write. + const result = this.tail.then(operation); + this.tail = result.catch(() => {}); + return result; + } + + materialize(images: AgentImageAttachment[]): Promise { + return this.exclusive(async () => { + const error = validateAgentImages(images); + if (error) throw new Error(error); + const extensions = images.map(image => { + const extension = MIME_EXTENSIONS[image.mimeType]; + if (!extension) throw new Error(`Unsupported image MIME type: ${image.mimeType}`); + return extension; + }); + this.touch(); + if (!images.length) return { paths: [], rollback: async () => {} }; + if (!this.directory) { + await this.collectOrphans(); + this.directory = await fs.mkdtemp(join(this.root, `pi-web-clipboard-v1-${process.pid}-`)); + await fs.chmod(this.directory, 0o700); + } + const batch = await fs.mkdtemp(join(this.directory, "batch-")); + const paths: string[] = []; + try { + await fs.chmod(batch, 0o700); + for (let index = 0; index < images.length; index++) { + const path = join(batch, `pi-web-clipboard-${randomUUID()}.${extensions[index]}`); + await fs.writeFile(path, Buffer.from(images[index].data, "base64"), { mode: 0o600, flag: "wx" }); + paths.push(path); + } + } catch (error) { + await fs.rm(batch, { recursive: true, force: true }).catch(reportCleanupError); + throw error; + } + this.touch(); + return { + paths, + rollback: () => this.exclusive(() => fs.rm(batch, { recursive: true, force: true })), + }; + }); + } + + collect(): Promise { + return this.exclusive(async () => { + if (this.isBusy()) this.touch(); + if (this.directory && this.now() - this.lastActivity >= IMAGE_IDLE_TTL_MS) { + await fs.rm(this.directory, { recursive: true, force: true }); + this.directory = undefined; + } + await this.collectOrphans(); + }); + } + + private async collectOrphans(): Promise { + for (const entry of await fs.readdir(this.root, { withFileTypes: true })) { + const match = OWNER_DIRECTORY.exec(entry.name); + if (!entry.isDirectory() || !match) continue; + const pid = Number(match[1]); + if (!Number.isSafeInteger(pid) || this.isProcessAlive(pid)) continue; + const path = join(this.root, entry.name); + try { + const stat = await fs.lstat(path); + if (!stat.isDirectory() || stat.isSymbolicLink()) continue; + if (process.getuid && stat.uid !== process.getuid()) continue; + if (this.now() - stat.mtimeMs < IMAGE_IDLE_TTL_MS) continue; + await fs.rm(path, { recursive: true, force: true }); + } catch (error) { + if ((error as NodeJS.ErrnoException).code !== "ENOENT") reportCleanupError(error); + } + } + } + + disposeSync(): void { + if (!this.directory) return; + rmSync(this.directory, { recursive: true, force: true }); + this.directory = undefined; + } +} + +function reportCleanupError(error: unknown): void { + console.error("[pi-web] image attachment cleanup failed:", error); +} + +export function appendImagePaths(message: string, paths: string[]): string { + if (!paths.length) return message; + // The space preserves the SDK's split-on-first-space /command behavior; JSON quoting keeps + // newlines inside a path from confusing the annotation. + return `${message} \n\n[Attached image files on this server (temporary):\n${paths.map((path, index) => `${index + 1}. ${JSON.stringify(path)}`).join("\n")}\nThe same images are also included as image inputs. Use these paths for file/image tools.]`; +} + +declare global { + var __piImageAttachmentStore: ImageAttachmentStore | undefined; +} + +export function getImageAttachmentStore(isBusy: () => boolean): ImageAttachmentStore { + if (!globalThis.__piImageAttachmentStore) { + const store = new ImageAttachmentStore({ isBusy }); + globalThis.__piImageAttachmentStore = store; + const timer = setInterval(() => { void store.collect().catch(reportCleanupError); }, GC_INTERVAL_MS); + timer.unref(); + process.once("exit", () => { + clearInterval(timer); + try { store.disposeSync(); } catch (error) { reportCleanupError(error); } + }); + } + return globalThis.__piImageAttachmentStore; +} diff --git a/lib/rpc-manager-images.test.mjs b/lib/rpc-manager-images.test.mjs new file mode 100644 index 000000000..12b09798a --- /dev/null +++ b/lib/rpc-manager-images.test.mjs @@ -0,0 +1,107 @@ +import assert from "node:assert/strict"; +import { mkdtemp, readFile, readdir, rm, stat } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { dirname, join } from "node:path"; +import test from "node:test"; +import { createReadTool } from "@earendil-works/pi-coding-agent"; +import { createJiti } from "jiti"; + +const jiti = createJiti(import.meta.url, { tsconfigPaths: true }); +const { AgentSessionWrapper } = await jiti.import("./rpc-manager.ts"); +const { ImageAttachmentStore } = await jiti.import("./image-materialization.ts"); +const png = { type: "image", mimeType: "image/png", data: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+jRZkAAAAASUVORK5CYII=" }; + +async function fixture(t, overrides = {}) { + const root = await mkdtemp(join(tmpdir(), "pi-web-rpc-images-")); + const previous = globalThis.__piImageAttachmentStore; + const store = new ImageAttachmentStore({ root }); + globalThis.__piImageAttachmentStore = store; + const calls = []; + const capture = async (message, images) => { + const paths = [...message.matchAll(/^\d+\. (".*")$/gm)].map(match => JSON.parse(match[1])); + assert.equal(paths.length, images?.length); + for (const path of paths) assert.deepEqual(await readFile(path), Buffer.from(png.data, "base64")); + calls.push({ message, images, paths }); + }; + const inner = { + sessionId: "image-test", sessionManager: { getCwd: () => root }, + isStreaming: false, isCompacting: false, isBashRunning: false, + agent: { state: {} }, extensionRunner: {}, dispose() {}, + prompt: async (message, options) => { + await capture(message, options.images); + options.preflightResult(true); + }, + steer: capture, followUp: capture, + ...overrides, + }; + const wrapper = new AgentSessionWrapper(inner); + t.after(async () => { + wrapper.destroy(); + store.disposeSync(); + globalThis.__piImageAttachmentStore = previous; + await rm(root, { recursive: true, force: true }); + }); + return { wrapper, calls, root, store, inner }; +} + +for (const command of [ + { type: "prompt" }, { type: "prompt", streamingBehavior: "steer" }, + { type: "prompt", streamingBehavior: "followUp" }, { type: "steer" }, { type: "follow_up" }, +]) test(`provides readable paths and unchanged images for ${JSON.stringify(command)}`, async (t) => { + const { wrapper, calls } = await fixture(t); + const images = [png, png]; + await wrapper.send({ ...command, message: "Inspect these", images }); + assert.equal(calls.length, 1); + assert.equal(calls[0].images, images); + assert.ok(calls[0].message.startsWith("Inspect these")); + const reader = createReadTool(tmpdir(), { autoResizeImages: false }); + const result = await reader.execute("read-fixture", { path: calls[0].paths[0] }); + assert.ok(result.content.some(block => block.type === "image" && block.mimeType === "image/png")); + wrapper.destroy(); + assert.deepEqual(await readFile(calls[0].paths[0]), Buffer.from(png.data, "base64")); +}); + +test("rejected preflight rolls back only its own image batch", async (t) => { + let rejectedPath; + const { wrapper, store } = await fixture(t, { + prompt: async (message, options) => { + rejectedPath = JSON.parse(message.match(/^1\. (".*")$/m)[1]); + assert.ok((await stat(rejectedPath)).isFile()); + options.preflightResult(false); + throw new Error("fixture rejected"); + }, + }); + const kept = await store.materialize([png]); + await assert.rejects(wrapper.send({ type: "prompt", message: "reject", images: [png] }), /fixture rejected/); + await assert.rejects(stat(rejectedPath), { code: "ENOENT" }); + assert.ok((await stat(kept.paths[0])).isFile()); + assert.equal((await readdir(dirname(dirname(kept.paths[0])))).length, 1); +}); + +test("failure after acceptance retains files referenced by the accepted message", async (t) => { + let path; + let fail; + const { wrapper } = await fixture(t, { + prompt: (message, options) => { + path = JSON.parse(message.match(/^1\. (".*")$/m)[1]); + options.preflightResult(true); + return new Promise((_, reject) => { fail = reject; }); + }, + }); + await wrapper.send({ type: "prompt", message: "accepted", images: [png] }); + fail(new Error("fixture provider failure")); + await new Promise(resolve => setImmediate(resolve)); + assert.ok((await stat(path)).isFile()); +}); + +test("text-only prompts are unchanged and allocate no image files", async (t) => { + let received; + const { wrapper, root } = await fixture(t, { prompt: async (message, options) => { + received = message; + assert.equal(options.images, undefined); + options.preflightResult(true); + } }); + await wrapper.send({ type: "prompt", message: "/command" }); + assert.equal(received, "/command"); + assert.deepEqual(await readdir(root), []); +}); diff --git a/lib/rpc-manager.ts b/lib/rpc-manager.ts index 9c257b172..f50786847 100644 --- a/lib/rpc-manager.ts +++ b/lib/rpc-manager.ts @@ -5,6 +5,7 @@ import { randomUUID } from "crypto"; import { existsSync, realpathSync, writeFileSync } from "fs"; import { resolve } from "path"; import { validateAgentImages } from "./image-attachments"; +import { appendImagePaths, getImageAttachmentStore, type MaterializedImages } from "./image-materialization"; import { invalidateModelsCache } from "./models-cache"; import { resolveVisibleModels, selectInitialModelScope } from "./model-scope"; import { @@ -551,6 +552,9 @@ export class AgentSessionWrapper { const tracksMutation = !allowedDuringReplacement; if (tracksMutation) this.activeMutatingCommands += 1; + let imageBatch: MaterializedImages | undefined; + let imagesAccepted = false; + let message = command.message as string; try { // Status reconciliation must not postpone forced cleanup after Stop. @@ -563,6 +567,17 @@ export class AgentSessionWrapper { if (type === "prompt" || type === "steer" || type === "follow_up") { const imageError = validateAgentImages(command.images); if (imageError) throw new Error(imageError); + globalThis.__piImageAttachmentStore?.touch(); + const images = command.images as Array<{ type: "image"; data: string; mimeType: string }> | undefined; + if (images?.length) { + if (typeof message !== "string") throw new Error("message must be a string"); + const store = getImageAttachmentStore(() => [...(globalThis.__piSessions?.values() ?? [])].some(session => + session.isRunning() || session.activeMutatingCommands > 0 + || Boolean(session.inner.getSteeringMessages?.().length || session.inner.getFollowUpMessages?.().length), + )); + imageBatch = await store.materialize(images); + message = appendImagePaths(message, imageBatch.paths); + } } switch (type) { @@ -587,6 +602,7 @@ export class AgentSessionWrapper { let rejectPreflight!: (error: unknown) => void; const preflight = new Promise((resolve, reject) => { acceptPreflight = () => { + imagesAccepted = true; preflightAccepted = true; this.agentRunNeedsCompletion = true; if (preflightSettled) return; @@ -610,7 +626,7 @@ export class AgentSessionWrapper { this.pendingPromptCount += 1; let prompt: Promise; try { - prompt = this.inner.prompt(command.message as string, { + prompt = this.inner.prompt(message, { ...(promptImages?.length ? { images: promptImages } : {}), ...(streamingBehavior ? { streamingBehavior } : {}), source: "rpc", @@ -867,13 +883,15 @@ export class AgentSessionWrapper { case "steer": { const steerImages = command.images as Array<{ type: "image"; data: string; mimeType: string }> | undefined; - await this.inner.steer(command.message as string, steerImages?.length ? steerImages : undefined); + await this.inner.steer(message, steerImages?.length ? steerImages : undefined); + imagesAccepted = true; return null; } case "follow_up": { const followImages = command.images as Array<{ type: "image"; data: string; mimeType: string }> | undefined; - await this.inner.followUp(command.message as string, followImages?.length ? followImages : undefined); + await this.inner.followUp(message, followImages?.length ? followImages : undefined); + imagesAccepted = true; return null; } @@ -993,6 +1011,13 @@ export class AgentSessionWrapper { default: throw new Error(`Unsupported command: ${type}`); } + } catch (error) { + if (imageBatch && !imagesAccepted) { + await imageBatch.rollback().catch(cleanupError => { + console.error("[pi-web] rejected image attachment cleanup failed:", cleanupError); + }); + } + throw error; } finally { if (tracksMutation) this.activeMutatingCommands = Math.max(0, this.activeMutatingCommands - 1); }