From b384f92bbfb0ee75250ddefc4b112c6c6885da9f Mon Sep 17 00:00:00 2001
From: ZhenYi <434836402@qq.com>
Date: Tue, 12 May 2026 13:05:20 +0800
Subject: [PATCH] refactor(chat): update frontend chat components

---
 src/app/chat/ChatHeader.tsx        |   8 --
 src/app/chat/ChatMessageBubble.tsx | 102 ++++++++---------
 src/app/chat/ChatMessageInput.tsx  | 121 +++++---------------
 src/app/chat/ChatMessageList.tsx   | 146 +++++++++++------------
 src/app/chat/ChatModelSelector.tsx | 178 +++++++++--------------------
 5 files changed, 197 insertions(+), 358 deletions(-)

diff --git a/src/app/chat/ChatHeader.tsx b/src/app/chat/ChatHeader.tsx
index dbe525e..c395d1a 100644
--- a/src/app/chat/ChatHeader.tsx
+++ b/src/app/chat/ChatHeader.tsx
@@ -1,7 +1,5 @@
 import {Loader2, MoreHorizontal, Pencil, PanelLeftOpen, PanelLeftClose, Share2} from "lucide-react";
 import {useConversationQuery} from "@/hooks/useAiChatQuery";
-import {useChatPage} from "./ChatPageContext";
-import {ChatModelSelector} from "./ChatModelSelector";
 
 interface ChatHeaderProps {
     conversationId: string | null;
@@ -12,7 +10,6 @@ interface ChatHeaderProps {
 
 export function ChatHeader({conversationId, isStreaming, isSidebarCollapsed, onToggleSidebar}: ChatHeaderProps) {
     const {data: conversation} = useConversationQuery(conversationId || "");
-    const {selectedModel, setSelectedModel} = useChatPage();
 
     const title = conversation?.title || "New Chat";
 
@@ -51,11 +48,6 @@ export function ChatHeader({conversationId, isStreaming, isSidebarCollapsed, onT
             </div>
 
             <div className="flex items-center gap-1 shrink-0">
-                <ChatModelSelector
-                    selectedModel={selectedModel}
-                    onSelect={setSelectedModel}
-                    conversationId={conversationId}
-                />
                 {conversationId && (
                     <>
                         <button
diff --git a/src/app/chat/ChatMessageBubble.tsx b/src/app/chat/ChatMessageBubble.tsx
index 9fe26f0..f153b10 100644
--- a/src/app/chat/ChatMessageBubble.tsx
+++ b/src/app/chat/ChatMessageBubble.tsx
@@ -4,7 +4,9 @@ import { useCurrentUserQuery } from "@/hooks/useAuth";
 import { useEditMessageMutation, useMessageVersionsQuery, useSwitchMessageVersionMutation } from "@/hooks/useAiChatQuery";
 import { MarkdownRenderer } from "@/components/ui/MarkdownRenderer";
 import { Avatar, AvatarFallback, AvatarImage } from "@/components/ui/avatar";
+import { Reasoning, ReasoningTrigger, ReasoningContent } from "@/components/ai-elements/reasoning";
 import type { MessageResponse } from "@/hooks/useAiChatQuery";
+import { getModelIcon } from "@/lib/icons/modelIcons";
 
 interface ChatMessageBubbleProps {
   message: MessageResponse;
@@ -97,49 +99,10 @@ function extractFullText(blocks: ContentBlock[]): string {
   }).join("\n\n");
 }
 
-/** DeepSeek-style collapsible thinking section with horizontal separator line. */
-function ThinkingSection({ content, isOpen, onToggle }: { content: string; isOpen: boolean; onToggle: () => void }) {
-  return (
-    <div className="mb-3">
-      {/* Horizontal separator with "Thinking" label */}
-      <div className="flex items-center gap-2 py-1.5 cursor-pointer select-none" onClick={onToggle}>
-        <div className="flex-1 h-px" style={{ backgroundColor: "var(--border-subtle)" }} />
-        <div className="flex items-center gap-1.5 shrink-0 px-1">
-          <svg
-            className={`w-3 h-3 transition-transform duration-200 ${isOpen ? "rotate-90" : ""}`}
-            fill="none" viewBox="0 0 24 24" stroke="currentColor" strokeWidth={2}
-            style={{ color: "var(--text-muted)" }}
-          >
-            <path strokeLinecap="round" strokeLinejoin="round" d="M9 5l7 7-7 7" />
-          </svg>
-          <span className="text-[12px] font-medium" style={{ color: "var(--text-muted)" }}>
-            {isOpen ? "Hide reasoning" : "Thinking"}
-          </span>
-        </div>
-        <div className="flex-1 h-px" style={{ backgroundColor: "var(--border-subtle)" }} />
-      </div>
-      {/* Collapsible content */}
-      {isOpen && (
-        <div
-          className="mt-1 p-3 rounded-xl text-sm leading-relaxed whitespace-pre-wrap max-h-[400px] overflow-y-auto"
-          style={{
-            backgroundColor: "var(--surface-elevated)",
-            border: "1px solid var(--border-subtle)",
-            color: "var(--text-muted)",
-            fontSize: "13px",
-          }}
-        >
-          {content}
-        </div>
-      )}
-    </div>
-  );
-}
 
 export const ChatMessageBubble = memo(function ChatMessageBubble({ message, conversationId, onRegenerate }: ChatMessageBubbleProps) {
   const isUser = message.role === "user";
   const [copied, setCopied] = useState<"answer" | "full" | false>(false);
-  const [openThinking, setOpenThinking] = useState<Record<number, boolean>>({});
   const [isEditing, setIsEditing] = useState(false);
   const [editText, setEditText] = useState("");
   const [showVersions, setShowVersions] = useState(false);
@@ -240,10 +203,6 @@ export const ChatMessageBubble = memo(function ChatMessageBubble({ message, conv
     }
   };
 
-  const toggleThinking = (idx: number) => {
-    setOpenThinking((prev) => ({ ...prev, [idx]: !prev[idx] }));
-  };
-
   return (
     <div className="flex gap-4 px-4 py-3 group max-w-3xl mx-auto w-full">
       {/* Avatar */}
@@ -262,12 +221,7 @@ export const ChatMessageBubble = memo(function ChatMessageBubble({ message, conv
             </AvatarFallback>
           </Avatar>
         ) : (
-          <div
-            className="w-7 h-7 rounded-full flex items-center justify-center"
-            style={{ backgroundColor: "var(--accent-muted)" }}
-          >
-            <Sparkles className="w-3.5 h-3.5" style={{ color: "var(--accent)" }} />
-          </div>
+          <ModelAvatar modelName={message.model} size={28} />
         )}
       </div>
 
@@ -276,7 +230,7 @@ export const ChatMessageBubble = memo(function ChatMessageBubble({ message, conv
         {/* Sender name */}
         <div className="flex items-center gap-2 mb-1">
           <span className="text-xs font-semibold" style={{ color: "var(--text-primary)" }}>
-            {isUser ? (user?.display_name || user?.username || "You") : "Assistant"}
+            {isUser ? (user?.display_name || user?.username || "You") : (message.model || "Assistant")}
           </span>
           <span className="text-[11px]" style={{ color: "var(--text-muted)" }}>
             {new Date(message.created_at).toLocaleTimeString([], { hour: "2-digit", minute: "2-digit" })}
@@ -337,12 +291,10 @@ export const ChatMessageBubble = memo(function ChatMessageBubble({ message, conv
             blocks.map((b, i) => {
               if (b.role === "thinking") {
                 return (
-                  <ThinkingSection
-                    key={i}
-                    content={b.content}
-                    isOpen={openThinking[i] ?? false}
-                    onToggle={() => toggleThinking(i)}
-                  />
+                  <Reasoning key={i} defaultOpen={false}>
+                    <ReasoningTrigger />
+                    <ReasoningContent>{b.content}</ReasoningContent>
+                  </Reasoning>
                 );
               }
               return (
@@ -472,4 +424,40 @@ export const ChatMessageBubble = memo(function ChatMessageBubble({ message, conv
       </div>
     </div>
   );
-});
\ No newline at end of file
+});
+
+// ─── Model Avatar ─────────────────────────────────────────────────────
+
+function ModelAvatar({ modelName, size = 28 }: { modelName?: string | null; size?: number }) {
+  if (!modelName) {
+    return (
+      <div
+        className="w-7 h-7 rounded-full flex items-center justify-center shrink-0"
+        style={{ width: size, height: size, backgroundColor: "var(--accent-muted)" }}
+      >
+        <Sparkles className="w-3.5 h-3.5" style={{ color: "var(--accent)", width: size * 0.5, height: size * 0.5 }} />
+      </div>
+    );
+  }
+  const Icon = getModelIcon(modelName);
+  if (Icon) {
+    return (
+      <div className="shrink-0" style={{ width: size, height: size }}>
+        <Icon.Avatar size={size} />
+      </div>
+    );
+  }
+  return (
+    <div
+      className="rounded-full flex items-center justify-center font-bold text-white shrink-0"
+      style={{
+        width: size,
+        height: size,
+        backgroundColor: hashColor(modelName),
+        fontSize: Math.max(10, size * 0.35),
+      }}
+    >
+      {modelName[0]?.toUpperCase() || "?"}
+    </div>
+  );
+}
\ No newline at end of file
diff --git a/src/app/chat/ChatMessageInput.tsx b/src/app/chat/ChatMessageInput.tsx
index faea829..21ec586 100644
--- a/src/app/chat/ChatMessageInput.tsx
+++ b/src/app/chat/ChatMessageInput.tsx
@@ -1,10 +1,19 @@
-import { useState, useRef } from "react";
-import { Square, ArrowUp, AlertCircle } from "lucide-react";
+import { useState } from "react";
+import { AlertCircle } from "lucide-react";
 import { useCreateMessageMutation, useCreateConversationMutation, streamChat } from "@/hooks/useAiChatQuery";
 import { useChatPage } from "./ChatPageContext";
 import { useQueryClient } from "@tanstack/react-query";
 import { useStreamingStore } from "@/store/streaming";
+import {
+  PromptInput,
+  PromptInputBody,
+  PromptInputTextarea,
+  PromptInputSubmit,
+  PromptInputFooter,
+} from "@/components/ai-elements/prompt-input";
 import type { CreateMessageParams } from "@/client/model";
+import type { PromptInputMessage } from "@/components/ai-elements/prompt-input";
+import { ChatModelSelector } from "./ChatModelSelector";
 
 interface ChatMessageInputProps {
   conversationId: string | null;
@@ -19,29 +28,23 @@ export function ChatMessageInput({
   setIsStreaming,
   onSelectConversation,
 }: ChatMessageInputProps) {
-  const [text, setText] = useState("");
-  const [isFocused, setIsFocused] = useState(false);
   const [showModelWarning, setShowModelWarning] = useState(false);
   const createMessageMutation = useCreateMessageMutation();
   const createConversationMutation = useCreateConversationMutation();
-  const { scope, projectId, selectedModel } = useChatPage();
+  const { scope, projectId, selectedModel, setSelectedModel } = useChatPage();
   const queryClient = useQueryClient();
   const streamingStore = useStreamingStore();
-  const textareaRef = useRef<HTMLTextAreaElement>(null);
 
-  const handleSend = async () => {
+  const handleSubmit = async ({ text }: PromptInputMessage) => {
     if (!text.trim()) return;
 
-    // Require model selection before sending
     if (!selectedModel) {
       setShowModelWarning(true);
-      // Auto-hide after 3 seconds
       setTimeout(() => setShowModelWarning(false), 3000);
       return;
     }
 
     const content = text.trim();
-    setText("");
 
     try {
       let activeConversationId = conversationId;
@@ -64,7 +67,7 @@ export function ChatMessageInput({
       const params: CreateMessageParams = {
         content: {
           role: "user",
-          content: content,
+          content,
         },
         model: selectedModel.model_name,
         parent_message_id: null,
@@ -80,7 +83,6 @@ export function ChatMessageInput({
 
       if (!messageResponse?.id) return;
 
-      // Clear any previous stream for this conversation
       streamingStore.clear(activeConversationId);
       setIsStreaming(true);
 
@@ -120,20 +122,6 @@ export function ChatMessageInput({
     }
   };
 
-  const handleKeyDown = (e: React.KeyboardEvent<HTMLTextAreaElement>) => {
-    if (e.key === "Enter" && !e.shiftKey) {
-      e.preventDefault();
-      handleSend();
-    }
-  };
-
-  const handleStop = () => {
-    setIsStreaming(false);
-  };
-
-  const hasText = text.trim().length > 0;
-  const canSend = hasText && selectedModel && !createConversationMutation.isPending;
-
   return (
     <div className="shrink-0 px-4 pb-4" style={{ backgroundColor: "var(--surface-ground)" }}>
       <div className="max-w-3xl mx-auto relative">
@@ -151,70 +139,23 @@ export function ChatMessageInput({
           </div>
         )}
 
-        <div
-          className="relative rounded-2xl transition-shadow duration-200"
-          style={{
-            backgroundColor: "var(--surface-ground)",
-            border: `1px solid ${isFocused ? "var(--accent)" : "var(--border-default)"}`,
-            boxShadow: isFocused
-              ? "0 4px 24px rgba(0,0,0,0.1), 0 1px 4px rgba(0,0,0,0.06)"
-              : "0 2px 8px rgba(0,0,0,0.04), 0 1px 2px rgba(0,0,0,0.03)",
-          }}
-        >
-          <textarea
-            ref={textareaRef}
-            value={text}
-            onChange={(e) => setText(e.target.value)}
-            onKeyDown={handleKeyDown}
-            onFocus={() => setIsFocused(true)}
-            onBlur={() => setIsFocused(false)}
-            placeholder="Message"
-            disabled={isStreaming}
-            rows={1}
-            className="w-full resize-none bg-transparent text-[15px] outline-none py-3.5 pl-4 pr-14"
-            style={{ color: "var(--text-primary)", lineHeight: "1.5" }}
-            onInput={(e) => {
-              const target = e.target as HTMLTextAreaElement;
-              target.style.height = "auto";
-              target.style.height = `${Math.min(target.scrollHeight, 200)}px`;
-            }}
-          />
-          <div className="absolute right-2 bottom-2">
-            {isStreaming ? (
-              <button
-                onClick={handleStop}
-                className="w-8 h-8 rounded-lg flex items-center justify-center transition-colors"
-                style={{ backgroundColor: "var(--text-primary)", color: "var(--text-inverse)" }}
-              >
-                <Square className="w-3.5 h-3.5" fill="currentColor" />
-              </button>
-            ) : (
-              <button
-                onClick={handleSend}
-                disabled={!canSend}
-                className="w-8 h-8 rounded-lg flex items-center justify-center transition-all duration-200"
-                style={{
-                  backgroundColor: canSend ? "var(--text-primary)" : "var(--muted)",
-                  color: canSend ? "var(--text-inverse)" : "var(--text-muted)",
-                  opacity: !canSend ? 0.4 : 1,
-                  transform: canSend ? "scale(1)" : "scale(0.95)",
-                }}
-              >
-                <ArrowUp className="w-4 h-4" strokeWidth={2.5} />
-              </button>
-            )}
-          </div>
-        </div>
-
-        <p className="text-[11px] text-center mt-2.5" style={{ color: "var(--text-muted)" }}>
-          {selectedModel?.model_name || (
-            <span className="font-medium" style={{ color: "var(--accent)" }}>
-              Please select a model to start chatting
-            </span>
-          )}
-          {selectedModel && " can make mistakes. Verify important info."}
-        </p>
+        <PromptInput onSubmit={handleSubmit}>
+          <PromptInputBody>
+            <PromptInputTextarea placeholder="Message" className="py-3" />
+          </PromptInputBody>
+          <PromptInputFooter>
+            <ChatModelSelector
+              selectedModel={selectedModel}
+              onSelect={setSelectedModel}
+              conversationId={conversationId}
+            />
+            <PromptInputSubmit
+              status={isStreaming ? "streaming" : "ready"}
+              onStop={() => setIsStreaming(false)}
+            />
+          </PromptInputFooter>
+        </PromptInput>
       </div>
     </div>
   );
-}
\ No newline at end of file
+}
diff --git a/src/app/chat/ChatMessageList.tsx b/src/app/chat/ChatMessageList.tsx
index fe48a46..6273da8 100644
--- a/src/app/chat/ChatMessageList.tsx
+++ b/src/app/chat/ChatMessageList.tsx
@@ -1,12 +1,16 @@
 import { useEffect, useRef, useState, useCallback } from "react";
-import { Loader2, Code, FileText, GitPullRequest, Brain, ChevronDown } from "lucide-react";
+import { Loader2, Code, FileText, GitPullRequest, Brain, ChevronDown, Sparkles } from "lucide-react";
 import { useVirtualizer } from "@tanstack/react-virtual";
 import { useQueryClient } from "@tanstack/react-query";
 import { useMessagesQuery } from "@/hooks/useAiChatQuery";
 import { useStreamingStore } from "@/store/streaming";
 import type { StreamPart } from "@/store/streaming";
 import { ChatMessageBubble } from "./ChatMessageBubble";
+import { useChatPage } from "./ChatPageContext";
+import { getModelIcon } from "@/lib/icons/modelIcons";
 import { MarkdownRenderer } from "@/components/ui/MarkdownRenderer";
+import { Shimmer } from "@/components/ai-elements/shimmer";
+import { Reasoning, ReasoningTrigger, ReasoningContent } from "@/components/ai-elements/reasoning";
 
 interface ChatMessageListProps {
   conversationId: string | null;
@@ -22,32 +26,6 @@ const PROMPT_SUGGESTIONS = [
 const OVERSCAN = 3;
 const ESTIMATED_SIZE = 200;
 
-/** Three-dot wave animation for streaming indicator */
-function WaveDots({ label }: { label?: string }) {
-  return (
-    <div className="flex items-center justify-center gap-2 py-1.5">
-      <div className="flex items-center gap-1">
-        {[0, 1, 2].map((i) => (
-          <span
-            key={i}
-            className="w-[5px] h-[5px] rounded-full inline-block animate-bounce"
-            style={{
-              backgroundColor: "var(--accent)",
-              animationDelay: `${i * 0.15}s`,
-              animationDuration: "0.8s",
-            }}
-          />
-        ))}
-      </div>
-      {label && (
-        <span className="text-[11px]" style={{ color: "var(--text-muted)" }}>
-          {label}
-        </span>
-      )}
-    </div>
-  );
-}
-
 export function ChatMessageList({ conversationId }: ChatMessageListProps) {
   const { data, isLoading } = useMessagesQuery(conversationId || "");
   const scrollRef = useRef<HTMLDivElement>(null);
@@ -64,7 +42,7 @@ export function ChatMessageList({ conversationId }: ChatMessageListProps) {
     const el = scrollRef.current;
     if (!el) return;
     const distance = el.scrollHeight - el.scrollTop - el.clientHeight;
-    const atBottom = distance < 80;
+    const atBottom = distance < 200;
     setIsAtBottom(atBottom);
     if (atBottom) setUserScrolledUp(false);
   }, []);
@@ -73,7 +51,7 @@ export function ChatMessageList({ conversationId }: ChatMessageListProps) {
     const el = scrollRef.current;
     if (!el) return;
     const distance = el.scrollHeight - el.scrollTop - el.clientHeight;
-    const atBottom = distance < 80;
+    const atBottom = distance < 200;
     setIsAtBottom(atBottom);
     if (!atBottom) setUserScrolledUp(true);
   }, []);
@@ -203,7 +181,7 @@ export function ChatMessageList({ conversationId }: ChatMessageListProps) {
           }}
         >
           <div className="flex items-center gap-2">
-            <WaveDots label="New response" />
+            <Shimmer duration={1}>New response</Shimmer>
             <ChevronDown className="w-3.5 h-3.5" style={{ color: "var(--text-muted)" }} />
           </div>
         </div>
@@ -212,7 +190,7 @@ export function ChatMessageList({ conversationId }: ChatMessageListProps) {
       {/* Virtualized message list — persisted messages only */}
       <div
         ref={scrollRef}
-        className="flex-1 overflow-y-auto"
+        className="flex-1 overflow-y-auto pb-24"
         onScroll={handleScroll}
         style={{ backgroundColor: "var(--surface-ground)" }}
       >
@@ -261,7 +239,7 @@ export function ChatMessageList({ conversationId }: ChatMessageListProps) {
 }
 
 function StreamingBubble({ parts, isDone }: { parts: StreamPart[]; isDone: boolean }) {
-  const [openThinking, setOpenThinking] = useState<Record<number, boolean>>({});
+  const { selectedModel } = useChatPage();
   // Display state synced at animation-frame rate so ReactMarkdown only
   // re-parses at ~60fps, not on every token from the SSE stream.
   const [displayParts, setDisplayParts] = useState<StreamPart[]>([]);
@@ -308,29 +286,18 @@ function StreamingBubble({ parts, isDone }: { parts: StreamPart[]; isDone: boole
     }
   }, [displayParts.length]);
 
-  const toggleThinking = (idx: number) => {
-    setOpenThinking((prev) => ({ ...prev, [idx]: !prev[idx] }));
-  };
-
   const firstThinkingIdx = displayParts.findIndex((p) => p.type === "thinking");
 
   return (
     <div ref={contentRef} className="flex gap-4 px-4 py-3 max-w-3xl mx-auto w-full">
-      {/* AI Avatar */}
+      {/* Model Avatar */}
       <div className="shrink-0 pt-0.5">
-        <div
-          className="w-7 h-7 rounded-full flex items-center justify-center"
-          style={{ backgroundColor: "var(--accent-muted)" }}
-        >
-          <svg width="14" height="14" viewBox="0 0 24 24" fill="none" stroke="currentColor" strokeWidth="2" style={{ color: "var(--accent)" }}>
-            <path d="M12 3l1.5 4.5L18 9l-4.5 1.5L12 15l-1.5-4.5L6 9l4.5-1.5L12 3z" />
-          </svg>
-        </div>
+        <StreamingModelAvatar modelName={selectedModel?.model_name} size={28} />
       </div>
       <div className="flex-1 min-w-0">
         <div className="flex items-center gap-2 mb-1">
           <span className="text-xs font-semibold" style={{ color: "var(--text-primary)" }}>
-            Assistant
+            {selectedModel?.model_name || "Assistant"}
           </span>
           {!displayDone && (
             <span className="text-[11px] animate-pulse" style={{ color: "var(--text-muted)" }}>
@@ -345,43 +312,12 @@ function StreamingBubble({ parts, isDone }: { parts: StreamPart[]; isDone: boole
         <div className="text-sm" style={{ color: "var(--text-primary)" }}>
           {displayParts.map((part, i) => {
             if (part.type === "thinking") {
-              const isOpen = openThinking[i] ?? false;
               const isActivelyThinking = !displayDone && i === firstThinkingIdx;
               return (
-                <div key={i} className="mb-3" style={{ whiteSpace: "normal" }}>
-                  <div
-                    className="flex items-center gap-2 py-1.5 cursor-pointer select-none"
-                    onClick={() => toggleThinking(i)}
-                  >
-                    <div className="flex-1 h-px" style={{ backgroundColor: "var(--border-subtle)" }} />
-                    <div className="flex items-center gap-1.5 shrink-0 px-1">
-                      <svg
-                        className={`w-3 h-3 transition-transform duration-200 ${isOpen ? "rotate-90" : ""}`}
-                        fill="none" viewBox="0 0 24 24" stroke="currentColor" strokeWidth={2}
-                        style={{ color: "var(--text-muted)" }}
-                      >
-                        <path strokeLinecap="round" strokeLinejoin="round" d="M9 5l7 7-7 7" />
-                      </svg>
-                      <span className="text-[12px] font-medium" style={{ color: "var(--text-muted)" }}>
-                        {isActivelyThinking ? "Thinking…" : isOpen ? "Hide reasoning" : "Thinking"}
-                      </span>
-                    </div>
-                    <div className="flex-1 h-px" style={{ backgroundColor: "var(--border-subtle)" }} />
-                  </div>
-                  {isOpen && (
-                    <div
-                      className="mt-1 p-3 rounded-xl text-sm leading-relaxed max-h-[400px] overflow-y-auto whitespace-pre-wrap"
-                      style={{
-                        backgroundColor: "var(--surface-elevated)",
-                        border: "1px solid var(--border-subtle)",
-                        color: "var(--text-muted)",
-                        fontSize: "13px",
-                      }}
-                    >
-                      {part.content}
-                    </div>
-                  )}
-                </div>
+                <Reasoning key={i} isStreaming={isActivelyThinking} defaultOpen={isActivelyThinking}>
+                  <ReasoningTrigger />
+                  <ReasoningContent>{part.content}</ReasoningContent>
+                </Reasoning>
               );
             }
             // Token content — rendered as full Markdown + safe HTML.
@@ -415,4 +351,52 @@ function StreamingCursor() {
       }}
     />
   );
+}
+
+const AVATAR_COLORS = [
+  "#6366f1", "#8b5cf6", "#d946ef", "#ec4899", "#f43f5e",
+  "#ef4444", "#f97316", "#eab308", "#22c55e", "#14b8a6",
+  "#06b6d4", "#3b82f6", "#2563eb", "#7c3aed", "#c026d3",
+];
+
+function hashColor(str: string): string {
+  let hash = 0;
+  for (let i = 0; i < str.length; i++) {
+    hash = str.charCodeAt(i) + ((hash << 5) - hash);
+  }
+  return AVATAR_COLORS[Math.abs(hash) % AVATAR_COLORS.length];
+}
+
+function StreamingModelAvatar({ modelName, size = 28 }: { modelName?: string | null; size?: number }) {
+  if (!modelName) {
+    return (
+      <div
+        className="rounded-full flex items-center justify-center shrink-0"
+        style={{ width: size, height: size, backgroundColor: "var(--accent-muted)" }}
+      >
+        <Sparkles className="w-3.5 h-3.5" style={{ color: "var(--accent)", width: size * 0.5, height: size * 0.5 }} />
+      </div>
+    );
+  }
+  const Icon = getModelIcon(modelName);
+  if (Icon) {
+    return (
+      <div className="shrink-0" style={{ width: size, height: size }}>
+        <Icon.Avatar size={size} />
+      </div>
+    );
+  }
+  return (
+    <div
+      className="rounded-full flex items-center justify-center font-bold text-white shrink-0"
+      style={{
+        width: size,
+        height: size,
+        backgroundColor: hashColor(modelName),
+        fontSize: Math.max(10, size * 0.35),
+      }}
+    >
+      {modelName[0]?.toUpperCase() || "?"}
+    </div>
+  );
 }
\ No newline at end of file
diff --git a/src/app/chat/ChatModelSelector.tsx b/src/app/chat/ChatModelSelector.tsx
index 27215e7..c858363 100644
--- a/src/app/chat/ChatModelSelector.tsx
+++ b/src/app/chat/ChatModelSelector.tsx
@@ -1,12 +1,23 @@
 import { useQueryClient } from "@tanstack/react-query";
-import { useState, useRef, useEffect } from "react";
-import { Loader2, Search, Check, ChevronDown, Cpu } from "lucide-react";
+import { useState } from "react";
+import { Check, ChevronDown, Cpu } from "lucide-react";
 import { useQuery } from "@tanstack/react-query";
 import { modelCatalog } from "@/client/api";
 import { updateConversation } from "@/client/aiChatApi";
 import type { ModelWithPricingResponse } from "@/client/model/modelWithPricingResponse";
 import { getModelIcon } from "@/lib/icons/modelIcons";
 import type { SelectedModel } from "./ChatPageContext";
+import {
+  ModelSelector,
+  ModelSelectorTrigger,
+  ModelSelectorContent,
+  ModelSelectorInput,
+  ModelSelectorList,
+  ModelSelectorGroup,
+  ModelSelectorItem,
+  ModelSelectorLogo,
+} from "@/components/ai-elements/model-selector";
+import type { ModelSelectorLogoProps } from "@/components/ai-elements/model-selector";
 
 const AVATAR_COLORS = [
   "#6366f1", "#8b5cf6", "#d946ef", "#ec4899", "#f43f5e",
@@ -46,14 +57,11 @@ interface ChatModelSelectorProps {
   selectedModel: SelectedModel | null;
   onSelect: (model: SelectedModel | null) => void;
   conversationId?: string | null;
-  /** Optional callback that receives the full model object for external use */
   onSelectModel?: (model: ModelWithPricingResponse) => void;
 }
 
 export function ChatModelSelector({ selectedModel, onSelect, conversationId, onSelectModel }: ChatModelSelectorProps) {
   const [open, setOpen] = useState(false);
-  const [search, setSearch] = useState("");
-  const containerRef = useRef<HTMLDivElement>(null);
   const queryClient = useQueryClient();
 
   const { data: catalog, isLoading } = useQuery({
@@ -65,33 +73,13 @@ export function ChatModelSelector({ selectedModel, onSelect, conversationId, onS
     enabled: open,
   });
 
-  // Close dropdown on click outside
-  useEffect(() => {
-    function handleClickOutside(e: MouseEvent) {
-      if (containerRef.current && !containerRef.current.contains(e.target as Node)) {
-        setOpen(false);
-        setSearch("");
-      }
-    }
-    document.addEventListener("mousedown", handleClickOutside);
-    return () => document.removeEventListener("mousedown", handleClickOutside);
-  }, []);
-
   const models = Array.isArray(catalog) ? catalog : [];
-  const filtered = search.trim()
-    ? models.filter(
-        (m) =>
-          m.name.toLowerCase().includes(search.toLowerCase()) ||
-          m.provider_id.toLowerCase().includes(search.toLowerCase())
-      )
-    : models;
 
   const handleSelect = async (model: ModelWithPricingResponse) => {
     const selected = { model_name: model.name };
     onSelect(selected);
     onSelectModel?.(model);
     setOpen(false);
-    setSearch("");
     if (conversationId) {
       try {
         await updateConversation(conversationId, {
@@ -105,106 +93,52 @@ export function ChatModelSelector({ selectedModel, onSelect, conversationId, onS
   };
 
   return (
-    <div ref={containerRef} className="relative">
-      <button
-        onClick={() => setOpen(!open)}
-        className="flex items-center gap-1.5 px-2.5 py-1 rounded-lg text-xs transition-colors cursor-pointer"
-        style={{
-          backgroundColor: open ? "var(--hover-bg-strong)" : "transparent",
-          color: "var(--text-secondary)",
-        }}
-        onMouseEnter={(e) => {
-          if (!open) e.currentTarget.style.backgroundColor = "var(--hover-bg)";
-        }}
-        onMouseLeave={(e) => {
-          if (!open) e.currentTarget.style.backgroundColor = "transparent";
-        }}
-      >
-        {selectedModel ? (
-          <ModelAvatar modelName={selectedModel.model_name} size={18} />
-        ) : (
-          <Cpu className="w-4 h-4" style={{ color: "var(--text-muted)" }} />
-        )}
-        <span className="max-w-[140px] truncate font-medium">
-          {selectedModel?.model_name || "Select Model"}
-        </span>
-        <ChevronDown
-          className={`w-3 h-3 shrink-0 transition-transform duration-200 ${
-            open ? "rotate-180" : ""
-          }`}
-          style={{ color: "var(--text-muted)" }}
-        />
-      </button>
-
-      {open && (
-        <div
-          className="absolute top-full right-0 mt-1.5 rounded-xl overflow-hidden z-50 w-80"
-          style={{
-            backgroundColor: "var(--surface-elevated)",
-            border: "1px solid var(--border-subtle)",
-            boxShadow: "0 8px 32px rgba(0,0,0,0.12)",
-          }}
+    <ModelSelector open={open} onOpenChange={setOpen}>
+      <ModelSelectorTrigger asChild>
+        <button
+          className="flex items-center gap-1.5 px-2.5 py-1 rounded-lg text-xs transition-colors cursor-pointer hover:bg-accent"
         >
-          {/* Search */}
-          <div className="p-3 border-b border-[var(--border-subtle)]">
-            <div className="relative">
-              <Search
-                className="absolute left-2.5 top-1/2 -translate-y-1/2 w-3.5 h-3.5"
-                style={{ color: "var(--text-muted)" }}
-              />
-              <input
-                type="text"
-                placeholder="Search models..."
-                value={search}
-                onChange={(e) => setSearch(e.target.value)}
-                className="w-full h-9 pl-8 pr-3 text-sm rounded-lg outline-none transition-colors"
-                style={{
-                  backgroundColor: "var(--input-bg)",
-                  color: "var(--text-primary)",
-                  border: "1px solid var(--border-subtle)",
-                }}
-              />
+          {selectedModel ? (
+            <ModelAvatar modelName={selectedModel.model_name} size={18} />
+          ) : (
+            <Cpu className="w-4 h-4 text-muted-foreground" />
+          )}
+          <span className="max-w-[140px] truncate font-medium">
+            {selectedModel?.model_name || "Select Model"}
+          </span>
+          <ChevronDown className="w-3 h-3 shrink-0 text-muted-foreground" />
+        </button>
+      </ModelSelectorTrigger>
+      <ModelSelectorContent>
+        <ModelSelectorInput placeholder="Search models..." />
+        <ModelSelectorList>
+          {isLoading ? (
+            <div className="flex justify-center py-6 text-sm text-muted-foreground">
+              Loading...
             </div>
-          </div>
-
-          {/* List */}
-          <div className="max-h-[320px] overflow-y-auto py-1">
-            {isLoading ? (
-              <div className="flex justify-center py-6">
-                <Loader2 className="w-4 h-4 animate-spin" style={{ color: "var(--text-muted)" }} />
-              </div>
-            ) : filtered.length === 0 ? (
-              <div className="text-center py-6 text-xs" style={{ color: "var(--text-muted)" }}>
-                No models found
-              </div>
-            ) : (
-              filtered.map((model) => (
-                <button
+          ) : models.length === 0 ? (
+            <div className="text-center py-6 text-sm text-muted-foreground">
+              No models found
+            </div>
+          ) : (
+            <ModelSelectorGroup>
+              {models.map((model) => (
+                <ModelSelectorItem
                   key={model.id}
-                  onClick={() => handleSelect(model)}
-                  className="w-full text-left px-3 py-2.5 transition-colors flex items-center gap-3 text-sm hover:bg-[var(--hover-bg)]"
-                  style={{ color: "var(--text-secondary)" }}
+                  onSelect={() => handleSelect(model)}
+                  className="flex items-center gap-3"
                 >
-                  <ModelAvatar modelName={model.name} size={28} />
-                  <div className="min-w-0 flex-1">
-                    <div className="flex items-center gap-1.5">
-                      <span className="truncate font-medium" style={{ color: "var(--text-primary)" }}>
-                        {model.name}
-                      </span>
-                      {selectedModel?.model_name === model.name && (
-                        <Check className="w-3.5 h-3.5 shrink-0" style={{ color: "var(--accent)" }} />
-                      )}
-                    </div>
-                    <div className="text-[11px] truncate" style={{ color: "var(--text-muted)" }}>
-                      {model.provider_id}
-                    </div>
-                  </div>
-                </button>
-              ))
-            )}
-          </div>
-        </div>
-      )}
-    </div>
+                  <ModelSelectorLogo provider={model.provider_id as ModelSelectorLogoProps["provider"]} />
+                  <span className="flex-1 truncate text-left">{model.name}</span>
+                  {selectedModel?.model_name === model.name && (
+                    <Check className="w-3.5 h-3.5 shrink-0 text-accent" />
+                  )}
+                </ModelSelectorItem>
+              ))}
+            </ModelSelectorGroup>
+          )}
+        </ModelSelectorList>
+      </ModelSelectorContent>
+    </ModelSelector>
   );
 }