refactor: update dependencies and streamline assistant-ui package usages

This commit is contained in:
Anish Sarkar 2026-03-24 02:22:51 +05:30
parent fed3a3b436
commit b8f3f41326
40 changed files with 886 additions and 1110 deletions

View file

@ -1,10 +1,9 @@
import {
ActionBarPrimitive,
AssistantIf,
AuiIf,
ErrorPrimitive,
MessagePrimitive,
useAssistantState,
useMessage,
useAuiState,
} from "@assistant-ui/react";
import { useAtomValue } from "jotai";
import { CheckIcon, CopyIcon, DownloadIcon, MessageSquare, RefreshCwIcon } from "lucide-react";
@ -21,6 +20,22 @@ import { ToolFallback } from "@/components/assistant-ui/tool-fallback";
import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
import { CommentPanelContainer } from "@/components/chat-comments/comment-panel-container/comment-panel-container";
import { CommentSheet } from "@/components/chat-comments/comment-sheet/comment-sheet";
import { CreateConfluencePageToolUI, DeleteConfluencePageToolUI, UpdateConfluencePageToolUI } from "@/components/tool-ui/confluence";
import { DeepAgentThinkingToolUI } from "@/components/tool-ui/deepagent-thinking";
import { DisplayImageToolUI } from "@/components/tool-ui/display-image";
import { GeneratePodcastToolUI } from "@/components/tool-ui/generate-podcast";
import { GenerateReportToolUI } from "@/components/tool-ui/generate-report";
import { GenerateVideoPresentationToolUI } from "@/components/tool-ui/video-presentation";
import { CreateGmailDraftToolUI, SendGmailEmailToolUI, TrashGmailEmailToolUI, UpdateGmailDraftToolUI } from "@/components/tool-ui/gmail";
import { CreateCalendarEventToolUI, DeleteCalendarEventToolUI, UpdateCalendarEventToolUI } from "@/components/tool-ui/google-calendar";
import { CreateGoogleDriveFileToolUI, DeleteGoogleDriveFileToolUI } from "@/components/tool-ui/google-drive";
import { CreateJiraIssueToolUI, DeleteJiraIssueToolUI, UpdateJiraIssueToolUI } from "@/components/tool-ui/jira";
import { CreateLinearIssueToolUI, DeleteLinearIssueToolUI, UpdateLinearIssueToolUI } from "@/components/tool-ui/linear";
import { LinkPreviewToolUI, MultiLinkPreviewToolUI } from "@/components/tool-ui/link-preview";
import { CreateNotionPageToolUI, DeleteNotionPageToolUI, UpdateNotionPageToolUI } from "@/components/tool-ui/notion";
import { SandboxExecuteToolUI } from "@/components/tool-ui/sandbox-execute";
import { ScrapeWebpageToolUI } from "@/components/tool-ui/scrape-webpage";
import { RecallMemoryToolUI, SaveMemoryToolUI } from "@/components/tool-ui/user-memory";
import { useComments } from "@/hooks/use-comments";
import { useMediaQuery } from "@/hooks/use-media-query";
import { cn } from "@/lib/utils";
@ -42,13 +57,13 @@ const ThinkingStepsPart: FC = () => {
const thinkingStepsMap = useContext(ThinkingStepsContext);
// Get the current message ID to look up thinking steps
const messageId = useAssistantState(({ message }) => message?.id);
const messageId = useAuiState(({ message }) => message?.id);
const thinkingSteps = thinkingStepsMap.get(messageId) || [];
// Check if this specific message is currently streaming
// A message is streaming if: thread is running AND this is the last assistant message
const isThreadRunning = useAssistantState(({ thread }) => thread.isRunning);
const isLastMessage = useAssistantState(({ message }) => message?.isLast ?? false);
const isThreadRunning = useAuiState(({ thread }) => thread.isRunning);
const isLastMessage = useAuiState(({ message }) => message?.isLast ?? false);
const isMessageStreaming = isThreadRunning && isLastMessage;
if (thinkingSteps.length === 0) return null;
@ -70,7 +85,43 @@ const AssistantMessageInner: FC = () => {
<MessagePrimitive.Parts
components={{
Text: MarkdownText,
tools: { Fallback: ToolFallback },
tools: {
by_name: {
deepagent_thinking: DeepAgentThinkingToolUI,
generate_report: GenerateReportToolUI,
generate_podcast: GeneratePodcastToolUI,
generate_video_presentation: GenerateVideoPresentationToolUI,
link_preview: LinkPreviewToolUI,
multi_link_preview: MultiLinkPreviewToolUI,
display_image: DisplayImageToolUI,
scrape_webpage: ScrapeWebpageToolUI,
save_memory: SaveMemoryToolUI,
recall_memory: RecallMemoryToolUI,
execute: SandboxExecuteToolUI,
create_notion_page: CreateNotionPageToolUI,
update_notion_page: UpdateNotionPageToolUI,
delete_notion_page: DeleteNotionPageToolUI,
create_linear_issue: CreateLinearIssueToolUI,
update_linear_issue: UpdateLinearIssueToolUI,
delete_linear_issue: DeleteLinearIssueToolUI,
create_google_drive_file: CreateGoogleDriveFileToolUI,
delete_google_drive_file: DeleteGoogleDriveFileToolUI,
create_calendar_event: CreateCalendarEventToolUI,
update_calendar_event: UpdateCalendarEventToolUI,
delete_calendar_event: DeleteCalendarEventToolUI,
create_gmail_draft: CreateGmailDraftToolUI,
update_gmail_draft: UpdateGmailDraftToolUI,
send_gmail_email: SendGmailEmailToolUI,
trash_gmail_email: TrashGmailEmailToolUI,
create_jira_issue: CreateJiraIssueToolUI,
update_jira_issue: UpdateJiraIssueToolUI,
delete_jira_issue: DeleteJiraIssueToolUI,
create_confluence_page: CreateConfluencePageToolUI,
update_confluence_page: UpdateConfluencePageToolUI,
delete_confluence_page: DeleteConfluencePageToolUI,
},
Fallback: ToolFallback,
},
}}
/>
<MessageError />
@ -95,7 +146,7 @@ export const AssistantMessage: FC = () => {
const messageRef = useRef<HTMLDivElement>(null);
const commentPanelRef = useRef<HTMLDivElement>(null);
const commentTriggerRef = useRef<HTMLButtonElement>(null);
const messageId = useAssistantState(({ message }) => message?.id);
const messageId = useAuiState(({ message }) => message?.id);
const searchSpaceId = useAtomValue(activeSearchSpaceIdAtom);
const dbMessageId = parseMessageId(messageId);
const commentsEnabled = useAtomValue(commentsEnabledAtom);
@ -104,8 +155,8 @@ export const AssistantMessage: FC = () => {
const isMediumScreen = useMediaQuery("(min-width: 768px) and (max-width: 1023px)");
const isDesktop = useMediaQuery("(min-width: 1024px)");
const isThreadRunning = useAssistantState(({ thread }) => thread.isRunning);
const isLastMessage = useAssistantState(({ message }) => message?.isLast ?? false);
const isThreadRunning = useAuiState(({ thread }) => thread.isRunning);
const isLastMessage = useAuiState(({ message }) => message?.isLast ?? false);
const isMessageStreaming = isThreadRunning && isLastMessage;
const { data: commentsData, isSuccess: commentsLoaded } = useComments({
@ -227,38 +278,38 @@ export const AssistantMessage: FC = () => {
};
const AssistantActionBar: FC = () => {
const { isLast } = useMessage();
const isLast = useAuiState((s) => s.message.isLast);
return (
<ActionBarPrimitive.Root
<ActionBarPrimitive.Root
hideWhenRunning
autohide="not-last"
autohideFloat="single-branch"
className="aui-assistant-action-bar-root -ml-1 col-start-3 row-start-2 flex gap-1 text-muted-foreground md:data-floating:absolute md:data-floating:rounded-md md:data-floating:p-1 [&>button]:opacity-100 md:[&>button]:opacity-[var(--aui-button-opacity,1)]"
>
<ActionBarPrimitive.Copy asChild>
<ActionBarPrimitive.Copy asChild>
<TooltipIconButton tooltip="Copy">
<AssistantIf condition={({ message }) => message.isCopied}>
<AuiIf condition={({ message }) => message.isCopied}>
<CheckIcon />
</AssistantIf>
<AssistantIf condition={({ message }) => !message.isCopied}>
</AuiIf>
<AuiIf condition={({ message }) => !message.isCopied}>
<CopyIcon />
</AssistantIf>
</AuiIf>
</TooltipIconButton>
</ActionBarPrimitive.Copy>
<ActionBarPrimitive.ExportMarkdown asChild>
<ActionBarPrimitive.ExportMarkdown asChild>
<TooltipIconButton tooltip="Download">
<DownloadIcon />
</TooltipIconButton>
</ActionBarPrimitive.ExportMarkdown>
{/* Only allow regenerating the last assistant message */}
{isLast && (
{/* Only allow regenerating the last assistant message */}
{isLast && (
<ActionBarPrimitive.Reload asChild>
<TooltipIconButton tooltip="Refresh">
<RefreshCwIcon />
</TooltipIconButton>
</ActionBarPrimitive.Reload>
)}
</ActionBarPrimitive.Root>
);
</ActionBarPrimitive.Root>
);
};

View file

@ -1,7 +1,6 @@
import { useAssistantState, useThreadViewport } from "@assistant-ui/react";
import { ChevronRightIcon } from "lucide-react";
import type { FC } from "react";
import { createContext, useCallback, useContext, useEffect, useRef, useState } from "react";
import { createContext, useCallback, useEffect, useState } from "react";
import { ChainOfThoughtItem } from "@/components/prompt-kit/chain-of-thought";
import { TextShimmerLoader } from "@/components/prompt-kit/loader";
import type { ThinkingStep } from "@/components/tool-ui/deepagent-thinking";
@ -154,52 +153,3 @@ export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?:
);
};
/**
* Component that handles auto-scroll when thinking steps update.
* Uses useThreadViewport to scroll to bottom when thinking steps change,
* ensuring the user always sees the latest content during streaming.
*/
export const ThinkingStepsScrollHandler: FC = () => {
const thinkingStepsMap = useContext(ThinkingStepsContext);
const viewport = useThreadViewport();
const isRunning = useAssistantState(({ thread }) => thread.isRunning);
// Track the serialized state to detect any changes
const prevStateRef = useRef<string>("");
useEffect(() => {
// Only act during streaming
if (!isRunning) {
prevStateRef.current = "";
return;
}
// Serialize the thinking steps state to detect any changes
// This catches new steps, status changes, and item additions
let stateString = "";
thinkingStepsMap.forEach((steps, msgId) => {
steps.forEach((step) => {
stateString += `${msgId}:${step.id}:${step.status}:${step.items?.length || 0};`;
});
});
// If state changed at all during streaming, scroll
if (stateString !== prevStateRef.current && stateString !== "") {
prevStateRef.current = stateString;
// Multiple attempts to ensure scroll happens after DOM updates
const scrollAttempt = () => {
try {
viewport.scrollToBottom();
} catch {
// Ignore errors - viewport might not be ready
}
};
// Delayed attempts to handle async DOM updates
requestAnimationFrame(scrollAttempt);
setTimeout(scrollAttempt, 100);
}
}, [thinkingStepsMap, viewport, isRunning]);
return null; // This component doesn't render anything
};

View file

@ -1,13 +1,13 @@
import {
ActionBarPrimitive,
AssistantIf,
AuiIf,
BranchPickerPrimitive,
ComposerPrimitive,
ErrorPrimitive,
MessagePrimitive,
ThreadPrimitive,
useAssistantState,
useComposerRuntime,
useAui,
useAuiState,
} from "@assistant-ui/react";
import { useAtom, useAtomValue, useSetAtom } from "jotai";
import {
@ -125,19 +125,19 @@ export const Thread: FC<ThreadProps> = ({ messageThinkingSteps = new Map() }) =>
const ThreadContent: FC = () => {
return (
<ThreadPrimitive.Root
<ThreadPrimitive.Root
className="aui-root aui-thread-root @container flex h-full min-h-0 flex-col bg-main-panel"
style={{
["--thread-max-width" as string]: "44rem",
}}
>
<ThreadPrimitive.Viewport
<ThreadPrimitive.Viewport
turnAnchor="top"
className="aui-thread-viewport relative flex flex-1 min-h-0 flex-col overflow-y-auto px-4 pt-4"
>
<AssistantIf condition={({ thread }) => thread.isEmpty}>
<AuiIf condition={({ thread }) => thread.isEmpty}>
<ThreadWelcome />
</AssistantIf>
</AuiIf>
<ThreadPrimitive.Messages
components={{
@ -152,15 +152,15 @@ const ThreadContent: FC = () => {
style={{ paddingBottom: "max(1rem, env(safe-area-inset-bottom))" }}
>
<ThreadScrollToBottom />
<AssistantIf condition={({ thread }) => !thread.isEmpty}>
<AuiIf condition={({ thread }) => !thread.isEmpty}>
<div className="fade-in slide-in-from-bottom-4 animate-in duration-500 ease-out fill-mode-both">
<Composer />
</div>
</AssistantIf>
</AuiIf>
</ThreadPrimitive.ViewportFooter>
</ThreadPrimitive.Viewport>
</ThreadPrimitive.Root>
);
</ThreadPrimitive.Root>
);
};
const ThreadScrollToBottom: FC = () => {
@ -327,11 +327,11 @@ const Composer: FC = () => {
const editorContainerRef = useRef<HTMLDivElement>(null);
const documentPickerRef = useRef<DocumentMentionPickerRef>(null);
const { search_space_id, chat_id } = useParams();
const composerRuntime = useComposerRuntime();
const aui = useAui();
const hasAutoFocusedRef = useRef(false);
const isThreadEmpty = useAssistantState(({ thread }) => thread.isEmpty);
const isThreadRunning = useAssistantState(({ thread }) => thread.isRunning);
const isThreadEmpty = useAuiState(({ thread }) => thread.isEmpty);
const isThreadRunning = useAuiState(({ thread }) => thread.isRunning);
// Cycling placeholder state - only cycles in new chats
const [placeholderIndex, setPlaceholderIndex] = useState(0);
@ -378,7 +378,7 @@ const Composer: FC = () => {
// hooks never fire their own network requests (eliminates N+1 API calls).
// Return a primitive string from the selector so useSyncExternalStore can
// compare snapshots by value and avoid infinite re-render loops.
const assistantIdsKey = useAssistantState(({ thread }) =>
const assistantIdsKey = useAuiState(({ thread }) =>
thread.messages
.filter((m) => m.role === "assistant" && m.id?.startsWith("msg-"))
.map((m) => m.id?.replace("msg-", ""))
@ -414,9 +414,9 @@ const Composer: FC = () => {
// Sync editor text with assistant-ui composer runtime
const handleEditorChange = useCallback(
(text: string) => {
composerRuntime.setText(text);
aui.composer().setText(text);
},
[composerRuntime]
[aui]
);
// Open document picker when @ mention is triggered
@ -469,7 +469,7 @@ const Composer: FC = () => {
return;
}
if (!showDocumentPopover) {
composerRuntime.send();
aui.composer().send();
editorRef.current?.clear();
setMentionedDocuments([]);
setSidebarDocs([]);
@ -478,7 +478,7 @@ const Composer: FC = () => {
showDocumentPopover,
isThreadRunning,
isBlockedByOtherUser,
composerRuntime,
aui,
setMentionedDocuments,
setSidebarDocs,
]);
@ -591,7 +591,7 @@ const ComposerAction: FC<ComposerActionProps> = ({ isBlockedByOtherUser = false
const atBottom = el.scrollHeight - el.scrollTop - el.clientHeight <= 2;
setToolsScrollPos(atTop ? "top" : atBottom ? "bottom" : "middle");
}, []);
const isComposerTextEmpty = useAssistantState(({ composer }) => {
const isComposerTextEmpty = useAuiState(({ composer }) => {
const text = composer.text?.trim() || "";
return text.length === 0;
});
@ -702,8 +702,8 @@ const ComposerAction: FC<ComposerActionProps> = ({ isBlockedByOtherUser = false
const isSendDisabled = isComposerEmpty || !hasModelConfigured || isBlockedByOtherUser;
return (
<div className="aui-composer-action-wrapper relative mx-3 mb-2 flex items-center justify-between">
<div className="flex items-center gap-1">
<div className="aui-composer-action-wrapper relative mx-3 mb-2 flex items-center justify-between">
<div className="flex items-center gap-1">
{!isDesktop ? (
<>
<DropdownMenu>
@ -1007,16 +1007,14 @@ const ComposerAction: FC<ComposerActionProps> = ({ isBlockedByOtherUser = false
</button>
)}
</div>
{!hasModelConfigured && (
{!hasModelConfigured && (
<div className="flex items-center gap-1.5 text-amber-600 dark:text-amber-400 text-xs">
<AlertCircle className="size-3" />
<span>Select a model</span>
</div>
)}
<div className="flex items-center gap-2">
<AssistantIf condition={({ thread }) => !thread.isRunning}>
<div className="flex items-center gap-2">
<AuiIf condition={({ thread }) => !thread.isRunning}>
<ComposerPrimitive.Send asChild disabled={isSendDisabled}>
<TooltipIconButton
tooltip={
@ -1042,9 +1040,9 @@ const ComposerAction: FC<ComposerActionProps> = ({ isBlockedByOtherUser = false
<ArrowUpIcon className="aui-composer-send-icon size-4" />
</TooltipIconButton>
</ComposerPrimitive.Send>
</AssistantIf>
</AuiIf>
<AssistantIf condition={({ thread }) => thread.isRunning}>
<AuiIf condition={({ thread }) => thread.isRunning}>
<ComposerPrimitive.Cancel asChild>
<Button
type="button"
@ -1056,10 +1054,10 @@ const ComposerAction: FC<ComposerActionProps> = ({ isBlockedByOtherUser = false
<SquareIcon className="aui-composer-cancel-icon size-3 fill-current" />
</Button>
</ComposerPrimitive.Cancel>
</AssistantIf>
</AuiIf>
</div>
</div>
);
</div>
);
};
/** Convert snake_case tool names to human-readable labels */
@ -1151,13 +1149,13 @@ const ThinkingStepsPart: FC = () => {
const thinkingStepsMap = useContext(ThinkingStepsContext);
// Get the current message ID to look up thinking steps
const messageId = useAssistantState(({ message }) => message?.id);
const messageId = useAuiState(({ message }) => message?.id);
const thinkingSteps = thinkingStepsMap.get(messageId) || [];
// Check if this specific message is currently streaming
// A message is streaming if: thread is running AND this is the last assistant message
const isThreadRunning = useAssistantState(({ thread }) => thread.isRunning);
const isLastMessage = useAssistantState(({ message }) => message?.isLast ?? false);
const isThreadRunning = useAuiState(({ thread }) => thread.isRunning);
const isLastMessage = useAuiState(({ message }) => message?.isLast ?? false);
const isMessageStreaming = isThreadRunning && isLastMessage;
if (thinkingSteps.length === 0) return null;
@ -1195,34 +1193,34 @@ const AssistantMessageInner: FC = () => {
const AssistantActionBar: FC = () => {
return (
<ActionBarPrimitive.Root
<ActionBarPrimitive.Root
hideWhenRunning
autohide="not-last"
autohideFloat="single-branch"
className="aui-assistant-action-bar-root -ml-1 col-start-3 row-start-2 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"
>
<ActionBarPrimitive.Copy asChild>
<ActionBarPrimitive.Copy asChild>
<TooltipIconButton tooltip="Copy">
<AssistantIf condition={({ message }) => message.isCopied}>
<AuiIf condition={({ message }) => message.isCopied}>
<CheckIcon />
</AssistantIf>
<AssistantIf condition={({ message }) => !message.isCopied}>
</AuiIf>
<AuiIf condition={({ message }) => !message.isCopied}>
<CopyIcon />
</AssistantIf>
</AuiIf>
</TooltipIconButton>
</ActionBarPrimitive.Copy>
<ActionBarPrimitive.ExportMarkdown asChild>
<ActionBarPrimitive.ExportMarkdown asChild>
<TooltipIconButton tooltip="Export as Markdown">
<DownloadIcon />
</TooltipIconButton>
</ActionBarPrimitive.ExportMarkdown>
<ActionBarPrimitive.Reload asChild>
<ActionBarPrimitive.Reload asChild>
<TooltipIconButton tooltip="Refresh">
<RefreshCwIcon />
</TooltipIconButton>
</ActionBarPrimitive.Reload>
</ActionBarPrimitive.Root>
);
</ActionBarPrimitive.Root>
);
};
const EditComposer: FC = () => {

View file

@ -1,4 +1,4 @@
import { ActionBarPrimitive, MessagePrimitive, useAssistantState } from "@assistant-ui/react";
import { ActionBarPrimitive, MessagePrimitive, useAuiState } from "@assistant-ui/react";
import { useAtomValue } from "jotai";
import { FileText, Pen } from "lucide-react";
import { type FC, useState } from "react";
@ -42,10 +42,10 @@ const UserAvatar: FC<AuthorMetadata> = ({ displayName, avatarUrl }) => {
};
export const UserMessage: FC = () => {
const messageId = useAssistantState(({ message }) => message?.id);
const messageId = useAuiState(({ message }) => message?.id);
const messageDocumentsMap = useAtomValue(messageDocumentsMapAtom);
const mentionedDocs = messageId ? messageDocumentsMap[messageId] : undefined;
const metadata = useAssistantState(({ message }) => message?.metadata);
const metadata = useAuiState(({ message }) => message?.metadata);
const author = metadata?.custom?.author as AuthorMetadata | undefined;
return (
@ -93,13 +93,13 @@ export const UserMessage: FC = () => {
};
const UserActionBar: FC = () => {
const isThreadRunning = useAssistantState(({ thread }) => thread.isRunning);
const isThreadRunning = useAuiState(({ thread }) => thread.isRunning);
// Get current message ID
const currentMessageId = useAssistantState(({ message }) => message?.id);
const currentMessageId = useAuiState(({ message }) => message?.id);
// Find the last user message ID in the thread (computed once, memoized by selector)
const lastUserMessageId = useAssistantState(({ thread }) => {
const lastUserMessageId = useAuiState(({ thread }) => {
const messages = thread.messages;
for (let i = messages.length - 1; i >= 0; i--) {
if (messages[i].role === "user") {