diff --git a/src/components/assistant-ui/assistant-loader.tsx b/src/components/assistant-ui/assistant-loader.tsx
index 0ae83746..48736fce 100644
--- a/src/components/assistant-ui/assistant-loader.tsx
+++ b/src/components/assistant-ui/assistant-loader.tsx
@@ -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>
     );
diff --git a/src/components/assistant-ui/markdown-text.tsx b/src/components/assistant-ui/markdown-text.tsx
index ba0518b9..ca4fce8b 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<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: "" };
@@ -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,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 =
@@ -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,
@@ -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>
@@ -163,9 +164,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<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";
 };
 
@@ -173,8 +199,39 @@ 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]);
 
   const animateConfig =
     streamingVariant === "reasoning"
@@ -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;
@@ -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;
@@ -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;
 
@@ -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");
@@ -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>
   );
 };
 
diff --git a/src/components/assistant-ui/reasoning.tsx b/src/components/assistant-ui/reasoning.tsx
index 052a58dc..049aa67e 100644
--- a/src/components/assistant-ui/reasoning.tsx
+++ b/src/components/assistant-ui/reasoning.tsx
@@ -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
diff --git a/src/components/assistant-ui/thread.tsx b/src/components/assistant-ui/thread.tsx
index d78e2b97..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";
@@ -400,10 +400,12 @@ const ComposerHoverWrapper: FC<ComposerHoverWrapperProps> = ({ 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) => {
@@ -1142,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 (
     <MessagePrimitive.Root asChild>
@@ -1150,21 +1165,13 @@ const AssistantMessage: FC = () => {
         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,
-              },
-            }}
-          />
+          <AuiIf condition={({ message }) =>
+            (message as any).status?.type === "running" &&
+            (!(message as any).content || (Array.isArray((message as any).content) && (message as any).content.length === 0))
+          }>
+            <AssistantLoaderVisual />
+          </AuiIf>
+          <MessagePrimitive.Parts components={ASSISTANT_MESSAGE_PARTS} />
           <MessageError />
         </div>
 
@@ -1243,6 +1250,11 @@ const UserMessageText: FC = () => {
   return <div className="whitespace-pre-wrap">{text}</div>;
 };
 
+const USER_MESSAGE_PARTS = {
+  Text: UserMessageText,
+  File: FileComponent,
+} as const;
+
 const UserMessage: FC = () => {
   const [expanded, setExpanded] = useState(false);
   const message = useMessage();
@@ -1282,12 +1294,7 @@ const UserMessage: FC = () => {
           <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,
-                }}
-              />
+              <MessagePrimitive.Parts components={USER_MESSAGE_PARTS} />
               {showExpand && (
                 <div className="flex justify-end pt-1.5 mt-1.5">
                   <button
diff --git a/src/components/assistant-ui/tool-group.tsx b/src/components/assistant-ui/tool-group.tsx
index dfa28481..1b7b8987 100644
--- a/src/components/assistant-ui/tool-group.tsx
+++ b/src/components/assistant-ui/tool-group.tsx
@@ -202,9 +202,11 @@ 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;
+    const lastId = messages[messages.length - 1]?.id;
+    if (!lastId || !message.id) return false;
+    return lastId === 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 (
     <div className="mb-2 flex flex-wrap gap-1.5">
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<StreamdownMarkdownProps> = ({
           },
         }}
       >
-        {preprocessLatex(children)}
+        {preprocessLatex(children).text}
       </Streamdown>
     </div>
   );
diff --git a/src/lib/utils/__tests__/preprocess-latex.test.ts b/src/lib/utils/__tests__/preprocess-latex.test.ts
index f9807806..41b1accb 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,45 @@ 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 <citation>https://example.com</citation> for details."
+      );
+      expect(citationUrls.size).toBe(1);
+      expect(text).toContain("<citation>urlcite0</citation>");
+      expect(citationUrls.get("urlcite0")).toBe("https://example.com");
+    });
+
+    it("returns empty citationUrls for non-URL citations", () => {
+      const { citationUrls } = preprocessLatex(
+        "See <citation>My Document</citation> for details."
+      );
+      expect(citationUrls.size).toBe(0);
+    });
+  });
+
+  describe("code block protection for citations", () => {
+    it("does not rewrite citations inside fenced code blocks", () => {
+      const { text, citationUrls } = preprocessLatex(
+        "```\n<citation>https://example.com</citation>\n```"
+      );
+      expect(text).toBe("```\n<citation>https://example.com</citation>\n```");
+      expect(citationUrls.size).toBe(0);
+    });
+
+    it("does not rewrite citations inside inline code", () => {
+      const { text, citationUrls } = preprocessLatex(
+        "Use `<citation>test</citation>` in code"
+      );
+      expect(text).toBe("Use `<citation>test</citation>` in code");
+      expect(citationUrls.size).toBe(0);
+    });
+  });
 });
diff --git a/src/lib/utils/preprocess-latex.ts b/src/lib/utils/preprocess-latex.ts
index b3119ad3..e433c797 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<string, string>();
-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, string>): string | undefined {
+  return citationUrls.get(placeholder);
 }
 
-/**
- * Handles citation markup.
- * - Model outputs <citation>X</citation> directly for workspace refs.
- * - For <citation>https://...</citation>, replaces URL with placeholder to avoid GFM autolinks.
- * - Fallback: [citation:X] → <citation>X</citation> 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<string, string> } {
+  const citationUrls = new Map<string, string>();
+  if (!markdown) return { text: markdown, citationUrls };
+  let urlCiteIdx = 0;
 
-  // URLs inside <citation>: replace with placeholder to avoid GFM autolinks
   let out = markdown.replace(
     /<citation>(https?:\/\/[^<]+)<\/citation>/g,
     (_, url: string) => {
-      const key = `urlcite${_urlCiteIdx++}`;
-      _pendingUrlCitations.set(key, url.trim());
+      const key = `urlcite${urlCiteIdx++}`;
+      citationUrls.set(key, url.trim());
       return `<citation>${key}</citation>`;
     }
   );
 
-  // Fallback: [citation:X] → <citation>X</citation> (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 `<citation>${key}</citation>`;
     }
     return `<citation>${trimmed}</citation>`;
   });
 
-  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,33 @@ 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<string, string> } {
+  if (!markdown) return { text: markdown, citationUrls: new Map() };
 
-  // 0. Convert [citation:X] to <citation>X</citation> (SurfSense-style inline)
-  markdown = preprocessCitations(markdown);
-
-  // 1. Protect code blocks and inline code from modification
   const preserved: string[] = [];
   let result = markdown.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 { text: withCitations, citationUrls } = preprocessCitations(result);
+  result = withCitations;
+
   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 };
 }