diff --git a/ui/src/app/chat-attachment-handoff.test.ts b/ui/src/app/chat-attachment-handoff.test.ts index 2d82762f5c56..a32196ba454c 100644 --- a/ui/src/app/chat-attachment-handoff.test.ts +++ b/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(); diff --git a/ui/src/app/chat-attachment-handoff.ts b/ui/src/app/chat-attachment-handoff.ts index c8338308ebce..8c4a5e8bcf9e 100644 --- a/ui/src/app/chat-attachment-handoff.ts +++ b/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(), + ) => { + 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)) { diff --git a/ui/src/pages/chat/chat-pane-attachment-handoff.test.ts b/ui/src/pages/chat/chat-pane-attachment-handoff.test.ts index f6c6f68e6689..8e0ccc632523 100644 --- a/ui/src/pages/chat/chat-pane-attachment-handoff.test.ts +++ b/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); + }); }); diff --git a/ui/src/pages/chat/chat-pane-attachment-handoff.ts b/ui/src/pages/chat/chat-pane-attachment-handoff.ts index b76baef31479..672565b4bd98 100644 --- a/ui/src/pages/chat/chat-pane-attachment-handoff.ts +++ b/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(