@narumitw/pi-webui 0.20.2 → 0.22.0

package/src/runtime.ts

@@ -8,11 +8,16 @@ import {
 	getSettingsListTheme,
 } from "@earendil-works/pi-coding-agent";
 import { Container, type SettingItem, SettingsList, Text } from "@earendil-works/pi-tui";
+import type { PreparedAttachment } from "./attachments.js";
 import { ConversationProjection, projectBranchMessages } from "./conversation.js";
+import { DEFAULT_IMAGE_LIMITS, type ImageLimits, imageLimits } from "./image-limits.js";
 import {
 	type BrowserImageInput,
 	type ProcessBrowserImageOptions,
+	type ProcessedBrowserImage,
 	processBrowserImages,
+	processStagedImage,
+	validateStagedBrowserImages,
 } from "./images.js";
 import { type EffectivePiImageSettings, readEffectivePiImageSettings } from "./pi-settings.js";
 import {
@@ -51,6 +56,12 @@ interface PendingBrowserInput {
 	resolve(): void;
 	reject(error: Error): void;
 	text: string;
+	retainedImageIds: string[];
+}
+
+interface AcceptedBrowserInput {
+	text: string;
+	retainedImageIds: string[];
 }
 
 export interface RuntimeDependencies {
@@ -63,6 +74,10 @@ export interface RuntimeDependencies {
 		inputs: BrowserImageInput[],
 		options?: ProcessBrowserImageOptions,
 	): Promise<ImageContent[]>;
+	processAttachment(
+		source: Uint8Array,
+		options?: ProcessBrowserImageOptions,
+	): Promise<ProcessedBrowserImage>;
 }
 
 const DEFAULT_DEPENDENCIES: RuntimeDependencies = {
@@ -72,6 +87,7 @@ const DEFAULT_DEPENDENCIES: RuntimeDependencies = {
 	startServer: (options) => WebUIServer.start(options),
 	readPiSettings: readEffectivePiImageSettings,
 	processImages: processBrowserImages,
+	processAttachment: processStagedImage,
 };
 
 export class WebUIRuntime {
@@ -88,8 +104,9 @@ export class WebUIRuntime {
 	private readonly activeMessageIds = new Map<string, string>();
 	private readonly finalMessageTimers = new Set<ReturnType<typeof setTimeout>>();
 	private readonly pendingBrowserInputs = new Map<string, PendingBrowserInput>();
-	private readonly acceptedBrowserInputs = new Map<string, string>();
+	private readonly acceptedBrowserInputs = new Map<string, AcceptedBrowserInput>();
 	private settings: WebUISettings = { ...DEFAULT_SETTINGS };
+	private effectiveImageLimits: Readonly<ImageLimits> = DEFAULT_IMAGE_LIMITS;
 	private settingsDocument?: Record<string, unknown> = {};
 	private settingsPath = "pi-webui.json";
 	private settingsSource: SettingsLoadResult["source"] = "defaults";
@@ -164,7 +181,10 @@ export class WebUIRuntime {
 			if (!wrapped) return;
 			const pending = this.pendingBrowserInputs.get(wrapped.nonce);
 			if (!pending) return;
-			this.acceptedBrowserInputs.set(wrapped.nonce, pending.text);
+			this.acceptedBrowserInputs.set(wrapped.nonce, {
+				text: pending.text,
+				retainedImageIds: [...pending.retainedImageIds],
+			});
 			this.settleBrowserInput(wrapped.nonce);
 		});
 		this.pi.on("message_start", async (event, ctx) => {
@@ -281,8 +301,11 @@ export class WebUIRuntime {
 			id = `web-live:${++this.nextLiveMessageId}`;
 			this.activeMessageIds.set(key, id);
 		}
+		const retainedImageIds = Array.isArray(event.retainedImageIds)
+			? event.retainedImageIds.filter((value): value is string => typeof value === "string")
+			: [];
 		if (phase !== "end") {
-			this.recordProjectedMessage(event.message, false, id);
+			this.recordProjectedMessage(event.message, false, id, retainedImageIds);
 			return;
 		}
 		this.activeMessageIds.delete(key);
@@ -290,14 +313,19 @@ export class WebUIRuntime {
 		const timer = setTimeout(() => {
 			this.finalMessageTimers.delete(timer);
 			if (generation !== this.generation || this.closed) return;
-			this.recordProjectedMessage(event.message, true, id);
+			this.recordProjectedMessage(event.message, true, id, retainedImageIds);
 		}, 0);
 		this.finalMessageTimers.add(timer);
 	}
 
-	private recordProjectedMessage(message: unknown, final: boolean, id: string): void {
+	private recordProjectedMessage(
+		message: unknown,
+		final: boolean,
+		id: string,
+		retainedImageIds: readonly string[] = [],
+	): void {
 		try {
-			this.conversation?.recordMessage(message, final, id);
+			this.conversation?.recordMessage(message, final, id, retainedImageIds);
 		} catch {
 			// Unknown custom message shapes do not block the supported transcript.
 		}
@@ -369,7 +397,7 @@ export class WebUIRuntime {
 							if (!this.settingsDocument) {
 								throw new Error("the invalid settings file must be repaired manually first");
 							}
-							const next = { startOnSessionStart: requested };
+							const next = { ...this.settings, startOnSessionStart: requested };
 							const document = await this.dependencies.saveSettings(
 								next,
 								this.settingsDocument,
@@ -409,6 +437,7 @@ export class WebUIRuntime {
 			[
 				"Pi WebUI status",
 				`startOnSessionStart: ${this.settings.startOnSessionStart} (${source})`,
+				`Image limits (${source}): ${this.effectiveImageLimits.maxImages} images, ${formatMib(this.effectiveImageLimits.maxImageBytes)}/image, ${formatMib(this.effectiveImageLimits.maxBatchBytes)}/batch, ${this.effectiveImageLimits.maxImagePixels.toLocaleString("en-US")} pixels/image`,
 				`Settings: ${this.settingsPath}`,
 				`Server: ${this.server ? "running" : "stopped"}`,
 			].join("\n"),
@@ -425,9 +454,10 @@ export class WebUIRuntime {
 				"settings: edit WebUI settings in TUI mode",
 				"status: show effective settings, source, path, and current server state",
 				"init: create the defaults file without overwriting existing content",
-				'Accepted JSON: { "startOnSessionStart": false }',
+				'Accepted JSON starts with { "startOnSessionStart": false } and may include advanced maxImages, maxImageBytes, maxBatchBytes, and maxImagePixels fields.',
 				`Settings path: ${this.settingsPath}`,
-				"The default is false. Changes apply on the next session initialization or reload.",
+				"Image byte/pixel limits stay in Advanced JSON; Pi provider-ready dimension/Base64 limits are fixed.",
+				"Changes apply on the next session initialization or reload.",
 			].join("\n"),
 			"info",
 		);
@@ -451,6 +481,7 @@ export class WebUIRuntime {
 
 	private applySettingsResult(result: SettingsLoadResult): void {
 		this.settings = { ...result.settings };
+		this.effectiveImageLimits = imageLimits(result.settings);
 		this.settingsDocument = result.kind === "invalid" ? undefined : { ...(result.document ?? {}) };
 		this.settingsPath = result.path;
 		this.settingsSource = result.source;
@@ -465,6 +496,14 @@ export class WebUIRuntime {
 			const starting = this.dependencies.startServer({
 				conversation,
 				send: (request) => this.sendBrowserMessage(request, generation),
+				imageLimits: this.effectiveImageLimits,
+				sentImageSettings: {
+					enabled: this.settings.retainSentImages,
+					maxImages: this.settings.maxRetainedImages,
+					maxBytes: this.settings.maxRetainedBytes,
+				},
+				processAttachment: (source, signal) =>
+					this.processStagedAttachment(source, generation, signal),
 			});
 			this.serverStarting = starting.then(async (server) => {
 				if (generation !== this.generation || this.closed || conversation !== this.conversation) {
@@ -483,6 +522,37 @@ export class WebUIRuntime {
 		}
 	}
 
+	private async processStagedAttachment(
+		source: Uint8Array,
+		generation: number,
+		signal?: AbortSignal,
+	): Promise<PreparedAttachment> {
+		const ctx = this.context;
+		if (!ctx || this.closed || generation !== this.generation) {
+			throw new Error("The Pi session has ended.");
+		}
+		const combinedSignal = signal
+			? AbortSignal.any([this.sessionAbort.signal, signal])
+			: this.sessionAbort.signal;
+		const settings = await this.dependencies.readPiSettings(ctx.cwd, ctx.isProjectTrusted());
+		this.notifySettingsWarnings(ctx, settings.warnings);
+		const image = await this.dependencies.processAttachment(source, {
+			maxImageBytes: this.effectiveImageLimits.maxImageBytes,
+			maxPixels: this.effectiveImageLimits.maxImagePixels,
+			autoResize: settings.autoResize,
+			blockImages: settings.blockImages,
+			supportsImages: ctx.model?.input.includes("image") ?? false,
+			signal: combinedSignal,
+		});
+		if (combinedSignal.aborted || generation !== this.generation || this.closed) {
+			throw new Error("Image processing was cancelled.");
+		}
+		return {
+			...image,
+			notes: attachmentNotes(image),
+		};
+	}
+
 	private async sendBrowserMessage(
 		request: WebSendRequest,
 		generation: number,
@@ -499,15 +569,16 @@ export class WebUIRuntime {
 		await this.preflightIdlePrompt(ctx, request, generation, signal);
 		let images: ImageContent[] = [];
 		if (request.images.length > 0) {
+			validateStagedBrowserImages(request.images, this.effectiveImageLimits);
 			const settings = await this.dependencies.readPiSettings(ctx.cwd, ctx.isProjectTrusted());
 			this.notifySettingsWarnings(ctx, settings.warnings);
-			images = await this.dependencies.processImages(request.images, {
-				autoResize: settings.autoResize,
-				blockImages: settings.blockImages,
-				supportsImages: ctx.model?.input.includes("image") ?? false,
-				signal,
-			});
+			if (settings.blockImages) throw new Error("Pi image sending is disabled.");
 			await this.validateCurrentModel(ctx, generation, signal, true);
+			images = request.images.map((image) => ({
+				type: "image" as const,
+				data: image.data,
+				mimeType: image.mimeType ?? "image/png",
+			}));
 		}
 		if (signal.aborted) throw new Error("The browser message was cancelled.");
 		if (!this.context || this.closed || generation !== this.generation) {
@@ -519,7 +590,7 @@ export class WebUIRuntime {
 			images.length === 0
 				? text
 				: [...(text.trim() ? ([{ type: "text", text }] satisfies TextContent[]) : []), ...images];
-		const wrapped = this.createBrowserInput(content);
+		const wrapped = this.createBrowserInput(content, request.retainedImageIds ?? []);
 		const delivery =
 			request.delivery === "steer"
 				? "steer"
@@ -574,7 +645,10 @@ export class WebUIRuntime {
 		if (ctx.model !== model) throw new Error("The Pi model changed; retry the browser message.");
 	}
 
-	private createBrowserInput(content: string | Array<TextContent | ImageContent>): {
+	private createBrowserInput(
+		content: string | Array<TextContent | ImageContent>,
+		retainedImageIds: readonly string[],
+	): {
 		nonce: string;
 		content: string | Array<TextContent | ImageContent>;
 		accepted: Promise<void>;
@@ -590,7 +664,12 @@ export class WebUIRuntime {
 						...content.filter((part): part is ImageContent => part.type === "image"),
 					];
 		const accepted = new Promise<void>((resolve, reject) => {
-			this.pendingBrowserInputs.set(nonce, { resolve, reject, text });
+			this.pendingBrowserInputs.set(nonce, {
+				resolve,
+				reject,
+				text,
+				retainedImageIds: [...retainedImageIds],
+			});
 		});
 		return { nonce, content: wrappedContent, accepted };
 	}
@@ -653,35 +732,56 @@ function contentText(content: Array<TextContent | ImageContent>): string {
 
 function sanitizeBrowserMessageEvent(
 	event: unknown,
-	acceptedNonces: Map<string, string>,
+	acceptedNonces: Map<string, AcceptedBrowserInput>,
 	consume: boolean,
 ): unknown {
 	if (!isRecord(event) || !isRecord(event.message) || event.message.role !== "user") return event;
 	const content = event.message.content;
 	if (typeof content === "string") {
 		const wrapped = parseBrowserInput(content);
-		if (!wrapped || !acceptedNonces.has(wrapped.nonce)) return event;
-		const text = acceptedNonces.get(wrapped.nonce) ?? "";
+		const accepted = wrapped ? acceptedNonces.get(wrapped.nonce) : undefined;
+		if (!wrapped || !accepted) return event;
 		if (consume) acceptedNonces.delete(wrapped.nonce);
-		return { ...event, message: { ...event.message, content: text } };
+		return {
+			...event,
+			retainedImageIds: [...accepted.retainedImageIds],
+			message: { ...event.message, content: accepted.text },
+		};
 	}
 	if (!Array.isArray(content)) return event;
-	let changed = false;
-	const consumed = new Set<string>();
-	const sanitized = content.flatMap((part) => {
+	let accepted: AcceptedBrowserInput | undefined;
+	let acceptedNonce: string | undefined;
+	const sanitizedText = content.flatMap((part) => {
 		if (!isRecord(part) || part.type !== "text" || typeof part.text !== "string") return [part];
 		const wrapped = parseBrowserInput(part.text);
-		if (!wrapped || !acceptedNonces.has(wrapped.nonce)) return [part];
-		const text = acceptedNonces.get(wrapped.nonce) ?? "";
-		changed = true;
-		consumed.add(wrapped.nonce);
-		return text ? [{ ...part, text }] : [];
+		const candidate = wrapped ? acceptedNonces.get(wrapped.nonce) : undefined;
+		if (!wrapped || !candidate) return [part];
+		accepted = candidate;
+		acceptedNonce = wrapped.nonce;
+		return candidate.text ? [{ ...part, text: candidate.text }] : [];
 	});
-	if (!changed) return event;
-	if (consume) {
-		for (const nonce of consumed) acceptedNonces.delete(nonce);
+	if (!accepted) return event;
+	if (consume && acceptedNonce) acceptedNonces.delete(acceptedNonce);
+	return {
+		...event,
+		retainedImageIds: [...accepted.retainedImageIds],
+		message: { ...event.message, content: sanitizedText },
+	};
+}
+
+function attachmentNotes(image: ProcessedBrowserImage): string[] {
+	const notes: string[] = [];
+	if (image.sourceFormat !== image.outputFormat) {
+		notes.push(
+			`Converted ${image.sourceFormat.toUpperCase()} to ${image.outputFormat.toUpperCase()}`,
+		);
+	}
+	if (image.resized) {
+		notes.push(
+			`Resized ${image.originalWidth}×${image.originalHeight} to ${image.width}×${image.height}`,
+		);
 	}
-	return { ...event, message: { ...event.message, content: sanitized } };
+	return notes;
 }
 
 function messageLifecycleKey(message: Record<string, unknown>): string {
@@ -692,6 +792,10 @@ function isRecord(value: unknown): value is Record<string, unknown> {
 	return typeof value === "object" && value !== null && !Array.isArray(value);
 }
 
+function formatMib(bytes: number): string {
+	return `${bytes / (1024 * 1024)} MiB`;
+}
+
 function formatError(error: unknown): string {
 	return error instanceof Error ? error.message : String(error);
 }

package/src/sent-images.ts

@@ -0,0 +1,279 @@
+import { createHmac, randomBytes } from "node:crypto";
+
+export interface SentImageSettings {
+	enabled: boolean;
+	maxImages: number;
+	maxBytes: number;
+}
+
+export interface RetainedImageInput {
+	name?: string;
+	mimeType?: string;
+	data: string;
+}
+
+export interface RetainedImageClone {
+	retainedId: string;
+	name: string;
+	bytes: Buffer;
+	mimeType: string;
+}
+
+export interface PublicSentImage {
+	id: string;
+	name: string;
+	mimeType: string;
+	size: number;
+}
+
+export interface PublicSentImageState {
+	revision: number;
+	enabled: boolean;
+	items: PublicSentImage[];
+	totalBytes: number;
+	maxImages: number;
+	maxBytes: number;
+}
+
+interface SentImageItem {
+	id: string;
+	name: string;
+	mimeType: string;
+	bytes: Buffer;
+	associations: Set<string>;
+}
+
+export class SentImageError extends Error {
+	constructor(
+		message: string,
+		readonly status = 409,
+	) {
+		super(message);
+		this.name = "SentImageError";
+	}
+}
+
+export class SentImageStore {
+	private readonly key = randomBytes(32);
+	private readonly items: SentImageItem[] = [];
+	private revision = 0;
+	private closed = false;
+
+	constructor(private readonly settings: SentImageSettings) {
+		for (const value of [settings.maxImages, settings.maxBytes]) {
+			if (!Number.isSafeInteger(value) || value <= 0) {
+				throw new Error("Sent-image limits must be positive integers.");
+			}
+		}
+	}
+
+	publicState(): PublicSentImageState {
+		return {
+			revision: this.revision,
+			enabled: this.settings.enabled,
+			items: this.items.map((item) => ({
+				id: item.id,
+				name: item.name,
+				mimeType: item.mimeType,
+				size: item.bytes.byteLength,
+			})),
+			totalBytes: this.residentBytes(),
+			maxImages: this.settings.maxImages,
+			maxBytes: this.settings.maxBytes,
+		};
+	}
+
+	residentBytes(): number {
+		return this.items.reduce((total, item) => total + item.bytes.byteLength, 0);
+	}
+
+	referencesFor(messageId: string, images: readonly RetainedImageInput[]): string[] {
+		this.assertOpen();
+		if (!this.settings.enabled) return [];
+		normalizeMessageId(messageId);
+		return images.map((image, index) => {
+			const normalized = normalizeImage(image, index);
+			return this.idFor(normalized.bytes);
+		});
+	}
+
+	commit(
+		messageId: string,
+		images: readonly RetainedImageInput[],
+		expectedReferences: readonly string[],
+	): PublicSentImageState {
+		this.assertOpen();
+		if (!this.settings.enabled) {
+			if (expectedReferences.length !== 0) {
+				throw new SentImageError("Sent-image retention is disabled.");
+			}
+			return this.publicState();
+		}
+		normalizeMessageId(messageId);
+		const normalized = images.map(normalizeImage);
+		const references = normalized.map((image) => this.idFor(image.bytes));
+		if (!sameIds(references, expectedReferences)) {
+			throw new SentImageError("Sent-image references are stale.");
+		}
+		let changed = false;
+		for (const [index, image] of normalized.entries()) {
+			const id = references[index];
+			if (!id) throw new SentImageError("Sent-image reference is missing.");
+			const existing = this.items.find((item) => item.id === id);
+			if (existing) {
+				existing.associations.add(messageId);
+				continue;
+			}
+			this.items.push({
+				id,
+				name: image.name,
+				mimeType: image.mimeType,
+				bytes: Buffer.from(image.bytes),
+				associations: new Set([messageId]),
+			});
+			changed = true;
+		}
+		if (this.evict()) changed = true;
+		if (changed) this.revision += 1;
+		return this.publicState();
+	}
+
+	preview(id: string): { bytes: Buffer; mimeType: string } {
+		const item = this.item(id);
+		return { bytes: Buffer.from(item.bytes), mimeType: item.mimeType };
+	}
+
+	clone(ids: readonly string[]): RetainedImageClone[] {
+		this.assertOpen();
+		if (!Array.isArray(ids) || ids.length === 0) {
+			throw new SentImageError("No retained images were selected.", 400);
+		}
+		if (new Set(ids).size !== ids.length) {
+			throw new SentImageError("Retained image selection contains a duplicate.", 400);
+		}
+		return ids.map((id) => {
+			const item = this.item(id);
+			return {
+				retainedId: item.id,
+				name: item.name,
+				bytes: Buffer.from(item.bytes),
+				mimeType: item.mimeType,
+			};
+		});
+	}
+
+	reconcile(externalBytes: number, totalBudget = this.settings.maxBytes): PublicSentImageState {
+		this.assertOpen();
+		if (
+			!Number.isSafeInteger(externalBytes) ||
+			externalBytes < 0 ||
+			!Number.isSafeInteger(totalBudget) ||
+			totalBudget <= 0
+		) {
+			throw new SentImageError("Image resident-byte accounting is invalid.", 500);
+		}
+		if (!this.evict(externalBytes, totalBudget)) return this.publicState();
+		this.revision += 1;
+		return this.publicState();
+	}
+
+	remove(id: string, expectedRevision: number): PublicSentImageState {
+		this.assertRevision(expectedRevision);
+		const index = this.items.findIndex((item) => item.id === id);
+		if (index === -1) throw new SentImageError("Retained image was not found.", 404);
+		this.items.splice(index, 1);
+		this.revision += 1;
+		return this.publicState();
+	}
+
+	clear(expectedRevision: number): PublicSentImageState {
+		this.assertRevision(expectedRevision);
+		if (this.items.length === 0) return this.publicState();
+		this.items.length = 0;
+		this.revision += 1;
+		return this.publicState();
+	}
+
+	close(): void {
+		if (this.closed) return;
+		this.closed = true;
+		this.items.length = 0;
+		this.revision += 1;
+	}
+
+	private item(id: string): SentImageItem {
+		this.assertOpen();
+		if (typeof id !== "string" || !/^[A-Za-z0-9_-]{1,128}$/.test(id)) {
+			throw new SentImageError("Retained image id is invalid.", 400);
+		}
+		const item = this.items.find((candidate) => candidate.id === id);
+		if (!item) throw new SentImageError("Retained image has expired.", 404);
+		return item;
+	}
+
+	private idFor(bytes: Buffer): string {
+		return `sent_${createHmac("sha256", this.key).update(bytes).digest("base64url").slice(0, 32)}`;
+	}
+
+	private evict(externalBytes = 0, totalBudget = this.settings.maxBytes): boolean {
+		let changed = false;
+		while (
+			this.items.length > 0 &&
+			(this.items.length > this.settings.maxImages ||
+				this.residentBytes() > this.settings.maxBytes ||
+				this.residentBytes() + externalBytes > totalBudget)
+		) {
+			this.items.shift();
+			changed = true;
+		}
+		return changed;
+	}
+
+	private assertRevision(expectedRevision: number): void {
+		this.assertOpen();
+		if (!Number.isSafeInteger(expectedRevision) || expectedRevision !== this.revision) {
+			throw new SentImageError(
+				`Sent-image revision mismatch; current revision is ${this.revision}.`,
+			);
+		}
+	}
+
+	private assertOpen(): void {
+		if (this.closed) throw new SentImageError("Sent-image store is closed.", 410);
+	}
+}
+
+function normalizeImage(
+	input: RetainedImageInput,
+	index = 0,
+): {
+	name: string;
+	mimeType: string;
+	bytes: Buffer;
+} {
+	if (!input || typeof input.data !== "string" || !validBase64(input.data)) {
+		throw new SentImageError(`Sent image ${index + 1} has invalid data.`, 400);
+	}
+	const bytes = Buffer.from(input.data, "base64");
+	if (bytes.length === 0) throw new SentImageError(`Sent image ${index + 1} is empty.`, 400);
+	const mimeType = input.mimeType;
+	if (typeof mimeType !== "string" || !/^image\/[a-z0-9.+-]{1,64}$/i.test(mimeType)) {
+		throw new SentImageError(`Sent image ${index + 1} has an invalid MIME type.`, 400);
+	}
+	const name = typeof input.name === "string" && input.name ? input.name.slice(0, 255) : "Image";
+	return { name, mimeType, bytes };
+}
+
+function normalizeMessageId(value: string): void {
+	if (typeof value !== "string" || !/^[A-Za-z0-9_-]{1,120}$/.test(value)) {
+		throw new SentImageError("Message association is invalid.", 400);
+	}
+}
+
+function validBase64(value: string): boolean {
+	return /^(?:[A-Za-z0-9+/]{4})*(?:[A-Za-z0-9+/]{2}==|[A-Za-z0-9+/]{3}=)?$/.test(value);
+}
+
+function sameIds(left: readonly string[], right: readonly string[]): boolean {
+	return left.length === right.length && left.every((id, index) => id === right[index]);
+}