@useatlas/create 0.0.5 → 0.0.7

package/templates/docker/src/ui/components/atlas-chat.tsx

@@ -1,24 +1,39 @@
 "use client";
 
 import { useChat } from "@ai-sdk/react";
-import { DefaultChatTransport, isToolUIPart } from "ai";
-import { useState, useRef, useEffect, useMemo, useCallback } from "react";
-import { AUTH_MODES, type AuthMode } from "../lib/types";
+import { isToolUIPart, getToolName } from "ai";
+import { useState, useRef, useEffect, useCallback } from "react";
+import type { PythonProgressData } from "./chat/python-result-card";
 import { useAtlasConfig, ActionAuthProvider } from "../context";
-import { DarkModeContext } from "../hooks/use-dark-mode";
-import { useDarkMode } from "../hooks/use-dark-mode";
+import { DarkModeContext, useDarkMode, useThemeMode, setTheme, type ThemeMode } from "../hooks/use-dark-mode";
+import { useAtlasTransport } from "../hooks/use-atlas-transport";
 import { useConversations } from "../hooks/use-conversations";
 import { ErrorBanner } from "./chat/error-banner";
 import { ApiKeyBar } from "./chat/api-key-bar";
-import { ManagedAuthCard } from "./chat/managed-auth-card";
 import { TypingIndicator } from "./chat/typing-indicator";
 import { ToolPart } from "./chat/tool-part";
 import { Markdown } from "./chat/markdown";
 import { STARTER_PROMPTS } from "./chat/starter-prompts";
+import { FollowUpChips } from "./chat/follow-up-chips";
+import { SuggestionChips } from "./chat/suggestion-chips";
+import type { QuerySuggestion } from "@/ui/lib/types";
+import { ShareDialog } from "./chat/share-dialog";
 import { ConversationSidebar } from "./conversations/conversation-sidebar";
 import { ChangePasswordDialog } from "./admin/change-password-dialog";
-
-const API_KEY_STORAGE_KEY = "atlas-api-key";
+import { Sun, Moon, Monitor, Star, TableProperties, BookOpen } from "lucide-react";
+import { SchemaExplorer } from "./schema-explorer/schema-explorer";
+import { PromptLibrary } from "./chat/prompt-library";
+import {
+  DropdownMenu,
+  DropdownMenuContent,
+  DropdownMenuItem,
+  DropdownMenuTrigger,
+} from "@/components/ui/dropdown-menu";
+import { Button } from "@/components/ui/button";
+import { Input } from "@/components/ui/input";
+import { ScrollArea } from "@/components/ui/scroll-area";
+import { parseSuggestions } from "../lib/helpers";
+import { ErrorBoundary } from "./error-boundary";
 
 /* Static SVG icons — hoisted to avoid recreation on every render */
 const MenuIcon = (
@@ -34,34 +49,128 @@ const AtlasLogo = (
   </svg>
 );
 
+const THEME_OPTIONS = [
+  { value: "light", label: "Light", icon: Sun },
+  { value: "dark", label: "Dark", icon: Moon },
+  { value: "system", label: "System", icon: Monitor },
+] as const satisfies readonly { value: ThemeMode; label: string; icon: typeof Sun }[];
+
+function ThemeToggle() {
+  const mode = useThemeMode();
+  const CurrentIcon = mode === "dark" ? Moon : mode === "light" ? Sun : Monitor;
+
+  return (
+    <DropdownMenu>
+      <DropdownMenuTrigger asChild>
+        <Button variant="ghost" size="icon" className="size-11 sm:size-8 text-zinc-500 dark:text-zinc-400">
+          <CurrentIcon className="size-4" />
+          <span className="sr-only">Toggle theme</span>
+        </Button>
+      </DropdownMenuTrigger>
+      <DropdownMenuContent align="end">
+        {THEME_OPTIONS.map(({ value, label, icon: Icon }) => (
+          <DropdownMenuItem
+            key={value}
+            onClick={() => setTheme(value)}
+            className={mode === value ? "bg-accent" : ""}
+          >
+            <Icon className="mr-2 size-4" />
+            {label}
+          </DropdownMenuItem>
+        ))}
+      </DropdownMenuContent>
+    </DropdownMenu>
+  );
+}
+
+function SaveButton({
+  conversationId,
+  conversations,
+  onStar,
+}: {
+  conversationId: string;
+  conversations: { id: string; starred: boolean }[];
+  onStar: (id: string, starred: boolean) => Promise<void>;
+}) {
+  const isStarred = conversations.find((c) => c.id === conversationId)?.starred ?? false;
+  const [pending, setPending] = useState(false);
+
+  async function handleToggle() {
+    setPending(true);
+    try {
+      await onStar(conversationId, !isStarred);
+    } catch (err: unknown) {
+      console.warn("Failed to update star:", err instanceof Error ? err.message : String(err));
+    } finally {
+      setPending(false);
+    }
+  }
+
+  return (
+    <Button
+      variant="ghost"
+      size="xs"
+      onClick={handleToggle}
+      disabled={pending}
+      className={
+        isStarred
+          ? "text-amber-500 hover:text-amber-600 dark:text-amber-400 dark:hover:text-amber-300"
+          : "text-zinc-400 hover:text-amber-500 dark:text-zinc-500 dark:hover:text-amber-400"
+      }
+      aria-label={isStarred ? "Unsave conversation" : "Save conversation"}
+    >
+      <Star className="h-3.5 w-3.5" fill={isStarred ? "currentColor" : "none"} />
+      <span>{isStarred ? "Saved" : "Save"}</span>
+    </Button>
+  );
+}
+
 export function AtlasChat() {
   const { apiUrl, isCrossOrigin, authClient } = useAtlasConfig();
   const dark = useDarkMode();
   const [input, setInput] = useState("");
-  const [authMode, setAuthMode] = useState<AuthMode | null>(null);
-  const [healthWarning, setHealthWarning] = useState("");
-  const [apiKey, setApiKey] = useState("");
   const [conversationId, setConversationId] = useState<string | null>(null);
+  const [transientWarning, setTransientWarning] = useState("");
   const [mobileMenuOpen, setMobileMenuOpen] = useState(false);
   const [loadingConversation, setLoadingConversation] = useState(false);
   const [passwordChangeRequired, setPasswordChangeRequired] = useState(false);
+  const [schemaExplorerOpen, setSchemaExplorerOpen] = useState(false);
+  const [promptLibraryOpen, setPromptLibraryOpen] = useState(false);
+  const [popularSuggestions, setPopularSuggestions] = useState<QuerySuggestion[]>([]);
+  const [suggestionsLoading, setSuggestionsLoading] = useState(false);
+  const [relatedSuggestions, setRelatedSuggestions] = useState<QuerySuggestion[]>([]);
   const scrollRef = useRef<HTMLDivElement>(null);
 
+  const {
+    transport,
+    authMode,
+    apiKey,
+    setApiKey,
+    getHeaders,
+    getCredentials,
+    healthWarning,
+    authResolved,
+  } = useAtlasTransport({
+    apiUrl,
+    isCrossOrigin,
+    getConversationId: () => conversationId,
+    onNewConversationId: (id) => {
+      setConversationId(id);
+      setTimeout(() => {
+        refreshConvosRef.current().catch((err: unknown) => {
+          console.warn(
+            "Sidebar refresh failed:",
+            err instanceof Error ? err.message : String(err),
+          );
+        });
+      }, 500);
+    },
+  });
+
   const managedSession = authClient.useSession();
-  const authResolved = authMode !== null;
   const isManaged = authMode === "managed";
   const isSignedIn = isManaged && !!managedSession.data?.user;
 
-  const getHeaders = useCallback(() => {
-    const headers: Record<string, string> = {};
-    if (apiKey) headers["Authorization"] = `Bearer ${apiKey}`;
-    return headers;
-  }, [apiKey]);
-
-  const getCredentials = useCallback((): RequestCredentials => {
-    return isCrossOrigin ? "include" : "same-origin";
-  }, [isCrossOrigin]);
-
   const convos = useConversations({
     apiUrl,
     enabled: true,
@@ -72,51 +181,6 @@ export function AtlasChat() {
   const refreshConvosRef = useRef(convos.refresh);
   refreshConvosRef.current = convos.refresh;
 
-  const conversationIdRef = useRef(conversationId);
-  conversationIdRef.current = conversationId;
-
-  // Load API key from sessionStorage on mount + fetch auth mode + conversations
-  useEffect(() => {
-    try {
-      const stored = sessionStorage.getItem(API_KEY_STORAGE_KEY);
-      if (stored) setApiKey(stored);
-    } catch (err) {
-      console.warn("Cannot read API key from sessionStorage:", err);
-    }
-
-    async function fetchHealth(attempt: number): Promise<void> {
-      try {
-        const res = await fetch(`${apiUrl}/api/health`, {
-          credentials: isCrossOrigin ? "include" : "same-origin",
-        });
-        if (!res.ok) {
-          console.warn(`Health check returned ${res.status}`);
-          if (attempt < 2) {
-            await new Promise((r) => setTimeout(r, 2000));
-            return fetchHealth(attempt + 1);
-          }
-          setHealthWarning("Health check failed — check server logs. Try refreshing the page.");
-          setAuthMode("none");
-          return;
-        }
-        const data = await res.json();
-        const mode = data?.checks?.auth?.mode;
-        if (typeof mode === "string" && AUTH_MODES.includes(mode as AuthMode)) {
-          setAuthMode(mode as AuthMode);
-        }
-      } catch (err) {
-        console.warn("Health endpoint unavailable:", err);
-        if (attempt < 2) {
-          await new Promise((r) => setTimeout(r, 2000));
-          return fetchHealth(attempt + 1);
-        }
-        setHealthWarning("Unable to reach the API server. Try refreshing the page.");
-        setAuthMode("none");
-      }
-    }
-    fetchHealth(1);
-  }, [apiUrl]);
-
   // Fetch conversation list after auth is resolved
   useEffect(() => {
     convos.fetchList();
@@ -131,57 +195,121 @@ export function AtlasChat() {
         const res = await fetch(`${apiUrl}/api/v1/admin/me/password-status`, {
           credentials: isCrossOrigin ? "include" : "same-origin",
         });
-        if (!res.ok) return;
+        if (!res.ok) {
+          console.warn("Password status check failed:", res.status, res.statusText);
+          return;
+        }
         const data = await res.json();
         if (data.passwordChangeRequired) setPasswordChangeRequired(true);
-      } catch {
-        // Non-critical — skip silently
+      } catch (err: unknown) {
+        console.warn("Failed to check password status:", err instanceof Error ? err.message : String(err));
       }
     }
     checkPasswordStatus();
   }, [isManaged, managedSession.data?.user, apiUrl, isCrossOrigin]);
 
-  const handleSaveApiKey = useCallback((key: string) => {
-    setApiKey(key);
-    try {
-      sessionStorage.setItem(API_KEY_STORAGE_KEY, key);
-    } catch (err) {
-      console.warn("Could not persist API key to sessionStorage:", err);
+  // Python streaming progress — keyed by tool invocation ID
+  const [pythonProgress, setPythonProgress] = useState<Map<string, PythonProgressData[]>>(new Map());
+
+  const onData = useCallback((dataPart: { type: string; id?: string; data: unknown }) => {
+    if (dataPart.type === "data-python-progress" && dataPart.id && dataPart.data) {
+      const d = dataPart.data as Record<string, unknown>;
+      // Minimal runtime validation — ensure the event has a known type
+      if (typeof d.type !== "string" || !["stdout", "chart", "recharts"].includes(d.type)) return;
+      const event = d as unknown as PythonProgressData;
+      setPythonProgress((prev) => {
+        const next = new Map(prev);
+        const events = next.get(dataPart.id!) ?? [];
+        next.set(dataPart.id!, [...events, event]);
+        return next;
+      });
     }
   }, []);
 
-  // Dynamic transport — captures x-conversation-id from response.
-  // conversationId is accessed via ref to avoid recreating the transport mid-stream
-  // (which causes an infinite re-render loop in useChat).
-  const transport = useMemo(() => {
-    const headers: Record<string, string> = {};
-    if (apiKey) {
-      headers["Authorization"] = `Bearer ${apiKey}`;
-    }
-    return new DefaultChatTransport({
-      api: `${apiUrl}/api/chat`,
-      headers,
-      credentials: isCrossOrigin ? "include" : undefined,
-      body: () => (conversationIdRef.current ? { conversationId: conversationIdRef.current } : {}),
-      fetch: (async (input, init) => {
-        const response = await globalThis.fetch(input, init);
-        const convId = response.headers.get("x-conversation-id");
-        if (convId && convId !== conversationIdRef.current) {
-          setConversationId(convId);
-          setTimeout(() => {
-            refreshConvosRef.current().catch((err) => {
-              console.warn("Sidebar refresh failed:", err);
-            });
-          }, 500);
+  // The AI SDK's onData expects DataUIPart<UIDataTypes> which structurally accepts
+  // { type: `data-${string}`; id?: string; data: unknown } — our callback matches.
+  // The cast is needed because the default UIMessage generic doesn't declare our custom
+  // data part type at compile time.
+  const { messages, setMessages, sendMessage, status, error } = useChat({
+    transport,
+    onData: onData as never,
+  });
+
+  const isLoading = status === "streaming" || status === "submitted";
+
+  // Fetch popular suggestions for the empty state
+  useEffect(() => {
+    if (messages.length > 0) return;
+    let cancelled = false;
+    setSuggestionsLoading(true);
+    fetch(`${apiUrl}/api/v1/suggestions/popular?limit=6`, {
+      credentials: isCrossOrigin ? "include" : "same-origin",
+      headers: getHeaders(),
+    })
+      .then((res) => (res.ok ? res.json() : null))
+      .then((data) => {
+        if (!cancelled && data?.suggestions) {
+          setPopularSuggestions(data.suggestions);
         }
-        return response;
-      }) as typeof fetch,
-    });
-  }, [apiKey, authMode, apiUrl, isCrossOrigin]);
+      })
+      .catch(() => {
+        // intentionally ignored: suggestions are non-critical
+      })
+      .finally(() => {
+        if (!cancelled) setSuggestionsLoading(false);
+      });
+    return () => { cancelled = true; };
+  }, [messages.length, apiUrl, isCrossOrigin, getHeaders]);
+
+  // Fetch related suggestions after a completed query with SQL results
+  useEffect(() => {
+    if (messages.length === 0 || isLoading) return;
+    const lastMsg = messages[messages.length - 1];
+    if (lastMsg.role !== "assistant") return;
+
+    // Extract tables from executeSQL tool invocations
+    const tables: string[] = [];
+    for (const part of lastMsg.parts ?? []) {
+      if (!isToolUIPart(part)) continue;
+      try {
+        const name = getToolName(part as Parameters<typeof getToolName>[0]);
+        if (name !== "executeSQL") continue;
+      } catch {
+        // intentionally ignored: getToolName throws if the part lacks a recognized toolName property
+        continue;
+      }
+      if ("result" in part && part.result != null) {
+        const result = part.result as Record<string, unknown>;
+        if (Array.isArray(result.tablesAccessed)) {
+          tables.push(...(result.tablesAccessed as string[]));
+        }
+      }
+    }
 
-  const { messages, setMessages, sendMessage, status, error } = useChat({ transport });
+    if (tables.length === 0) {
+      setRelatedSuggestions([]);
+      return;
+    }
 
-  const isLoading = status === "streaming" || status === "submitted";
+    const uniqueTables = [...new Set(tables)];
+    const params = uniqueTables.map((t) => `table=${encodeURIComponent(t)}`).join("&");
+
+    let cancelled = false;
+    fetch(`${apiUrl}/api/v1/suggestions?${params}&limit=3`, {
+      credentials: isCrossOrigin ? "include" : "same-origin",
+      headers: getHeaders(),
+    })
+      .then((res) => (res.ok ? res.json() : null))
+      .then((data) => {
+        if (!cancelled && data?.suggestions) {
+          setRelatedSuggestions(data.suggestions);
+        }
+      })
+      .catch(() => {
+        // intentionally ignored: suggestions are non-critical
+      });
+    return () => { cancelled = true; };
+  }, [messages.length, isLoading, apiUrl, isCrossOrigin, getHeaders]);
 
   useEffect(() => {
     const el = scrollRef.current;
@@ -194,12 +322,23 @@ export function AtlasChat() {
     if (!text.trim()) return;
     const saved = text;
     setInput("");
-    sendMessage({ text: saved }).catch((err) => {
-      console.error("Failed to send message:", err);
+    sendMessage({ text: saved }).catch((err: unknown) => {
+      console.error("Failed to send message:", err instanceof Error ? err.message : String(err));
       setInput(saved);
-      setHealthWarning("Failed to send message. Please try again.");
-      setTimeout(() => setHealthWarning(""), 5000);
+      setTransientWarning("Failed to send message. Please try again.");
+      setTimeout(() => setTransientWarning(""), 5000);
+    });
+  }
+
+  function handleSuggestionSelect(text: string, id: string) {
+    fetch(`${apiUrl}/api/v1/suggestions/${id}/click`, {
+      method: "POST",
+      credentials: isCrossOrigin ? "include" : "same-origin",
+      headers: getHeaders(),
+    }).catch(() => {
+      // intentionally ignored: click tracking is non-critical
     });
+    handleSend(text);
   }
 
   async function handleSelectConversation(id: string) {
@@ -207,15 +346,14 @@ export function AtlasChat() {
     setLoadingConversation(true);
     try {
       const uiMessages = await convos.loadConversation(id);
-      if (uiMessages) {
-        setMessages(uiMessages);
-        setConversationId(id);
-        convos.setSelectedId(id);
-        setMobileMenuOpen(false);
-      } else {
-        setHealthWarning("Could not load conversation. It may have been deleted.");
-        setTimeout(() => setHealthWarning(""), 5000);
-      }
+      setMessages(uiMessages);
+      setConversationId(id);
+      convos.setSelectedId(id);
+      setMobileMenuOpen(false);
+    } catch (err: unknown) {
+      console.warn("Failed to load conversation:", err instanceof Error ? err.message : String(err));
+      setTransientWarning("Failed to load conversation. Please try again.");
+      setTimeout(() => setTransientWarning(""), 5000);
     } finally {
       setLoadingConversation(false);
     }
@@ -227,6 +365,7 @@ export function AtlasChat() {
     convos.setSelectedId(null);
     setInput("");
     setMobileMenuOpen(false);
+    setPythonProgress(new Map());
   }
 
   // Wait for auth mode detection before rendering — prevents flash of chat UI
@@ -256,19 +395,21 @@ export function AtlasChat() {
           />
         )}
 
-        <main className="flex flex-1 flex-col overflow-hidden">
-          <div className="mx-auto flex w-full max-w-4xl flex-1 flex-col p-4">
+        <main id="main" tabIndex={-1} className="flex flex-1 flex-col overflow-hidden">
+          <div className="mx-auto flex w-full max-w-4xl flex-1 flex-col overflow-hidden p-4">
             <header className="mb-4 flex-none border-b border-zinc-100 pb-3 dark:border-zinc-800">
               <div className="flex items-center justify-between">
                 <div className="flex items-center gap-3">
                   {convos.available && (
-                    <button
+                    <Button
+                      variant="ghost"
+                      size="icon"
                       onClick={() => setMobileMenuOpen(true)}
-                      className="rounded p-1 text-zinc-400 hover:text-zinc-700 md:hidden dark:hover:text-zinc-200"
+                      className="size-11 text-zinc-400 hover:text-zinc-700 md:hidden dark:hover:text-zinc-200"
                       aria-label="Open conversation history"
                     >
                       {MenuIcon}
-                    </button>
+                    </Button>
                   )}
                   <div className="flex items-center gap-2.5">
                     {AtlasLogo}
@@ -278,71 +419,118 @@ export function AtlasChat() {
                     </div>
                   </div>
                 </div>
-                {isSignedIn && (
-                  <div className="flex items-center gap-3">
-                    <span className="text-xs text-zinc-500 dark:text-zinc-400">
-                      {managedSession.data?.user?.email}
-                    </span>
-                    <button
-                      onClick={() => {
-                        authClient.signOut().catch((err: unknown) => {
-                          console.error("Sign out failed:", err);
-                          setHealthWarning("Sign out failed. Please try again.");
-                          setTimeout(() => setHealthWarning(""), 5000);
-                        });
-                      }}
-                      className="rounded border border-zinc-200 px-2 py-1 text-xs text-zinc-500 transition-colors hover:border-zinc-400 hover:text-zinc-800 dark:border-zinc-700 dark:text-zinc-400 dark:hover:border-zinc-500 dark:hover:text-zinc-200"
-                    >
-                      Sign out
-                    </button>
-                  </div>
-                )}
+                <div className="flex items-center gap-2">
+                  <Button
+                    variant="ghost"
+                    size="icon"
+                    className="size-11 sm:size-8 text-zinc-500 dark:text-zinc-400"
+                    onClick={() => setPromptLibraryOpen(true)}
+                    aria-label="Open prompt library"
+                  >
+                    <BookOpen className="size-4" />
+                  </Button>
+                  <Button
+                    variant="ghost"
+                    size="icon"
+                    className="size-11 sm:size-8 text-zinc-500 dark:text-zinc-400"
+                    onClick={() => setSchemaExplorerOpen(true)}
+                    aria-label="Open schema explorer"
+                  >
+                    <TableProperties className="size-4" />
+                  </Button>
+                  <ThemeToggle />
+                  {isSignedIn && (
+                    <>
+                      <span className="hidden text-xs text-zinc-500 sm:inline dark:text-zinc-400">
+                        {managedSession.data?.user?.email}
+                      </span>
+                      <Button
+                        variant="outline"
+                        size="sm"
+                        onClick={() => {
+                          authClient.signOut().catch((err: unknown) => {
+                            console.error("Sign out failed:", err instanceof Error ? err.message : String(err));
+                            setTransientWarning("Sign out failed. Please try again.");
+                            setTimeout(() => setTransientWarning(""), 5000);
+                          });
+                        }}
+                        className="text-xs text-zinc-500 dark:text-zinc-400"
+                      >
+                        Sign out
+                      </Button>
+                    </>
+                  )}
+                </div>
               </div>
             </header>
 
-            {healthWarning && (
-              <p className="mb-2 text-xs text-zinc-400 dark:text-zinc-500">{healthWarning}</p>
+            {(healthWarning || transientWarning || convos.fetchError) && (
+              <p className="mb-2 text-xs text-zinc-400 dark:text-zinc-500">{healthWarning || transientWarning || convos.fetchError}</p>
             )}
 
             {isManaged && !isSignedIn ? (
-              <ManagedAuthCard />
+              null /* proxy redirects unauthenticated users to /signup */
             ) : (
               <ActionAuthProvider getHeaders={getHeaders} getCredentials={getCredentials}>
                 {authMode === "simple-key" && (
                   <div className="mb-3 flex-none">
-                    <ApiKeyBar apiKey={apiKey} onSave={handleSaveApiKey} />
+                    <ApiKeyBar apiKey={apiKey} onSave={setApiKey} />
                   </div>
                 )}
 
-                <div ref={scrollRef} className="flex-1 space-y-4 overflow-y-auto pb-4">
+                <ScrollArea viewportRef={scrollRef} className="min-h-0 flex-1">
+                <ErrorBoundary
+                  fallbackRender={(_error, reset) => (
+                    <div className="flex flex-col items-center justify-center gap-2 p-6 text-sm text-red-600 dark:text-red-400">
+                      <p>Failed to render messages.</p>
+                      <Button variant="link" size="sm" onClick={reset} className="text-xs">Try again</Button>
+                    </div>
+                  )}
+                >
+                <div className="space-y-4 pb-4 pr-3">
                   {messages.length === 0 && !error && (
                     <div className="flex h-full flex-col items-center justify-center gap-6">
                       <div className="text-center">
                         <p className="text-lg font-medium text-zinc-500 dark:text-zinc-400">
                           What would you like to know?
                         </p>
-                        <p className="mt-1 text-sm text-zinc-400 dark:text-zinc-600">
+                        <p className="mt-1 text-sm text-zinc-500 dark:text-zinc-500">
                           Ask a question about your data to get started
                         </p>
                       </div>
-                      <div className="grid w-full max-w-lg grid-cols-2 gap-2">
+                      <SuggestionChips
+                        suggestions={popularSuggestions}
+                        onSelect={handleSuggestionSelect}
+                        loading={suggestionsLoading}
+                        label="Popular queries"
+                      />
+                      <div className="grid w-full max-w-lg grid-cols-1 gap-2 sm:grid-cols-2">
                         {STARTER_PROMPTS.map((prompt) => (
-                          <button
+                          <Button
                             key={prompt}
+                            variant="outline"
                             onClick={() => handleSend(prompt)}
-                            className="rounded-lg border border-zinc-200 bg-zinc-50 px-3 py-2.5 text-left text-sm text-zinc-500 transition-colors hover:border-zinc-400 hover:bg-zinc-100 hover:text-zinc-800 dark:border-zinc-800 dark:bg-zinc-900 dark:text-zinc-400 dark:hover:border-zinc-600 dark:hover:bg-zinc-800 dark:hover:text-zinc-200"
+                            className="h-auto whitespace-normal justify-start rounded-lg bg-zinc-50 px-3 py-2.5 text-left text-sm text-zinc-500 hover:text-zinc-800 dark:bg-zinc-900 dark:text-zinc-400 dark:hover:text-zinc-200"
                           >
                             {prompt}
-                          </button>
+                          </Button>
                         ))}
                       </div>
+                      <Button
+                        variant="link"
+                        onClick={() => setPromptLibraryOpen(true)}
+                        className="text-xs text-zinc-400 dark:text-zinc-500"
+                      >
+                        <BookOpen className="mr-1.5 size-3.5" />
+                        Browse prompt library
+                      </Button>
                     </div>
                   )}
 
-                  {messages.map((m) => {
+                  {messages.map((m, msgIndex) => {
                     if (m.role === "user") {
                       return (
-                        <div key={m.id} className="flex justify-end">
+                        <div key={m.id} className="flex justify-end" role="article" aria-label="Message from you">
                           <div className="max-w-[85%] rounded-xl bg-blue-600 px-4 py-3 text-sm text-white">
                             {m.parts?.map((part, i) =>
                               part.type === "text" ? (
@@ -356,14 +544,35 @@ export function AtlasChat() {
                       );
                     }
 
+                    const isLastAssistant =
+                      m.role === "assistant" &&
+                      msgIndex === messages.length - 1;
+
+                    // Skip rendering assistant messages with no visible content
+                    // (happens when stream errors before producing any text)
+                    const hasVisibleParts = m.parts?.some(
+                      (p) => (p.type === "text" && p.text.trim()) || isToolUIPart(p),
+                    );
+                    if (!hasVisibleParts && !isLastAssistant) return null;
+
+                    // Extract suggestions from the last text part that contains them
+                    const lastTextWithSuggestions = m.parts
+                      ?.filter((p): p is typeof p & { type: "text"; text: string } => p.type === "text" && !!p.text.trim())
+                      .findLast((p) => parseSuggestions(p.text).suggestions.length > 0);
+                    const suggestions = lastTextWithSuggestions
+                      ? parseSuggestions(lastTextWithSuggestions.text).suggestions
+                      : [];
+
                     return (
-                      <div key={m.id} className="space-y-2">
+                      <div key={m.id} className="space-y-2" role="article" aria-label="Message from Atlas">
                         {m.parts?.map((part, i) => {
                           if (part.type === "text" && part.text.trim()) {
+                            const displayText = parseSuggestions(part.text).text;
+                            if (!displayText.trim()) return null;
                             return (
                               <div key={i} className="max-w-[90%]">
                                 <div className="rounded-xl bg-zinc-100 px-4 py-3 text-sm text-zinc-800 dark:bg-zinc-800 dark:text-zinc-200">
-                                  <Markdown content={part.text} />
+                                  <Markdown content={displayText} />
                                 </div>
                               </div>
                             );
@@ -371,20 +580,78 @@ export function AtlasChat() {
                           if (isToolUIPart(part)) {
                             return (
                               <div key={i} className="max-w-[95%]">
-                                <ToolPart part={part} />
+                                <ToolPart part={part} pythonProgress={pythonProgress} />
                               </div>
                             );
                           }
                           return null;
                         })}
+                        {/* Show inline error when the last assistant message is empty (stream failed before producing content) */}
+                        {isLastAssistant && !hasVisibleParts && !isLoading && error && (
+                          <div className="max-w-[90%]">
+                            <div className="rounded-xl bg-red-50 px-4 py-3 text-sm text-red-700 dark:bg-red-950/30 dark:text-red-400">
+                              The response stream was interrupted before producing content. Try sending your message again.
+                            </div>
+                          </div>
+                        )}
+                        {isLastAssistant && !isLoading && hasVisibleParts && (
+                          <>
+                            <FollowUpChips
+                              suggestions={suggestions}
+                              onSelect={handleSend}
+                            />
+                            {relatedSuggestions.length > 0 && (
+                              <SuggestionChips
+                                suggestions={relatedSuggestions}
+                                onSelect={handleSuggestionSelect}
+                                label="Related queries"
+                              />
+                            )}
+                            {conversationId && convos.available && (
+                              <div className="flex items-center gap-1">
+                                <SaveButton
+                                  conversationId={conversationId}
+                                  conversations={convos.conversations}
+                                  onStar={convos.starConversation}
+                                />
+                                <ShareDialog
+                                  key={conversationId}
+                                  conversationId={conversationId}
+                                  onShare={convos.shareConversation}
+                                  onUnshare={convos.unshareConversation}
+                                  onGetShareStatus={convos.getShareStatus}
+                                />
+                              </div>
+                            )}
+                          </>
+                        )}
                       </div>
                     );
                   })}
 
                   {isLoading && messages.length > 0 && <TypingIndicator />}
                 </div>
-
-                {error && <ErrorBanner error={error} authMode={authMode} />}
+                </ErrorBoundary>
+                </ScrollArea>
+
+                {error && (
+                  <ErrorBanner
+                    error={error}
+                    authMode={authMode ?? "none"}
+                    onRetry={
+                      messages.some((m) => m.role === "user")
+                        ? () => {
+                            const lastUserMsg = messages.toReversed().find((m) => m.role === "user");
+                            const text = lastUserMsg?.parts
+                              ?.filter((p): p is { type: "text"; text: string } => p.type === "text")
+                              .map((p) => p.text)
+                              .join(" ");
+                            if (text) handleSend(text);
+                          }
+                        : undefined
+                    }
+                  />
+                )}
 
                 <form
                   onSubmit={(e) => {
@@ -393,30 +660,46 @@ export function AtlasChat() {
                   }}
                   className="flex flex-none gap-2 border-t border-zinc-100 pt-4 dark:border-zinc-800"
                 >
-                  <input
+                  <Input
                     value={input}
                     onChange={(e) => setInput(e.target.value)}
                     placeholder="Ask a question about your data..."
-                    className="flex-1 rounded-lg border border-zinc-200 bg-zinc-50 px-4 py-3 text-sm text-zinc-900 placeholder-zinc-400 outline-none focus:border-blue-500 dark:border-zinc-700 dark:bg-zinc-900 dark:text-zinc-100 dark:placeholder-zinc-600"
+                    className="min-w-0 flex-1 py-3 text-base sm:text-sm"
                     disabled={isLoading}
+                    aria-label="Chat message"
                   />
-                  <button
+                  <Button
                     type="submit"
-                    disabled={isLoading || !input.trim()}
-                    className="rounded-lg bg-blue-600 px-5 py-3 text-sm font-medium text-white transition-colors hover:bg-blue-500 disabled:opacity-40"
+                    disabled={isLoading}
+                    aria-disabled={!isLoading && !input.trim() ? true : undefined}
+                    className="shrink-0 px-5"
                   >
                     Ask
-                  </button>
+                  </Button>
                 </form>
               </ActionAuthProvider>
             )}
           </div>
         </main>
       </div>
+      <SchemaExplorer
+        open={schemaExplorerOpen}
+        onOpenChange={setSchemaExplorerOpen}
+        onInsertQuery={(text) => setInput(text)}
+        getHeaders={getHeaders}
+        getCredentials={getCredentials}
+      />
+      <PromptLibrary
+        open={promptLibraryOpen}
+        onOpenChange={setPromptLibraryOpen}
+        onSendPrompt={handleSend}
+        getHeaders={getHeaders}
+        getCredentials={getCredentials}
+      />
       <ChangePasswordDialog
         open={passwordChangeRequired}
         onComplete={() => setPasswordChangeRequired(false)}
       />
     </DarkModeContext.Provider>
   );
-}
+}