fix: retain replacement attachment payload ownership

ui/src/app/chat-attachment-handoff.test.ts

@@ -106,6 +106,33 @@ describe("chat attachment route handoff", () => {
     }
   });
 
+  it("keeps payloads reused by a replacement prepare", () => {
+    const owner = {} as GatewayBrowserClient;
+    const retained = storedAttachment("replacement-retained", "image/png", false);
+    const removed = storedAttachment("replacement-removed", "image/png", false);
+    const handoff = createChatAttachmentHandoff();
+    handoff.prepare({
+      owner,
+      paneId: "p1",
+      scopeKey: "one",
+      attachments: [retained, removed],
+      fallbacks: {},
+    });
+    handoff.prepare({
+      owner,
+      paneId: "p1",
+      scopeKey: "one",
+      attachments: [retained],
+      fallbacks: {},
+    });
+
+    expect(getChatAttachmentDataUrl(retained)).not.toBeNull();
+    expect(getChatAttachmentDataUrl(removed)).toBeNull();
+    expect(handoff.consume({ owner, paneId: "p1", scopeKey: "one" })?.attachments).toEqual([
+      retained,
+    ]);
+  });
+
   it("bounds abandoned entries and releases pane-clear and application disposal", () => {
     const owner = {} as GatewayBrowserClient;
     const handoff = createChatAttachmentHandoff();

ui/src/app/chat-attachment-handoff.ts

@@ -19,17 +19,23 @@ export function createChatAttachmentHandoff(): ApplicationChatAttachmentHandoff
 
   const release = (attachments: readonly ChatAttachment[] = []) =>
     releaseChatAttachmentPayloads(attachments);
-  const releaseHandoff = (handoff: PendingChatAttachmentHandoff | undefined) => {
-    if (!handoff) {
-      return;
-    }
+  const handoffAttachments = (handoff: PendingChatAttachmentHandoff) => {
     const byId = new Map(handoff.attachments.map((attachment) => [attachment.id, attachment]));
     for (const fallback of Object.values(handoff.fallbacks)) {
       for (const attachment of fallback.attachments) {
         byId.set(attachment.id, attachment);
       }
     }
-    release([...byId.values()]);
+    return [...byId.values()];
+  };
+  const releaseHandoff = (
+    handoff: PendingChatAttachmentHandoff | undefined,
+    retainedIds = new Set<string>(),
+  ) => {
+    if (!handoff) {
+      return;
+    }
+    release(handoffAttachments(handoff).filter((attachment) => !retainedIds.has(attachment.id)));
   };
   const take = (paneId: string) => {
     const handoff = pending.get(paneId);
@@ -47,7 +53,13 @@ export function createChatAttachmentHandoff(): ApplicationChatAttachmentHandoff
         releaseHandoff(previous);
         return;
       }
-      releaseHandoff(previous);
+      const retainedIds = new Set(attachments.map((attachment) => attachment.id));
+      for (const fallback of Object.values(fallbacks)) {
+        for (const attachment of fallback.attachments) {
+          retainedIds.add(attachment.id);
+        }
+      }
+      releaseHandoff(previous, retainedIds);
       if (!owner || disposed) {
         release(attachments);
         for (const fallback of Object.values(fallbacks)) {

ui/src/pages/chat/chat-pane-attachment-handoff.test.ts

@@ -115,4 +115,42 @@ describe("staged chat attachment pane handoff", () => {
     );
     discardStateStagedAttachments(remount);
   });
+
+  it("releases a restored fallback displaced by mounted state", () => {
+    const owner = {} as GatewayBrowserClient;
+    const handoff = createChatAttachmentHandoff();
+    const context = { chatAttachmentHandoff: handoff } as unknown as ApplicationContext;
+    const displaced = storedAttachment("displaced");
+    const mounted = storedAttachment("mounted");
+    handoff.prepare({
+      owner,
+      paneId: "p1",
+      scopeKey: "active",
+      attachments: [],
+      fallbacks: {
+        collision: {
+          attachments: [displaced],
+          message: "old",
+          sequence: 1,
+          storageFailed: false,
+        },
+      },
+    });
+    const remount = state([]);
+    remount.chatComposerFallbackByScope = {
+      collision: {
+        attachments: [mounted],
+        message: "new",
+        sequence: 2,
+        storageFailed: false,
+      },
+    };
+
+    restorePaneStagedAttachments(context, "p1", remount, owner);
+
+    expect(remount.chatComposerFallbackByScope.collision?.attachments).toEqual([mounted]);
+    expect(getChatAttachmentDataUrl(displaced)).toBeNull();
+    expect(getChatAttachmentDataUrl(mounted)).not.toBeNull();
+    discardStateStagedAttachments(remount);
+  });
 });

ui/src/pages/chat/chat-pane-attachment-handoff.ts

@@ -44,10 +44,20 @@ export function restorePaneStagedAttachments(
     ...state.chatAttachments,
     ...restored.attachments.filter((attachment) => !currentIds.has(attachment.id)),
   ];
+  const displaced = Object.entries(restored.fallbacks)
+    .filter(([scopeKey]) => Object.hasOwn(state.chatComposerFallbackByScope, scopeKey))
+    .flatMap(([, fallback]) => fallback.attachments);
   state.chatComposerFallbackByScope = {
     ...restored.fallbacks,
     ...state.chatComposerFallbackByScope,
   };
+  const retainedIds = new Set(state.chatAttachments.map((attachment) => attachment.id));
+  for (const fallback of Object.values(state.chatComposerFallbackByScope)) {
+    for (const attachment of fallback.attachments) {
+      retainedIds.add(attachment.id);
+    }
+  }
+  releaseAttachments(displaced, retainedIds);
 }
 
 export function preparePaneStagedAttachments(