Skip to content

Fix/assistant-chat-performance-lag #347

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
urjitc wants to merge 5 commits into from
Closed

Fix/assistant-chat-performance-lag #347

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
27f0136
Fix assistant chat performance by eliminating unnecessary re-renders,…
urjitc Apr 11, 2026
1a089c4
fix: resolve conditional hook call and stabilize citationUrls Map ref…
urjitc Apr 11, 2026
05d3810
fix: make Streamdown mode conditional on streaming status
urjitc Apr 11, 2026
944ea7f
refactor: replace AssistantLoader self-gating with AuiIf, hoist Parts…
urjitc Apr 11, 2026
6428c2c
fix: guard isLastMessage against undefined IDs, protect code blocks b…
urjitc Apr 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 15 additions & 21 deletions src/components/assistant-ui/assistant-loader.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,33 +1,27 @@
"use client";

import { useAuiState } from "@assistant-ui/react";
import { DotLottieReact } from "@lottiefiles/dotlottie-react";
import { lazy, Suspense } from "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 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 DotLottieReact = lazy(() =>
import("@lottiefiles/dotlottie-react").then((m) => ({ default: m.DotLottieReact }))
);

export const AssistantLoaderVisual = () => {
const { resolvedTheme } = useTheme();
const lottieSrc = resolvedTheme === 'light' ? '/thinkexlight.lottie' : '/logo.lottie';

return (
<div className="flex items-center gap-3 py-2">
<DotLottieReact
src={lottieSrc}
loop
autoplay
mode="bounce"
className="w-4 h-4 self-center"
/>
<Suspense fallback={<div className="w-4 h-4" />}>
<DotLottieReact
src={lottieSrc}
loop
autoplay
mode="bounce"
className="w-4 h-4 self-center"
/>
</Suspense>
<span className="text-base text-muted-foreground">Thinking...</span>
</div>
);
Expand Down
181 changes: 107 additions & 74 deletions src/components/assistant-ui/markdown-text.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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<typeof useNavigateToItem>;
openWorkspaceItem: (id: string) => void;
setCitationHighlightQuery: (query: { itemId: string; query: string; pageNumber?: number }) => void;
citationUrls: Map<string, string>;
};

const CitationContext = createContext<CitationContextValue | null>(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: "" };
Expand All @@ -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 "";
Expand All @@ -74,43 +86,34 @@ 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 ctx = useContext(CitationContext);
const ref = extractCitationText(children);
if (!ref) return null;

// URL placeholder (from preprocess): render as direct link
const url = getCitationUrl(ref);
const url = ctx ? getCitationUrl(ref, ctx.citationUrls) : undefined;
if (url) {
return <UrlCitation url={url} />;
}

// Direct URL (in case preprocess didn't run, e.g. edge case)
if (ref.startsWith("http://") || ref.startsWith("https://")) {
return <UrlCitation url={ref} />;
}

// 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 <workspace>
const items = workspaceState;
const byPath = resolveItemByPath(items, title);
const item =
Expand All @@ -122,7 +125,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,
Expand All @@ -134,11 +136,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 (
<InlineCitation>
Expand All @@ -163,18 +164,74 @@ 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<HTMLAnchorElement> & { node?: any }) => (
<MarkdownLink {...props} />
);

const MarkdownCitation = (props: any) => (
<CitationRenderer>{props.children}</CitationRenderer>
);

const MarkdownOl = ({ children, node, ...props }: HTMLAttributes<HTMLOListElement> & { node?: any }) => (
<ol className="ml-4 list-outside list-decimal whitespace-normal" {...props}>
{children}
</ol>
);

const MarkdownUl = ({ children, node, ...props }: HTMLAttributes<HTMLUListElement> & { node?: any }) => (
<ul className="ml-4 list-outside list-disc whitespace-normal" {...props}>
{children}
</ul>
);

const STREAMDOWN_COMPONENTS = {
a: MarkdownA,
citation: MarkdownCitation,
ol: MarkdownOl,
ul: MarkdownUl,
} as const;

type MarkdownTextProps = Partial<TextMessagePartProps> & {
/** 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 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 citationUrlsRef = useRef<Map<string, string>>(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,
openWorkspaceItem,
setCitationHighlightQuery,
citationUrls,
}), [workspaceState, navigateToItem, openWorkspaceItem, setCitationHighlightQuery, citationUrls]);
Comment on lines +228 to +234

@greptile-apps greptile-apps Bot Apr 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 citationUrls Map reference changes on every text update, defeating context stability

preprocessLatex returns a new Map on every call, so citationUrls is always a new reference. Because it sits in the citationCtx useMemo dependency array, the context value is recreated on every streaming token — even when no citation URLs actually changed — re-rendering every CitationRenderer consumer. Stabilising by comparing map content (or storing citationUrls in a useRef and only updating it when entries differ) would prevent unnecessary propagation during mid-stream tokens.

Fix in Cursor


const animateConfig =
streamingVariant === "reasoning"
Expand All @@ -183,10 +240,6 @@ const MarkdownTextImpl = (props: MarkdownTextProps) => {

const containerRef = useRef<HTMLDivElement>(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;
Expand All @@ -197,11 +250,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;
Expand All @@ -216,9 +267,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<HTMLDivElement>) => {
if (typeof window === "undefined") return;

Expand All @@ -228,11 +278,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<HTMLElement>("*").forEach((el) => {
el.style?.removeProperty("background");
el.style?.removeProperty("background-color");
Expand All @@ -253,44 +301,29 @@ const MarkdownTextImpl = (props: MarkdownTextProps) => {
};

return (
<div key={key} ref={containerRef} className="aui-md" onCopy={handleCopy}>
<Streamdown
allowedTags={{ citation: [] }}
animated={animateConfig}
isAnimating={status.type === "running"}
className={cn(
"streamdown-content size-full [&>*:first-child]:mt-0 [&>*:last-child]:mb-0"
)}
linkSafety={{ enabled: false }}
plugins={{ code, mermaid, math }}
components={{
a: (props: AnchorHTMLAttributes<HTMLAnchorElement> & { node?: any }) => (
<MarkdownLink {...props} />
),

citation: (props: any) => (
<CitationRenderer>{props.children}</CitationRenderer>
),
ol: ({ children, node, ...props }: HTMLAttributes<HTMLOListElement> & { node?: any }) => (
<ol className="ml-4 list-outside list-decimal whitespace-normal" {...props}>
{children}
</ol>
),
ul: ({ children, node, ...props }: HTMLAttributes<HTMLUListElement> & { node?: any }) => (
<ul className="ml-4 list-outside list-disc whitespace-normal" {...props}>
{children}
</ul>
),
}}
mermaid={{
config: {
theme: 'dark',
},
}}
>
{preprocessLatex(text)}
</Streamdown>
</div>
<CitationContext.Provider value={citationCtx}>
<div ref={containerRef} className="aui-md" onCopy={handleCopy}>
<Streamdown
mode={status.type === "running" ? "streaming" : "static"}
allowedTags={{ citation: [] }}
animated={animateConfig}
isAnimating={status.type === "running"}
className={cn(
"streamdown-content size-full [&>*:first-child]:mt-0 [&>*:last-child]:mb-0"
)}
linkSafety={{ enabled: false }}
plugins={{ code, mermaid, math }}
components={STREAMDOWN_COMPONENTS}
mermaid={{
config: {
theme: 'dark',
},
}}
>
{processedText}
</Streamdown>
</div>
</CitationContext.Provider>
);
};

Expand Down
8 changes: 5 additions & 3 deletions src/components/assistant-ui/reasoning.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -236,9 +236,11 @@ 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;
const lastId = messages[messages.length - 1]?.id;
if (!lastId || !message.id) return false;
return lastId === message.id;
});

// Subscribe to reasoning text length so we re-run scroll effect on each stream chunk
Expand Down
Loading
Loading