🤖 fix: persist chat draft images (#1187)

Persist pasted image attachments in ChatInput drafts so they survive
workspace switches and app restarts.

Key points:
- Added `inputImages:<scopeId>` localStorage key and included it in
workspace storage copy/delete logic.
- ChatInput now restores persisted image attachments on mount and
persists updates on paste/drop/remove/send.
- Creation flow clears pending image drafts after successful workspace
creation.

Tests:
- `make static-check`

---

<details>
<summary>📋 Implementation Plan</summary>

# Persist pasted image attachments in ChatInput drafts

## Why images are currently lost
- Draft **text** is persisted via `usePersistedState(getInputKey(...))`
in `src/browser/components/ChatInput/index.tsx`.
- Draft **images** are held only in component memory
(`useState<ImageAttachment[]>([])`), so they reset whenever:
  - the workspace view remounts (workspace switch), or
  - the renderer reloads (app restart).

## Goal
Make pasted/drag-dropped image attachments persist the same way draft
text does:
- per-workspace (and per “creation” project scope)
- survives workspace switches and app restarts
- doesn’t regress the ability to attach large images (quota issues
should not break the UI)

## Recommended approach (minimal + safe)
Persist image attachments **best-effort** to localStorage under a new
key, while keeping the authoritative in-memory state in React.

Rationale:
- Matches existing draft persistence strategy (localStorage) without
introducing new backend APIs.
- Avoids a regression where attaching a large image fails to show in the
UI if localStorage quota is exceeded (which would happen if we used
`usePersistedState` as the source of truth).

## Implementation plan

### 1) Add a workspace-scoped storage key for draft images
**Files:**
- `src/common/constants/storage.ts`

Actions:
- Add a new helper:
- `getInputImagesKey(scopeId: string): string` →
`inputImages:${scopeId}`
- (Use the same `scopeId` inputs as `getInputKey`: workspaceId for
workspace variant; `getPendingScopeId(projectPath)` for creation
variant.)
- Add `getInputImagesKey` to `PERSISTENT_WORKSPACE_KEY_FUNCTIONS` so
images are:
  - copied on workspace fork (`copyWorkspaceStorage`)
  - deleted on workspace removal (`deleteWorkspaceStorage`)
  - migrated on workspace ID migration (`migrateWorkspaceStorage`)

### 2) Load persisted images when ChatInput mounts
**Files:**
- `src/browser/components/ChatInput/index.tsx`

Actions:
- Extend `storageKeys` to include `imagesKey`:
  - workspace: `getInputImagesKey(props.workspaceId)`
  - creation: `getInputImagesKey(getPendingScopeId(props.projectPath))`
- Initialize `imageAttachments` from localStorage:
- `useState(() =>
readPersistedState<ImageAttachment[]>(storageKeys.imagesKey, []))`
- Add a small runtime validator (defensive programming):
- if the persisted value isn’t an array of `{id,url,mediaType}` strings,
log and fall back to `[]` (and consider clearing the bad key).

### 3) Persist image changes back to localStorage (best-effort)
**Files:**
- `src/browser/components/ChatInput/index.tsx`

Actions:
- Keep `imageAttachments` as a normal `useState`.
- Wrap updates so they also persist:
- When setting images, call `updatePersistedState(storageKeys.imagesKey,
nextImages.length ? nextImages : undefined)`.
  - Ensure all mutation paths persist:
    - paste (`handlePaste`)
    - drop (`handleDrop`)
    - remove (`handleRemoveImage`)
    - clear on send success
    - restoreImages API
    - edit-mode transitions via `setDraft`

Optional but recommended (quota UX):
- Add a simple size guard before persisting:
  - compute `JSON.stringify(nextImages).length`
- if above a conservative threshold (e.g. 6–8MB chars), skip persistence
and show a toast like:
    - “Image draft too large to save; it will be lost on restart.”
- Still keep `imageAttachments` in memory so the user can send the
message.

### 4) Clear pending image drafts on successful workspace creation
**Files:**
- `src/browser/components/ChatInput/useCreationWorkspace.ts`

Actions:
- When creation succeeds (where we already clear `pendingInputKey`),
also clear:
-
`updatePersistedState(getInputImagesKey(getPendingScopeId(projectPath)),
undefined)`

### 5) Tests
**Goal:** prevent regressions in the future.

Recommended tests:
- `src/common/constants/storage.test.ts`
  - verifies `getInputImagesKey()` format
- verifies `copyWorkspaceStorage()` copies the image key (using a fake
localStorage)
  - verifies `deleteWorkspaceStorage()` removes the image key
- `src/browser/components/ChatInput` test (new):
- pre-populate localStorage with `inputImages:<workspaceId>` containing
one `ImageAttachment`
- render `ChatInput` for that workspace and assert the thumbnail `<img
src=...>` is present
- update state (e.g., remove image) and assert localStorage key is
removed/updated

### 6) Manual QA checklist
- Workspace draft persistence:
  1. Open workspace A → paste an image + type text.
  2. Switch to workspace B → switch back to A.
  3. Confirm text + image thumbnail are still present.
- Restart persistence:
  1. With an image attached in workspace A draft, fully restart Mux.
  2. Confirm the draft image is restored.
- Creation flow:
  1. Go to creation mode → paste an image.
  2. Create workspace.
3. Confirm the pending draft is cleared (image doesn’t “stick around” in
future creation attempts).
- Fork/removal:
  - Fork a workspace with a draft image: image draft should copy.
  - Delete a workspace: image draft key should be removed.

---

<details>
<summary>Alternative: store draft images on disk (more robust, more
work)</summary>

If localStorage size limits or sync write jank are a concern, store
draft images under `~/.mux/sessions/<workspaceId>/draft-images.json`
using `SessionFileManager` (server-side) and expose
`workspace.getDraft()` / `workspace.setDraft()` ORPC endpoints.

Pros:
- avoids localStorage quota
- avoids large synchronous localStorage writes in the renderer

Cons:
- new backend API surface
- more wiring + tests (IPC/orpc + session file persistence)

</details>

</details>

---
_Generated with `mux` • Model: `openai:gpt-5.2` • Thinking: `xhigh`_

---------

Signed-off-by: Thomas Kosiewski <tk@coder.com>

Author: ThomasK33

src/browser/components/ChatInput/draftImagesStorage.test.ts

@@ -0,0 +1,34 @@
+import { describe, expect, test } from "bun:test";
+
+import {
+  estimatePersistedImageAttachmentsChars,
+  parsePersistedImageAttachments,
+} from "./draftImagesStorage";
+
+describe("draftImagesStorage", () => {
+  test("parsePersistedImageAttachments returns [] for non-arrays", () => {
+    expect(parsePersistedImageAttachments(null)).toEqual([]);
+    expect(parsePersistedImageAttachments({})).toEqual([]);
+    expect(parsePersistedImageAttachments("nope")).toEqual([]);
+  });
+
+  test("parsePersistedImageAttachments returns [] for invalid array items", () => {
+    expect(parsePersistedImageAttachments([{}])).toEqual([]);
+    expect(
+      parsePersistedImageAttachments([{ id: "img", url: 123, mediaType: "image/png" }])
+    ).toEqual([]);
+  });
+
+  test("parsePersistedImageAttachments returns attachments for valid items", () => {
+    expect(
+      parsePersistedImageAttachments([
+        { id: "img-1", url: "data:image/png;base64,AAA", mediaType: "image/png" },
+      ])
+    ).toEqual([{ id: "img-1", url: "data:image/png;base64,AAA", mediaType: "image/png" }]);
+  });
+
+  test("estimatePersistedImageAttachmentsChars matches JSON length", () => {
+    const images = [{ id: "img-1", url: "data:image/png;base64,AAA", mediaType: "image/png" }];
+    expect(estimatePersistedImageAttachmentsChars(images)).toBe(JSON.stringify(images).length);
+  });
+});

src/browser/components/ChatInput/draftImagesStorage.ts

@@ -0,0 +1,39 @@
+import type { ImageAttachment } from "@/browser/components/ImageAttachments";
+import { readPersistedState } from "@/browser/hooks/usePersistedState";
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return typeof value === "object" && value !== null;
+}
+
+function isImageAttachment(value: unknown): value is ImageAttachment {
+  if (!isRecord(value)) return false;
+  return (
+    typeof value.id === "string" &&
+    typeof value.url === "string" &&
+    typeof value.mediaType === "string"
+  );
+}
+
+export function parsePersistedImageAttachments(raw: unknown): ImageAttachment[] {
+  if (!Array.isArray(raw)) {
+    return [];
+  }
+
+  const attachments: ImageAttachment[] = [];
+  for (const item of raw) {
+    if (!isImageAttachment(item)) {
+      return [];
+    }
+    attachments.push({ id: item.id, url: item.url, mediaType: item.mediaType });
+  }
+
+  return attachments;
+}
+
+export function readPersistedImageAttachments(imagesKey: string): ImageAttachment[] {
+  return parsePersistedImageAttachments(readPersistedState<unknown>(imagesKey, []));
+}
+
+export function estimatePersistedImageAttachmentsChars(images: ImageAttachment[]): number {
+  return JSON.stringify(images).length;
+}

src/browser/components/ChatInput/index.tsx

@@ -24,6 +24,7 @@ import { useSendMessageOptions } from "@/browser/hooks/useSendMessageOptions";
 import {
   getModelKey,
   getInputKey,
+  getInputImagesKey,
   VIM_ENABLED_KEY,
   getProjectScopeId,
   getPendingScopeId,
@@ -78,9 +79,16 @@ import { useCreationWorkspace } from "./useCreationWorkspace";
 import { useTutorial } from "@/browser/contexts/TutorialContext";
 import { useVoiceInput } from "@/browser/hooks/useVoiceInput";
 import { VoiceInputButton } from "./VoiceInputButton";
+import {
+  estimatePersistedImageAttachmentsChars,
+  readPersistedImageAttachments,
+} from "./draftImagesStorage";
 import { RecordingOverlay } from "./RecordingOverlay";
 import { ReviewBlockFromData } from "../shared/ReviewBlock";
 
+// localStorage quotas are environment-dependent and relatively small.
+// Be conservative here so we can warn the user before writes start failing.
+const MAX_PERSISTED_IMAGE_DRAFT_CHARS = 4_000_000;
 type TokenCountReader = () => number;
 
 function createTokenCountResource(promise: Promise<number>): TokenCountReader {
@@ -132,13 +140,16 @@ const ChatInputInner: React.FC<ChatInputProps> = (props) => {
   // Storage keys differ by variant
   const storageKeys = (() => {
     if (variant === "creation") {
+      const pendingScopeId = getPendingScopeId(props.projectPath);
       return {
-        inputKey: getInputKey(getPendingScopeId(props.projectPath)),
+        inputKey: getInputKey(pendingScopeId),
+        imagesKey: getInputImagesKey(pendingScopeId),
         modelKey: getModelKey(getProjectScopeId(props.projectPath)),
       };
     }
     return {
       inputKey: getInputKey(props.workspaceId),
+      imagesKey: getInputImagesKey(props.workspaceId),
       modelKey: getModelKey(props.workspaceId),
     };
   })();
@@ -150,10 +161,6 @@ const ChatInputInner: React.FC<ChatInputProps> = (props) => {
   const [commandSuggestions, setCommandSuggestions] = useState<SlashSuggestion[]>([]);
   const [providerNames, setProviderNames] = useState<string[]>([]);
   const [toast, setToast] = useState<Toast | null>(null);
-  const [imageAttachments, setImageAttachments] = useState<ImageAttachment[]>([]);
-  // Attached reviews come from parent via props (persisted in pendingReviews state)
-  const attachedReviews = variant === "workspace" ? (props.attachedReviews ?? []) : [];
-
   const pushToast = useCallback(
     (nextToast: Omit<Toast, "id">) => {
       setToast({ id: Date.now().toString(), ...nextToast });
@@ -163,6 +170,60 @@ const ChatInputInner: React.FC<ChatInputProps> = (props) => {
   const handleToastDismiss = useCallback(() => {
     setToast(null);
   }, []);
+
+  const imageDraftTooLargeToastKeyRef = useRef<string | null>(null);
+
+  const [imageAttachments, setImageAttachmentsState] = useState<ImageAttachment[]>(() => {
+    return readPersistedImageAttachments(storageKeys.imagesKey);
+  });
+  const persistImageAttachments = useCallback(
+    (nextImages: ImageAttachment[]) => {
+      if (nextImages.length === 0) {
+        imageDraftTooLargeToastKeyRef.current = null;
+        updatePersistedState<ImageAttachment[] | undefined>(storageKeys.imagesKey, undefined);
+        return;
+      }
+
+      const estimatedChars = estimatePersistedImageAttachmentsChars(nextImages);
+      if (estimatedChars > MAX_PERSISTED_IMAGE_DRAFT_CHARS) {
+        // Clear persisted value to avoid restoring stale images on restart.
+        updatePersistedState<ImageAttachment[] | undefined>(storageKeys.imagesKey, undefined);
+
+        if (imageDraftTooLargeToastKeyRef.current !== storageKeys.imagesKey) {
+          imageDraftTooLargeToastKeyRef.current = storageKeys.imagesKey;
+          pushToast({
+            type: "error",
+            message:
+              "This draft image is too large to save. It will be lost when you switch workspaces or restart.",
+            duration: 5000,
+          });
+        }
+        return;
+      }
+
+      imageDraftTooLargeToastKeyRef.current = null;
+      updatePersistedState<ImageAttachment[] | undefined>(storageKeys.imagesKey, nextImages);
+    },
+    [storageKeys.imagesKey, pushToast]
+  );
+
+  // Keep image drafts in sync when the storage scope changes (e.g. switching creation projects).
+  useEffect(() => {
+    imageDraftTooLargeToastKeyRef.current = null;
+    setImageAttachmentsState(readPersistedImageAttachments(storageKeys.imagesKey));
+  }, [storageKeys.imagesKey]);
+  const setImageAttachments = useCallback(
+    (value: ImageAttachment[] | ((prev: ImageAttachment[]) => ImageAttachment[])) => {
+      setImageAttachmentsState((prev) => {
+        const next = value instanceof Function ? value(prev) : value;
+        persistImageAttachments(next);
+        return next;
+      });
+    },
+    [persistImageAttachments]
+  );
+  // Attached reviews come from parent via props (persisted in pendingReviews state)
+  const attachedReviews = variant === "workspace" ? (props.attachedReviews ?? []) : [];
   const inputRef = useRef<HTMLTextAreaElement>(null);
   const modelSelectorRef = useRef<ModelSelectorRef>(null);
 
@@ -181,7 +242,7 @@ const ChatInputInner: React.FC<ChatInputProps> = (props) => {
       setInput(draft.text);
       setImageAttachments(draft.images);
     },
-    [setInput]
+    [setInput, setImageAttachments]
   );
   const preEditDraftRef = useRef<DraftState>({ text: "", images: [] });
   const { open } = useSettings();
@@ -387,14 +448,17 @@ const ChatInputInner: React.FC<ChatInputProps> = (props) => {
   );
 
   // Method to restore images to input (used by queued message edit)
-  const restoreImages = useCallback((images: ImagePart[]) => {
-    const attachments: ImageAttachment[] = images.map((img, index) => ({
-      id: `restored-${Date.now()}-${index}`,
-      url: img.url,
-      mediaType: img.mediaType,
-    }));
-    setImageAttachments(attachments);
-  }, []);
+  const restoreImages = useCallback(
+    (images: ImagePart[]) => {
+      const attachments: ImageAttachment[] = images.map((img, index) => ({
+        id: `restored-${Date.now()}-${index}`,
+        url: img.url,
+        mediaType: img.mediaType,
+      }));
+      setImageAttachments(attachments);
+    },
+    [setImageAttachments]
+  );
 
   // Provide API to parent via callback
   useEffect(() => {
@@ -678,24 +742,30 @@ const ChatInputInner: React.FC<ChatInputProps> = (props) => {
   }, [variant, workspaceIdForFocus, focusMessageInput, setChatInputAutoFocusState]);
 
   // Handle paste events to extract images
-  const handlePaste = useCallback((e: React.ClipboardEvent<HTMLTextAreaElement>) => {
-    const items = e.clipboardData?.items;
-    if (!items) return;
+  const handlePaste = useCallback(
+    (e: React.ClipboardEvent<HTMLTextAreaElement>) => {
+      const items = e.clipboardData?.items;
+      if (!items) return;
 
-    const imageFiles = extractImagesFromClipboard(items);
-    if (imageFiles.length === 0) return;
+      const imageFiles = extractImagesFromClipboard(items);
+      if (imageFiles.length === 0) return;
 
-    e.preventDefault(); // Prevent default paste behavior for images
+      e.preventDefault(); // Prevent default paste behavior for images
 
-    void processImageFiles(imageFiles).then((attachments) => {
-      setImageAttachments((prev) => [...prev, ...attachments]);
-    });
-  }, []);
+      void processImageFiles(imageFiles).then((attachments) => {
+        setImageAttachments((prev) => [...prev, ...attachments]);
+      });
+    },
+    [setImageAttachments]
+  );
 
   // Handle removing an image attachment
-  const handleRemoveImage = useCallback((id: string) => {
-    setImageAttachments((prev) => prev.filter((img) => img.id !== id));
-  }, []);
+  const handleRemoveImage = useCallback(
+    (id: string) => {
+      setImageAttachments((prev) => prev.filter((img) => img.id !== id));
+    },
+    [setImageAttachments]
+  );
 
   // Handle drag over to allow drop
   const handleDragOver = useCallback((e: React.DragEvent<HTMLTextAreaElement>) => {
@@ -707,16 +777,19 @@ const ChatInputInner: React.FC<ChatInputProps> = (props) => {
   }, []);
 
   // Handle drop to extract images
-  const handleDrop = useCallback((e: React.DragEvent<HTMLTextAreaElement>) => {
-    e.preventDefault();
+  const handleDrop = useCallback(
+    (e: React.DragEvent<HTMLTextAreaElement>) => {
+      e.preventDefault();
 
-    const imageFiles = extractImagesFromDrop(e.dataTransfer);
-    if (imageFiles.length === 0) return;
+      const imageFiles = extractImagesFromDrop(e.dataTransfer);
+      if (imageFiles.length === 0) return;
 
-    void processImageFiles(imageFiles).then((attachments) => {
-      setImageAttachments((prev) => [...prev, ...attachments]);
-    });
-  }, []);
+      void processImageFiles(imageFiles).then((attachments) => {
+        setImageAttachments((prev) => [...prev, ...attachments]);
+      });
+    },
+    [setImageAttachments]
+  );
 
   // Handle command selection
   const handleCommandSelect = useCallback(

src/browser/components/ChatInput/useCreationWorkspace.test.tsx

@@ -2,6 +2,7 @@ import type { APIClient } from "@/browser/contexts/API";
 import type { DraftWorkspaceSettings } from "@/browser/hooks/useDraftWorkspaceSettings";
 import {
   getInputKey,
+  getInputImagesKey,
   getModelKey,
   getModeKey,
   getPendingScopeId,
@@ -461,10 +462,13 @@ describe("useCreationWorkspace", () => {
     expect(readPersistedStateCalls).toContainEqual([projectModeKey, null]);
 
     const modeKey = getModeKey(TEST_WORKSPACE_ID);
-    const pendingInputKey = getInputKey(getPendingScopeId(TEST_PROJECT_PATH));
+    const pendingScopeId = getPendingScopeId(TEST_PROJECT_PATH);
+    const pendingInputKey = getInputKey(pendingScopeId);
+    const pendingImagesKey = getInputImagesKey(pendingScopeId);
     expect(updatePersistedStateCalls).toContainEqual([modeKey, "plan"]);
     // Note: thinking level is no longer synced per-workspace, it's stored per-model globally
     expect(updatePersistedStateCalls).toContainEqual([pendingInputKey, ""]);
+    expect(updatePersistedStateCalls).toContainEqual([pendingImagesKey, undefined]);
   });
 
   test("handleSend surfaces backend errors and resets state", async () => {

src/browser/components/ChatInput/useCreationWorkspace.ts

@@ -8,6 +8,7 @@ import { readPersistedState, updatePersistedState } from "@/browser/hooks/usePer
 import { getSendOptionsFromStorage } from "@/browser/utils/messages/sendOptions";
 import {
   getInputKey,
+  getInputImagesKey,
   getModelKey,
   getModeKey,
   getPendingScopeId,
@@ -198,8 +199,9 @@ export function useCreationWorkspace({
         // Sync preferences immediately (before switching)
         syncCreationPreferences(projectPath, metadata.id);
         if (projectPath) {
-          const pendingInputKey = getInputKey(getPendingScopeId(projectPath));
-          updatePersistedState(pendingInputKey, "");
+          const pendingScopeId = getPendingScopeId(projectPath);
+          updatePersistedState(getInputKey(pendingScopeId), "");
+          updatePersistedState(getInputImagesKey(pendingScopeId), undefined);
         }
 
         // Switch to the workspace IMMEDIATELY after creation to exit splash faster.

src/browser/components/ChatInputToast.test.tsx

@@ -0,0 +1,50 @@
+import React from "react";
+import { afterEach, beforeEach, describe, expect, test } from "bun:test";
+import { GlobalWindow } from "happy-dom";
+import { cleanup, fireEvent, render, waitFor } from "@testing-library/react";
+
+import { ChatInputToast, type Toast } from "./ChatInputToast";
+
+describe("ChatInputToast", () => {
+  beforeEach(() => {
+    globalThis.window = new GlobalWindow() as unknown as Window & typeof globalThis;
+    globalThis.document = globalThis.window.document;
+  });
+
+  afterEach(() => {
+    cleanup();
+    globalThis.window = undefined as unknown as Window & typeof globalThis;
+    globalThis.document = undefined as unknown as Document;
+  });
+
+  test("resets leaving state when a new toast is shown", async () => {
+    const toast1: Toast = { id: "toast-1", type: "error", message: "first" };
+    const toast2: Toast = { id: "toast-2", type: "error", message: "second" };
+
+    function Harness() {
+      const [toast, setToast] = React.useState<Toast | null>(toast1);
+      return (
+        <div>
+          <ChatInputToast toast={toast} onDismiss={() => undefined} />
+          <button onClick={() => setToast(toast2)}>Next toast</button>
+        </div>
+      );
+    }
+
+    const { getByLabelText, getByRole, getByText } = render(<Harness />);
+
+    fireEvent.click(getByLabelText("Dismiss"));
+
+    await waitFor(() => {
+      expect(getByRole("alert").className).toContain("toastFadeOut");
+    });
+
+    fireEvent.click(getByText("Next toast"));
+
+    await waitFor(() => {
+      const className = getByRole("alert").className;
+      expect(className).toContain("toastSlideIn");
+      expect(className).not.toContain("toastFadeOut");
+    });
+  });
+});