@narumitw/pi-image-drop 0.18.0 → 0.20.1

package/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "@narumitw/pi-image-drop",
-	"version": "0.18.0",
+	"version": "0.20.1",
 	"description": "Pi extension for staging browser images and attaching them to the next message.",
 	"type": "module",
 	"license": "MIT",

package/src/batch.ts

@@ -30,6 +30,7 @@ interface BatchItem extends ItemReservation {
 	source?: Buffer;
 	processed?: ProcessedImage;
 	processedAutoResize?: boolean;
+	processingOwner?: "browser" | "runtime";
 	error?: string;
 }
 
@@ -142,6 +143,9 @@ export class BatchStore {
 		this.assertOpen();
 		if (this.reservation) throw new BatchError("Batch is frozen", "frozen");
 		const item = this.item(id);
+		if (item.status !== "uploading" && item.status !== "error") {
+			throw new BatchError("Image is already processing or ready", "not-ready");
+		}
 		if (source.byteLength !== item.size) {
 			throw new BatchError("Uploaded size does not match the reservation", "invalid");
 		}
@@ -150,6 +154,7 @@ export class BatchStore {
 		item.processedAutoResize = undefined;
 		item.error = undefined;
 		item.status = "processing";
+		item.processingOwner = "browser";
 		return this.bump();
 	}
 
@@ -173,6 +178,7 @@ export class BatchStore {
 		item.processed = cloneProcessed(processed);
 		item.processedAutoResize = autoResize;
 		item.status = "ready";
+		item.processingOwner = undefined;
 		item.error = undefined;
 		if (duplicates.length > 0) {
 			const laterIds = new Set(duplicates.map(({ candidate }) => candidate.id));
@@ -192,6 +198,7 @@ export class BatchStore {
 		item.error = sanitizeError(error);
 		item.processed = undefined;
 		item.processedAutoResize = undefined;
+		item.processingOwner = undefined;
 		return this.bump();
 	}
 
@@ -207,6 +214,26 @@ export class BatchStore {
 		return this.bump();
 	}
 
+	cancelInFlight(error: string): boolean {
+		this.assertOpen();
+		if (this.reservation) return false;
+		const inFlight = this.items.filter(
+			(item) =>
+				item.status === "uploading" ||
+				(item.status === "processing" && item.processingOwner === "browser"),
+		);
+		if (inFlight.length === 0) return false;
+		for (const item of inFlight) {
+			item.status = "error";
+			item.error = sanitizeError(error);
+			item.processed = undefined;
+			item.processedAutoResize = undefined;
+			item.processingOwner = undefined;
+		}
+		this.bump();
+		return true;
+	}
+
 	beginAutoResizeReprocessing(autoResize: boolean): Array<{ id: string; source: Buffer }> {
 		this.assertOpen();
 		if (this.reservation) throw new BatchError("Batch is frozen", "frozen");
@@ -225,6 +252,7 @@ export class BatchStore {
 			item.status = "processing";
 			item.processed = undefined;
 			item.processedAutoResize = undefined;
+			item.processingOwner = "runtime";
 			item.error = undefined;
 		}
 		this.bump();
@@ -239,6 +267,7 @@ export class BatchStore {
 			throw new BatchError("Image cannot be retried without uploaded source bytes", "not-ready");
 		}
 		item.status = "processing";
+		item.processingOwner = "browser";
 		item.error = undefined;
 		item.processedAutoResize = undefined;
 		this.bump();

package/src/runtime.ts

@@ -83,17 +83,21 @@ export class ImageDropRuntime {
 		this.pi.on("before_agent_start", async () => this.batch?.markPreflightStarted());
 		this.pi.on("message_start", async (event, ctx) => this.handleMessageStart(event, ctx));
 		(this.pi as LatestExtensionAPI).on("agent_settled", async (_event, ctx) => {
+			if (!ctx.isIdle() || ctx.hasPendingMessages()) return;
 			await this.recoverReservation(ctx, "Queued image message was not delivered; restored it.");
 		});
 	}
 
 	async start(ctx: ExtensionContext): Promise<void> {
-		this.generation += 1;
+		const generation = ++this.generation;
+		const previousBatch = this.batch;
 		this.closed = true;
 		this.sessionAbort.abort();
 		await this.releaseServer();
-		this.batch?.close();
+		previousBatch?.close();
+		if (generation !== this.generation) return;
 		const result = await this.dependencies.loadSettings();
+		if (generation !== this.generation) return;
 		this.settings = result.settings;
 		this.batch = new BatchStore(result.settings);
 		this.processor = this.dependencies.createProcessor();
@@ -109,11 +113,13 @@ export class ImageDropRuntime {
 	}
 
 	async shutdown(ctx: ExtensionContext): Promise<void> {
-		this.generation += 1;
+		const generation = ++this.generation;
+		const previousBatch = this.batch;
 		this.closed = true;
 		this.sessionAbort.abort();
 		await this.releaseServer();
-		this.batch?.close();
+		previousBatch?.close();
+		if (generation !== this.generation) return;
 		this.batch = undefined;
 		this.settings = undefined;
 		this.processor = undefined;
@@ -150,7 +156,12 @@ export class ImageDropRuntime {
 			ctx.ui.notify("The current model does not support image input.", "warning");
 			return { action: "handled" };
 		}
+		const generation = this.generation;
+		const batch = this.batch;
 		const piSettings = await this.dependencies.readPiSettings(ctx.cwd, ctx.isProjectTrusted());
+		if (generation !== this.generation || batch !== this.batch || this.closed) {
+			return { action: "handled" };
+		}
 		this.notifyPiSettingsWarnings(ctx, piSettings.warnings);
 		if (piSettings.blockImages) {
 			this.restoreEditor(ctx, event.text);
@@ -162,7 +173,7 @@ export class ImageDropRuntime {
 		}
 
 		try {
-			const reservation = this.batch.reserveMessage(event.text, event.streamingBehavior);
+			const reservation = batch.reserveMessage(event.text, event.streamingBehavior);
 			this.server?.broadcastState();
 			this.updateWidget(ctx);
 			return {
@@ -230,19 +241,18 @@ export class ImageDropRuntime {
 						if (!(failure instanceof BatchError) || failure.code !== "not-found") throw failure;
 					}
 				} finally {
-					this.server?.broadcastState();
+					if (generation === this.generation) this.server?.broadcastState();
 				}
 			}),
 		);
-		if (generation !== this.generation || batch.publicState().phase !== "ready") {
+		if (generation !== this.generation) return false;
+		if (batch.publicState().phase !== "ready") {
 			this.restoreEditor(ctx, text);
-			if (generation === this.generation) {
-				ctx.ui.notify(
-					"Images could not be updated for the current auto-resize setting.",
-					"warning",
-				);
-				this.updateWidget(ctx);
-			}
+			ctx.ui.notify(
+				"Images could not be updated for the current auto-resize setting.",
+				"warning",
+			);
+			this.updateWidget(ctx);
 			return false;
 		}
 		this.updateWidget(ctx);
@@ -300,10 +310,10 @@ export class ImageDropRuntime {
 
 	private async releaseServer(): Promise<void> {
 		const server = this.server;
-		this.server = undefined;
-		if (server) await server.close();
 		const starting = this.serverStarting;
+		this.server = undefined;
 		this.serverStarting = undefined;
+		if (server) await server.close();
 		if (starting) {
 			try {
 				await (await starting).close();
@@ -318,9 +328,7 @@ export class ImageDropRuntime {
 		const reservation = this.batch?.currentReservation();
 		if (!reservation) return;
 		const images = userMessageImages(event);
-		if (images.length < reservation.images.length) return;
-		const attached = images.slice(-reservation.images.length);
-		if (digestImages(attached) !== reservation.digest) return;
+		if (!containsImageSequence(images, reservation.images.length, reservation.digest)) return;
 		this.batch?.commitReservation(reservation.digest);
 		this.server?.broadcastState();
 		this.updateWidget(ctx);
@@ -398,6 +406,13 @@ function userMessageImages(event: unknown): ImageContent[] {
 	return content.filter(isImageContent);
 }
 
+function containsImageSequence(images: ImageContent[], length: number, digest: string): boolean {
+	for (let start = 0; start + length <= images.length; start += 1) {
+		if (digestImages(images.slice(start, start + length)) === digest) return true;
+	}
+	return false;
+}
+
 function isImageContent(value: unknown): value is ImageContent {
 	return (
 		isRecord(value) &&

package/src/server.ts

@@ -28,6 +28,11 @@ interface SseClient {
 	response: ServerResponse;
 }
 
+interface ClientLease {
+	clientId: string;
+	generation: number;
+}
+
 class HttpError extends Error {
 	constructor(
 		readonly status: number,
@@ -47,6 +52,7 @@ export class ImageDropServer {
 	private readonly sockets = new Set<Socket>();
 	private readonly sseClients = new Set<SseClient>();
 	private activeClientId?: string;
+	private activeClientGeneration = 0;
 	private closed = false;
 	private closePromise?: Promise<void>;
 
@@ -183,13 +189,15 @@ export class ImageDropServer {
 					throw new HttpError(400, "Invalid client id");
 				}
 				this.takeLease(clientId);
+				this.broadcastState();
 				this.json(response, 200, this.statePayload());
 				return;
 			}
 			this.assertMutation(request);
-			this.assertActiveClient(request);
+			const lease = this.assertActiveClient(request);
 			if (request.method === "POST" && url.pathname === "/api/items") {
 				const body = await readJson(request, JSON_LIMIT);
+				this.assertLease(lease);
 				const revision = integerField(body, "revision");
 				const inputs = arrayField(body, "items").map((item) => {
 					if (!isRecord(item)) throw new HttpError(400, "Invalid item reservation");
@@ -210,6 +218,7 @@ export class ImageDropServer {
 				try {
 					source = await readBody(request, this.options.settings.maxImageBytes);
 				} catch (error) {
+					this.assertLease(lease);
 					try {
 						this.options.batch.failUpload(id, `Upload failed: ${formatError(error)}`);
 					} catch (failure) {
@@ -218,9 +227,10 @@ export class ImageDropServer {
 					this.broadcastState();
 					throw error;
 				}
+				this.assertLease(lease);
 				this.options.batch.startProcessing(id, source);
 				this.broadcastState();
-				const completed = await this.processItem(id, source, request);
+				const completed = await this.processItem(id, source, request, lease);
 				this.broadcastState();
 				this.json(response, 200, { ...this.statePayload(), duplicateOf: completed.duplicateOf });
 				return;
@@ -228,6 +238,7 @@ export class ImageDropServer {
 			const failureMatch = /^\/api\/items\/([A-Za-z0-9_-]{1,80})\/fail$/.exec(url.pathname);
 			if (request.method === "POST" && failureMatch) {
 				const body = await readJson(request, JSON_LIMIT);
+				this.assertLease(lease);
 				this.options.batch.failUpload(failureMatch[1] as string, stringField(body, "error"));
 				this.respondWithState(response);
 				return;
@@ -237,7 +248,7 @@ export class ImageDropServer {
 				const id = retryMatch[1] as string;
 				const source = this.options.batch.retrySource(id);
 				this.broadcastState();
-				const completed = await this.processItem(id, source, request);
+				const completed = await this.processItem(id, source, request, lease);
 				this.broadcastState();
 				this.json(response, 200, { ...this.statePayload(), duplicateOf: completed.duplicateOf });
 				return;
@@ -250,6 +261,7 @@ export class ImageDropServer {
 			}
 			if (request.method === "PUT" && url.pathname === "/api/order") {
 				const body = await readJson(request, JSON_LIMIT);
+				this.assertLease(lease);
 				const ids = arrayField(body, "ids").map((id) => {
 					if (typeof id !== "string") throw new HttpError(400, "Invalid image order");
 					return id;
@@ -260,6 +272,7 @@ export class ImageDropServer {
 			}
 			if (request.method === "POST" && url.pathname === "/api/clear") {
 				const body = await readJson(request, JSON_LIMIT);
+				this.assertLease(lease);
 				this.options.batch.clear(integerField(body, "revision"));
 				this.respondWithState(response);
 				return;
@@ -295,17 +308,31 @@ export class ImageDropServer {
 		if (request.headers.origin !== this.origin) throw new HttpError(403, "Unexpected Origin");
 	}
 
-	private assertActiveClient(request: IncomingMessage): void {
-		const client = request.headers["x-image-drop-client"];
-		if (typeof client !== "string" || client !== this.activeClientId) {
+	private assertActiveClient(request: IncomingMessage): ClientLease {
+		const clientId = request.headers["x-image-drop-client"];
+		if (typeof clientId !== "string" || clientId !== this.activeClientId) {
+			throw new HttpError(409, "This page no longer owns the editing lease");
+		}
+		return { clientId, generation: this.activeClientGeneration };
+	}
+
+	private assertLease(lease: ClientLease): void {
+		if (
+			lease.clientId !== this.activeClientId ||
+			lease.generation !== this.activeClientGeneration
+		) {
 			throw new HttpError(409, "This page no longer owns the editing lease");
 		}
 	}
 
 	private takeLease(clientId: string): void {
 		const previous = this.activeClientId;
+		if (previous !== clientId) this.activeClientGeneration += 1;
 		this.activeClientId = clientId;
 		if (previous && previous !== clientId) {
+			this.options.batch.cancelInFlight(
+				"Upload cancelled because the Image Drop page was replaced.",
+			);
 			for (const client of [...this.sseClients]) {
 				if (client.id !== previous) continue;
 				writeSse(client.response, "stale", { message: "Image Drop opened in another tab" });
@@ -335,21 +362,26 @@ export class ImageDropServer {
 		id: string,
 		source: Buffer,
 		request: IncomingMessage,
+		lease: ClientLease,
 	): Promise<{ duplicateOf?: string }> {
 		const requestAbort = new AbortController();
 		request.once("aborted", () => requestAbort.abort());
 		const signal = AbortSignal.any([this.abortController.signal, requestAbort.signal]);
 		try {
 			const autoResize = await this.options.getAutoResize();
+			this.assertLease(lease);
 			const processed = await this.options.process(source, {
 				autoResize,
 				maxImagePixels: this.options.settings.maxImagePixels,
 				signal,
 			});
+			this.assertLease(lease);
 			const completion = this.options.batch.complete(id, processed, autoResize);
 			return completion.kind === "duplicate" ? { duplicateOf: completion.existingId } : {};
 		} catch (error) {
+			if (error instanceof HttpError) throw error;
 			if (this.closed || signal.aborted || isDiscardedItemError(error)) return {};
+			this.assertLease(lease);
 			if (!this.failItem(id, formatError(error))) return {};
 			throw new HttpError(422, formatError(error));
 		} finally {

package/src/web/app.js

@@ -1,4 +1,5 @@
 import {
+	attemptMutation,
 	canMutate,
 	formatBytes,
 	moveItem,
@@ -198,8 +199,9 @@ async function retry(id) {
 }
 
 async function remove(id) {
-	await mutate(`/api/items/${id}?revision=${state.batch.revision}`, { method: "DELETE" });
-	pendingFiles.delete(id);
+	if (await mutate(`/api/items/${id}?revision=${state.batch.revision}`, { method: "DELETE" })) {
+		pendingFiles.delete(id);
+	}
 }
 
 async function reorder(ids) {
@@ -208,18 +210,21 @@ async function reorder(ids) {
 }
 
 async function clearAll() {
-	await mutate("/api/clear", { method: "POST", json: { revision: state.batch.revision } });
-	pendingFiles.clear();
+	if (await mutate("/api/clear", { method: "POST", json: { revision: state.batch.revision } })) {
+		pendingFiles.clear();
+	}
 }
 
 async function mutate(path, options) {
-	try {
-		applyState(await request(path, options));
-		clearError();
-		render();
-	} catch (error) {
-		showError(errorMessage(error));
+	const result = await attemptMutation(() => request(path, options));
+	if (!result.ok) {
+		showError(errorMessage(result.error));
+		return false;
 	}
+	applyState(result.value);
+	clearError();
+	render();
+	return true;
 }
 
 function render() {

package/src/web/state.js

@@ -49,6 +49,14 @@ export function preferNewestState(current, next) {
 	return current;
 }
 
+export async function attemptMutation(operation) {
+	try {
+		return { ok: true, value: await operation() };
+	} catch (error) {
+		return { ok: false, error };
+	}
+}
+
 export function formatBytes(value) {
 	const bytes = Math.max(0, Number(value) || 0);
 	if (bytes < 1024) return `${bytes} B`;