diff --git a/frontend/src/components/ai-elements/inline-citation.tsx b/frontend/src/components/ai-elements/inline-citation.tsx
index d60ac07..506e428 100644
--- a/frontend/src/components/ai-elements/inline-citation.tsx
+++ b/frontend/src/components/ai-elements/inline-citation.tsx
@@ -29,35 +29,31 @@ import {
 import { Shimmer } from "./shimmer";
 import { useI18n } from "@/core/i18n/hooks";
 
-export type InlineCitationCardProps = ComponentProps<typeof HoverCard>;
+type InlineCitationCardProps = ComponentProps<typeof HoverCard>;
 
-export const InlineCitationCard = (props: InlineCitationCardProps) => (
+const InlineCitationCard = (props: InlineCitationCardProps) => (
   <HoverCard closeDelay={0} openDelay={0} {...props} />
 );
 
-export type InlineCitationCardBodyProps = ComponentProps<"div">;
-
-export const InlineCitationCardBody = ({
+const InlineCitationCardBody = ({
   className,
   ...props
-}: InlineCitationCardBodyProps) => (
+}: ComponentProps<"div">) => (
   <HoverCardContent className={cn("relative w-80 p-0", className)} {...props} />
 );
 
-export type InlineCitationSourceProps = ComponentProps<"div"> & {
-  title?: string;
-  url?: string;
-  description?: string;
-};
-
-export const InlineCitationSource = ({
+const InlineCitationSource = ({
   title,
   url,
   description,
   className,
   children,
   ...props
-}: InlineCitationSourceProps) => (
+}: ComponentProps<"div"> & {
+  title?: string;
+  url?: string;
+  description?: string;
+}) => (
   <div className={cn("space-y-1", className)} {...props}>
     {title && (
       <h4 className="truncate font-medium text-sm leading-tight">{title}</h4>
diff --git a/frontend/src/components/workspace/artifacts/artifact-file-detail.tsx b/frontend/src/components/workspace/artifacts/artifact-file-detail.tsx
index d459e6a..2bfae75 100644
--- a/frontend/src/components/workspace/artifacts/artifact-file-detail.tsx
+++ b/frontend/src/components/workspace/artifacts/artifact-file-detail.tsx
@@ -32,13 +32,11 @@ import {
 import { ToggleGroup, ToggleGroupItem } from "@/components/ui/toggle-group";
 import { CodeEditor } from "@/components/workspace/code-editor";
 import { useArtifactContent } from "@/core/artifacts/hooks";
-import { CitationsLoadingIndicator } from "@/components/ai-elements/inline-citation";
 import { urlOfArtifact } from "@/core/artifacts/utils";
 import type { Citation } from "@/core/citations";
 import {
   contentWithoutCitationsFromParsed,
   removeAllCitations,
-  shouldShowCitationLoading,
   useParsedCitations,
 } from "@/core/citations";
 import { useI18n } from "@/core/i18n/hooks";
@@ -50,6 +48,7 @@ import { cn } from "@/lib/utils";
 
 import { Tooltip } from "../tooltip";
 
+import { SafeCitationContent } from "../messages/safe-citation-content";
 import { useThread } from "../messages/context";
 
 import { useArtifacts } from "./context";
@@ -252,31 +251,27 @@ export function ArtifactFileDetail({
         </div>
       </ArtifactHeader>
       <ArtifactContent className="p-0">
-        {previewable && viewMode === "preview" && (
+        {previewable &&
+          viewMode === "preview" &&
           language === "markdown" &&
-          content &&
-          shouldShowCitationLoading(
-            content,
-            parsed.cleanContent,
-            thread.isLoading,
-          ) ? (
-            <div className="flex size-full items-center justify-center p-4">
-              <CitationsLoadingIndicator
-                citations={parsed.citations}
-                className="my-0"
-              />
-            </div>
-          ) : (
-            <ArtifactFilePreview
-              filepath={filepath}
-              threadId={threadId}
+          content && (
+            <SafeCitationContent
               content={content}
-              language={language ?? "text"}
-              cleanContent={parsed.cleanContent}
-              citationMap={parsed.citationMap}
+              isLoading={thread.isLoading}
+              rehypePlugins={streamdownPlugins.rehypePlugins}
+              className="flex size-full items-center justify-center p-4 my-0"
+              renderBody={(p) => (
+                <ArtifactFilePreview
+                  filepath={filepath}
+                  threadId={threadId}
+                  content={content}
+                  language={language ?? "text"}
+                  cleanContent={p.cleanContent}
+                  citationMap={p.citationMap}
+                />
+              )}
             />
-          )
-        )}
+          )}
         {isCodeFile && viewMode === "code" && (
           <CodeEditor
             className="size-full resize-none rounded-none border-none"
diff --git a/frontend/src/components/workspace/messages/message-group.tsx b/frontend/src/components/workspace/messages/message-group.tsx
index 1aeb7fe..3a2bb94 100644
--- a/frontend/src/components/workspace/messages/message-group.tsx
+++ b/frontend/src/components/workspace/messages/message-group.tsx
@@ -22,10 +22,8 @@ import {
   ChainOfThoughtStep,
 } from "@/components/ai-elements/chain-of-thought";
 import { CodeBlock } from "@/components/ai-elements/code-block";
-import { CitationsLoadingIndicator } from "@/components/ai-elements/inline-citation";
 import { MessageResponse } from "@/components/ai-elements/message";
 import { Button } from "@/components/ui/button";
-import { shouldShowCitationLoading, useParsedCitations } from "@/core/citations";
 import { useI18n } from "@/core/i18n/hooks";
 import {
   extractReasoningContentFromMessage,
@@ -130,7 +128,12 @@ export function MessageGroup({
                   }
                 ></ChainOfThoughtStep>
               ) : (
-                <ToolCall key={step.id} {...step} isLoading={isLoading} />
+                <ToolCall
+                  key={step.id}
+                  {...step}
+                  isLoading={isLoading}
+                  rehypePlugins={rehypePlugins}
+                />
               ),
             )}
           {lastToolCallStep && (
@@ -140,6 +143,7 @@ export function MessageGroup({
                 {...lastToolCallStep}
                 isLast={true}
                 isLoading={isLoading}
+                rehypePlugins={rehypePlugins}
               />
             </FlipDisplay>
           )}
@@ -197,6 +201,7 @@ function ToolCall({
   result,
   isLast = false,
   isLoading = false,
+  rehypePlugins,
 }: {
   id?: string;
   messageId?: string;
@@ -205,6 +210,7 @@ function ToolCall({
   result?: string | Record<string, unknown>;
   isLast?: boolean;
   isLoading?: boolean;
+  rehypePlugins: ReturnType<typeof useRehypeSplitWordsIntoSpans>;
 }) {
   const { t } = useI18n();
   const { setOpen, autoOpen, autoSelect, selectedArtifact, select } =
@@ -213,7 +219,6 @@ function ToolCall({
   const threadIsLoading = thread.isLoading;
 
   const fileContent = typeof args.content === "string" ? args.content : "";
-  const { citations, cleanContent } = useParsedCitations(fileContent);
 
   if (name === "web_search") {
     let label: React.ReactNode = t.toolCalls.searchForRelatedInfo;
@@ -362,13 +367,6 @@ function ToolCall({
     const isMarkdown =
       path?.toLowerCase().endsWith(".md") ||
       path?.toLowerCase().endsWith(".markdown");
-    const showCitationsLoading =
-      isMarkdown &&
-      shouldShowCitationLoading(
-        fileContent,
-        cleanContent,
-        threadIsLoading && isLast,
-      );
 
     return (
       <>
@@ -392,10 +390,14 @@ function ToolCall({
             </ChainOfThoughtSearchResult>
           )}
         </ChainOfThoughtStep>
-        {showCitationsLoading && (
-          <div className="mt-2 ml-8">
-            <CitationsLoadingIndicator citations={citations} />
-          </div>
+        {isMarkdown && (
+          <SafeCitationContent
+            content={fileContent}
+            isLoading={threadIsLoading && isLast}
+            rehypePlugins={rehypePlugins}
+            loadingOnly
+            className="mt-2 ml-8"
+          />
         )}
       </>
     );
diff --git a/frontend/src/components/workspace/messages/safe-citation-content.tsx b/frontend/src/components/workspace/messages/safe-citation-content.tsx
index 36d8fec..e453b89 100644
--- a/frontend/src/components/workspace/messages/safe-citation-content.tsx
+++ b/frontend/src/components/workspace/messages/safe-citation-content.tsx
@@ -1,6 +1,7 @@
 "use client";
 
 import type { ImgHTMLAttributes } from "react";
+import type { ReactNode } from "react";
 import { useMemo } from "react";
 
 import {
@@ -14,6 +15,7 @@ import {
 import {
   shouldShowCitationLoading,
   useParsedCitations,
+  type UseParsedCitationsResult,
 } from "@/core/citations";
 import { streamdownPlugins } from "@/core/streamdown";
 import { cn } from "@/lib/utils";
@@ -25,7 +27,11 @@ export type SafeCitationContentProps = {
   className?: string;
   remarkPlugins?: MessageResponseProps["remarkPlugins"];
   isHuman?: boolean;
-  img?: (props: ImgHTMLAttributes<HTMLImageElement> & { threadId?: string; maxWidth?: string }) => React.ReactNode;
+  img?: (props: ImgHTMLAttributes<HTMLImageElement> & { threadId?: string; maxWidth?: string }) => ReactNode;
+  /** When true, only show loading indicator or null (e.g. write_file step). */
+  loadingOnly?: boolean;
+  /** When set, use instead of default MessageResponse (e.g. artifact preview). */
+  renderBody?: (parsed: UseParsedCitationsResult) => ReactNode;
 };
 
 /** Single place for citation-aware body: shows loading until citations complete (no half-finished refs), else body. */
@@ -37,20 +43,12 @@ export function SafeCitationContent({
   remarkPlugins = streamdownPlugins.remarkPlugins,
   isHuman = false,
   img,
+  loadingOnly = false,
+  renderBody,
 }: SafeCitationContentProps) {
-  const { citations, cleanContent, citationMap } = useParsedCitations(content);
+  const parsed = useParsedCitations(content);
+  const { citations, cleanContent, citationMap } = parsed;
   const showLoading = shouldShowCitationLoading(content, cleanContent, isLoading);
-
-  if (showLoading) {
-    return (
-      <CitationsLoadingIndicator
-        citations={citations}
-        className={cn("my-2", className)}
-      />
-    );
-  }
-  if (!cleanContent) return null;
-
   const components = useMemo(
     () =>
       createCitationMarkdownComponents({
@@ -61,6 +59,19 @@ export function SafeCitationContent({
       }),
     [citationMap, isHuman, img],
   );
+
+  if (showLoading) {
+    return (
+      <CitationsLoadingIndicator
+        citations={citations}
+        className={cn("my-2", className)}
+      />
+    );
+  }
+  if (loadingOnly) return null;
+  if (renderBody) return renderBody(parsed);
+  if (!cleanContent) return null;
+
   return (
     <MessageResponse
       className={className}
diff --git a/frontend/src/core/citations/index.ts b/frontend/src/core/citations/index.ts
index 65b01a4..ac41316 100644
--- a/frontend/src/core/citations/index.ts
+++ b/frontend/src/core/citations/index.ts
@@ -1,7 +1,6 @@
 export {
   contentWithoutCitationsFromParsed,
   extractDomainFromUrl,
-  hasCitationsBlock,
   isExternalUrl,
   parseCitations,
   removeAllCitations,
diff --git a/frontend/src/core/citations/use-parsed-citations.ts b/frontend/src/core/citations/use-parsed-citations.ts
index 1dc8463..7be76de 100644
--- a/frontend/src/core/citations/use-parsed-citations.ts
+++ b/frontend/src/core/citations/use-parsed-citations.ts
@@ -2,7 +2,7 @@
 
 import { useMemo } from "react";
 
-import { buildCitationMap, parseCitations } from "./utils";
+import { parseCitations } from "./utils";
 import type { Citation } from "./utils";
 
 export interface UseParsedCitationsResult {
@@ -13,12 +13,12 @@ export interface UseParsedCitationsResult {
 
 /**
  * Parse content for citations and build citation map. Memoized by content.
- * Use in message and artifact components to avoid repeating parseCitations + buildCitationMap.
  */
 export function useParsedCitations(content: string): UseParsedCitationsResult {
   return useMemo(() => {
     const parsed = parseCitations(content ?? "");
-    const citationMap = buildCitationMap(parsed.citations);
+    const citationMap = new Map<string, Citation>();
+    for (const c of parsed.citations) citationMap.set(c.url, c);
     return {
       citations: parsed.citations,
       cleanContent: parsed.cleanContent,
diff --git a/frontend/src/core/citations/utils.ts b/frontend/src/core/citations/utils.ts
index bcca859..fb5cded 100644
--- a/frontend/src/core/citations/utils.ts
+++ b/frontend/src/core/citations/utils.ts
@@ -120,22 +120,6 @@ export function parseCitations(content: string): ParseCitationsResult {
   return { citations, cleanContent };
 }
 
-/**
- * Build a map from URL to Citation for quick lookup
- *
- * @param citations - Array of citations
- * @returns Map with URL as key and Citation as value
- */
-export function buildCitationMap(
-  citations: Citation[],
-): Map<string, Citation> {
-  const map = new Map<string, Citation>();
-  for (const citation of citations) {
-    map.set(citation.url, citation);
-  }
-  return map;
-}
-
 /**
  * Whether the URL is external (http/https).
  */