fix(web): cancelled queued-message edits release their staged uploads

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

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

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<void>((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 = () => {};

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();