diff --git a/apps/web/src/components/chat/queuedMessageEdit.test.ts b/apps/web/src/components/chat/queuedMessageEdit.test.ts index 0f839f414f74f..f67d1843c38be 100644 --- a/apps/web/src/components/chat/queuedMessageEdit.test.ts +++ b/apps/web/src/components/chat/queuedMessageEdit.test.ts @@ -8,6 +8,7 @@ import { type ComposerFileAttachment, type ComposerImageAttachment, } from "../../composerDraftStore"; +import { readAttachmentUpload, useAttachmentUploadStore } from "../../lib/attachmentUploadQueue"; import { environmentContextRecord, responseContextRecord } from "../../lib/composerContextRecords"; import { prepareQueuedEditAttachments, recoverQueuedMessageEdit } from "./queuedMessageEdit"; @@ -166,10 +167,16 @@ describe("queued message file edits", () => { const store = useComposerDraftStore.getState(); store.setPrompt(threadTarget, "Separate draft"); store.addFiles(editTarget, [file]); + useAttachmentUploadStore.setState({ + uploadsByImageId: { + [file.id]: { status: "ready", environmentId, attachmentId: uploadedFile.id }, + }, + }); expect(recoverQueuedMessageEdit({ editTarget, threadTarget, originalText: "" })).toBe( "discarded", ); expect(store.getComposerDraft(threadTarget)?.prompt).toBe("Separate draft"); expect(store.getComposerDraft(editTarget)).toBeNull(); + expect(readAttachmentUpload(file.id)).toBeUndefined(); }); }); diff --git a/apps/web/src/components/chat/queuedMessageEdit.ts b/apps/web/src/components/chat/queuedMessageEdit.ts index b5e576a36bff7..48f7ee5d1ae10 100644 --- a/apps/web/src/components/chat/queuedMessageEdit.ts +++ b/apps/web/src/components/chat/queuedMessageEdit.ts @@ -7,6 +7,20 @@ import { type ComposerImageAttachment, type ComposerThreadTarget, } from "../../composerDraftStore"; +import { releaseDraftAttachments } from "../../lib/attachmentUploadQueue"; + +/** Staged attachments upload as soon as they are added, so a dropped edit must release them. */ +export function discardQueuedEditDraft(target: ComposerThreadTarget): void { + const store = useComposerDraftStore.getState(); + const draft = store.getComposerDraft(target); + if (draft) { + releaseDraftAttachments([...draft.images, ...draft.files]); + for (const image of draft.images) { + if (image.previewUrl.startsWith("blob:")) URL.revokeObjectURL(image.previewUrl); + } + } + store.clearComposerContent(target); +} export function recoverQueuedMessageEdit(input: { readonly editTarget: ComposerThreadTarget; @@ -28,8 +42,10 @@ export function recoverQueuedMessageEdit(input: { }); if (recovery === "kept") { store.moveComposerContent(input.editTarget, input.threadTarget); + store.clearComposerContent(input.editTarget); + } else { + discardQueuedEditDraft(input.editTarget); } - store.clearComposerContent(input.editTarget); return recovery; } diff --git a/apps/web/src/components/chat/useQueuedRunEdit.browser.test.tsx b/apps/web/src/components/chat/useQueuedRunEdit.browser.test.tsx index 3118cd3ec7100..28785ecd768af 100644 --- a/apps/web/src/components/chat/useQueuedRunEdit.browser.test.tsx +++ b/apps/web/src/components/chat/useQueuedRunEdit.browser.test.tsx @@ -28,6 +28,7 @@ import { deriveComposerSendState } from "../ChatView.logic"; import { ComposerPrimaryActions } from "./ComposerPrimaryActions"; import { AppAtomRegistryProvider, appAtomRegistry } from "../../rpc/atomRegistry"; import { useQueuedRunEdit } from "./useQueuedRunEdit"; +import { readAttachmentUpload, useAttachmentUploadStore } from "../../lib/attachmentUploadQueue"; interface RunStatus { readonly id: RunId; @@ -150,6 +151,7 @@ function view( function setup() { mocks.editQueuedRun.mockReset(); + useAttachmentUploadStore.setState({ uploadsByImageId: {} }); useComposerDraftStore.setState({ draftsByThreadKey: {}, draftThreadsByThreadKey: {} }); setRuns(THREAD_A, "queued"); setRuns(THREAD_B, "queued"); @@ -406,6 +408,78 @@ describe("useQueuedRunEdit", () => { expect(promptOf(THREAD_A)).toBe(""); }); + it.each([ + { transition: "cancel", result: "Success" }, + { transition: "recover", result: "Failure" }, + { transition: "replace", result: "Success" }, + ] as const)( + "settles saved attachments after $transition and $result", + async ({ transition, result }) => { + const { hook } = setup(); + let finishSave!: (result: { readonly _tag: "Success" | "Failure" }) => void; + let started!: () => void; + const saving = new Promise<{ readonly _tag: "Success" | "Failure" }>((resolve) => { + finishSave = resolve; + }); + const requested = new Promise((resolve) => { + started = resolve; + }); + mocks.editQueuedRun.mockImplementation(() => { + started(); + return saving; + }); + await using _screen = await render(view(THREAD_A, "queued", hook)); + await beginChangedEdit(hook); + const image = { + type: "image" as const, + id: "saved-image", + name: "image.png", + mimeType: "image/png", + sizeBytes: 5, + file: new File(["image"], "image.png", { type: "image/png" }), + previewUrl: "blob:saved-image", + }; + using revoke = vi.spyOn(URL, "revokeObjectURL").mockImplementation(() => {}); + await act(() => + useComposerDraftStore.getState().addImage(requireHook(hook).draftTarget, image), + ); + useAttachmentUploadStore.setState({ + uploadsByImageId: { + [image.id]: { status: "ready", environmentId: ENV, attachmentId: "saved-upload" }, + }, + }); + const target = requireHook(hook).draftTarget; + const saved = requireHook(hook).saveEdit( + { ...savedContent("changed"), images: [image] }, + () => {}, + ); + await act(() => requested); + expect(requireHook(hook).isSaving).toBe(true); + if (transition === "cancel") await act(() => requireHook(hook).cancelEdit()); + if (transition === "recover") await act(() => setRuns(THREAD_A, "running")); + if (transition === "replace") { + await beginChangedEdit(hook); + await act(() => + useComposerDraftStore.getState().addImage(requireHook(hook).draftTarget, image), + ); + } + expect(readAttachmentUpload(image.id)).toBeDefined(); + expect(revoke).not.toHaveBeenCalled(); + finishSave({ _tag: result }); + await act(() => saved); + if (transition === "cancel") { + expect(readAttachmentUpload(image.id)).toBeUndefined(); + expect(revoke).toHaveBeenCalledWith(image.previewUrl); + expect(useComposerDraftStore.getState().getComposerDraft(target)).toBeNull(); + } else { + expect(readAttachmentUpload(image.id)).toBeDefined(); + expect(revoke).not.toHaveBeenCalled(); + const owner = transition === "recover" ? THREAD_A : target; + expect(useComposerDraftStore.getState().getComposerDraft(owner)?.images).toEqual([image]); + } + }, + ); + it("leaves a newer edit alone when a save from before navigation finishes", async () => { const { hook } = setup(); let finishSave: (result: { readonly _tag: "Success" }) => void = () => {}; diff --git a/apps/web/src/components/chat/useQueuedRunEdit.ts b/apps/web/src/components/chat/useQueuedRunEdit.ts index 0ed3f9228940f..e7f321bec9dc2 100644 --- a/apps/web/src/components/chat/useQueuedRunEdit.ts +++ b/apps/web/src/components/chat/useQueuedRunEdit.ts @@ -60,7 +60,11 @@ import { useAtomCommand } from "../../state/use-atom-command"; import { stackedThreadToast, toastManager } from "../ui/toast"; import type { ChatComposerHandle } from "./ChatComposer"; import { fileAttachmentCapabilityBlockReason } from "./composerAttachmentFiles"; -import { prepareQueuedEditAttachments, recoverQueuedMessageEdit } from "./queuedMessageEdit"; +import { + discardQueuedEditDraft, + prepareQueuedEditAttachments, + recoverQueuedMessageEdit, +} from "./queuedMessageEdit"; import type { EditQueuedRunRequest } from "./QueuedRunsControl"; interface EditingQueuedRun { @@ -232,10 +236,12 @@ export function useQueuedRunEdit(input: { const focusComposerSoon = () => { window.requestAnimationFrame(() => composerRef.current?.focusAtEnd()); }; - const clearEditDraft = (runId: RunId) => - useComposerDraftStore - .getState() - .clearComposerContent(queuedRunEditDraftId(routeThreadRef, runId)); + const clearEditDraft = (runId: RunId) => { + const target = queuedRunEditDraftId(routeThreadRef, runId); + // A save in flight still sends these uploads. + if (isSaving) useComposerDraftStore.getState().clearComposerContent(target); + else discardQueuedEditDraft(target); + }; const beginEdit = useStableCallback((request: EditQueuedRunRequest) => { if (activeThreadId === null) return; @@ -366,10 +372,29 @@ export function useQueuedRunEdit(input: { }); setThreadError(editing.threadId, error); const current = finishSave(error === null); + if (!current) { + // A failed save can recover into the thread draft; a newer edit can also retain resources. + const drafts = Object.values(useComposerDraftStore.getState().draftsByThreadKey); + const retainedIds = new Set( + drafts.flatMap((draft) => + [...draft.images, ...draft.files].map((attachment) => attachment.id), + ), + ); + releaseDraftAttachments( + newAttachments.filter((attachment) => !retainedIds.has(attachment.id)), + ); + const retainedPreviews = new Set( + drafts.flatMap((draft) => draft.images.map((image) => image.previewUrl)), + ); + for (const image of content.images) { + if (image.previewUrl.startsWith("blob:") && !retainedPreviews.has(image.previewUrl)) { + URL.revokeObjectURL(image.previewUrl); + } + } + return; + } if (error !== null) return; releaseDraftAttachments(newAttachments); - // The edit was cancelled or begun again while this save ran; its drafts are the newer edit's. - if (!current) return; promptRef.current = ""; clearEditDraft(editing.runId); composerRef.current?.resetCursorState();