fix(image-paste): preserve attachment identity across draft restoration

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<Attachment>) {
 export class ImageAttachmentStore {
   private nextAttachmentId = 1;
   private readonly draft = new Map<number, Attachment>();
+  private readonly removed = new Map<number, Attachment>();
 
   get hasDraft() {
     return this.draft.size > 0;
@@ -202,23 +204,34 @@ export class ImageAttachmentStore {
     return placeholder;
   }
 
-  reconcileDraft(text: string) {
+  reconcileDraft(text: string, undo = false) {
     const occurrenceCounts = new Map<number, number>();
     for (const occurrence of placeholderOccurrences(text)) {
       occurrenceCounts.set(
         occurrence.id,
         (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);
   });
 

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}`);

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<void>;
+    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]"]);
+});