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 4a5672e44567..672054bf3342 100644 --- a/ui/src/pages/chat/chat-pane-attachment-handoff.test.ts +++ b/ui/src/pages/chat/chat-pane-attachment-handoff.test.ts @@ -15,6 +15,7 @@ import { closeStagedPane, discardStateStagedAttachments, preparePaneStagedAttachments, + replacePaneStagedAttachmentGatewayOwner, restorePaneStagedAttachments, } from "./chat-pane-attachment-handoff.ts"; import { createTestChatPane } from "./chat-pane.test-support.ts"; @@ -180,6 +181,37 @@ describe("staged chat attachment pane handoff", () => { expect(current.chatComposerFallbackByScope.fallback?.attachments).toEqual([]); }); + it("keeps plain staged attachments across a gateway client rotation", () => { + const previousOwner = {} as GatewayBrowserClient; + const nextOwner = {} as GatewayBrowserClient; + const handoff = createChatAttachmentHandoff(); + const context = { chatAttachmentHandoff: handoff } as unknown as ApplicationContext; + const plainImage = storedAttachment("rotation-image"); + const plainFile = storedAttachment("rotation-file", "application/pdf"); + const annotated: ChatAttachment = { + ...storedAttachment("rotation-annotation"), + browserAnnotation: { pageUrl: "https://example.test" } as never, + }; + const current = state([plainImage, plainFile, annotated]); + + const returned = replacePaneStagedAttachmentGatewayOwner( + context, + "p1", + current, + previousOwner, + nextOwner, + ); + + expect(returned).toBe(nextOwner); + // Plain payloads are client-local; rotation must not silently discard them. + expect(current.chatAttachments).toEqual([plainImage, plainFile]); + expect(getChatAttachmentDataUrl(plainImage)).not.toBeNull(); + expect(getChatAttachmentDataUrl(plainFile)).not.toBeNull(); + // Annotation Undo context dies with the old client; its payload is released. + expect(getChatAttachmentDataUrl(annotated)).toBeNull(); + discardStateStagedAttachments(current); + }); + it("restores a mixed package only to the exact mounted owner", () => { const owner = {} as GatewayBrowserClient; const otherOwner = {} as GatewayBrowserClient; diff --git a/ui/src/pages/chat/chat-pane-attachment-handoff.ts b/ui/src/pages/chat/chat-pane-attachment-handoff.ts index 698e878967f7..9a6842b58ea7 100644 --- a/ui/src/pages/chat/chat-pane-attachment-handoff.ts +++ b/ui/src/pages/chat/chat-pane-attachment-handoff.ts @@ -97,10 +97,21 @@ export function replacePaneStagedAttachmentGatewayOwner( if (!nextOwner || previousOwner === nextOwner) { return previousOwner; } - discardStateStagedAttachments(state); - state?.requestUpdate?.(); + // Rotating the client invalidates annotation Undo context owned by the old + // client, but plain file/image payloads are client-local data URLs — a gap + // reconnect or plugin-install rotation must not silently discard them. + if (state) { + const dropAnnotations = (attachments: readonly ChatAttachment[]) => { + releaseAttachments(attachments.filter((attachment) => attachment.browserAnnotation)); + return attachments.filter((attachment) => !attachment.browserAnnotation); + }; + state.chatAttachments = dropAnnotations(state.chatAttachments); + for (const fallback of Object.values(state.chatComposerFallbackByScope)) { + fallback.attachments = dropAnnotations(fallback.attachments); + } + state.requestUpdate?.(); + } context.chatAttachmentHandoff.clearPane(paneId); - // Rotating the token also invalidates any pending annotation Undo owned by the old client. return nextOwner; } diff --git a/ui/src/pages/chat/chat-pane-browser-annotation-lifecycle.test.ts b/ui/src/pages/chat/chat-pane-browser-annotation-lifecycle.test.ts index 5b04d7d18c51..10c1211a6bd7 100644 --- a/ui/src/pages/chat/chat-pane-browser-annotation-lifecycle.test.ts +++ b/ui/src/pages/chat/chat-pane-browser-annotation-lifecycle.test.ts @@ -403,11 +403,13 @@ describe("staged attachment composer adoption", () => { phase: "reconnecting", hello: null, }); - expect(state.chatAttachments).toEqual([]); - expect(state.chatComposerFallbackByScope.fallback?.attachments).toEqual([]); + // Rotation invalidates annotation-backed attachments only; plain payloads + // are client-local and survive (live state and fallbacks alike). + expect(state.chatAttachments).toEqual([ordinary]); + expect(state.chatComposerFallbackByScope.fallback?.attachments).toEqual([ordinary]); expect(getChatAttachmentDataUrl(current)).toBeNull(); expect(getChatAttachmentDataUrl(fallback)).toBeNull(); - expect(getChatAttachmentDataUrl(ordinary)).toBeNull(); + expect(getChatAttachmentDataUrl(ordinary)).not.toBeNull(); }); it("discards a staged package captured without a client when the first client arrives", () => { @@ -433,9 +435,9 @@ describe("staged attachment composer adoption", () => { hello: null, }); - expect(state.chatAttachments).toEqual([]); + expect(state.chatAttachments).toEqual([ordinary]); expect(getChatAttachmentDataUrl(annotation)).toBeNull(); - expect(getChatAttachmentDataUrl(ordinary)).toBeNull(); + expect(getChatAttachmentDataUrl(ordinary)).not.toBeNull(); }); it("invalidates annotation Undo when the logical Gateway client is replaced", async () => { @@ -474,9 +476,9 @@ describe("staged attachment composer adoption", () => { }); toastHost.querySelector(".app-toast__action")?.click(); - expect(state.chatAttachments).toEqual([]); + expect(state.chatAttachments).toEqual([ordinary]); expect(getChatAttachmentDataUrl(annotation)).toBeNull(); - expect(getChatAttachmentDataUrl(ordinary)).toBeNull(); + expect(getChatAttachmentDataUrl(ordinary)).not.toBeNull(); toastHost.remove(); }); });