@davideasden/pi-undo 0.1.2 → 0.2.0

package/extensions/pi-undo.ts

@@ -27,13 +27,32 @@ export type PiUndoRuntimeFactory = (
 	pi: ExtensionAPI,
 ) => Promise<PiUndoRuntime>;
 
+type DeferredImage = NonNullable<InputEvent["images"]>[number];
+
+interface DeferredPrompt {
+	readonly text: string;
+	readonly images?: readonly DeferredImage[];
+}
+
 export function createPiUndoExtension(runtimeFactory: PiUndoRuntimeFactory): (pi: ExtensionAPI) => void {
 	return (pi) => {
 		let runtime: PiUndoRuntime | undefined;
+		let runtimeContext: ExtensionContext | undefined;
 		let generation = 0;
+		let deferredPrompts: DeferredPrompt[] = [];
+		let replaying: DeferredPrompt | undefined;
+		let acceptedReplay: DeferredPrompt | undefined;
+		let activeCommands = new Set<symbol>();
+		let activeAction: "undo" | "redo" | undefined;
 
 		const initialize = async (context: ExtensionContext): Promise<void> => {
 			const currentGeneration = ++generation;
+			runtimeContext = context;
+			deferredPrompts = [];
+			replaying = undefined;
+			acceptedReplay = undefined;
+			activeCommands = new Set<symbol>();
+			activeAction = undefined;
 			try {
 				const next = await runtimeFactory(context, pi);
 				if (currentGeneration !== generation) return;
@@ -49,6 +68,57 @@ export function createPiUndoExtension(runtimeFactory: PiUndoRuntimeFactory): (pi
 			}
 		};
 
+		const dispatchDeferredPrompt = (active: PiUndoRuntime, expectedGeneration: number): void => {
+			if (
+				expectedGeneration !== generation || runtime !== active || activeCommands.size > 0 ||
+				replaying !== undefined || deferredPrompts.length === 0 || active.controller.history().locked
+			) return;
+			const prompt = deferredPrompts[0]!;
+			replaying = prompt;
+			acceptedReplay = undefined;
+			queueMicrotask(() => {
+				if (expectedGeneration !== generation || runtime !== active || replaying !== prompt) return;
+				try {
+					pi.sendUserMessage(prompt.images === undefined || prompt.images.length === 0
+						? prompt.text
+						: [{ type: "text" as const, text: prompt.text }, ...prompt.images]);
+				} catch (error) {
+					replaying = undefined;
+					acceptedReplay = undefined;
+					deferredPrompts.shift();
+					restoreEditorText(runtimeContext, prompt.text);
+					runtimeContext?.ui.notify(`Unable to replay queued prompt: ${errorMessage(error)}`, "warning");
+				}
+			});
+		};
+
+		const restoreDeferredPrompts = (context: ExtensionContext | undefined): void => {
+			if (deferredPrompts.length === 0) return;
+			const prompts = deferredPrompts.splice(0);
+			replaying = undefined;
+			acceptedReplay = undefined;
+			restoreEditorText(context, prompts.map((prompt) => prompt.text).join("\n\n"));
+			if (prompts.some((prompt) => (prompt.images?.length ?? 0) > 0)) {
+				context?.ui.notify("Queued prompt text restored; image attachments must be reattached", "warning");
+			}
+		};
+
+		const resumeDeferredPrompts = (
+			active: PiUndoRuntime,
+			expectedGeneration: number,
+			lockedReason: string,
+		): void => {
+			if (expectedGeneration !== generation || runtime !== active) return;
+			const history = active.controller.history();
+			if (history.locked) {
+				active.reporter.setRecoveryRequired(lockedReason);
+				restoreDeferredPrompts(runtimeContext);
+			} else {
+				active.reporter.setReady(history.undoCount, history.redoCount);
+				dispatchDeferredPrompt(active, expectedGeneration);
+			}
+		};
+
 		const runCommand = async (
 			action: "undo" | "redo",
 			context: ExtensionCommandContext,
@@ -59,6 +129,10 @@ export function createPiUndoExtension(runtimeFactory: PiUndoRuntimeFactory): (pi
 				context.ui.notify("pi-undo session unavailable", "warning");
 				return;
 			}
+			const commandToken = Symbol(action);
+			const commandSet = activeCommands;
+			commandSet.add(commandToken);
+			activeAction = action;
 			active.reporter.setPhase(action === "undo" ? "undoing" : "redoing");
 			active.setCommandContext?.(context);
 			let result;
@@ -68,15 +142,17 @@ export function createPiUndoExtension(runtimeFactory: PiUndoRuntimeFactory): (pi
 					: await active.controller.redo();
 			} finally {
 				active.setCommandContext?.(undefined);
+				commandSet.delete(commandToken);
+				if (commandSet === activeCommands && commandSet.size === 0) activeAction = undefined;
 			}
 			if (commandGeneration !== generation || runtime !== active) return;
 			active.reporter.result(result);
-			if (action === "undo" && result.code === "ok" && result.refillPrompt !== undefined) {
-				active.reporter.refillPrompt(result.refillPrompt);
-			}
-			const history = active.controller.history();
-			if (history.locked) active.reporter.setRecoveryRequired(result.message ?? result.code);
-			else active.reporter.setReady(history.undoCount, history.redoCount);
+			const hasDeferredPrompt = deferredPrompts.length > 0 || replaying !== undefined;
+			if (
+				action === "undo" && result.code === "ok" && result.refillPrompt !== undefined &&
+				!hasDeferredPrompt
+			) active.reporter.refillPrompt(result.refillPrompt);
+			resumeDeferredPrompts(active, commandGeneration, result.message ?? result.code);
 		};
 
 		const runDiff = async (args: string, context: ExtensionCommandContext): Promise<void> => {
@@ -132,19 +208,69 @@ export function createPiUndoExtension(runtimeFactory: PiUndoRuntimeFactory): (pi
 		});
 
 		pi.on("session_start", async (_event: unknown, context: ExtensionContext) => initialize(context));
-		pi.on("input", async (event: InputEvent) => {
-			if (runtime === undefined) return { action: "handled" as const };
-			return runtime.controller.prepareInput(event.text, { streaming: event.streamingBehavior !== undefined });
+		pi.on("input", async (event: InputEvent, context: ExtensionContext) => {
+			const active = runtime;
+			if (active === undefined) return { action: "handled" as const };
+			const result = await active.controller.prepareInput(event.text, {
+				streaming: event.streamingBehavior !== undefined,
+			});
+			const replay = replaying;
+			if (result.action === "defer") {
+				if (replay !== undefined && event.source === "extension" && samePrompt(event.text, event.images, replay)) {
+					replaying = undefined;
+					acceptedReplay = undefined;
+					active.reporter.setPhase(`${activeAction ?? "operation"} queued:${deferredPrompts.length}`);
+					return { action: "handled" as const };
+				}
+				if (event.text.trimStart().startsWith("/")) {
+					restoreEditorText(context, event.text);
+					context.ui.notify("Command input preserved until undo/redo completes", "info");
+					return { action: "handled" as const };
+				}
+				deferredPrompts.push({
+					text: event.text,
+					...(event.images === undefined ? {} : { images: event.images.map((image) => ({ ...image })) }),
+				});
+				active.reporter.setPhase(`${activeAction ?? "operation"} queued:${deferredPrompts.length}`);
+				return { action: "handled" as const };
+			}
+			if (replay !== undefined && event.source === "extension" && samePrompt(event.text, event.images, replay)) {
+				if (result.action === "continue") {
+					acceptedReplay = replay;
+				} else {
+					removeDeferredPrompt(deferredPrompts, replay);
+					replaying = undefined;
+					acceptedReplay = undefined;
+					restoreEditorText(context, replay.text);
+					if ((replay.images?.length ?? 0) > 0) {
+						context.ui.notify("Queued prompt text restored; image attachments must be reattached", "warning");
+					}
+				}
+			}
+			return result;
+		});
+		pi.on("before_agent_start", async () => {
+			const active = runtime;
+			const startGeneration = generation;
+			if (active === undefined) return;
+			await active.controller.beforeAgentStart();
+			const replay = replaying;
+			if (
+				runtime === active && generation === startGeneration && replay !== undefined &&
+				acceptedReplay === replay
+			) {
+				removeDeferredPrompt(deferredPrompts, replay);
+				replaying = undefined;
+				acceptedReplay = undefined;
+			}
 		});
-		pi.on("before_agent_start", async () => { await runtime?.controller.beforeAgentStart(); });
 		pi.on("agent_settled", async () => {
 			const active = runtime;
+			const settledGeneration = generation;
 			if (active === undefined) return;
 			await active.controller.agentSettled();
-			if (runtime !== active) return;
-			const history = active.controller.history();
-			if (history.locked) active.reporter.setRecoveryRequired("session state ambiguous");
-			else active.reporter.setReady(history.undoCount, history.redoCount);
+			if (runtime !== active || generation !== settledGeneration) return;
+			resumeDeferredPrompts(active, settledGeneration, "session state ambiguous");
 		});
 		pi.on("session_before_tree", async (event: PiSessionBeforeTreeEvent) => {
 			if (runtime === undefined) return { cancel: true };
@@ -152,7 +278,7 @@ export function createPiUndoExtension(runtimeFactory: PiUndoRuntimeFactory): (pi
 			const active = runtime;
 			const result = await active.controller.beforeTree({ targetLeafId: event.preparation.targetId });
 			if (result === undefined) {
-				 event.signal?.addEventListener("abort", () => { void active.controller.cancelTree?.(); }, { once: true });
+				event.signal?.addEventListener("abort", () => { void active.controller.cancelTree?.(); }, { once: true });
 			}
 			return result;
 		});
@@ -165,6 +291,12 @@ export function createPiUndoExtension(runtimeFactory: PiUndoRuntimeFactory): (pi
 		});
 		pi.on("session_shutdown", async () => {
 			generation += 1;
+			deferredPrompts = [];
+			replaying = undefined;
+			acceptedReplay = undefined;
+			activeCommands = new Set<symbol>();
+			activeAction = undefined;
+			runtimeContext = undefined;
 			await runtime?.controller.cancelTree?.();
 			runtime?.reporter.clear();
 			runtime = undefined;
@@ -172,6 +304,27 @@ export function createPiUndoExtension(runtimeFactory: PiUndoRuntimeFactory): (pi
 	};
 }
 
+function samePrompt(
+	text: string,
+	images: readonly DeferredImage[] | undefined,
+	prompt: DeferredPrompt,
+): boolean {
+	if (text !== prompt.text || (images?.length ?? 0) !== (prompt.images?.length ?? 0)) return false;
+	return (images ?? []).every((image, index) => JSON.stringify(image) === JSON.stringify(prompt.images?.[index]));
+}
+
+function removeDeferredPrompt(prompts: DeferredPrompt[], prompt: DeferredPrompt): void {
+	const index = prompts.indexOf(prompt);
+	if (index >= 0) prompts.splice(index, 1);
+}
+
+function restoreEditorText(context: ExtensionContext | undefined, text: string): void {
+	if (context === undefined || text.length === 0) return;
+	const current = context.ui.getEditorText();
+	if (current === text || current.startsWith(`${text}\n\n`)) return;
+	context.ui.setEditorText(current.length === 0 ? text : `${text}\n\n${current}`);
+}
+
 function errorMessage(error: unknown): string {
 	return error instanceof Error ? error.message : String(error);
 }

package/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "@davideasden/pi-undo",
-	"version": "0.1.2",
+	"version": "0.2.0",
 	"description": "Persistent workspace undo and redo for Pi",
 	"type": "module",
 	"keywords": [

package/src/controller.ts

@@ -80,9 +80,10 @@ export interface OperationResult {
 	readonly refillPrompt?: string;
 }
 
-export interface InputEventResult {
-	readonly action: "continue" | "handled";
-}
+export type InputEventResult =
+	| { readonly action: "continue" }
+	| { readonly action: "handled" }
+	| { readonly action: "defer" };
 
 export interface InputContext {
 	readonly streaming: boolean;
@@ -163,6 +164,7 @@ export class UndoControllerImpl implements UndoController {
 	private locked = false;
 	private historyPaused = false;
 	private operationInFlight = false;
+	private promptDeferralInFlight = false;
 	private lastSafetyManifestId: ManifestId | null = null;
 
 	constructor(dependencies: ControllerDependencies, initialState: ControllerInitialState = {}) {
@@ -182,6 +184,7 @@ export class UndoControllerImpl implements UndoController {
 	}
 
 	async prepareInput(text: string, context: InputContext): Promise<InputEventResult> {
+		if (this.promptDeferralInFlight) return { action: "defer" };
 		if (this.locked || this.operationInFlight) return { action: "handled" };
 		if (context.streaming || text.length === 0) return { action: "continue" };
 		try {
@@ -382,6 +385,7 @@ export class UndoControllerImpl implements UndoController {
 	): Promise<OperationResult> {
 		if (this.locked || this.operationInFlight) return { code: "busy", changedFiles: 0 };
 		this.operationInFlight = true;
+		this.promptDeferralInFlight = true;
 		this.lastSafetyManifestId = null;
 		let lease: { release(): Promise<void> } | undefined;
 		try {
@@ -454,6 +458,7 @@ export class UndoControllerImpl implements UndoController {
 			if (lease !== undefined) {
 				await lease.release().catch(() => { this.locked = true; });
 			}
+			this.promptDeferralInFlight = false;
 			this.operationInFlight = false;
 		}
 	}