Refactor: Split thread components and add chat runtime ACL

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 }) => {

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'],

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(() => {

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) => {

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) => {

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) => {

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 }) => {

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 }) => {

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,

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 }) => {

src/components/assistant-ui/SearchWorkspaceToolUI.tsx

@@ -1,15 +1,15 @@
 "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";
 
 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 }) => {

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 });

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 }) => {

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"] = ({

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 }) => {

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;