Fix queued chat editing to link component selections and annotations (#3029)

e2e-tests/queued_message.spec.ts

-import { test } from "./helpers/test_helper";
+import { test, testSkipIfWindows, Timeout } from "./helpers/test_helper";
 import { expect, Locator } from "@playwright/test";
 
 test.describe("queued messages", () => {
@@ -150,3 +150,145 @@ test.describe("queued messages", () => {
     await expect(messagesList.getByText("tc=2")).toBeVisible();
   });
 });
+
+testSkipIfWindows(
+  "editing queued message restores attachments and selected components",
+  async ({ po }) => {
+    await po.setUp();
+    const chatInput = po.chatActions.getChatInput();
+
+    // Build an app so we have a preview with selectable components
+    await po.sendPrompt("tc=basic");
+    await po.previewPanel.clickTogglePreviewPanel();
+
+    // Start a slow streaming response so subsequent messages get queued
+    await po.sendPrompt("tc=1 [sleep=medium]", {
+      skipWaitForCompletion: true,
+    });
+    await expect(chatInput).toBeVisible();
+
+    // While streaming, select a component
+    await po.previewPanel.clickPreviewPickElement();
+    await po.previewPanel
+      .getPreviewIframeElement()
+      .contentFrame()
+      .getByRole("heading", { name: "Welcome to Your Blank App" })
+      .click();
+    await expect(po.previewPanel.getSelectedComponentsDisplay()).toBeVisible({
+      timeout: Timeout.SHORT,
+    });
+
+    // Attach a file
+    await po.chatActions
+      .getChatInputContainer()
+      .getByTestId("auxiliary-actions-menu")
+      .click();
+    await po.page.getByRole("menuitem", { name: "Attach files" }).click();
+    const chatContextItem = po.page.getByText("Attach file as chat context");
+    await expect(chatContextItem).toBeVisible();
+    const fileChooserPromise = po.page.waitForEvent("filechooser");
+    await chatContextItem.click();
+    const fileChooser = await fileChooserPromise;
+    await fileChooser.setFiles("e2e-tests/fixtures/images/logo.png");
+    await expect(po.page.getByText("logo.png")).toBeVisible();
+
+    // Queue a message with both attachment and selected component
+    await chatInput.fill("queued with extras");
+    await chatInput.press("Enter");
+
+    // After queuing, both should be cleared
+    await expect(
+      po.previewPanel.getSelectedComponentsDisplay(),
+    ).not.toBeVisible();
+    await expect(po.page.getByText("logo.png")).not.toBeVisible();
+    await expect(po.page.getByText(/\d+ Queued/)).toBeVisible();
+
+    // Edit the queued message
+    const queuedRow = po.page.locator("li", {
+      hasText: "queued with extras",
+    });
+    await queuedRow.hover();
+    await queuedRow.getByTitle("Edit").click();
+
+    // The input should contain the queued message text
+    await expect(chatInput).toContainText("queued with extras");
+
+    // Both attachment and selected components should be restored
+    await expect(po.page.getByText("logo.png")).toBeVisible();
+    await expect(po.previewPanel.getSelectedComponentsDisplay()).toBeVisible({
+      timeout: Timeout.SHORT,
+    });
+
+    // Submit the edit — both should clear again
+    await chatInput.press("Enter");
+    await expect(po.page.getByText("logo.png")).not.toBeVisible();
+    await expect(
+      po.previewPanel.getSelectedComponentsDisplay(),
+    ).not.toBeVisible();
+
+    // Wait for all messages to complete
+    await po.chatActions.waitForChatCompletion();
+    await po.chatActions.waitForChatCompletion();
+
+    // Verify both messages were sent
+    const messagesList = po.page.locator('[data-testid="messages-list"]');
+    await expect(messagesList.getByText("tc=1 [sleep=medium]")).toBeVisible();
+    await expect(messagesList.getByText("queued with extras")).toBeVisible();
+  },
+);
+
+testSkipIfWindows(
+  "canceling queued message edit clears restored components",
+  async ({ po }) => {
+    await po.setUp();
+    const chatInput = po.chatActions.getChatInput();
+
+    // Build an app so we have a preview with selectable components
+    await po.sendPrompt("tc=basic");
+    await po.previewPanel.clickTogglePreviewPanel();
+
+    // Start a slow streaming response so subsequent messages get queued
+    await po.sendPrompt("tc=1 [sleep=medium]", {
+      skipWaitForCompletion: true,
+    });
+    await expect(chatInput).toBeVisible();
+
+    // While streaming, select a component and queue a message with it
+    await po.previewPanel.clickPreviewPickElement();
+    await po.previewPanel
+      .getPreviewIframeElement()
+      .contentFrame()
+      .getByRole("heading", { name: "Welcome to Your Blank App" })
+      .click();
+    await expect(po.previewPanel.getSelectedComponentsDisplay()).toBeVisible({
+      timeout: Timeout.SHORT,
+    });
+
+    await chatInput.fill("queued with component");
+    await chatInput.press("Enter");
+    await expect(po.page.getByText(/\d+ Queued/)).toBeVisible();
+
+    // Edit the queued message — components should be restored
+    const queuedRow = po.page.locator("li", {
+      hasText: "queued with component",
+    });
+    await queuedRow.hover();
+    await queuedRow.getByTitle("Edit").click();
+    await expect(po.previewPanel.getSelectedComponentsDisplay()).toBeVisible({
+      timeout: Timeout.SHORT,
+    });
+
+    // Cancel the edit — components should be cleared
+    await po.page.getByText("Cancel", { exact: true }).click();
+    await expect(
+      po.previewPanel.getSelectedComponentsDisplay(),
+    ).not.toBeVisible();
+
+    // Input should be empty after cancel
+    await expect(chatInput).toBeEmpty();
+
+    // Wait for the in-flight chat and the queued message to finish before ending the test
+    await po.chatActions.waitForChatCompletion();
+    await po.chatActions.waitForChatCompletion();
+  },
+);

e2e-tests/snapshots/select_component.spec.ts_deselect-component-1.aria.yml

@@ -4,5 +4,4 @@
 - heading "Welcome to Your Blank App" [level=1]
 - paragraph: Start building your amazing project here!
 - link "Made with Dyad":
-  - /url: https://www.dyad.sh/
-- text: h1 src/pages/Index.tsx
\ No newline at end of file
+  - /url: https://www.dyad.sh/
\ No newline at end of file

e2e-tests/snapshots/select_component.spec.ts_deselect-individual-component-from-multiple-2.aria.yml

@@ -4,5 +4,4 @@
 - heading "Welcome to Your Blank App" [level=1]
 - paragraph: Start building your amazing project here!
 - link "Made with Dyad":
-  - /url: https://www.dyad.sh/
-- text: a src/components/made-with-dyad.tsx
\ No newline at end of file
+  - /url: https://www.dyad.sh/
\ No newline at end of file

e2e-tests/snapshots/select_component.spec.ts_select-component-1.aria.yml

@@ -4,5 +4,4 @@
 - heading "Welcome to Your Blank App" [level=1]
 - paragraph: Start building your amazing project here!
 - link "Made with Dyad":
-  - /url: https://www.dyad.sh/
-- text: h1 src/pages/Index.tsx
\ No newline at end of file
+  - /url: https://www.dyad.sh/
\ No newline at end of file

e2e-tests/snapshots/select_component.spec.ts_select-multiple-components-1.aria.yml

@@ -4,5 +4,4 @@
 - heading "Welcome to Your Blank App" [level=1]
 - paragraph: Start building your amazing project here!
 - link "Made with Dyad":
-  - /url: https://www.dyad.sh/
-- text: a src/components/made-with-dyad.tsx
\ No newline at end of file
+  - /url: https://www.dyad.sh/
\ No newline at end of file

src/atoms/previewAtoms.ts

@@ -17,6 +17,8 @@ export const previewIframeRefAtom = atom<HTMLIFrameElement | null>(null);
 
 export const annotatorModeAtom = atom<boolean>(false);
 
+export const isRestoringQueuedSelectionAtom = atom<boolean>(false);
+
 export const screenshotDataUrlAtom = atom<string | null>(null);
 export const pendingVisualChangesAtom = atom<Map<string, VisualEditingChange>>(
   new Map(),

src/components/chat/ChatInput.tsx

@@ -72,6 +72,7 @@ import {
   visualEditingSelectedComponentAtom,
   currentComponentCoordinatesAtom,
   pendingVisualChangesAtom,
+  isRestoringQueuedSelectionAtom,
 } from "@/atoms/previewAtoms";
 import { SelectedComponentsDisplay } from "./SelectedComponentDisplay";
 import { useCheckProblems } from "@/hooks/useCheckProblems";
@@ -147,6 +148,9 @@ export function ChatInput({ chatId }: { chatId?: number }) {
     currentComponentCoordinatesAtom,
   );
   const setPendingVisualChanges = useSetAtom(pendingVisualChangesAtom);
+  const setIsRestoringQueuedSelection = useSetAtom(
+    isRestoringQueuedSelectionAtom,
+  );
   const [pendingAgentConsents, setPendingAgentConsents] = useAtom(
     pendingAgentConsentsAtom,
   );
@@ -175,6 +179,7 @@ export function ChatInput({ chatId }: { chatId?: number }) {
     handleDragLeave,
     handleDrop,
     clearAttachments,
+    replaceAttachments,
     handlePaste,
     confirmPendingFiles,
     cancelPendingFiles,
@@ -276,17 +281,95 @@ export function ChatInput({ chatId }: { chatId?: number }) {
     });
   }, [chatId, setMessagesById]);
 
+  // Shared cleanup for exiting queued message editing state
+  const resetEditingState = useCallback(() => {
+    setEditingQueuedMessageId(null);
+    setInputValue("");
+    clearAttachments();
+    setSelectedComponents([]);
+    setVisualEditingSelectedComponent(null);
+    if (previewIframeRef?.contentWindow) {
+      previewIframeRef.contentWindow.postMessage(
+        { type: "clear-dyad-component-overlays" },
+        "*",
+      );
+    }
+  }, [
+    setInputValue,
+    clearAttachments,
+    setSelectedComponents,
+    setVisualEditingSelectedComponent,
+    previewIframeRef,
+  ]);
+
+  // Clear editing state if the edited queued message is auto-dequeued
+  useEffect(() => {
+    if (!editingQueuedMessageId) return;
+    const stillInQueue = queuedMessages.some(
+      (m) => m.id === editingQueuedMessageId,
+    );
+    if (!stillInQueue) {
+      resetEditingState();
+    }
+  }, [editingQueuedMessageId, queuedMessages, resetEditingState]);
+
+  // Track editing state in a ref for unmount cleanup
+  const editingQueuedMessageIdRef = useRef(editingQueuedMessageId);
+  editingQueuedMessageIdRef.current = editingQueuedMessageId;
+
+  // Clear editing extras on unmount to avoid leaking state across navigations
+  useEffect(() => {
+    return () => {
+      if (editingQueuedMessageIdRef.current) {
+        clearAttachments();
+        setSelectedComponents([]);
+        setVisualEditingSelectedComponent(null);
+      }
+    };
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, []);
+
   // Queue management handlers
   const handleEditQueuedMessage = useCallback(
     (id: string) => {
       const msg = queuedMessages.find((m) => m.id === id);
       if (!msg) return;
+      // Auto-save current edits if switching between queued messages
+      if (editingQueuedMessageId && editingQueuedMessageId !== id) {
+        const componentsToSave =
+          selectedComponents && selectedComponents.length > 0
+            ? selectedComponents
+            : [];
+        updateQueuedMessage(editingQueuedMessageId, {
+          prompt: inputValue,
+          attachments,
+          selectedComponents: componentsToSave,
+        });
+      }
       // Load the message content into the input
       setInputValue(msg.prompt);
+      // Restore attachments and selected components from the queued message
+      replaceAttachments(msg.attachments ?? []);
+      setIsRestoringQueuedSelection(true);
+      setSelectedComponents(msg.selectedComponents ?? []);
+      // Reset visual editing target to avoid stale toolbar state
+      setVisualEditingSelectedComponent(null);
       // Set editing mode
       setEditingQueuedMessageId(id);
     },
-    [queuedMessages, setInputValue],
+    [
+      queuedMessages,
+      editingQueuedMessageId,
+      inputValue,
+      attachments,
+      selectedComponents,
+      setInputValue,
+      replaceAttachments,
+      setSelectedComponents,
+      setVisualEditingSelectedComponent,
+      setIsRestoringQueuedSelection,
+      updateQueuedMessage,
+    ],
   );
 
   const handleMoveUp = useCallback(
@@ -313,12 +396,11 @@ export function ChatInput({ chatId }: { chatId?: number }) {
     (id: string) => {
       // Clear editing state if deleting the message being edited
       if (editingQueuedMessageId === id) {
-        setEditingQueuedMessageId(null);
-        setInputValue("");
+        resetEditingState();
       }
       removeQueuedMessage(id);
     },
-    [editingQueuedMessageId, removeQueuedMessage, setInputValue],
+    [editingQueuedMessageId, removeQueuedMessage, resetEditingState],
   );
 
   const handleSubmit = async () => {
@@ -370,9 +452,10 @@ export function ChatInput({ chatId }: { chatId?: number }) {
     if (editingQueuedMessageId) {
       updateQueuedMessage(editingQueuedMessageId, {
         prompt: currentInput,
+        attachments,
+        selectedComponents: componentsToSend,
       });
-      setInputValue("");
-      setEditingQueuedMessageId(null);
+      resetEditingState();
       return;
     }
 
@@ -433,9 +516,9 @@ export function ChatInput({ chatId }: { chatId?: number }) {
     // could potentially run if the backend responds before queue clearing.
     clearAllQueuedMessages();
     // Reset editing state so the "Editing queued message" banner is dismissed
+    // and restored attachments/components are cleared
     if (editingQueuedMessageId) {
-      setEditingQueuedMessageId(null);
-      setInputValue("");
+      resetEditingState();
     }
     if (chatId) {
       ipc.chat.cancelStream(chatId);
@@ -630,10 +713,7 @@ export function ChatInput({ chatId }: { chatId?: number }) {
               </span>
               <button
                 type="button"
-                onClick={() => {
-                  setEditingQueuedMessageId(null);
-                  setInputValue("");
-                }}
+                onClick={() => resetEditingState()}
                 className="text-xs text-muted-foreground hover:text-foreground cursor-pointer"
               >
                 Cancel

src/components/preview_panel/PreviewIframe.tsx

@@ -48,6 +48,7 @@ import {
   annotatorModeAtom,
   screenshotDataUrlAtom,
   pendingVisualChangesAtom,
+  isRestoringQueuedSelectionAtom,
 } from "@/atoms/previewAtoms";
 import { isChatPanelHiddenAtom } from "@/atoms/viewAtoms";
 import { ComponentSelection } from "@/ipc/types";
@@ -211,9 +212,12 @@ export const PreviewIframe = ({ loading }: { loading: boolean }) => {
     }
     return 0;
   });
-  const setSelectedComponentsPreview = useSetAtom(
+  const [selectedComponentsPreview, setSelectedComponentsPreview] = useAtom(
     selectedComponentsPreviewAtom,
   );
+  const [isRestoringQueuedSelection, setIsRestoringQueuedSelection] = useAtom(
+    isRestoringQueuedSelectionAtom,
+  );
   const [visualEditingSelectedComponent, setVisualEditingSelectedComponent] =
     useAtom(visualEditingSelectedComponentAtom);
   const setCurrentComponentCoordinates = useSetAtom(
@@ -378,6 +382,37 @@ export const PreviewIframe = ({ loading }: { loading: boolean }) => {
     }
   }, [isProMode, isComponentSelectorInitialized]);
 
+  // Restore component overlays in iframe only during queued-message edit restoration.
+  // Normal interactive selections are already handled by the iframe's own click handler,
+  // so we guard this effect to avoid redundant clear-and-restore round-trips.
+  useEffect(() => {
+    if (!isRestoringQueuedSelection) return;
+    if (!iframeRef.current?.contentWindow || !isComponentSelectorInitialized) {
+      return;
+    }
+    // Clear the flag before sending so it doesn't re-trigger
+    setIsRestoringQueuedSelection(false);
+    if (selectedComponentsPreview.length === 0) {
+      iframeRef.current.contentWindow.postMessage(
+        { type: "clear-dyad-component-overlays" },
+        "*",
+      );
+      return;
+    }
+    iframeRef.current.contentWindow.postMessage(
+      {
+        type: "restore-dyad-component-overlays",
+        componentIds: selectedComponentsPreview.map((c) => c.id),
+      },
+      "*",
+    );
+  }, [
+    isRestoringQueuedSelection,
+    selectedComponentsPreview,
+    isComponentSelectorInitialized,
+    setIsRestoringQueuedSelection,
+  ]);
+
   // Add message listener for iframe errors and navigation events
   useEffect(() => {
     const handleMessage = (event: MessageEvent) => {

src/hooks/useAttachments.ts

@@ -98,6 +98,11 @@ export function useAttachments() {
     setPendingFiles(null);
   };
 
+  const replaceAttachments = (newAttachments: FileAttachment[]) => {
+    setAttachments(newAttachments);
+    setPendingFiles(null);
+  };
+
   const handlePaste = async (e: React.ClipboardEvent) => {
     if (pendingFiles) return;
 
@@ -153,6 +158,7 @@ export function useAttachments() {
     clearAttachments,
     handlePaste,
     addAttachments,
+    replaceAttachments,
     confirmPendingFiles,
     cancelPendingFiles,
   };

worker/dyad-component-selector-client.js

@@ -605,6 +605,21 @@
         removeOverlayById(e.data.componentId);
       }
     }
+    if (e.data.type === "restore-dyad-component-overlays") {
+      const componentIds = e.data.componentIds;
+      if (Array.isArray(componentIds)) {
+        clearOverlays();
+        for (const id of componentIds) {
+          const el = document.querySelector(
+            `[data-dyad-id="${CSS.escape(id)}"]`,
+          );
+          if (el) {
+            updateOverlay(el, true);
+          }
+        }
+        requestAnimationFrame(updateAllOverlayPositions);
+      }
+    }
   });
 
   // Always listen for keyboard shortcuts