From 27f0136479e9b4ce9a411b10843eeef52e0c6606 Mon Sep 17 00:00:00 2001 From: urjitc <135136842+urjitc@users.noreply.github.com> Date: Sat, 11 Apr 2026 07:35:51 +0000 Subject: [PATCH 1/5] Fix assistant chat performance by eliminating unnecessary re-renders, optimizing selectors, and memo --- .../assistant-ui/assistant-loader.tsx | 26 ++- src/components/assistant-ui/markdown-text.tsx | 167 ++++++++++-------- src/components/assistant-ui/reasoning.tsx | 6 +- src/components/assistant-ui/thread.tsx | 8 +- src/components/assistant-ui/tool-group.tsx | 6 +- src/components/chat/MessageContextBadges.tsx | 16 +- src/components/ui/streamdown-markdown.tsx | 2 +- .../utils/__tests__/preprocess-latex.test.ts | 45 +++-- src/lib/utils/preprocess-latex.ts | 53 ++---- 9 files changed, 178 insertions(+), 151 deletions(-) diff --git a/src/components/assistant-ui/assistant-loader.tsx b/src/components/assistant-ui/assistant-loader.tsx index 0ae83746..19a00353 100644 --- a/src/components/assistant-ui/assistant-loader.tsx +++ b/src/components/assistant-ui/assistant-loader.tsx @@ -1,10 +1,14 @@ "use client"; +import { memo, lazy, Suspense } from "react"; import { useAuiState } from "@assistant-ui/react"; -import { DotLottieReact } from "@lottiefiles/dotlottie-react"; import { useTheme } from "next-themes"; -export const AssistantLoader = () => { +const DotLottieReact = lazy(() => + import("@lottiefiles/dotlottie-react").then((m) => ({ default: m.DotLottieReact })) +); + +const AssistantLoaderImpl = () => { const { resolvedTheme } = useTheme(); const isRunning = useAuiState( ({ message }) => (message as { status?: { type: string } })?.status?.type === "running" @@ -21,14 +25,18 @@ export const AssistantLoader = () => { return (
- + }> + + Thinking...
); }; + +export const AssistantLoader = memo(AssistantLoaderImpl); diff --git a/src/components/assistant-ui/markdown-text.tsx b/src/components/assistant-ui/markdown-text.tsx index ba0518b9..0bae69c6 100644 --- a/src/components/assistant-ui/markdown-text.tsx +++ b/src/components/assistant-ui/markdown-text.tsx @@ -5,13 +5,16 @@ 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 { useMessagePartText, type TextMessagePartProps } from "@assistant-ui/react"; import { Children, + createContext, isValidElement, memo, + useContext, useRef, useEffect, + useMemo, type ReactNode, } from "react"; import type { @@ -30,13 +33,23 @@ import { getCitationUrl } from "@/lib/utils/preprocess-latex"; import { resolveItemByPath } from "@/lib/ai/tools/workspace-search-utils"; import { preprocessLatex } from "@/lib/utils/preprocess-latex"; import { cn } from "@/lib/utils"; +import type { Item } from "@/lib/workspace-state/types"; const math = createMathPlugin({ singleDollarTextMath: true }); const code = createCodePlugin({ themes: ["one-dark-pro", "one-dark-pro"], }) as CodeHighlighterPlugin; -/** Parse page number from citation ref (e.g. "Title | quote | p. 5" or "Title | p. 5"). */ +type CitationContextValue = { + workspaceState: Item[]; + navigateToItem: ReturnType; + openWorkspaceItem: (id: string) => void; + setCitationHighlightQuery: (query: { itemId: string; query: string; pageNumber?: number }) => void; + citationUrls: Map; +}; + +const CitationContext = createContext(null); + function parseCitationPage(ref: string): { title: string; quote?: string; pageNumber?: number } { const segments = ref.split(/\s*\|\s*/).map((s) => s.trim()).filter(Boolean); if (segments.length === 0) return { title: "" }; @@ -56,7 +69,6 @@ function parseCitationPage(ref: string): { title: string; quote?: string; pageNu return { title, quote: quote || undefined, pageNumber }; } -/** Extract raw text from citation element children (handles nested elements). */ function extractCitationText(children: ReactNode): string { if (typeof children === "string") return children.trim(); if (children == null) return ""; @@ -74,43 +86,35 @@ function extractCitationText(children: ReactNode): string { .trim(); } -/** - * SurfSense-style citation renderer. - * Content is the ref itself: urlciteN (URL), or Title, or Title|quote (workspace). - */ const CitationRenderer = memo( ({ children }: { children?: ReactNode }) => { const ref = extractCitationText(children); if (!ref) return null; - // URL placeholder (from preprocess): render as direct link - const url = getCitationUrl(ref); + const ctx = useContext(CitationContext); + + const url = ctx ? getCitationUrl(ref, ctx.citationUrls) : undefined; if (url) { return ; } - // Direct URL (in case preprocess didn't run, e.g. edge case) if (ref.startsWith("http://") || ref.startsWith("https://")) { return ; } - // Workspace citation: Title, Title|quote, or Title|quote|p. 5 (with optional page) + if (!ctx) return null; + + const { workspaceState, navigateToItem, openWorkspaceItem, setCitationHighlightQuery } = ctx; + const parsed = parseCitationPage(ref); const { title: parsedTitle, quote: parsedQuote, pageNumber } = parsed; const title = parsedTitle; const quote = parsedQuote; - const workspaceId = useWorkspaceStore((s) => s.currentWorkspaceId); - const { state: workspaceState } = useWorkspaceState(workspaceId); - const navigateToItem = useNavigateToItem(); - const openWorkspaceItem = useUIStore((s) => s.openWorkspaceItem); - // Normalize for matching: trim, lowercase, strip .pdf (model may include it; items often don't) const titleNorm = (s: string) => s.trim().toLowerCase().replace(/\.pdf$/i, ""); - const setCitationHighlightQuery = useUIStore((s) => s.setCitationHighlightQuery); const handleWorkspaceItemClick = () => { if (!title) return; - // Resolve by virtual path first (e.g. "pdfs/Syllabus.pdf") — AI may cite using paths from const items = workspaceState; const byPath = resolveItemByPath(items, title); const item = @@ -122,7 +126,6 @@ const CitationRenderer = memo( titleNorm(i.name) === titleNorm(title) ); if (!item) return; - // Set citation highlight: for PDFs with page, or when we have a quote to search if (quote?.trim() || (pageNumber != null && item.type === "pdf")) { setCitationHighlightQuery({ itemId: item.id, @@ -134,11 +137,10 @@ const CitationRenderer = memo( openWorkspaceItem(item.id); }; - // For path-like refs (e.g. pdfs/Syllabus.pdf), show basename in badge const displayTitle = title.includes("/") ? title.split("/").pop() ?? title : title; const badgeLabel = - displayTitle.slice(0, 20) + (displayTitle.length > 20 ? "…" : "") + - (pageNumber != null ? ` · p.${pageNumber}` : ""); + displayTitle.slice(0, 20) + (displayTitle.length > 20 ? "\u2026" : "") + + (pageNumber != null ? ` \u00b7 p.${pageNumber}` : ""); return ( @@ -163,9 +165,34 @@ const CitationRenderer = memo( ); CitationRenderer.displayName = "CitationRenderer"; -/** Props from assistant-ui when used as Text component, or optional when used directly (e.g. in Reasoning) */ +const MarkdownA = (props: AnchorHTMLAttributes & { node?: any }) => ( + +); + +const MarkdownCitation = (props: any) => ( + {props.children} +); + +const MarkdownOl = ({ children, node, ...props }: HTMLAttributes & { node?: any }) => ( +
    + {children} +
+); + +const MarkdownUl = ({ children, node, ...props }: HTMLAttributes & { node?: any }) => ( +
    + {children} +
+); + +const STREAMDOWN_COMPONENTS = { + a: MarkdownA, + citation: MarkdownCitation, + ol: MarkdownOl, + ul: MarkdownUl, +} as const; + type MarkdownTextProps = Partial & { - /** Use "reasoning" for smoother streaming in reasoning blocks (blurIn, longer duration) */ streamingVariant?: "default" | "reasoning"; }; @@ -173,8 +200,24 @@ const MarkdownTextImpl = (props: MarkdownTextProps) => { const streamingVariant = props.streamingVariant ?? "default"; const { text, status } = useMessagePartText(); - const threadId = useAuiState(({ threads }) => (threads as any)?.mainThreadId); - const messageId = useAuiState(({ message }) => (message as any)?.id); + const workspaceId = useWorkspaceStore((s) => s.currentWorkspaceId); + const { state: workspaceState } = useWorkspaceState(workspaceId); + const navigateToItem = useNavigateToItem(); + const openWorkspaceItem = useUIStore((s) => s.openWorkspaceItem); + const setCitationHighlightQuery = useUIStore((s) => s.setCitationHighlightQuery); + + const { text: processedText, citationUrls } = useMemo( + () => preprocessLatex(text), + [text] + ); + + const citationCtx = useMemo(() => ({ + workspaceState, + navigateToItem, + openWorkspaceItem, + setCitationHighlightQuery, + citationUrls, + }), [workspaceState, navigateToItem, openWorkspaceItem, setCitationHighlightQuery, citationUrls]); const animateConfig = streamingVariant === "reasoning" @@ -183,10 +226,6 @@ const MarkdownTextImpl = (props: MarkdownTextProps) => { const containerRef = useRef(null); - // Combine thread and message ID for unique key per message - const key = `${threadId || 'no-thread'}-${messageId || 'no-message'}`; - - // Set up wheel event handlers for all code blocks to prevent vertical scroll trapping useEffect(() => { const container = containerRef.current; if (!container) return; @@ -197,11 +236,9 @@ const MarkdownTextImpl = (props: MarkdownTextProps) => { if (!preElement) return; - // If user is primarily scrolling vertically (more vertical than horizontal movement) const isVerticalScroll = Math.abs(e.deltaY) > Math.abs(e.deltaX); if (isVerticalScroll) { - // Always let vertical scroll bubble up to parent - don't let code block trap it const scrollParent = preElement.closest('.aui-thread-viewport') as HTMLElement; if (scrollParent) { scrollParent.scrollTop += e.deltaY; @@ -216,9 +253,8 @@ const MarkdownTextImpl = (props: MarkdownTextProps) => { return () => { container.removeEventListener('wheel', handleWheel, { capture: true }); }; - }, [threadId, messageId]); + }, []); - // Ensure copied content keeps rich HTML but strips background/highlight styles const handleCopy = (event: ReactClipboardEvent) => { if (typeof window === "undefined") return; @@ -228,11 +264,9 @@ const MarkdownTextImpl = (props: MarkdownTextProps) => { const range = selection.getRangeAt(0); const fragment = range.cloneContents(); - // Wrap in a temporary container so we can serialize + sanitize const wrapper = document.createElement("div"); wrapper.appendChild(fragment); - // Strip background/highlight styles from copied HTML wrapper.querySelectorAll("*").forEach((el) => { el.style?.removeProperty("background"); el.style?.removeProperty("background-color"); @@ -253,44 +287,29 @@ const MarkdownTextImpl = (props: MarkdownTextProps) => { }; return ( -
- *:first-child]:mt-0 [&>*:last-child]:mb-0" - )} - linkSafety={{ enabled: false }} - plugins={{ code, mermaid, math }} - components={{ - a: (props: AnchorHTMLAttributes & { node?: any }) => ( - - ), - - citation: (props: any) => ( - {props.children} - ), - ol: ({ children, node, ...props }: HTMLAttributes & { node?: any }) => ( -
    - {children} -
- ), - ul: ({ children, node, ...props }: HTMLAttributes & { node?: any }) => ( -
    - {children} -
- ), - }} - mermaid={{ - config: { - theme: 'dark', - }, - }} - > - {preprocessLatex(text)} - -
+ +
+ *:first-child]:mt-0 [&>*:last-child]:mb-0" + )} + linkSafety={{ enabled: false }} + plugins={{ code, mermaid, math }} + components={STREAMDOWN_COMPONENTS} + mermaid={{ + config: { + theme: 'dark', + }, + }} + > + {processedText} + +
+ ); }; diff --git a/src/components/assistant-ui/reasoning.tsx b/src/components/assistant-ui/reasoning.tsx index 052a58dc..f9ca7bee 100644 --- a/src/components/assistant-ui/reasoning.tsx +++ b/src/components/assistant-ui/reasoning.tsx @@ -236,9 +236,9 @@ const ReasoningGroupImpl: ReasoningGroupComponent = ({ }); 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 messages = (thread as unknown as { messages?: Array<{ id?: string }> })?.messages; + if (!messages || messages.length === 0) return false; + return messages[messages.length - 1]?.id === message.id; }); // Subscribe to reasoning text length so we re-run scroll effect on each stream chunk diff --git a/src/components/assistant-ui/thread.tsx b/src/components/assistant-ui/thread.tsx index d78e2b97..2c0143d1 100644 --- a/src/components/assistant-ui/thread.tsx +++ b/src/components/assistant-ui/thread.tsx @@ -400,10 +400,12 @@ const ComposerHoverWrapper: FC = ({ items }) => { null, ); const isThreadEmpty = useAuiState(({ thread }) => thread?.isEmpty ?? true); - const composerText = useAuiState( - (s) => (s as { composer?: { text?: string } })?.composer?.text ?? "", + const hasComposerText = useAuiState( + (s) => { + const text = (s as { composer?: { text?: string } })?.composer?.text; + return Boolean(text?.trim()); + }, ); - const hasComposerText = Boolean(composerText?.trim()); const handleDirectFill = useCallback( (fill: string) => { diff --git a/src/components/assistant-ui/tool-group.tsx b/src/components/assistant-ui/tool-group.tsx index dfa28481..a9d41802 100644 --- a/src/components/assistant-ui/tool-group.tsx +++ b/src/components/assistant-ui/tool-group.tsx @@ -202,9 +202,9 @@ const ToolGroupImpl: FC< }); 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 messages = (thread as unknown as { messages?: Array<{ id?: string }> })?.messages; + if (!messages || messages.length === 0) return false; + return messages[messages.length - 1]?.id === message.id; }); const [isManuallyOpen, setIsManuallyOpen] = useState(isLastMessage); diff --git a/src/components/chat/MessageContextBadges.tsx b/src/components/chat/MessageContextBadges.tsx index 54f2735a..6ee8e57d 100644 --- a/src/components/chat/MessageContextBadges.tsx +++ b/src/components/chat/MessageContextBadges.tsx @@ -17,15 +17,17 @@ type MessageCustomMetadata = { * when viewing user messages in history. */ function MessageContextBadgesImpl() { - const message = useAuiState((s) => s.message); - if (!message || message.role !== "user") return null; + const badgeData = useAuiState((s) => { + const msg = s.message; + if (!msg || msg.role !== "user") return null; + const custom = (msg.metadata as { custom?: MessageCustomMetadata } | undefined)?.custom; + const selections = custom?.replySelections; + return selections && selections.length > 0 ? selections : null; + }); - const custom = (message.metadata as { custom?: MessageCustomMetadata } | undefined)?.custom; - if (!custom) return null; + if (!badgeData) return null; - const { replySelections } = custom; - const hasAny = replySelections && replySelections.length > 0; - if (!hasAny) return null; + const replySelections = badgeData; return (
diff --git a/src/components/ui/streamdown-markdown.tsx b/src/components/ui/streamdown-markdown.tsx index 246b658e..3a07fdf2 100644 --- a/src/components/ui/streamdown-markdown.tsx +++ b/src/components/ui/streamdown-markdown.tsx @@ -57,7 +57,7 @@ const StreamdownMarkdownImpl: React.FC = ({ }, }} > - {preprocessLatex(children)} + {preprocessLatex(children).text}
); diff --git a/src/lib/utils/__tests__/preprocess-latex.test.ts b/src/lib/utils/__tests__/preprocess-latex.test.ts index f9807806..3144050a 100644 --- a/src/lib/utils/__tests__/preprocess-latex.test.ts +++ b/src/lib/utils/__tests__/preprocess-latex.test.ts @@ -4,51 +4,50 @@ import { preprocessLatex } from "@/lib/utils/preprocess-latex"; describe("preprocessLatex - currency and math", () => { describe("currency protection", () => { it("preserves $5 as text (not math)", () => { - const result = preprocessLatex("The cost is $5 today."); + const { text: result } = preprocessLatex("The cost is $5 today."); expect(result).toBe("The cost is $5 today."); }); it("preserves $19.99 as text", () => { - const result = preprocessLatex("Price: $19.99"); + const { text: result } = preprocessLatex("Price: $19.99"); expect(result).toBe("Price: $19.99"); }); it("preserves $1,000.50 with commas", () => { - const result = preprocessLatex("Total: $1,000.50"); + const { text: result } = preprocessLatex("Total: $1,000.50"); expect(result).toBe("Total: $1,000.50"); }); it("preserves multiple currency values in one string", () => { - const result = preprocessLatex("Items: $5, $19.99, and $1,000"); + const { text: result } = preprocessLatex("Items: $5, $19.99, and $1,000"); expect(result).toBe("Items: $5, $19.99, and $1,000"); }); it("preserves $100k and $100M (k/M suffixes)", () => { - const result = preprocessLatex("Budget: $100k revenue $100M target."); + const { text: result } = preprocessLatex("Budget: $100k revenue $100M target."); expect(result).toBe("Budget: $100k revenue $100M target."); }); it("does NOT protect $127$ (math - dollar on both sides)", () => { const input = "The formula $127$ is wrong"; - const result = preprocessLatex(input); - // $127$ is not currency (dollar on both sides), so we leave it unchanged + const { text: result } = preprocessLatex(input); expect(result).toBe(input); }); }); describe("LaTeX delimiter conversion", () => { it("converts \\(...\\) to $...$ (inline math)", () => { - const result = preprocessLatex("Formula: \\(x^2\\) here"); + const { text: result } = preprocessLatex("Formula: \\(x^2\\) here"); expect(result).toBe("Formula: $x^2$ here"); }); it("converts \\[...\\] to $$...$$ (display math)", () => { - const result = preprocessLatex("Display: \\[E=mc^2\\]"); + const { text: result } = preprocessLatex("Display: \\[E=mc^2\\]"); expect(result).toBe("Display: $$E=mc^2$$"); }); it("handles multiline display math", () => { - const result = preprocessLatex("\\[\\int_0^1 x \\, dx\\]"); + const { text: result } = preprocessLatex("\\[\\int_0^1 x \\, dx\\]"); expect(result).toBe("$$\\int_0^1 x \\, dx$$"); }); }); @@ -56,19 +55,19 @@ describe("preprocessLatex - currency and math", () => { describe("code block preservation", () => { it("does not modify content inside code blocks", () => { const md = "```\n$5 and \\(x\\)\n```"; - const result = preprocessLatex(md); + const { text: result } = preprocessLatex(md); expect(result).toBe("```\n$5 and \\(x\\)\n```"); }); it("does not modify inline code", () => { - const result = preprocessLatex("Use `$var` in code"); + const { text: result } = preprocessLatex("Use `$var` in code"); expect(result).toBe("Use `$var` in code"); }); }); describe("mixed content", () => { it("handles currency and math in same string", () => { - const result = preprocessLatex( + const { text: result } = preprocessLatex( "Cost is $19.99. The formula \\(E=mc^2\\) is famous." ); expect(result).toContain("$19.99"); @@ -76,9 +75,27 @@ describe("preprocessLatex - currency and math", () => { }); it("handles currency adjacent to math", () => { - const result = preprocessLatex("Price $5. The formula $x^2$ works."); + const { text: result } = preprocessLatex("Price $5. The formula $x^2$ works."); expect(result).toContain("$5."); expect(result).toContain("$x^2$"); }); }); + + describe("citation URL extraction", () => { + it("returns citationUrls map for URL citations", () => { + const { text, citationUrls } = preprocessLatex( + "See https://example.com for details." + ); + expect(citationUrls.size).toBe(1); + expect(text).toContain("urlcite0"); + expect(citationUrls.get("urlcite0")).toBe("https://example.com"); + }); + + it("returns empty citationUrls for non-URL citations", () => { + const { citationUrls } = preprocessLatex( + "See My Document for details." + ); + expect(citationUrls.size).toBe(0); + }); + }); }); diff --git a/src/lib/utils/preprocess-latex.ts b/src/lib/utils/preprocess-latex.ts index b3119ad3..460df3e1 100644 --- a/src/lib/utils/preprocess-latex.ts +++ b/src/lib/utils/preprocess-latex.ts @@ -12,51 +12,37 @@ * Preserves code blocks (``` and inline `) so their contents are never modified. */ -// Storage for URL citations replaced during preprocess to avoid GFM autolink interference. -// Populated in preprocessCitations, consumed by CitationRenderer. -let _pendingUrlCitations = new Map(); -let _urlCiteIdx = 0; - -/** Get a stored URL by placeholder key (urlcite0, urlcite1, etc.). */ -export function getCitationUrl(placeholder: string): string | undefined { - return _pendingUrlCitations.get(placeholder); +export function getCitationUrl(placeholder: string, citationUrls: Map): string | undefined { + return citationUrls.get(placeholder); } -/** - * Handles citation markup. - * - Model outputs X directly for workspace refs. - * - For https://..., replaces URL with placeholder to avoid GFM autolinks. - * - Fallback: [citation:X] → X for legacy or model slip-ups. - */ -function preprocessCitations(markdown: string): string { - if (!markdown) return markdown; - _pendingUrlCitations = new Map(); - _urlCiteIdx = 0; +function preprocessCitations(markdown: string): { text: string; citationUrls: Map } { + const citationUrls = new Map(); + if (!markdown) return { text: markdown, citationUrls }; + let urlCiteIdx = 0; - // URLs inside : replace with placeholder to avoid GFM autolinks let out = markdown.replace( /(https?:\/\/[^<]+)<\/citation>/g, (_, url: string) => { - const key = `urlcite${_urlCiteIdx++}`; - _pendingUrlCitations.set(key, url.trim()); + const key = `urlcite${urlCiteIdx++}`; + citationUrls.set(key, url.trim()); return `${key}`; } ); - // Fallback: [citation:X] → X (legacy format; URL in brackets also gets placeholder) out = out.replace(/\[citation:\s*([^\]]+)\s*\]/g, (_, content: string) => { const trimmed = content.trim(); if (!trimmed) return ""; const urlMatch = trimmed.match(/^(https?:\/\/\S+)$/); if (urlMatch) { - const key = `urlcite${_urlCiteIdx++}`; - _pendingUrlCitations.set(key, urlMatch[1].trim()); + const key = `urlcite${urlCiteIdx++}`; + citationUrls.set(key, urlMatch[1].trim()); return `${key}`; } return `${trimmed}`; }); - return out; + return { text: out, citationUrls }; } // Currency pattern: $ followed by digits, optional commas/decimals, optional k/M/B — e.g. $5, $19.99, $100k, $100M @@ -72,39 +58,32 @@ export function convertTexDelimitersToDollars(s: string): string { .replace(/\\\[([\s\S]*?)\\\]/g, (_match, math: string) => `$$${math}$$`); } -export function preprocessLatex(markdown: string): string { - if (!markdown) return markdown; +export function preprocessLatex(markdown: string): { text: string; citationUrls: Map } { + if (!markdown) return { text: markdown, citationUrls: new Map() }; - // 0. Convert [citation:X] to X (SurfSense-style inline) - markdown = preprocessCitations(markdown); + const { text: withCitations, citationUrls } = preprocessCitations(markdown); - // 1. Protect code blocks and inline code from modification const preserved: string[] = []; - let result = markdown.replace(/```[\s\S]*?```|`[^`\n]+`/g, (match: string) => { + let result = withCitations.replace(/```[\s\S]*?```|`[^`\n]+`/g, (match: string) => { preserved.push(match); return `\x00CODE${preserved.length - 1}\x00`; }); - // 2. Protect currency values from being parsed as math delimiters - // e.g. "$19.99" → placeholder, restored at the end const currencies: string[] = []; result = result.replace(CURRENCY_REGEX, (match) => { currencies.push(match); return `\x00CUR${currencies.length - 1}\x00`; }); - // 3–4. Convert \(...\) / \[...\] → $...$ / $$...$$ result = convertTexDelimitersToDollars(result); - // 5. Restore protected currency values result = result.replace(/\x00CUR(\d+)\x00/g, (_match, idx) => { return currencies[Number(idx)]; }); - // 6. Restore protected code blocks result = result.replace(/\x00CODE(\d+)\x00/g, (_match, idx) => { return preserved[Number(idx)]; }); - return result; + return { text: result, citationUrls }; } From 1a089c4913c89c177b21f37476378c08af95541b Mon Sep 17 00:00:00 2001 From: urjitc <135136842+urjitc@users.noreply.github.com> Date: Sat, 11 Apr 2026 07:45:35 +0000 Subject: [PATCH 2/5] fix: resolve conditional hook call and stabilize citationUrls Map reference MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fix A: Move useContext(CitationContext) before early return in CitationRenderer to comply with Rules of Hooks. Fix B: Use useRef to stabilize the citationUrls Map reference across streaming renders — only update when map content actually changes, preventing unnecessary CitationContext rebuilds per token. Co-authored-by: capy-ai[bot] <230910855+capy-ai[bot]@users.noreply.github.com> --- src/components/assistant-ui/markdown-text.tsx | 20 ++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/src/components/assistant-ui/markdown-text.tsx b/src/components/assistant-ui/markdown-text.tsx index 0bae69c6..2180ed7d 100644 --- a/src/components/assistant-ui/markdown-text.tsx +++ b/src/components/assistant-ui/markdown-text.tsx @@ -88,11 +88,10 @@ function extractCitationText(children: ReactNode): string { const CitationRenderer = memo( ({ children }: { children?: ReactNode }) => { + const ctx = useContext(CitationContext); const ref = extractCitationText(children); if (!ref) return null; - const ctx = useContext(CitationContext); - const url = ctx ? getCitationUrl(ref, ctx.citationUrls) : undefined; if (url) { return ; @@ -206,11 +205,26 @@ const MarkdownTextImpl = (props: MarkdownTextProps) => { const openWorkspaceItem = useUIStore((s) => s.openWorkspaceItem); const setCitationHighlightQuery = useUIStore((s) => s.setCitationHighlightQuery); - const { text: processedText, citationUrls } = useMemo( + const citationUrlsRef = useRef>(new Map()); + + const { text: processedText, citationUrls: rawCitationUrls } = useMemo( () => preprocessLatex(text), [text] ); + const citationUrls = useMemo(() => { + const prev = citationUrlsRef.current; + if (prev.size === rawCitationUrls.size) { + let same = true; + for (const [k, v] of rawCitationUrls) { + if (prev.get(k) !== v) { same = false; break; } + } + if (same) return prev; + } + citationUrlsRef.current = rawCitationUrls; + return rawCitationUrls; + }, [rawCitationUrls]); + const citationCtx = useMemo(() => ({ workspaceState, navigateToItem, From 05d38101f15483afc5153e47199a8292c61e5060 Mon Sep 17 00:00:00 2001 From: urjitc <135136842+urjitc@users.noreply.github.com> Date: Sat, 11 Apr 2026 19:49:10 +0000 Subject: [PATCH 3/5] fix: make Streamdown mode conditional on streaming status MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Use 'streaming' mode only while message is running, 'static' for completed messages — skips remend processing and startTransition wrapping for finished messages. Co-authored-by: capy-ai[bot] <230910855+capy-ai[bot]@users.noreply.github.com> --- src/components/assistant-ui/markdown-text.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/components/assistant-ui/markdown-text.tsx b/src/components/assistant-ui/markdown-text.tsx index 2180ed7d..ca4fce8b 100644 --- a/src/components/assistant-ui/markdown-text.tsx +++ b/src/components/assistant-ui/markdown-text.tsx @@ -304,7 +304,7 @@ const MarkdownTextImpl = (props: MarkdownTextProps) => {
Date: Sat, 11 Apr 2026 19:59:47 +0000 Subject: [PATCH 4/5] refactor: replace AssistantLoader self-gating with AuiIf, hoist Parts components to module scope Simplification A: AssistantLoader stripped to AssistantLoaderVisual (pure visual, no useAuiState hooks). Condition moved to AuiIf in AssistantMessage so the component never mounts for completed messages. Simplification B: Extract MessagePrimitive.Parts components objects to module-scope constants (ASSISTANT_MESSAGE_PARTS, USER_MESSAGE_PARTS) to avoid creating new object references on every render. Co-authored-by: capy-ai[bot] <230910855+capy-ai[bot]@users.noreply.github.com> --- .../assistant-ui/assistant-loader.tsx | 18 +------ src/components/assistant-ui/thread.tsx | 49 ++++++++++--------- 2 files changed, 29 insertions(+), 38 deletions(-) diff --git a/src/components/assistant-ui/assistant-loader.tsx b/src/components/assistant-ui/assistant-loader.tsx index 19a00353..48736fce 100644 --- a/src/components/assistant-ui/assistant-loader.tsx +++ b/src/components/assistant-ui/assistant-loader.tsx @@ -1,26 +1,14 @@ "use client"; -import { memo, lazy, Suspense } from "react"; -import { useAuiState } from "@assistant-ui/react"; +import { lazy, Suspense } from "react"; import { useTheme } from "next-themes"; const DotLottieReact = lazy(() => import("@lottiefiles/dotlottie-react").then((m) => ({ default: m.DotLottieReact })) ); -const AssistantLoaderImpl = () => { +export const AssistantLoaderVisual = () => { const { resolvedTheme } = useTheme(); - const isRunning = useAuiState( - ({ message }) => (message as { status?: { type: string } })?.status?.type === "running" - ); - - const isMessageEmpty = useAuiState(({ message }) => { - const msg = message as any; - return !msg?.content || (Array.isArray(msg.content) && msg.content.length === 0); - }); - - if (!isRunning || !isMessageEmpty) return null; - const lottieSrc = resolvedTheme === 'light' ? '/thinkexlight.lottie' : '/logo.lottie'; return ( @@ -38,5 +26,3 @@ const AssistantLoaderImpl = () => {
); }; - -export const AssistantLoader = memo(AssistantLoaderImpl); diff --git a/src/components/assistant-ui/thread.tsx b/src/components/assistant-ui/thread.tsx index 2c0143d1..f7db8ee1 100644 --- a/src/components/assistant-ui/thread.tsx +++ b/src/components/assistant-ui/thread.tsx @@ -69,7 +69,7 @@ import { ComposerAddAttachment, UserMessageAttachments, } from "@/components/assistant-ui/attachment"; -import { AssistantLoader } from "@/components/assistant-ui/assistant-loader"; +import { AssistantLoaderVisual } 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"; @@ -1144,6 +1144,19 @@ const MessageError: FC = () => { ); }; +const ASSISTANT_MESSAGE_PARTS = { + Text: MarkdownText, + File: FileComponent, + Source: Sources, + Image, + Reasoning: Reasoning, + ReasoningGroup: ReasoningGroup, + ToolGroup, + tools: { + Fallback: ToolFallback, + }, +} as const; + const AssistantMessage: FC = () => { return ( @@ -1152,21 +1165,13 @@ const AssistantMessage: FC = () => { data-role="assistant" >
- - + + (message as any).status?.type === "running" && + (!(message as any).content || (Array.isArray((message as any).content) && (message as any).content.length === 0)) + }> + + +
@@ -1245,6 +1250,11 @@ const UserMessageText: FC = () => { return
{text}
; }; +const USER_MESSAGE_PARTS = { + Text: UserMessageText, + File: FileComponent, +} as const; + const UserMessage: FC = () => { const [expanded, setExpanded] = useState(false); const message = useMessage(); @@ -1284,12 +1294,7 @@ const UserMessage: FC = () => {
- + {showExpand && (