From b1ef65c07a252e0aabe0986e596b24cf03971721 Mon Sep 17 00:00:00 2001
From: Wasim Yousef Said <wasimysdev@gmail.com>
Date: Tue, 26 May 2026 13:17:43 +0200
Subject: [PATCH] Improve image generation UI (#5784)

* Improve image generation UI

* Polish generated image edit UI

* Tune generated image UI polish

* Soften generated image UI

* Refine generated image loading surface

* Align generated image caption clamp
---
 .../src/components/assistant-ui/thread.tsx    |  18 +--
 .../assistant-ui/tool-ui-image-generation.tsx | 108 ++++++++++++++++--
 .../src/features/chat/api/chat-adapter.ts     |  21 +++-
 .../src/features/chat/shared-composer.tsx     |  21 +++-
 4 files changed, 143 insertions(+), 25 deletions(-)

diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx
index 6a99f30508..da243f37e1 100644
--- a/studio/frontend/src/components/assistant-ui/thread.tsx
+++ b/studio/frontend/src/components/assistant-ui/thread.tsx
@@ -75,7 +75,6 @@ import {
   DownloadIcon,
   GlobeIcon,
   HeadphonesIcon,
-  ImageIcon,
   LightbulbIcon,
   LightbulbOffIcon,
   MicIcon,
@@ -89,6 +88,7 @@ import {
   Copy01Icon,
   Delete02Icon,
   Edit03Icon,
+  Image03Icon,
   Tick02Icon,
 } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
@@ -271,17 +271,17 @@ const GeneratedImageViewportOverlay: FC<{ hideComposer?: boolean }> = ({
             className="w-full max-w-[min(100%,46rem)] shrink-0 text-center"
             title={overlay.title}
           >
-            <p className="truncate font-medium text-foreground/70 text-xs">
+            <p className="truncate text-xs font-semibold text-foreground/80">
               Generated image
             </p>
             {overlay.metadata ? (
-              <p className="truncate text-[11px] text-muted-foreground/75">
+              <p className="truncate text-[11px] font-medium text-muted-foreground">
                 {overlay.metadata}
               </p>
             ) : null}
             {hideComposer ? null : (
-              <p className="mt-1 text-muted-foreground text-xs">
-                Type edits below, then send.
+              <p className="mx-auto mt-2 inline-flex rounded-full bg-primary/10 px-3 py-1 text-xs font-medium text-primary">
+                Type edits below, then send
               </p>
             )}
           </div>
@@ -525,13 +525,15 @@ const Composer: FC<{
       <PendingAudioChip />
       <ToolStatusDisplay />
       <ComposerPrimitive.Input
-        placeholder="Send a message..."
+        placeholder={
+          overlay ? "Type your edits for your image" : "Send a message..."
+        }
         className="aui-composer-input composer-input"
         minRows={1}
         maxRows={12}
         autoFocus={!disabled}
         disabled={disabled}
-        aria-label="Message input"
+        aria-label={overlay ? "Image edit instructions" : "Message input"}
         // dir="auto": browser picks LTR/RTL from the first strong char;
         // no effect on Latin / CJK / Devanagari.
         dir="auto"
@@ -1107,7 +1109,7 @@ const ImagesToggle: FC = () => {
           : "Enable image generation"
       }
     >
-      <ImageIcon className="size-3.5" />
+      <HugeiconsIcon icon={Image03Icon} className="size-3.5" strokeWidth={2} />
       <span>Images</span>
     </button>
   );
diff --git a/studio/frontend/src/components/assistant-ui/tool-ui-image-generation.tsx b/studio/frontend/src/components/assistant-ui/tool-ui-image-generation.tsx
index e32d23acf6..7dfdd903fe 100644
--- a/studio/frontend/src/components/assistant-ui/tool-ui-image-generation.tsx
+++ b/studio/frontend/src/components/assistant-ui/tool-ui-image-generation.tsx
@@ -8,7 +8,7 @@ import { cn } from "@/lib/utils";
 import type { ToolCallMessagePartComponent } from "@assistant-ui/react";
 import { DownloadIcon, ImageIcon, PencilIcon } from "lucide-react";
 import type { CSSProperties, MouseEvent } from "react";
-import { memo, useState } from "react";
+import { memo, useCallback, useEffect, useRef, useState } from "react";
 import { useGeneratedImageOverlay } from "./generated-image-overlay-context";
 import { Image, downloadImagePart } from "./image";
 import {
@@ -63,6 +63,8 @@ type GeneratedImagePart = {
   filename?: string;
 };
 
+const CAPTION_COLLAPSED_LINES = 4;
+
 const extensionForMime = (mime: string): string => {
   switch (mime.toLowerCase()) {
     case "image/jpeg":
@@ -119,7 +121,7 @@ function GeneratedImagePlaceholder({ label }: { label: string }) {
   return (
     <div
       className={cn(
-        "generated-image-loading-card flex aspect-square w-[480px] max-w-full items-center justify-center rounded-2xl border border-border/70 bg-muted/15",
+        "generated-image-loading-card flex aspect-square w-[480px] max-w-full items-center justify-center rounded-2xl bg-muted/20 shadow-[0_0_12px_rgba(15,23,42,0.05),0_6px_18px_rgba(15,23,42,0.04)] dark:shadow-[0_0_12px_rgba(0,0,0,0.18),0_6px_18px_rgba(0,0,0,0.12)]",
       )}
       aria-busy="true"
       aria-label={label}
@@ -154,6 +156,8 @@ const ImageGenerationToolUIImpl: ToolCallMessagePartComponent = ({
     : null;
   const imageTitle =
     imageResult?.prompt?.trim() || prompt.trim() || "Generated image";
+  const captionPrompt = imageResult?.prompt?.trim() || prompt.trim();
+  const promptLikelyNeedsExpansion = captionPrompt.length > 220;
   const imageMetadata = [imageResult?.size, imageResult?.quality, mime]
     .filter(Boolean)
     .join(" · ");
@@ -174,8 +178,63 @@ const ImageGenerationToolUIImpl: ToolCallMessagePartComponent = ({
     : null;
 
   const [open, setOpen] = useState(true);
+  const [expandedCaptionPrompt, setExpandedCaptionPrompt] = useState<
+    string | null
+  >(null);
+  const [promptOverflow, setPromptOverflow] = useState<{
+    prompt: string;
+    canExpand: boolean;
+  } | null>(null);
+  const captionRef = useRef<HTMLDivElement | null>(null);
   const isPendingImage = !imagePart && status?.type === "running";
 
+  const promptOverflowMeasured = promptOverflow?.prompt === captionPrompt;
+  const promptCanExpand = promptOverflowMeasured
+    ? promptOverflow.canExpand
+    : false;
+  const promptExpanded = expandedCaptionPrompt === captionPrompt;
+
+  const updatePromptOverflow = useCallback(() => {
+    const captionElement = captionRef.current;
+    if (!captionElement || !captionPrompt) {
+      return;
+    }
+    const computedStyle = window.getComputedStyle(captionElement);
+    const lineHeight = Number.parseFloat(computedStyle.lineHeight);
+    const collapsedHeight =
+      (Number.isFinite(lineHeight) ? lineHeight : 20) *
+      CAPTION_COLLAPSED_LINES;
+    const hasOverflow = captionElement.scrollHeight > collapsedHeight + 1;
+    setPromptOverflow((current) =>
+      current?.prompt === captionPrompt && current.canExpand === hasOverflow
+        ? current
+        : { prompt: captionPrompt, canExpand: hasOverflow },
+    );
+  }, [captionPrompt]);
+
+  useEffect(() => {
+    const captionElement = captionRef.current;
+    if (!captionElement || !captionPrompt) {
+      return;
+    }
+    const frame = window.requestAnimationFrame(updatePromptOverflow);
+    const resizeObserver =
+      typeof ResizeObserver === "undefined"
+        ? null
+        : new ResizeObserver(updatePromptOverflow);
+    resizeObserver?.observe(captionElement);
+    window.addEventListener("resize", updatePromptOverflow);
+    return () => {
+      window.cancelAnimationFrame(frame);
+      resizeObserver?.disconnect();
+      window.removeEventListener("resize", updatePromptOverflow);
+    };
+  }, [captionPrompt, updatePromptOverflow]);
+
+  const shouldClampPrompt =
+    (promptOverflowMeasured ? promptCanExpand : promptLikelyNeedsExpansion) &&
+    !promptExpanded;
+
   const runningLabel = "Generating image…";
   const completedLabel = formatGeneratedImageLabel(prompt);
 
@@ -231,21 +290,28 @@ const ImageGenerationToolUIImpl: ToolCallMessagePartComponent = ({
       <ToolFallbackContent>
         {imagePart ? (
           <figure className="m-0 flex flex-col gap-2">
-            <div className="group/generated-image relative aspect-square w-[480px] max-w-full overflow-hidden rounded-2xl border border-border/70 bg-muted/30 shadow-sm">
+            <div className="group/generated-image relative aspect-square w-[480px] max-w-full overflow-hidden rounded-2xl bg-muted/25 shadow-lg shadow-foreground/5 dark:shadow-black/25">
+              <img
+                src={imagePart.image}
+                alt=""
+                aria-hidden={true}
+                className="pointer-events-none absolute inset-0 size-full scale-110 object-cover opacity-25 blur-2xl saturate-125"
+              />
+              <div className="pointer-events-none absolute inset-0 bg-background/45" />
               <button
                 type="button"
-                className="block size-full cursor-zoom-in overflow-hidden rounded-2xl focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-2 focus-visible:ring-offset-background"
+                className="relative z-10 block size-full cursor-zoom-in overflow-hidden rounded-2xl focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-2 focus-visible:ring-offset-background"
                 onClick={showPreview}
                 aria-label="Open generated image preview"
               >
                 <Image.Preview
                   src={imagePart.image}
                   alt={imageTitle}
-                  containerClassName="size-full min-h-0 bg-background"
-                  className="size-full object-cover"
+                  containerClassName="flex size-full min-h-0 items-center justify-center bg-transparent"
+                  className="size-full object-contain"
                 />
               </button>
-              <div className="pointer-events-none absolute inset-x-0 bottom-0 flex items-end justify-between gap-2 bg-gradient-to-t from-black/55 via-black/20 to-transparent p-3 opacity-100 transition-opacity sm:opacity-0 sm:group-hover/generated-image:opacity-100 sm:group-focus-within/generated-image:opacity-100">
+              <div className="pointer-events-none absolute inset-x-0 bottom-0 z-20 flex items-end justify-between gap-2 bg-gradient-to-t from-black/55 via-black/20 to-transparent p-3 opacity-100 transition-opacity sm:opacity-0 sm:group-hover/generated-image:opacity-100 sm:group-focus-within/generated-image:opacity-100">
                 <Button
                   type="button"
                   variant="dark"
@@ -268,9 +334,31 @@ const ImageGenerationToolUIImpl: ToolCallMessagePartComponent = ({
                 </Button>
               </div>
             </div>
-            {prompt ? (
-              <figcaption className="max-w-[480px] text-xs leading-snug text-muted-foreground">
-                {prompt}
+            {captionPrompt ? (
+              <figcaption className="max-w-[480px] text-xs leading-5 text-muted-foreground">
+                <div
+                  ref={captionRef}
+                  className={cn(
+                    "whitespace-pre-wrap break-words",
+                    shouldClampPrompt && "max-h-20 overflow-hidden",
+                  )}
+                >
+                  {captionPrompt}
+                </div>
+                {promptCanExpand ? (
+                  <button
+                    type="button"
+                    className="mt-2 inline-flex text-xs font-medium text-foreground/80 underline-offset-4 hover:text-foreground hover:underline focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-2 focus-visible:ring-offset-background"
+                    onClick={() =>
+                      setExpandedCaptionPrompt((value) =>
+                        value === captionPrompt ? null : captionPrompt,
+                      )
+                    }
+                    aria-expanded={promptExpanded}
+                  >
+                    {promptExpanded ? "Show less" : "Show more"}
+                  </button>
+                ) : null}
               </figcaption>
             ) : null}
           </figure>
diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts
index 49a5eebd6b..49622c8090 100644
--- a/studio/frontend/src/features/chat/api/chat-adapter.ts
+++ b/studio/frontend/src/features/chat/api/chat-adapter.ts
@@ -1433,6 +1433,17 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
       // Tool call content parts — accumulated and yielded cumulatively.
       // result is set directly on the tool-call part when tool_end arrives.
       const toolCallParts: ToolCallMessagePart[] = [];
+      const orderAssistantContent = (
+        textParts: ReturnType<typeof parseAssistantContent>,
+      ) => {
+        const imageToolParts = toolCallParts.filter(
+          (part) => part.toolName === "image_generation",
+        );
+        const otherToolParts = toolCallParts.filter(
+          (part) => part.toolName !== "image_generation",
+        );
+        return [...otherToolParts, ...textParts, ...imageToolParts];
+      };
       // Anthropic document_citations tool_event payload, converted to
       // Sources-panel source parts at end-of-stream so the inline [N]
       // markers have matching entries.
@@ -2036,10 +2047,11 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
                     };
                   }
                 }
-                // Yield cumulative state so tool UI updates (tools first, text after)
+                // Yield cumulative state so tool UI updates. Search/code tools stay
+                // before the text, while generated images sit after the answer.
                 const textParts = parseAssistantContent(cumulativeText);
                 yield {
-                  content: [...toolCallParts, ...textParts],
+                  content: orderAssistantContent(textParts),
                   metadata: {
                     timing: buildTiming(
                       streamStartTime,
@@ -2187,7 +2199,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
 
               if (parts.length > 0 || toolCallParts.length > 0) {
                 yield {
-                  content: [...toolCallParts, ...parts],
+                  content: orderAssistantContent(parts),
                   metadata: {
                     timing: buildTiming(
                       streamStartTime,
@@ -2283,8 +2295,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
 
         yield {
           content: [
-            ...toolCallParts,
-            ...parseAssistantContent(cumulativeText),
+            ...orderAssistantContent(parseAssistantContent(cumulativeText)),
             ...sourceParts,
             ...documentCitationParts,
           ],
diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx
index 76e77d1288..ba4a897f63 100644
--- a/studio/frontend/src/features/chat/shared-composer.tsx
+++ b/studio/frontend/src/features/chat/shared-composer.tsx
@@ -21,7 +21,20 @@ import { isTauri } from "@/lib/api-base";
 import { isMultimodalResponse } from "./types/api";
 import { getImageInputUnavailableReason } from "./utils/image-input-support";
 import { useAui } from "@assistant-ui/react";
-import { ArrowUpIcon, DownloadIcon, GlobeIcon, HeadphonesIcon, ImageIcon, LightbulbIcon, LightbulbOffIcon, MicIcon, PlusIcon, SquareIcon, XIcon } from "lucide-react";
+import {
+  ArrowUpIcon,
+  DownloadIcon,
+  GlobeIcon,
+  HeadphonesIcon,
+  LightbulbIcon,
+  LightbulbOffIcon,
+  MicIcon,
+  PlusIcon,
+  SquareIcon,
+  XIcon,
+} from "lucide-react";
+import { Image03Icon } from "@hugeicons/core-free-icons";
+import { HugeiconsIcon } from "@hugeicons/react";
 import { toast } from "@/lib/toast";
 import { loadModel, validateModel } from "./api/chat-api";
 import { parseExternalModelId, providerTypeSupportsVision } from "./external-providers";
@@ -1115,7 +1128,11 @@ export function SharedComposer({
                 imageToolsEnabled ? "Disable image generation" : "Enable image generation"
               }
             >
-              <ImageIcon className="size-3.5" />
+              <HugeiconsIcon
+                icon={Image03Icon}
+                className="size-3.5"
+                strokeWidth={2}
+              />
               <span>Images</span>
             </button>
           )}