From 75cafd5fbc3b1e272bc7c571923693cdadf5c68a Mon Sep 17 00:00:00 2001
From: urjitc <135136842+urjitc@users.noreply.github.com>
Date: Sun, 19 Apr 2026 21:37:07 +0000
Subject: [PATCH 1/7] Split thread.tsx into components, rename Composer to
 PromptInput, and add chat runtime ACL

---
 src/components/assistant-ui/attachment.tsx    |    6 +-
 src/components/assistant-ui/thread.tsx        | 1511 +----------------
 .../thread/AssistantActionBar.tsx             |   45 +
 .../assistant-ui/thread/AssistantMessage.tsx  |   50 +
 .../assistant-ui/thread/BranchPicker.tsx      |   40 +
 .../assistant-ui/thread/EditPromptInput.tsx   |   81 +
 .../assistant-ui/thread/MessageError.tsx      |   14 +
 .../assistant-ui/thread/PromptInput.tsx       |  159 ++
 .../assistant-ui/thread/PromptInputShell.tsx  |  167 ++
 .../thread/PromptInputToolbar.tsx             |  187 ++
 src/components/assistant-ui/thread/Thread.tsx |   46 +
 .../thread/ThreadLoadingSkeleton.tsx          |   34 +
 .../assistant-ui/thread/ThreadSuggestions.tsx |   92 +
 .../assistant-ui/thread/ThreadWelcome.tsx     |   27 +
 .../assistant-ui/thread/UserActionBar.tsx     |   46 +
 .../assistant-ui/thread/UserMessage.tsx       |   92 +
 .../thread/UserMessageTruncateContext.tsx     |   30 +
 .../thread/VirtualizedMessages.tsx            |   55 +
 .../thread/hooks/use-mention-menu.ts          |  153 ++
 .../thread/hooks/use-prompt-input-paste.ts    |   53 +
 src/components/assistant-ui/thread/index.ts   |    1 +
 .../assistant-ui/thread/message-components.ts |    9 +
 .../thread/prompt-input-floating-actions.ts   |   53 +
 .../assistant-ui/thread/suggestion-actions.ts |   47 +
 src/lib/chat/runtime/hooks.ts                 |   70 +
 src/lib/chat/runtime/index.ts                 |   11 +
 src/lib/chat/runtime/primitives.ts            |   20 +
 src/lib/chat/runtime/types.ts                 |   45 +
 .../process-pdf-attachments-in-background.ts  |   67 +
 29 files changed, 1698 insertions(+), 1513 deletions(-)
 create mode 100644 src/components/assistant-ui/thread/AssistantActionBar.tsx
 create mode 100644 src/components/assistant-ui/thread/AssistantMessage.tsx
 create mode 100644 src/components/assistant-ui/thread/BranchPicker.tsx
 create mode 100644 src/components/assistant-ui/thread/EditPromptInput.tsx
 create mode 100644 src/components/assistant-ui/thread/MessageError.tsx
 create mode 100644 src/components/assistant-ui/thread/PromptInput.tsx
 create mode 100644 src/components/assistant-ui/thread/PromptInputShell.tsx
 create mode 100644 src/components/assistant-ui/thread/PromptInputToolbar.tsx
 create mode 100644 src/components/assistant-ui/thread/Thread.tsx
 create mode 100644 src/components/assistant-ui/thread/ThreadLoadingSkeleton.tsx
 create mode 100644 src/components/assistant-ui/thread/ThreadSuggestions.tsx
 create mode 100644 src/components/assistant-ui/thread/ThreadWelcome.tsx
 create mode 100644 src/components/assistant-ui/thread/UserActionBar.tsx
 create mode 100644 src/components/assistant-ui/thread/UserMessage.tsx
 create mode 100644 src/components/assistant-ui/thread/UserMessageTruncateContext.tsx
 create mode 100644 src/components/assistant-ui/thread/VirtualizedMessages.tsx
 create mode 100644 src/components/assistant-ui/thread/hooks/use-mention-menu.ts
 create mode 100644 src/components/assistant-ui/thread/hooks/use-prompt-input-paste.ts
 create mode 100644 src/components/assistant-ui/thread/index.ts
 create mode 100644 src/components/assistant-ui/thread/message-components.ts
 create mode 100644 src/components/assistant-ui/thread/prompt-input-floating-actions.ts
 create mode 100644 src/components/assistant-ui/thread/suggestion-actions.ts
 create mode 100644 src/lib/chat/runtime/hooks.ts
 create mode 100644 src/lib/chat/runtime/index.ts
 create mode 100644 src/lib/chat/runtime/primitives.ts
 create mode 100644 src/lib/chat/runtime/types.ts
 create mode 100644 src/lib/uploads/process-pdf-attachments-in-background.ts

diff --git a/src/components/assistant-ui/attachment.tsx b/src/components/assistant-ui/attachment.tsx
index 241b5eea..d1031709 100644
--- a/src/components/assistant-ui/attachment.tsx
+++ b/src/components/assistant-ui/attachment.tsx
@@ -385,7 +385,7 @@ export const UserMessageAttachments: FC = () => {
   );
 };
 
-export const ComposerAttachments: FC = () => {
+export const PromptInputAttachments: FC = () => {
   return (
     <div className="aui-composer-attachments mb-2 flex w-full flex-row items-center gap-2 overflow-x-auto pt-0.5 pb-1 empty:hidden">
       <ComposerPrimitive.Attachments
@@ -395,7 +395,7 @@ export const ComposerAttachments: FC = () => {
   );
 };
 
-export const ComposerAddAttachment: FC = () => {
+export const PromptInputAddAttachment: FC = () => {
   const fileInputRef = useRef<HTMLInputElement>(null);
   const containerRef = useRef<HTMLDivElement>(null);
   const aui = useAui();
@@ -480,7 +480,7 @@ export const ComposerAddAttachment: FC = () => {
     }
   };
 
-  const uploadInputId = "composer-file-upload";
+  const uploadInputId = "prompt-input-file-upload";
 
   return (
     <>
diff --git a/src/components/assistant-ui/thread.tsx b/src/components/assistant-ui/thread.tsx
index e2efadc8..afb9c1f3 100644
--- a/src/components/assistant-ui/thread.tsx
+++ b/src/components/assistant-ui/thread.tsx
@@ -1,1510 +1 @@
-import {
-  ArrowUpIcon,
-  CheckIcon,
-  ChevronDown,
-  ChevronUp,
-  Loader2,
-  CheckCircle2,
-  ChevronLeftIcon,
-  ChevronRightIcon,
-  CopyIcon,
-  FileText,
-  Upload,
-  PencilIcon,
-  PlusSquareIcon,
-  RefreshCwIcon,
-  Square,
-  GalleryHorizontalEnd,
-  AlertTriangle,
-  Sparkles,
-  Bug,
-  Brain,
-  Play,
-  Search,
-} from "lucide-react";
-import { FaWandMagicSparkles } from "react-icons/fa6";
-import { LuBook } from "react-icons/lu";
-import { PiCardsThreeBold } from "react-icons/pi";
-import { cn } from "@/lib/utils";
-import {
-  ActionBarPrimitive,
-  AuiIf,
-  BranchPickerPrimitive,
-  ComposerPrimitive,
-  ErrorPrimitive,
-  MessagePrimitive,
-  ThreadPrimitive,
-  useAui,
-  useMessage,
-  useMessagePartText,
-  useAuiState,
-} from "@assistant-ui/react";
-import { useVirtualizer } from "@tanstack/react-virtual";
-
-import type { FC, RefObject } from "react";
-import { createContext, useContext } from "react";
-import { useEffect, useRef, useState, useMemo, useCallback } from "react";
-
-import { Button } from "@/components/ui/button";
-import { Skeleton } from "@/components/ui/skeleton";
-import {
-  DropdownMenu,
-  DropdownMenuContent,
-  DropdownMenuItem,
-  DropdownMenuTrigger,
-} from "@/components/ui/dropdown-menu";
-import {
-  Popover,
-  PopoverContent,
-  PopoverTrigger,
-} from "@/components/ui/popover";
-import Link from "next/link";
-import { MarkdownText } from "@/components/assistant-ui/markdown-text";
-import { ToolFallback } from "@/components/assistant-ui/tool-fallback";
-import { AIFeedbackDialog } from "@/components/assistant-ui/AIFeedbackDialog";
-
-import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
-import { useAttachmentUploadStore } from "@/lib/stores/attachment-upload-store";
-import {
-  ComposerAttachments,
-  ComposerAddAttachment,
-  UserMessageAttachments,
-} from "@/components/assistant-ui/attachment";
-import { AssistantLoader } from "@/components/assistant-ui/assistant-loader";
-import { File as FileComponent } from "@/components/assistant-ui/file";
-import { isOfficeDocument } from "@/lib/uploads/office-document-validation";
-import { Sources } from "@/components/assistant-ui/sources";
-import { Image } from "@/components/assistant-ui/image";
-import { Reasoning, ReasoningGroup } from "@/components/assistant-ui/reasoning";
-import { ToolGroup } from "@/components/assistant-ui/tool-group";
-
-import type { Item } from "@/lib/workspace-state/types";
-import { CardContextDisplay } from "@/components/chat/CardContextDisplay";
-import { ReplyContextDisplay } from "@/components/chat/ReplyContextDisplay";
-import { MessageContextBadges } from "@/components/chat/MessageContextBadges";
-import { MentionMenu } from "@/components/chat/MentionMenu";
-import { useWorkspaceStore } from "@/lib/stores/workspace-store";
-import { useUIStore, selectReplySelections } from "@/lib/stores/ui-store";
-import { useSelectedCardIds } from "@/hooks/ui/use-selected-card-ids";
-import { useShallow } from "zustand/react/shallow";
-import { useWorkspaceState } from "@/hooks/workspace/use-workspace-state";
-import { useWorkspaceOperations } from "@/hooks/workspace/use-workspace-operations";
-import { useFeatureFlagEnabled } from "posthog-js/react";
-import { toast } from "sonner";
-import { filterItems } from "@/lib/workspace-state/search";
-import { useSession } from "@/lib/auth-client";
-import { focusComposerInput } from "@/lib/utils/composer-utils";
-import { buildWorkspaceItemDefinitionsFromAssets } from "@/lib/uploads/uploaded-asset";
-import { uploadSelectedFiles } from "@/lib/uploads/upload-selection";
-import { startAssetProcessing } from "@/lib/uploads/start-asset-processing";
-import {
-  getDocumentUploadFailureMessage,
-  getDocumentUploadPartialMessage,
-  getDocumentUploadSuccessMessage,
-} from "@/lib/uploads/upload-feedback";
-import { SpeechToTextButton } from "@/components/assistant-ui/SpeechToTextButton";
-import {
-  PromptBuilderDialog,
-  type PromptBuilderAction,
-} from "@/components/assistant-ui/PromptBuilderDialog";
-import { ModelPicker } from "@/components/assistant-ui/ModelPicker";
-import { ModelSettingsMenu } from "@/components/assistant-ui/ModelSettingsMenu";
-import { ThinkExLogo } from "@/components/ui/thinkex-logo";
-
-interface ThreadProps {
-  items?: Item[];
-}
-
-export const Thread: FC<ThreadProps> = ({ items = [] }) => {
-  const viewportRef = useRef<HTMLDivElement>(null);
-
-  return (
-    <ThreadPrimitive.Root
-      className="aui-root aui-thread-root @container flex h-full flex-col bg-sidebar"
-      style={{
-        ["--thread-max-width" as string]: "50rem",
-      }}
-    >
-      <ThreadPrimitive.Viewport
-        ref={viewportRef}
-        turnAnchor="top"
-        autoScroll={false}
-        className="aui-thread-viewport relative flex min-h-0 flex-1 flex-col overflow-x-auto overflow-y-scroll px-4"
-      >
-        <AuiIf condition={({ thread }) => thread.isLoading}>
-          <ThreadLoadingSkeleton />
-        </AuiIf>
-        <AuiIf condition={({ thread }) => thread.isEmpty && !thread.isLoading}>
-          <ThreadWelcome items={items} />
-        </AuiIf>
-
-        <VirtualizedMessages scrollRef={viewportRef} />
-      </ThreadPrimitive.Viewport>
-
-      <div className="aui-thread-composer-wrapper mx-auto flex w-full max-w-[var(--thread-max-width)] flex-shrink-0 flex-col gap-4 overflow-visible rounded-t-3xl bg-sidebar px-4 pb-3 md:pb-4">
-        <ComposerHoverWrapper items={items} />
-      </div>
-    </ThreadPrimitive.Root>
-  );
-};
-
-interface VirtualizedMessagesProps {
-  scrollRef: RefObject<HTMLDivElement | null>;
-}
-
-const VirtualizedMessages: FC<VirtualizedMessagesProps> = ({ scrollRef }) => {
-  const messageCount = useAuiState((s) => s.thread.messages.length);
-
-  const virtualizer = useVirtualizer({
-    count: messageCount,
-    getScrollElement: () => scrollRef.current,
-    estimateSize: () => 350,
-    overscan: 5,
-  });
-
-  if (messageCount === 0) return null;
-
-  return (
-    <div
-      style={{
-        height: `${virtualizer.getTotalSize()}px`,
-        width: "100%",
-        position: "relative",
-      }}
-    >
-      {virtualizer.getVirtualItems().map((virtualRow) => (
-        <div
-          key={virtualRow.key}
-          data-index={virtualRow.index}
-          ref={virtualizer.measureElement}
-          style={{
-            position: "absolute",
-            top: 0,
-            left: 0,
-            width: "100%",
-            transform: `translateY(${virtualRow.start}px)`,
-          }}
-        >
-          <ThreadPrimitive.MessageByIndex
-            index={virtualRow.index}
-            components={MESSAGE_COMPONENTS}
-          />
-        </div>
-      ))}
-    </div>
-  );
-};
-
-const ThreadLoadingSkeleton: FC = () => {
-  return (
-    <div
-      role="status"
-      aria-label="Loading chat"
-      className="aui-thread-loading-skeleton mx-auto flex w-full max-w-[var(--thread-max-width)] flex-col gap-6 px-2 py-8"
-    >
-      {/* User message skeleton (right-aligned) */}
-      <div className="flex justify-end">
-        <div className="max-w-[85%] space-y-2">
-          <Skeleton className="h-12 w-48 rounded-lg" />
-        </div>
-      </div>
-      {/* Assistant message skeleton (left-aligned) */}
-      <div className="space-y-2">
-        <Skeleton className="h-4 w-full max-w-[90%] rounded" />
-        <Skeleton className="h-4 w-full max-w-[70%] rounded" />
-        <Skeleton className="h-4 w-32 rounded" />
-      </div>
-      {/* User message skeleton */}
-      <div className="flex justify-end">
-        <Skeleton className="h-10 w-64 rounded-lg" />
-      </div>
-      {/* Assistant message skeleton - taller */}
-      <div className="space-y-2">
-        <Skeleton className="h-4 w-full max-w-[95%] rounded" />
-        <Skeleton className="h-4 w-full max-w-[80%] rounded" />
-        <Skeleton className="h-4 w-full max-w-[60%] rounded" />
-        <Skeleton className="h-4 w-24 rounded" />
-      </div>
-    </div>
-  );
-};
-
-interface ThreadWelcomeProps {
-  items: Item[];
-}
-
-const ThreadWelcome: FC<ThreadWelcomeProps> = ({ items }) => {
-  return (
-    <div className="aui-thread-welcome-root mx-auto my-auto flex w-full max-w-[var(--thread-max-width)] flex-grow flex-col">
-      <div className="aui-thread-welcome-center flex w-full flex-grow flex-col items-center justify-center">
-        <div className="aui-thread-welcome-message flex size-full flex-col items-center justify-center px-8">
-          <div className="aui-thread-welcome-message-motion-0 mb-1 flex justify-center">
-            <div className="relative flex h-12 w-12 items-center justify-center">
-              <ThinkExLogo size={48} priority />
-            </div>
-          </div>
-        </div>
-      </div>
-      <ThreadSuggestions items={items} />
-    </div>
-  );
-};
-
-interface ThreadSuggestionsProps {
-  items: Item[];
-}
-
-const SUGGESTION_ACTIONS = [
-  {
-    title: "Search",
-    icon: Search,
-    iconClassName: "size-4 shrink-0 text-sky-500",
-    action: "search" as PromptBuilderAction,
-    useDialog: true,
-  },
-  {
-    title: "Flashcards",
-    icon: PiCardsThreeBold,
-    iconClassName: "size-4 shrink-0 text-purple-400 rotate-180",
-    action: "flashcards" as PromptBuilderAction,
-    useDialog: true,
-  },
-  {
-    title: "YouTube",
-    icon: Play,
-    iconClassName: "size-4 shrink-0 text-red-500",
-    action: "youtube" as PromptBuilderAction,
-    useDialog: true,
-  },
-  {
-    title: "Upload",
-    icon: Upload,
-    iconClassName: "size-4 shrink-0 text-red-400",
-    triggerFileInput: true,
-  },
-  {
-    title: "Quiz",
-    icon: Brain,
-    iconClassName: "size-4 shrink-0 text-green-400",
-    action: "quiz" as PromptBuilderAction,
-    useDialog: true,
-  },
-  {
-    title: "Document",
-    icon: FileText,
-    iconClassName: "size-4 shrink-0 text-sky-400",
-    action: "document" as PromptBuilderAction,
-    useDialog: true,
-  },
-];
-
-const ThreadSuggestions: FC<ThreadSuggestionsProps> = ({ items }) => {
-  const aui = useAui();
-  const [dialogAction, setDialogAction] = useState<PromptBuilderAction | null>(
-    null,
-  );
-
-  const handleDirectFill = useCallback(
-    (action: string) => {
-      aui?.composer()?.setText(action);
-      focusComposerInput(true);
-    },
-    [aui],
-  );
-
-  const handleTriggerFileInput = useCallback(() => {
-    document.getElementById("composer-file-upload")?.click();
-    focusComposerInput(true);
-  }, []);
-
-  return (
-    <>
-      <div className="aui-thread-welcome-suggestions grid w-full grid-cols-2 gap-2 pb-4 sm:grid-cols-3">
-        {SUGGESTION_ACTIONS.map((suggestedAction, index) => {
-          const Icon = suggestedAction.icon;
-          return (
-            <div
-              key={`suggested-action-${suggestedAction.title}-${index}`}
-              className="aui-thread-welcome-suggestion-display"
-            >
-              <Button
-                type="button"
-                variant="ghost"
-                className="aui-thread-welcome-suggestion h-auto w-full flex-1 flex-wrap items-center justify-start gap-2 rounded-lg border border-sidebar-border px-5 py-4 text-left text-sm dark:hover:bg-accent/60"
-                aria-label={suggestedAction.title}
-                onClick={() => {
-                  if (
-                    "triggerFileInput" in suggestedAction &&
-                    suggestedAction.triggerFileInput
-                  ) {
-                    handleTriggerFileInput();
-                  } else if (
-                    suggestedAction.useDialog &&
-                    "action" in suggestedAction &&
-                    suggestedAction.action
-                  ) {
-                    setDialogAction(suggestedAction.action);
-                  } else if (
-                    "composerFill" in suggestedAction &&
-                    typeof suggestedAction.composerFill === "string"
-                  ) {
-                    handleDirectFill(suggestedAction.composerFill);
-                  }
-                }}
-              >
-                <Icon className={suggestedAction.iconClassName} />
-                <span className="aui-thread-welcome-suggestion-text-1 font-medium">
-                  {suggestedAction.title}
-                </span>
-              </Button>
-            </div>
-          );
-        })}
-      </div>
-
-      {dialogAction && (
-        <PromptBuilderDialog
-          open={!!dialogAction}
-          onOpenChange={(open) => !open && setDialogAction(null)}
-          action={dialogAction}
-          items={items}
-        />
-      )}
-    </>
-  );
-};
-
-// Floating action buttons shown above composer on hover
-const COMPOSER_FLOATING_ACTIONS = [
-  {
-    id: "document",
-    label: "Document",
-    icon: FileText,
-    iconClassName: "size-3.5 shrink-0 text-sky-400",
-    action: "document" as PromptBuilderAction,
-    useDialog: true,
-  },
-  {
-    id: "learn",
-    label: "Learn",
-    icon: LuBook,
-    iconClassName: "size-3.5 shrink-0 text-amber-500",
-    subActions: [
-      {
-        id: "flashcards",
-        label: "Flashcards",
-        icon: PiCardsThreeBold,
-        iconClassName: "size-4 text-purple-400 rotate-180",
-        action: "flashcards" as PromptBuilderAction,
-      },
-      {
-        id: "quiz",
-        label: "Quiz",
-        icon: Brain,
-        iconClassName: "size-4 text-green-400",
-        action: "quiz" as PromptBuilderAction,
-      },
-    ],
-  },
-  {
-    id: "youtube",
-    label: "YouTube",
-    icon: Play,
-    iconClassName: "size-3.5 text-red-500",
-    action: "youtube" as PromptBuilderAction,
-    useDialog: true,
-  },
-  {
-    id: "search",
-    label: "Search",
-    icon: Search,
-    iconClassName: "size-3.5 text-teal-500",
-    action: "search" as PromptBuilderAction,
-    useDialog: true,
-  },
-];
-
-interface ComposerHoverWrapperProps {
-  items: Item[];
-}
-
-const FLOATING_MENU_HIDE_DELAY_MS = 400;
-
-const ComposerHoverWrapper: FC<ComposerHoverWrapperProps> = ({ items }) => {
-  const [isHovered, setIsHovered] = useState(false);
-  const hideTimeoutRef = useRef<NodeJS.Timeout | null>(null);
-  const aui = useAui();
-  const [dialogAction, setDialogAction] = useState<PromptBuilderAction | null>(
-    null,
-  );
-  const isThreadEmpty = useAuiState(({ thread }) => thread?.isEmpty ?? true);
-  const hasComposerText = useAuiState(
-    (s) =>
-      Boolean(
-        ((s as { composer?: { text?: string } })?.composer?.text ?? "").trim(),
-      ),
-  );
-
-  const handleDirectFill = useCallback(
-    (fill: string) => {
-      aui?.composer()?.setText(fill);
-      focusComposerInput(true);
-    },
-    [aui],
-  );
-
-  const handleMouseEnter = useCallback(() => {
-    if (hideTimeoutRef.current) {
-      clearTimeout(hideTimeoutRef.current);
-      hideTimeoutRef.current = null;
-    }
-    setIsHovered(true);
-  }, []);
-
-  const handleMouseLeave = useCallback(() => {
-    hideTimeoutRef.current = setTimeout(() => {
-      setIsHovered(false);
-    }, FLOATING_MENU_HIDE_DELAY_MS);
-  }, []);
-
-  useEffect(() => {
-    return () => {
-      if (hideTimeoutRef.current) clearTimeout(hideTimeoutRef.current);
-    };
-  }, []);
-
-  return (
-    <div className="relative">
-      {/* Composer + floating menu - main hover zone */}
-      <div
-        className="relative"
-        onMouseEnter={handleMouseEnter}
-        onMouseLeave={handleMouseLeave}
-      >
-        {/* Floating buttons - appear above composer on hover */}
-        <div
-          className={cn(
-            "absolute bottom-full left-0 right-0 z-20 flex justify-center gap-0.5 pb-2",
-            "transition-opacity duration-150 ease-out",
-            !isThreadEmpty && isHovered && !hasComposerText
-              ? "opacity-100 pointer-events-auto"
-              : "opacity-0 pointer-events-none",
-          )}
-        >
-          <div className="flex flex-wrap items-center justify-center gap-0.5 rounded-xl border border-sidebar-border bg-sidebar-accent px-1.5 py-1 shadow-md dark:border-sidebar-border/15">
-            {COMPOSER_FLOATING_ACTIONS.map((action) => {
-              if ("subActions" in action) {
-                // Learn button with dropdown
-                const Icon = action.icon;
-                return (
-                  <DropdownMenu key={action.id}>
-                    <DropdownMenuTrigger asChild>
-                      <button
-                        type="button"
-                        className={cn(
-                          "inline-flex cursor-pointer items-center gap-1.5 rounded-lg px-2 py-1 text-xs font-normal",
-                          "text-sidebar-foreground transition-colors",
-                          "hover:bg-sidebar-foreground/10 dark:hover:bg-sidebar-foreground/15",
-                        )}
-                        aria-label={action.label}
-                      >
-                        <Icon className={action.iconClassName} />
-                        <span>{action.label}</span>
-                      </button>
-                    </DropdownMenuTrigger>
-                    <DropdownMenuContent
-                      align="center"
-                      side="top"
-                      className="min-w-[140px]"
-                      sideOffset={4}
-                    >
-                      {(action.subActions ?? []).map((sub) => {
-                        const SubIcon = sub.icon;
-                        return (
-                          <DropdownMenuItem
-                            key={sub.id}
-                            onSelect={() => setDialogAction(sub.action)}
-                            className="flex cursor-pointer items-center gap-2"
-                          >
-                            <SubIcon className={sub.iconClassName} />
-                            {sub.label}
-                          </DropdownMenuItem>
-                        );
-                      })}
-                    </DropdownMenuContent>
-                  </DropdownMenu>
-                );
-              }
-              const Icon = action.icon;
-              return (
-                <button
-                  key={action.id}
-                  type="button"
-                  onClick={() => {
-                    if (
-                      action.useDialog &&
-                      "action" in action &&
-                      action.action
-                    ) {
-                      setDialogAction(action.action);
-                    } else if (
-                      "composerFill" in action &&
-                      typeof action.composerFill === "string"
-                    ) {
-                      handleDirectFill(action.composerFill);
-                    }
-                  }}
-                  className={cn(
-                    "inline-flex cursor-pointer items-center gap-1.5 rounded-lg px-2 py-1 text-xs font-normal",
-                    "text-sidebar-foreground transition-colors",
-                    "hover:bg-sidebar-foreground/10 dark:hover:bg-sidebar-foreground/15",
-                  )}
-                  aria-label={action.label}
-                >
-                  <Icon className={action.iconClassName} />
-                  <span>{action.label}</span>
-                </button>
-              );
-            })}
-          </div>
-        </div>
-
-        <Composer items={items} />
-
-        {dialogAction && (
-          <PromptBuilderDialog
-            open={!!dialogAction}
-            onOpenChange={(open) => !open && setDialogAction(null)}
-            action={dialogAction}
-            items={items}
-          />
-        )}
-      </div>
-    </div>
-  );
-};
-
-interface ComposerProps {
-  items: Item[];
-}
-
-const Composer: FC<ComposerProps> = ({ items }) => {
-  const currentWorkspaceId = useWorkspaceStore(
-    (state) => state.currentWorkspaceId,
-  );
-  const aui = useAui();
-  const replySelections = useUIStore(useShallow(selectReplySelections));
-  const clearReplySelections = useUIStore(
-    (state) => state.clearReplySelections,
-  );
-  const { selectedCardIds } = useSelectedCardIds();
-
-  // Get workspace state and operations for PDF card creation
-  const { state: workspaceState } = useWorkspaceState(currentWorkspaceId);
-  const operations = useWorkspaceOperations(currentWorkspaceId, workspaceState);
-
-  // Watch for thread changes to auto-focus composer (built-in assistant-ui behavior)
-  const mainThreadId = useAuiState(
-    ({ threads }) => (threads as any)?.mainThreadId,
-  );
-  const inputRef = useRef<HTMLTextAreaElement | null>(null);
-
-  // Auto-focus composer when thread changes
-  useEffect(() => {
-    if (mainThreadId && inputRef.current) {
-      // Small delay to ensure DOM is ready after thread switch
-      const timeoutId = setTimeout(() => {
-        inputRef.current?.focus();
-      }, 100);
-      return () => clearTimeout(timeoutId);
-    }
-  }, [mainThreadId]);
-
-  // Mention menu state
-  const [mentionMenuOpen, setMentionMenuOpen] = useState(false);
-  const [mentionQuery, setMentionQuery] = useState("");
-  const [mentionStartIndex, setMentionStartIndex] = useState<number | null>(
-    null,
-  );
-  const toggleCardSelection = useUIStore((state) => state.toggleCardSelection);
-
-  // Handle input changes for @ mention detection
-  const handleInput = useCallback(
-    (e: React.FormEvent<HTMLTextAreaElement>) => {
-      const textarea = e.currentTarget;
-      const value = textarea.value;
-      const cursorPos = textarea.selectionStart ?? 0;
-
-      // Check if we're currently tracking a mention
-      if (mentionStartIndex !== null) {
-        // Extract the query from @ to cursor
-        const query = value.slice(mentionStartIndex + 1, cursorPos);
-
-        // Check if we've moved before the @ or if there's a space/newline after @
-        if (
-          cursorPos <= mentionStartIndex ||
-          query.includes(" ") ||
-          query.includes("\n")
-        ) {
-          setMentionMenuOpen(false);
-          setMentionStartIndex(null);
-          setMentionQuery("");
-        } else {
-          setMentionQuery(query);
-        }
-      }
-    },
-    [mentionStartIndex],
-  );
-
-  // Handle keydown for @ detection and menu control
-  const handleKeyDown = useCallback(
-    (e: React.KeyboardEvent<HTMLTextAreaElement>) => {
-      const textarea = e.currentTarget;
-
-      // Detect @ key
-      if (e.key === "@" && !mentionMenuOpen) {
-        const cursorPos = textarea.selectionStart ?? 0;
-        // Check if @ is at start or preceded by whitespace
-        const charBefore = cursorPos > 0 ? textarea.value[cursorPos - 1] : " ";
-        if (charBefore === " " || charBefore === "\n" || cursorPos === 0) {
-          setMentionMenuOpen(true);
-          setMentionStartIndex(cursorPos);
-          setMentionQuery("");
-        }
-      }
-
-      // Close menu on Escape
-      if (e.key === "Escape" && mentionMenuOpen) {
-        e.preventDefault();
-        setMentionMenuOpen(false);
-        setMentionStartIndex(null);
-        setMentionQuery("");
-      }
-
-      // Prevent default behavior when menu is open for arrow keys and Enter
-      if (
-        mentionMenuOpen &&
-        ["ArrowUp", "ArrowDown", "Enter"].includes(e.key)
-      ) {
-        e.preventDefault();
-      }
-    },
-    [mentionMenuOpen],
-  );
-
-  // Clear the @query from input (extracted for reuse)
-  const clearMentionQuery = useCallback(() => {
-    if (mentionStartIndex !== null && inputRef.current) {
-      const textarea = inputRef.current;
-      const currentValue = textarea.value;
-
-      // Calculate what to remove: from the @ symbol to current cursor/end of query
-      const atSymbolIndex = mentionStartIndex;
-
-      // Find where the query ends (current text after @ until space/newline or end)
-      let queryEndIndex = mentionStartIndex;
-      while (
-        queryEndIndex < currentValue.length &&
-        currentValue[queryEndIndex] !== " " &&
-        currentValue[queryEndIndex] !== "\n"
-      ) {
-        queryEndIndex++;
-      }
-
-      const textBefore = currentValue.substring(0, atSymbolIndex);
-      const textAfter = currentValue.substring(queryEndIndex);
-
-      // Set the new value without the @query
-      const newValue = textBefore + textAfter;
-
-      // Update the textarea value
-      aui?.composer()?.setText(newValue);
-
-      // Reset mention state
-      setMentionQuery("");
-      setMentionStartIndex(null);
-
-      // Focus and position cursor at where the @ was
-      setTimeout(() => {
-        if (inputRef.current) {
-          inputRef.current.focus();
-          const newCursorPos = textBefore.length;
-          inputRef.current.setSelectionRange(newCursorPos, newCursorPos);
-        }
-      }, 0);
-    }
-  }, [mentionStartIndex, aui]);
-
-  // Handle mention selection - toggle item, keep menu open for multi-select
-  const handleMentionSelect = useCallback(
-    (item: Item) => {
-      toggleCardSelection(item.id);
-      // Don't close menu or clear query - allow selecting multiple items
-    },
-    [toggleCardSelection],
-  );
-
-  // Handle mention menu close - remove the @query from input
-  const handleMentionMenuClose = useCallback(
-    (open: boolean) => {
-      if (!open) {
-        clearMentionQuery();
-      }
-      setMentionMenuOpen(open);
-    },
-    [clearMentionQuery],
-  );
-
-  const handlePaste = async (e: React.ClipboardEvent<HTMLTextAreaElement>) => {
-    const clipboardData = e.clipboardData;
-    if (!clipboardData || !currentWorkspaceId) return;
-
-    // Check if clipboard contains files (images or other file types)
-    const files = Array.from(clipboardData.files) as File[];
-
-    if (files.length > 0) {
-      e.preventDefault();
-      // Add the first file (or prioritize images if multiple files) as an attachment
-      // It will be uploaded when the message is sent
-      const imageFile = files.find((file: File) =>
-        file.type.startsWith("image/"),
-      );
-      const fileToUpload = imageFile || files[0];
-
-      if (fileToUpload) {
-        try {
-          await aui?.composer()?.addAttachment(fileToUpload);
-        } catch (error) {
-          console.error("Failed to add file attachment:", error);
-        }
-      }
-      return;
-    }
-
-    // Also check clipboard items for image data (e.g., screenshots)
-    const clipboardItems = Array.from(
-      clipboardData.items,
-    ) as DataTransferItem[];
-    const imageItem = clipboardItems.find((item: DataTransferItem) =>
-      item.type.startsWith("image/"),
-    );
-
-    if (imageItem) {
-      e.preventDefault();
-      const file = imageItem.getAsFile();
-      if (file) {
-        try {
-          // Add as attachment - will be uploaded when message is sent
-          await aui?.composer()?.addAttachment(file);
-        } catch (error) {
-          console.error("Failed to add image attachment:", error);
-        }
-      }
-      return;
-    }
-
-    // Check if clipboard contains a URL (logic removed)
-  };
-
-  /**
-   * Process PDF attachments in the background without blocking message sending
-   */
-  const processPdfAttachmentsInBackground = async (
-    pdfAttachments: any[],
-    workspaceId: string,
-    operations: any,
-  ) => {
-    let files: File[] = [];
-    try {
-      files = pdfAttachments
-        .map((attachment) => attachment.file)
-        .filter((file): file is File => !!file);
-      const { uploads, failedFiles } = await uploadSelectedFiles(files);
-
-      if (uploads.length > 0) {
-        const pdfCardDefinitions =
-          buildWorkspaceItemDefinitionsFromAssets(uploads);
-        const createdIds = operations.createItems(pdfCardDefinitions, {
-          showSuccessToast: false,
-        });
-
-        void startAssetProcessing({
-          workspaceId,
-          assets: uploads,
-          itemIds: createdIds,
-          onOcrError: (error) => {
-            console.error("Error starting assistant file processing:", error);
-          },
-        });
-
-        // Show success toast
-        if (failedFiles.length === 0) {
-          toast.success(getDocumentUploadSuccessMessage(uploads.length));
-        } else {
-          toast.warning(
-            getDocumentUploadPartialMessage(uploads.length, failedFiles.length),
-          );
-        }
-      } else {
-        toast.error(
-          getDocumentUploadFailureMessage(failedFiles.length || files.length),
-        );
-      }
-    } catch (error) {
-      console.error("Error creating PDF cards in background:", error);
-      toast.error(
-        getDocumentUploadFailureMessage(files.length || pdfAttachments.length),
-      );
-    }
-  };
-
-  return (
-    <ComposerPrimitive.Root
-      className="aui-composer-root relative flex w-full flex-col rounded-lg border border-sidebar-border bg-sidebar-accent px-3.5 pt-2 pb-1 shadow-[0_9px_9px_0px_rgba(0,0,0,0.01),0_2px_5px_0px_rgba(0,0,0,0.06)] dark:border-sidebar-border/15"
-      onClick={(e) => {
-        // Focus the input when clicking anywhere in the composer area
-        // This allows users to easily return focus after interacting with quizzes or other cards
-        if (inputRef.current && !e.defaultPrevented) {
-          inputRef.current.focus();
-        }
-      }}
-      onSubmit={async (e) => {
-        e.preventDefault();
-
-        // Wait for attachment uploads before sending (same UX as home input)
-        if (useAttachmentUploadStore.getState().uploadingIds.size > 0) {
-          toast.info("Please wait for uploads to finish before sending");
-          return;
-        }
-
-        // Get the current composer state
-        const composerState = aui?.composer()?.getState();
-        if (!composerState) return;
-
-        const currentText = composerState.text;
-        const attachments = composerState.attachments || [];
-
-        // Prevent empty messages when reply context already provides content.
-        const hasReplyContext = replySelections.length > 0;
-        if (
-          !currentText.trim() &&
-          attachments.length === 0 &&
-          !hasReplyContext
-        ) {
-          return;
-        }
-
-        // Detect PDF attachments for background processing
-        const pdfAttachments = attachments.filter((att) => {
-          const file = att.file;
-          return (
-            file &&
-            (file.type === "application/pdf" ||
-              file.name.toLowerCase().endsWith(".pdf") ||
-              isOfficeDocument(file))
-          );
-        });
-
-        // Process PDFs in background - don't block message sending
-        if (pdfAttachments.length > 0 && currentWorkspaceId) {
-          processPdfAttachmentsInBackground(
-            pdfAttachments,
-            currentWorkspaceId,
-            operations,
-          );
-        }
-
-        // Get selected cards for context
-        const selectedItems = items.filter((item) =>
-          selectedCardIds.has(item.id),
-        );
-
-        // Combine all context: selected cards, reply texts, and user message
-        // Use placeholder when empty so AI SDK accepts (backend injects reply context into message)
-        let modifiedText =
-          currentText.trim() || (hasReplyContext ? "Empty message" : "");
-
-        // Attach per-request context as metadata via runConfig
-        // This flows through as body.metadata.custom on the server
-        // IMPORTANT: Always set runConfig (even empty) to clear stale data from previous sends,
-        // because the composer does NOT reset runConfig after send().
-        const customMetadata: Record<string, unknown> = {};
-        if (replySelections.length > 0) {
-          customMetadata.replySelections = replySelections;
-        }
-        aui
-          ?.composer()
-          ?.setRunConfig(
-            Object.keys(customMetadata).length > 0
-              ? { custom: customMetadata }
-              : {},
-          );
-
-        // Set the modified text and send
-        aui?.composer()?.setText(modifiedText);
-        aui?.composer()?.send();
-
-        // Clear all per-request state immediately — captured in runConfig before send()
-        clearReplySelections();
-      }}
-    >
-      {/* Attachment Display - shows uploaded files */}
-      <ComposerAttachments />
-      {/* Card Context Display - shows selected cards inside input area */}
-      <CardContextDisplay items={items} />
-      {/* Ask AI context chips (text selections from chat / PDF / document) */}
-      <ReplyContextDisplay />
-      <div className="relative">
-        <ComposerPrimitive.Input
-          ref={inputRef}
-          placeholder="Ask anything or @mention items"
-          className="aui-composer-input max-h-32 w-full resize-none bg-transparent py-1.5 text-base text-sidebar-foreground outline-none placeholder:text-sidebar-foreground/60 focus:outline-none"
-          rows={1}
-          autoFocus
-          aria-label="Message input"
-          maxLength={10000}
-          cancelOnEscape={false}
-          onPaste={handlePaste}
-          onKeyDown={handleKeyDown}
-          onInput={handleInput}
-        />
-        {/* Mention Menu */}
-        <MentionMenu
-          open={mentionMenuOpen}
-          onOpenChange={handleMentionMenuClose}
-          query={mentionQuery}
-          items={items}
-          onSelect={handleMentionSelect}
-          selectedCardIds={selectedCardIds}
-          selectedIndicator={(isSelected) =>
-            isSelected ? (
-              <CheckCircle2 className="size-4 text-primary flex-shrink-0" />
-            ) : undefined
-          }
-        />
-      </div>
-      <ComposerAction items={items} />
-    </ComposerPrimitive.Root>
-  );
-};
-
-interface ComposerActionProps {
-  items: Item[];
-}
-
-const ComposerAction: FC<ComposerActionProps> = ({ items }) => {
-  const { data: session } = useSession();
-  useAui();
-  const hasUploading = useAttachmentUploadStore((s) => s.uploadingIds.size > 0);
-  const isAnonymous = session?.user?.isAnonymous ?? false;
-  const { selectedCardIds } = useSelectedCardIds();
-  const toggleCardSelection = useUIStore((state) => state.toggleCardSelection);
-
-  const [isWarningPopoverOpen, setIsWarningPopoverOpen] = useState(false);
-  const hoverTimeoutRef = useRef<NodeJS.Timeout | null>(null);
-  const [isFeedbackDialogOpen, setIsFeedbackDialogOpen] = useState(false);
-
-  // AI Debug button: always show in dev, only when feature flag enabled in prod
-  const isDev = process.env.NODE_ENV === "development";
-  const aiDebugFlagEnabled = useFeatureFlagEnabled("ai-debug-feedback");
-  const showAiDebugButton = isDev || aiDebugFlagEnabled === true;
-
-  // Cleanup timeout on unmount
-  useEffect(() => {
-    return () => {
-      if (hoverTimeoutRef.current) {
-        clearTimeout(hoverTimeoutRef.current);
-      }
-    };
-  }, []);
-
-  // Filter items for the dropdown based on selected tags
-  const filteredItems = useMemo(() => {
-    return filterItems(items, "");
-  }, [items]);
-
-  return (
-    <div className="aui-composer-action-wrapper relative mb-2 flex items-center justify-between">
-      {/* Attachment buttons on the left */}
-      <div className="relative z-0 flex items-center gap-0">
-        <div className="relative z-0">
-          <ComposerAddAttachment />
-        </div>
-        {!isAnonymous && <ModelSettingsMenu />}
-        <div className={cn(!isAnonymous && "ml-0.5")}>
-          <ModelPicker />
-        </div>
-        {/* AI Debug Button - feature flag in prod, localhost debugger in dev */}
-        {showAiDebugButton && (
-          <button
-            type="button"
-            className="ml-2 flex items-center gap-1.5 px-1.5 py-1 rounded-md bg-red-500/10 hover:bg-red-500/20 transition-colors flex-shrink-0 text-xs font-normal text-red-500 hover:text-red-400 cursor-pointer border border-red-500/20"
-            onClick={() => {
-              if (process.env.NODE_ENV === "development") {
-                window.open("http://localhost:4983", "_blank");
-              } else {
-                setIsFeedbackDialogOpen(true);
-              }
-            }}
-          >
-            <Bug className="w-3.5 h-3.5" />
-            <span>AI Debug</span>
-          </button>
-        )}
-        <AIFeedbackDialog
-          open={isFeedbackDialogOpen}
-          onOpenChange={setIsFeedbackDialogOpen}
-        />
-        {/* Warning icon for anonymous users */}
-        {isAnonymous && (
-          <Popover
-            open={isWarningPopoverOpen}
-            onOpenChange={(open) => {
-              setIsWarningPopoverOpen(open);
-              if (!open) {
-                focusComposerInput();
-              }
-            }}
-          >
-            <PopoverTrigger asChild>
-              <button
-                type="button"
-                onMouseEnter={() => {
-                  if (hoverTimeoutRef.current) {
-                    clearTimeout(hoverTimeoutRef.current);
-                  }
-                  setIsWarningPopoverOpen(true);
-                }}
-                onMouseLeave={() => {
-                  hoverTimeoutRef.current = setTimeout(() => {
-                    setIsWarningPopoverOpen(false);
-                  }, 100);
-                }}
-                className="flex items-center justify-center w-5 h-5 rounded-md hover:bg-accent transition-colors flex-shrink-0 cursor-pointer"
-                aria-label="Warning: AI chats won't save unless logged in"
-              >
-                <AlertTriangle className="w-4 h-4 text-yellow-500 animate-[pulse-scale_3.5s_ease-in-out_infinite]" />
-              </button>
-            </PopoverTrigger>
-            <PopoverContent
-              side="top"
-              align="start"
-              sideOffset={8}
-              onMouseEnter={() => {
-                if (hoverTimeoutRef.current) {
-                  clearTimeout(hoverTimeoutRef.current);
-                }
-                setIsWarningPopoverOpen(true);
-              }}
-              onMouseLeave={() => {
-                hoverTimeoutRef.current = setTimeout(() => {
-                  setIsWarningPopoverOpen(false);
-                }, 100);
-              }}
-              className="w-64 p-3"
-            >
-              <div className="space-y-3">
-                <p className="text-sm text-foreground">
-                  Your AI chats won't save unless you are logged in.
-                </p>
-                <div className="flex items-center gap-2">
-                  <Link href="/auth/sign-in" className="flex-1">
-                    <Button
-                      variant="outline"
-                      size="sm"
-                      className="w-full h-7 text-xs"
-                      onClick={() => setIsWarningPopoverOpen(false)}
-                    >
-                      Sign in
-                    </Button>
-                  </Link>
-                  <Link href="/auth/sign-up" className="flex-1">
-                    <Button
-                      size="sm"
-                      className="w-full h-7 text-xs"
-                      onClick={() => setIsWarningPopoverOpen(false)}
-                    >
-                      Sign up
-                    </Button>
-                  </Link>
-                </div>
-              </div>
-            </PopoverContent>
-          </Popover>
-        )}
-      </div>
-      {/* Right side: speech/send/cancel button */}
-      <div className="flex items-center gap-2">
-        {!isAnonymous && <SpeechToTextButton />}
-        <AuiIf condition={({ thread }) => !thread.isRunning}>
-          <TooltipIconButton
-            tooltip={hasUploading ? "Uploading attachments..." : "Send message"}
-            side="bottom"
-            type="submit"
-            variant="default"
-            size="icon"
-            className="aui-composer-send size-[34px] rounded-full p-1"
-            aria-label="Send message"
-            disabled={hasUploading}
-          >
-            {hasUploading ? (
-              <Loader2 className="aui-composer-send-icon size-4 text-background animate-spin" />
-            ) : (
-              <ArrowUpIcon className="aui-composer-send-icon size-4 text-background" />
-            )}
-          </TooltipIconButton>
-        </AuiIf>
-
-        <AuiIf condition={({ thread }) => thread.isRunning}>
-          <ComposerPrimitive.Cancel asChild>
-            <Button
-              type="button"
-              variant="default"
-              size="icon"
-              className="aui-composer-cancel size-[34px] rounded-full border border-muted-foreground/60 hover:bg-primary/75 dark:border-muted-foreground/90"
-              aria-label="Stop generating"
-            >
-              <Square className="aui-composer-cancel-icon size-3 text-background fill-current" />
-            </Button>
-          </ComposerPrimitive.Cancel>
-        </AuiIf>
-      </div>
-    </div>
-  );
-};
-
-const MessageError: FC = () => {
-  return (
-    <MessagePrimitive.Error>
-      <ErrorPrimitive.Root className="aui-message-error-root mt-2 rounded-md border border-destructive bg-destructive/10 p-3 text-sm text-destructive dark:bg-destructive/5 dark:text-red-200">
-        <ErrorPrimitive.Message className="aui-message-error-message line-clamp-2" />
-      </ErrorPrimitive.Root>
-    </MessagePrimitive.Error>
-  );
-};
-
-const AssistantMessage: FC = () => {
-  return (
-    <MessagePrimitive.Root asChild>
-      <div
-        className="aui-assistant-message-root relative mx-auto w-full max-w-[var(--thread-max-width)] animate-in pb-4 duration-150 ease-out fade-in slide-in-from-bottom-1 last:mb-4"
-        data-role="assistant"
-      >
-        <div className="aui-assistant-message-content mx-2 leading-7 break-words text-foreground">
-          <AssistantLoader />
-          <MessagePrimitive.Parts
-            components={{
-              Text: MarkdownText,
-              File: FileComponent,
-              Source: Sources,
-              Image,
-              Reasoning: Reasoning,
-              ReasoningGroup: ReasoningGroup,
-              ToolGroup,
-              tools: {
-                Fallback: ToolFallback,
-              },
-            }}
-          />
-          <MessageError />
-        </div>
-
-        <div className="aui-assistant-message-footer mt-2 ml-2 flex">
-          <BranchPicker />
-          <AssistantActionBar />
-        </div>
-      </div>
-    </MessagePrimitive.Root>
-  );
-};
-
-const AssistantActionBar: FC = () => {
-  const { content } = useMessage();
-
-  const textContent = useMemo(() => {
-    const textParts = content.filter(
-      (part): part is { type: "text"; text: string } => part.type === "text",
-    );
-    return textParts.map((part) => part.text ?? "").join("\n\n");
-  }, [content]);
-
-  const [copied, setCopied] = useState(false);
-  const copyTimeoutRef = useRef<NodeJS.Timeout | null>(null);
-
-  const handleCopy = useCallback(() => {
-    if (!textContent) return;
-    navigator.clipboard.writeText(textContent);
-    setCopied(true);
-    if (copyTimeoutRef.current) clearTimeout(copyTimeoutRef.current);
-    copyTimeoutRef.current = setTimeout(() => setCopied(false), 2000);
-  }, [textContent]);
-
-  return (
-    <ActionBarPrimitive.Root
-      hideWhenRunning
-      autohide="never"
-      className="aui-assistant-action-bar-root col-start-3 row-start-2 -ml-1 flex gap-0.5 text-muted-foreground data-floating:absolute data-floating:rounded-md data-floating:border data-floating:bg-background data-floating:p-1 data-floating:shadow-sm"
-    >
-      <TooltipIconButton tooltip="Copy" onClick={handleCopy}>
-        {copied ? <CheckIcon /> : <CopyIcon />}
-      </TooltipIconButton>
-      <ActionBarPrimitive.Reload asChild>
-        <TooltipIconButton tooltip="Refresh">
-          <RefreshCwIcon />
-        </TooltipIconButton>
-      </ActionBarPrimitive.Reload>
-    </ActionBarPrimitive.Root>
-  );
-};
-
-const USER_MESSAGE_MAX_CHARS = 250;
-
-const UserMessageTruncateContext = createContext<{
-  maxChars: number;
-  expanded: boolean;
-  showExpand: boolean;
-} | null>(null);
-
-// Custom Text component for UserMessage (plain text + truncation)
-const UserMessageText: FC = () => {
-  const { text: rawText } = useMessagePartText();
-  const truncateCtx = useContext(UserMessageTruncateContext);
-
-  let text = rawText;
-
-  if (
-    truncateCtx &&
-    !truncateCtx.expanded &&
-    truncateCtx.maxChars < Infinity &&
-    text.length > truncateCtx.maxChars
-  ) {
-    text = text.slice(0, truncateCtx.maxChars).trim() + "...";
-  }
-
-  return <div className="whitespace-pre-wrap">{text}</div>;
-};
-
-const UserMessage: FC = () => {
-  const [expanded, setExpanded] = useState(false);
-  const message = useMessage();
-
-  const textLength = useMemo(
-    () =>
-      message.content
-        .filter(
-          (part): part is { type: "text"; text: string } =>
-            part.type === "text",
-        )
-        .reduce((sum, part) => sum + (part.text?.length ?? 0), 0),
-    [message.content],
-  );
-
-  const showExpand = textLength > USER_MESSAGE_MAX_CHARS;
-
-  const truncateCtxValue = useMemo(
-    () => ({
-      maxChars: USER_MESSAGE_MAX_CHARS,
-      expanded,
-      showExpand,
-    }),
-    [expanded, showExpand],
-  );
-
-  return (
-    <MessagePrimitive.Root asChild>
-      <div
-        className="aui-user-message-root mx-auto grid w-full max-w-[var(--thread-max-width)] animate-breathe-in auto-rows-auto grid-cols-[minmax(72px,1fr)_auto] gap-y-2 px-2 pt-4 pb-1 last:mb-5 [&:where(>*)]:col-start-2"
-        data-role="user"
-      >
-        {/* Attachments display */}
-        <UserMessageAttachments />
-
-        <div className="aui-user-message-content-wrapper relative col-start-2 min-w-0">
-          <MessageContextBadges />
-          <UserMessageTruncateContext.Provider value={truncateCtxValue}>
-            <div className="aui-user-message-content relative rounded-lg bg-muted px-3 py-2 break-words text-foreground text-sm">
-              <MessagePrimitive.Parts
-                components={{
-                  Text: UserMessageText,
-                  File: FileComponent,
-                }}
-              />
-              {showExpand && (
-                <div className="flex justify-end pt-1.5 mt-1.5">
-                  <button
-                    type="button"
-                    onClick={() => setExpanded((e) => !e)}
-                    className="flex items-center gap-1 text-xs text-muted-foreground hover:text-foreground transition-colors"
-                    aria-label={expanded ? "Show less" : "Show more"}
-                  >
-                    {expanded ? (
-                      <ChevronUp className="size-3.5" />
-                    ) : (
-                      <ChevronDown className="size-3.5" />
-                    )}
-                    {expanded ? "Show less" : "Show more"}
-                  </button>
-                </div>
-              )}
-            </div>
-          </UserMessageTruncateContext.Provider>
-        </div>
-
-        <div className="aui-user-message-footer ml-2 flex justify-end col-start-2 relative min-h-[20px]">
-          <div className="absolute right-0">
-            <UserActionBar />
-          </div>
-        </div>
-
-        <BranchPicker className="aui-user-branch-picker col-span-full col-start-1 row-start-3 -mr-1 justify-end" />
-      </div>
-    </MessagePrimitive.Root>
-  );
-};
-
-const UserActionBar: FC = () => {
-  const message = useMessage();
-  const [copied, setCopied] = useState(false);
-  const copyTimeoutRef = useRef<NodeJS.Timeout | null>(null);
-
-  const textContent = useMemo(() => {
-    return message.content
-      .filter(
-        (part): part is { type: "text"; text: string } => part.type === "text",
-      )
-      .map((part) => part.text ?? "")
-      .join("\n\n");
-  }, [message.content]);
-
-  const handleCopy = useCallback(() => {
-    if (!textContent) return;
-    navigator.clipboard.writeText(textContent);
-    setCopied(true);
-    if (copyTimeoutRef.current) clearTimeout(copyTimeoutRef.current);
-    copyTimeoutRef.current = setTimeout(() => setCopied(false), 2000);
-  }, [textContent]);
-
-  return (
-    <ActionBarPrimitive.Root
-      hideWhenRunning
-      autohide="not-last"
-      className="aui-user-action-bar-root flex gap-1 text-muted-foreground data-floating:absolute data-floating:rounded-md data-floating:border data-floating:bg-background data-floating:p-1 data-floating:shadow-sm"
-    >
-      <TooltipIconButton tooltip="Copy" onClick={handleCopy}>
-        {copied ? <CheckIcon /> : <CopyIcon />}
-      </TooltipIconButton>
-      <ActionBarPrimitive.Edit asChild>
-        <TooltipIconButton tooltip="Edit">
-          <PencilIcon />
-        </TooltipIconButton>
-      </ActionBarPrimitive.Edit>
-    </ActionBarPrimitive.Root>
-  );
-};
-
-const EditComposer: FC = () => {
-  const aui = useAui();
-  const hasUploading = useAttachmentUploadStore((s) => s.uploadingIds.size > 0);
-  const [originalText, setOriginalText] = useState<string>(
-    () => aui?.composer()?.getState()?.text ?? "",
-  );
-  const [currentText, setCurrentText] = useState<string>(
-    () => aui?.composer()?.getState()?.text ?? "",
-  );
-
-  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">
-      <ComposerPrimitive.Root
-        className="aui-edit-composer-root ml-auto flex w-full max-w-7/8 flex-col rounded-xl bg-sidebar-accent border border-sidebar-border"
-        onSubmit={(e) => {
-          e.preventDefault();
-
-          if (useAttachmentUploadStore.getState().uploadingIds.size > 0) {
-            toast.info("Please wait for uploads to finish before sending");
-            return;
-          }
-
-          aui?.composer()?.send();
-        }}
-      >
-        <ComposerAttachments />
-
-        <ComposerPrimitive.Input
-          className="aui-edit-composer-input flex min-h-[60px] w-full resize-none bg-transparent p-4 text-sidebar-foreground outline-none"
-          autoFocus
-          maxLength={10000}
-          onChange={(e) => setCurrentText(e.target.value)}
-        />
-
-        <div className="aui-edit-composer-footer mx-3 mb-3 flex items-center justify-between">
-          <div className="flex items-center gap-2">
-            <ComposerAddAttachment />
-          </div>
-          <div className="flex items-center gap-2">
-            <ComposerPrimitive.Cancel asChild>
-              <Button variant="ghost" size="sm" aria-label="Cancel edit">
-                Cancel
-              </Button>
-            </ComposerPrimitive.Cancel>
-            <Button
-              type="submit"
-              size="sm"
-              aria-label="Update message"
-              disabled={currentText === originalText || hasUploading}
-              className={cn(
-                (currentText === originalText || hasUploading) &&
-                  "opacity-50 cursor-not-allowed",
-              )}
-            >
-              {hasUploading ? (
-                <>
-                  <Loader2 className="size-3.5 animate-spin mr-1.5" />
-                  Uploading...
-                </>
-              ) : (
-                "Update"
-              )}
-            </Button>
-          </div>
-        </div>
-      </ComposerPrimitive.Root>
-    </div>
-  );
-};
-
-const MESSAGE_COMPONENTS = {
-  UserMessage,
-  EditComposer,
-  AssistantMessage,
-};
-
-const BranchPicker: FC<BranchPickerPrimitive.Root.Props> = ({
-  className,
-  ...rest
-}) => {
-  return (
-    <BranchPickerPrimitive.Root
-      hideWhenSingleBranch
-      className={cn(
-        "aui-branch-picker-root mr-2 -ml-2 inline-flex items-center text-xs text-muted-foreground",
-        className,
-      )}
-      {...rest}
-    >
-      <BranchPickerPrimitive.Previous asChild>
-        <TooltipIconButton tooltip="Previous">
-          <ChevronLeftIcon />
-        </TooltipIconButton>
-      </BranchPickerPrimitive.Previous>
-      <span className="aui-branch-picker-state font-medium">
-        <BranchPickerPrimitive.Number /> / <BranchPickerPrimitive.Count />
-      </span>
-      <BranchPickerPrimitive.Next asChild>
-        <TooltipIconButton tooltip="Next">
-          <ChevronRightIcon />
-        </TooltipIconButton>
-      </BranchPickerPrimitive.Next>
-    </BranchPickerPrimitive.Root>
-  );
-};
+export { Thread } from "./thread/index";
diff --git a/src/components/assistant-ui/thread/AssistantActionBar.tsx b/src/components/assistant-ui/thread/AssistantActionBar.tsx
new file mode 100644
index 00000000..84a23a62
--- /dev/null
+++ b/src/components/assistant-ui/thread/AssistantActionBar.tsx
@@ -0,0 +1,45 @@
+"use client";
+
+import { CheckIcon, CopyIcon, RefreshCwIcon } from "lucide-react";
+import { useCallback, useMemo, useRef, useState, type FC } from "react";
+import { ChatActionBar, useChatMessage } from "@/lib/chat/runtime";
+import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
+
+export const AssistantActionBar: FC = () => {
+  const { content } = useChatMessage();
+
+  const textContent = useMemo(() => {
+    const textParts = content.filter(
+      (part): part is { type: "text"; text: string } => part.type === "text",
+    );
+    return textParts.map((part) => part.text ?? "").join("\n\n");
+  }, [content]);
+
+  const [copied, setCopied] = useState(false);
+  const copyTimeoutRef = useRef<NodeJS.Timeout | null>(null);
+
+  const handleCopy = useCallback(() => {
+    if (!textContent) return;
+    navigator.clipboard.writeText(textContent);
+    setCopied(true);
+    if (copyTimeoutRef.current) clearTimeout(copyTimeoutRef.current);
+    copyTimeoutRef.current = setTimeout(() => setCopied(false), 2000);
+  }, [textContent]);
+
+  return (
+    <ChatActionBar.Root
+      hideWhenRunning
+      autohide="never"
+      className="aui-assistant-action-bar-root col-start-3 row-start-2 -ml-1 flex gap-0.5 text-muted-foreground data-floating:absolute data-floating:rounded-md data-floating:border data-floating:bg-background data-floating:p-1 data-floating:shadow-sm"
+    >
+      <TooltipIconButton tooltip="Copy" onClick={handleCopy}>
+        {copied ? <CheckIcon /> : <CopyIcon />}
+      </TooltipIconButton>
+      <ChatActionBar.Reload asChild>
+        <TooltipIconButton tooltip="Refresh">
+          <RefreshCwIcon />
+        </TooltipIconButton>
+      </ChatActionBar.Reload>
+    </ChatActionBar.Root>
+  );
+};
diff --git a/src/components/assistant-ui/thread/AssistantMessage.tsx b/src/components/assistant-ui/thread/AssistantMessage.tsx
new file mode 100644
index 00000000..5d726c65
--- /dev/null
+++ b/src/components/assistant-ui/thread/AssistantMessage.tsx
@@ -0,0 +1,50 @@
+"use client";
+
+import type { FC } from "react";
+import { AssistantLoader } from "@/components/assistant-ui/assistant-loader";
+import { File as FileComponent } from "@/components/assistant-ui/file";
+import { Image } from "@/components/assistant-ui/image";
+import { MarkdownText } from "@/components/assistant-ui/markdown-text";
+import { Reasoning, ReasoningGroup } from "@/components/assistant-ui/reasoning";
+import { Sources } from "@/components/assistant-ui/sources";
+import { ToolFallback } from "@/components/assistant-ui/tool-fallback";
+import { ToolGroup } from "@/components/assistant-ui/tool-group";
+import { ChatMessage } from "@/lib/chat/runtime";
+import { AssistantActionBar } from "./AssistantActionBar";
+import { BranchPicker } from "./BranchPicker";
+import { MessageError } from "./MessageError";
+
+export const AssistantMessage: FC = () => {
+  return (
+    <ChatMessage.Root asChild>
+      <div
+        className="aui-assistant-message-root relative mx-auto w-full max-w-[var(--thread-max-width)] animate-in pb-4 duration-150 ease-out fade-in slide-in-from-bottom-1 last:mb-4"
+        data-role="assistant"
+      >
+        <div className="aui-assistant-message-content mx-2 leading-7 break-words text-foreground">
+          <AssistantLoader />
+          <ChatMessage.Parts
+            components={{
+              Text: MarkdownText,
+              File: FileComponent,
+              Source: Sources,
+              Image,
+              Reasoning,
+              ReasoningGroup,
+              ToolGroup,
+              tools: {
+                Fallback: ToolFallback,
+              },
+            }}
+          />
+          <MessageError />
+        </div>
+
+        <div className="aui-assistant-message-footer mt-2 ml-2 flex">
+          <BranchPicker />
+          <AssistantActionBar />
+        </div>
+      </div>
+    </ChatMessage.Root>
+  );
+};
diff --git a/src/components/assistant-ui/thread/BranchPicker.tsx b/src/components/assistant-ui/thread/BranchPicker.tsx
new file mode 100644
index 00000000..0dea7fd8
--- /dev/null
+++ b/src/components/assistant-ui/thread/BranchPicker.tsx
@@ -0,0 +1,40 @@
+"use client";
+
+import { ChevronLeftIcon, ChevronRightIcon } from "lucide-react";
+import type { FC } from "react";
+import {
+  ChatBranchPicker,
+  type ChatBranchPickerRootProps,
+} from "@/lib/chat/runtime";
+import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
+import { cn } from "@/lib/utils";
+
+export const BranchPicker: FC<ChatBranchPickerRootProps> = ({
+  className,
+  ...rest
+}) => {
+  return (
+    <ChatBranchPicker.Root
+      hideWhenSingleBranch
+      className={cn(
+        "aui-branch-picker-root mr-2 -ml-2 inline-flex items-center text-xs text-muted-foreground",
+        className,
+      )}
+      {...rest}
+    >
+      <ChatBranchPicker.Previous asChild>
+        <TooltipIconButton tooltip="Previous">
+          <ChevronLeftIcon />
+        </TooltipIconButton>
+      </ChatBranchPicker.Previous>
+      <span className="aui-branch-picker-state font-medium">
+        <ChatBranchPicker.Number /> / <ChatBranchPicker.Count />
+      </span>
+      <ChatBranchPicker.Next asChild>
+        <TooltipIconButton tooltip="Next">
+          <ChevronRightIcon />
+        </TooltipIconButton>
+      </ChatBranchPicker.Next>
+    </ChatBranchPicker.Root>
+  );
+};
diff --git a/src/components/assistant-ui/thread/EditPromptInput.tsx b/src/components/assistant-ui/thread/EditPromptInput.tsx
new file mode 100644
index 00000000..ce66d50b
--- /dev/null
+++ b/src/components/assistant-ui/thread/EditPromptInput.tsx
@@ -0,0 +1,81 @@
+"use client";
+
+import { Loader2 } from "lucide-react";
+import { useState, type FC } from "react";
+import {
+  PromptInputAddAttachment,
+  PromptInputAttachments,
+} from "@/components/assistant-ui/attachment";
+import { Button } from "@/components/ui/button";
+import { ChatPromptInput, usePromptInput } from "@/lib/chat/runtime";
+import { useAttachmentUploadStore } from "@/lib/stores/attachment-upload-store";
+import { cn } from "@/lib/utils";
+import { toast } from "sonner";
+
+export const EditPromptInput: FC = () => {
+  const promptInput = usePromptInput();
+  const hasUploading = useAttachmentUploadStore((s) => s.uploadingIds.size > 0);
+  const [originalText] = useState<string>(() => promptInput?.getState()?.text ?? "");
+  const [currentText, setCurrentText] = useState<string>(
+    () => promptInput?.getState()?.text ?? "",
+  );
+
+  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">
+      <ChatPromptInput.Root
+        className="aui-edit-composer-root ml-auto flex w-full max-w-7/8 flex-col rounded-xl bg-sidebar-accent border border-sidebar-border"
+        onSubmit={(e) => {
+          e.preventDefault();
+
+          if (useAttachmentUploadStore.getState().uploadingIds.size > 0) {
+            toast.info("Please wait for uploads to finish before sending");
+            return;
+          }
+
+          promptInput?.send();
+        }}
+      >
+        <PromptInputAttachments />
+
+        <ChatPromptInput.Input
+          className="aui-edit-composer-input flex min-h-[60px] w-full resize-none bg-transparent p-4 text-sidebar-foreground outline-none"
+          autoFocus
+          maxLength={10000}
+          onChange={(e) => setCurrentText(e.target.value)}
+        />
+
+        <div className="aui-edit-composer-footer mx-3 mb-3 flex items-center justify-between">
+          <div className="flex items-center gap-2">
+            <PromptInputAddAttachment />
+          </div>
+          <div className="flex items-center gap-2">
+            <ChatPromptInput.Cancel asChild>
+              <Button variant="ghost" size="sm" aria-label="Cancel edit">
+                Cancel
+              </Button>
+            </ChatPromptInput.Cancel>
+            <Button
+              type="submit"
+              size="sm"
+              aria-label="Update message"
+              disabled={currentText === originalText || hasUploading}
+              className={cn(
+                (currentText === originalText || hasUploading) &&
+                  "opacity-50 cursor-not-allowed",
+              )}
+            >
+              {hasUploading ? (
+                <>
+                  <Loader2 className="size-3.5 animate-spin mr-1.5" />
+                  Uploading...
+                </>
+              ) : (
+                "Update"
+              )}
+            </Button>
+          </div>
+        </div>
+      </ChatPromptInput.Root>
+    </div>
+  );
+};
diff --git a/src/components/assistant-ui/thread/MessageError.tsx b/src/components/assistant-ui/thread/MessageError.tsx
new file mode 100644
index 00000000..44f5551c
--- /dev/null
+++ b/src/components/assistant-ui/thread/MessageError.tsx
@@ -0,0 +1,14 @@
+"use client";
+
+import type { FC } from "react";
+import { ChatError, ChatMessage } from "@/lib/chat/runtime";
+
+export const MessageError: FC = () => {
+  return (
+    <ChatMessage.Error>
+      <ChatError.Root className="aui-message-error-root mt-2 rounded-md border border-destructive bg-destructive/10 p-3 text-sm text-destructive dark:bg-destructive/5 dark:text-red-200">
+        <ChatError.Message className="aui-message-error-message line-clamp-2" />
+      </ChatError.Root>
+    </ChatMessage.Error>
+  );
+};
diff --git a/src/components/assistant-ui/thread/PromptInput.tsx b/src/components/assistant-ui/thread/PromptInput.tsx
new file mode 100644
index 00000000..07431367
--- /dev/null
+++ b/src/components/assistant-ui/thread/PromptInput.tsx
@@ -0,0 +1,159 @@
+"use client";
+
+import { CheckCircle2 } from "lucide-react";
+import { useEffect, useRef, type FC } from "react";
+import {
+  PromptInputAttachments,
+} from "@/components/assistant-ui/attachment";
+import { CardContextDisplay } from "@/components/chat/CardContextDisplay";
+import { MentionMenu } from "@/components/chat/MentionMenu";
+import { ReplyContextDisplay } from "@/components/chat/ReplyContextDisplay";
+import { ChatPromptInput, useMainThreadId, usePromptInput } from "@/lib/chat/runtime";
+import { useAttachmentUploadStore } from "@/lib/stores/attachment-upload-store";
+import { selectReplySelections, useUIStore } from "@/lib/stores/ui-store";
+import { useWorkspaceStore } from "@/lib/stores/workspace-store";
+import { isOfficeDocument } from "@/lib/uploads/office-document-validation";
+import { processPdfAttachmentsInBackground } from "@/lib/uploads/process-pdf-attachments-in-background";
+import type { Item } from "@/lib/workspace-state/types";
+import { useSelectedCardIds } from "@/hooks/ui/use-selected-card-ids";
+import { useWorkspaceOperations } from "@/hooks/workspace/use-workspace-operations";
+import { useWorkspaceState } from "@/hooks/workspace/use-workspace-state";
+import { toast } from "sonner";
+import { useShallow } from "zustand/react/shallow";
+import { PromptInputToolbar } from "./PromptInputToolbar";
+import { useMentionMenu } from "./hooks/use-mention-menu";
+import { usePromptInputPaste } from "./hooks/use-prompt-input-paste";
+
+interface PromptInputProps {
+  items: Item[];
+}
+
+export const PromptInput: FC<PromptInputProps> = ({ items }) => {
+  const currentWorkspaceId = useWorkspaceStore(
+    (state) => state.currentWorkspaceId,
+  );
+  const promptInput = usePromptInput();
+  const replySelections = useUIStore(useShallow(selectReplySelections));
+  const clearReplySelections = useUIStore(
+    (state) => state.clearReplySelections,
+  );
+  const { selectedCardIds } = useSelectedCardIds();
+  const { state: workspaceState } = useWorkspaceState(currentWorkspaceId);
+  const operations = useWorkspaceOperations(currentWorkspaceId, workspaceState);
+  const mainThreadId = useMainThreadId();
+  const inputRef = useRef<HTMLTextAreaElement | null>(null);
+  const toggleCardSelection = useUIStore((state) => state.toggleCardSelection);
+  const mention = useMentionMenu({
+    inputRef,
+    promptInput,
+    onSelectItem: (item) => toggleCardSelection(item.id),
+  });
+  const handlePaste = usePromptInputPaste({
+    promptInput,
+    workspaceId: currentWorkspaceId,
+  });
+
+  useEffect(() => {
+    if (mainThreadId && inputRef.current) {
+      const timeoutId = setTimeout(() => {
+        inputRef.current?.focus();
+      }, 100);
+      return () => clearTimeout(timeoutId);
+    }
+  }, [mainThreadId]);
+
+  return (
+    <ChatPromptInput.Root
+      className="aui-composer-root relative flex w-full flex-col rounded-lg border border-sidebar-border bg-sidebar-accent px-3.5 pt-2 pb-1 shadow-[0_9px_9px_0px_rgba(0,0,0,0.01),0_2px_5px_0px_rgba(0,0,0,0.06)] dark:border-sidebar-border/15"
+      onClick={(e) => {
+        if (inputRef.current && !e.defaultPrevented) {
+          inputRef.current.focus();
+        }
+      }}
+      onSubmit={async (e) => {
+        e.preventDefault();
+
+        if (useAttachmentUploadStore.getState().uploadingIds.size > 0) {
+          toast.info("Please wait for uploads to finish before sending");
+          return;
+        }
+
+        const composerState = promptInput?.getState();
+        if (!composerState) return;
+
+        const currentText = composerState.text;
+        const attachments = composerState.attachments || [];
+        const hasReplyContext = replySelections.length > 0;
+        if (!currentText.trim() && attachments.length === 0 && !hasReplyContext) {
+          return;
+        }
+
+        const pdfAttachments = attachments.filter((att) => {
+          const file = att.file;
+          return (
+            file &&
+            (file.type === "application/pdf" ||
+              file.name.toLowerCase().endsWith(".pdf") ||
+              isOfficeDocument(file))
+          );
+        });
+
+        if (pdfAttachments.length > 0 && currentWorkspaceId) {
+          void processPdfAttachmentsInBackground(
+            pdfAttachments,
+            currentWorkspaceId,
+            operations,
+          );
+        }
+
+        const modifiedText =
+          currentText.trim() || (hasReplyContext ? "Empty message" : "");
+
+        const customMetadata: Record<string, unknown> = {};
+        if (replySelections.length > 0) {
+          customMetadata.replySelections = replySelections;
+        }
+        promptInput?.setRunConfig(
+          Object.keys(customMetadata).length > 0 ? { custom: customMetadata } : {},
+        );
+
+        promptInput?.setText(modifiedText);
+        promptInput?.send();
+        clearReplySelections();
+      }}
+    >
+      <PromptInputAttachments />
+      <CardContextDisplay items={items} />
+      <ReplyContextDisplay />
+      <div className="relative">
+        <ChatPromptInput.Input
+          ref={inputRef}
+          placeholder="Ask anything or @mention items"
+          className="aui-composer-input max-h-32 w-full resize-none bg-transparent py-1.5 text-base text-sidebar-foreground outline-none placeholder:text-sidebar-foreground/60 focus:outline-none"
+          rows={1}
+          autoFocus
+          aria-label="Message input"
+          maxLength={10000}
+          cancelOnEscape={false}
+          onPaste={handlePaste}
+          onKeyDown={mention.handleKeyDown}
+          onInput={mention.handleInput}
+        />
+        <MentionMenu
+          open={mention.mentionMenuOpen}
+          onOpenChange={mention.handleMentionMenuClose}
+          query={mention.mentionQuery}
+          items={items}
+          onSelect={mention.handleMentionSelect}
+          selectedCardIds={selectedCardIds}
+          selectedIndicator={(isSelected) =>
+            isSelected ? (
+              <CheckCircle2 className="size-4 text-primary flex-shrink-0" />
+            ) : undefined
+          }
+        />
+      </div>
+      <PromptInputToolbar />
+    </ChatPromptInput.Root>
+  );
+};
diff --git a/src/components/assistant-ui/thread/PromptInputShell.tsx b/src/components/assistant-ui/thread/PromptInputShell.tsx
new file mode 100644
index 00000000..7fb9b651
--- /dev/null
+++ b/src/components/assistant-ui/thread/PromptInputShell.tsx
@@ -0,0 +1,167 @@
+"use client";
+
+import { useCallback, useEffect, useRef, useState, type FC } from "react";
+import {
+  PromptBuilderDialog,
+  type PromptBuilderAction,
+} from "@/components/assistant-ui/PromptBuilderDialog";
+import {
+  DropdownMenu,
+  DropdownMenuContent,
+  DropdownMenuItem,
+  DropdownMenuTrigger,
+} from "@/components/ui/dropdown-menu";
+import { useHasPromptInputText, useIsThreadEmpty, usePromptInput } from "@/lib/chat/runtime";
+import { cn } from "@/lib/utils";
+import { focusComposerInput } from "@/lib/utils/composer-utils";
+import type { Item } from "@/lib/workspace-state/types";
+import { PROMPT_INPUT_FLOATING_ACTIONS } from "./prompt-input-floating-actions";
+import { PromptInput } from "./PromptInput";
+
+interface PromptInputShellProps {
+  items: Item[];
+}
+
+const FLOATING_MENU_HIDE_DELAY_MS = 400;
+
+export const PromptInputShell: FC<PromptInputShellProps> = ({ items }) => {
+  const [isHovered, setIsHovered] = useState(false);
+  const hideTimeoutRef = useRef<NodeJS.Timeout | null>(null);
+  const promptInput = usePromptInput();
+  const [dialogAction, setDialogAction] = useState<PromptBuilderAction | null>(
+    null,
+  );
+  const isThreadEmpty = useIsThreadEmpty();
+  const hasPromptInputText = useHasPromptInputText();
+
+  const handleDirectFill = useCallback(
+    (fill: string) => {
+      promptInput?.setText(fill);
+      focusComposerInput(true);
+    },
+    [promptInput],
+  );
+
+  const handleMouseEnter = useCallback(() => {
+    if (hideTimeoutRef.current) {
+      clearTimeout(hideTimeoutRef.current);
+      hideTimeoutRef.current = null;
+    }
+    setIsHovered(true);
+  }, []);
+
+  const handleMouseLeave = useCallback(() => {
+    hideTimeoutRef.current = setTimeout(() => {
+      setIsHovered(false);
+    }, FLOATING_MENU_HIDE_DELAY_MS);
+  }, []);
+
+  useEffect(() => {
+    return () => {
+      if (hideTimeoutRef.current) clearTimeout(hideTimeoutRef.current);
+    };
+  }, []);
+
+  return (
+    <div className="relative">
+      <div
+        className="relative"
+        onMouseEnter={handleMouseEnter}
+        onMouseLeave={handleMouseLeave}
+      >
+        <div
+          className={cn(
+            "absolute bottom-full left-0 right-0 z-20 flex justify-center gap-0.5 pb-2",
+            "transition-opacity duration-150 ease-out",
+            !isThreadEmpty && isHovered && !hasPromptInputText
+              ? "opacity-100 pointer-events-auto"
+              : "opacity-0 pointer-events-none",
+          )}
+        >
+          <div className="flex flex-wrap items-center justify-center gap-0.5 rounded-xl border border-sidebar-border bg-sidebar-accent px-1.5 py-1 shadow-md dark:border-sidebar-border/15">
+            {PROMPT_INPUT_FLOATING_ACTIONS.map((action) => {
+              if ("subActions" in action) {
+                const Icon = action.icon;
+                return (
+                  <DropdownMenu key={action.id}>
+                    <DropdownMenuTrigger asChild>
+                      <button
+                        type="button"
+                        className={cn(
+                          "inline-flex cursor-pointer items-center gap-1.5 rounded-lg px-2 py-1 text-xs font-normal",
+                          "text-sidebar-foreground transition-colors",
+                          "hover:bg-sidebar-foreground/10 dark:hover:bg-sidebar-foreground/15",
+                        )}
+                        aria-label={action.label}
+                      >
+                        <Icon className={action.iconClassName} />
+                        <span>{action.label}</span>
+                      </button>
+                    </DropdownMenuTrigger>
+                    <DropdownMenuContent
+                      align="center"
+                      side="top"
+                      className="min-w-[140px]"
+                      sideOffset={4}
+                    >
+                      {(action.subActions ?? []).map((sub) => {
+                        const SubIcon = sub.icon;
+                        return (
+                          <DropdownMenuItem
+                            key={sub.id}
+                            onSelect={() => setDialogAction(sub.action)}
+                            className="flex cursor-pointer items-center gap-2"
+                          >
+                            <SubIcon className={sub.iconClassName} />
+                            {sub.label}
+                          </DropdownMenuItem>
+                        );
+                      })}
+                    </DropdownMenuContent>
+                  </DropdownMenu>
+                );
+              }
+              const Icon = action.icon;
+              return (
+                <button
+                  key={action.id}
+                  type="button"
+                  onClick={() => {
+                    if (action.useDialog && "action" in action && action.action) {
+                      setDialogAction(action.action);
+                    } else if (
+                      "composerFill" in action &&
+                      typeof action.composerFill === "string"
+                    ) {
+                      handleDirectFill(action.composerFill);
+                    }
+                  }}
+                  className={cn(
+                    "inline-flex cursor-pointer items-center gap-1.5 rounded-lg px-2 py-1 text-xs font-normal",
+                    "text-sidebar-foreground transition-colors",
+                    "hover:bg-sidebar-foreground/10 dark:hover:bg-sidebar-foreground/15",
+                  )}
+                  aria-label={action.label}
+                >
+                  <Icon className={action.iconClassName} />
+                  <span>{action.label}</span>
+                </button>
+              );
+            })}
+          </div>
+        </div>
+
+        <PromptInput items={items} />
+
+        {dialogAction && (
+          <PromptBuilderDialog
+            open={!!dialogAction}
+            onOpenChange={(open) => !open && setDialogAction(null)}
+            action={dialogAction}
+            items={items}
+          />
+        )}
+      </div>
+    </div>
+  );
+};
diff --git a/src/components/assistant-ui/thread/PromptInputToolbar.tsx b/src/components/assistant-ui/thread/PromptInputToolbar.tsx
new file mode 100644
index 00000000..a76b43b2
--- /dev/null
+++ b/src/components/assistant-ui/thread/PromptInputToolbar.tsx
@@ -0,0 +1,187 @@
+"use client";
+
+import { PromptInputAddAttachment } from "@/components/assistant-ui/attachment";
+import { AIFeedbackDialog } from "@/components/assistant-ui/AIFeedbackDialog";
+import { ModelPicker } from "@/components/assistant-ui/ModelPicker";
+import { ModelSettingsMenu } from "@/components/assistant-ui/ModelSettingsMenu";
+import { SpeechToTextButton } from "@/components/assistant-ui/SpeechToTextButton";
+import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
+import { Button } from "@/components/ui/button";
+import {
+  Popover,
+  PopoverContent,
+  PopoverTrigger,
+} from "@/components/ui/popover";
+import { useSession } from "@/lib/auth-client";
+import { ChatIf, ChatPromptInput } from "@/lib/chat/runtime";
+import { useAttachmentUploadStore } from "@/lib/stores/attachment-upload-store";
+import { focusComposerInput } from "@/lib/utils/composer-utils";
+import { AlertTriangle, ArrowUpIcon, Bug, Loader2, Square } from "lucide-react";
+import Link from "next/link";
+import { useEffect, useRef, useState, type FC } from "react";
+import { useFeatureFlagEnabled } from "posthog-js/react";
+
+export const PromptInputToolbar: FC = () => {
+  const { data: session } = useSession();
+  const hasUploading = useAttachmentUploadStore((s) => s.uploadingIds.size > 0);
+  const isAnonymous = session?.user?.isAnonymous ?? false;
+  const [isWarningPopoverOpen, setIsWarningPopoverOpen] = useState(false);
+  const hoverTimeoutRef = useRef<NodeJS.Timeout | null>(null);
+  const [isFeedbackDialogOpen, setIsFeedbackDialogOpen] = useState(false);
+  const isDev = process.env.NODE_ENV === "development";
+  const aiDebugFlagEnabled = useFeatureFlagEnabled("ai-debug-feedback");
+  const showAiDebugButton = isDev || aiDebugFlagEnabled === true;
+
+  useEffect(() => {
+    return () => {
+      if (hoverTimeoutRef.current) {
+        clearTimeout(hoverTimeoutRef.current);
+      }
+    };
+  }, []);
+
+  return (
+    <div className="aui-composer-action-wrapper relative mb-2 flex items-center justify-between">
+      <div className="relative z-0 flex items-center gap-0">
+        <div className="relative z-0">
+          <PromptInputAddAttachment />
+        </div>
+        {!isAnonymous && <ModelSettingsMenu />}
+        <div className={!isAnonymous ? "ml-0.5" : undefined}>
+          <ModelPicker />
+        </div>
+        {showAiDebugButton && (
+          <button
+            type="button"
+            className="ml-2 flex items-center gap-1.5 px-1.5 py-1 rounded-md bg-red-500/10 hover:bg-red-500/20 transition-colors flex-shrink-0 text-xs font-normal text-red-500 hover:text-red-400 cursor-pointer border border-red-500/20"
+            onClick={() => {
+              if (process.env.NODE_ENV === "development") {
+                window.open("http://localhost:4983", "_blank");
+              } else {
+                setIsFeedbackDialogOpen(true);
+              }
+            }}
+          >
+            <Bug className="w-3.5 h-3.5" />
+            <span>AI Debug</span>
+          </button>
+        )}
+        <AIFeedbackDialog
+          open={isFeedbackDialogOpen}
+          onOpenChange={setIsFeedbackDialogOpen}
+        />
+        {isAnonymous && (
+          <Popover
+            open={isWarningPopoverOpen}
+            onOpenChange={(open) => {
+              setIsWarningPopoverOpen(open);
+              if (!open) {
+                focusComposerInput();
+              }
+            }}
+          >
+            <PopoverTrigger asChild>
+              <button
+                type="button"
+                onMouseEnter={() => {
+                  if (hoverTimeoutRef.current) {
+                    clearTimeout(hoverTimeoutRef.current);
+                  }
+                  setIsWarningPopoverOpen(true);
+                }}
+                onMouseLeave={() => {
+                  hoverTimeoutRef.current = setTimeout(() => {
+                    setIsWarningPopoverOpen(false);
+                  }, 100);
+                }}
+                className="flex items-center justify-center w-5 h-5 rounded-md hover:bg-accent transition-colors flex-shrink-0 cursor-pointer"
+                aria-label="Warning: AI chats won't save unless logged in"
+              >
+                <AlertTriangle className="w-4 h-4 text-yellow-500 animate-[pulse-scale_3.5s_ease-in-out_infinite]" />
+              </button>
+            </PopoverTrigger>
+            <PopoverContent
+              side="top"
+              align="start"
+              sideOffset={8}
+              onMouseEnter={() => {
+                if (hoverTimeoutRef.current) {
+                  clearTimeout(hoverTimeoutRef.current);
+                }
+                setIsWarningPopoverOpen(true);
+              }}
+              onMouseLeave={() => {
+                hoverTimeoutRef.current = setTimeout(() => {
+                  setIsWarningPopoverOpen(false);
+                }, 100);
+              }}
+              className="w-64 p-3"
+            >
+              <div className="space-y-3">
+                <p className="text-sm text-foreground">
+                  Your AI chats won't save unless you are logged in.
+                </p>
+                <div className="flex items-center gap-2">
+                  <Link href="/auth/sign-in" className="flex-1">
+                    <Button
+                      variant="outline"
+                      size="sm"
+                      className="w-full h-7 text-xs"
+                      onClick={() => setIsWarningPopoverOpen(false)}
+                    >
+                      Sign in
+                    </Button>
+                  </Link>
+                  <Link href="/auth/sign-up" className="flex-1">
+                    <Button
+                      size="sm"
+                      className="w-full h-7 text-xs"
+                      onClick={() => setIsWarningPopoverOpen(false)}
+                    >
+                      Sign up
+                    </Button>
+                  </Link>
+                </div>
+              </div>
+            </PopoverContent>
+          </Popover>
+        )}
+      </div>
+      <div className="flex items-center gap-2">
+        {!isAnonymous && <SpeechToTextButton />}
+        <ChatIf condition={({ thread }) => !thread.isRunning}>
+          <TooltipIconButton
+            tooltip={hasUploading ? "Uploading attachments..." : "Send message"}
+            side="bottom"
+            type="submit"
+            variant="default"
+            size="icon"
+            className="aui-composer-send size-[34px] rounded-full p-1"
+            aria-label="Send message"
+            disabled={hasUploading}
+          >
+            {hasUploading ? (
+              <Loader2 className="aui-composer-send-icon size-4 text-background animate-spin" />
+            ) : (
+              <ArrowUpIcon className="aui-composer-send-icon size-4 text-background" />
+            )}
+          </TooltipIconButton>
+        </ChatIf>
+
+        <ChatIf condition={({ thread }) => thread.isRunning}>
+          <ChatPromptInput.Cancel asChild>
+            <Button
+              type="button"
+              variant="default"
+              size="icon"
+              className="aui-composer-cancel size-[34px] rounded-full border border-muted-foreground/60 hover:bg-primary/75 dark:border-muted-foreground/90"
+              aria-label="Stop generating"
+            >
+              <Square className="aui-composer-cancel-icon size-3 text-background fill-current" />
+            </Button>
+          </ChatPromptInput.Cancel>
+        </ChatIf>
+      </div>
+    </div>
+  );
+};
diff --git a/src/components/assistant-ui/thread/Thread.tsx b/src/components/assistant-ui/thread/Thread.tsx
new file mode 100644
index 00000000..a41089fa
--- /dev/null
+++ b/src/components/assistant-ui/thread/Thread.tsx
@@ -0,0 +1,46 @@
+"use client";
+
+import { useRef, type FC } from "react";
+import { ChatIf, ChatThread } from "@/lib/chat/runtime";
+import type { Item } from "@/lib/workspace-state/types";
+import { PromptInputShell } from "./PromptInputShell";
+import { ThreadLoadingSkeleton } from "./ThreadLoadingSkeleton";
+import { ThreadWelcome } from "./ThreadWelcome";
+import { VirtualizedMessages } from "./VirtualizedMessages";
+
+interface ThreadProps {
+  items?: Item[];
+}
+
+export const Thread: FC<ThreadProps> = ({ items = [] }) => {
+  const viewportRef = useRef<HTMLDivElement>(null);
+
+  return (
+    <ChatThread.Root
+      className="aui-root aui-thread-root @container flex h-full flex-col bg-sidebar"
+      style={{
+        ["--thread-max-width" as string]: "50rem",
+      }}
+    >
+      <ChatThread.Viewport
+        ref={viewportRef}
+        turnAnchor="top"
+        autoScroll={false}
+        className="aui-thread-viewport relative flex min-h-0 flex-1 flex-col overflow-x-auto overflow-y-scroll px-4"
+      >
+        <ChatIf condition={({ thread }) => thread.isLoading}>
+          <ThreadLoadingSkeleton />
+        </ChatIf>
+        <ChatIf condition={({ thread }) => thread.isEmpty && !thread.isLoading}>
+          <ThreadWelcome items={items} />
+        </ChatIf>
+
+        <VirtualizedMessages scrollRef={viewportRef} />
+      </ChatThread.Viewport>
+
+      <div className="aui-thread-composer-wrapper mx-auto flex w-full max-w-[var(--thread-max-width)] flex-shrink-0 flex-col gap-4 overflow-visible rounded-t-3xl bg-sidebar px-4 pb-3 md:pb-4">
+        <PromptInputShell items={items} />
+      </div>
+    </ChatThread.Root>
+  );
+};
diff --git a/src/components/assistant-ui/thread/ThreadLoadingSkeleton.tsx b/src/components/assistant-ui/thread/ThreadLoadingSkeleton.tsx
new file mode 100644
index 00000000..9d943617
--- /dev/null
+++ b/src/components/assistant-ui/thread/ThreadLoadingSkeleton.tsx
@@ -0,0 +1,34 @@
+"use client";
+
+import type { FC } from "react";
+import { Skeleton } from "@/components/ui/skeleton";
+
+export const ThreadLoadingSkeleton: FC = () => {
+  return (
+    <div
+      role="status"
+      aria-label="Loading chat"
+      className="aui-thread-loading-skeleton mx-auto flex w-full max-w-[var(--thread-max-width)] flex-col gap-6 px-2 py-8"
+    >
+      <div className="flex justify-end">
+        <div className="max-w-[85%] space-y-2">
+          <Skeleton className="h-12 w-48 rounded-lg" />
+        </div>
+      </div>
+      <div className="space-y-2">
+        <Skeleton className="h-4 w-full max-w-[90%] rounded" />
+        <Skeleton className="h-4 w-full max-w-[70%] rounded" />
+        <Skeleton className="h-4 w-32 rounded" />
+      </div>
+      <div className="flex justify-end">
+        <Skeleton className="h-10 w-64 rounded-lg" />
+      </div>
+      <div className="space-y-2">
+        <Skeleton className="h-4 w-full max-w-[95%] rounded" />
+        <Skeleton className="h-4 w-full max-w-[80%] rounded" />
+        <Skeleton className="h-4 w-full max-w-[60%] rounded" />
+        <Skeleton className="h-4 w-24 rounded" />
+      </div>
+    </div>
+  );
+};
diff --git a/src/components/assistant-ui/thread/ThreadSuggestions.tsx b/src/components/assistant-ui/thread/ThreadSuggestions.tsx
new file mode 100644
index 00000000..d399728e
--- /dev/null
+++ b/src/components/assistant-ui/thread/ThreadSuggestions.tsx
@@ -0,0 +1,92 @@
+"use client";
+
+import { useCallback, useState, type FC } from "react";
+import {
+  PromptBuilderDialog,
+  type PromptBuilderAction,
+} from "@/components/assistant-ui/PromptBuilderDialog";
+import { Button } from "@/components/ui/button";
+import { usePromptInput } from "@/lib/chat/runtime";
+import { focusComposerInput } from "@/lib/utils/composer-utils";
+import type { Item } from "@/lib/workspace-state/types";
+import { SUGGESTION_ACTIONS } from "./suggestion-actions";
+
+interface ThreadSuggestionsProps {
+  items: Item[];
+}
+
+export const ThreadSuggestions: FC<ThreadSuggestionsProps> = ({ items }) => {
+  const promptInput = usePromptInput();
+  const [dialogAction, setDialogAction] = useState<PromptBuilderAction | null>(
+    null,
+  );
+
+  const handleDirectFill = useCallback(
+    (action: string) => {
+      promptInput?.setText(action);
+      focusComposerInput(true);
+    },
+    [promptInput],
+  );
+
+  const handleTriggerFileInput = useCallback(() => {
+    document.getElementById("prompt-input-file-upload")?.click();
+    focusComposerInput(true);
+  }, []);
+
+  return (
+    <>
+      <div className="aui-thread-welcome-suggestions grid w-full grid-cols-2 gap-2 pb-4 sm:grid-cols-3">
+        {SUGGESTION_ACTIONS.map((suggestedAction, index) => {
+          const Icon = suggestedAction.icon;
+          return (
+            <div
+              key={`suggested-action-${suggestedAction.title}-${index}`}
+              className="aui-thread-welcome-suggestion-display"
+            >
+              <Button
+                type="button"
+                variant="ghost"
+                className="aui-thread-welcome-suggestion h-auto w-full flex-1 flex-wrap items-center justify-start gap-2 rounded-lg border border-sidebar-border px-5 py-4 text-left text-sm dark:hover:bg-accent/60"
+                aria-label={suggestedAction.title}
+                onClick={() => {
+                  if (
+                    "triggerFileInput" in suggestedAction &&
+                    suggestedAction.triggerFileInput
+                  ) {
+                    handleTriggerFileInput();
+                  } else if (
+                    suggestedAction.useDialog &&
+                    "action" in suggestedAction &&
+                    suggestedAction.action
+                  ) {
+                    setDialogAction(suggestedAction.action);
+                  } else if (
+                    "composerFill" in suggestedAction &&
+                    typeof suggestedAction.composerFill === "string"
+                  ) {
+                    handleDirectFill(suggestedAction.composerFill);
+                  }
+                }}
+              >
+                <Icon className={suggestedAction.iconClassName} />
+                <span className="aui-thread-welcome-suggestion-text-1 font-medium">
+                  {suggestedAction.title}
+                </span>
+              </Button>
+            </div>
+          );
+        })}
+      </div>
+
+      {dialogAction && (
+        <PromptBuilderDialog
+          open={!!dialogAction}
+          onOpenChange={(open) => !open && setDialogAction(null)}
+          action={dialogAction}
+          items={items}
+        />
+      )}
+    </>
+  );
+};
diff --git a/src/components/assistant-ui/thread/ThreadWelcome.tsx b/src/components/assistant-ui/thread/ThreadWelcome.tsx
new file mode 100644
index 00000000..b70fca34
--- /dev/null
+++ b/src/components/assistant-ui/thread/ThreadWelcome.tsx
@@ -0,0 +1,27 @@
+"use client";
+
+import type { FC } from "react";
+import { ThinkExLogo } from "@/components/ui/thinkex-logo";
+import type { Item } from "@/lib/workspace-state/types";
+import { ThreadSuggestions } from "./ThreadSuggestions";
+
+interface ThreadWelcomeProps {
+  items: Item[];
+}
+
+export const ThreadWelcome: FC<ThreadWelcomeProps> = ({ items }) => {
+  return (
+    <div className="aui-thread-welcome-root mx-auto my-auto flex w-full max-w-[var(--thread-max-width)] flex-grow flex-col">
+      <div className="aui-thread-welcome-center flex w-full flex-grow flex-col items-center justify-center">
+        <div className="aui-thread-welcome-message flex size-full flex-col items-center justify-center px-8">
+          <div className="aui-thread-welcome-message-motion-0 mb-1 flex justify-center">
+            <div className="relative flex h-12 w-12 items-center justify-center">
+              <ThinkExLogo size={48} priority />
+            </div>
+          </div>
+        </div>
+      </div>
+      <ThreadSuggestions items={items} />
+    </div>
+  );
+};
diff --git a/src/components/assistant-ui/thread/UserActionBar.tsx b/src/components/assistant-ui/thread/UserActionBar.tsx
new file mode 100644
index 00000000..5fe6b16a
--- /dev/null
+++ b/src/components/assistant-ui/thread/UserActionBar.tsx
@@ -0,0 +1,46 @@
+"use client";
+
+import { CheckIcon, CopyIcon, PencilIcon } from "lucide-react";
+import { useCallback, useMemo, useRef, useState, type FC } from "react";
+import { ChatActionBar, useChatMessage } from "@/lib/chat/runtime";
+import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
+
+export const UserActionBar: FC = () => {
+  const message = useChatMessage();
+  const [copied, setCopied] = useState(false);
+  const copyTimeoutRef = useRef<NodeJS.Timeout | null>(null);
+
+  const textContent = useMemo(() => {
+    return message.content
+      .filter(
+        (part): part is { type: "text"; text: string } => part.type === "text",
+      )
+      .map((part) => part.text ?? "")
+      .join("\n\n");
+  }, [message.content]);
+
+  const handleCopy = useCallback(() => {
+    if (!textContent) return;
+    navigator.clipboard.writeText(textContent);
+    setCopied(true);
+    if (copyTimeoutRef.current) clearTimeout(copyTimeoutRef.current);
+    copyTimeoutRef.current = setTimeout(() => setCopied(false), 2000);
+  }, [textContent]);
+
+  return (
+    <ChatActionBar.Root
+      hideWhenRunning
+      autohide="not-last"
+      className="aui-user-action-bar-root flex gap-1 text-muted-foreground data-floating:absolute data-floating:rounded-md data-floating:border data-floating:bg-background data-floating:p-1 data-floating:shadow-sm"
+    >
+      <TooltipIconButton tooltip="Copy" onClick={handleCopy}>
+        {copied ? <CheckIcon /> : <CopyIcon />}
+      </TooltipIconButton>
+      <ChatActionBar.Edit asChild>
+        <TooltipIconButton tooltip="Edit">
+          <PencilIcon />
+        </TooltipIconButton>
+      </ChatActionBar.Edit>
+    </ChatActionBar.Root>
+  );
+};
diff --git a/src/components/assistant-ui/thread/UserMessage.tsx b/src/components/assistant-ui/thread/UserMessage.tsx
new file mode 100644
index 00000000..928db68d
--- /dev/null
+++ b/src/components/assistant-ui/thread/UserMessage.tsx
@@ -0,0 +1,92 @@
+"use client";
+
+import { ChevronDown, ChevronUp } from "lucide-react";
+import { useMemo, useState, type FC } from "react";
+import { UserMessageAttachments } from "@/components/assistant-ui/attachment";
+import { File as FileComponent } from "@/components/assistant-ui/file";
+import { MessageContextBadges } from "@/components/chat/MessageContextBadges";
+import { ChatMessage, useChatMessage } from "@/lib/chat/runtime";
+import { BranchPicker } from "./BranchPicker";
+import {
+  USER_MESSAGE_MAX_CHARS,
+  UserMessageText,
+  UserMessageTruncateContext,
+} from "./UserMessageTruncateContext";
+import { UserActionBar } from "./UserActionBar";
+
+export const UserMessage: FC = () => {
+  const [expanded, setExpanded] = useState(false);
+  const message = useChatMessage();
+
+  const textLength = useMemo(
+    () =>
+      message.content
+        .filter(
+          (part): part is { type: "text"; text: string } =>
+            part.type === "text",
+        )
+        .reduce((sum, part) => sum + (part.text?.length ?? 0), 0),
+    [message.content],
+  );
+
+  const showExpand = textLength > USER_MESSAGE_MAX_CHARS;
+
+  const truncateCtxValue = useMemo(
+    () => ({
+      maxChars: USER_MESSAGE_MAX_CHARS,
+      expanded,
+      showExpand,
+    }),
+    [expanded, showExpand],
+  );
+
+  return (
+    <ChatMessage.Root asChild>
+      <div
+        className="aui-user-message-root mx-auto grid w-full max-w-[var(--thread-max-width)] animate-breathe-in auto-rows-auto grid-cols-[minmax(72px,1fr)_auto] gap-y-2 px-2 pt-4 pb-1 last:mb-5 [&:where(>*)]:col-start-2"
+        data-role="user"
+      >
+        <UserMessageAttachments />
+
+        <div className="aui-user-message-content-wrapper relative col-start-2 min-w-0">
+          <MessageContextBadges />
+          <UserMessageTruncateContext.Provider value={truncateCtxValue}>
+            <div className="aui-user-message-content relative rounded-lg bg-muted px-3 py-2 break-words text-foreground text-sm">
+              <ChatMessage.Parts
+                components={{
+                  Text: UserMessageText,
+                  File: FileComponent,
+                }}
+              />
+              {showExpand && (
+                <div className="flex justify-end pt-1.5 mt-1.5">
+                  <button
+                    type="button"
+                    onClick={() => setExpanded((e) => !e)}
+                    className="flex items-center gap-1 text-xs text-muted-foreground hover:text-foreground transition-colors"
+                    aria-label={expanded ? "Show less" : "Show more"}
+                  >
+                    {expanded ? (
+                      <ChevronUp className="size-3.5" />
+                    ) : (
+                      <ChevronDown className="size-3.5" />
+                    )}
+                    {expanded ? "Show less" : "Show more"}
+                  </button>
+                </div>
+              )}
+            </div>
+          </UserMessageTruncateContext.Provider>
+        </div>
+
+        <div className="aui-user-message-footer ml-2 flex justify-end col-start-2 relative min-h-[20px]">
+          <div className="absolute right-0">
+            <UserActionBar />
+          </div>
+        </div>
+
+        <BranchPicker className="aui-user-branch-picker col-span-full col-start-1 row-start-3 -mr-1 justify-end" />
+      </div>
+    </ChatMessage.Root>
+  );
+};
diff --git a/src/components/assistant-ui/thread/UserMessageTruncateContext.tsx b/src/components/assistant-ui/thread/UserMessageTruncateContext.tsx
new file mode 100644
index 00000000..d7304e4e
--- /dev/null
+++ b/src/components/assistant-ui/thread/UserMessageTruncateContext.tsx
@@ -0,0 +1,30 @@
+"use client";
+
+import { createContext, useContext, type FC } from "react";
+import { useChatMessagePartText } from "@/lib/chat/runtime";
+
+export const USER_MESSAGE_MAX_CHARS = 250;
+
+export const UserMessageTruncateContext = createContext<{
+  maxChars: number;
+  expanded: boolean;
+  showExpand: boolean;
+} | null>(null);
+
+export const UserMessageText: FC = () => {
+  const { text: rawText } = useChatMessagePartText();
+  const truncateCtx = useContext(UserMessageTruncateContext);
+
+  let text = rawText;
+
+  if (
+    truncateCtx &&
+    !truncateCtx.expanded &&
+    truncateCtx.maxChars < Infinity &&
+    text.length > truncateCtx.maxChars
+  ) {
+    text = text.slice(0, truncateCtx.maxChars).trim() + "...";
+  }
+
+  return <div className="whitespace-pre-wrap">{text}</div>;
+};
diff --git a/src/components/assistant-ui/thread/VirtualizedMessages.tsx b/src/components/assistant-ui/thread/VirtualizedMessages.tsx
new file mode 100644
index 00000000..df32bcca
--- /dev/null
+++ b/src/components/assistant-ui/thread/VirtualizedMessages.tsx
@@ -0,0 +1,55 @@
+"use client";
+
+import type { FC, RefObject } from "react";
+import { useVirtualizer } from "@tanstack/react-virtual";
+import { ChatThread, useThreadMessageCount } from "@/lib/chat/runtime";
+import { MESSAGE_COMPONENTS } from "./message-components";
+
+interface VirtualizedMessagesProps {
+  scrollRef: RefObject<HTMLDivElement | null>;
+}
+
+export const VirtualizedMessages: FC<VirtualizedMessagesProps> = ({
+  scrollRef,
+}) => {
+  const messageCount = useThreadMessageCount();
+
+  const virtualizer = useVirtualizer({
+    count: messageCount,
+    getScrollElement: () => scrollRef.current,
+    estimateSize: () => 350,
+    overscan: 5,
+  });
+
+  if (messageCount === 0) return null;
+
+  return (
+    <div
+      style={{
+        height: `${virtualizer.getTotalSize()}px`,
+        width: "100%",
+        position: "relative",
+      }}
+    >
+      {virtualizer.getVirtualItems().map((virtualRow) => (
+        <div
+          key={virtualRow.key}
+          data-index={virtualRow.index}
+          ref={virtualizer.measureElement}
+          style={{
+            position: "absolute",
+            top: 0,
+            left: 0,
+            width: "100%",
+            transform: `translateY(${virtualRow.start}px)`,
+          }}
+        >
+          <ChatThread.MessageByIndex
+            index={virtualRow.index}
+            components={MESSAGE_COMPONENTS}
+          />
+        </div>
+      ))}
+    </div>
+  );
+};
diff --git a/src/components/assistant-ui/thread/hooks/use-mention-menu.ts b/src/components/assistant-ui/thread/hooks/use-mention-menu.ts
new file mode 100644
index 00000000..d3167dfd
--- /dev/null
+++ b/src/components/assistant-ui/thread/hooks/use-mention-menu.ts
@@ -0,0 +1,153 @@
+"use client";
+
+import {
+  useCallback,
+  useState,
+  type FormEvent,
+  type KeyboardEvent,
+  type RefObject,
+} from "react";
+import type { ComposerActions } from "@/lib/chat/runtime";
+import type { Item } from "@/lib/workspace-state/types";
+
+interface UseMentionMenuArgs {
+  inputRef: RefObject<HTMLTextAreaElement | null>;
+  promptInput: ComposerActions | null;
+  onSelectItem: (item: Item) => void;
+}
+
+interface UseMentionMenuResult {
+  mentionMenuOpen: boolean;
+  mentionQuery: string;
+  handleInput: (e: FormEvent<HTMLTextAreaElement>) => void;
+  handleKeyDown: (e: KeyboardEvent<HTMLTextAreaElement>) => void;
+  handleMentionMenuClose: (open: boolean) => void;
+  handleMentionSelect: (item: Item) => void;
+}
+
+export function useMentionMenu({
+  inputRef,
+  promptInput,
+  onSelectItem,
+}: UseMentionMenuArgs): UseMentionMenuResult {
+  const [mentionMenuOpen, setMentionMenuOpen] = useState(false);
+  const [mentionQuery, setMentionQuery] = useState("");
+  const [mentionStartIndex, setMentionStartIndex] = useState<number | null>(
+    null,
+  );
+
+  const handleInput = useCallback(
+    (e: FormEvent<HTMLTextAreaElement>) => {
+      const textarea = e.currentTarget;
+      const value = textarea.value;
+      const cursorPos = textarea.selectionStart ?? 0;
+
+      if (mentionStartIndex !== null) {
+        const query = value.slice(mentionStartIndex + 1, cursorPos);
+
+        if (
+          cursorPos <= mentionStartIndex ||
+          query.includes(" ") ||
+          query.includes("\n")
+        ) {
+          setMentionMenuOpen(false);
+          setMentionStartIndex(null);
+          setMentionQuery("");
+        } else {
+          setMentionQuery(query);
+        }
+      }
+    },
+    [mentionStartIndex],
+  );
+
+  const handleKeyDown = useCallback(
+    (e: KeyboardEvent<HTMLTextAreaElement>) => {
+      const textarea = e.currentTarget;
+
+      if (e.key === "@" && !mentionMenuOpen) {
+        const cursorPos = textarea.selectionStart ?? 0;
+        const charBefore = cursorPos > 0 ? textarea.value[cursorPos - 1] : " ";
+        if (charBefore === " " || charBefore === "\n" || cursorPos === 0) {
+          setMentionMenuOpen(true);
+          setMentionStartIndex(cursorPos);
+          setMentionQuery("");
+        }
+      }
+
+      if (e.key === "Escape" && mentionMenuOpen) {
+        e.preventDefault();
+        setMentionMenuOpen(false);
+        setMentionStartIndex(null);
+        setMentionQuery("");
+      }
+
+      if (
+        mentionMenuOpen &&
+        ["ArrowUp", "ArrowDown", "Enter"].includes(e.key)
+      ) {
+        e.preventDefault();
+      }
+    },
+    [mentionMenuOpen],
+  );
+
+  const clearMentionQuery = useCallback(() => {
+    if (mentionStartIndex !== null && inputRef.current) {
+      const textarea = inputRef.current;
+      const currentValue = textarea.value;
+      const atSymbolIndex = mentionStartIndex;
+
+      let queryEndIndex = mentionStartIndex;
+      while (
+        queryEndIndex < currentValue.length &&
+        currentValue[queryEndIndex] !== " " &&
+        currentValue[queryEndIndex] !== "\n"
+      ) {
+        queryEndIndex++;
+      }
+
+      const textBefore = currentValue.substring(0, atSymbolIndex);
+      const textAfter = currentValue.substring(queryEndIndex);
+      const newValue = textBefore + textAfter;
+
+      promptInput?.setText(newValue);
+      setMentionQuery("");
+      setMentionStartIndex(null);
+
+      setTimeout(() => {
+        if (inputRef.current) {
+          inputRef.current.focus();
+          const newCursorPos = textBefore.length;
+          inputRef.current.setSelectionRange(newCursorPos, newCursorPos);
+        }
+      }, 0);
+    }
+  }, [inputRef, mentionStartIndex, promptInput]);
+
+  const handleMentionSelect = useCallback(
+    (item: Item) => {
+      onSelectItem(item);
+    },
+    [onSelectItem],
+  );
+
+  const handleMentionMenuClose = useCallback(
+    (open: boolean) => {
+      if (!open) {
+        clearMentionQuery();
+      }
+      setMentionMenuOpen(open);
+    },
+    [clearMentionQuery],
+  );
+
+  return {
+    mentionMenuOpen,
+    mentionQuery,
+    handleInput,
+    handleKeyDown,
+    handleMentionMenuClose,
+    handleMentionSelect,
+  };
+}
diff --git a/src/components/assistant-ui/thread/hooks/use-prompt-input-paste.ts b/src/components/assistant-ui/thread/hooks/use-prompt-input-paste.ts
new file mode 100644
index 00000000..190fa681
--- /dev/null
+++ b/src/components/assistant-ui/thread/hooks/use-prompt-input-paste.ts
@@ -0,0 +1,53 @@
+"use client";
+
+import type { ClipboardEvent } from "react";
+import type { ComposerActions } from "@/lib/chat/runtime";
+
+interface UsePromptInputPasteArgs {
+  promptInput: ComposerActions | null;
+  workspaceId: string | null;
+}
+
+export function usePromptInputPaste({
+  promptInput,
+  workspaceId,
+}: UsePromptInputPasteArgs) {
+  return async function handlePaste(e: ClipboardEvent<HTMLTextAreaElement>) {
+    const clipboardData = e.clipboardData;
+    if (!clipboardData || !workspaceId) return;
+
+    const files = Array.from(clipboardData.files) as File[];
+
+    if (files.length > 0) {
+      e.preventDefault();
+      const imageFile = files.find((file: File) => file.type.startsWith("image/"));
+      const fileToUpload = imageFile || files[0];
+
+      if (fileToUpload) {
+        try {
+          await promptInput?.addAttachment(fileToUpload);
+        } catch (error) {
+          console.error("Failed to add file attachment:", error);
+        }
+      }
+      return;
+    }
+
+    const clipboardItems = Array.from(clipboardData.items) as DataTransferItem[];
+    const imageItem = clipboardItems.find((item: DataTransferItem) =>
+      item.type.startsWith("image/"),
+    );
+
+    if (imageItem) {
+      e.preventDefault();
+      const file = imageItem.getAsFile();
+      if (file) {
+        try {
+          await promptInput?.addAttachment(file);
+        } catch (error) {
+          console.error("Failed to add image attachment:", error);
+        }
+      }
+    }
+  };
+}
diff --git a/src/components/assistant-ui/thread/index.ts b/src/components/assistant-ui/thread/index.ts
new file mode 100644
index 00000000..a6166749
--- /dev/null
+++ b/src/components/assistant-ui/thread/index.ts
@@ -0,0 +1 @@
+export { Thread } from "./Thread";
diff --git a/src/components/assistant-ui/thread/message-components.ts b/src/components/assistant-ui/thread/message-components.ts
new file mode 100644
index 00000000..bfc2cb17
--- /dev/null
+++ b/src/components/assistant-ui/thread/message-components.ts
@@ -0,0 +1,9 @@
+import { AssistantMessage } from "./AssistantMessage";
+import { EditPromptInput } from "./EditPromptInput";
+import { UserMessage } from "./UserMessage";
+
+export const MESSAGE_COMPONENTS = {
+  UserMessage,
+  EditComposer: EditPromptInput,
+  AssistantMessage,
+};
diff --git a/src/components/assistant-ui/thread/prompt-input-floating-actions.ts b/src/components/assistant-ui/thread/prompt-input-floating-actions.ts
new file mode 100644
index 00000000..4795a167
--- /dev/null
+++ b/src/components/assistant-ui/thread/prompt-input-floating-actions.ts
@@ -0,0 +1,53 @@
+import { Brain, FileText, Play, Search } from "lucide-react";
+import { LuBook } from "react-icons/lu";
+import { PiCardsThreeBold } from "react-icons/pi";
+import type { PromptBuilderAction } from "@/components/assistant-ui/PromptBuilderDialog";
+
+export const PROMPT_INPUT_FLOATING_ACTIONS = [
+  {
+    id: "document",
+    label: "Document",
+    icon: FileText,
+    iconClassName: "size-3.5 shrink-0 text-sky-400",
+    action: "document" as PromptBuilderAction,
+    useDialog: true,
+  },
+  {
+    id: "learn",
+    label: "Learn",
+    icon: LuBook,
+    iconClassName: "size-3.5 shrink-0 text-amber-500",
+    subActions: [
+      {
+        id: "flashcards",
+        label: "Flashcards",
+        icon: PiCardsThreeBold,
+        iconClassName: "size-4 text-purple-400 rotate-180",
+        action: "flashcards" as PromptBuilderAction,
+      },
+      {
+        id: "quiz",
+        label: "Quiz",
+        icon: Brain,
+        iconClassName: "size-4 text-green-400",
+        action: "quiz" as PromptBuilderAction,
+      },
+    ],
+  },
+  {
+    id: "youtube",
+    label: "YouTube",
+    icon: Play,
+    iconClassName: "size-3.5 text-red-500",
+    action: "youtube" as PromptBuilderAction,
+    useDialog: true,
+  },
+  {
+    id: "search",
+    label: "Search",
+    icon: Search,
+    iconClassName: "size-3.5 text-teal-500",
+    action: "search" as PromptBuilderAction,
+    useDialog: true,
+  },
+];
diff --git a/src/components/assistant-ui/thread/suggestion-actions.ts b/src/components/assistant-ui/thread/suggestion-actions.ts
new file mode 100644
index 00000000..0b135af8
--- /dev/null
+++ b/src/components/assistant-ui/thread/suggestion-actions.ts
@@ -0,0 +1,47 @@
+import { Brain, FileText, Play, Search, Upload } from "lucide-react";
+import { PiCardsThreeBold } from "react-icons/pi";
+import type { PromptBuilderAction } from "@/components/assistant-ui/PromptBuilderDialog";
+
+export const SUGGESTION_ACTIONS = [
+  {
+    title: "Search",
+    icon: Search,
+    iconClassName: "size-4 shrink-0 text-sky-500",
+    action: "search" as PromptBuilderAction,
+    useDialog: true,
+  },
+  {
+    title: "Flashcards",
+    icon: PiCardsThreeBold,
+    iconClassName: "size-4 shrink-0 text-purple-400 rotate-180",
+    action: "flashcards" as PromptBuilderAction,
+    useDialog: true,
+  },
+  {
+    title: "YouTube",
+    icon: Play,
+    iconClassName: "size-4 shrink-0 text-red-500",
+    action: "youtube" as PromptBuilderAction,
+    useDialog: true,
+  },
+  {
+    title: "Upload",
+    icon: Upload,
+    iconClassName: "size-4 shrink-0 text-red-400",
+    triggerFileInput: true,
+  },
+  {
+    title: "Quiz",
+    icon: Brain,
+    iconClassName: "size-4 shrink-0 text-green-400",
+    action: "quiz" as PromptBuilderAction,
+    useDialog: true,
+  },
+  {
+    title: "Document",
+    icon: FileText,
+    iconClassName: "size-4 shrink-0 text-sky-400",
+    action: "document" as PromptBuilderAction,
+    useDialog: true,
+  },
+];
diff --git a/src/lib/chat/runtime/hooks.ts b/src/lib/chat/runtime/hooks.ts
new file mode 100644
index 00000000..cf5e9652
--- /dev/null
+++ b/src/lib/chat/runtime/hooks.ts
@@ -0,0 +1,70 @@
+"use client";
+
+import {
+  useAui,
+  useAuiState,
+  useMessage,
+  useMessagePartText,
+} from "@assistant-ui/react";
+import type {
+  ChatMessage,
+  ComposerActions,
+  ComposerStateSnapshot,
+  ThreadState,
+} from "./types";
+
+export function useThreadState(): ThreadState {
+  return useAuiState((s: any) => ({
+    messageCount: s.thread?.messages?.length ?? 0,
+    isLoading: !!s.thread?.isLoading,
+    isEmpty: s.thread?.isEmpty ?? true,
+    isRunning: !!s.thread?.isRunning,
+  }));
+}
+
+export function useIsThreadLoading(): boolean {
+  return useAuiState((s: any) => !!s.thread?.isLoading);
+}
+
+export function useIsThreadEmpty(): boolean {
+  return useAuiState((s: any) => s.thread?.isEmpty ?? true);
+}
+
+export function useIsThreadRunning(): boolean {
+  return useAuiState((s: any) => !!s.thread?.isRunning);
+}
+
+export function useThreadMessageCount(): number {
+  return useAuiState((s: any) => s.thread?.messages?.length ?? 0);
+}
+
+export function useMainThreadId(): string | null {
+  return useAuiState((s: any) => s.threads?.mainThreadId ?? null);
+}
+
+export function useHasPromptInputText(): boolean {
+  return useAuiState((s: any) => Boolean((s.composer?.text ?? "").trim()));
+}
+
+export function useChatMessage(): ChatMessage {
+  const msg = useMessage();
+  return msg as unknown as ChatMessage;
+}
+
+export function useChatMessagePartText() {
+  return useMessagePartText();
+}
+
+export function usePromptInput(): ComposerActions | null {
+  const aui = useAui();
+  if (!aui) return null;
+  const composer = aui.composer?.();
+  if (!composer) return null;
+  return {
+    setText: (t) => composer.setText(t),
+    send: () => composer.send(),
+    addAttachment: (f) => composer.addAttachment(f),
+    setRunConfig: (cfg) => composer.setRunConfig(cfg as any),
+    getState: () => composer.getState() as unknown as ComposerStateSnapshot | undefined,
+  };
+}
diff --git a/src/lib/chat/runtime/index.ts b/src/lib/chat/runtime/index.ts
new file mode 100644
index 00000000..6e2e3156
--- /dev/null
+++ b/src/lib/chat/runtime/index.ts
@@ -0,0 +1,11 @@
+export type {
+  ChatMessage as ChatMessageData,
+  ChatMessagePart,
+  ChatMessageRole,
+  ChatTextPart,
+  ComposerActions,
+  ComposerStateSnapshot,
+  ThreadState,
+} from "./types";
+export * from "./primitives";
+export * from "./hooks";
diff --git a/src/lib/chat/runtime/primitives.ts b/src/lib/chat/runtime/primitives.ts
new file mode 100644
index 00000000..7f85ad6d
--- /dev/null
+++ b/src/lib/chat/runtime/primitives.ts
@@ -0,0 +1,20 @@
+import {
+  ActionBarPrimitive,
+  AuiIf,
+  BranchPickerPrimitive,
+  ComposerPrimitive,
+  ErrorPrimitive,
+  MessagePrimitive,
+  ThreadPrimitive,
+} from "@assistant-ui/react";
+import type { BranchPickerPrimitive as _BPP } from "@assistant-ui/react";
+
+export const ChatThread = ThreadPrimitive;
+export const ChatMessage = MessagePrimitive;
+export const ChatPromptInput = ComposerPrimitive;
+export const ChatActionBar = ActionBarPrimitive;
+export const ChatBranchPicker = BranchPickerPrimitive;
+export const ChatError = ErrorPrimitive;
+export const ChatIf = AuiIf;
+
+export type ChatBranchPickerRootProps = _BPP.Root.Props;
diff --git a/src/lib/chat/runtime/types.ts b/src/lib/chat/runtime/types.ts
new file mode 100644
index 00000000..1907ca84
--- /dev/null
+++ b/src/lib/chat/runtime/types.ts
@@ -0,0 +1,45 @@
+export type ChatMessageRole = "user" | "assistant" | "system";
+
+export type ChatTextPart = { type: "text"; text: string };
+export type ChatMessagePart =
+  | ChatTextPart
+  | { type: "file"; [k: string]: unknown }
+  | { type: "image"; [k: string]: unknown }
+  | { type: "source"; [k: string]: unknown }
+  | { type: "reasoning"; [k: string]: unknown }
+  | { type: "tool-call" | "tool-result" | "tool"; [k: string]: unknown }
+  | { type: string; [k: string]: unknown };
+
+export interface ChatMessage {
+  id?: string;
+  role: ChatMessageRole;
+  content: ChatMessagePart[];
+  metadata?: { custom?: Record<string, unknown>; [k: string]: unknown };
+}
+
+export interface ThreadState {
+  messageCount: number;
+  isLoading: boolean;
+  isEmpty: boolean;
+  isRunning: boolean;
+}
+
+export interface ComposerStateSnapshot {
+  text: string;
+  attachments: Array<{
+    id?: string;
+    file?: File;
+    name?: string;
+    type?: string;
+    [k: string]: unknown;
+  }>;
+  runConfig?: { custom?: Record<string, unknown> };
+}
+
+export interface ComposerActions {
+  setText(text: string): void;
+  send(): void;
+  addAttachment(file: File): Promise<unknown>;
+  setRunConfig(config: { custom?: Record<string, unknown> }): void;
+  getState(): ComposerStateSnapshot | undefined;
+}
diff --git a/src/lib/uploads/process-pdf-attachments-in-background.ts b/src/lib/uploads/process-pdf-attachments-in-background.ts
new file mode 100644
index 00000000..89f77204
--- /dev/null
+++ b/src/lib/uploads/process-pdf-attachments-in-background.ts
@@ -0,0 +1,67 @@
+import { toast } from "sonner";
+import { buildWorkspaceItemDefinitionsFromAssets } from "@/lib/uploads/uploaded-asset";
+import { uploadSelectedFiles } from "@/lib/uploads/upload-selection";
+import { startAssetProcessing } from "@/lib/uploads/start-asset-processing";
+import {
+  getDocumentUploadFailureMessage,
+  getDocumentUploadPartialMessage,
+  getDocumentUploadSuccessMessage,
+} from "@/lib/uploads/upload-feedback";
+
+interface PdfAttachmentLike {
+  file?: File;
+}
+
+interface WorkspaceOperationsLike {
+  createItems: (
+    defs: ReturnType<typeof buildWorkspaceItemDefinitionsFromAssets>,
+    options?: { showSuccessToast?: boolean },
+  ) => string[];
+}
+
+export async function processPdfAttachmentsInBackground(
+  pdfAttachments: PdfAttachmentLike[],
+  workspaceId: string,
+  operations: WorkspaceOperationsLike,
+) {
+  let files: File[] = [];
+  try {
+    files = pdfAttachments
+      .map((attachment) => attachment.file)
+      .filter((file): file is File => !!file);
+    const { uploads, failedFiles } = await uploadSelectedFiles(files);
+
+    if (uploads.length > 0) {
+      const pdfCardDefinitions = buildWorkspaceItemDefinitionsFromAssets(uploads);
+      const createdIds = operations.createItems(pdfCardDefinitions, {
+        showSuccessToast: false,
+      });
+
+      void startAssetProcessing({
+        workspaceId,
+        assets: uploads,
+        itemIds: createdIds,
+        onOcrError: (error) => {
+          console.error("Error starting assistant file processing:", error);
+        },
+      });
+
+      if (failedFiles.length === 0) {
+        toast.success(getDocumentUploadSuccessMessage(uploads.length));
+      } else {
+        toast.warning(
+          getDocumentUploadPartialMessage(uploads.length, failedFiles.length),
+        );
+      }
+    } else {
+      toast.error(
+        getDocumentUploadFailureMessage(failedFiles.length || files.length),
+      );
+    }
+  } catch (error) {
+    console.error("Error creating PDF cards in background:", error);
+    toast.error(
+      getDocumentUploadFailureMessage(files.length || pdfAttachments.length),
+    );
+  }
+}

From 57d1c58f9c50acc7fed1912dda2457a3278b67ca Mon Sep 17 00:00:00 2001
From: urjitc <135136842+urjitc@users.noreply.github.com>
Date: Sun, 19 Apr 2026 21:50:02 +0000
Subject: [PATCH 2/7] Phase 2a: migrate rendering primitives (markdown-text,
 reasoning, tool-group, assistant-loader, file, image, sources, tool-fallback)
 to chat-runtime ACL

---
 .../assistant-ui/assistant-loader.tsx         | 11 +--
 src/components/assistant-ui/file.tsx          |  6 +-
 src/components/assistant-ui/image.tsx         |  6 +-
 src/components/assistant-ui/markdown-text.tsx | 15 ++--
 src/components/assistant-ui/reasoning.tsx     | 44 +++-------
 src/components/assistant-ui/sources.tsx       |  6 +-
 src/components/assistant-ui/tool-fallback.tsx |  4 +-
 src/components/assistant-ui/tool-group.tsx    | 23 ++---
 src/lib/chat/runtime/hooks.ts                 | 85 +++++++++++++++++++
 src/lib/chat/runtime/index.ts                 |  7 ++
 src/lib/chat/runtime/types.ts                 | 21 +++++
 11 files changed, 159 insertions(+), 69 deletions(-)

diff --git a/src/components/assistant-ui/assistant-loader.tsx b/src/components/assistant-ui/assistant-loader.tsx
index 0ae83746..dc7bfc93 100644
--- a/src/components/assistant-ui/assistant-loader.tsx
+++ b/src/components/assistant-ui/assistant-loader.tsx
@@ -1,19 +1,14 @@
 "use client";
 
-import { useAuiState } from "@assistant-ui/react";
+import { useIsMessageEmpty, useIsMessageRunning } from "@/lib/chat/runtime";
 import { DotLottieReact } from "@lottiefiles/dotlottie-react";
 import { useTheme } from "next-themes";
 
 export const AssistantLoader = () => {
     const { resolvedTheme } = useTheme();
-    const isRunning = useAuiState(
-        ({ message }) => (message as { status?: { type: string } })?.status?.type === "running"
-    );
+    const isRunning = useIsMessageRunning();
 
-    const isMessageEmpty = useAuiState(({ message }) => {
-        const msg = message as any;
-        return !msg?.content || (Array.isArray(msg.content) && msg.content.length === 0);
-    });
+    const isMessageEmpty = useIsMessageEmpty();
 
     if (!isRunning || !isMessageEmpty) return null;
 
diff --git a/src/components/assistant-ui/file.tsx b/src/components/assistant-ui/file.tsx
index 92d93a1b..cf8db319 100644
--- a/src/components/assistant-ui/file.tsx
+++ b/src/components/assistant-ui/file.tsx
@@ -11,7 +11,7 @@ import {
   BracesIcon,
   DownloadIcon,
 } from "lucide-react";
-import type { FileMessagePartComponent } from "@assistant-ui/react";
+import type { ChatFilePartComponent } from "@/lib/chat/runtime";
 import { cn } from "@/lib/utils";
 
 const fileVariants = cva(
@@ -187,7 +187,7 @@ function FileDownload({
   );
 }
 
-const FileImpl: FileMessagePartComponent = ({ filename, data, mimeType }) => {
+const FileImpl: ChatFilePartComponent = ({ filename, data, mimeType }) => {
   const bytes = getBase64Size(data);
 
   return (
@@ -206,7 +206,7 @@ const FileImpl: FileMessagePartComponent = ({ filename, data, mimeType }) => {
   );
 };
 
-const File = memo(FileImpl) as unknown as FileMessagePartComponent & {
+const File = memo(FileImpl) as unknown as ChatFilePartComponent & {
   Root: typeof FileRoot;
   Icon: typeof FileIconDisplay;
   Name: typeof FileName;
diff --git a/src/components/assistant-ui/image.tsx b/src/components/assistant-ui/image.tsx
index 6f7b6ee6..da93c891 100644
--- a/src/components/assistant-ui/image.tsx
+++ b/src/components/assistant-ui/image.tsx
@@ -10,7 +10,7 @@ import {
 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 type { ChatImagePartComponent } from "@/lib/chat/runtime";
 import { cn } from "@/lib/utils";
 
 const imageVariants = cva(
@@ -230,7 +230,7 @@ function ImageZoom({ src, alt = "Image preview", children }: ImageZoomProps) {
   );
 }
 
-const ImageImpl: ImageMessagePartComponent = ({ image, filename }) => {
+const ImageImpl: ChatImagePartComponent = ({ image, filename }) => {
   return (
     <ImageRoot>
       <ImageZoom src={image} alt={filename || "Image content"}>
@@ -241,7 +241,7 @@ const ImageImpl: ImageMessagePartComponent = ({ image, filename }) => {
   );
 };
 
-const Image = memo(ImageImpl) as unknown as ImageMessagePartComponent & {
+const Image = memo(ImageImpl) as unknown as ChatImagePartComponent & {
   Root: typeof ImageRoot;
   Preview: typeof ImagePreview;
   Filename: typeof ImageFilename;
diff --git a/src/components/assistant-ui/markdown-text.tsx b/src/components/assistant-ui/markdown-text.tsx
index ba0518b9..f9f13ec6 100644
--- a/src/components/assistant-ui/markdown-text.tsx
+++ b/src/components/assistant-ui/markdown-text.tsx
@@ -5,7 +5,12 @@ import "streamdown/styles.css";
 import { createCodePlugin } from "@streamdown/code";
 import { mermaid } from "@streamdown/mermaid";
 import { createMathPlugin } from "@streamdown/math";
-import { useMessagePartText, useAuiState, type TextMessagePartProps } from "@assistant-ui/react";
+import {
+  useChatMessagePartText,
+  useCurrentMessageId,
+  useMainThreadId,
+  type ChatTextPartProps,
+} from "@/lib/chat/runtime";
 import {
   Children,
   isValidElement,
@@ -164,17 +169,17 @@ const CitationRenderer = memo(
 CitationRenderer.displayName = "CitationRenderer";
 
 /** Props from assistant-ui when used as Text component, or optional when used directly (e.g. in Reasoning) */
-type MarkdownTextProps = Partial<TextMessagePartProps> & {
+type MarkdownTextProps = Partial<ChatTextPartProps> & {
   /** Use "reasoning" for smoother streaming in reasoning blocks (blurIn, longer duration) */
   streamingVariant?: "default" | "reasoning";
 };
 
 const MarkdownTextImpl = (props: MarkdownTextProps) => {
   const streamingVariant = props.streamingVariant ?? "default";
-  const { text, status } = useMessagePartText();
+  const { text, status } = useChatMessagePartText();
 
-  const threadId = useAuiState(({ threads }) => (threads as any)?.mainThreadId);
-  const messageId = useAuiState(({ message }) => (message as any)?.id);
+  const threadId = useMainThreadId();
+  const messageId = useCurrentMessageId();
 
   const animateConfig =
     streamingVariant === "reasoning"
diff --git a/src/components/assistant-ui/reasoning.tsx b/src/components/assistant-ui/reasoning.tsx
index 052a58dc..9314e47b 100644
--- a/src/components/assistant-ui/reasoning.tsx
+++ b/src/components/assistant-ui/reasoning.tsx
@@ -4,11 +4,13 @@ import { memo, useCallback, useRef, useState, useEffect, useLayoutEffect, forwar
 import { cva, type VariantProps } from "class-variance-authority";
 import { ChevronDownIcon } from "lucide-react";
 import {
-  useScrollLock,
-  useAuiState,
-  type ReasoningMessagePartComponent,
-  type ReasoningGroupComponent,
-} from "@assistant-ui/react";
+  useChatScrollLock,
+  useIsLastMessage,
+  useIsMessagePartStreaming,
+  useMessagePartTextLengthSnapshot,
+  type ChatReasoningGroupComponent,
+  type ChatReasoningPartComponent,
+} from "@/lib/chat/runtime";
 import { MarkdownText } from "@/components/assistant-ui/markdown-text";
 import {
   Collapsible,
@@ -53,7 +55,7 @@ function ReasoningRoot({
 }: ReasoningRootProps) {
   const collapsibleRef = useRef<HTMLDivElement>(null);
   const [uncontrolledOpen, setUncontrolledOpen] = useState(defaultOpen);
-  const lockScroll = useScrollLock(collapsibleRef, ANIMATION_DURATION);
+  const lockScroll = useChatScrollLock(collapsibleRef, ANIMATION_DURATION);
 
   const isControlled = controlledOpen !== undefined;
   const isOpen = isControlled ? controlledOpen : uncontrolledOpen;
@@ -216,40 +218,22 @@ const ReasoningText = forwardRef<HTMLDivElement, React.ComponentProps<"div">>(
   }
 );
 
-const ReasoningImpl: ReasoningMessagePartComponent = () => (
+const ReasoningImpl: ChatReasoningPartComponent = () => (
   <MarkdownText streamingVariant="reasoning" />
 );
 
-const ReasoningGroupImpl: ReasoningGroupComponent = ({
+const ReasoningGroupImpl: ChatReasoningGroupComponent = ({
   children,
   startIndex,
   endIndex,
 }) => {
   const textContainerRef = useRef<HTMLDivElement>(null);
-  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;
-  });
+  const isReasoningStreaming = useIsMessagePartStreaming("reasoning", startIndex, endIndex);
 
-  const isLastMessage = useAuiState(({ thread, message }) => {
-    const messages = (thread as unknown as { messages?: Array<{ id?: string }> })?.messages ?? [];
-    const idx = messages.findIndex((m) => m.id === message.id);
-    return idx >= 0 && idx === messages.length - 1;
-  });
+  const isLastMessage = useIsLastMessage();
 
   // Subscribe to reasoning text length so we re-run scroll effect on each stream chunk
-  const reasoningTextSnapshot = useAuiState(({ message }) => {
-    let len = 0;
-    for (let i = startIndex; i <= endIndex && i < message.parts.length; i++) {
-      const p = message.parts[i] as { type?: string; text?: string } | undefined;
-      if (p?.type === "reasoning" && typeof p.text === "string") len += p.text.length;
-    }
-    return len;
-  });
+  const reasoningTextSnapshot = useMessagePartTextLengthSnapshot("reasoning", startIndex, endIndex);
 
   const [isManuallyOpen, setIsManuallyOpen] = useState(false);
   const isOpen = isReasoningStreaming || isManuallyOpen;
@@ -296,7 +280,7 @@ const ReasoningGroupImpl: ReasoningGroupComponent = ({
 
 const Reasoning = memo(
   ReasoningImpl,
-) as unknown as ReasoningMessagePartComponent & {
+) as unknown as ChatReasoningPartComponent & {
   Root: typeof ReasoningRoot;
   Trigger: typeof ReasoningTrigger;
   Content: typeof ReasoningContent;
diff --git a/src/components/assistant-ui/sources.tsx b/src/components/assistant-ui/sources.tsx
index 51273b67..09e5501d 100644
--- a/src/components/assistant-ui/sources.tsx
+++ b/src/components/assistant-ui/sources.tsx
@@ -3,7 +3,7 @@
 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 type { ChatSourcePartComponent } from "@/lib/chat/runtime";
 import { cn } from "@/lib/utils";
 
 const extractDomain = (url: string): string => {
@@ -116,7 +116,7 @@ function Source({
   );
 }
 
-const SourcesImpl: SourceMessagePartComponent = ({
+const SourcesImpl: ChatSourcePartComponent = ({
   url,
   title,
   sourceType,
@@ -134,7 +134,7 @@ const SourcesImpl: SourceMessagePartComponent = ({
   );
 };
 
-const Sources = memo(SourcesImpl) as unknown as SourceMessagePartComponent & {
+const Sources = memo(SourcesImpl) as unknown as ChatSourcePartComponent & {
   Root: typeof Source;
   Icon: typeof SourceIcon;
   Title: typeof SourceTitle;
diff --git a/src/components/assistant-ui/tool-fallback.tsx b/src/components/assistant-ui/tool-fallback.tsx
index aca40305..5da1bf67 100644
--- a/src/components/assistant-ui/tool-fallback.tsx
+++ b/src/components/assistant-ui/tool-fallback.tsx
@@ -1,9 +1,9 @@
-import type { ToolCallMessagePartComponent } from "@assistant-ui/react";
+import type { ChatToolCallPartComponent } from "@/lib/chat/runtime";
 import { CheckIcon, ChevronDownIcon, ChevronUpIcon } from "lucide-react";
 import { useState } from "react";
 import { Button } from "@/components/ui/button";
 
-export const ToolFallback: ToolCallMessagePartComponent = ({
+export const ToolFallback: ChatToolCallPartComponent = ({
   toolName,
   argsText,
   result,
diff --git a/src/components/assistant-ui/tool-group.tsx b/src/components/assistant-ui/tool-group.tsx
index dfa28481..8ab66bfb 100644
--- a/src/components/assistant-ui/tool-group.tsx
+++ b/src/components/assistant-ui/tool-group.tsx
@@ -11,7 +11,11 @@ import {
 } from "react";
 import { ChevronDownIcon, LoaderIcon } from "lucide-react";
 import { cva, type VariantProps } from "class-variance-authority";
-import { useAuiState, useScrollLock } from "@assistant-ui/react";
+import {
+  useChatScrollLock,
+  useIsLastMessage,
+  useIsMessagePartStreaming,
+} from "@/lib/chat/runtime";
 import {
   Collapsible,
   CollapsibleContent,
@@ -53,7 +57,7 @@ function ToolGroupRoot({
 }: ToolGroupRootProps) {
   const collapsibleRef = useRef<HTMLDivElement>(null);
   const [uncontrolledOpen, setUncontrolledOpen] = useState(defaultOpen);
-  const lockScroll = useScrollLock(collapsibleRef, ANIMATION_DURATION);
+  const lockScroll = useChatScrollLock(collapsibleRef, ANIMATION_DURATION);
 
   const isControlled = controlledOpen !== undefined;
   const isOpen = isControlled ? controlledOpen : uncontrolledOpen;
@@ -192,20 +196,9 @@ const ToolGroupImpl: FC<
 
   // Match `ReasoningGroup` behavior: mark active while the *current streaming part*
   // is a tool-call within this group's index range.
-  const isToolGroupStreaming = 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 !== "tool-call") return false;
-    return lastIndex >= startIndex && lastIndex <= endIndex;
-  });
+  const isToolGroupStreaming = useIsMessagePartStreaming("tool-call", startIndex, endIndex);
 
-  const isLastMessage = useAuiState(({ thread, message }) => {
-    const messages = (thread as unknown as { messages?: Array<{ id?: string }> })?.messages ?? [];
-    const idx = messages.findIndex((m) => m.id === message.id);
-    return idx >= 0 && idx === messages.length - 1;
-  });
+  const isLastMessage = useIsLastMessage();
 
   const [isManuallyOpen, setIsManuallyOpen] = useState(isLastMessage);
   const isOpen = isToolGroupStreaming || isManuallyOpen;
diff --git a/src/lib/chat/runtime/hooks.ts b/src/lib/chat/runtime/hooks.ts
index cf5e9652..41ded0d3 100644
--- a/src/lib/chat/runtime/hooks.ts
+++ b/src/lib/chat/runtime/hooks.ts
@@ -5,6 +5,7 @@ import {
   useAuiState,
   useMessage,
   useMessagePartText,
+  useScrollLock,
 } from "@assistant-ui/react";
 import type {
   ChatMessage,
@@ -68,3 +69,87 @@ export function usePromptInput(): ComposerActions | null {
     getState: () => composer.getState() as unknown as ComposerStateSnapshot | undefined,
   };
 }
+/**
+ * Lock the given element's scroll position for a fixed duration. Used when
+ * collapsing streaming content to prevent scroll jumps.
+ */
+export const useChatScrollLock = useScrollLock;
+
+/** Current assistant/user message id from MessageRuntime context. */
+export function useCurrentMessageId(): string | null {
+  return useAuiState((s: any) => (s.message as { id?: string } | undefined)?.id ?? null);
+}
+
+/** True while this message is actively streaming. */
+export function useIsMessageRunning(): boolean {
+  return useAuiState((s: any) => (s.message as { status?: { type?: string } } | undefined)?.status?.type === "running");
+}
+
+/**
+ * True when the current message has no content parts — used to show loading
+ * states before the first token arrives.
+ */
+export function useIsMessageEmpty(): boolean {
+  return useAuiState((s: any) => {
+    const msg = s.message as { content?: unknown[] } | undefined;
+    return !msg?.content || (Array.isArray(msg.content) && msg.content.length === 0);
+  });
+}
+
+/**
+ * True when the current message is the last message in the thread.
+ * Used to gate old reasoning/tool groups from rendering once newer messages exist.
+ */
+export function useIsLastMessage(): boolean {
+  return useAuiState((s: any) => {
+    const thread = s.thread as { messages?: Array<{ id?: string }> } | undefined;
+    const messageId = (s.message as { id?: string } | undefined)?.id;
+    const messages = thread?.messages ?? [];
+    const idx = messages.findIndex((m) => m.id === messageId);
+    return idx >= 0 && idx === messages.length - 1;
+  });
+}
+
+/**
+ * True when the current streaming part is of `partType` and falls within
+ * [startIndex, endIndex]. Shared by ReasoningGroup and ToolGroup to detect
+ * whether the group is actively streaming.
+ */
+export function useIsMessagePartStreaming(
+  partType: string,
+  startIndex: number,
+  endIndex: number,
+): boolean {
+  return useAuiState((s: any) => {
+    const msg = s.message as
+      | { status?: { type?: string }; parts?: Array<{ type?: string }> }
+      | undefined;
+    if (msg?.status?.type !== "running") return false;
+    const parts = msg.parts ?? [];
+    const lastIndex = parts.length - 1;
+    if (lastIndex < 0) return false;
+    const lastType = parts[lastIndex]?.type;
+    if (lastType !== partType) return false;
+    return lastIndex >= startIndex && lastIndex <= endIndex;
+  });
+}
+
+/**
+ * Sum the `.text` length of all parts of `partType` within [startIndex, endIndex].
+ * Used to force re-renders of scroll effects as streaming text grows.
+ */
+export function useMessagePartTextLengthSnapshot(
+  partType: string,
+  startIndex: number,
+  endIndex: number,
+): number {
+  return useAuiState((s: any) => {
+    const parts = (s.message as { parts?: Array<{ type?: string; text?: string }> } | undefined)?.parts ?? [];
+    let len = 0;
+    for (let i = startIndex; i <= endIndex && i < parts.length; i++) {
+      const p = parts[i];
+      if (p?.type === partType && typeof p.text === "string") len += p.text.length;
+    }
+    return len;
+  });
+}
diff --git a/src/lib/chat/runtime/index.ts b/src/lib/chat/runtime/index.ts
index 6e2e3156..e6e790d1 100644
--- a/src/lib/chat/runtime/index.ts
+++ b/src/lib/chat/runtime/index.ts
@@ -6,6 +6,13 @@ export type {
   ComposerActions,
   ComposerStateSnapshot,
   ThreadState,
+  ChatFilePartComponent,
+  ChatImagePartComponent,
+  ChatReasoningGroupComponent,
+  ChatReasoningPartComponent,
+  ChatSourcePartComponent,
+  ChatTextPartProps,
+  ChatToolCallPartComponent,
 } from "./types";
 export * from "./primitives";
 export * from "./hooks";
diff --git a/src/lib/chat/runtime/types.ts b/src/lib/chat/runtime/types.ts
index 1907ca84..e0df171e 100644
--- a/src/lib/chat/runtime/types.ts
+++ b/src/lib/chat/runtime/types.ts
@@ -43,3 +43,24 @@ export interface ComposerActions {
   setRunConfig(config: { custom?: Record<string, unknown> }): void;
   getState(): ComposerStateSnapshot | undefined;
 }
+import type {
+  TextMessagePartProps,
+  FileMessagePartComponent,
+  ImageMessagePartComponent,
+  SourceMessagePartComponent,
+  ToolCallMessagePartComponent,
+  ReasoningMessagePartComponent,
+  ReasoningGroupComponent,
+} from "@assistant-ui/react";
+
+/**
+ * Type re-exports — the ACL owns the identity of these types so consumers never
+ * import them directly from @assistant-ui/react.
+ */
+export type ChatTextPartProps = TextMessagePartProps;
+export type ChatFilePartComponent = FileMessagePartComponent;
+export type ChatImagePartComponent = ImageMessagePartComponent;
+export type ChatSourcePartComponent = SourceMessagePartComponent;
+export type ChatToolCallPartComponent = ToolCallMessagePartComponent;
+export type ChatReasoningPartComponent = ReasoningMessagePartComponent;
+export type ChatReasoningGroupComponent = ReasoningGroupComponent;

From bb4a10b27f200ddd5e9ed6e2324742bbbf40c275 Mon Sep 17 00:00:00 2001
From: urjitc <135136842+urjitc@users.noreply.github.com>
Date: Sun, 19 Apr 2026 22:05:46 +0000
Subject: [PATCH 3/7] Phase 2b: migrate attachment.tsx to chat-runtime ACL
 (adds useAttachmentSnapshot, useAttachmentScope, useAttachmentId,
 useIsAttachmentImage, ChatAttachment primitive)

---
 src/components/assistant-ui/attachment.tsx | 157 +++++++--------------
 src/lib/chat/runtime/hooks.ts              |  44 ++++++
 src/lib/chat/runtime/primitives.ts         |   2 +
 src/lib/chat/runtime/types.ts              |  38 +++--
 4 files changed, 128 insertions(+), 113 deletions(-)

diff --git a/src/components/assistant-ui/attachment.tsx b/src/components/assistant-ui/attachment.tsx
index d1031709..91df9317 100644
--- a/src/components/assistant-ui/attachment.tsx
+++ b/src/components/assistant-ui/attachment.tsx
@@ -10,23 +10,21 @@ import React, {
 import {
   XIcon,
   Link as LinkIcon,
-  SearchIcon,
-  Plus,
-  Code as CodeIcon,
-  GalleryHorizontalEnd,
   FileText,
   Loader2,
 } from "lucide-react";
 import { LuPaperclip } from "react-icons/lu";
 import { toast } from "sonner";
 import {
-  AttachmentPrimitive,
-  ComposerPrimitive,
-  MessagePrimitive,
-  useAui,
-} from "@assistant-ui/react";
-import { useAuiState } from "@assistant-ui/react";
-import { useShallow } from "zustand/shallow";
+  ChatAttachment,
+  ChatMessage,
+  ChatPromptInput,
+  useAttachmentId,
+  useAttachmentScope,
+  useAttachmentSnapshot,
+  useIsAttachmentImage,
+  usePromptInput,
+} from "@/lib/chat/runtime";
 import {
   Tooltip,
   TooltipContent,
@@ -42,10 +40,8 @@ import { Avatar, AvatarImage, AvatarFallback } from "@/components/ui/avatar";
 import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
 import { cn } from "@/lib/utils";
 import { useAttachmentUploadStore } from "@/lib/stores/attachment-upload-store";
-import { useUIStore } from "@/lib/stores/ui-store";
 import { emitPasswordProtectedPdf } from "@/components/modals/PasswordProtectedPdfDialog";
 import { filterPasswordProtectedPdfs } from "@/lib/uploads/pdf-validation";
-import { FaCheck } from "react-icons/fa";
 
 const useFileSrc = (file: File | undefined) => {
   const [src, setSrc] = useState<string | undefined>(undefined);
@@ -109,46 +105,29 @@ function isUrlAttachment(att: UrlLikeAttachment | undefined): boolean {
 }
 
 const useAttachmentSrc = () => {
-  const attachmentState = useAuiState(
-    useShallow(
-      ({
-        attachment,
-      }): { file?: File; src?: string; isUrl?: boolean; url?: string } => {
-        const att = attachment as
-          | {
-              type?: string;
-              name?: string;
-              file?: File & { name: string };
-              content?: Array<{ type: string; text?: string; image?: string }>;
-            }
-          | undefined;
-        if (!att)
-          return {
-            file: undefined,
-            src: undefined,
-            isUrl: false,
-            url: undefined,
-          };
-        if (isUrlAttachment(att)) {
-          const url = getHttpAttachmentUrl(att);
-          if (url) {
-            return { isUrl: true, src: getFaviconUrl(url), url };
-          }
-          return { isUrl: true };
-        }
+  const att = useAttachmentSnapshot();
 
-        if (att.type !== "image") return {};
-        if (att.file) return { file: att.file };
-        const imageContent = att.content?.find(
-          (c: { type: string }) => c.type === "image",
-        ) as { type: "image"; image: string } | undefined;
-        if (imageContent?.image) {
-          return { src: imageContent.image };
-        }
-        return {};
-      },
-    ),
-  );
+  const attachmentState: { file?: File; src?: string; isUrl?: boolean; url?: string } = (() => {
+    if (!att) {
+      return { file: undefined, src: undefined, isUrl: false, url: undefined };
+    }
+    if (isUrlAttachment(att)) {
+      const url = getHttpAttachmentUrl(att);
+      if (url) {
+        return { isUrl: true, src: getFaviconUrl(url), url };
+      }
+      return { isUrl: true };
+    }
+    if (att.type !== "image") return {};
+    if (att.file) return { file: att.file };
+    const imageContent = att.content?.find(
+      (c: { type: string }) => c.type === "image",
+    ) as { type: "image"; image: string } | undefined;
+    if (imageContent?.image) {
+      return { src: imageContent.image };
+    }
+    return {};
+  })();
 
   return {
     src: useFileSrc(attachmentState.file) ?? attachmentState.src,
@@ -205,9 +184,7 @@ const AttachmentPreviewDialog: FC<PropsWithChildren> = ({ children }) => {
 };
 
 const AttachmentThumb: FC = () => {
-  const isImage = useAuiState(
-    ({ attachment }) => (attachment as { type?: string })?.type === "image",
-  );
+  const isImage = useIsAttachmentImage();
   const attachmentSrc = useAttachmentSrc();
   const { src, isUrl } = attachmentSrc;
 
@@ -253,34 +230,21 @@ const AttachmentThumb: FC = () => {
 };
 
 const AttachmentUI: FC = () => {
-  const aui = useAui();
-  const isComposer = aui.attachment.source === "composer";
-  const attachmentId = useAuiState(
-    ({ attachment }) => (attachment as { id?: string })?.id,
-  );
+  const scope = useAttachmentScope();
+  const isComposer = scope === "composer";
+  const attachmentId = useAttachmentId();
   const isUploading = useAttachmentUploadStore(
     (s) => attachmentId != null && s.uploadingIds.has(attachmentId),
   );
+  const isImage = useIsAttachmentImage();
 
-  const isImage = useAuiState(
-    ({ attachment }) => (attachment as { type?: string })?.type === "image",
-  );
+  // Snapshot is memoized in the ACL, so deriving multiple values from it is cheap
+  const attSnapshot = useAttachmentSnapshot();
 
-  // Split into separate selectors to avoid creating new objects on each render
-  const typeLabel = useAuiState(({ attachment }) => {
-    const att = attachment as
-      | {
-          type?: string;
-          name?: string;
-          file?: { name: string };
-          content?: Array<{ type: string; text?: string }>;
-        }
-      | undefined;
+  const typeLabel = (() => {
+    const att = attSnapshot;
     if (!att) return "File";
-    if (isUrlAttachment(att)) {
-      return "URL";
-    }
-
+    if (isUrlAttachment(att)) return "URL";
     const type = att.type;
     switch (type) {
       case "image":
@@ -292,24 +256,11 @@ const AttachmentUI: FC = () => {
       default:
         return "File";
     }
-  });
-
-  const isUrl = useAuiState(({ attachment }) => {
-    const att = attachment as
-      | {
-          type?: string;
-          name?: string;
-          file?: { name: string };
-          content?: Array<{ type: string; text?: string }>;
-        }
-      | undefined;
-    if (!att) return false;
-    return isUrlAttachment(att);
-  });
+  })();
 
   return (
     <Tooltip>
-      <AttachmentPrimitive.Root
+      <ChatAttachment.Root
         className={cn(
           "aui-attachment-root relative flex flex-col items-center gap-1.5 max-w-[100px]",
           isImage &&
@@ -350,12 +301,12 @@ const AttachmentUI: FC = () => {
         </div>
         {!(isComposer && isUploading) && (
           <div className="text-[11px] text-muted-foreground w-full truncate text-center px-1 leading-tight">
-            <AttachmentPrimitive.Name />
+            <ChatAttachment.Name />
           </div>
         )}
-      </AttachmentPrimitive.Root>
+      </ChatAttachment.Root>
       <TooltipContent side="top">
-        <AttachmentPrimitive.Name />
+        <ChatAttachment.Name />
       </TooltipContent>
     </Tooltip>
   );
@@ -365,7 +316,7 @@ const AttachmentRemove: FC = () => {
   const { isUrl } = useAttachmentSrc();
 
   return (
-    <AttachmentPrimitive.Remove asChild>
+    <ChatAttachment.Remove asChild>
       <TooltipIconButton
         tooltip={isUrl ? "Remove URL" : "Remove file"}
         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"
@@ -373,14 +324,14 @@ const AttachmentRemove: FC = () => {
       >
         <XIcon className="aui-attachment-remove-icon size-3 dark:stroke-[2.5px]" />
       </TooltipIconButton>
-    </AttachmentPrimitive.Remove>
+    </ChatAttachment.Remove>
   );
 };
 
 export const UserMessageAttachments: FC = () => {
   return (
     <div className="aui-user-message-attachments-end col-span-full col-start-1 row-start-1 flex w-full flex-row justify-end gap-2">
-      <MessagePrimitive.Attachments components={{ Attachment: AttachmentUI }} />
+      <ChatMessage.Attachments components={{ Attachment: AttachmentUI }} />
     </div>
   );
 };
@@ -388,9 +339,7 @@ export const UserMessageAttachments: FC = () => {
 export const PromptInputAttachments: FC = () => {
   return (
     <div className="aui-composer-attachments mb-2 flex w-full flex-row items-center gap-2 overflow-x-auto pt-0.5 pb-1 empty:hidden">
-      <ComposerPrimitive.Attachments
-        components={{ Attachment: AttachmentUI }}
-      />
+      <ChatPromptInput.Attachments components={{ Attachment: AttachmentUI }} />
     </div>
   );
 };
@@ -398,7 +347,7 @@ export const PromptInputAttachments: FC = () => {
 export const PromptInputAddAttachment: FC = () => {
   const fileInputRef = useRef<HTMLInputElement>(null);
   const containerRef = useRef<HTMLDivElement>(null);
-  const aui = useAui();
+  const promptInput = usePromptInput();
 
   const handleFileChange = async (e: React.ChangeEvent<HTMLInputElement>) => {
     const files = e.target.files;
@@ -465,7 +414,7 @@ export const PromptInputAddAttachment: FC = () => {
     // Add valid files (non–Office-doc, non–password-protected) — others still upload successfully
     if (filesToAdd.length > 0) {
       filesToAdd.forEach((file) => {
-        aui.composer().addAttachment(file);
+        promptInput?.addAttachment(file);
       });
 
       if (filesToAdd.length < fileArray.length) {
@@ -480,7 +429,7 @@ export const PromptInputAddAttachment: FC = () => {
     }
   };
 
-  const uploadInputId = "prompt-input-file-upload";
+  const uploadInputId = "composer-file-upload";
 
   return (
     <>
diff --git a/src/lib/chat/runtime/hooks.ts b/src/lib/chat/runtime/hooks.ts
index 41ded0d3..9ca0bb3f 100644
--- a/src/lib/chat/runtime/hooks.ts
+++ b/src/lib/chat/runtime/hooks.ts
@@ -7,10 +7,13 @@ import {
   useMessagePartText,
   useScrollLock,
 } from "@assistant-ui/react";
+import { useShallow } from "zustand/shallow";
 import type {
+  ChatAttachmentSnapshot,
   ChatMessage,
   ComposerActions,
   ComposerStateSnapshot,
+  AttachmentScope,
   ThreadState,
 } from "./types";
 
@@ -153,3 +156,44 @@ export function useMessagePartTextLengthSnapshot(
     return len;
   });
 }
+
+/**
+ * Scope of the current attachment ("composer" for user uploads, "message" for rendered attachments in a chat message).
+ * Reads `aui.attachment.source` from the AttachmentPrimitive context.
+ */
+export function useAttachmentScope(): AttachmentScope {
+  const aui = useAui();
+  const source = (aui as unknown as { attachment?: { source?: string } } | null)?.attachment?.source;
+  return source === "composer" ? "composer" : "message";
+}
+
+/** Attachment id from AttachmentPrimitive context. */
+export function useAttachmentId(): string | undefined {
+  return useAuiState((s: any) => (s.attachment as { id?: string } | undefined)?.id);
+}
+
+/** Whether the current attachment is an image. */
+export function useIsAttachmentImage(): boolean {
+  return useAuiState((s: any) => (s.attachment as { type?: string } | undefined)?.type === "image");
+}
+
+/**
+ * Shallow-equal snapshot of the current attachment object. The returned object
+ * is stable across renders when the underlying assistant-ui attachment state
+ * hasn't changed.
+ */
+export function useAttachmentSnapshot(): ChatAttachmentSnapshot | undefined {
+  return useAuiState(
+    useShallow((s: any) => {
+      const att = s.attachment as ChatAttachmentSnapshot | undefined;
+      if (!att) return undefined;
+      return {
+        id: att.id,
+        type: att.type,
+        name: att.name,
+        file: att.file,
+        content: att.content,
+      };
+    }),
+  );
+}
diff --git a/src/lib/chat/runtime/primitives.ts b/src/lib/chat/runtime/primitives.ts
index 7f85ad6d..b7e312dc 100644
--- a/src/lib/chat/runtime/primitives.ts
+++ b/src/lib/chat/runtime/primitives.ts
@@ -1,5 +1,6 @@
 import {
   ActionBarPrimitive,
+  AttachmentPrimitive,
   AuiIf,
   BranchPickerPrimitive,
   ComposerPrimitive,
@@ -12,6 +13,7 @@ import type { BranchPickerPrimitive as _BPP } from "@assistant-ui/react";
 export const ChatThread = ThreadPrimitive;
 export const ChatMessage = MessagePrimitive;
 export const ChatPromptInput = ComposerPrimitive;
+export const ChatAttachment = AttachmentPrimitive;
 export const ChatActionBar = ActionBarPrimitive;
 export const ChatBranchPicker = BranchPickerPrimitive;
 export const ChatError = ErrorPrimitive;
diff --git a/src/lib/chat/runtime/types.ts b/src/lib/chat/runtime/types.ts
index e0df171e..2cd22608 100644
--- a/src/lib/chat/runtime/types.ts
+++ b/src/lib/chat/runtime/types.ts
@@ -1,3 +1,14 @@
+import type {
+  AttachmentPrimitive as _AP,
+  TextMessagePartProps,
+  FileMessagePartComponent,
+  ImageMessagePartComponent,
+  SourceMessagePartComponent,
+  ToolCallMessagePartComponent,
+  ReasoningMessagePartComponent,
+  ReasoningGroupComponent,
+} from '@assistant-ui/react';
+
 export type ChatMessageRole = "user" | "assistant" | "system";
 
 export type ChatTextPart = { type: "text"; text: string };
@@ -43,15 +54,6 @@ export interface ComposerActions {
   setRunConfig(config: { custom?: Record<string, unknown> }): void;
   getState(): ComposerStateSnapshot | undefined;
 }
-import type {
-  TextMessagePartProps,
-  FileMessagePartComponent,
-  ImageMessagePartComponent,
-  SourceMessagePartComponent,
-  ToolCallMessagePartComponent,
-  ReasoningMessagePartComponent,
-  ReasoningGroupComponent,
-} from "@assistant-ui/react";
 
 /**
  * Type re-exports — the ACL owns the identity of these types so consumers never
@@ -64,3 +66,21 @@ export type ChatSourcePartComponent = SourceMessagePartComponent;
 export type ChatToolCallPartComponent = ToolCallMessagePartComponent;
 export type ChatReasoningPartComponent = ReasoningMessagePartComponent;
 export type ChatReasoningGroupComponent = ReasoningGroupComponent;
+
+/** Scope an attachment belongs to — either the user composer or a rendered message. */
+export type AttachmentScope = "composer" | "message";
+
+/**
+ * Snapshot of the current attachment in AttachmentPrimitive context. Read by
+ * AttachmentThumb / AttachmentUI to render previews, labels, and upload states.
+ */
+export interface ChatAttachmentSnapshot {
+  id?: string;
+  type?: string;
+  name?: string;
+  file?: File & { name: string };
+  content?: Array<{ type: string; text?: string; image?: string }>;
+}
+
+/** Props for the neutral ChatAttachment primitive — mirrors AttachmentPrimitive.Root. */
+export type ChatAttachmentRootProps = _AP.Root.Props;

From 14478d712bb1defc62c107cd3aaad815b9318d19 Mon Sep 17 00:00:00 2001
From: urjitc <135136842+urjitc@users.noreply.github.com>
Date: Sun, 19 Apr 2026 22:18:51 +0000
Subject: [PATCH 4/7] Phase 2c: migrate 11 tool-UI files to chat-runtime ACL;
 fix Welcome Upload button id drift
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- AddYoutubeVideoToolUI, CreateDocumentToolUI, CreateFlashcardToolUI, CreateQuizToolUI,
  EditItemToolUI, ExecuteCodeToolUI, ReadWorkspaceToolUI, SearchWorkspaceToolUI,
  URLContextToolUI, WebSearchToolUI, YouTubeSearchToolUI now use ChatToolUIProps
  and useChatScrollLock from @/lib/chat/runtime instead of importing from
  @assistant-ui/react directly.
- Added ChatToolUIProps<Args, Result> type alias to the ACL.
- Fixed bug: Welcome-screen Upload suggestion clicked
  getElementById("prompt-input-file-upload") but attachment.tsx rendered
  id="composer-file-upload" — IDs now aligned.
- Dropped unused useAui import from YouTubeSearchToolUI.
---
 src/components/assistant-ui/AddYoutubeVideoToolUI.tsx | 4 ++--
 src/components/assistant-ui/CreateDocumentToolUI.tsx  | 4 ++--
 src/components/assistant-ui/CreateFlashcardToolUI.tsx | 4 ++--
 src/components/assistant-ui/CreateQuizToolUI.tsx      | 4 ++--
 src/components/assistant-ui/EditItemToolUI.tsx        | 4 ++--
 src/components/assistant-ui/ExecuteCodeToolUI.tsx     | 4 ++--
 src/components/assistant-ui/ReadWorkspaceToolUI.tsx   | 4 ++--
 src/components/assistant-ui/SearchWorkspaceToolUI.tsx | 4 ++--
 src/components/assistant-ui/URLContextToolUI.tsx      | 6 +++---
 src/components/assistant-ui/WebSearchToolUI.tsx       | 6 +++---
 src/components/assistant-ui/YouTubeSearchToolUI.tsx   | 6 +++---
 src/components/assistant-ui/attachment.tsx            | 2 +-
 src/lib/chat/runtime/index.ts                         | 1 +
 src/lib/chat/runtime/types.ts                         | 2 ++
 14 files changed, 29 insertions(+), 26 deletions(-)

diff --git a/src/components/assistant-ui/AddYoutubeVideoToolUI.tsx b/src/components/assistant-ui/AddYoutubeVideoToolUI.tsx
index 77702116..535d421c 100644
--- a/src/components/assistant-ui/AddYoutubeVideoToolUI.tsx
+++ b/src/components/assistant-ui/AddYoutubeVideoToolUI.tsx
@@ -1,7 +1,7 @@
 "use client";
 
 import type { ReactNode } from "react";
-import type { AssistantToolUIProps } from "@assistant-ui/react";
+import type { ChatToolUIProps } from "@/lib/chat/runtime";
 import { X, Eye, Play } from "lucide-react";
 import { Button } from "@/components/ui/button";
 import { cn } from "@/lib/utils";
@@ -107,7 +107,7 @@ const AddYoutubeVideoReceipt = ({
   );
 };
 
-export const renderAddYoutubeVideoToolUI: AssistantToolUIProps<
+export const renderAddYoutubeVideoToolUI: ChatToolUIProps<
   AddYoutubeVideoArgs,
   WorkspaceResult
 >["render"] = ({ args, result, status }) => {
diff --git a/src/components/assistant-ui/CreateDocumentToolUI.tsx b/src/components/assistant-ui/CreateDocumentToolUI.tsx
index 3203300b..dfae445c 100644
--- a/src/components/assistant-ui/CreateDocumentToolUI.tsx
+++ b/src/components/assistant-ui/CreateDocumentToolUI.tsx
@@ -2,7 +2,7 @@
 
 import { useState, useMemo } from "react";
 import { useWorkspaceState } from "@/hooks/workspace/use-workspace-state";
-import type { AssistantToolUIProps } from "@assistant-ui/react";
+import type { ChatToolUIProps } from "@/lib/chat/runtime";
 import { X, Eye, FolderInput, FileText } from "lucide-react";
 import { useWorkspaceStore } from "@/lib/stores/workspace-store";
 import { Button } from "@/components/ui/button";
@@ -179,7 +179,7 @@ const CreateDocumentReceipt = ({
   );
 };
 
-export const renderCreateDocumentToolUI: AssistantToolUIProps<
+export const renderCreateDocumentToolUI: ChatToolUIProps<
   CreateDocumentArgs,
   WorkspaceResult
 >["render"] = (props) => {
diff --git a/src/components/assistant-ui/CreateFlashcardToolUI.tsx b/src/components/assistant-ui/CreateFlashcardToolUI.tsx
index 52d46cc3..1992e187 100644
--- a/src/components/assistant-ui/CreateFlashcardToolUI.tsx
+++ b/src/components/assistant-ui/CreateFlashcardToolUI.tsx
@@ -2,7 +2,7 @@
 
 import { useEffect, useState, useMemo } from "react";
 import { useWorkspaceState } from "@/hooks/workspace/use-workspace-state";
-import type { AssistantToolUIProps } from "@assistant-ui/react";
+import type { ChatToolUIProps } from "@/lib/chat/runtime";
 import { X, Eye, FolderInput } from "lucide-react";
 import { PiCardsThreeBold } from "react-icons/pi";
 import { logger } from "@/lib/utils/logger";
@@ -219,7 +219,7 @@ const CreateFlashcardReceipt = ({
   );
 };
 
-export const renderCreateFlashcardToolUI: AssistantToolUIProps<
+export const renderCreateFlashcardToolUI: ChatToolUIProps<
   CreateFlashcardArgs,
   FlashcardResult
 >["render"] = (props) => {
diff --git a/src/components/assistant-ui/CreateQuizToolUI.tsx b/src/components/assistant-ui/CreateQuizToolUI.tsx
index f32bd049..bfb10b22 100644
--- a/src/components/assistant-ui/CreateQuizToolUI.tsx
+++ b/src/components/assistant-ui/CreateQuizToolUI.tsx
@@ -3,7 +3,7 @@
 import type { MouseEvent, ReactNode } from "react";
 import { useEffect, useState, useMemo } from "react";
 import { useWorkspaceState } from "@/hooks/workspace/use-workspace-state";
-import type { AssistantToolUIProps } from "@assistant-ui/react";
+import type { ChatToolUIProps } from "@/lib/chat/runtime";
 import { X, Eye, FolderInput, Brain } from "lucide-react";
 import { logger } from "@/lib/utils/logger";
 import { useWorkspaceStore } from "@/lib/stores/workspace-store";
@@ -198,7 +198,7 @@ const CreateQuizReceipt = ({
   );
 };
 
-export const renderCreateQuizToolUI: AssistantToolUIProps<
+export const renderCreateQuizToolUI: ChatToolUIProps<
   CreateQuizInput,
   QuizResult
 >["render"] = (props) => {
diff --git a/src/components/assistant-ui/EditItemToolUI.tsx b/src/components/assistant-ui/EditItemToolUI.tsx
index 938986ef..5310e6d8 100644
--- a/src/components/assistant-ui/EditItemToolUI.tsx
+++ b/src/components/assistant-ui/EditItemToolUI.tsx
@@ -3,7 +3,7 @@
 import type { ReactNode } from "react";
 import { useMemo } from "react";
 import { useWorkspaceState } from "@/hooks/workspace/use-workspace-state";
-import type { AssistantToolUIProps } from "@assistant-ui/react";
+import type { ChatToolUIProps } from "@/lib/chat/runtime";
 import { X, Eye } from "lucide-react";
 import { Pencil } from "lucide-react";
 import { useWorkspaceStore } from "@/lib/stores/workspace-store";
@@ -152,7 +152,7 @@ const EditItemReceipt = ({ args, result, status }: EditItemReceiptProps) => {
   );
 };
 
-export const renderEditItemToolUI: AssistantToolUIProps<
+export const renderEditItemToolUI: ChatToolUIProps<
   EditItemArgs,
   WorkspaceResult
 >["render"] = ({ args, result, status }) => {
diff --git a/src/components/assistant-ui/ExecuteCodeToolUI.tsx b/src/components/assistant-ui/ExecuteCodeToolUI.tsx
index 3004682d..a2f974e1 100644
--- a/src/components/assistant-ui/ExecuteCodeToolUI.tsx
+++ b/src/components/assistant-ui/ExecuteCodeToolUI.tsx
@@ -7,7 +7,7 @@ import {
   ChevronDownIcon,
   ChevronRightIcon,
 } from "lucide-react";
-import type { AssistantToolUIProps } from "@assistant-ui/react";
+import type { ChatToolUIProps } from "@/lib/chat/runtime";
 
 import { ToolUIErrorBoundary } from "@/components/tool-ui/shared";
 import { ToolUILoadingShell } from "@/components/assistant-ui/tool-ui-loading-shell";
@@ -31,7 +31,7 @@ function chartLabel(chartType: string, index: number): string {
   return `Chart ${n}`;
 }
 
-export const renderExecuteCodeToolUI: AssistantToolUIProps<
+export const renderExecuteCodeToolUI: ChatToolUIProps<
   { code: string },
   CodeExecuteResult
 >["render"] = ({ status, args, result }) => {
diff --git a/src/components/assistant-ui/ReadWorkspaceToolUI.tsx b/src/components/assistant-ui/ReadWorkspaceToolUI.tsx
index 7a30f969..95f21d22 100644
--- a/src/components/assistant-ui/ReadWorkspaceToolUI.tsx
+++ b/src/components/assistant-ui/ReadWorkspaceToolUI.tsx
@@ -1,7 +1,7 @@
 "use client";
 
 import { Eye } from "lucide-react";
-import type { AssistantToolUIProps } from "@assistant-ui/react";
+import type { ChatToolUIProps } from "@/lib/chat/runtime";
 import { ToolUIErrorBoundary } from "@/components/tool-ui/shared";
 import { ToolUILoadingShell } from "@/components/assistant-ui/tool-ui-loading-shell";
 import { ToolUIErrorShell } from "@/components/assistant-ui/tool-ui-error-shell";
@@ -22,7 +22,7 @@ function stripExtension(s: string): string {
     return s.replace(/\.[^.]+$/, "");
 }
 
-export const renderReadWorkspaceToolUI: AssistantToolUIProps<
+export const renderReadWorkspaceToolUI: ChatToolUIProps<
   ReadArgs,
   ReadResult
 >["render"] = ({ args, status, result }) => {
diff --git a/src/components/assistant-ui/SearchWorkspaceToolUI.tsx b/src/components/assistant-ui/SearchWorkspaceToolUI.tsx
index e7f67383..aafd57eb 100644
--- a/src/components/assistant-ui/SearchWorkspaceToolUI.tsx
+++ b/src/components/assistant-ui/SearchWorkspaceToolUI.tsx
@@ -1,7 +1,7 @@
 "use client";
 
 import { Search } from "lucide-react";
-import type { AssistantToolUIProps } from "@assistant-ui/react";
+import type { ChatToolUIProps } from "@/lib/chat/runtime";
 import { ToolUIErrorBoundary } from "@/components/tool-ui/shared";
 import { ToolUILoadingShell } from "@/components/assistant-ui/tool-ui-loading-shell";
 import { ToolUIErrorShell } from "@/components/assistant-ui/tool-ui-error-shell";
@@ -9,7 +9,7 @@ import { ToolUIErrorShell } from "@/components/assistant-ui/tool-ui-error-shell"
 type GrepArgs = { pattern: string; include?: string; path?: string };
 type GrepResult = { success: boolean; matches?: number; output?: string; message?: string };
 
-export const renderSearchWorkspaceToolUI: AssistantToolUIProps<
+export const renderSearchWorkspaceToolUI: ChatToolUIProps<
   GrepArgs,
   GrepResult
 >["render"] = ({ status, result }) => {
diff --git a/src/components/assistant-ui/URLContextToolUI.tsx b/src/components/assistant-ui/URLContextToolUI.tsx
index 3a32a790..6023596f 100644
--- a/src/components/assistant-ui/URLContextToolUI.tsx
+++ b/src/components/assistant-ui/URLContextToolUI.tsx
@@ -9,7 +9,7 @@ import {
   type PropsWithChildren,
 } from "react";
 
-import { useScrollLock, type AssistantToolUIProps } from "@assistant-ui/react";
+import { useChatScrollLock, type ChatToolUIProps } from "@/lib/chat/runtime";
 
 import { ToolUIErrorBoundary } from "@/components/tool-ui/shared";
 import { parseURLContextResult } from "@/lib/ai/tool-result-schemas";
@@ -36,7 +36,7 @@ const ToolRoot: FC<
 > = ({ className, children }) => {
   const collapsibleRef = useRef<HTMLDivElement>(null);
   const [isOpen, setIsOpen] = useState(false);
-  const lockScroll = useScrollLock(collapsibleRef, ANIMATION_DURATION);
+  const lockScroll = useChatScrollLock(collapsibleRef, ANIMATION_DURATION);
 
   const handleOpenChange = useCallback(
     (open: boolean) => {
@@ -210,7 +210,7 @@ type ProcessUrlsResult =
     };
   };
 
-export const renderURLContextToolUI: AssistantToolUIProps<{
+export const renderURLContextToolUI: ChatToolUIProps<{
   urls?: string[];
   jsonInput?: string;
 }, ProcessUrlsResult>["render"] = ({ args, status, result }) => {
diff --git a/src/components/assistant-ui/WebSearchToolUI.tsx b/src/components/assistant-ui/WebSearchToolUI.tsx
index cc73358a..576b31c2 100644
--- a/src/components/assistant-ui/WebSearchToolUI.tsx
+++ b/src/components/assistant-ui/WebSearchToolUI.tsx
@@ -9,7 +9,7 @@ import {
     type PropsWithChildren,
 } from "react";
 
-import { useScrollLock, type AssistantToolUIProps } from "@assistant-ui/react";
+import { useChatScrollLock, type ChatToolUIProps } from "@/lib/chat/runtime";
 
 import { ToolUIErrorBoundary } from "@/components/tool-ui/shared";
 import { parseWebSearchResult } from "@/lib/ai/tool-result-schemas";
@@ -35,7 +35,7 @@ const ToolRoot: FC<
 > = ({ className, children }) => {
     const collapsibleRef = useRef<HTMLDivElement>(null);
     const [isOpen, setIsOpen] = useState(false);
-    const lockScroll = useScrollLock(collapsibleRef, ANIMATION_DURATION);
+    const lockScroll = useChatScrollLock(collapsibleRef, ANIMATION_DURATION);
 
     const handleOpenChange = useCallback(
         (open: boolean) => {
@@ -288,7 +288,7 @@ WebSearchContent.displayName = "WebSearchContent";
  * Tool UI component for web_search tool.
  * Displays search query and results in a collapsible format similar to Reasoning.
  */
-export const renderWebSearchToolUI: AssistantToolUIProps<
+export const renderWebSearchToolUI: ChatToolUIProps<
   { query: string },
   WebSearchResult
 >["render"] = ({
diff --git a/src/components/assistant-ui/YouTubeSearchToolUI.tsx b/src/components/assistant-ui/YouTubeSearchToolUI.tsx
index e8c4254f..d8688337 100644
--- a/src/components/assistant-ui/YouTubeSearchToolUI.tsx
+++ b/src/components/assistant-ui/YouTubeSearchToolUI.tsx
@@ -1,6 +1,6 @@
 "use client";
 
-import { useAui, useScrollLock, type AssistantToolUIProps } from "@assistant-ui/react";
+import { useChatScrollLock, type ChatToolUIProps } from "@/lib/chat/runtime";
 import { Loader2, Plus, Check, ChevronDownIcon } from "lucide-react";
 import { YouTubeMark } from "@/components/icons/YouTubeMark";
 import { Button } from "@/components/ui/button";
@@ -52,7 +52,7 @@ const ToolRoot: FC<
 > = ({ className, children, defaultOpen = false }) => {
     const collapsibleRef = useRef<HTMLDivElement>(null);
     const [isOpen, setIsOpen] = useState(defaultOpen);
-    const lockScroll = useScrollLock(collapsibleRef, ANIMATION_DURATION);
+    const lockScroll = useChatScrollLock(collapsibleRef, ANIMATION_DURATION);
 
     const handleOpenChange = useCallback(
         (open: boolean) => {
@@ -380,7 +380,7 @@ const YouTubeSearchContent: FC<{
     );
 };
 
-export const renderYouTubeSearchToolUI: AssistantToolUIProps<
+export const renderYouTubeSearchToolUI: ChatToolUIProps<
   SearchYoutubeArgs,
   SearchYoutubeResult
 >["render"] = ({ args, status, result }) => {
diff --git a/src/components/assistant-ui/attachment.tsx b/src/components/assistant-ui/attachment.tsx
index 91df9317..cbe02bc1 100644
--- a/src/components/assistant-ui/attachment.tsx
+++ b/src/components/assistant-ui/attachment.tsx
@@ -429,7 +429,7 @@ export const PromptInputAddAttachment: FC = () => {
     }
   };
 
-  const uploadInputId = "composer-file-upload";
+  const uploadInputId = "prompt-input-file-upload";
 
   return (
     <>
diff --git a/src/lib/chat/runtime/index.ts b/src/lib/chat/runtime/index.ts
index e6e790d1..2d93f77f 100644
--- a/src/lib/chat/runtime/index.ts
+++ b/src/lib/chat/runtime/index.ts
@@ -13,6 +13,7 @@ export type {
   ChatSourcePartComponent,
   ChatTextPartProps,
   ChatToolCallPartComponent,
+  ChatToolUIProps,
 } from "./types";
 export * from "./primitives";
 export * from "./hooks";
diff --git a/src/lib/chat/runtime/types.ts b/src/lib/chat/runtime/types.ts
index 2cd22608..72953d40 100644
--- a/src/lib/chat/runtime/types.ts
+++ b/src/lib/chat/runtime/types.ts
@@ -7,6 +7,7 @@ import type {
   ToolCallMessagePartComponent,
   ReasoningMessagePartComponent,
   ReasoningGroupComponent,
+  AssistantToolUIProps,
 } from '@assistant-ui/react';
 
 export type ChatMessageRole = "user" | "assistant" | "system";
@@ -66,6 +67,7 @@ export type ChatSourcePartComponent = SourceMessagePartComponent;
 export type ChatToolCallPartComponent = ToolCallMessagePartComponent;
 export type ChatReasoningPartComponent = ReasoningMessagePartComponent;
 export type ChatReasoningGroupComponent = ReasoningGroupComponent;
+export type ChatToolUIProps<Args = unknown, Result = unknown> = AssistantToolUIProps<Args, Result>;
 
 /** Scope an attachment belongs to — either the user composer or a rendered message. */
 export type AttachmentScope = "composer" | "message";

From 1fb50a59503fa163affcbead8f27c37cf83bfe29 Mon Sep 17 00:00:00 2001
From: urjitc <135136842+urjitc@users.noreply.github.com>
Date: Sun, 19 Apr 2026 22:35:01 +0000
Subject: [PATCH 5/7] Phase 2d: migrate final 10 consumer files to chat-runtime
 ACL

- AssistantDropzone, PromptBuilderDialog, SpeechToTextButton,
  AssistantTextSelectionManager, thread-list-dropdown, AppChatHeader,
  MessageContextBadges, PdfPanelHeader, QuizContent, and
  use-workspace-context-provider now import from @/lib/chat/runtime
  instead of @assistant-ui/react.
- Added to the ACL:
  - Primitives: ChatThreadList, ChatThreadListItem
  - Hooks: useCurrentChatMessage, useThreadListItemId,
    useChatThreadListItem, usePromptInputThreadActions,
    useChatAssistantContext
  - Types: ChatThreadListItem, CurrentChatMessage,
    ChatAssistantContextOptions, PromptInputThreadActions

After this, only the runtime boundary files still import directly from
@assistant-ui/react: WorkspaceRuntimeProvider, chat-toolkit,
custom-thread-history-adapter, custom-thread-list-adapter,
supabase-attachment-adapter, and toCreateMessageWithContext.
---
 .../assistant-ui/AssistantDropzone.tsx        | 12 ++---
 .../AssistantTextSelectionManager.tsx         | 12 ++---
 .../assistant-ui/PromptBuilderDialog.tsx      |  8 +--
 .../assistant-ui/SpeechToTextButton.tsx       | 10 ++--
 .../assistant-ui/thread-list-dropdown.tsx     | 51 +++++++++----------
 src/components/chat/AppChatHeader.tsx         |  9 ++--
 src/components/chat/MessageContextBadges.tsx  |  4 +-
 src/components/pdf/PdfPanelHeader.tsx         | 10 ++--
 .../workspace-canvas/QuizContent.tsx          | 10 ++--
 .../ai/use-workspace-context-provider.ts      |  4 +-
 src/lib/chat/runtime/hooks.ts                 | 49 ++++++++++++++++++
 src/lib/chat/runtime/primitives.ts            |  4 ++
 src/lib/chat/runtime/types.ts                 | 31 +++++++++++
 13 files changed, 146 insertions(+), 68 deletions(-)

diff --git a/src/components/assistant-ui/AssistantDropzone.tsx b/src/components/assistant-ui/AssistantDropzone.tsx
index 30f7d129..21c1accb 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 { useAui } from "@assistant-ui/react";
+import { usePromptInput } from "@/lib/chat/runtime";
 import { useWorkspaceStore } from "@/lib/stores/workspace-store";
 import { Upload } from "lucide-react";
 import { useCallback, useState, useRef } from "react";
@@ -19,7 +19,7 @@ interface AssistantDropzoneProps {
  * Accepts all supported file types and adds them as attachments to the chat composer.
  */
 export function AssistantDropzone({ children }: AssistantDropzoneProps) {
-  const aui = useAui();
+  const promptInput = usePromptInput();
   const currentWorkspaceId = useWorkspaceStore((state) => state.currentWorkspaceId);
   const [isDragging, setIsDragging] = useState(false);
 
@@ -34,7 +34,7 @@ export function AssistantDropzone({ children }: AssistantDropzoneProps) {
 
   const onDrop = useCallback(
     async (acceptedFiles: File[]) => {
-      if (!currentWorkspaceId || !aui) return;
+      if (!currentWorkspaceId || !promptInput) return;
 
       const MAX_FILES = 10;
       const MAX_FILE_SIZE_MB = 50;
@@ -113,7 +113,7 @@ export function AssistantDropzone({ children }: AssistantDropzoneProps) {
         // Add each file to the composer
         const addPromises = filesToAdd.map(async (file) => {
           try {
-            await aui.composer().addAttachment(file);
+            await promptInput.addAttachment(file);
           } catch (error) {
             console.error("Failed to add attachment:", error);
             // Remove from processing set on error so it can be retried
@@ -142,7 +142,7 @@ export function AssistantDropzone({ children }: AssistantDropzoneProps) {
         }, 200);
       }
     },
-    [aui, currentWorkspaceId]
+    [promptInput, currentWorkspaceId]
   );
 
   // Clear processing state when drag ends (user drags away or cancels)
@@ -157,7 +157,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 || !aui, // Disable if no workspace is selected or api is not available
+    disabled: !currentWorkspaceId || !promptInput, // Disable if no workspace is selected or api is not available
     accept: {
       'image/*': ['.png', '.jpg', '.jpeg', '.gif', '.webp', '.svg', '.bmp', '.heic', '.heif', '.avif', '.tiff', '.tif'],
       'application/pdf': ['.pdf'],
diff --git a/src/components/assistant-ui/AssistantTextSelectionManager.tsx b/src/components/assistant-ui/AssistantTextSelectionManager.tsx
index 905e6d6e..e9064b6a 100644
--- a/src/components/assistant-ui/AssistantTextSelectionManager.tsx
+++ b/src/components/assistant-ui/AssistantTextSelectionManager.tsx
@@ -8,9 +8,9 @@ import {
   AssistantThreadSelection,
   type SelectionInfo,
 } from "@/components/assistant-ui/assistant-thread-selection";
+import { useMainThreadId, useThreadListItemId } from "@/lib/chat/runtime";
 import { useUIStore } from "@/lib/stores/ui-store";
 import { useWorkspaceStore } from "@/lib/stores/workspace-store";
-import { useAuiState } from "@assistant-ui/react";
 
 export default function AssistantTextSelectionManager({
   className,
@@ -28,13 +28,9 @@ export default function AssistantTextSelectionManager({
 
   const workspaceId = useWorkspaceStore((state) => state.currentWorkspaceId);
 
-  const threadListItemId = useAuiState(
-    ({ threadListItem }) => (threadListItem as any)?.id,
-  );
-  const mainThreadId = useAuiState(
-    ({ threads }) => (threads as any)?.mainThreadId,
-  );
-  const currentThreadId = threadListItemId || mainThreadId;
+  const threadListItemId = useThreadListItemId();
+  const mainThreadId = useMainThreadId();
+  const currentThreadId = (threadListItemId || mainThreadId) ?? undefined;
 
   const prevWorkspaceIdRef = useRef<string | null>(workspaceId);
   useEffect(() => {
diff --git a/src/components/assistant-ui/PromptBuilderDialog.tsx b/src/components/assistant-ui/PromptBuilderDialog.tsx
index aac90229..386c1523 100644
--- a/src/components/assistant-ui/PromptBuilderDialog.tsx
+++ b/src/components/assistant-ui/PromptBuilderDialog.tsx
@@ -38,7 +38,7 @@ import { Checkbox } from "@/components/ui/checkbox";
 import { useUIStore } from "@/lib/stores/ui-store";
 import { useSelectedCardIds } from "@/hooks/ui/use-selected-card-ids";
 import type { Item } from "@/lib/workspace-state/types";
-import { useAui } from "@assistant-ui/react";
+import { usePromptInput } from "@/lib/chat/runtime";
 import { focusComposerInput } from "@/lib/utils/composer-utils";
 import {
   Brain,
@@ -234,7 +234,7 @@ export function PromptBuilderDialog({
 }: PromptBuilderDialogProps) {
   const config = ACTION_CONFIG[action];
   const Icon = config.icon;
-  const aui = useAui();
+  const promptInput = usePromptInput();
   const formId = useId();
 
   const { selectedCardIds } = useSelectedCardIds();
@@ -399,7 +399,7 @@ export function PromptBuilderDialog({
       if (action !== "search" && selectedContextIds.size > 0) {
         selectMultipleCards(Array.from(selectedContextIds));
       }
-      aui?.composer().setText(builtPrompt);
+      promptInput?.setText(builtPrompt);
       focusComposerInput();
     }
     onOpenChange(false);
@@ -409,7 +409,7 @@ export function PromptBuilderDialog({
     hasValidTopic,
     onBeforeSubmit,
     onBuild,
-    aui,
+    promptInput,
     onOpenChange,
     selectedContextIds,
     selectMultipleCards,
diff --git a/src/components/assistant-ui/SpeechToTextButton.tsx b/src/components/assistant-ui/SpeechToTextButton.tsx
index 3cbc3aa6..fa8e4a94 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 { useAui } from "@assistant-ui/react";
+import { usePromptInput } from "@/lib/chat/runtime";
 import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
 import { cn } from "@/lib/utils";
 
 export const SpeechToTextButton: FC = () => {
-    const aui = useAui();
+    const promptInput = usePromptInput();
     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(' ') ? ' ' : '';
-            aui.composer().setText(originalText + separator + transcript);
+            promptInput?.setText(originalText + separator + transcript);
         }
-    }, [transcript, listening, originalText, aui]);
+    }, [transcript, listening, originalText, promptInput]);
 
     const handleStartListening = () => {
         // Reset transcript before starting a new session
         resetTranscript();
 
         // Capture current text before starting
-        const currentText = aui.composer().getState().text;
+        const currentText = promptInput?.getState()?.text ?? "";
         setOriginalText(currentText);
 
         SpeechRecognition.startListening({ continuous: true });
diff --git a/src/components/assistant-ui/thread-list-dropdown.tsx b/src/components/assistant-ui/thread-list-dropdown.tsx
index 4db92c8d..33d0c10f 100644
--- a/src/components/assistant-ui/thread-list-dropdown.tsx
+++ b/src/components/assistant-ui/thread-list-dropdown.tsx
@@ -3,13 +3,13 @@
 import type { FC } from "react";
 import { useState, useRef, useEffect } from "react";
 import {
-  AuiIf,
-  ThreadListItemPrimitive,
-  ThreadListPrimitive,
-  useAui,
-} from "@assistant-ui/react";
+  ChatIf,
+  ChatThreadList,
+  ChatThreadListItem,
+  useChatThreadListItem,
+  usePromptInputThreadActions,
+} from "@/lib/chat/runtime";
 import { Trash2Icon, PencilIcon } from "lucide-react";
-import { useThreadListItem } from "@assistant-ui/react";
 import { PiNotePencilBold } from "react-icons/pi";
 import { toast } from "sonner";
 
@@ -55,18 +55,18 @@ export const ThreadListDropdown: FC<ThreadListDropdownProps> = ({ trigger }) =>
         align="end"
         className="w-80 bg-sidebar border-sidebar-border max-h-[500px] p-0 overflow-hidden"
       >
-        <ThreadListPrimitive.Root className="aui-root aui-thread-list-root flex flex-col items-stretch">
+        <ChatThreadList.Root className="aui-root aui-thread-list-root flex flex-col items-stretch">
           <ThreadListNew onSelect={() => setOpen(false)} />
           <DropdownMenuSeparator className="bg-sidebar-border m-0" />
           <div className="flex-1 overflow-y-auto max-h-[400px] p-1">
-            <AuiIf condition={({ threads }) => threads.isLoading}>
+            <ChatIf condition={({ threads }) => threads.isLoading}>
               <ThreadListSkeleton />
-            </AuiIf>
-            <AuiIf condition={({ threads }) => !threads.isLoading}>
-              <ThreadListPrimitive.Items components={{ ThreadListItem: () => <ThreadListItem onSelect={() => setOpen(false)} /> }} />
-            </AuiIf>
+            </ChatIf>
+            <ChatIf condition={({ threads }) => !threads.isLoading}>
+              <ChatThreadList.Items components={{ ThreadListItem: () => <ThreadListItem onSelect={() => setOpen(false)} /> }} />
+            </ChatIf>
           </div>
-        </ThreadListPrimitive.Root>
+        </ChatThreadList.Root>
       </DropdownMenuContent>
     </DropdownMenu>
   );
@@ -74,7 +74,7 @@ export const ThreadListDropdown: FC<ThreadListDropdownProps> = ({ trigger }) =>
 
 const ThreadListNew: FC<{ onSelect?: () => void }> = ({ onSelect }) => {
   return (
-    <ThreadListPrimitive.New asChild>
+    <ChatThreadList.New asChild>
       <Button
         className="aui-thread-list-new flex items-center justify-start gap-2 rounded-none px-4 py-3 text-start hover:bg-muted/50 transition-all duration-200 font-medium w-full"
         variant="ghost"
@@ -83,7 +83,7 @@ const ThreadListNew: FC<{ onSelect?: () => void }> = ({ onSelect }) => {
         <PiNotePencilBold className="h-4 w-4 text-primary" />
         New Chat
       </Button>
-    </ThreadListPrimitive.New>
+    </ChatThreadList.New>
   );
 };
 
@@ -106,9 +106,9 @@ const ThreadListSkeleton: FC = () => {
 
 const ThreadListItem: FC<{ onSelect?: () => void }> = ({ onSelect }) => {
   return (
-    <ThreadListItemPrimitive.Root className="aui-thread-list-item group flex items-center gap-2 rounded-lg transition-all hover:bg-muted focus-visible:bg-muted focus-visible:ring-2 focus-visible:ring-ring focus-visible:outline-none data-active:bg-muted">
+    <ChatThreadListItem.Root className="aui-thread-list-item group flex items-center gap-2 rounded-lg transition-all hover:bg-muted focus-visible:bg-muted focus-visible:ring-2 focus-visible:ring-ring focus-visible:outline-none data-active:bg-muted">
       <ThreadListItemContent onSelect={onSelect} />
-    </ThreadListItemPrimitive.Root>
+    </ChatThreadListItem.Root>
   );
 };
 
@@ -116,11 +116,11 @@ const ThreadListItemContent: FC<{ onSelect?: () => void }> = ({ onSelect }) => {
   const [isEditing, setIsEditing] = useState(false);
   const [editValue, setEditValue] = useState("");
   const inputRef = useRef<HTMLInputElement>(null);
-  const aui = useAui();
+  const threadActions = usePromptInputThreadActions();
 
   // 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)
-  const threadListItem = useThreadListItem();
+  const threadListItem = useChatThreadListItem();
   const title = threadListItem?.title || "New Chat";
   const isThreadInitialized = !!threadListItem?.remoteId;
 
@@ -162,7 +162,7 @@ const ThreadListItemContent: FC<{ onSelect?: () => void }> = ({ onSelect }) => {
 
     if (trimmedValue && trimmedValue !== title) {
       try {
-        await aui?.threadListItem().rename(trimmedValue);
+        await threadActions?.rename(trimmedValue);
         toast.success("Title updated");
       } catch (error) {
         console.error("Failed to rename thread:", error);
@@ -208,12 +208,12 @@ const ThreadListItemContent: FC<{ onSelect?: () => void }> = ({ onSelect }) => {
           />
         ) : (
           <>
-            <ThreadListItemPrimitive.Trigger
+            <ChatThreadListItem.Trigger
               className="aui-thread-list-item-trigger flex-1 text-start cursor-pointer min-w-0"
               onClick={onSelect}
             >
               <ThreadListItemTitle onStartEdit={handleStartEdit} />
-            </ThreadListItemPrimitive.Trigger>
+            </ChatThreadListItem.Trigger>
             <Tooltip>
               <TooltipTrigger asChild>
                 <button
@@ -257,7 +257,7 @@ const ThreadListItemTitle: FC<{
       onDoubleClick={handleDoubleClick}
     >
       <span className="aui-thread-list-item-title text-sm break-words block">
-        <ThreadListItemPrimitive.Title fallback="New Chat" />
+        <ChatThreadListItem.Title fallback="New Chat" />
       </span>
     </div>
   );
@@ -284,14 +284,14 @@ const ThreadListItemArchive: FC = () => {
 
   return (
     <>
-      <ThreadListItemPrimitive.Archive asChild>
+      <ChatThreadListItem.Archive asChild>
         <button
           ref={archiveButtonRef}
           className="hidden"
           aria-hidden="true"
           tabIndex={-1}
         />
-      </ThreadListItemPrimitive.Archive>
+      </ChatThreadListItem.Archive>
       <TooltipIconButton
         onClick={handleDeleteClick}
         className="aui-thread-list-item-archive mr-3 ml-auto size-4 p-0 text-foreground hover:text-destructive opacity-0 group-hover:opacity-100 transition-opacity"
@@ -323,4 +323,3 @@ const ThreadListItemArchive: FC = () => {
     </>
   );
 };
-
diff --git a/src/components/chat/AppChatHeader.tsx b/src/components/chat/AppChatHeader.tsx
index 3b173701..08142149 100644
--- a/src/components/chat/AppChatHeader.tsx
+++ b/src/components/chat/AppChatHeader.tsx
@@ -4,7 +4,7 @@ import Link from "next/link";
 import { X, ChevronDown, Trash2, Edit2, Check, MessageSquarePlus } from "lucide-react";
 import { RiChatHistoryLine } from "react-icons/ri";
 import { LuMaximize2, LuMinimize2, LuPanelRightClose } from "react-icons/lu";
-import { useAui, useThreadListItem } from "@assistant-ui/react";
+import { useChatThreadListItem, usePromptInputThreadActions } from "@/lib/chat/runtime";
 import {
   DropdownMenu,
   DropdownMenuContent,
@@ -60,11 +60,11 @@ export function AppChatHeader({
   const [isEditingCurrentTitle, setIsEditingCurrentTitle] = useState(false);
   const [currentTitleEditValue, setCurrentTitleEditValue] = useState("");
   const currentTitleTextareaRef = useRef<HTMLTextAreaElement>(null);
-  const aui = useAui();
+  const threadActions = usePromptInputThreadActions();
 
   // Get current thread title and state from assistant-ui
   // Using safe hook to handle race condition during thread switching (GitHub issue #2722)
-  const threadListItem = useThreadListItem();
+  const threadListItem = useChatThreadListItem();
   const currentThreadTitle = threadListItem?.title || "New Chat";
   const isThreadInitialized = !!threadListItem?.remoteId;
 
@@ -153,7 +153,7 @@ export function AppChatHeader({
 
     if (trimmedValue && trimmedValue !== currentThreadTitle) {
       try {
-        await aui?.threadListItem().rename(trimmedValue);
+        await threadActions?.rename(trimmedValue);
         toast.success("Title updated");
       } catch (error) {
         console.error("Failed to rename thread:", error);
@@ -445,4 +445,3 @@ export default AppChatHeader;
 
 
 
-
diff --git a/src/components/chat/MessageContextBadges.tsx b/src/components/chat/MessageContextBadges.tsx
index 48df7055..9866aa86 100644
--- a/src/components/chat/MessageContextBadges.tsx
+++ b/src/components/chat/MessageContextBadges.tsx
@@ -1,7 +1,7 @@
 "use client";
 
 import { memo } from "react";
-import { useAuiState } from "@assistant-ui/react";
+import { useCurrentChatMessage } from "@/lib/chat/runtime";
 import { BsArrowReturnRight } from "react-icons/bs";
 import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
 import type { ReplySelection } from "@/lib/stores/ui-store";
@@ -16,7 +16,7 @@ type MessageCustomMetadata = {
  * when viewing user messages in history.
  */
 function MessageContextBadgesImpl() {
-  const message = useAuiState((s) => s.message);
+  const message = useCurrentChatMessage();
   if (!message || message.role !== "user") return null;
 
   const custom = (message.metadata as { custom?: MessageCustomMetadata } | undefined)?.custom;
diff --git a/src/components/pdf/PdfPanelHeader.tsx b/src/components/pdf/PdfPanelHeader.tsx
index d053ea77..64a0d54f 100644
--- a/src/components/pdf/PdfPanelHeader.tsx
+++ b/src/components/pdf/PdfPanelHeader.tsx
@@ -17,7 +17,7 @@ import {
 
 import ChatFloatingButton from "@/components/chat/ChatFloatingButton";
 import { useUIStore } from "@/lib/stores/ui-store";
-import { useAui } from "@assistant-ui/react";
+import { usePromptInput } from "@/lib/chat/runtime";
 import { toast } from "sonner";
 import { focusComposerInput } from "@/lib/utils/composer-utils";
 
@@ -83,9 +83,9 @@ export const PdfPanelHeader = memo(function PdfPanelHeader({
     const captureRef = useRef(capture);
     captureRef.current = capture;
 
-    const aui = useAui();
-    const auiRef = useRef(aui);
-    auiRef.current = aui;
+    const promptInput = usePromptInput();
+    const promptInputRef = useRef(promptInput);
+    promptInputRef.current = promptInput;
 
     const isChatExpanded = useUIStore((state) => state.isChatExpanded);
     const setIsChatExpanded = useUIStore((state) => state.setIsChatExpanded);
@@ -102,7 +102,7 @@ export const PdfPanelHeader = memo(function PdfPanelHeader({
                 const file = new File([result.blob], filename, { type: result.imageType });
 
                 // Add attachment to composer
-                await auiRef.current.composer().addAttachment(file);
+                await promptInputRef.current?.addAttachment(file);
                 toast.success("Screenshot added to chat");
 
                 // Turn off capture mode
diff --git a/src/components/workspace-canvas/QuizContent.tsx b/src/components/workspace-canvas/QuizContent.tsx
index 18ef0237..608f7776 100644
--- a/src/components/workspace-canvas/QuizContent.tsx
+++ b/src/components/workspace-canvas/QuizContent.tsx
@@ -21,7 +21,7 @@ import {
 } from "lucide-react";
 import { StreamdownMarkdown } from "@/components/ui/streamdown-markdown";
 import { toast } from "sonner";
-import { useAui } from "@assistant-ui/react";
+import { usePromptInput } from "@/lib/chat/runtime";
 import { useUIStore } from "@/lib/stores/ui-store";
 import { focusComposerInput } from "@/lib/utils/composer-utils";
 
@@ -40,7 +40,7 @@ export function QuizContent({
 }: QuizContentProps) {
   const quizData = item.data as QuizData;
   const questions = quizData.questions || [];
-  const aui = useAui();
+  const promptInput = usePromptInput();
 
   // UI store for card selection
   const selectedCardIds = useUIStore((state) => state.selectedCardIds);
@@ -191,7 +191,7 @@ export function QuizContent({
     }
 
     // Then send the message via composer
-    const composer = aui?.composer?.();
+    const composer = promptInput;
     if (composer) {
       try {
         composer.setText("Add 5 more questions to this quiz");
@@ -210,7 +210,7 @@ export function QuizContent({
       toggleCardSelection(item.id);
     }
 
-    const composer = aui?.composer?.();
+    const composer = promptInput;
     if (composer) {
       try {
         composer.setText(
@@ -231,7 +231,7 @@ export function QuizContent({
       toggleCardSelection(item.id);
     }
 
-    const composer = aui?.composer?.();
+    const composer = promptInput;
     if (composer) {
       try {
         const userAnswer =
diff --git a/src/hooks/ai/use-workspace-context-provider.ts b/src/hooks/ai/use-workspace-context-provider.ts
index 6195a330..80338cf6 100644
--- a/src/hooks/ai/use-workspace-context-provider.ts
+++ b/src/hooks/ai/use-workspace-context-provider.ts
@@ -1,4 +1,4 @@
-import { useAssistantContext } from "@assistant-ui/react";
+import { useChatAssistantContext } from "@/lib/chat/runtime";
 import { useCallback, useMemo } from "react";
 import type { Item } from "@/lib/workspace-state/types";
 import { formatWorkspaceContext } from "@/lib/utils/format-workspace-context";
@@ -21,7 +21,7 @@ export function useWorkspaceContextProvider(
 
   const getContext = useCallback(() => contextInstructions, [contextInstructions]);
 
-  useAssistantContext({
+  useChatAssistantContext({
     disabled: !workspaceId,
     getContext,
   });
diff --git a/src/lib/chat/runtime/hooks.ts b/src/lib/chat/runtime/hooks.ts
index 9ca0bb3f..a67e2736 100644
--- a/src/lib/chat/runtime/hooks.ts
+++ b/src/lib/chat/runtime/hooks.ts
@@ -3,17 +3,23 @@
 import {
   useAui,
   useAuiState,
+  useAssistantContext,
   useMessage,
   useMessagePartText,
   useScrollLock,
+  useThreadListItem,
 } from "@assistant-ui/react";
 import { useShallow } from "zustand/shallow";
 import type {
   ChatAttachmentSnapshot,
+  ChatAssistantContextOptions,
   ChatMessage,
+  ChatThreadListItem,
   ComposerActions,
   ComposerStateSnapshot,
+  CurrentChatMessage,
   AttachmentScope,
+  PromptInputThreadActions,
   ThreadState,
 } from "./types";
 
@@ -197,3 +203,46 @@ export function useAttachmentSnapshot(): ChatAttachmentSnapshot | undefined {
     }),
   );
 }
+
+/** Current chat message from MessageRuntime context — returns the entire message object. */
+export function useCurrentChatMessage(): CurrentChatMessage | undefined {
+  return useAuiState((s: any) => s.message as CurrentChatMessage | undefined);
+}
+
+/** id of the current thread-list item (if open inside a ThreadListItemPrimitive.Root). */
+export function useThreadListItemId(): string | undefined {
+  return useAuiState((s: any) => (s.threadListItem as { id?: string } | undefined)?.id);
+}
+
+/**
+ * Safe wrapper around assistant-ui's useThreadListItem() — returns the current thread list item.
+ * Used by AppChatHeader and thread-list-dropdown to read the current thread title and initialization state.
+ */
+export function useChatThreadListItem(): ChatThreadListItem | undefined {
+  return useThreadListItem() as unknown as ChatThreadListItem | undefined;
+}
+
+/**
+ * Actions exposed on the current thread-list item (e.g. rename). Returns null when runtime isn't ready.
+ * Consumers previously wrote `aui?.threadListItem().rename(title)` — this facade moves that call behind the ACL.
+ */
+export function usePromptInputThreadActions(): PromptInputThreadActions | null {
+  const aui = useAui();
+  if (!aui) return null;
+  const item = (aui as unknown as { threadListItem?: () => { rename?: (t: string) => Promise<unknown> } }).threadListItem?.();
+  if (!item) return null;
+  return {
+    rename: (newTitle: string) => {
+      if (!item.rename) return Promise.resolve();
+      return item.rename(newTitle);
+    },
+  };
+}
+
+/**
+ * Inject workspace/user context into the assistant's prompt.
+ * Thin wrapper around assistant-ui's useAssistantContext — same signature.
+ */
+export function useChatAssistantContext(options: ChatAssistantContextOptions): void {
+  useAssistantContext(options);
+}
diff --git a/src/lib/chat/runtime/primitives.ts b/src/lib/chat/runtime/primitives.ts
index b7e312dc..3f41be5d 100644
--- a/src/lib/chat/runtime/primitives.ts
+++ b/src/lib/chat/runtime/primitives.ts
@@ -6,6 +6,8 @@ import {
   ComposerPrimitive,
   ErrorPrimitive,
   MessagePrimitive,
+  ThreadListItemPrimitive,
+  ThreadListPrimitive,
   ThreadPrimitive,
 } from "@assistant-ui/react";
 import type { BranchPickerPrimitive as _BPP } from "@assistant-ui/react";
@@ -18,5 +20,7 @@ export const ChatActionBar = ActionBarPrimitive;
 export const ChatBranchPicker = BranchPickerPrimitive;
 export const ChatError = ErrorPrimitive;
 export const ChatIf = AuiIf;
+export const ChatThreadList = ThreadListPrimitive;
+export const ChatThreadListItem = ThreadListItemPrimitive;
 
 export type ChatBranchPickerRootProps = _BPP.Root.Props;
diff --git a/src/lib/chat/runtime/types.ts b/src/lib/chat/runtime/types.ts
index 72953d40..ac4749cf 100644
--- a/src/lib/chat/runtime/types.ts
+++ b/src/lib/chat/runtime/types.ts
@@ -86,3 +86,34 @@ export interface ChatAttachmentSnapshot {
 
 /** Props for the neutral ChatAttachment primitive — mirrors AttachmentPrimitive.Root. */
 export type ChatAttachmentRootProps = _AP.Root.Props;
+
+/**
+ * Thread list item handle (returned by useChatThreadListItem).
+ * The `remoteId` is populated after the thread has been initialized on the server.
+ */
+export interface ChatThreadListItem {
+  id?: string;
+  title?: string;
+  remoteId?: string;
+  [key: string]: unknown;
+}
+
+/** The current chat message as exposed by useCurrentChatMessage. Used by MessageContextBadges. */
+export interface CurrentChatMessage {
+  id?: string;
+  role?: "user" | "assistant" | "system";
+  metadata?: { custom?: Record<string, unknown>; [k: string]: unknown };
+  content?: unknown[];
+  [key: string]: unknown;
+}
+
+/** Options for useChatAssistantContext, mirrors assistant-ui's useAssistantContext. */
+export interface ChatAssistantContextOptions {
+  disabled?: boolean;
+  getContext: () => string;
+}
+
+/** Handle returned by usePromptInputThreadActions — lets consumers rename the current thread. */
+export interface PromptInputThreadActions {
+  rename(newTitle: string): Promise<unknown>;
+}

From 10ca7705c963292421046e2b57cf1aae76e8ebad Mon Sep 17 00:00:00 2001
From: urjitc <135136842+urjitc@users.noreply.github.com>
Date: Sun, 19 Apr 2026 22:55:15 +0000
Subject: [PATCH 6/7] PR #406 review follow-ups: 6 targeted fixes

- types.ts: drop catch-all `type: string` variant from ChatMessagePart
  so discriminated-union narrowing works on part.type === "text"
- use-prompt-input-paste: wrap handler in useCallback and early-return
  when promptInput is null (don't preventDefault if we can't handle it)
- BranchPicker: destructure hideWhenSingleBranch out of rest so callers
  can't override the forced default via prop spread
- PdfPanelHeader: throw explicit "Chat composer not ready" error when
  promptInputRef.current is null so the catch fires with the proper
  error toast instead of silently reporting success
- hooks.ts: delete dead useThreadState hook (zero consumers; narrow
  hooks like useIsThreadLoading handle all callers)
- hooks.ts: replace (s: any) casts with a local AuiStateShape interface
  describing the state fields the ACL actually reads

Pre-existing issues on main that this PR preserves intentionally
(clipboard.writeText not awaited, edit-onSubmit guard, nested Button-in-
Link, surrogate-pair slicing, process-pdf unhandled rejection) are
explicitly out of scope and not touched here.
---
 .../assistant-ui/thread/BranchPicker.tsx      |  1 +
 .../thread/hooks/use-prompt-input-paste.ts    | 73 +++++++-------
 src/components/pdf/PdfPanelHeader.tsx         |  6 +-
 src/lib/chat/runtime/hooks.ts                 | 95 ++++++++++++-------
 src/lib/chat/runtime/types.ts                 |  3 +-
 5 files changed, 106 insertions(+), 72 deletions(-)

diff --git a/src/components/assistant-ui/thread/BranchPicker.tsx b/src/components/assistant-ui/thread/BranchPicker.tsx
index 0dea7fd8..527edcf0 100644
--- a/src/components/assistant-ui/thread/BranchPicker.tsx
+++ b/src/components/assistant-ui/thread/BranchPicker.tsx
@@ -11,6 +11,7 @@ import { cn } from "@/lib/utils";
 
 export const BranchPicker: FC<ChatBranchPickerRootProps> = ({
   className,
+  hideWhenSingleBranch: _hideWhenSingleBranch,
   ...rest
 }) => {
   return (
diff --git a/src/components/assistant-ui/thread/hooks/use-prompt-input-paste.ts b/src/components/assistant-ui/thread/hooks/use-prompt-input-paste.ts
index 190fa681..c5d3b48e 100644
--- a/src/components/assistant-ui/thread/hooks/use-prompt-input-paste.ts
+++ b/src/components/assistant-ui/thread/hooks/use-prompt-input-paste.ts
@@ -1,6 +1,6 @@
 "use client";
 
-import type { ClipboardEvent } from "react";
+import { useCallback, type ClipboardEvent } from "react";
 import type { ComposerActions } from "@/lib/chat/runtime";
 
 interface UsePromptInputPasteArgs {
@@ -12,42 +12,45 @@ export function usePromptInputPaste({
   promptInput,
   workspaceId,
 }: UsePromptInputPasteArgs) {
-  return async function handlePaste(e: ClipboardEvent<HTMLTextAreaElement>) {
-    const clipboardData = e.clipboardData;
-    if (!clipboardData || !workspaceId) return;
-
-    const files = Array.from(clipboardData.files) as File[];
-
-    if (files.length > 0) {
-      e.preventDefault();
-      const imageFile = files.find((file: File) => file.type.startsWith("image/"));
-      const fileToUpload = imageFile || files[0];
-
-      if (fileToUpload) {
-        try {
-          await promptInput?.addAttachment(fileToUpload);
-        } catch (error) {
-          console.error("Failed to add file attachment:", error);
+  return useCallback(
+    async (e: ClipboardEvent<HTMLTextAreaElement>) => {
+      const clipboardData = e.clipboardData;
+      if (!clipboardData || !workspaceId || !promptInput) return;
+
+      const files = Array.from(clipboardData.files) as File[];
+
+      if (files.length > 0) {
+        e.preventDefault();
+        const imageFile = files.find((file: File) => file.type.startsWith("image/"));
+        const fileToUpload = imageFile || files[0];
+
+        if (fileToUpload) {
+          try {
+            await promptInput.addAttachment(fileToUpload);
+          } catch (error) {
+            console.error("Failed to add file attachment:", error);
+          }
         }
+        return;
       }
-      return;
-    }
-
-    const clipboardItems = Array.from(clipboardData.items) as DataTransferItem[];
-    const imageItem = clipboardItems.find((item: DataTransferItem) =>
-      item.type.startsWith("image/"),
-    );
-
-    if (imageItem) {
-      e.preventDefault();
-      const file = imageItem.getAsFile();
-      if (file) {
-        try {
-          await promptInput?.addAttachment(file);
-        } catch (error) {
-          console.error("Failed to add image attachment:", error);
+
+      const clipboardItems = Array.from(clipboardData.items) as DataTransferItem[];
+      const imageItem = clipboardItems.find((item: DataTransferItem) =>
+        item.type.startsWith("image/"),
+      );
+
+      if (imageItem) {
+        e.preventDefault();
+        const file = imageItem.getAsFile();
+        if (file) {
+          try {
+            await promptInput.addAttachment(file);
+          } catch (error) {
+            console.error("Failed to add image attachment:", error);
+          }
         }
       }
-    }
-  };
+    },
+    [promptInput, workspaceId],
+  );
 }
diff --git a/src/components/pdf/PdfPanelHeader.tsx b/src/components/pdf/PdfPanelHeader.tsx
index 64a0d54f..f22d660e 100644
--- a/src/components/pdf/PdfPanelHeader.tsx
+++ b/src/components/pdf/PdfPanelHeader.tsx
@@ -102,7 +102,11 @@ export const PdfPanelHeader = memo(function PdfPanelHeader({
                 const file = new File([result.blob], filename, { type: result.imageType });
 
                 // Add attachment to composer
-                await promptInputRef.current?.addAttachment(file);
+                const promptInput = promptInputRef.current;
+                if (!promptInput) {
+                    throw new Error("Chat composer not ready");
+                }
+                await promptInput.addAttachment(file);
                 toast.success("Screenshot added to chat");
 
                 // Turn off capture mode
diff --git a/src/lib/chat/runtime/hooks.ts b/src/lib/chat/runtime/hooks.ts
index a67e2736..f202a5fc 100644
--- a/src/lib/chat/runtime/hooks.ts
+++ b/src/lib/chat/runtime/hooks.ts
@@ -20,40 +20,69 @@ import type {
   CurrentChatMessage,
   AttachmentScope,
   PromptInputThreadActions,
-  ThreadState,
 } from "./types";
 
-export function useThreadState(): ThreadState {
-  return useAuiState((s: any) => ({
-    messageCount: s.thread?.messages?.length ?? 0,
-    isLoading: !!s.thread?.isLoading,
-    isEmpty: s.thread?.isEmpty ?? true,
-    isRunning: !!s.thread?.isRunning,
-  }));
+/**
+ * Local shape describing the @assistant-ui/react state fields we actually read.
+ * This is intentionally minimal — assistant-ui does not export its state type,
+ * so the ACL defines what it needs. Expand as new selectors are added.
+ */
+interface AuiStateShape {
+  thread?: {
+    messages?: ReadonlyArray<{ id?: string }>;
+    isLoading?: boolean;
+    isEmpty?: boolean;
+    isRunning?: boolean;
+  };
+  threads?: {
+    mainThreadId?: string;
+  };
+  threadListItem?: {
+    id?: string;
+  };
+  composer?: {
+    text?: string;
+  };
+  message?: {
+    id?: string;
+    role?: "user" | "assistant" | "system";
+    status?: { type?: string };
+    parts?: ReadonlyArray<{ type?: string; text?: string }>;
+    content?: readonly unknown[];
+    metadata?: { custom?: Record<string, unknown>; [k: string]: unknown };
+  };
+  attachment?: {
+    id?: string;
+    type?: string;
+    name?: string;
+    file?: File & { name: string };
+    content?: ReadonlyArray<{ type: string; text?: string; image?: string }>;
+  };
 }
 
+
 export function useIsThreadLoading(): boolean {
-  return useAuiState((s: any) => !!s.thread?.isLoading);
+  return useAuiState((s: AuiStateShape) => !!s.thread?.isLoading);
 }
 
 export function useIsThreadEmpty(): boolean {
-  return useAuiState((s: any) => s.thread?.isEmpty ?? true);
+  return useAuiState((s: AuiStateShape) => s.thread?.isEmpty ?? true);
 }
 
 export function useIsThreadRunning(): boolean {
-  return useAuiState((s: any) => !!s.thread?.isRunning);
+  return useAuiState((s: AuiStateShape) => !!s.thread?.isRunning);
 }
 
 export function useThreadMessageCount(): number {
-  return useAuiState((s: any) => s.thread?.messages?.length ?? 0);
+  return useAuiState((s: AuiStateShape) => s.thread?.messages?.length ?? 0);
 }
 
 export function useMainThreadId(): string | null {
-  return useAuiState((s: any) => s.threads?.mainThreadId ?? null);
+  return useAuiState((s: AuiStateShape) => s.threads?.mainThreadId ?? null);
 }
 
 export function useHasPromptInputText(): boolean {
-  return useAuiState((s: any) => Boolean((s.composer?.text ?? "").trim()));
+  return useAuiState((s: AuiStateShape) => Boolean((s.composer?.text ?? "").trim()));
 }
 
 export function useChatMessage(): ChatMessage {
@@ -86,12 +115,12 @@ export const useChatScrollLock = useScrollLock;
 
 /** Current assistant/user message id from MessageRuntime context. */
 export function useCurrentMessageId(): string | null {
-  return useAuiState((s: any) => (s.message as { id?: string } | undefined)?.id ?? null);
+  return useAuiState((s: AuiStateShape) => s.message?.id ?? null);
 }
 
 /** True while this message is actively streaming. */
 export function useIsMessageRunning(): boolean {
-  return useAuiState((s: any) => (s.message as { status?: { type?: string } } | undefined)?.status?.type === "running");
+  return useAuiState((s: AuiStateShape) => s.message?.status?.type === "running");
 }
 
 /**
@@ -99,8 +128,8 @@ export function useIsMessageRunning(): boolean {
  * states before the first token arrives.
  */
 export function useIsMessageEmpty(): boolean {
-  return useAuiState((s: any) => {
-    const msg = s.message as { content?: unknown[] } | undefined;
+  return useAuiState((s: AuiStateShape) => {
+    const msg = s.message;
     return !msg?.content || (Array.isArray(msg.content) && msg.content.length === 0);
   });
 }
@@ -110,9 +139,9 @@ export function useIsMessageEmpty(): boolean {
  * Used to gate old reasoning/tool groups from rendering once newer messages exist.
  */
 export function useIsLastMessage(): boolean {
-  return useAuiState((s: any) => {
-    const thread = s.thread as { messages?: Array<{ id?: string }> } | undefined;
-    const messageId = (s.message as { id?: string } | undefined)?.id;
+  return useAuiState((s: AuiStateShape) => {
+    const thread = s.thread;
+    const messageId = s.message?.id;
     const messages = thread?.messages ?? [];
     const idx = messages.findIndex((m) => m.id === messageId);
     return idx >= 0 && idx === messages.length - 1;
@@ -129,10 +158,8 @@ export function useIsMessagePartStreaming(
   startIndex: number,
   endIndex: number,
 ): boolean {
-  return useAuiState((s: any) => {
-    const msg = s.message as
-      | { status?: { type?: string }; parts?: Array<{ type?: string }> }
-      | undefined;
+  return useAuiState((s: AuiStateShape) => {
+    const msg = s.message;
     if (msg?.status?.type !== "running") return false;
     const parts = msg.parts ?? [];
     const lastIndex = parts.length - 1;
@@ -152,8 +179,8 @@ export function useMessagePartTextLengthSnapshot(
   startIndex: number,
   endIndex: number,
 ): number {
-  return useAuiState((s: any) => {
-    const parts = (s.message as { parts?: Array<{ type?: string; text?: string }> } | undefined)?.parts ?? [];
+  return useAuiState((s: AuiStateShape) => {
+    const parts = s.message?.parts ?? [];
     let len = 0;
     for (let i = startIndex; i <= endIndex && i < parts.length; i++) {
       const p = parts[i];
@@ -175,12 +202,12 @@ export function useAttachmentScope(): AttachmentScope {
 
 /** Attachment id from AttachmentPrimitive context. */
 export function useAttachmentId(): string | undefined {
-  return useAuiState((s: any) => (s.attachment as { id?: string } | undefined)?.id);
+  return useAuiState((s: AuiStateShape) => s.attachment?.id);
 }
 
 /** Whether the current attachment is an image. */
 export function useIsAttachmentImage(): boolean {
-  return useAuiState((s: any) => (s.attachment as { type?: string } | undefined)?.type === "image");
+  return useAuiState((s: AuiStateShape) => s.attachment?.type === "image");
 }
 
 /**
@@ -190,8 +217,8 @@ export function useIsAttachmentImage(): boolean {
  */
 export function useAttachmentSnapshot(): ChatAttachmentSnapshot | undefined {
   return useAuiState(
-    useShallow((s: any) => {
-      const att = s.attachment as ChatAttachmentSnapshot | undefined;
+    useShallow((s: AuiStateShape) => {
+      const att = s.attachment;
       if (!att) return undefined;
       return {
         id: att.id,
@@ -199,19 +226,19 @@ export function useAttachmentSnapshot(): ChatAttachmentSnapshot | undefined {
         name: att.name,
         file: att.file,
         content: att.content,
-      };
+      } as ChatAttachmentSnapshot;
     }),
   );
 }
 
 /** Current chat message from MessageRuntime context — returns the entire message object. */
 export function useCurrentChatMessage(): CurrentChatMessage | undefined {
-  return useAuiState((s: any) => s.message as CurrentChatMessage | undefined);
+  return useAuiState((s: AuiStateShape) => s.message as CurrentChatMessage | undefined);
 }
 
 /** id of the current thread-list item (if open inside a ThreadListItemPrimitive.Root). */
 export function useThreadListItemId(): string | undefined {
-  return useAuiState((s: any) => (s.threadListItem as { id?: string } | undefined)?.id);
+  return useAuiState((s: AuiStateShape) => s.threadListItem?.id);
 }
 
 /**
diff --git a/src/lib/chat/runtime/types.ts b/src/lib/chat/runtime/types.ts
index ac4749cf..b1a15229 100644
--- a/src/lib/chat/runtime/types.ts
+++ b/src/lib/chat/runtime/types.ts
@@ -19,8 +19,7 @@ export type ChatMessagePart =
   | { type: "image"; [k: string]: unknown }
   | { type: "source"; [k: string]: unknown }
   | { type: "reasoning"; [k: string]: unknown }
-  | { type: "tool-call" | "tool-result" | "tool"; [k: string]: unknown }
-  | { type: string; [k: string]: unknown };
+  | { type: "tool-call" | "tool-result" | "tool"; [k: string]: unknown };
 
 export interface ChatMessage {
   id?: string;

From c94a10475fa93ba24418e908be80a769e84c6787 Mon Sep 17 00:00:00 2001
From: urjitc <135136842+urjitc@users.noreply.github.com>
Date: Sun, 19 Apr 2026 23:18:42 +0000
Subject: [PATCH 7/7] Fix usePromptInput / usePromptInputThreadActions crash
 outside AuiProvider

PromptBuilderDialog is mounted from WorkspaceHeader and WorkspaceSection,
which render outside WorkspaceRuntimeProvider when no workspace is
active. useAui() there returns a DefaultAssistantClient proxy whose
.composer and .threadListItem getters return functions that throw when
invoked ("You are using a component or hook that requires an
AuiProvider..."). The ACL's usePromptInput() and
usePromptInputThreadActions() called those getters during render,
crashing the component via ErrorBoundary.

Defer proxy resolution to call time, wrapped in try/catch. Matches the
pre-refactor behavior where `aui?.composer()?.setText(...)` only
resolved at event time (e.g. inside a submit handler, by which point
the component is definitely inside a provider).
---
 src/lib/chat/runtime/hooks.ts | 53 ++++++++++++++++++++++++++++-------
 1 file changed, 43 insertions(+), 10 deletions(-)

diff --git a/src/lib/chat/runtime/hooks.ts b/src/lib/chat/runtime/hooks.ts
index f202a5fc..adeb546a 100644
--- a/src/lib/chat/runtime/hooks.ts
+++ b/src/lib/chat/runtime/hooks.ts
@@ -97,14 +97,35 @@ export function useChatMessagePartText() {
 export function usePromptInput(): ComposerActions | null {
   const aui = useAui();
   if (!aui) return null;
-  const composer = aui.composer?.();
-  if (!composer) return null;
+
+  const resolveComposer = () => {
+    try {
+      return aui.composer?.() ?? null;
+    } catch {
+      return null;
+    }
+  };
+
   return {
-    setText: (t) => composer.setText(t),
-    send: () => composer.send(),
-    addAttachment: (f) => composer.addAttachment(f),
-    setRunConfig: (cfg) => composer.setRunConfig(cfg as any),
-    getState: () => composer.getState() as unknown as ComposerStateSnapshot | undefined,
+    setText: (t) => {
+      resolveComposer()?.setText(t);
+    },
+    send: () => {
+      resolveComposer()?.send();
+    },
+    addAttachment: async (f) => {
+      const composer = resolveComposer();
+      if (!composer) return;
+      return composer.addAttachment(f);
+    },
+    setRunConfig: (cfg) => {
+      resolveComposer()?.setRunConfig(cfg as any);
+    },
+    getState: () => {
+      const composer = resolveComposer();
+      if (!composer) return undefined;
+      return composer.getState() as unknown as ComposerStateSnapshot | undefined;
+    },
   };
 }
 /**
@@ -256,11 +277,23 @@ export function useChatThreadListItem(): ChatThreadListItem | undefined {
 export function usePromptInputThreadActions(): PromptInputThreadActions | null {
   const aui = useAui();
   if (!aui) return null;
-  const item = (aui as unknown as { threadListItem?: () => { rename?: (t: string) => Promise<unknown> } }).threadListItem?.();
-  if (!item) return null;
+
+  const resolveItem = () => {
+    try {
+      return (
+        (aui as unknown as {
+          threadListItem?: () => { rename?: (t: string) => Promise<unknown> };
+        }).threadListItem?.() ?? null
+      );
+    } catch {
+      return null;
+    }
+  };
+
   return {
     rename: (newTitle: string) => {
-      if (!item.rename) return Promise.resolve();
+      const item = resolveItem();
+      if (!item?.rename) return Promise.resolve();
       return item.rename(newTitle);
     },
   };