feat: add image generation tool and WebUI mode

Co-authored-by: Cursor <cursoragent@cursor.com>

webui/src/components/thread/ThreadShell.tsx

@@ -4,9 +4,12 @@ import {
   BookOpen,
   ChevronRight,
   Code2,
+  ImageIcon,
   LayoutGrid,
   Lightbulb,
   MoreHorizontal,
+  Palette,
+  Sparkles,
 } from "lucide-react";
 import { useTranslation } from "react-i18next";
 
@@ -15,7 +18,7 @@ import { ThreadComposer } from "@/components/thread/ThreadComposer";
 import { ThreadHeader } from "@/components/thread/ThreadHeader";
 import { StreamErrorNotice } from "@/components/thread/StreamErrorNotice";
 import { ThreadViewport } from "@/components/thread/ThreadViewport";
-import { useNanobotStream } from "@/hooks/useNanobotStream";
+import { useNanobotStream, type SendImage, type SendOptions } from "@/hooks/useNanobotStream";
 import { useSessionHistory } from "@/hooks/useSessions";
 import { listSlashCommands } from "@/lib/api";
 import type { ChatSummary, SlashCommand, UIMessage } from "@/lib/types";
@@ -52,6 +55,21 @@ const QUICK_ACTION_KEYS = [
   { key: "more", icon: MoreHorizontal, tone: "text-muted-foreground/65" },
 ] as const;
 
+const IMAGE_QUICK_ACTION_KEYS = [
+  { key: "icon", icon: ImageIcon, tone: "text-[#4f9de8]" },
+  { key: "sticker", icon: Sparkles, tone: "text-[#f25b8f]" },
+  { key: "poster", icon: Palette, tone: "text-[#eba45d]" },
+  { key: "product", icon: LayoutGrid, tone: "text-[#53c59d]" },
+  { key: "portrait", icon: ImageIcon, tone: "text-[#a877e7]" },
+  { key: "edit", icon: MoreHorizontal, tone: "text-muted-foreground/65" },
+] as const;
+
+interface PendingFirstMessage {
+  content: string;
+  images?: SendImage[];
+  options?: SendOptions;
+}
+
 export function ThreadShell({
   session,
   title,
@@ -67,10 +85,11 @@ export function ThreadShell({
   const chatId = session?.chatId ?? null;
   const historyKey = session?.key ?? null;
   const { messages: historical, loading, hasPendingToolCalls } = useSessionHistory(historyKey);
-  const { client, modelName, token } = useClient();
+  const { modelName, token } = useClient();
   const [booting, setBooting] = useState(false);
   const [slashCommands, setSlashCommands] = useState<SlashCommand[]>([]);
-  const pendingFirstRef = useRef<string | null>(null);
+  const [heroImageMode, setHeroImageMode] = useState(false);
+  const pendingFirstRef = useRef<PendingFirstMessage | null>(null);
   const messageCacheRef = useRef<Map<string, UIMessage[]>>(new Map());
   const lastCachedChatIdRef = useRef<string | null>(null);
 
@@ -82,6 +101,7 @@ export function ThreadShell({
     messages,
     isStreaming,
     send,
+    stop,
     setMessages,
     streamError,
     dismissStreamError,
@@ -109,7 +129,11 @@ export function ThreadShell({
     // When the user switches away and back, keep the local in-memory thread
     // state (including not-yet-persisted messages) instead of replacing it with
     // whatever the history endpoint currently knows about.
-    setMessages(cached && cached.length > 0 ? cached : historical);
+    setMessages((prev) => {
+      if (cached && cached.length > 0) return cached;
+      if (historical.length === 0 && prev.length > 0) return prev;
+      return historical;
+    });
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [loading, chatId, historical]);
 
@@ -142,18 +166,9 @@ export function ThreadShell({
     const pending = pendingFirstRef.current;
     if (!pending) return;
     pendingFirstRef.current = null;
-    client.sendMessage(chatId, pending);
-    setMessages((prev) => [
-      ...prev,
-      {
-        id: crypto.randomUUID(),
-        role: "user",
-        content: pending,
-        createdAt: Date.now(),
-      },
-    ]);
+    send(pending.content, pending.images, pending.options);
     setBooting(false);
-  }, [chatId, client, setMessages]);
+  }, [chatId, send]);
 
   useEffect(() => {
     let cancelled = false;
@@ -171,10 +186,10 @@ export function ThreadShell({
   }, [token]);
 
   const handleWelcomeSend = useCallback(
-    async (content: string) => {
+    async (content: string, images?: SendImage[], options?: SendOptions) => {
       if (booting) return;
       setBooting(true);
-      pendingFirstRef.current = content;
+      pendingFirstRef.current = { content, images, options };
       const newId = await onCreateChat?.();
       if (!newId) {
         pendingFirstRef.current = null;
@@ -186,20 +201,27 @@ export function ThreadShell({
 
   const handleQuickAction = useCallback(
     (prompt: string) => {
+      const options: SendOptions | undefined = heroImageMode
+        ? { imageGeneration: { enabled: true, aspect_ratio: null } }
+        : undefined;
       if (session) {
-        send(prompt);
+        send(prompt, undefined, options);
         return;
       }
-      void handleWelcomeSend(prompt);
+      void handleWelcomeSend(prompt, undefined, options);
     },
-    [handleWelcomeSend, send, session],
+    [handleWelcomeSend, heroImageMode, send, session],
   );
 
+  const quickActionItems = heroImageMode ? IMAGE_QUICK_ACTION_KEYS : QUICK_ACTION_KEYS;
+  const quickActionPrefix = heroImageMode
+    ? "thread.empty.imageQuickActions"
+    : "thread.empty.quickActions";
   const quickActions = (
     <div className="mx-auto grid w-full max-w-[58rem] grid-cols-2 gap-3 pt-4 sm:grid-cols-3 lg:grid-cols-6 lg:gap-4">
-      {QUICK_ACTION_KEYS.map(({ key, icon: Icon, tone }) => {
-        const title = t(`thread.empty.quickActions.${key}.title`);
-        const prompt = t(`thread.empty.quickActions.${key}.prompt`);
+      {quickActionItems.map(({ key, icon: Icon, tone }) => {
+        const title = t(`${quickActionPrefix}.${key}.title`);
+        const prompt = t(`${quickActionPrefix}.${key}.prompt`);
         return (
           <button
             key={key}
@@ -247,6 +269,9 @@ export function ThreadShell({
           modelLabel={toModelBadgeLabel(modelName)}
           variant={showHeroComposer ? "hero" : "thread"}
           slashCommands={slashCommands}
+          imageMode={showHeroComposer ? heroImageMode : undefined}
+          onImageModeChange={showHeroComposer ? setHeroImageMode : undefined}
+          onStop={stop}
         />
       ) : (
         <ThreadComposer
@@ -260,6 +285,8 @@ export function ThreadShell({
           }
           modelLabel={toModelBadgeLabel(modelName)}
           variant="hero"
+          imageMode={heroImageMode}
+          onImageModeChange={setHeroImageMode}
         />
       )}
       {showHeroComposer ? quickActions : null}