diff --git a/extensions/image-paste/index.ts b/extensions/image-paste/index.ts index 214da422..6d33500e 100644 --- a/extensions/image-paste/index.ts +++ b/extensions/image-paste/index.ts @@ -16,6 +16,7 @@ import { removeEditorLayer, } from "../shared/editor-layers.ts"; +const MAX_UNDO_ATTACHMENTS = 128; const IMAGE_PLACEHOLDER = /\[Image #(\d+)\]/g; const PI_CLIPBOARD_IMAGE = /^pi-clipboard-[0-9a-f]{8}-[0-9a-f]{4}-[1-5][0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}\.(gif|jpe?g|png|webp)$/i; @@ -186,6 +187,7 @@ function expandWith(text: string, attachments: Iterable) { export class ImageAttachmentStore { private nextAttachmentId = 1; private readonly draft = new Map(); + private readonly removed = new Map(); get hasDraft() { return this.draft.size > 0; @@ -202,7 +204,7 @@ export class ImageAttachmentStore { return placeholder; } - reconcileDraft(text: string) { + reconcileDraft(text: string, undo = false) { const occurrenceCounts = new Map(); for (const occurrence of placeholderOccurrences(text)) { occurrenceCounts.set( @@ -210,15 +212,26 @@ export class ImageAttachmentStore { (occurrenceCounts.get(occurrence.id) ?? 0) + 1, ); } - for (const id of [...this.draft.keys()]) { - // Placeholder text is user-editable, so only an unambiguous single - // occurrence keeps its mapping. Duplicated or removed tokens degrade to - // ordinary text; the underlying file is left alone either way. - if (occurrenceCounts.get(id) === 1) continue; + for (const [id, attachment] of this.removed) { + const count = occurrenceCounts.get(id) ?? 0; + if (count === 0) continue; + // Only native Undo can recover ownership. Typing a matching literal or + // duplicating a token revokes it, even if that edit is later undone. + this.removed.delete(id); + if (undo && count === 1) this.draft.set(id, attachment); + } + for (const [id, attachment] of this.draft) { + const count = occurrenceCounts.get(id) ?? 0; + if (count === 1) continue; this.draft.delete(id); + // Undoing an insertion discards it; there is no native redo to recover. + if (count === 0 && !undo) this.removed.set(id, attachment); + } + // Keep only a bounded recovery window, not another editor undo stack. + while (this.removed.size > MAX_UNDO_ATTACHMENTS) { + const oldest = this.removed.keys().next().value; + if (oldest !== undefined) this.removed.delete(oldest); } - this.nextAttachmentId = - this.draft.size === 0 ? 1 : Math.max(...this.draft.keys()) + 1; } /** @@ -244,6 +257,7 @@ export class ImageAttachmentStore { clearDraft() { this.draft.clear(); + this.removed.clear(); this.nextAttachmentId = 1; } @@ -264,6 +278,7 @@ export class ImageAttachmentStore { if (matches.length === 0) return undefined; this.clearDraft(); + const reserved = new Set(placeholderOccurrences(text).map(({ id }) => id)); // Build replacements in reverse document order so earlier indices stay // valid while we splice. Each occurrence — even of the same path — gets // its own placeholder identity so that expandWith's single-occurrence @@ -272,6 +287,7 @@ export class ImageAttachmentStore { []; for (const match of matches) { const path = match[0]; + while (reserved.has(this.nextAttachmentId)) this.nextAttachmentId += 1; const id = this.nextAttachmentId++; const placeholder = `[Image #${id}]`; this.draft.set(id, { id, placeholder, path } satisfies Attachment); @@ -338,6 +354,7 @@ export class ImageAttachmentEditor extends BelowEditorNavigationEditor { private downstreamSubmit?: (text: string) => void; private submittedDraft?: readonly Attachment[]; private settingText = false; + private restoringUndo = false; constructor( base: EditorComponent, @@ -367,12 +384,11 @@ export class ImageAttachmentEditor extends BelowEditorNavigationEditor { this.downstreamChange = value; super.onChange = (text) => { if (!this.settingText) { - // submitValue() clears the buffer and reports it here before invoking - // onSubmit, so an empty buffer ends the draft and restarts numbering. - // Keep a snapshot so the submit that caused it can still expand. + // Native deletion and submit both report an empty buffer. Retain + // recoverable mappings until onSubmit confirms the draft ended. if (text.length === 0) { this.submittedDraft = this.attachments.snapshotDraft(); - this.attachments.clearDraft(); + this.attachments.reconcileDraft(text, this.restoringUndo); } else { this.submittedDraft = undefined; // History recall uses setTextInternal (bypasses setText) and only @@ -389,7 +405,7 @@ export class ImageAttachmentEditor extends BelowEditorNavigationEditor { this.settingText = false; } } else { - this.attachments.reconcileDraft(text); + this.attachments.reconcileDraft(text, this.restoringUndo); } } } @@ -411,11 +427,11 @@ export class ImageAttachmentEditor extends BelowEditorNavigationEditor { // an already expanded string is a no-op. const snapshot = this.submittedDraft; this.submittedDraft = undefined; - value( - snapshot - ? expandWith(text, snapshot) - : this.attachments.expandPlaceholders(text), - ); + const expanded = snapshot + ? expandWith(text, snapshot) + : this.attachments.expandPlaceholders(text); + this.attachments.clearDraft(); + value(expanded); } : undefined; } @@ -435,6 +451,7 @@ export class ImageAttachmentEditor extends BelowEditorNavigationEditor { // Pi restores queued messages and history entries as expanded paths. Adopt // them so the buffer shows compact tokens and the next submit can expand // again; ordinary text falls through to plain reconciliation. + this.submittedDraft = undefined; const adopted = text.length > 0 ? this.attachments.adoptExpandedText(text) : undefined; const next = adopted ?? text; @@ -498,7 +515,15 @@ export class ImageAttachmentEditor extends BelowEditorNavigationEditor { ) { return; } - super.handleInput(data); + this.restoringUndo = this.editorKeybindings.matches( + data, + "tui.editor.undo", + ); + try { + super.handleInput(data); + } finally { + this.restoringUndo = false; + } } } @@ -531,6 +556,7 @@ export default function imagePaste( ); pi.on("session_start", (_event, ctx) => { + attachments.cleanup(); installImagePasteEditor(pi, ctx, attachments); }); diff --git a/tests/extensions/image-paste/index.test.ts b/tests/extensions/image-paste/index.test.ts index 8c509c8a..6b9ccdf3 100644 --- a/tests/extensions/image-paste/index.test.ts +++ b/tests/extensions/image-paste/index.test.ts @@ -222,7 +222,7 @@ test("duplicating a placeholder makes both copies ordinary text", () => { assert.equal(existsSync(path), true); }); -test("deleting the last image makes its number available to the next paste", () => { +test("deleted image numbers stay reserved for native undo", () => { const first = temporaryImage("png", "first"); const second = temporaryImage("png", "second"); const third = temporaryImage("png", "third"); @@ -239,7 +239,7 @@ test("deleting the last image makes its number available to the next paste", () base.cursor = editor.getText().length; editor.insertTextAtCursor(third); - assert.equal(editor.getText(), "[Image #1] [Image #2]"); + assert.equal(editor.getText(), "[Image #1] [Image #3]"); base.submit(); assert.equal(submitted[0], `${first} ${third}`); diff --git a/tests/extensions/image-paste/native-editor.test.ts b/tests/extensions/image-paste/native-editor.test.ts new file mode 100644 index 00000000..27ef47e0 --- /dev/null +++ b/tests/extensions/image-paste/native-editor.test.ts @@ -0,0 +1,188 @@ +import assert from "node:assert/strict"; +import { randomUUID } from "node:crypto"; +import { rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import test from "node:test"; +import { + InteractiveMode, + type KeybindingsManager, +} from "@earendil-works/pi-coding-agent"; +import { + Editor, + getKeybindings, + ProcessTerminal, + TuiMainScreen, +} from "@earendil-works/pi-tui"; +import { + ImageAttachmentEditor, + ImageAttachmentStore, +} from "../../../extensions/image-paste/index.ts"; + +const files: string[] = []; +test.after(() => { + for (const path of files) rmSync(path, { force: true }); +}); +function image() { + const path = join(tmpdir(), `pi-clipboard-${randomUUID()}.png`); + writeFileSync( + path, + Buffer.from( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+a7X8AAAAASUVORK5CYII=", + "base64", + ), + ); + files.push(path); + return path; +} +function harness() { + const identity = (text: string) => text; + const base = new Editor(new TuiMainScreen(new ProcessTerminal()), { + borderColor: identity, + selectList: { + selectedPrefix: identity, + selectedText: identity, + description: identity, + scrollInfo: identity, + noMatch: identity, + }, + }); + const store = new ImageAttachmentStore(); + const editor = new ImageAttachmentEditor( + base, + getKeybindings() as unknown as KeybindingsManager, + store, + ); + const submitted: string[] = []; + editor.onSubmit = (text) => submitted.push(text); + return { editor, store, submitted }; +} +async function roundTrip(editor: ImageAttachmentEditor) { + const queue: string[] = []; + const host = { + editor, + ui: { requestRender() {} }, + session: { + isCompacting: false, + isStreaming: true, + async prompt(text: string) { + queue.push(text); + }, + }, + updatePendingMessagesDisplay() {}, + clearAllQueues() { + return { steering: [], followUp: queue.splice(0) }; + }, + }; + // Exercise the locked native private hooks without constructing a provider session. + const hooks = InteractiveMode.prototype as unknown as { + handleFollowUp(this: typeof host): Promise; + restoreQueuedMessagesToEditor(this: typeof host): number; + }; + await hooks.handleFollowUp.call(host); + const queued = queue[0]; + assert.equal(editor.getText(), ""); + assert.equal(hooks.restoreQueuedMessagesToEditor.call(host), 1); + return queued; +} + +test("native queue restoration preserves literal tokens and repeated image paths", async () => { + const { editor, submitted } = harness(); + const path = image(); + editor.handleInput("[Image #1] [Image #3] "); + editor.insertTextAtCursor(path); + editor.handleInput(" and "); + editor.insertTextAtCursor(path); + const queued = await roundTrip(editor); + assert.equal(queued, `[Image #1] [Image #3] ${path} and ${path}`); + assert.equal( + editor.getText(), + "[Image #1] [Image #3] [Image #2] and [Image #4]", + ); + editor.handleInput("\r"); + assert.equal(submitted[0], queued); +}); + +test("native Undo recovers a removed attachment without reusing its identity", () => { + const { editor, submitted } = harness(); + const first = image(), + second = image(); + editor.insertTextAtCursor(first); + editor.handleInput("\x15"); // Ctrl+U + assert.equal(editor.getText(), ""); + editor.insertTextAtCursor(second); + assert.equal(editor.getText(), "[Image #2]"); + editor.handleInput("\x1f"); // Undo second insertion + editor.handleInput("\x1f"); // Undo Ctrl+U + assert.equal(editor.getExpandedText(), first); + editor.handleInput("\r"); + assert.deepEqual(submitted, [first]); +}); + +test("typing a removed token remains literal even after native Undo", () => { + const { editor, submitted } = harness(); + editor.insertTextAtCursor(image()); + editor.handleInput("\x15"); + editor.handleInput("[Image #1]"); + assert.equal(editor.getExpandedText(), "[Image #1]"); + editor.handleInput("\x15"); + editor.handleInput("\x1f"); + editor.handleInput("\r"); + assert.deepEqual(submitted, ["[Image #1]"]); +}); + +test("explicit clear and cleanup cannot resurrect attachments through native Undo", () => { + for (const reset of ["clear", "cleanup"] as const) { + const { editor, store, submitted } = harness(); + editor.insertTextAtCursor(image()); + if (reset === "clear") editor.setText(""); + else { + editor.handleInput("\x15"); + store.cleanup(); + } + editor.handleInput("\x1f"); + editor.handleInput("\r"); + assert.deepEqual(submitted, ["[Image #1]"]); + } +}); + +test("submission releases removed mappings before the next native draft", () => { + const { editor, submitted } = harness(); + editor.insertTextAtCursor(image()); + editor.handleInput("\x15"); + editor.handleInput("ordinary"); + editor.handleInput("\r"); + editor.handleInput("[Image #1]"); + editor.handleInput("\r"); + assert.deepEqual(submitted, ["ordinary", "[Image #1]"]); +}); + +test("attachment recovery keeps a bounded window of removed identities", () => { + const store = new ImageAttachmentStore(); + const path = image(); + const first = store.attachClipboardPath(path, ""); + let last = first; + for (let i = 0; i < 256; i += 1) { + store.reconcileDraft(""); + last = store.attachClipboardPath(path, ""); + } + store.reconcileDraft(""); + assert(first && last); + store.reconcileDraft(`${first} ${last}`, true); + assert.equal( + store.expandPlaceholders(`${first} ${last}`), + `${first} ${path}`, + ); +}); + +test("undoing a new insertion cannot bind an older cleared token to that image", () => { + const { editor, submitted } = harness(); + editor.insertTextAtCursor(image()); + editor.handleInput("\x15"); + editor.setText(""); + editor.insertTextAtCursor(image()); + editor.handleInput("\x1f"); + editor.handleInput("\x1f"); + editor.handleInput("\r"); + assert.deepEqual(submitted, ["[Image #1]"]); +});