Add image generation from chat (#3055)

e2e-tests/chat_image_generation.spec.ts

+import { expect } from "@playwright/test";
+import { test, Timeout } from "./helpers/test_helper";
+
+/**
+ * E2E tests for generating an image from the chat via the auxiliary actions menu.
+ * This tests the flow: + menu → Generate Image → fill prompt → Generate → image appears in strip → send auto-adds to chat.
+ */
+
+test("generate image from chat - full flow", async ({ po }) => {
+  await po.setUpDyadPro();
+  await po.importApp("minimal");
+
+  // Approve the code proposal from the import so the send button is unblocked
+  await po.approveProposal();
+
+  // Open auxiliary actions menu in the chat input
+  await po.chatActions
+    .getChatInputContainer()
+    .getByTestId("auxiliary-actions-menu")
+    .click();
+
+  // Click "Generate Image" menu item
+  const generateImageItem = po.page.getByTestId("generate-image-menu-item");
+  await expect(generateImageItem).toBeVisible();
+  await generateImageItem.click();
+
+  // The Image Generator dialog should be open
+  const dialog = po.page.getByRole("dialog");
+  await expect(dialog).toBeVisible();
+
+  // Fill in the prompt
+  const promptTextarea = dialog.getByPlaceholder(
+    "Describe the image you want to create...",
+  );
+  await expect(promptTextarea).toBeVisible();
+  await promptTextarea.fill("A beautiful sunset over mountains");
+
+  // Click Generate (app is auto-selected since there's only one)
+  const generateButton = dialog.getByRole("button", { name: "Generate" });
+  await expect(generateButton).toBeEnabled();
+  await generateButton.click();
+
+  // Dialog should close after clicking Generate
+  await expect(dialog).not.toBeVisible();
+
+  // Wait for the generated image to appear in the strip (thumbnail appears on success)
+  const imageStrip = po.chatActions.getChatInputContainer();
+  const generatedImage = imageStrip.locator(
+    "img[alt='A beautiful sunset over mountains']",
+  );
+  await expect(generatedImage).toBeVisible({ timeout: Timeout.LONG });
+
+  // The send button should be enabled even without text input
+  const sendButton = po.page.getByRole("button", { name: "Send message" });
+  await expect(sendButton).toBeEnabled();
+
+  // Click send - images are automatically added to the message
+  await sendButton.click();
+
+  // The image strip entry should be dismissed after sending
+  await expect(generatedImage).not.toBeVisible();
+
+  // Verify the sent message contains the generated image (rendered as an image element)
+  const messagesList = po.page.locator('[data-testid="messages-list"]');
+  await expect(
+    messagesList.locator("img[alt*='generated_a_beautiful_sunset']"),
+  ).toBeVisible({ timeout: Timeout.LONG });
+});

src/atoms/imageGenerationAtoms.ts

@@ -17,6 +17,7 @@ export interface ImageGenerationJob {
   startedAt: number;
   result?: GenerateImageResponse;
   error?: string;
+  source?: "chat" | "media-library";
 }
 
 const THIRTY_MINUTES_MS = 30 * 60 * 1000;
@@ -47,3 +48,16 @@ export const pendingImageGenerationsCountAtom = atom((get) => {
   const jobs = get(imageGenerationJobsAtom);
   return jobs.filter((job) => job.status === "pending").length;
 });
+
+export const chatImageGenerationJobsAtom = atom((get) => {
+  const jobs = get(imageGenerationJobsAtom);
+  // Only jobs with source === "chat" appear in the chat strip.
+  // Jobs from media.tsx / library-home.tsx intentionally omit `source`
+  // and therefore never appear here.
+  return jobs.filter((job) => job.source === "chat");
+});
+
+/** Tracks dismissed job IDs globally so dismissals persist across mounts. */
+export const dismissedImageGenerationJobIdsAtom = atom<Set<string>>(
+  new Set<string>(),
+);

src/components/ImageGeneratorDialog.tsx

-import { useState } from "react";
+import { useState, useEffect } from "react";
 import {
   ImageIcon,
   Box,
@@ -61,9 +61,13 @@ const THEME_MODES: {
 export function ImageGeneratorDialog({
   open,
   onOpenChange,
+  defaultAppId,
+  source,
 }: {
   open: boolean;
   onOpenChange: (open: boolean) => void;
+  defaultAppId?: number;
+  source?: "chat" | "media-library";
 }) {
   const [prompt, setPrompt] = useState("");
   const [themeMode, setThemeMode] = useState<ImageThemeMode>("plain");
@@ -74,6 +78,14 @@ export function ImageGeneratorDialog({
   const { userBudget, isLoadingUserBudget: isBudgetLoading } =
     useUserBudgetInfo();
 
+  // Sync defaultAppId only when dialog opens (not while already open)
+  useEffect(() => {
+    if (open && defaultAppId != null) {
+      setTargetAppId(defaultAppId);
+    }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [open]);
+
   const effectiveTargetAppId =
     targetAppId ?? (apps.length === 1 ? apps[0].id : null);
 
@@ -89,6 +101,7 @@ export function ImageGeneratorDialog({
       themeMode,
       targetAppId: effectiveTargetAppId,
       targetAppName: targetApp.name,
+      source,
     });
 
     // Auto-close dialog immediately after starting generation
@@ -205,7 +218,7 @@ export function ImageGeneratorDialog({
         </div>
 
         <DialogFooter>
-          <Button variant="outline" onClick={() => onOpenChange(false)}>
+          <Button variant="outline" onClick={() => handleOpenChange(false)}>
             Cancel
           </Button>
           <div className="flex items-center gap-2">

src/components/chat/AttachmentsList.tsx

@@ -34,7 +34,7 @@ export function AttachmentsList({
                 <img
                   src={URL.createObjectURL(attachment.file)}
                   alt={attachment.file.name}
-                  className="w-5 h-5 object-cover rounded"
+                  className="w-12 h-12 object-cover rounded-md"
                   onLoad={(e) =>
                     URL.revokeObjectURL((e.target as HTMLImageElement).src)
                   }
@@ -42,7 +42,7 @@ export function AttachmentsList({
                     URL.revokeObjectURL((e.target as HTMLImageElement).src)
                   }
                 />
-                <div className="absolute hidden group-hover:block top-6 left-0 z-10">
+                <div className="absolute hidden group-hover:block top-14 left-0 z-10">
                   <img
                     src={URL.createObjectURL(attachment.file)}
                     alt={attachment.file.name}

src/components/chat/AuxiliaryActionsMenu.tsx

@@ -9,6 +9,7 @@ import {
   Brush,
   PlusCircle,
   MoreHorizontal,
+  ImageIcon,
 } from "lucide-react";
 import {
   DropdownMenu,
@@ -46,6 +47,7 @@ interface AuxiliaryActionsMenuProps {
   toggleShowTokenBar?: () => void;
   hideContextFilesPicker?: boolean;
   appId?: number;
+  onGenerateImage?: () => void;
 }
 
 export function AuxiliaryActionsMenu({
@@ -54,6 +56,7 @@ export function AuxiliaryActionsMenu({
   toggleShowTokenBar,
   hideContextFilesPicker,
   appId,
+  onGenerateImage,
 }: AuxiliaryActionsMenuProps) {
   const [isOpen, setIsOpen] = useState(false);
   const [customThemeDialogOpen, setCustomThemeDialogOpen] = useState(false);
@@ -277,6 +280,21 @@ export function AuxiliaryActionsMenu({
             </DropdownMenuSubContent>
           </DropdownMenuSub>
 
+          {/* Generate Image */}
+          {onGenerateImage && (
+            <DropdownMenuItem
+              onClick={() => {
+                setIsOpen(false);
+                onGenerateImage();
+              }}
+              className="py-2 px-3"
+              data-testid="generate-image-menu-item"
+            >
+              <ImageIcon size={16} className="mr-2" />
+              Generate Image
+            </DropdownMenuItem>
+          )}
+
           {toggleShowTokenBar && (
             <>
               <DropdownMenuSeparator />

src/components/chat/ChatImageGenerationStrip.tsx

+import { useEffect, useState } from "react";
+import { useAtom, useAtomValue } from "jotai";
+import { X, Loader2, Plus, AlertCircle, RotateCcw } from "lucide-react";
+import {
+  chatImageGenerationJobsAtom,
+  dismissedImageGenerationJobIdsAtom,
+} from "@/atoms/imageGenerationAtoms";
+import { selectedAppIdAtom } from "@/atoms/appAtoms";
+import {
+  useCancelImageGeneration,
+  useGenerateImage,
+} from "@/hooks/useGenerateImage";
+import { buildDyadMediaUrl } from "@/lib/dyadMediaUrl";
+import { ImageLightbox } from "./ImageLightbox";
+import type { ImageGenerationJob } from "@/atoms/imageGenerationAtoms";
+
+interface ChatImageGenerationStripProps {
+  onGenerateImage: () => void;
+}
+
+export function ChatImageGenerationStrip({
+  onGenerateImage,
+}: ChatImageGenerationStripProps) {
+  const jobs = useAtomValue(chatImageGenerationJobsAtom);
+  const selectedAppId = useAtomValue(selectedAppIdAtom);
+  const cancelImageGeneration = useCancelImageGeneration();
+  const generateImage = useGenerateImage();
+  const [dismissedJobIds, setDismissedJobIds] = useAtom(
+    dismissedImageGenerationJobIdsAtom,
+  );
+  const [lightboxJob, setLightboxJob] = useState<ImageGenerationJob | null>(
+    null,
+  );
+
+  // Prune stale dismissed IDs that no longer correspond to active jobs
+  useEffect(() => {
+    const validJobIds = new Set(jobs.map((j) => j.id));
+    if ([...dismissedJobIds].some((id) => !validJobIds.has(id))) {
+      setDismissedJobIds(
+        new Set([...dismissedJobIds].filter((id) => validJobIds.has(id))),
+      );
+    }
+  }, [jobs, dismissedJobIds, setDismissedJobIds]);
+
+  // Only show jobs for the currently selected app
+  const appJobs = selectedAppId
+    ? jobs.filter((job) => job.targetAppId === selectedAppId)
+    : jobs;
+
+  const visibleJobs = appJobs.filter(
+    (job) =>
+      !dismissedJobIds.has(job.id) &&
+      (job.status === "pending" ||
+        job.status === "success" ||
+        job.status === "error"),
+  );
+
+  if (visibleJobs.length === 0) return null;
+
+  const handleDismiss = (jobId: string) => {
+    setDismissedJobIds((prev: Set<string>) => new Set(prev).add(jobId));
+  };
+
+  const handleRetry = (job: ImageGenerationJob) => {
+    setDismissedJobIds((prev: Set<string>) => new Set(prev).add(job.id));
+    generateImage.mutate({
+      requestId: crypto.randomUUID(),
+      prompt: job.prompt,
+      themeMode: job.themeMode,
+      targetAppId: job.targetAppId,
+      targetAppName: job.targetAppName,
+      source: job.source,
+    });
+  };
+
+  const handleCancel = (jobId: string) => {
+    void cancelImageGeneration(jobId);
+    setDismissedJobIds((prev: Set<string>) => new Set(prev).add(jobId));
+  };
+
+  return (
+    <>
+      <div className="px-2 pt-2 flex flex-wrap items-center gap-2">
+        {visibleJobs.map((job) => (
+          <div
+            key={job.id}
+            className="flex items-center bg-muted rounded-lg px-2 py-1.5 text-xs gap-2"
+          >
+            {job.status === "pending" ? (
+              <>
+                <div className="w-12 h-12 rounded-md bg-muted-foreground/10 animate-pulse flex items-center justify-center shrink-0">
+                  <Loader2
+                    size={16}
+                    className="animate-spin text-muted-foreground"
+                  />
+                </div>
+                <div className="min-w-0 flex-1">
+                  <span className="text-muted-foreground truncate block max-w-[120px]">
+                    {job.prompt}
+                  </span>
+                  <span className="text-muted-foreground/60 text-[10px]">
+                    Generating...
+                  </span>
+                </div>
+                <button
+                  onClick={() => handleCancel(job.id)}
+                  className="hover:bg-muted-foreground/20 rounded-full p-1.5 shrink-0"
+                  aria-label="Cancel generation"
+                >
+                  <X size={12} />
+                </button>
+              </>
+            ) : job.status === "error" ? (
+              <>
+                <div className="w-12 h-12 rounded-md bg-destructive/15 flex items-center justify-center shrink-0">
+                  <AlertCircle
+                    size={16}
+                    className="text-destructive-foreground"
+                  />
+                </div>
+                <div className="min-w-0 flex-1">
+                  <span
+                    className="text-destructive-foreground truncate block max-w-[120px]"
+                    title={job.error ?? "Generation failed"}
+                  >
+                    {job.error ?? "Generation failed"}
+                  </span>
+                </div>
+                <button
+                  onClick={() => handleRetry(job)}
+                  className="hover:bg-muted-foreground/20 rounded-full p-1.5 shrink-0"
+                  aria-label="Retry generation"
+                  title="Retry"
+                >
+                  <RotateCcw size={12} />
+                </button>
+                <button
+                  onClick={() => handleDismiss(job.id)}
+                  className="hover:bg-muted-foreground/20 rounded-full p-1.5 shrink-0"
+                  aria-label="Dismiss"
+                >
+                  <X size={12} />
+                </button>
+              </>
+            ) : (
+              <>
+                {job.result && (
+                  <img
+                    src={buildDyadMediaUrl(
+                      job.result.appPath,
+                      job.result.fileName,
+                    )}
+                    alt={job.prompt}
+                    className="w-12 h-12 rounded-md object-cover shrink-0 cursor-pointer hover:opacity-80 transition-opacity"
+                    onClick={() => setLightboxJob(job)}
+                  />
+                )}
+                <div className="min-w-0 flex-1">
+                  <span className="truncate block max-w-[120px]">
+                    {job.result?.fileName ?? "Generated image"}
+                  </span>
+                </div>
+                <button
+                  onClick={() => handleDismiss(job.id)}
+                  className="hover:bg-muted-foreground/20 rounded-full p-1.5 shrink-0"
+                  aria-label="Dismiss"
+                >
+                  <X size={12} />
+                </button>
+              </>
+            )}
+          </div>
+        ))}
+        <button
+          onClick={onGenerateImage}
+          className="group flex items-center justify-center w-12 h-12 shrink-0 cursor-pointer"
+          aria-label="Generate another image"
+          title="Generate another image"
+        >
+          <Plus
+            size={18}
+            className="text-muted-foreground group-hover:text-foreground transition-colors"
+          />
+        </button>
+      </div>
+
+      {lightboxJob?.result && (
+        <ImageLightbox
+          imageUrl={buildDyadMediaUrl(
+            lightboxJob.result.appPath,
+            lightboxJob.result.fileName,
+          )}
+          alt={lightboxJob.prompt}
+          filePath={lightboxJob.result.filePath}
+          onClose={() => setLightboxJob(null)}
+        />
+      )}
+    </>
+  );
+}

src/components/chat/ChatInput.tsx

@@ -78,6 +78,12 @@ import { SelectedComponentsDisplay } from "./SelectedComponentDisplay";
 import { useCheckProblems } from "@/hooks/useCheckProblems";
 import { LexicalChatInput } from "./LexicalChatInput";
 import { AuxiliaryActionsMenu } from "./AuxiliaryActionsMenu";
+import { ChatImageGenerationStrip } from "./ChatImageGenerationStrip";
+import {
+  chatImageGenerationJobsAtom,
+  dismissedImageGenerationJobIdsAtom,
+} from "@/atoms/imageGenerationAtoms";
+import { ImageGeneratorDialog } from "@/components/ImageGeneratorDialog";
 import { useChatModeToggle } from "@/hooks/useChatModeToggle";
 import { VisualEditingChangesDialog } from "@/components/preview_panel/VisualEditingChangesDialog";
 import { useUserBudgetInfo } from "@/hooks/useUserBudgetInfo";
@@ -168,6 +174,29 @@ export function ChatInput({ chatId }: { chatId?: number }) {
   const { navigate } = useRouter();
   const setSelectedChatId = useSetAtom(selectedChatIdAtom);
   const { invalidateChats } = useChats(appId);
+  const [imageGeneratorOpen, setImageGeneratorOpen] = useState(false);
+  const handleOpenImageGenerator = useCallback(() => {
+    setImageGeneratorOpen(true);
+  }, []);
+
+  // Image generation jobs for auto-adding to chat on send
+  const chatImageJobs = useAtomValue(chatImageGenerationJobsAtom);
+  const [dismissedImageJobIds, setDismissedImageJobIds] = useAtom(
+    dismissedImageGenerationJobIdsAtom,
+  );
+  const visibleSuccessfulImageJobs = useMemo(() => {
+    const appJobs = appId
+      ? chatImageJobs.filter((job) => job.targetAppId === appId)
+      : chatImageJobs;
+    return appJobs.filter(
+      (job) =>
+        !dismissedImageJobIds.has(job.id) &&
+        job.status === "success" &&
+        job.result,
+    );
+  }, [chatImageJobs, dismissedImageJobIds, appId]);
+  const hasSuccessfulImageJobs = visibleSuccessfulImageJobs.length > 0;
+
   // Use the attachments hook
   const {
     attachments,
@@ -405,7 +434,9 @@ export function ChatInput({ chatId }: { chatId?: number }) {
 
   const handleSubmit = async () => {
     if (
-      (!inputValue.trim() && attachments.length === 0) ||
+      (!inputValue.trim() &&
+        attachments.length === 0 &&
+        !hasSuccessfulImageJobs) ||
       !chatId ||
       pendingFiles
     ) {
@@ -416,10 +447,30 @@ export function ChatInput({ chatId }: { chatId?: number }) {
       await toggleRecording();
     }
 
+    // Build prompt with auto-added image mentions
+    const imageMentions = visibleSuccessfulImageJobs
+      .map((job) => `@media:${encodeURIComponent(job.result!.fileName)}`)
+      .join(" ");
+    const promptWithImages = inputValue.trim()
+      ? imageMentions
+        ? `${inputValue} ${imageMentions}`
+        : inputValue
+      : imageMentions;
+
+    // Dismiss image jobs that were auto-added
+    if (visibleSuccessfulImageJobs.length > 0) {
+      setDismissedImageJobIds((prev) => {
+        const next = new Set(prev);
+        for (const job of visibleSuccessfulImageJobs) {
+          next.add(job.id);
+        }
+        return next;
+      });
+    }
+
     // If switching to plan mode from another mode in a chat with messages,
     // create a new chat for a clean context.
     if (needsFreshPlanChat && settings?.selectedChatMode === "plan" && appId) {
-      const currentInput = inputValue;
       setInputValue("");
       setNeedsFreshPlanChat(false);
 
@@ -430,7 +481,7 @@ export function ChatInput({ chatId }: { chatId?: number }) {
       showInfo("We've switched you to a new chat for a clean context");
 
       await streamMessage({
-        prompt: currentInput,
+        prompt: promptWithImages,
         chatId: newChatId,
         attachments,
         redo: false,
@@ -440,7 +491,7 @@ export function ChatInput({ chatId }: { chatId?: number }) {
       return;
     }
 
-    const currentInput = inputValue;
+    const currentInput = promptWithImages;
 
     // Use all selected components for multi-component editing
     const componentsToSend =
@@ -800,6 +851,11 @@ export function ChatInput({ chatId }: { chatId?: number }) {
             onRemove={removeAttachment}
           />
 
+          {/* Chat image generation strip */}
+          <ChatImageGenerationStrip
+            onGenerateImage={handleOpenImageGenerator}
+          />
+
           {/* Use the DragDropOverlay component */}
           <DragDropOverlay isDraggingOver={isDraggingOver} />
 
@@ -906,7 +962,9 @@ export function ChatInput({ chatId }: { chatId?: number }) {
                     <button
                       onClick={handleSubmit}
                       disabled={
-                        (!inputValue.trim() && attachments.length === 0) ||
+                        (!inputValue.trim() &&
+                          attachments.length === 0 &&
+                          !hasSuccessfulImageJobs) ||
                         disableSendButton
                       }
                       aria-label={t("sendMessage")}
@@ -930,12 +988,21 @@ export function ChatInput({ chatId }: { chatId?: number }) {
               showTokenBar={showTokenBar}
               toggleShowTokenBar={toggleShowTokenBar}
               appId={appId ?? undefined}
+              onGenerateImage={handleOpenImageGenerator}
             />
           </div>
           {/* TokenBar is only displayed when showTokenBar is true */}
           {showTokenBar && <TokenBar chatId={chatId} />}
         </div>
       </div>
+
+      {/* Image Generator Dialog */}
+      <ImageGeneratorDialog
+        open={imageGeneratorOpen}
+        onOpenChange={setImageGeneratorOpen}
+        defaultAppId={appId ?? undefined}
+        source="chat"
+      />
     </>
   );
 }

src/hooks/useGenerateImage.ts

@@ -20,6 +20,7 @@ interface StartGenerationParams {
   themeMode: ImageThemeMode;
   targetAppId: number;
   targetAppName: string;
+  source?: "chat" | "media-library";
 }
 
 // Track cancelled job IDs so onError can skip them when the abort error arrives.
@@ -69,6 +70,7 @@ export function useGenerateImage() {
         targetAppName: params.targetAppName,
         status: "pending",
         startedAt: Date.now(),
+        source: params.source,
       });
 
       // Show / update the single generating toast with the new pending count