diff --git a/package.json b/package.json
index 42b00c54..b07c594e 100644
--- a/package.json
+++ b/package.json
@@ -27,7 +27,7 @@
   },
   "packageManager": "pnpm@10.14.0",
   "dependencies": {
-    "@ai-sdk/google": "^2.0.11",
+    "@ai-sdk/google": "^3.0.13",
     "@assistant-ui/react": "^0.12.1",
     "@assistant-ui/react-ai-sdk": "^1.3.1",
     "@assistant-ui/react-devtools": "^0.2.0",
@@ -88,7 +88,7 @@
     "@tanstack/react-query-devtools": "^5.91.2",
     "@tanstack/react-virtual": "^3.13.14",
     "@vercel/speed-insights": "^1.3.1",
-    "ai": "^5.0.116",
+    "ai": "^6.0.50",
     "assistant-cloud": "^0.1.13",
     "better-auth": "^1.4.10",
     "class-variance-authority": "^0.7.1",
@@ -129,6 +129,7 @@
     "sonner": "^2.0.7",
     "streamdown": "^2.1.0",
     "tailwind-merge": "^3.4.0",
+    "tw-shimmer": "^0.4.3",
     "typewriter-effect": "^2.22.0",
     "zod": "^4.3.4",
     "zustand": "^5.0.9"
diff --git a/src/app/api/chat/route.ts b/src/app/api/chat/route.ts
index f49dd48a..2fcb0278 100644
--- a/src/app/api/chat/route.ts
+++ b/src/app/api/chat/route.ts
@@ -1,5 +1,6 @@
 import { google } from "@ai-sdk/google";
-import { streamText, convertToModelMessages, stepCountIs } from "ai";
+import { streamText, convertToModelMessages, stepCountIs, tool, zodSchema } from "ai";
+import type { UIMessage } from "ai";
 import { logger } from "@/lib/utils/logger";
 import { auth } from "@/lib/auth";
 import { headers } from "next/headers";
@@ -175,7 +176,7 @@ export async function POST(req: Request) {
     const session = await auth.api.getSession({ headers: headersObj });
     const userId = session?.user?.id || null;
 
-    const messages = body.messages || [];
+    const { messages = [] }: { messages?: UIMessage[] } = body;
     const system = body.system || "";
     workspaceId = extractWorkspaceId(body);
     activeFolderId = body.activeFolderId;
diff --git a/src/app/globals.css b/src/app/globals.css
index b32e7bfc..1224d48f 100644
--- a/src/app/globals.css
+++ b/src/app/globals.css
@@ -1,6 +1,7 @@
 @import "tailwindcss";
 @import "tw-animate-css";
 @import "@daveyplate/better-auth-ui/css";
+@import "tw-shimmer";
 
 /* Path to your installed `@blocknote/shadcn` package. */
 @source "../node_modules/@blocknote/shadcn";
@@ -247,8 +248,6 @@
   margin-top: -0.125rem !important;
 }
 
-
-
 .bn-block-content[data-content-type="checkListItem"] input[type="checkbox"] {
   margin-top: 0.25rem !important;
   flex-shrink: 0 !important;
@@ -257,7 +256,6 @@
   cursor: pointer !important;
 }
 
-
 .bn-block-content[data-content-type="checkListItem"] .bn-inline-content {
   margin-top: 0 !important;
   flex: 1 !important;
diff --git a/src/components/assistant-ui/AssistantDropzone.tsx b/src/components/assistant-ui/AssistantDropzone.tsx
index bf5e7ee7..150f44d5 100644
--- a/src/components/assistant-ui/AssistantDropzone.tsx
+++ b/src/components/assistant-ui/AssistantDropzone.tsx
@@ -1,7 +1,7 @@
 "use client";
 
 import { useDropzone } from "react-dropzone";
-import { useAssistantApi } from "@assistant-ui/react";
+import { useAui } from "@assistant-ui/react";
 import { useWorkspaceStore } from "@/lib/stores/workspace-store";
 import { Upload } from "lucide-react";
 import { useCallback, useState, useRef } from "react";
@@ -16,7 +16,7 @@ interface AssistantDropzoneProps {
  * Accepts all supported file types and adds them as attachments to the chat composer.
  */
 export function AssistantDropzone({ children }: AssistantDropzoneProps) {
-  const api = useAssistantApi();
+  const aui = useAui();
   const currentWorkspaceId = useWorkspaceStore((state) => state.currentWorkspaceId);
   const [isDragging, setIsDragging] = useState(false);
 
@@ -31,7 +31,7 @@ export function AssistantDropzone({ children }: AssistantDropzoneProps) {
 
   const onDrop = useCallback(
     async (acceptedFiles: File[]) => {
-      if (!currentWorkspaceId || !api) return;
+      if (!currentWorkspaceId || !aui) return;
 
       const MAX_FILES = 10;
       const MAX_FILE_SIZE_MB = 10;
@@ -93,7 +93,7 @@ export function AssistantDropzone({ children }: AssistantDropzoneProps) {
         // Add each file to the composer
         const addPromises = validFiles.map(async (file) => {
           try {
-            await api.composer().addAttachment(file);
+            await aui.composer().addAttachment(file);
           } catch (error) {
             console.error("Failed to add attachment:", error);
             // Remove from processing set on error so it can be retried
@@ -124,7 +124,7 @@ export function AssistantDropzone({ children }: AssistantDropzoneProps) {
         }, 200);
       }
     },
-    [api, currentWorkspaceId]
+    [aui, currentWorkspaceId]
   );
 
   // Clear processing state when drag ends (user drags away or cancels)
@@ -139,7 +139,7 @@ export function AssistantDropzone({ children }: AssistantDropzoneProps) {
     onDrop,
     noClick: true, // Don't trigger on click, only drag and drop
     noKeyboard: true, // Don't trigger on keyboard
-    disabled: !currentWorkspaceId || !api, // Disable if no workspace is selected or api is not available
+    disabled: !currentWorkspaceId || !aui, // Disable if no workspace is selected or api is not available
     accept: {
       'image/*': ['.png', '.jpg', '.jpeg', '.gif', '.webp', '.svg', '.bmp'],
       'video/*': ['.mp4', '.webm', '.avi', '.mov', '.mkv'],
diff --git a/src/components/assistant-ui/AssistantTextSelectionManager.tsx b/src/components/assistant-ui/AssistantTextSelectionManager.tsx
index cf701b64..8ecb4d03 100644
--- a/src/components/assistant-ui/AssistantTextSelectionManager.tsx
+++ b/src/components/assistant-ui/AssistantTextSelectionManager.tsx
@@ -15,7 +15,7 @@ import SelectableText, {
 import { useUIStore, selectReplySelections } from "@/lib/stores/ui-store";
 import { getHighlightColorById } from "@/lib/utils/highlight-colors";
 import { useWorkspaceStore } from "@/lib/stores/workspace-store";
-import { useAssistantState } from "@assistant-ui/react";
+import { useAuiState } from "@assistant-ui/react";
 import { useQueryClient } from "@tanstack/react-query";
 import {
   Dialog,
@@ -105,8 +105,8 @@ export default function AssistantTextSelectionManager({
   // This component is always rendered within AssistantRuntimeProvider, so the hook is safe to use
   // Use threadListItem.id if available, otherwise fall back to mainThreadId
   // Using safe hooks to handle race condition during thread switching (GitHub issue #2722)
-  const threadListItemId = useAssistantState(({ threadListItem }) => (threadListItem as any)?.id);
-  const mainThreadId = useAssistantState(({ threads }) => (threads as any)?.mainThreadId);
+  const threadListItemId = useAuiState(({ threadListItem }) => (threadListItem as any)?.id);
+  const mainThreadId = useAuiState(({ threads }) => (threads as any)?.mainThreadId);
   const currentThreadId = threadListItemId || mainThreadId;
 
 
diff --git a/src/components/assistant-ui/SafeAssistantRuntimeProvider.tsx b/src/components/assistant-ui/SafeAssistantRuntimeProvider.tsx
deleted file mode 100644
index 468bcef4..00000000
--- a/src/components/assistant-ui/SafeAssistantRuntimeProvider.tsx
+++ /dev/null
@@ -1,102 +0,0 @@
-"use client";
-
-import { Component, ReactNode } from "react";
-import { AssistantRuntimeProvider } from "@assistant-ui/react";
-import { AssistantRuntime } from "@assistant-ui/react";
-
-interface Props {
-    runtime: AssistantRuntime;
-    children: ReactNode;
-}
-
-interface State {
-    hasError: boolean;
-    remountKey: number;
-}
-
-export class SafeAssistantRuntimeProvider extends Component<Props, State> {
-    constructor(props: Props) {
-        super(props);
-        this.state = { hasError: false, remountKey: 0 };
-    }
-
-    static getDerivedStateFromError(error: unknown) {
-        return { hasError: true };
-    }
-
-    componentDidCatch(error: unknown, errorInfo: unknown) {
-        const errorMessage = error instanceof Error ? error.message : String(error);
-
-        // Check if this is a known transient runtime error
-        const isTransientError = 
-            errorMessage.includes("tapLookupResources") || 
-            errorMessage.includes("Resource not found") ||
-            errorMessage.includes("ThreadListItemRuntime is not available");
-
-        if (isTransientError) {
-            console.warn("[SafeAssistantRuntimeProvider] Caught transient runtime error, attempting graceful recovery:", error);
-            // Attempt to recover by resetting error state and incrementing remount key after a short delay
-            // This allows the tree to re-mount with a fresh key, hopefully after the race condition passes
-            setTimeout(() => {
-                if (this.state.hasError) {
-                    this.setState({ 
-                        hasError: false,
-                        remountKey: this.state.remountKey + 1
-                    });
-                }
-            }, 100);
-        } else {
-            console.error("[SafeAssistantRuntimeProvider] Unhandled error:", error, errorInfo);
-            // For non-transient errors, still try to recover after a longer delay
-            setTimeout(() => {
-                if (this.state.hasError) {
-                    this.setState({ 
-                        hasError: false,
-                        remountKey: this.state.remountKey + 1
-                    });
-                }
-            }, 500);
-        }
-    }
-
-    componentDidUpdate(prevProps: Props) {
-        // If runtime changed, try to recover from error state
-        if (prevProps.runtime !== this.props.runtime && this.state.hasError) {
-            this.setState({ 
-                hasError: false,
-                remountKey: this.state.remountKey + 1
-            });
-        }
-    }
-
-    render() {
-        if (this.state.hasError) {
-            // During error state, render a minimal invisible placeholder
-            // This prevents gray screen while avoiding hook errors from children
-            // The setTimeout in componentDidCatch will trigger recovery
-            return (
-                <div 
-                    style={{ 
-                        position: 'absolute', 
-                        width: '100%', 
-                        height: '100%',
-                        opacity: 0,
-                        pointerEvents: 'none'
-                    }} 
-                    aria-hidden="true"
-                />
-            );
-        }
-
-        // Use remountKey to force complete remount of AssistantRuntimeProvider on recovery
-        // This ensures a clean state after error recovery
-        return (
-            <AssistantRuntimeProvider 
-                key={this.state.remountKey}
-                runtime={this.props.runtime}
-            >
-                {this.props.children}
-            </AssistantRuntimeProvider>
-        );
-    }
-}
diff --git a/src/components/assistant-ui/SpeechToTextButton.tsx b/src/components/assistant-ui/SpeechToTextButton.tsx
index 3b0311be..3cbc3aa6 100644
--- a/src/components/assistant-ui/SpeechToTextButton.tsx
+++ b/src/components/assistant-ui/SpeechToTextButton.tsx
@@ -2,12 +2,12 @@ import "regenerator-runtime/runtime";
 import { Mic, MicOff } from "lucide-react";
 import { FC, useEffect, useState } from "react";
 import SpeechRecognition, { useSpeechRecognition } from "react-speech-recognition";
-import { useAssistantApi } from "@assistant-ui/react";
+import { useAui } from "@assistant-ui/react";
 import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
 import { cn } from "@/lib/utils";
 
 export const SpeechToTextButton: FC = () => {
-    const api = useAssistantApi();
+    const aui = useAui();
     const {
         transcript,
         listening,
@@ -23,16 +23,16 @@ export const SpeechToTextButton: FC = () => {
             // We append the transcript to the original text
             // Ideally we would insert at cursor, but appending is safer for now
             const separator = originalText && !originalText.endsWith(' ') ? ' ' : '';
-            api.composer().setText(originalText + separator + transcript);
+            aui.composer().setText(originalText + separator + transcript);
         }
-    }, [transcript, listening, originalText, api]);
+    }, [transcript, listening, originalText, aui]);
 
     const handleStartListening = () => {
         // Reset transcript before starting a new session
         resetTranscript();
 
         // Capture current text before starting
-        const currentText = api.composer().getState().text;
+        const currentText = aui.composer().getState().text;
         setOriginalText(currentText);
 
         SpeechRecognition.startListening({ continuous: true });
diff --git a/src/components/assistant-ui/WorkspaceRuntimeProvider.tsx b/src/components/assistant-ui/WorkspaceRuntimeProvider.tsx
index 750f57c7..85caa045 100644
--- a/src/components/assistant-ui/WorkspaceRuntimeProvider.tsx
+++ b/src/components/assistant-ui/WorkspaceRuntimeProvider.tsx
@@ -1,7 +1,6 @@
 "use client";
 
-import { AssistantCloud } from "@assistant-ui/react";
-import { SafeAssistantRuntimeProvider } from "@/components/assistant-ui/SafeAssistantRuntimeProvider";
+import { AssistantCloud, AssistantRuntimeProvider } from "@assistant-ui/react";
 import { useChatRuntime, AssistantChatTransport } from "@assistant-ui/react-ai-sdk";
 import { useMemo, useCallback } from "react";
 import { AssistantAvailableProvider } from "@/contexts/AssistantAvailabilityContext";
@@ -151,10 +150,10 @@ export function WorkspaceRuntimeProvider({
   });
 
   return (
-    <SafeAssistantRuntimeProvider runtime={runtime}>
+    <AssistantRuntimeProvider runtime={runtime}>
       <AssistantAvailableProvider>
         {children}
       </AssistantAvailableProvider>
-    </SafeAssistantRuntimeProvider>
+    </AssistantRuntimeProvider>
   );
 }
diff --git a/src/components/assistant-ui/assistant-loader.tsx b/src/components/assistant-ui/assistant-loader.tsx
index af2e67b4..d8e913c9 100644
--- a/src/components/assistant-ui/assistant-loader.tsx
+++ b/src/components/assistant-ui/assistant-loader.tsx
@@ -1,6 +1,6 @@
 "use client";
 
-import { useAssistantState } from "@assistant-ui/react";
+import { useAuiState } from "@assistant-ui/react";
 import { useState, useEffect } from "react";
 import ShinyText from "@/components/ShinyText";
 
@@ -14,11 +14,11 @@ const thinkingMessages = [
 ];
 
 export const AssistantLoader = () => {
-    const isRunning = useAssistantState(
+    const isRunning = useAuiState(
         ({ message }) => (message as { status?: { type: string } })?.status?.type === "running"
     );
 
-    const isMessageEmpty = useAssistantState(({ message }) => {
+    const isMessageEmpty = useAuiState(({ message }) => {
         const msg = message as any;
         return !msg?.content || (Array.isArray(msg.content) && msg.content.length === 0);
     });
diff --git a/src/components/assistant-ui/attachment.tsx b/src/components/assistant-ui/attachment.tsx
index a676417b..5454e99a 100644
--- a/src/components/assistant-ui/attachment.tsx
+++ b/src/components/assistant-ui/attachment.tsx
@@ -1,7 +1,6 @@
 "use client";
 
 import React, { PropsWithChildren, useEffect, useState, useRef, type FC } from "react";
-import Image from "next/image";
 import { XIcon, FileText, Link as LinkIcon, Upload, Link2, SearchIcon, Plus, Code as CodeIcon, GalleryHorizontalEnd } from "lucide-react";
 import { LuPaperclip } from "react-icons/lu";
 import { toast } from "sonner";
@@ -9,9 +8,9 @@ import {
   AttachmentPrimitive,
   ComposerPrimitive,
   MessagePrimitive,
-  useAssistantApi,
+  useAui,
 } from "@assistant-ui/react";
-import { useAssistantState } from "@assistant-ui/react";
+import { useAuiState } from "@assistant-ui/react";
 import { useShallow } from "zustand/shallow";
 import {
   Tooltip,
@@ -78,7 +77,7 @@ function getFaviconUrl(url: string): string {
 }
 
 const useAttachmentSrc = () => {
-  const attachmentState = useAssistantState(
+  const attachmentState = useAuiState(
     useShallow(({ attachment }): { file?: File; src?: string; isUrl?: boolean; url?: string } => {
       const att = attachment as {
         type?: string;
@@ -166,18 +165,16 @@ type AttachmentPreviewProps = {
 const AttachmentPreview: FC<AttachmentPreviewProps> = ({ src }) => {
   const [isLoaded, setIsLoaded] = useState(false);
   return (
-    <Image
+    <img
       src={src}
       alt="Image Preview"
-      width={1}
-      height={1}
-      className={
+      className={cn(
+        "block h-auto max-h-[80vh] w-auto max-w-full object-contain",
         isLoaded
-          ? "aui-attachment-preview-image-loaded block h-auto max-h-[80vh] w-auto max-w-full object-contain"
-          : "aui-attachment-preview-image-loading hidden"
-      }
-      onLoadingComplete={() => setIsLoaded(true)}
-      priority={false}
+          ? "aui-attachment-preview-image-loaded"
+          : "aui-attachment-preview-image-loading invisible"
+      )}
+      onLoad={() => setIsLoaded(true)}
     />
   );
 };
@@ -196,7 +193,7 @@ const AttachmentPreviewDialog: FC<PropsWithChildren> = ({ children }) => {
       >
         {children}
       </DialogTrigger>
-      <DialogContent className="aui-attachment-preview-dialog-content p-2 sm:max-w-3xl [&_svg]:text-background [&>button]:rounded-full [&>button]:bg-foreground/60 [&>button]:p-1 [&>button]:opacity-100 [&>button]:!ring-0 [&>button]:hover:[&_svg]:text-destructive">
+      <DialogContent className="aui-attachment-preview-dialog-content p-2 sm:max-w-3xl [&>button]:rounded-full [&>button]:bg-foreground/60 [&>button]:p-1 [&>button]:opacity-100 [&>button]:ring-0! [&_svg]:text-background [&>button]:hover:[&_svg]:text-destructive">
         <DialogTitle className="aui-sr-only sr-only">
           Image Attachment Preview
         </DialogTitle>
@@ -209,7 +206,7 @@ const AttachmentPreviewDialog: FC<PropsWithChildren> = ({ children }) => {
 };
 
 const AttachmentThumb: FC = () => {
-  const isImage = useAssistantState(
+  const isImage = useAuiState(
     ({ attachment }) => (attachment as { type?: string })?.type === "image"
   );
   const attachmentSrc = useAttachmentSrc();
@@ -257,15 +254,15 @@ const AttachmentThumb: FC = () => {
 };
 
 const AttachmentUI: FC = () => {
-  const api = useAssistantApi();
-  const isComposer = api.attachment.source === "composer";
+  const aui = useAui();
+  const isComposer = aui.attachment.source === "composer";
 
-  const isImage = useAssistantState(
+  const isImage = useAuiState(
     ({ attachment }) => (attachment as { type?: string })?.type === "image"
   );
 
   // Split into separate selectors to avoid creating new objects on each render
-  const typeLabel = useAssistantState(({ attachment }) => {
+  const typeLabel = useAuiState(({ attachment }) => {
     const att = attachment as {
       type?: string;
       name?: string;
@@ -324,7 +321,7 @@ const AttachmentUI: FC = () => {
     }
   });
 
-  const isUrl = useAssistantState(({ attachment }) => {
+  const isUrl = useAuiState(({ attachment }) => {
     const att = attachment as {
       type?: string;
       name?: string;
@@ -382,7 +379,7 @@ const AttachmentUI: FC = () => {
           "aui-attachment-root-composer only:[&>#attachment-tile]:size-24",
         )}
       >
-        <div className="relative group/attachment">
+        <div className="relative">
           <AttachmentPreviewDialog>
             <TooltipTrigger asChild>
               <div
@@ -419,10 +416,10 @@ const AttachmentRemove: FC = () => {
     <AttachmentPrimitive.Remove asChild>
       <TooltipIconButton
         tooltip={isUrl ? "Remove URL" : "Remove file"}
-        className="aui-attachment-tile-remove absolute top-1.5 right-1.5 size-4 rounded-full bg-black/50 text-white opacity-0 group-hover/attachment:opacity-100 transition-all shadow-sm hover:!bg-zinc-200 hover:text-red-500 hover:scale-110"
+        className="aui-attachment-tile-remove absolute top-1.5 right-1.5 size-3.5 rounded-full bg-white text-muted-foreground opacity-100 shadow-sm hover:bg-white! [&_svg]:text-black hover:[&_svg]:text-destructive"
         side="top"
       >
-        <XIcon className="aui-attachment-remove-icon size-2.5 dark:stroke-[2.5px]" />
+        <XIcon className="aui-attachment-remove-icon size-3 dark:stroke-[2.5px]" />
       </TooltipIconButton>
     </AttachmentPrimitive.Remove>
   );
@@ -454,7 +451,7 @@ export const ComposerAddAttachment: FC = () => {
   const [isDropdownOpen, setIsDropdownOpen] = useState(false);
   const fileInputRef = useRef<HTMLInputElement>(null);
   const containerRef = useRef<HTMLDivElement>(null);
-  const api = useAssistantApi();
+  const aui = useAui();
 
 
 
@@ -511,7 +508,7 @@ export const ComposerAddAttachment: FC = () => {
     // Add valid files
     if (validFiles.length > 0) {
       validFiles.forEach((file) => {
-        api.composer().addAttachment(file);
+        aui.composer().addAttachment(file);
       });
 
       if (validFiles.length < fileArray.length) {
@@ -535,7 +532,7 @@ export const ComposerAddAttachment: FC = () => {
     if (trimmedUrl) {
       try {
         const urlFile = createUrlFile(trimmedUrl);
-        api.composer().addAttachment(urlFile);
+        aui.composer().addAttachment(urlFile);
         setUrlInput("");
         setShowUrlDialog(false);
         focusComposerInput();
diff --git a/src/components/assistant-ui/file.tsx b/src/components/assistant-ui/file.tsx
new file mode 100644
index 00000000..92d93a1b
--- /dev/null
+++ b/src/components/assistant-ui/file.tsx
@@ -0,0 +1,235 @@
+"use client";
+
+import { memo, type FC } from "react";
+import { cva, type VariantProps } from "class-variance-authority";
+import {
+  FileIcon,
+  FileTextIcon,
+  ImageIcon,
+  MusicIcon,
+  VideoIcon,
+  BracesIcon,
+  DownloadIcon,
+} from "lucide-react";
+import type { FileMessagePartComponent } from "@assistant-ui/react";
+import { cn } from "@/lib/utils";
+
+const fileVariants = cva(
+  "aui-file-root inline-flex items-center gap-3 rounded-lg transition-colors",
+  {
+    variants: {
+      variant: {
+        outline: "border border-border hover:bg-muted/50",
+        ghost: "hover:bg-muted/50",
+        muted: "bg-muted/50 hover:bg-muted/70",
+      },
+      size: {
+        sm: "px-2.5 py-1.5 text-xs",
+        default: "px-3 py-2 text-sm",
+        lg: "px-4 py-3 text-base",
+      },
+    },
+    defaultVariants: {
+      variant: "outline",
+      size: "default",
+    },
+  },
+);
+
+function getMimeTypeIcon(mimeType: string): FC<{ className?: string }> {
+  if (mimeType.startsWith("image/")) {
+    return ImageIcon;
+  }
+  if (mimeType === "application/pdf") {
+    return FileTextIcon;
+  }
+  if (mimeType === "application/json") {
+    return BracesIcon;
+  }
+  if (mimeType.startsWith("text/")) {
+    return FileTextIcon;
+  }
+  if (mimeType.startsWith("audio/")) {
+    return MusicIcon;
+  }
+  if (mimeType.startsWith("video/")) {
+    return VideoIcon;
+  }
+  return FileIcon;
+}
+
+function getBase64Size(base64: string): number {
+  const commaIndex = base64.indexOf(",");
+  const base64Data = commaIndex >= 0 ? base64.slice(commaIndex + 1) : base64;
+  const padding = (base64Data.match(/=/g) || []).length;
+  return Math.floor((base64Data.length * 3) / 4) - padding;
+}
+
+function formatFileSize(bytes: number): string {
+  if (bytes < 1024) {
+    return `${bytes} B`;
+  }
+  if (bytes < 1024 * 1024) {
+    return `${(bytes / 1024).toFixed(1)} KB`;
+  }
+  return `${(bytes / (1024 * 1024)).toFixed(1)} MB`;
+}
+
+export type FileRootProps = React.ComponentProps<"div"> &
+  VariantProps<typeof fileVariants>;
+
+function FileRoot({
+  className,
+  variant,
+  size,
+  children,
+  ...props
+}: FileRootProps) {
+  return (
+    <div
+      data-slot="file-root"
+      data-variant={variant}
+      data-size={size}
+      className={cn(fileVariants({ variant, size, className }))}
+      {...props}
+    >
+      {children}
+    </div>
+  );
+}
+
+type FileIconDisplayProps = React.ComponentProps<"span"> & {
+  mimeType?: string;
+};
+
+function FileIconDisplay({
+  mimeType,
+  className,
+  children,
+  ...props
+}: FileIconDisplayProps) {
+  const IconComponent = mimeType ? getMimeTypeIcon(mimeType) : FileIcon;
+
+  return (
+    <span
+      data-slot="file-icon"
+      className={cn("shrink-0 text-muted-foreground", className)}
+      {...props}
+    >
+      {children ?? <IconComponent className="size-5" />}
+    </span>
+  );
+}
+
+function FileName({
+  className,
+  children,
+  ...props
+}: React.ComponentProps<"span">) {
+  return (
+    <span
+      data-slot="file-name"
+      className={cn("min-w-0 flex-1 truncate font-medium", className)}
+      {...props}
+    >
+      {children || "Unnamed file"}
+    </span>
+  );
+}
+
+type FileSizeProps = React.ComponentProps<"span"> & {
+  bytes: number;
+};
+
+function FileSize({ bytes, className, ...props }: FileSizeProps) {
+  return (
+    <span
+      data-slot="file-size"
+      className={cn("shrink-0 text-muted-foreground", className)}
+      {...props}
+    >
+      {formatFileSize(bytes)}
+    </span>
+  );
+}
+
+type FileDownloadProps = Omit<React.ComponentProps<"a">, "href"> & {
+  data: string;
+  mimeType: string;
+  filename?: string;
+};
+
+function FileDownload({
+  data,
+  mimeType,
+  filename,
+  className,
+  children,
+  ...props
+}: FileDownloadProps) {
+  const href = data.startsWith("data:")
+    ? data
+    : `data:${mimeType};base64,${data}`;
+
+  return (
+    <a
+      data-slot="file-download"
+      href={href}
+      download={filename || "download"}
+      className={cn(
+        "shrink-0 rounded-md p-1 text-muted-foreground transition-colors hover:bg-accent hover:text-accent-foreground",
+        className,
+      )}
+      {...props}
+    >
+      {children || <DownloadIcon className="size-4" />}
+    </a>
+  );
+}
+
+const FileImpl: FileMessagePartComponent = ({ filename, data, mimeType }) => {
+  const bytes = getBase64Size(data);
+
+  return (
+    <FileRoot>
+      <FileIconDisplay mimeType={mimeType} />
+      <div className="flex min-w-0 flex-1 flex-col gap-0.5">
+        <FileName>{filename}</FileName>
+        <FileSize bytes={bytes} className="text-xs" />
+      </div>
+      <FileDownload
+        data={data}
+        mimeType={mimeType}
+        {...(filename !== undefined && { filename })}
+      />
+    </FileRoot>
+  );
+};
+
+const File = memo(FileImpl) as unknown as FileMessagePartComponent & {
+  Root: typeof FileRoot;
+  Icon: typeof FileIconDisplay;
+  Name: typeof FileName;
+  Size: typeof FileSize;
+  Download: typeof FileDownload;
+};
+
+File.displayName = "File";
+File.Root = FileRoot;
+File.Icon = FileIconDisplay;
+File.Name = FileName;
+File.Size = FileSize;
+File.Download = FileDownload;
+
+export {
+  File,
+  FileRoot,
+  FileIconDisplay,
+  FileName,
+  FileSize,
+  FileDownload,
+  fileVariants,
+  getMimeTypeIcon,
+  getBase64Size,
+  formatFileSize,
+};
diff --git a/src/components/assistant-ui/image.tsx b/src/components/assistant-ui/image.tsx
new file mode 100644
index 00000000..6f7b6ee6
--- /dev/null
+++ b/src/components/assistant-ui/image.tsx
@@ -0,0 +1,264 @@
+"use client";
+
+import {
+  memo,
+  useState,
+  useEffect,
+  useRef,
+  type PropsWithChildren,
+} from "react";
+import { createPortal } from "react-dom";
+import { cva, type VariantProps } from "class-variance-authority";
+import { ImageIcon, ImageOffIcon } from "lucide-react";
+import type { ImageMessagePartComponent } from "@assistant-ui/react";
+import { cn } from "@/lib/utils";
+
+const imageVariants = cva(
+  "aui-image-root relative overflow-hidden rounded-lg",
+  {
+    variants: {
+      variant: {
+        outline: "border border-border",
+        ghost: "",
+        muted: "bg-muted/50",
+      },
+      size: {
+        sm: "max-w-64",
+        default: "max-w-96",
+        lg: "max-w-[512px]",
+        full: "w-full",
+      },
+    },
+    defaultVariants: {
+      variant: "outline",
+      size: "default",
+    },
+  },
+);
+
+export type ImageRootProps = React.ComponentProps<"div"> &
+  VariantProps<typeof imageVariants>;
+
+function ImageRoot({
+  className,
+  variant,
+  size,
+  children,
+  ...props
+}: ImageRootProps) {
+  return (
+    <div
+      data-slot="image-root"
+      data-variant={variant}
+      data-size={size}
+      className={cn(imageVariants({ variant, size, className }))}
+      {...props}
+    >
+      {children}
+    </div>
+  );
+}
+
+type ImagePreviewProps = Omit<React.ComponentProps<"img">, "children"> & {
+  containerClassName?: string;
+};
+
+function ImagePreview({
+  className,
+  containerClassName,
+  onLoad,
+  onError,
+  alt = "Image content",
+  src,
+  ...props
+}: ImagePreviewProps) {
+  const imgRef = useRef<HTMLImageElement>(null);
+  const [loadedSrc, setLoadedSrc] = useState<string | undefined>(undefined);
+  const [errorSrc, setErrorSrc] = useState<string | undefined>(undefined);
+
+  const loaded = loadedSrc === src;
+  const error = errorSrc === src;
+
+  useEffect(() => {
+    if (
+      typeof src === "string" &&
+      imgRef.current?.complete &&
+      imgRef.current.naturalWidth > 0
+    ) {
+      setLoadedSrc(src);
+    }
+  }, [src]);
+
+  return (
+    <div
+      data-slot="image-preview"
+      className={cn("relative min-h-32", containerClassName)}
+    >
+      {!loaded && !error && (
+        <div
+          data-slot="image-preview-loading"
+          className="absolute inset-0 flex items-center justify-center bg-muted/50"
+        >
+          <ImageIcon className="size-8 animate-pulse text-muted-foreground" />
+        </div>
+      )}
+      {error ? (
+        <div
+          data-slot="image-preview-error"
+          className="flex min-h-32 items-center justify-center bg-muted/50 p-4"
+        >
+          <ImageOffIcon className="size-8 text-muted-foreground" />
+        </div>
+      ) : (
+        <img
+          ref={imgRef}
+          src={src}
+          alt={alt}
+          className={cn(
+            "block h-auto w-full object-contain",
+            !loaded && "invisible",
+            className,
+          )}
+          onLoad={(e) => {
+            if (typeof src === "string") setLoadedSrc(src);
+            onLoad?.(e);
+          }}
+          onError={(e) => {
+            if (typeof src === "string") setErrorSrc(src);
+            onError?.(e);
+          }}
+          {...props}
+        />
+      )}
+    </div>
+  );
+}
+
+function ImageFilename({
+  className,
+  children,
+  ...props
+}: React.ComponentProps<"span">) {
+  if (!children) return null;
+
+  return (
+    <span
+      data-slot="image-filename"
+      className={cn(
+        "block truncate px-2 py-1.5 text-muted-foreground text-xs",
+        className,
+      )}
+      {...props}
+    >
+      {children}
+    </span>
+  );
+}
+
+type ImageZoomProps = PropsWithChildren<{
+  src: string;
+  alt?: string;
+}>;
+
+function ImageZoom({ src, alt = "Image preview", children }: ImageZoomProps) {
+  const [isMounted, setIsMounted] = useState(false);
+  const [isOpen, setIsOpen] = useState(false);
+
+  useEffect(() => {
+    setIsMounted(true);
+  }, []);
+
+  const handleOpen = () => setIsOpen(true);
+  const handleClose = () => setIsOpen(false);
+
+  useEffect(() => {
+    if (!isOpen) return;
+    const handleKeyDown = (e: KeyboardEvent) => {
+      if (e.key === "Escape") setIsOpen(false);
+    };
+    document.addEventListener("keydown", handleKeyDown);
+    return () => document.removeEventListener("keydown", handleKeyDown);
+  }, [isOpen]);
+
+  useEffect(() => {
+    if (!isOpen) return;
+    const originalOverflow = document.body.style.overflow;
+    document.body.style.overflow = "hidden";
+    return () => {
+      document.body.style.overflow = originalOverflow;
+    };
+  }, [isOpen]);
+
+  return (
+    <>
+      <div
+        onClick={handleOpen}
+        onKeyDown={(e) => e.key === "Enter" && handleOpen()}
+        role="button"
+        tabIndex={0}
+        className="aui-image-zoom-trigger cursor-zoom-in"
+        aria-label="Click to zoom image"
+      >
+        {children}
+      </div>
+      {isMounted &&
+        isOpen &&
+        createPortal(
+          <div
+            data-slot="image-zoom-overlay"
+            role="button"
+            tabIndex={0}
+            className="aui-image-zoom-overlay fade-in fixed inset-0 z-50 flex animate-in items-center justify-center bg-black/80 duration-200"
+            onClick={handleClose}
+            onKeyDown={(e) => e.key === "Enter" && handleClose()}
+            aria-label="Close zoomed image"
+          >
+            <img
+              data-slot="image-zoom-content"
+              src={src}
+              alt={alt}
+              className="aui-image-zoom-content fade-in zoom-in-95 max-h-[90vh] max-w-[90vw] animate-in cursor-zoom-out object-contain duration-200"
+              onClick={(e) => {
+                e.stopPropagation();
+                handleClose();
+              }}
+            />
+          </div>,
+          document.body,
+        )}
+    </>
+  );
+}
+
+const ImageImpl: ImageMessagePartComponent = ({ image, filename }) => {
+  return (
+    <ImageRoot>
+      <ImageZoom src={image} alt={filename || "Image content"}>
+        <ImagePreview src={image} alt={filename || "Image content"} />
+      </ImageZoom>
+      <ImageFilename>{filename}</ImageFilename>
+    </ImageRoot>
+  );
+};
+
+const Image = memo(ImageImpl) as unknown as ImageMessagePartComponent & {
+  Root: typeof ImageRoot;
+  Preview: typeof ImagePreview;
+  Filename: typeof ImageFilename;
+  Zoom: typeof ImageZoom;
+};
+
+Image.displayName = "Image";
+Image.Root = ImageRoot;
+Image.Preview = ImagePreview;
+Image.Filename = ImageFilename;
+Image.Zoom = ImageZoom;
+
+export {
+  Image,
+  ImageRoot,
+  ImagePreview,
+  ImageFilename,
+  ImageZoom,
+  imageVariants,
+};
diff --git a/src/components/assistant-ui/markdown-text.tsx b/src/components/assistant-ui/markdown-text.tsx
index bf67c06c..385e5604 100644
--- a/src/components/assistant-ui/markdown-text.tsx
+++ b/src/components/assistant-ui/markdown-text.tsx
@@ -10,7 +10,7 @@ const code = createCodePlugin({
   themes: ['one-dark-pro', 'one-dark-pro'],
 });
 import { useMessagePartText } from "@assistant-ui/react";
-import { useAssistantState } from "@assistant-ui/react";
+import { useAuiState } from "@assistant-ui/react";
 import { cn } from "@/lib/utils";
 import React, { memo, useRef, useEffect } from "react";
 
@@ -19,11 +19,11 @@ const MarkdownTextImpl = () => {
   const { text } = useMessagePartText();
 
   // Get thread and message ID for unique key per message
-  const threadId = useAssistantState(({ threads }) => (threads as any)?.mainThreadId);
-  const messageId = useAssistantState(({ message }) => (message as any)?.id);
+  const threadId = useAuiState(({ threads }) => (threads as any)?.mainThreadId);
+  const messageId = useAuiState(({ message }) => (message as any)?.id);
 
   // Check if the message is currently streaming
-  const isRunning = useAssistantState(({ thread }) => (thread as any)?.isRunning ?? false);
+  const isRunning = useAuiState(({ thread }) => (thread as any)?.isRunning ?? false);
 
   const containerRef = useRef<HTMLDivElement>(null);
 
@@ -144,11 +144,12 @@ const MarkdownTextImpl = () => {
 export const MarkdownText = memo(MarkdownTextImpl);
 
 function normalizeCustomMathTags(input: string): string {
-  return (
-    input
-      // Convert \( ... \) to $$...$$ (inline math) - streamdown uses $$ for both, inline has no newlines
-      .replace(/\\{1,2}\(([\s\S]*?)\\{1,2}\)/g, (_, content) => `$$${content.trim()}$$`)
-      // Convert \[ ... \] to $$...$$ (block math) - streamdown needs newlines for block triggers
-      .replace(/\\{1,2}\[([\s\S]*?)\\{1,2}\]/g, (_, content) => `$$\n${content.trim()}\n$$`)
-  );
+  return (input
+    // Convert \( ... \) to $$...$$ (inline math) - streamdown uses $$ for both, inline has no newlines
+    .replace(/\\{1,2}\(([\s\S]*?)\\{1,2}\)/g, (_, content) => `$$${content.trim()}$$`)
+    // Convert \[ ... \] to $$...$$ (block math) - streamdown needs newlines for block triggers
+    .replace(
+    /\\{1,2}\[([\s\S]*?)\\{1,2}\]/g,
+    (_, content) => `$$\n${content.trim()}\n$$`
+  ));
 }
\ No newline at end of file
diff --git a/src/components/assistant-ui/reasoning.tsx b/src/components/assistant-ui/reasoning.tsx
new file mode 100644
index 00000000..022cfbdb
--- /dev/null
+++ b/src/components/assistant-ui/reasoning.tsx
@@ -0,0 +1,275 @@
+"use client";
+
+import { memo, useCallback, useRef, useState } from "react";
+import { cva, type VariantProps } from "class-variance-authority";
+import { BrainIcon, ChevronDownIcon } from "lucide-react";
+import {
+  useScrollLock,
+  useAuiState,
+  type ReasoningMessagePartComponent,
+  type ReasoningGroupComponent,
+} from "@assistant-ui/react";
+import { MarkdownText } from "@/components/assistant-ui/markdown-text";
+import {
+  Collapsible,
+  CollapsibleContent,
+  CollapsibleTrigger,
+} from "@/components/ui/collapsible";
+import { cn } from "@/lib/utils";
+
+const ANIMATION_DURATION = 200;
+
+const reasoningVariants = cva("aui-reasoning-root mb-4 w-full", {
+  variants: {
+    variant: {
+      outline: "rounded-lg border px-3 py-2",
+      ghost: "",
+      muted: "rounded-lg bg-muted/50 px-3 py-2",
+    },
+  },
+  defaultVariants: {
+    variant: "outline",
+  },
+});
+
+export type ReasoningRootProps = Omit<
+  React.ComponentProps<typeof Collapsible>,
+  "open" | "onOpenChange"
+> &
+  VariantProps<typeof reasoningVariants> & {
+    open?: boolean;
+    onOpenChange?: (open: boolean) => void;
+    defaultOpen?: boolean;
+  };
+
+function ReasoningRoot({
+  className,
+  variant,
+  open: controlledOpen,
+  onOpenChange: controlledOnOpenChange,
+  defaultOpen = false,
+  children,
+  ...props
+}: ReasoningRootProps) {
+  const collapsibleRef = useRef<HTMLDivElement>(null);
+  const [uncontrolledOpen, setUncontrolledOpen] = useState(defaultOpen);
+  const lockScroll = useScrollLock(collapsibleRef, ANIMATION_DURATION);
+
+  const isControlled = controlledOpen !== undefined;
+  const isOpen = isControlled ? controlledOpen : uncontrolledOpen;
+
+  const handleOpenChange = useCallback(
+    (open: boolean) => {
+      if (!open) {
+        lockScroll();
+      }
+      if (!isControlled) {
+        setUncontrolledOpen(open);
+      }
+      controlledOnOpenChange?.(open);
+    },
+    [lockScroll, isControlled, controlledOnOpenChange],
+  );
+
+  return (
+    <Collapsible
+      ref={collapsibleRef}
+      data-slot="reasoning-root"
+      data-variant={variant}
+      open={isOpen}
+      onOpenChange={handleOpenChange}
+      className={cn(
+        "group/reasoning-root",
+        reasoningVariants({ variant, className }),
+      )}
+      style={
+        {
+          "--animation-duration": `${ANIMATION_DURATION}ms`,
+        } as React.CSSProperties
+      }
+      {...props}
+    >
+      {children}
+    </Collapsible>
+  );
+}
+
+function ReasoningFade({ className, ...props }: React.ComponentProps<"div">) {
+  return (
+    <div
+      data-slot="reasoning-fade"
+      className={cn(
+        "aui-reasoning-fade pointer-events-none absolute inset-x-0 bottom-0 z-10 h-8",
+        "bg-[linear-gradient(to_top,var(--color-background),transparent)]",
+        "group-data-[variant=muted]/reasoning-root:bg-[linear-gradient(to_top,hsl(var(--muted)/0.5),transparent)]",
+        "fade-in-0 animate-in",
+        "group-data-[state=open]/collapsible-content:animate-out",
+        "group-data-[state=open]/collapsible-content:fade-out-0",
+        "group-data-[state=open]/collapsible-content:delay-[calc(var(--animation-duration)*0.75)]",
+        "group-data-[state=open]/collapsible-content:fill-mode-forwards",
+        "duration-(--animation-duration)",
+        "group-data-[state=open]/collapsible-content:duration-(--animation-duration)",
+        className,
+      )}
+      {...props}
+    />
+  );
+}
+
+function ReasoningTrigger({
+  active,
+  duration,
+  className,
+  ...props
+}: React.ComponentProps<typeof CollapsibleTrigger> & {
+  active?: boolean;
+  duration?: number;
+}) {
+  const durationText = duration ? ` (${duration}s)` : "";
+
+  return (
+    <CollapsibleTrigger
+      data-slot="reasoning-trigger"
+      className={cn(
+        "aui-reasoning-trigger group/trigger flex max-w-[75%] items-center gap-2 py-1 text-muted-foreground text-sm transition-colors hover:text-foreground",
+        className,
+      )}
+      {...props}
+    >
+      <BrainIcon
+        data-slot="reasoning-trigger-icon"
+        className="aui-reasoning-trigger-icon size-4 shrink-0"
+      />
+      <span
+        data-slot="reasoning-trigger-label"
+        className="aui-reasoning-trigger-label-wrapper relative inline-block leading-none"
+      >
+        <span>Reasoning{durationText}</span>
+        {active ? (
+          <span
+            aria-hidden
+            data-slot="reasoning-trigger-shimmer"
+            className="aui-reasoning-trigger-shimmer shimmer pointer-events-none absolute inset-0 motion-reduce:animate-none"
+          >
+            Reasoning{durationText}
+          </span>
+        ) : null}
+      </span>
+      <ChevronDownIcon
+        data-slot="reasoning-trigger-chevron"
+        className={cn(
+          "aui-reasoning-trigger-chevron mt-0.5 size-4 shrink-0",
+          "transition-transform duration-(--animation-duration) ease-out",
+          "group-data-[state=closed]/trigger:-rotate-90",
+          "group-data-[state=open]/trigger:rotate-0",
+        )}
+      />
+    </CollapsibleTrigger>
+  );
+}
+
+function ReasoningContent({
+  className,
+  children,
+  ...props
+}: React.ComponentProps<typeof CollapsibleContent>) {
+  return (
+    <CollapsibleContent
+      data-slot="reasoning-content"
+      className={cn(
+        "aui-reasoning-content relative overflow-hidden text-muted-foreground text-sm outline-none",
+        "group/collapsible-content ease-out",
+        "data-[state=closed]:animate-collapsible-up",
+        "data-[state=open]:animate-collapsible-down",
+        "data-[state=closed]:fill-mode-forwards",
+        "data-[state=closed]:pointer-events-none",
+        "data-[state=open]:duration-(--animation-duration)",
+        "data-[state=closed]:duration-(--animation-duration)",
+        className,
+      )}
+      {...props}
+    >
+      {children}
+      <ReasoningFade />
+    </CollapsibleContent>
+  );
+}
+
+function ReasoningText({ className, ...props }: React.ComponentProps<"div">) {
+  return (
+    <div
+      data-slot="reasoning-text"
+      className={cn(
+        "aui-reasoning-text relative z-0 max-h-64 overflow-y-auto pt-2 pb-2 pl-6 leading-relaxed",
+        "transform-gpu transition-[transform,opacity]",
+        "group-data-[state=open]/collapsible-content:animate-in",
+        "group-data-[state=closed]/collapsible-content:animate-out",
+        "group-data-[state=open]/collapsible-content:fade-in-0",
+        "group-data-[state=closed]/collapsible-content:fade-out-0",
+        "group-data-[state=open]/collapsible-content:slide-in-from-top-4",
+        "group-data-[state=closed]/collapsible-content:slide-out-to-top-4",
+        "group-data-[state=open]/collapsible-content:duration-(--animation-duration)",
+        "group-data-[state=closed]/collapsible-content:duration-(--animation-duration)",
+        className,
+      )}
+      {...props}
+    />
+  );
+}
+
+const ReasoningImpl: ReasoningMessagePartComponent = () => <MarkdownText />;
+
+const ReasoningGroupImpl: ReasoningGroupComponent = ({
+  children,
+  startIndex,
+  endIndex,
+}) => {
+  const isReasoningStreaming = useAuiState(({ message }) => {
+    if (message.status?.type !== "running") return false;
+    const lastIndex = message.parts.length - 1;
+    if (lastIndex < 0) return false;
+    const lastType = message.parts[lastIndex]?.type;
+    if (lastType !== "reasoning") return false;
+    return lastIndex >= startIndex && lastIndex <= endIndex;
+  });
+
+  return (
+    <ReasoningRoot defaultOpen={isReasoningStreaming}>
+      <ReasoningTrigger active={isReasoningStreaming} />
+      <ReasoningContent aria-busy={isReasoningStreaming}>
+        <ReasoningText>{children}</ReasoningText>
+      </ReasoningContent>
+    </ReasoningRoot>
+  );
+};
+
+const Reasoning = memo(
+  ReasoningImpl,
+) as unknown as ReasoningMessagePartComponent & {
+  Root: typeof ReasoningRoot;
+  Trigger: typeof ReasoningTrigger;
+  Content: typeof ReasoningContent;
+  Text: typeof ReasoningText;
+  Fade: typeof ReasoningFade;
+};
+
+Reasoning.displayName = "Reasoning";
+Reasoning.Root = ReasoningRoot;
+Reasoning.Trigger = ReasoningTrigger;
+Reasoning.Content = ReasoningContent;
+Reasoning.Text = ReasoningText;
+Reasoning.Fade = ReasoningFade;
+
+const ReasoningGroup = memo(ReasoningGroupImpl);
+ReasoningGroup.displayName = "ReasoningGroup";
+
+export {
+  Reasoning,
+  ReasoningGroup,
+  ReasoningRoot,
+  ReasoningTrigger,
+  ReasoningContent,
+  ReasoningText,
+  ReasoningFade,
+  reasoningVariants,
+};
diff --git a/src/components/assistant-ui/sources.tsx b/src/components/assistant-ui/sources.tsx
new file mode 100644
index 00000000..51273b67
--- /dev/null
+++ b/src/components/assistant-ui/sources.tsx
@@ -0,0 +1,148 @@
+"use client";
+
+import { memo, useState } from "react";
+import { Slot } from "@radix-ui/react-slot";
+import { cva, type VariantProps } from "class-variance-authority";
+import type { SourceMessagePartComponent } from "@assistant-ui/react";
+import { cn } from "@/lib/utils";
+
+const extractDomain = (url: string): string => {
+  try {
+    return new URL(url).hostname.replace(/^www\./, "");
+  } catch {
+    return url;
+  }
+};
+
+const getDomainInitial = (url: string): string => {
+  const domain = extractDomain(url);
+  return domain.charAt(0).toUpperCase();
+};
+
+const sourceVariants = cva(
+  "inline-flex cursor-pointer items-center gap-1.5 rounded-md text-xs outline-none transition-colors focus-visible:border-ring focus-visible:ring-[3px] focus-visible:ring-ring/50",
+  {
+    variants: {
+      variant: {
+        outline:
+          "border bg-transparent text-muted-foreground hover:bg-accent hover:text-accent-foreground",
+        ghost: "hover:bg-accent hover:text-accent-foreground",
+        muted: "bg-secondary text-secondary-foreground hover:bg-secondary/70",
+      },
+      size: {
+        default: "px-2 py-1",
+        sm: "px-1.5 py-0.5",
+        lg: "px-2.5 py-1.5",
+      },
+    },
+    defaultVariants: {
+      variant: "outline",
+      size: "default",
+    },
+  },
+);
+
+function SourceIcon({
+  url,
+  className,
+  ...props
+}: React.ComponentProps<"span"> & { url: string }) {
+  const [hasError, setHasError] = useState(false);
+  const domain = extractDomain(url);
+
+  if (hasError) {
+    return (
+      <span
+        data-slot="source-icon-fallback"
+        className={cn(
+          "flex size-3.5 shrink-0 items-center justify-center rounded-sm bg-muted font-medium text-[10px]",
+          className,
+        )}
+        {...props}
+      >
+        {getDomainInitial(url)}
+      </span>
+    );
+  }
+
+  return (
+    <img
+      data-slot="source-icon"
+      src={`https://www.google.com/s2/favicons?domain=${domain}&sz=32`}
+      alt=""
+      className={cn("size-3.5 shrink-0 rounded-sm", className)}
+      onError={() => setHasError(true)}
+      {...(props as React.ComponentProps<"img">)}
+    />
+  );
+}
+
+function SourceTitle({ className, ...props }: React.ComponentProps<"span">) {
+  return (
+    <span
+      data-slot="source-title"
+      className={cn("max-w-37.5 truncate", className)}
+      {...props}
+    />
+  );
+}
+
+export type SourceProps = React.ComponentProps<"a"> &
+  VariantProps<typeof sourceVariants> & {
+    asChild?: boolean;
+  };
+
+function Source({
+  className,
+  variant,
+  size,
+  asChild = false,
+  target = "_blank",
+  rel = "noopener noreferrer",
+  ...props
+}: SourceProps) {
+  const Comp = asChild ? Slot : "a";
+
+  return (
+    <Comp
+      data-slot="source"
+      data-variant={variant}
+      data-size={size}
+      target={target}
+      rel={rel}
+      className={cn(sourceVariants({ variant, size, className }))}
+      {...props}
+    />
+  );
+}
+
+const SourcesImpl: SourceMessagePartComponent = ({
+  url,
+  title,
+  sourceType,
+}) => {
+  if (sourceType !== "url" || !url) return null;
+
+  const domain = extractDomain(url);
+  const displayTitle = title || domain;
+
+  return (
+    <Source href={url} target="_blank" rel="noopener noreferrer">
+      <SourceIcon url={url} />
+      <SourceTitle>{displayTitle}</SourceTitle>
+    </Source>
+  );
+};
+
+const Sources = memo(SourcesImpl) as unknown as SourceMessagePartComponent & {
+  Root: typeof Source;
+  Icon: typeof SourceIcon;
+  Title: typeof SourceTitle;
+};
+
+Sources.displayName = "Sources";
+Sources.Root = Source;
+Sources.Icon = SourceIcon;
+Sources.Title = SourceTitle;
+
+export { Sources, Source, SourceIcon, SourceTitle, sourceVariants };
diff --git a/src/components/assistant-ui/thread-list-dropdown.tsx b/src/components/assistant-ui/thread-list-dropdown.tsx
index b39c6608..4db92c8d 100644
--- a/src/components/assistant-ui/thread-list-dropdown.tsx
+++ b/src/components/assistant-ui/thread-list-dropdown.tsx
@@ -3,10 +3,10 @@
 import type { FC } from "react";
 import { useState, useRef, useEffect } from "react";
 import {
-  AssistantIf,
+  AuiIf,
   ThreadListItemPrimitive,
   ThreadListPrimitive,
-  useAssistantApi,
+  useAui,
 } from "@assistant-ui/react";
 import { Trash2Icon, PencilIcon } from "lucide-react";
 import { useThreadListItem } from "@assistant-ui/react";
@@ -59,12 +59,12 @@ export const ThreadListDropdown: FC<ThreadListDropdownProps> = ({ trigger }) =>
           <ThreadListNew onSelect={() => setOpen(false)} />
           <DropdownMenuSeparator className="bg-sidebar-border m-0" />
           <div className="flex-1 overflow-y-auto max-h-[400px] p-1">
-            <AssistantIf condition={({ threads }) => threads.isLoading}>
+            <AuiIf condition={({ threads }) => threads.isLoading}>
               <ThreadListSkeleton />
-            </AssistantIf>
-            <AssistantIf condition={({ threads }) => !threads.isLoading}>
+            </AuiIf>
+            <AuiIf condition={({ threads }) => !threads.isLoading}>
               <ThreadListPrimitive.Items components={{ ThreadListItem: () => <ThreadListItem onSelect={() => setOpen(false)} /> }} />
-            </AssistantIf>
+            </AuiIf>
           </div>
         </ThreadListPrimitive.Root>
       </DropdownMenuContent>
@@ -116,7 +116,7 @@ const ThreadListItemContent: FC<{ onSelect?: () => void }> = ({ onSelect }) => {
   const [isEditing, setIsEditing] = useState(false);
   const [editValue, setEditValue] = useState("");
   const inputRef = useRef<HTMLInputElement>(null);
-  const api = useAssistantApi();
+  const aui = useAui();
 
   // Get the current title and thread state - this is now inside the Root context
   // Using safe hook to handle race condition during thread switching (GitHub issue #2722)
@@ -162,7 +162,7 @@ const ThreadListItemContent: FC<{ onSelect?: () => void }> = ({ onSelect }) => {
 
     if (trimmedValue && trimmedValue !== title) {
       try {
-        await api?.threadListItem().rename(trimmedValue);
+        await aui?.threadListItem().rename(trimmedValue);
         toast.success("Title updated");
       } catch (error) {
         console.error("Failed to rename thread:", error);
diff --git a/src/components/assistant-ui/thread.tsx b/src/components/assistant-ui/thread.tsx
index dd4447f9..5180f8c6 100644
--- a/src/components/assistant-ui/thread.tsx
+++ b/src/components/assistant-ui/thread.tsx
@@ -24,15 +24,15 @@ import { LuSparkle } from "react-icons/lu";
 import { cn } from "@/lib/utils";
 import {
   ActionBarPrimitive,
-  AssistantIf,
+  AuiIf,
   BranchPickerPrimitive,
   ComposerPrimitive,
   ErrorPrimitive,
   MessagePrimitive,
   ThreadPrimitive,
-  useAssistantApi,
+  useAui,
   useMessagePartText,
-  useAssistantState,
+  useAuiState,
 } from "@assistant-ui/react";
 
 
@@ -77,6 +77,10 @@ import {
 } from "@/components/assistant-ui/attachment";
 import { SelectCardsToolUI } from "@/components/assistant-ui/SelectCardsToolUI";
 import { AssistantLoader } from "@/components/assistant-ui/assistant-loader";
+import { File as FileComponent } from "@/components/assistant-ui/file";
+import { Sources } from "@/components/assistant-ui/sources";
+import { Image } from "@/components/assistant-ui/image";
+import { Reasoning, ReasoningGroup } from "@/components/assistant-ui/reasoning";
 
 import type { Item, PdfData } from "@/lib/workspace-state/types";
 import { CardContextDisplay } from "@/components/chat/CardContextDisplay";
@@ -143,9 +147,9 @@ export const Thread: FC<ThreadProps> = ({ items = [] }) => {
             autoScroll={false}
             className="aui-thread-viewport relative flex flex-1 flex-col overflow-x-auto overflow-y-scroll px-4"
           >
-            <AssistantIf condition={({ thread }) => thread.isEmpty}>
+            <AuiIf condition={({ thread }) => thread.isEmpty}>
               <ThreadWelcome />
-            </AssistantIf>
+            </AuiIf>
 
             <ThreadPrimitive.Messages
               components={{
@@ -287,7 +291,7 @@ interface ComposerProps {
 
 const Composer: FC<ComposerProps> = ({ items }) => {
   const currentWorkspaceId = useWorkspaceStore((state) => state.currentWorkspaceId);
-  const api = useAssistantApi();
+  const aui = useAui();
   const replySelections = useUIStore(useShallow(selectReplySelections));
   const clearReplySelections = useUIStore((state) => state.clearReplySelections);
   const clearSelectedActions = useUIStore((state) => state.clearSelectedActions);
@@ -389,7 +393,7 @@ const Composer: FC<ComposerProps> = ({ items }) => {
       const newValue = textBefore + textAfter;
 
       // Update the textarea value
-      api.composer().setText(newValue);
+      aui?.composer()?.setText(newValue);
 
       // Reset mention state
       setMentionQuery("");
@@ -404,7 +408,7 @@ const Composer: FC<ComposerProps> = ({ items }) => {
         }
       }, 0);
     }
-  }, [mentionStartIndex, api]);
+  }, [mentionStartIndex, aui]);
 
   // Handle mention selection - select item and close menu
   const handleMentionSelect = useCallback((item: Item) => {
@@ -439,7 +443,7 @@ const Composer: FC<ComposerProps> = ({ items }) => {
 
       if (fileToUpload) {
         try {
-          await api.composer().addAttachment(fileToUpload);
+          await aui?.composer()?.addAttachment(fileToUpload);
         } catch (error) {
           console.error("Failed to add file attachment:", error);
         }
@@ -457,7 +461,7 @@ const Composer: FC<ComposerProps> = ({ items }) => {
       if (file) {
         try {
           // Add as attachment - will be uploaded when message is sent
-          await api.composer().addAttachment(file);
+          await aui?.composer()?.addAttachment(file);
         } catch (error) {
           console.error("Failed to add image attachment:", error);
         }
@@ -481,7 +485,7 @@ const Composer: FC<ComposerProps> = ({ items }) => {
             // Create a virtual File object from the URL
             const urlFile = createUrlFile(url);
             // Add it as an attachment
-            await api.composer().addAttachment(urlFile);
+            await aui?.composer()?.addAttachment(urlFile);
           } catch (error) {
             console.error("Failed to add URL attachment:", error);
           }
@@ -536,7 +540,9 @@ const Composer: FC<ComposerProps> = ({ items }) => {
         e.preventDefault();
 
         // Get the current composer state
-        const composerState = api.composer().getState();
+        const composerState = aui?.composer()?.getState();
+        if (!composerState) return;
+        
         const currentText = composerState.text;
         const attachments = composerState.attachments || [];
 
@@ -632,8 +638,8 @@ const Composer: FC<ComposerProps> = ({ items }) => {
         let modifiedText = currentText;
 
         // Set the modified text and send
-        api.composer().setText(modifiedText);
-        api.composer().send();
+        aui?.composer()?.setText(modifiedText);
+        aui?.composer()?.send();
 
         // Clear reply selections immediately (these are added to the message itself)
         clearReplySelections();
@@ -687,7 +693,7 @@ interface ComposerActionProps {
 
 const ComposerAction: FC<ComposerActionProps> = ({ items }) => {
   const { data: session } = useSession();
-  const api = useAssistantApi();
+  useAui();
   const isAnonymous = session?.user?.isAnonymous ?? false;
   const selectedCardIdsArray = useUIStore(
     useShallow(selectSelectedCardIdsArray)
@@ -964,11 +970,10 @@ const ComposerAction: FC<ComposerActionProps> = ({ items }) => {
           </Popover>
         )}
       </div>
-
       {/* Right side: speech/send/cancel button */}
       <div className="flex items-center gap-2">
         {!isAnonymous && <SpeechToTextButton />}
-        <AssistantIf condition={({ thread }) => !thread.isRunning}>
+        <AuiIf condition={({ thread }) => !thread.isRunning}>
           <ComposerPrimitive.Send asChild>
             <TooltipIconButton
               tooltip="Send message"
@@ -982,9 +987,9 @@ const ComposerAction: FC<ComposerActionProps> = ({ items }) => {
               <ArrowUpIcon className="aui-composer-send-icon size-4 text-gray-900 dark:text-gray-600" />
             </TooltipIconButton>
           </ComposerPrimitive.Send>
-        </AssistantIf>
+        </AuiIf>
 
-        <AssistantIf condition={({ thread }) => thread.isRunning}>
+        <AuiIf condition={({ thread }) => thread.isRunning}>
           <ComposerPrimitive.Cancel asChild>
             <Button
               type="button"
@@ -996,7 +1001,7 @@ const ComposerAction: FC<ComposerActionProps> = ({ items }) => {
               <Square className="aui-composer-cancel-icon size-3 text-gray-900 dark:text-gray-600 fill-current" />
             </Button>
           </ComposerPrimitive.Cancel>
-        </AssistantIf>
+        </AuiIf>
       </div>
     </div >
   );
@@ -1024,6 +1029,11 @@ const AssistantMessage: FC = () => {
           <MessagePrimitive.Parts
             components={{
               Text: MarkdownText,
+              File: FileComponent,
+              Source: Sources,
+              Image,
+              Reasoning: Reasoning,
+              ReasoningGroup: ReasoningGroup,
               tools: {
                 Fallback: ToolFallback,
               },
@@ -1407,6 +1417,7 @@ const UserMessage: FC = () => {
             <MessagePrimitive.Parts
               components={{
                 Text: UserMessageText,
+                File: FileComponent,
               }}
             />
           </div>
@@ -1626,8 +1637,8 @@ const truncateText = (text: string, maxLength: number = 30) => {
 };
 
 const EditComposer: FC = () => {
-  const api = useAssistantApi();
-  const messageAttachments = useAssistantState(
+  const aui = useAui();
+  const messageAttachments = useAuiState(
     useShallow(({ message }) => (message as { attachments?: unknown[] })?.attachments || [])
   );
   const hasParsedRef = useRef(false);
@@ -1644,7 +1655,7 @@ const EditComposer: FC = () => {
     hasParsedRef.current = false;
     hasAttachmentsRestoredRef.current = false;
 
-    const composerState = api.composer().getState();
+    const composerState = aui?.composer()?.getState();
 
     if (!composerState || !composerState.text) return;
 
@@ -1671,7 +1682,7 @@ const EditComposer: FC = () => {
 
     // Update input text to show only clean text (without markers)
     if (cleanText !== composerState.text) {
-      api.composer().setText(cleanText);
+      aui?.composer()?.setText(cleanText);
     }
 
     hasParsedRef.current = true;
@@ -1711,7 +1722,7 @@ const EditComposer: FC = () => {
                 }
               }
 
-              api.composer().addAttachment(file);
+              aui?.composer()?.addAttachment(file);
             } catch (error) {
               console.error("Failed to restore file attachment:", error);
             }
@@ -1742,7 +1753,7 @@ const EditComposer: FC = () => {
               if (urlMatch && urlMatch[1]) {
                 try {
                   const urlFile = createUrlFile(urlMatch[1]);
-                  api.composer().addAttachment(urlFile);
+                  aui?.composer()?.addAttachment(urlFile);
                 } catch (error) {
                   console.error("Failed to restore URL attachment:", error);
                 }
@@ -1752,7 +1763,7 @@ const EditComposer: FC = () => {
         });
       }
     }
-  }, [api, messageAttachments]);
+  }, [aui, messageAttachments]);
 
   return (
     <div className="aui-edit-composer-wrapper mx-auto flex w-full max-w-[var(--thread-max-width)] flex-col gap-4 px-2 first:mt-4 mb-4">
@@ -1762,7 +1773,9 @@ const EditComposer: FC = () => {
           e.preventDefault();
 
           // Get the current composer state
-          const composerState = api.composer().getState();
+          const composerState = aui?.composer()?.getState();
+          if (!composerState) return;
+          
           const currentText = composerState.text;
 
           // Re-add URL markers from parsed URLs (stored in state)
@@ -1787,8 +1800,8 @@ const EditComposer: FC = () => {
           }
 
           // Set the modified text and send
-          api.composer().setText(modifiedText);
-          api.composer().send();
+          aui?.composer()?.setText(modifiedText);
+          aui?.composer()?.send();
         }}
       >
         {/* Attachment Display - shows restored and new attachments */}
diff --git a/src/components/chat/AppChatHeader.tsx b/src/components/chat/AppChatHeader.tsx
index b17877ba..421b0bcb 100644
--- a/src/components/chat/AppChatHeader.tsx
+++ b/src/components/chat/AppChatHeader.tsx
@@ -3,7 +3,7 @@
 import { X, ChevronDown, Trash2, Edit2, Check, MessageSquarePlus } from "lucide-react";
 import { RiChatHistoryLine } from "react-icons/ri";
 import { LuMaximize, LuMinimize, LuPanelRightClose } from "react-icons/lu";
-import { useAssistantApi, useThreadListItem } from "@assistant-ui/react";
+import { useAui, useThreadListItem } from "@assistant-ui/react";
 import {
   DropdownMenu,
   DropdownMenuContent,
@@ -59,7 +59,7 @@ export function AppChatHeader({
   const [isEditingCurrentTitle, setIsEditingCurrentTitle] = useState(false);
   const [currentTitleEditValue, setCurrentTitleEditValue] = useState("");
   const currentTitleTextareaRef = useRef<HTMLTextAreaElement>(null);
-  const api = useAssistantApi();
+  const aui = useAui();
 
   // Get UI store state for sidebar behavior
   const openPanelIds = useUIStore((state) => state.openPanelIds);
@@ -158,7 +158,7 @@ export function AppChatHeader({
 
     if (trimmedValue && trimmedValue !== currentThreadTitle) {
       try {
-        await api?.threadListItem().rename(trimmedValue);
+        await aui?.threadListItem().rename(trimmedValue);
         toast.success("Title updated");
       } catch (error) {
         console.error("Failed to rename thread:", error);
diff --git a/src/components/pdf/PdfPanelHeader.tsx b/src/components/pdf/PdfPanelHeader.tsx
index 56b252c2..cc610b22 100644
--- a/src/components/pdf/PdfPanelHeader.tsx
+++ b/src/components/pdf/PdfPanelHeader.tsx
@@ -15,7 +15,7 @@ import {
 } from '@/components/ui/dropdown-menu';
 import ChatFloatingButton from "@/components/chat/ChatFloatingButton";
 import { useUIStore } from "@/lib/stores/ui-store";
-import { useAssistantApi } from "@assistant-ui/react";
+import { useAui } from "@assistant-ui/react";
 import { toast } from "sonner";
 
 // PDF Plugin imports
@@ -193,7 +193,7 @@ export function PdfPanelHeader({
     const { provides: searchProvides, state: searchState } = useSearch(documentId);
     const { provides: capture, state: captureState } = useCapture(documentId);
 
-    const api = useAssistantApi();
+    const aui = useAui();
 
     const isChatExpanded = useUIStore((state) => state.isChatExpanded);
     const setIsChatExpanded = useUIStore((state) => state.setIsChatExpanded);
@@ -232,7 +232,7 @@ export function PdfPanelHeader({
                 const file = new File([result.blob], filename, { type: result.imageType });
 
                 // Add attachment to composer
-                await api.composer().addAttachment(file);
+                await aui.composer().addAttachment(file);
                 toast.success("Screenshot added to chat");
 
                 // Turn off capture mode
@@ -249,7 +249,7 @@ export function PdfPanelHeader({
         return () => {
             unsubscribe();
         };
-    }, [capture, api, captureState.isMarqueeCaptureActive]);
+    }, [capture, aui, captureState.isMarqueeCaptureActive]);
 
 
     const zoomPercent = zoomState?.currentZoomLevel
diff --git a/src/components/workspace-canvas/SidebarQuickActions.tsx b/src/components/workspace-canvas/SidebarQuickActions.tsx
index f60e941d..b783747a 100644
--- a/src/components/workspace-canvas/SidebarQuickActions.tsx
+++ b/src/components/workspace-canvas/SidebarQuickActions.tsx
@@ -13,7 +13,7 @@ import {
 import { useWorkspaceOperations } from "@/hooks/workspace/use-workspace-operations";
 import { useWorkspaceState } from "@/hooks/workspace/use-workspace-state";
 import { useUIStore } from "@/lib/stores/ui-store";
-import { useAssistantApi } from "@assistant-ui/react";
+import { useAui } from "@assistant-ui/react";
 import { useRef, useState, useCallback } from "react";
 import { toast } from "sonner";
 import type { PdfData } from "@/lib/workspace-state/types";
@@ -28,7 +28,7 @@ interface SidebarQuickActionsProps {
 
 export function SidebarQuickActions({ currentWorkspaceId, isChatExpanded, setIsChatExpanded }: SidebarQuickActionsProps) {
     // Get assistant API
-    const api = useAssistantApi();
+    const aui = useAui();
     // Get workspace state and operations
     const { state: workspaceState } = useWorkspaceState(currentWorkspaceId);
     const operations = useWorkspaceOperations(currentWorkspaceId, workspaceState);
@@ -89,7 +89,7 @@ export function SidebarQuickActions({ currentWorkspaceId, isChatExpanded, setIsC
         setSelectedActions(['deep-research']);
 
 
-    }, [api, isChatExpanded, setIsChatExpanded, setSelectedActions]);
+    }, [aui, isChatExpanded, setIsChatExpanded, setSelectedActions]);
 
     const handleQuizClick = useCallback(() => {
         // Open chat if closed
@@ -98,13 +98,13 @@ export function SidebarQuickActions({ currentWorkspaceId, isChatExpanded, setIsC
         }
 
         // Fill composer with quiz creation prompt
-        api.composer().setText("Create a quiz about ");
+        aui.composer().setText("Create a quiz about ");
 
         // Focus the composer input
         focusComposerInput();
 
         toast.success("Quiz creation started");
-    }, [api, isChatExpanded, setIsChatExpanded]);
+    }, [aui, isChatExpanded, setIsChatExpanded]);
 
     // PDF Upload Logic (Adapted from WorkspaceSection/BottomActionBar)
     const handlePDFUpload = async (event: React.ChangeEvent<HTMLInputElement>) => {
diff --git a/src/components/workspace-canvas/WorkspaceContent.tsx b/src/components/workspace-canvas/WorkspaceContent.tsx
index 361dd702..9665ecde 100644
--- a/src/components/workspace-canvas/WorkspaceContent.tsx
+++ b/src/components/workspace-canvas/WorkspaceContent.tsx
@@ -9,7 +9,7 @@ import { WorkspaceGrid } from "./WorkspaceGrid";
 import type { LayoutItem } from "react-grid-layout";
 import { useUIStore, selectSelectedCardIdsArray } from "@/lib/stores/ui-store";
 import { useShallow } from "zustand/react/shallow";
-import { useAssistantApi } from "@assistant-ui/react";
+import { useAui } from "@assistant-ui/react";
 import { toast } from "sonner";
 
 interface WorkspaceContentProps {
@@ -93,7 +93,7 @@ export default function WorkspaceContent({
 
   // File upload for empty state
   const fileInputRef = useRef<HTMLInputElement>(null);
-  const api = useAssistantApi();
+  const aui = useAui();
 
   // Handle copy JSON to clipboard
   const handleCopyJson = useCallback(async () => {
@@ -227,7 +227,7 @@ export default function WorkspaceContent({
       // Add all valid files to composer
       for (const file of validFiles) {
         try {
-          await api.composer().addAttachment(file);
+          await aui.composer().addAttachment(file);
         } catch (error) {
           console.error("Failed to add attachment:", error);
         }
@@ -238,7 +238,7 @@ export default function WorkspaceContent({
         fileInputRef.current.value = "";
       }
     },
-    [api]
+    [aui]
   );
 
   // Handle drag start
diff --git a/src/components/workspace-canvas/WorkspaceHeader.tsx b/src/components/workspace-canvas/WorkspaceHeader.tsx
index a3de7c89..03d7389d 100644
--- a/src/components/workspace-canvas/WorkspaceHeader.tsx
+++ b/src/components/workspace-canvas/WorkspaceHeader.tsx
@@ -11,7 +11,7 @@ import WorkspaceSaveIndicator from "@/components/workspace/WorkspaceSaveIndicato
 import ChatFloatingButton from "@/components/chat/ChatFloatingButton";
 import { useUIStore } from "@/lib/stores/ui-store";
 import { IconRenderer } from "@/hooks/use-icon-picker";
-import { useAssistantApi } from "@assistant-ui/react";
+import { useAui } from "@assistant-ui/react";
 import { focusComposerInput } from "@/lib/utils/composer-utils";
 import {
   DropdownMenu,
@@ -125,7 +125,7 @@ export default function WorkspaceHeader({
   const [ellipsisDropdownOpen, setEllipsisDropdownOpen] = useState(false);
 
   // Assistant API for Deep Research action
-  const api = useAssistantApi();
+  const aui = useAui();
   const setSelectedActions = useUIStore((state) => state.setSelectedActions);
 
 
@@ -409,65 +409,61 @@ export default function WorkspaceHeader({
               </button>
             ) : !activeFolderId ? (
               // When at root level, show dropdown menu on click
-              (onOpenSettings || onOpenShare) ? (
-                <DropdownMenu>
-                  <DropdownMenuTrigger asChild>
-                    <button 
-                      data-breadcrumb-target="root"
-                      className={cn(
-                        "flex items-center gap-1.5 min-w-0 rounded transition-colors hover:bg-sidebar-accent cursor-pointer px-1 py-0.5 -mx-1 -my-0.5",
-                        hoveredBreadcrumbTarget === 'root' && "border-2 border-blue-500 bg-blue-500/10 rounded"
-                      )}
-                    >
-                      <IconRenderer
-                        icon={workspaceIcon}
-                        className="h-4 w-4 shrink-0"
-                        style={{ color: workspaceColor || undefined }}
-                      />
-                      <span className="truncate text-sidebar-foreground font-medium max-w-[200px]" title={workspaceName}>
-                        {workspaceName || "Untitled"}
-                      </span>
-                    </button>
-                  </DropdownMenuTrigger>
-                  <DropdownMenuContent align="start" className="w-48">
-                    {onOpenSettings && (
-                      <DropdownMenuItem
-                        onClick={onOpenSettings}
-                        className="flex items-center gap-2 cursor-pointer"
-                      >
-                        <Settings className="h-4 w-4" />
-                        Settings
-                      </DropdownMenuItem>
-                    )}
-                    {onOpenShare && (
-                      <DropdownMenuItem
-                        onClick={onOpenShare}
-                        className="flex items-center gap-2 cursor-pointer"
-                      >
-                        <Share2 className="h-4 w-4" />
-                        Share
-                      </DropdownMenuItem>
+              ((onOpenSettings || onOpenShare) ? (<DropdownMenu>
+                <DropdownMenuTrigger asChild>
+                  <button 
+                    data-breadcrumb-target="root"
+                    className={cn(
+                      "flex items-center gap-1.5 min-w-0 rounded transition-colors hover:bg-sidebar-accent cursor-pointer px-1 py-0.5 -mx-1 -my-0.5",
+                      hoveredBreadcrumbTarget === 'root' && "border-2 border-blue-500 bg-blue-500/10 rounded"
                     )}
-                  </DropdownMenuContent>
-                </DropdownMenu>
-              ) : (
-                <div 
-                  data-breadcrumb-target="root"
-                  className={cn(
-                    "flex items-center gap-1.5 min-w-0",
-                    hoveredBreadcrumbTarget === 'root' && "border-2 border-blue-500 bg-blue-500/10 rounded px-1 py-0.5 -mx-1 -my-0.5"
+                  >
+                    <IconRenderer
+                      icon={workspaceIcon}
+                      className="h-4 w-4 shrink-0"
+                      style={{ color: workspaceColor || undefined }}
+                    />
+                    <span className="truncate text-sidebar-foreground font-medium max-w-[200px]" title={workspaceName}>
+                      {workspaceName || "Untitled"}
+                    </span>
+                  </button>
+                </DropdownMenuTrigger>
+                <DropdownMenuContent align="start" className="w-48">
+                  {onOpenSettings && (
+                    <DropdownMenuItem
+                      onClick={onOpenSettings}
+                      className="flex items-center gap-2 cursor-pointer"
+                    >
+                      <Settings className="h-4 w-4" />
+                      Settings
+                    </DropdownMenuItem>
                   )}
-                >
-                  <IconRenderer
-                    icon={workspaceIcon}
-                    className="h-4 w-4 shrink-0"
-                    style={{ color: workspaceColor || undefined }}
-                  />
-                  <span className="truncate text-sidebar-foreground font-medium max-w-[200px]" title={workspaceName}>
-                    {workspaceName || "Untitled"}
-                  </span>
-                </div>
-              )
+                  {onOpenShare && (
+                    <DropdownMenuItem
+                      onClick={onOpenShare}
+                      className="flex items-center gap-2 cursor-pointer"
+                    >
+                      <Share2 className="h-4 w-4" />
+                      Share
+                    </DropdownMenuItem>
+                  )}
+                </DropdownMenuContent>
+              </DropdownMenu>) : (<div 
+                data-breadcrumb-target="root"
+                className={cn(
+                  "flex items-center gap-1.5 min-w-0",
+                  hoveredBreadcrumbTarget === 'root' && "border-2 border-blue-500 bg-blue-500/10 rounded px-1 py-0.5 -mx-1 -my-0.5"
+                )}
+              >
+                <IconRenderer
+                  icon={workspaceIcon}
+                  className="h-4 w-4 shrink-0"
+                  style={{ color: workspaceColor || undefined }}
+                />
+                <span className="truncate text-sidebar-foreground font-medium max-w-[200px]" title={workspaceName}>
+                  {workspaceName || "Untitled"}
+                </span>
+              </div>))
             ) : null}
 
             {/* Folder path breadcrumbs - compact mode shows dropdown only */}
@@ -475,7 +471,7 @@ export default function WorkspaceHeader({
               <>
                 {isCompactMode ? (
                   /* Compact mode: Show dropdown with current folder only, full path in dropdown */
-                  <>
+                  (<>
                     <span className="text-sidebar-foreground/50 mx-1 font-bold">/</span>
                     <DropdownMenu>
                       <DropdownMenuTrigger asChild>
@@ -537,7 +533,7 @@ export default function WorkspaceHeader({
                         ))}
                       </DropdownMenuContent>
                     </DropdownMenu>
-                  </>
+                  </>)
                 ) : folderPath.length === 1 ? (
                   folderPath.map((folder) => (
                     <span key={folder.id} className="flex items-center gap-1.5">
@@ -563,7 +559,7 @@ export default function WorkspaceHeader({
                   ))
                 ) : (
                   /* Show root, dropdown with all middle folders, and last for 2+ levels */
-                  <>
+                  (<>
                     <span className="text-sidebar-foreground/50 mx-1 font-bold">/</span>
                     <HoverCard 
                       open={ellipsisDropdownOpen} 
@@ -624,7 +620,7 @@ export default function WorkspaceHeader({
                         {folderPath[folderPath.length - 1].name}
                       </span>
                     </button>
-                  </>
+                  </>)
                 )}
               </>
             )}
@@ -799,7 +795,7 @@ export default function WorkspaceHeader({
                           setIsChatExpanded(true);
                         }
                         // Fill composer with quiz creation prompt
-                        api?.composer().setText("Create a quiz about ");
+                        aui?.composer().setText("Create a quiz about ");
                         // Focus the composer input
                         focusComposerInput();
                         toast.success("Quiz creation started");
@@ -830,7 +826,7 @@ export default function WorkspaceHeader({
                       onClick={() => {
                         toast.success("Deep Research action selected");
                         setSelectedActions(["deep-research"]);
-                        api?.composer().setText("I want to do research on ");
+                        aui?.composer().setText("I want to do research on ");
                         if (setIsChatExpanded && !isChatExpanded) {
                           setIsChatExpanded(true);
                         }
@@ -867,7 +863,6 @@ export default function WorkspaceHeader({
           ) : null}
         </div>
       </div>
-
       {/* Rename Folder Dialog */}
       {
         onRenameFolder && (
@@ -906,7 +901,6 @@ export default function WorkspaceHeader({
           </Dialog>
         )
       }
-
       {/* YouTube Dialog */}
       <CreateYouTubeDialog
         open={showYouTubeDialog}
diff --git a/src/components/workspace-canvas/WorkspaceSection.tsx b/src/components/workspace-canvas/WorkspaceSection.tsx
index 31b86617..2f6c258a 100644
--- a/src/components/workspace-canvas/WorkspaceSection.tsx
+++ b/src/components/workspace-canvas/WorkspaceSection.tsx
@@ -41,7 +41,7 @@ import ShareWorkspaceDialog from "@/components/workspace/ShareWorkspaceDialog";
 import { CreateYouTubeDialog } from "@/components/modals/CreateYouTubeDialog";
 import { useWorkspaceContext } from "@/contexts/WorkspaceContext";
 import type { WorkspaceWithState } from "@/lib/workspace-state/types";
-import { useAssistantApi } from "@assistant-ui/react";
+import { useAui } from "@assistant-ui/react";
 import { focusComposerInput } from "@/lib/utils/composer-utils";
 
 interface WorkspaceSectionProps {
@@ -156,7 +156,7 @@ export function WorkspaceSection({
 
   // Assistant API for Deep Research action
   // Note: WorkspaceSection is inside WorkspaceRuntimeProvider in DashboardLayout, so this hook works
-  const api = useAssistantApi();
+  const aui = useAui();
 
 
   // Get active folder info from UI store
@@ -442,7 +442,6 @@ export function WorkspaceSection({
           isItemPanelOpen={isItemPanelOpen}
         />
       )}
-
       <ContextMenu>
         <ContextMenuTrigger asChild>
           <div ref={scrollAreaRef} className="flex-1 overflow-auto scrollbar-thin scrollbar-thumb-gray-600 scrollbar-track-transparent">
@@ -476,7 +475,7 @@ export function WorkspaceSection({
                 <WorkspaceSkeleton />
               ) : (
                 /* Workspace content - assumes workspace exists (home route handles no-workspace state) */
-                <WorkspaceContent
+                (<WorkspaceContent
                   key={`workspace-content-${state.workspaceId || 'none'}`}
                   viewState={state}
                   showJsonView={showJsonView}
@@ -498,7 +497,7 @@ export function WorkspaceSection({
                   onMoveItems={operations?.moveItemsToFolder}
                   onDeleteFolderWithContents={operations?.deleteFolderWithContents}
                   onPDFUpload={handlePDFUpload}
-                />
+                />)
               )}
 
               {/* Marquee selector for rectangular card selection - inside scroll container to capture all events */}
@@ -579,7 +578,7 @@ export function WorkspaceSection({
                       setIsChatExpanded(true);
                     }
                     // Fill composer with quiz creation prompt
-                    api.composer().setText("Create a quiz about ");
+                    aui.composer().setText("Create a quiz about ");
                     // Focus the composer input
                     focusComposerInput();
                     toast.success("Quiz creation started");
@@ -608,7 +607,7 @@ export function WorkspaceSection({
                   onSelect={() => {
                     toast.success("Deep Research action selected");
                     setSelectedActions(["deep-research"]);
-                    api.composer().setText("I want to do research on ");
+                    aui.composer().setText("I want to do research on ");
                     if (setIsChatExpanded && !isChatExpanded && isDesktop) {
                       setIsChatExpanded(true);
                     }
@@ -623,7 +622,6 @@ export function WorkspaceSection({
           </ContextMenuContent>
         )}
       </ContextMenu>
-
       {/* Selection Action Bar - show when cards are selected */}
       {(state.items ?? []).length > 0 && !isChatMaximized && selectedCardIds.size > 0 && (
         <SelectionActionBar
@@ -635,12 +633,8 @@ export function WorkspaceSection({
           isCompactMode={isItemPanelOpen && isChatExpanded}
         />
       )}
-
       {/* Modal Manager - scoped to workspace panel only */}
       {modalManager}
-
-
-
       {/* Move To Dialog */}
       {showMoveDialog && selectedCardIdsArray.length > 0 && (
         <MoveToDialog
@@ -655,7 +649,6 @@ export function WorkspaceSection({
           onMoveMultiple={handleMoveConfirm}
         />
       )}
-
       {/* Delete Confirmation Dialog */}
       <AlertDialog open={showDeleteDialog} onOpenChange={setShowDeleteDialog}>
         <AlertDialogContent>
@@ -679,22 +672,18 @@ export function WorkspaceSection({
           </AlertDialogFooter>
         </AlertDialogContent>
       </AlertDialog>
-
       {/* Workspace Settings Modal */}
       <WorkspaceSettingsModal
         workspace={currentWorkspace}
         open={showSettingsModal}
         onOpenChange={setShowSettingsModal}
       />
-
       {/* Share Workspace Dialog */}
       <ShareWorkspaceDialog
         workspace={currentWorkspace}
         open={showShareDialog}
         onOpenChange={setShowShareDialog}
       />
-
-
       {/* YouTube Dialog */}
       <CreateYouTubeDialog
         open={showYouTubeDialog}
diff --git a/src/hooks/ai/use-blocknote-selection-context-provider.ts b/src/hooks/ai/use-blocknote-selection-context-provider.ts
index f6965159..c398762f 100644
--- a/src/hooks/ai/use-blocknote-selection-context-provider.ts
+++ b/src/hooks/ai/use-blocknote-selection-context-provider.ts
@@ -1,4 +1,4 @@
-import { useAssistantApi } from "@assistant-ui/react";
+import { useAui } from "@assistant-ui/react";
 import { useEffect, useMemo } from "react";
 import { useUIStore, selectBlockNoteSelection } from "@/lib/stores/ui-store";
 
@@ -7,7 +7,7 @@ import { useUIStore, selectBlockNoteSelection } from "@/lib/stores/ui-store";
  * Automatically updates when BlockNote selection changes and cleans up on unmount
  */
 export function useBlockNoteSelectionContextProvider() {
-  const api = useAssistantApi();
+  const aui = useAui();
   const blockNoteSelection = useUIStore(selectBlockNoteSelection);
 
   // Format context for BlockNote selection - memoized to avoid recalculation
@@ -43,10 +43,10 @@ export function useBlockNoteSelectionContextProvider() {
       return;
     }
 
-    return api.modelContext().register({
+    return aui.modelContext().register({
       getModelContext: () => ({
         system: contextInstructions,
       }),
     });
-  }, [api, contextInstructions, blockNoteSelection]);
+  }, [aui, contextInstructions, blockNoteSelection]);
 }
diff --git a/src/hooks/ai/use-selected-actions-context-provider.ts b/src/hooks/ai/use-selected-actions-context-provider.ts
index bc5a5308..a927d61a 100644
--- a/src/hooks/ai/use-selected-actions-context-provider.ts
+++ b/src/hooks/ai/use-selected-actions-context-provider.ts
@@ -1,4 +1,4 @@
-import { useAssistantApi } from "@assistant-ui/react";
+import { useAui } from "@assistant-ui/react";
 import { useEffect, useMemo } from "react";
 import { formatSelectedActionsContext } from "@/lib/utils/format-workspace-context";
 
@@ -9,7 +9,7 @@ import { formatSelectedActionsContext } from "@/lib/utils/format-workspace-conte
 export function useSelectedActionsContextProvider(
   selectedActions: string[]
 ) {
-  const api = useAssistantApi();
+  const aui = useAui();
 
   // Format context for selected actions - memoized to avoid recalculation
   const contextInstructions = useMemo(
@@ -24,11 +24,11 @@ export function useSelectedActionsContextProvider(
       return;
     }
 
-    return api.modelContext().register({
+    return aui.modelContext().register({
       getModelContext: () => ({
         system: contextInstructions,
       }),
     });
-  }, [api, contextInstructions, selectedActions.length]);
+  }, [aui, contextInstructions, selectedActions.length]);
 }
 
diff --git a/src/hooks/ai/use-workspace-context-provider.ts b/src/hooks/ai/use-workspace-context-provider.ts
index 4a0b93af..0607b146 100644
--- a/src/hooks/ai/use-workspace-context-provider.ts
+++ b/src/hooks/ai/use-workspace-context-provider.ts
@@ -1,4 +1,4 @@
-import { useAssistantApi } from "@assistant-ui/react";
+import { useAui } from "@assistant-ui/react";
 import { useEffect, useMemo } from "react";
 import type { AgentState } from "@/lib/workspace-state/types";
 import { formatWorkspaceContext } from "@/lib/utils/format-workspace-context";
@@ -12,7 +12,7 @@ export function useWorkspaceContextProvider(
   workspaceId: string | null, 
   state: AgentState
 ) {
-  const api = useAssistantApi();
+  const aui = useAui();
 
   // Format workspace context - memoized to avoid recalculation
   const contextInstructions = useMemo(
@@ -26,10 +26,10 @@ export function useWorkspaceContextProvider(
       return;
     }
 
-    return api.modelContext().register({
+    return aui.modelContext().register({
       getModelContext: () => ({
         system: contextInstructions,
       }),
     });
-  }, [api, contextInstructions, workspaceId]);
+  }, [aui, contextInstructions, workspaceId]);
 }
diff --git a/src/lib/ai/tools/deep-research.ts b/src/lib/ai/tools/deep-research.ts
index 30ebf132..e65d577b 100644
--- a/src/lib/ai/tools/deep-research.ts
+++ b/src/lib/ai/tools/deep-research.ts
@@ -1,4 +1,5 @@
 import { GoogleGenAI } from "@google/genai";
+import { tool, zodSchema } from "ai";
 import { z } from "zod";
 import { logger } from "@/lib/utils/logger";
 import { workspaceWorker } from "@/lib/ai/workers";
@@ -8,12 +9,14 @@ import type { WorkspaceToolContext } from "./workspace-tools";
  * Create the deepResearch tool
  */
 export function createDeepResearchTool(ctx: WorkspaceToolContext) {
-    return {
+    return tool({
         description: "Perform deep, multi-step research on a complex topic. Use this when the user explicitly asks for 'deep research' or when a simple web search is insufficient for the depth required. This tool IMMEDIATELY creates a special research card in the workspace that will stream progress and display the final report. You should ask clarifying questions BEFORE calling this tool if the request is vague. Once ready, call this tool with the refined topic/prompt.",
-        inputSchema: z.object({
-            prompt: z.string().describe("The detailed research topic and instructions."),
-        }),
-        execute: async ({ prompt }: { prompt: string }) => {
+        inputSchema: zodSchema(
+            z.object({
+                prompt: z.string().describe("The detailed research topic and instructions."),
+            })
+        ),
+        execute: async ({ prompt }) => {
             logger.debug("🎯 [DEEP-RESEARCH] Starting deep research for:", prompt);
 
             try {
@@ -68,5 +71,5 @@ export function createDeepResearchTool(ctx: WorkspaceToolContext) {
                 };
             }
         },
-    };
+    });
 }
diff --git a/src/lib/ai/tools/process-files.ts b/src/lib/ai/tools/process-files.ts
index 47f9f947..a9a0da44 100644
--- a/src/lib/ai/tools/process-files.ts
+++ b/src/lib/ai/tools/process-files.ts
@@ -1,5 +1,5 @@
 import { google } from "@ai-sdk/google";
-import { generateText } from "ai";
+import { generateText, tool, zodSchema } from "ai";
 import { z } from "zod";
 import { logger } from "@/lib/utils/logger";
 import { readFile } from "fs/promises";
@@ -206,12 +206,14 @@ Provide a clear, comprehensive analysis of the video content.`;
  * Create the processFiles tool
  */
 export function createProcessFilesTool() {
-    return {
+    return tool({
         description: "Process and analyze files including PDFs, images, documents, and videos. Handles local file URLs (/api/files/...), Supabase storage URLs (files uploaded to your workspace), and YouTube videos. Files are downloaded and analyzed directly by Gemini. You can provide custom instructions for what to extract or focus on. Use this for file URLs and video URLs, NOT for regular web pages.",
-        inputSchema: z.object({
-            jsonInput: z.string().describe("JSON string containing an object with 'urls' (array of file/video URLs) and optional 'instruction' (string for custom analysis). Example: '{\"urls\": [\"https://...storage.../file.pdf\"], \"instruction\": \"summarize key points\"}'"),
-        }),
-        execute: async ({ jsonInput }: { jsonInput: string }) => {
+        inputSchema: zodSchema(
+            z.object({
+                jsonInput: z.string().describe("JSON string containing an object with 'urls' (array of file/video URLs) and optional 'instruction' (string for custom analysis). Example: '{\"urls\": [\"https://...storage.../file.pdf\"], \"instruction\": \"summarize key points\"}'"),
+            })
+        ),
+        execute: async ({ jsonInput }) => {
             let parsed;
             try {
                 parsed = JSON.parse(jsonInput);
@@ -294,5 +296,5 @@ export function createProcessFilesTool() {
 
             return fileResults.join('\n\n---\n\n');
         },
-    };
+    });
 }
diff --git a/src/lib/ai/tools/process-urls.ts b/src/lib/ai/tools/process-urls.ts
index 253ec44d..5f6584af 100644
--- a/src/lib/ai/tools/process-urls.ts
+++ b/src/lib/ai/tools/process-urls.ts
@@ -1,5 +1,5 @@
 import { google } from "@ai-sdk/google";
-import { generateText } from "ai";
+import { generateText, tool, zodSchema } from "ai";
 import { z } from "zod";
 import { logger } from "@/lib/utils/logger";
 
@@ -7,12 +7,14 @@ import { logger } from "@/lib/utils/logger";
  * Create the processUrls tool for analyzing web pages
  */
 export function createProcessUrlsTool() {
-    return {
+    return tool({
         description: "Analyze web pages using Google's URL Context API. Extracts content, key information, and metadata from regular web URLs (http/https). Use this for web pages, articles, documentation, and other web content. For files (PDFs, images, documents) or videos, use the processFiles tool instead.",
-        inputSchema: z.object({
-            jsonInput: z.string().describe("JSON string containing an object with 'urls' (array of web URLs). Example: '{\"urls\": [\"https://example.com\"]}'"),
-        }),
-        execute: async ({ jsonInput }: { jsonInput: string }) => {
+        inputSchema: zodSchema(
+            z.object({
+                jsonInput: z.string().describe("JSON string containing an object with 'urls' (array of web URLs). Example: '{\"urls\": [\"https://example.com\"]}'"),
+            })
+        ),
+        execute: async ({ jsonInput }) => {
             let parsed;
             try {
                 parsed = JSON.parse(jsonInput);
@@ -108,5 +110,5 @@ Provide a clear, accurate answer based on the URL content.`;
                 };
             }
         },
-    };
+    });
 }
diff --git a/src/lib/ai/tools/search-code.ts b/src/lib/ai/tools/search-code.ts
index 893505ed..c0102fea 100644
--- a/src/lib/ai/tools/search-code.ts
+++ b/src/lib/ai/tools/search-code.ts
@@ -1,4 +1,5 @@
 import { z } from "zod";
+import { tool, zodSchema } from "ai";
 import { logger } from "@/lib/utils/logger";
 import { searchWorker, codeExecutionWorker } from "@/lib/ai/workers";
 
@@ -6,29 +7,33 @@ import { searchWorker, codeExecutionWorker } from "@/lib/ai/workers";
  * Create the searchWeb tool
  */
 export function createSearchWebTool() {
-    return {
+    return tool({
         description: "Search the web for current information, facts, news, or research. Use this when you need up-to-date information from the internet.",
-        inputSchema: z.object({
-            query: z.string().describe("The search query"),
-        }),
-        execute: async ({ query }: { query: string }) => {
+        inputSchema: zodSchema(
+            z.object({
+                query: z.string().describe("The search query"),
+            })
+        ),
+        execute: async ({ query }) => {
             return await searchWorker(query);
         },
-    };
+    });
 }
 
 /**
  * Create the executeCode tool
  */
 export function createExecuteCodeTool() {
-    return {
+    return tool({
         description: "Execute Python code for calculations, data processing, algorithms, or mathematical computations.",
-        inputSchema: z.object({
-            task: z.string().describe("Description of the task to solve with code"),
-        }),
-        execute: async ({ task }: { task: string }) => {
+        inputSchema: zodSchema(
+            z.object({
+                task: z.string().describe("Description of the task to solve with code"),
+            })
+        ),
+        execute: async ({ task }) => {
             logger.debug("🎯 [ORCHESTRATOR] Delegating to Code Execution Worker:", task);
             return await codeExecutionWorker(task);
         },
-    };
+    });
 }
diff --git a/src/lib/ai/tools/workspace-tools.ts b/src/lib/ai/tools/workspace-tools.ts
index d65c4b22..235151b7 100644
--- a/src/lib/ai/tools/workspace-tools.ts
+++ b/src/lib/ai/tools/workspace-tools.ts
@@ -1,3 +1,4 @@
+import { tool, zodSchema } from "ai";
 import { z } from "zod";
 import { logger } from "@/lib/utils/logger";
 import { workspaceWorker } from "@/lib/ai/workers";
@@ -15,12 +16,15 @@ export interface WorkspaceToolContext {
  * Create the createNote tool
  */
 export function createNoteTool(ctx: WorkspaceToolContext) {
-    return {
+    return tool({
         description: "Create a note card. returns success message.\n\nCRITICAL CONSTRAINTS:\n1. 'content' MUST NOT start with the title.\n2. Start directly with body text.\n3. NO Mermaid diagrams.\n\nCRITICAL - MATHEMATICAL EXPRESSIONS: Use LaTeX with DOUBLE DOLLAR SIGNS ($$) for ALL math:\n- Use $$...$$ for ALL math expressions (both inline and block)\n- Single $ is for CURRENCY only (e.g., $19.99). NEVER use single $ for math\n- For inline math: $$E = mc^2$$ (same line as text)\n- For block math (separate lines): Put $$ delimiters on separate lines\n- Always ensure math blocks are properly closed with matching $$",
-        inputSchema: z.any().describe(
-            "JSON {title, content}. 'content': markdown body. DO NOT repeat title in content. Start with subheadings/text. No Mermaid. Use $$...$$ for ALL math (both inline and block). Single $ is for currency only."
+        inputSchema: zodSchema(
+            z.object({
+                title: z.string().describe("The title of the note card"),
+                content: z.string().describe("The markdown body content. DO NOT repeat title in content. Start with subheadings/text. No Mermaid. Use $$...$$ for ALL math (both inline and block). Single $ is for currency only."),
+            })
         ),
-        execute: async ({ title, content }: { title: string; content: string }) => {
+        execute: async ({ title, content }) => {
             // Validate inputs before use
             if (!title || typeof title !== 'string') {
                 return {
@@ -51,19 +55,25 @@ export function createNoteTool(ctx: WorkspaceToolContext) {
                 folderId: ctx.activeFolderId,
             });
         },
-    };
+    });
 }
 
 /**
  * Create the updateCard tool
  */
 export function createUpdateCardTool(ctx: WorkspaceToolContext) {
-    return {
+    return tool({
         description: "Update the content of an existing card. This tool COMPLETELY REPLACES the existing content. You must synthesize the FULL new content by combining the existing card content (from your context) with the user's requested changes. Do not just provide the diff; provide the complete new markdown content.\n\nCRITICAL - MATHEMATICAL EXPRESSIONS: Use LaTeX with DOUBLE DOLLAR SIGNS ($$) for ALL math:\n- Use $$...$$ for ALL math expressions (both inline and block)\n- Single $ is for CURRENCY only (e.g., $19.99). NEVER use single $ for math\n- For inline math: $$E = mc^2$$ (same line as text)\n- For block math (separate lines): Put $$ delimiters on separate lines\n- Always ensure math blocks are properly closed with matching $$\n- Add spaces around $$ symbols when math appears in lists or tables\n- Do not place punctuation immediately after math expressions.",
-        inputSchema: z.any().describe(
-            "A JSON object with 'id' (string) and 'markdown' (string) or 'content' (string) fields. The 'id' uniquely identifies the note to update. The 'markdown' or 'content' field contains the full note body ONLY (do not include the title as a header). Use $$...$$ for ALL math (both inline and block). Single $ is for currency only. Ensure math inside lists and tables has spaces around the $$ symbols. Do not place punctuation immediately after math expressions."
+        inputSchema: zodSchema(
+            z.object({
+                id: z.string().describe("The ID of the card to update"),
+                markdown: z.string().optional().describe("The full note body ONLY (do not include the title as a header). Use $$...$$ for ALL math (both inline and block). Single $ is for currency only."),
+                content: z.string().optional().describe("Alternative to 'markdown' - the full note body ONLY (do not include the title as a header). Use $$...$$ for ALL math (both inline and block). Single $ is for currency only."),
+            }).refine((data) => data.markdown !== undefined || data.content !== undefined, {
+                message: "Either 'markdown' or 'content' must be provided",
+            })
         ),
-        execute: async (input: any) => {
+        execute: async (input) => {
             logger.group("🎯 [UPDATE-CARD] Tool execution started", true);
             logger.debug("Raw input received:", {
                 inputType: typeof input,
@@ -131,19 +141,21 @@ export function createUpdateCardTool(ctx: WorkspaceToolContext) {
                 };
             }
         },
-    };
+    });
 }
 
 /**
  * Create the clearCardContent tool
  */
 export function createClearCardContentTool(ctx: WorkspaceToolContext) {
-    return {
+    return tool({
         description: "Clear/delete the content of a card while preserving its title. Use this when the user wants to delete the contents of a card.",
-        inputSchema: z.object({
-            id: z.string().describe("The ID of the card to clear"),
-        }),
-        execute: async ({ id }: { id: string }) => {
+        inputSchema: zodSchema(
+            z.object({
+                id: z.string().describe("The ID of the card to clear"),
+            })
+        ),
+        execute: async ({ id }) => {
             logger.debug("🎯 [ORCHESTRATOR] Delegating to Workspace Worker (clear):", { id });
 
             if (!ctx.workspaceId) {
@@ -159,19 +171,21 @@ export function createClearCardContentTool(ctx: WorkspaceToolContext) {
                 content: "",
             });
         },
-    };
+    });
 }
 
 /**
  * Create the deleteCard tool
  */
 export function createDeleteCardTool(ctx: WorkspaceToolContext) {
-    return {
+    return tool({
         description: "Permanently delete a card/note from the workspace. Use this when the user explicitly asks to delete or remove a card.",
-        inputSchema: z.object({
-            id: z.string().describe("The ID of the card to delete"),
-        }),
-        execute: async ({ id }: { id: string }) => {
+        inputSchema: zodSchema(
+            z.object({
+                id: z.string().describe("The ID of the card to delete"),
+            })
+        ),
+        execute: async ({ id }) => {
             logger.debug("🎯 [ORCHESTRATOR] Delegating to Workspace Worker (delete):", { id });
 
             if (!ctx.workspaceId) {
@@ -186,21 +200,22 @@ export function createDeleteCardTool(ctx: WorkspaceToolContext) {
                 itemId: id,
             });
         },
-    };
+    });
 }
 
 /**
  * Create the selectCards tool
  */
 export function createSelectCardsTool(ctx: WorkspaceToolContext) {
-    return {
+    return tool({
         description:
             "Select one or more cards by their TITLES and add them to the conversation context. This tool helps you surface specific cards when the user refers to them. The tool will perform fuzzy matching to find the best matching cards and return their full content immediately.",
-        inputSchema: z.object({
-            cardTitles: z.array(z.string()).describe("Array of card titles to search for and select"),
-        }),
-        execute: async (input: { cardTitles: string[] }) => {
-            const { cardTitles } = input;
+        inputSchema: zodSchema(
+            z.object({
+                cardTitles: z.array(z.string()).describe("Array of card titles to search for and select"),
+            })
+        ),
+        execute: async ({ cardTitles }) => {
 
             if (!cardTitles || cardTitles.length === 0) {
                 return {
@@ -286,5 +301,5 @@ export function createSelectCardsTool(ctx: WorkspaceToolContext) {
                 };
             }
         },
-    };
+    });
 }