feat(frontend): unify citation logic and prevent half-finished citations

- Add SafeCitationContent as single component for citation-aware body:
  useParsedCitations + shouldShowCitationLoading; show loading until
  citations complete, then render body with createCitationMarkdownComponents.
  Supports optional remarkPlugins, rehypePlugins, isHuman, img.

- Refactor MessageListItem: assistant message body now uses
  SafeCitationContent only; remove duplicate useParsedCitations,
  shouldShowCitationLoading, createCitationMarkdownComponents and
  CitationsLoadingIndicator logic. Human messages keep plain
  AIElementMessageResponse (no citation parsing).

- Use SafeCitationContent for clarification, present-files (message-list),
  thinking steps and write_file loading (message-group), subtask result
  (subtask-card). Artifact markdown preview keeps same guard
  (shouldShowCitationLoading) with ArtifactFilePreview.

- Unify loading condition: shouldShowCitationLoading(rawContent,
  cleanContent, isLoading) is the single source of truth. Show loading when
  (isLoading && hasCitationsBlock(rawContent)) or when
  (hasCitationsBlock(rawContent) && hasUnreplacedCitationRefs(cleanContent))
  so Pro/Ultra modes also show "loading citations" and half-finished
  [cite-N] never appear.

- message-group write_file: replace hasCitationsBlock + threadIsLoading
  with shouldShowCitationLoading(fileContent, cleanContent,
  threadIsLoading && isLast) for consistency.

- citations/utils: parse incomplete <citations> during streaming;
  remove isCitationsBlockIncomplete; keep hasUnreplacedCitationRefs
  internal; document display rule in file header.

Co-authored-by: Cursor <cursoragent@cursor.com>

---
feat(前端): 统一引用逻辑并杜绝半成品引用

- 新增 SafeCitationContent 作为引用正文的唯一出口：内部使用
  useParsedCitations + shouldShowCitationLoading，在引用未就绪时只显示
  「正在整理引用」，就绪后用 createCitationMarkdownComponents 渲染正文；
  支持可选 remarkPlugins、rehypePlugins、isHuman、img。

- 重构 MessageListItem：助手消息正文仅通过 SafeCitationContent 渲染，
  删除重复的 useParsedCitations、shouldShowCitationLoading、
  createCitationMarkdownComponents、CitationsLoadingIndicator 等逻辑；
  用户消息仍用 AIElementMessageResponse，不做引用解析。

- 澄清、present-files（message-list）、思考步骤与 write_file 加载
  （message-group）、子任务结果（subtask-card）均使用
  SafeCitationContent；Artifact 的 markdown 预览仍用同一 guard
  shouldShowCitationLoading，正文由 ArtifactFilePreview 渲染。

- 统一加载条件：shouldShowCitationLoading(rawContent, cleanContent,
  isLoading) 为唯一判断。在「流式中且已有引用块」或「有引用块且
  cleanContent 中仍有未替换的 [cite-N]」时仅显示加载，从而在 Pro/Ultra
  下也能看到「正在整理引用」，且永不出现半成品 [cite-N]。

- message-group 的 write_file：用 shouldShowCitationLoading(
  fileContent, cleanContent, threadIsLoading && isLast) 替代
  hasCitationsBlock + threadIsLoading，与其他场景一致。

- citations/utils：流式时解析未闭合的 <citations>；移除
  isCitationsBlockIncomplete；hasUnreplacedCitationRefs 保持内部使用；
  在文件头注释中说明展示规则。

frontend/src/components/ai-elements/inline-citation.tsx

@@ -12,7 +12,14 @@ import {
   externalLinkClassNoUnderline,
 } from "@/lib/utils";
 import { ExternalLinkIcon } from "lucide-react";
-import { type ComponentProps, Children } from "react";
+import {
+  type AnchorHTMLAttributes,
+  type ComponentProps,
+  type ImgHTMLAttributes,
+  type ReactElement,
+  type ReactNode,
+  Children,
+} from "react";
 import type { Citation } from "@/core/citations";
 import {
   extractDomainFromUrl,
@@ -202,6 +209,48 @@ export const CitationAwareLink = ({
   );
 };
 
+/**
+ * Options for creating markdown components that render links as citations.
+ * Used by message list (all modes: Flash/Thinking/Pro/Ultra), artifact preview, and CoT.
+ */
+export type CreateCitationMarkdownComponentsOptions = {
+  citationMap: Map<string, Citation>;
+  isHuman?: boolean;
+  isLoadingCitations?: boolean;
+  syntheticExternal?: boolean;
+  /** Optional custom img component (e.g. MessageImage with threadId). Omit for artifact. */
+  img?: (props: ImgHTMLAttributes<HTMLImageElement> & { threadId?: string; maxWidth?: string }) => ReactNode;
+};
+
+/**
+ * Create markdown `components` (a, optional img) that use CitationAwareLink.
+ * Reused across message-list-item (all modes), artifact-file-detail, and any CoT markdown.
+ */
+export function createCitationMarkdownComponents(
+  options: CreateCitationMarkdownComponentsOptions,
+): {
+  a: (props: AnchorHTMLAttributes<HTMLAnchorElement>) => ReactElement;
+  img?: (props: ImgHTMLAttributes<HTMLImageElement> & { threadId?: string; maxWidth?: string }) => ReactNode;
+} {
+  const {
+    citationMap,
+    isHuman = false,
+    isLoadingCitations = false,
+    syntheticExternal = false,
+    img,
+  } = options;
+  const a = (props: AnchorHTMLAttributes<HTMLAnchorElement>) => (
+    <CitationAwareLink
+      {...props}
+      citationMap={citationMap}
+      isHuman={isHuman}
+      isLoadingCitations={isLoadingCitations}
+      syntheticExternal={syntheticExternal}
+    />
+  );
+  return img ? { a, img } : { a };
+}
+
 /**
  * Shared CitationsLoadingIndicator component
  * Used across message-list-item and message-group to show loading citations

frontend/src/components/workspace/artifacts/artifact-file-detail.tsx

@@ -21,7 +21,7 @@ import {
   ArtifactHeader,
   ArtifactTitle,
 } from "@/components/ai-elements/artifact";
-import { CitationAwareLink } from "@/components/ai-elements/inline-citation";
+import { createCitationMarkdownComponents } from "@/components/ai-elements/inline-citation";
 import { Select, SelectItem } from "@/components/ui/select";
 import {
   SelectContent,
@@ -32,11 +32,13 @@ 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";
@@ -48,6 +50,8 @@ import { cn } from "@/lib/utils";
 
 import { Tooltip } from "../tooltip";
 
+import { useThread } from "../messages/context";
+
 import { useArtifacts } from "./context";
 
 export function ArtifactFileDetail({
@@ -89,6 +93,7 @@ export function ArtifactFileDetail({
   const previewable = useMemo(() => {
     return (language === "html" && !isWriteFile) || language === "markdown";
   }, [isWriteFile, language]);
+  const { thread } = useThread();
   const { content } = useArtifactContent({
     threadId,
     filepath: filepathFromProps,
@@ -248,14 +253,29 @@ export function ArtifactFileDetail({
       </ArtifactHeader>
       <ArtifactContent className="p-0">
         {previewable && viewMode === "preview" && (
-          <ArtifactFilePreview
+          language === "markdown" &&
-            filepath={filepath}
+          content &&
-            threadId={threadId}
+          shouldShowCitationLoading(
-            content={content}
+            content,
-            language={language ?? "text"}
+            parsed.cleanContent,
-            cleanContent={parsed.cleanContent}
+            thread.isLoading,
-            citationMap={parsed.citationMap}
+          ) ? (
-          />
+            <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={content}
+              language={language ?? "text"}
+              cleanContent={parsed.cleanContent}
+              citationMap={parsed.citationMap}
+            />
+          )
         )}
         {isCodeFile && viewMode === "code" && (
           <CodeEditor
@@ -291,20 +311,16 @@ export function ArtifactFilePreview({
   citationMap: Map<string, Citation>;
 }) {
   if (language === "markdown") {
+    const components = createCitationMarkdownComponents({
+      citationMap,
+      syntheticExternal: true,
+    });
     return (
       <div className="size-full px-4">
         <Streamdown
           className="size-full"
           {...streamdownPlugins}
-          components={{
+          components={components}
-            a: (props: React.AnchorHTMLAttributes<HTMLAnchorElement>) => (
-              <CitationAwareLink
-                {...props}
-                citationMap={citationMap}
-                syntheticExternal
-              />
-            ),
-          }}
         >
           {cleanContent ?? ""}
         </Streamdown>

frontend/src/components/workspace/messages/message-group.tsx

@@ -25,11 +25,7 @@ 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 {
+import { shouldShowCitationLoading, useParsedCitations } from "@/core/citations";
-  getCleanContent,
-  hasCitationsBlock,
-  useParsedCitations,
-} from "@/core/citations";
 import { useI18n } from "@/core/i18n/hooks";
 import {
   extractReasoningContentFromMessage,
@@ -47,6 +43,8 @@ import { Tooltip } from "../tooltip";
 
 import { useThread } from "./context";
 
+import { SafeCitationContent } from "./safe-citation-content";
+
 export function MessageGroup({
   className,
   messages,
@@ -124,12 +122,11 @@ export function MessageGroup({
                 <ChainOfThoughtStep
                   key={step.id}
                   label={
-                    <MessageResponse
+                    <SafeCitationContent
-                      remarkPlugins={streamdownPlugins.remarkPlugins}
+                      content={step.reasoning ?? ""}
+                      isLoading={isLoading}
                       rehypePlugins={rehypePlugins}
-                    >
+                    />
-                      {getCleanContent(step.reasoning ?? "")}
-                    </MessageResponse>
                   }
                 ></ChainOfThoughtStep>
               ) : (
@@ -177,12 +174,11 @@ export function MessageGroup({
               <ChainOfThoughtStep
                 key={lastReasoningStep.id}
                 label={
-                  <MessageResponse
+                  <SafeCitationContent
-                    remarkPlugins={streamdownPlugins.remarkPlugins}
+                    content={lastReasoningStep.reasoning ?? ""}
+                    isLoading={isLoading}
                     rehypePlugins={rehypePlugins}
-                  >
+                  />
-                    {getCleanContent(lastReasoningStep.reasoning ?? "")}
-                  </MessageResponse>
                 }
               ></ChainOfThoughtStep>
             </ChainOfThoughtContent>
@@ -217,7 +213,7 @@ function ToolCall({
   const threadIsLoading = thread.isLoading;
 
   const fileContent = typeof args.content === "string" ? args.content : "";
-  const { citations } = useParsedCitations(fileContent);
+  const { citations, cleanContent } = useParsedCitations(fileContent);
 
   if (name === "web_search") {
     let label: React.ReactNode = t.toolCalls.searchForRelatedInfo;
@@ -363,12 +359,16 @@ function ToolCall({
       }, 100);
     }
 
-    // Check if this is a markdown file with citations
     const isMarkdown =
       path?.toLowerCase().endsWith(".md") ||
       path?.toLowerCase().endsWith(".markdown");
     const showCitationsLoading =
-      isMarkdown && threadIsLoading && hasCitationsBlock(fileContent) && isLast;
+      isMarkdown &&
+      shouldShowCitationLoading(
+        fileContent,
+        cleanContent,
+        threadIsLoading && isLast,
+      );
 
     return (
       <>

frontend/src/components/workspace/messages/message-list-item.tsx

@@ -4,10 +4,6 @@ import { useParams } from "next/navigation";
 import { memo, useMemo } from "react";
 import rehypeKatex from "rehype-katex";
 
-import {
-  CitationAwareLink,
-  CitationsLoadingIndicator,
-} from "@/components/ai-elements/inline-citation";
 import {
   Message as AIElementMessage,
   MessageContent as AIElementMessageContent,
@@ -16,11 +12,7 @@ import {
 } from "@/components/ai-elements/message";
 import { Badge } from "@/components/ui/badge";
 import { resolveArtifactURL } from "@/core/artifacts/utils";
-import {
+import { removeAllCitations } from "@/core/citations";
-  isCitationsBlockIncomplete,
-  removeAllCitations,
-  useParsedCitations,
-} from "@/core/citations";
 import {
   extractContentFromMessage,
   extractReasoningContentFromMessage,
@@ -28,10 +20,11 @@ import {
   type UploadedFile,
 } from "@/core/messages/utils";
 import { useRehypeSplitWordsIntoSpans } from "@/core/rehype";
-import { humanMessagePlugins, streamdownPlugins } from "@/core/streamdown";
+import { humanMessagePlugins } from "@/core/streamdown";
 import { cn } from "@/lib/utils";
 
 import { CopyButton } from "../copy-button";
+import { SafeCitationContent } from "./safe-citation-content";
 
 export function MessageListItem({
   className,
@@ -116,79 +109,36 @@ function MessageContent_({
   const isHuman = message.type === "human";
   const { thread_id } = useParams<{ thread_id: string }>();
 
-  // Content to parse for citations (and optionally uploaded files)
+  const rawContent = extractContentFromMessage(message);
-  const { contentToParse, uploadedFiles, isLoadingCitations } = useMemo(() => {
+  const reasoningContent = extractReasoningContentFromMessage(message);
-    const reasoningContent = extractReasoningContentFromMessage(message);
+  const { contentToParse, uploadedFiles } = useMemo(() => {
-    const rawContent = extractContentFromMessage(message);
-
     if (!isLoading && reasoningContent && !rawContent) {
-      return {
+      return { contentToParse: reasoningContent, uploadedFiles: [] as UploadedFile[] };
-        contentToParse: reasoningContent,
-        uploadedFiles: [] as UploadedFile[],
-        isLoadingCitations: false,
-      };
     }
-
     if (isHuman && rawContent) {
       const { files, cleanContent: contentWithoutFiles } =
         parseUploadedFiles(rawContent);
-      return {
+      return { contentToParse: contentWithoutFiles, uploadedFiles: files };
-        contentToParse: contentWithoutFiles,
-        uploadedFiles: files,
-        isLoadingCitations: false,
-      };
     }
-
     return {
       contentToParse: rawContent ?? "",
       uploadedFiles: [] as UploadedFile[],
-      isLoadingCitations:
-        isLoading && isCitationsBlockIncomplete(rawContent ?? ""),
     };
-  }, [isLoading, message, isHuman]);
+  }, [isLoading, rawContent, reasoningContent, isHuman]);
 
-  const { citations, cleanContent, citationMap } =
-    useParsedCitations(contentToParse);
-
-  // Shared markdown components
-  const markdownComponents = useMemo(() => ({
-    a: (props: React.AnchorHTMLAttributes<HTMLAnchorElement>) => (
-      <CitationAwareLink
-        {...props}
-        citationMap={citationMap}
-        isHuman={isHuman}
-        isLoadingCitations={isLoadingCitations}
-      />
-    ),
-    img: (props: React.ImgHTMLAttributes<HTMLImageElement>) => (
-      <MessageImage {...props} threadId={thread_id} maxWidth={isHuman ? "full" : "90%"} />
-    ),
-  }), [citationMap, thread_id, isHuman, isLoadingCitations]);
-
-  // Render message response
-  // Human messages use humanMessagePlugins (no autolink) to prevent URL bleeding into adjacent text
-  const messageResponse = cleanContent ? (
-    <AIElementMessageResponse
-      remarkPlugins={isHuman ? humanMessagePlugins.remarkPlugins : streamdownPlugins.remarkPlugins}
-      rehypePlugins={isHuman ? humanMessagePlugins.rehypePlugins : [...rehypePlugins, [rehypeKatex, { output: "html" }]]}
-      components={markdownComponents}
-    >
-      {cleanContent}
-    </AIElementMessageResponse>
-  ) : null;
-
-  // Uploaded files list
   const filesList = uploadedFiles.length > 0 && thread_id ? (
     <UploadedFilesList files={uploadedFiles} threadId={thread_id} />
   ) : null;
 
-  // Citations loading indicator
+  if (isHuman) {
-  const citationsLoadingIndicator = isLoadingCitations ? (
+    const messageResponse = contentToParse ? (
-    <CitationsLoadingIndicator citations={citations} className="my-3" />
+      <AIElementMessageResponse
-  ) : null;
+        remarkPlugins={humanMessagePlugins.remarkPlugins}
-
+        rehypePlugins={humanMessagePlugins.rehypePlugins}
-  // Human messages with uploaded files: render outside bubble
+      >
-  if (isHuman && uploadedFiles.length > 0) {
+        {contentToParse}
+      </AIElementMessageResponse>
+    ) : null;
     return (
       <div className={cn("ml-auto flex flex-col gap-2", className)}>
         {filesList}
@@ -201,12 +151,23 @@ function MessageContent_({
     );
   }
 
-  // Default rendering
   return (
     <AIElementMessageContent className={className}>
       {filesList}
-      {messageResponse}
+      <SafeCitationContent
-      {citationsLoadingIndicator}
+        content={contentToParse}
+        isLoading={isLoading}
+        rehypePlugins={[...rehypePlugins, [rehypeKatex, { output: "html" }]]}
+        className="my-3"
+        isHuman={false}
+        img={(props) => (
+          <MessageImage
+            {...props}
+            threadId={thread_id}
+            maxWidth="90%"
+          />
+        )}
+      />
     </AIElementMessageContent>
   );
 }

frontend/src/components/workspace/messages/message-list.tsx

@@ -4,7 +4,6 @@ import {
   Conversation,
   ConversationContent,
 } from "@/components/ai-elements/conversation";
-import { MessageResponse } from "@/components/ai-elements/message";
 import { useI18n } from "@/core/i18n/hooks";
 import {
   extractContentFromMessage,
@@ -26,6 +25,7 @@ import { StreamingIndicator } from "../streaming-indicator";
 
 import { MessageGroup } from "./message-group";
 import { MessageListItem } from "./message-list-item";
+import { SafeCitationContent } from "./safe-citation-content";
 import { MessageListSkeleton } from "./skeleton";
 import { SubtaskCard } from "./subtask-card";
 
@@ -64,9 +64,12 @@ export function MessageList({
             const message = group.messages[0];
             if (message && hasContent(message)) {
               return (
-                <MessageResponse key={group.id} rehypePlugins={rehypePlugins}>
+                <SafeCitationContent
-                  {extractContentFromMessage(message)}
+                  key={group.id}
-                </MessageResponse>
+                  content={extractContentFromMessage(message)}
+                  isLoading={thread.isLoading}
+                  rehypePlugins={rehypePlugins}
+                />
               );
             }
             return null;
@@ -81,12 +84,12 @@ export function MessageList({
             return (
               <div className="w-full" key={group.id}>
                 {group.messages[0] && hasContent(group.messages[0]) && (
-                  <MessageResponse
+                  <SafeCitationContent
-                    className="mb-4"
+                    content={extractContentFromMessage(group.messages[0])}
+                    isLoading={thread.isLoading}
                     rehypePlugins={rehypePlugins}
-                  >
+                    className="mb-4"
-                    {extractContentFromMessage(group.messages[0])}
+                  />
-                  </MessageResponse>
                 )}
                 <ArtifactFileList files={files} threadId={threadId} />
               </div>

frontend/src/components/workspace/messages/safe-citation-content.tsx

@@ -0,0 +1,74 @@
+"use client";
+
+import type { ImgHTMLAttributes } from "react";
+import { useMemo } from "react";
+
+import {
+  CitationsLoadingIndicator,
+  createCitationMarkdownComponents,
+} from "@/components/ai-elements/inline-citation";
+import {
+  MessageResponse,
+  type MessageResponseProps,
+} from "@/components/ai-elements/message";
+import {
+  shouldShowCitationLoading,
+  useParsedCitations,
+} from "@/core/citations";
+import { streamdownPlugins } from "@/core/streamdown";
+import { cn } from "@/lib/utils";
+
+export type SafeCitationContentProps = {
+  content: string;
+  isLoading: boolean;
+  rehypePlugins: MessageResponseProps["rehypePlugins"];
+  className?: string;
+  remarkPlugins?: MessageResponseProps["remarkPlugins"];
+  isHuman?: boolean;
+  img?: (props: ImgHTMLAttributes<HTMLImageElement> & { threadId?: string; maxWidth?: string }) => React.ReactNode;
+};
+
+/** Single place for citation-aware body: shows loading until citations complete (no half-finished refs), else body. */
+export function SafeCitationContent({
+  content,
+  isLoading,
+  rehypePlugins,
+  className,
+  remarkPlugins = streamdownPlugins.remarkPlugins,
+  isHuman = false,
+  img,
+}: SafeCitationContentProps) {
+  const { citations, cleanContent, citationMap } = useParsedCitations(content);
+  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({
+        citationMap,
+        isHuman,
+        isLoadingCitations: false,
+        img,
+      }),
+    [citationMap, isHuman, img],
+  );
+  return (
+    <MessageResponse
+      className={className}
+      remarkPlugins={remarkPlugins}
+      rehypePlugins={rehypePlugins}
+      components={components}
+    >
+      {cleanContent}
+    </MessageResponse>
+  );
+}

frontend/src/components/workspace/messages/subtask-card.tsx

@@ -18,6 +18,7 @@ import { Button } from "@/components/ui/button";
 import { ShineBorder } from "@/components/ui/shine-border";
 import { useI18n } from "@/core/i18n/hooks";
 import { hasToolCalls } from "@/core/messages/utils";
+import { useRehypeSplitWordsIntoSpans } from "@/core/rehype";
 import {
   streamdownPlugins,
   streamdownPluginsWithWordAnimation,
@@ -28,9 +29,12 @@ import { cn } from "@/lib/utils";
 
 import { FlipDisplay } from "../flip-display";
 
+import { SafeCitationContent } from "./safe-citation-content";
+
 export function SubtaskCard({
   className,
   taskId,
+  isLoading,
 }: {
   className?: string;
   taskId: string;
@@ -38,6 +42,7 @@ export function SubtaskCard({
 }) {
   const { t } = useI18n();
   const [collapsed, setCollapsed] = useState(true);
+  const rehypePlugins = useRehypeSplitWordsIntoSpans(isLoading);
   const task = useSubtask(taskId)!;
   const icon = useMemo(() => {
     if (task.status === "completed") {
@@ -147,7 +152,13 @@ export function SubtaskCard({
               ></ChainOfThoughtStep>
               <ChainOfThoughtStep
                 label={
-                  <Streamdown {...streamdownPlugins}>{task.result}</Streamdown>
+                  task.result ? (
+                    <SafeCitationContent
+                      content={task.result}
+                      isLoading={false}
+                      rehypePlugins={rehypePlugins}
+                    />
+                  ) : null
                 }
               ></ChainOfThoughtStep>
             </>

frontend/src/core/citations/index.ts

@@ -1,12 +1,11 @@
 export {
   contentWithoutCitationsFromParsed,
   extractDomainFromUrl,
-  getCleanContent,
   hasCitationsBlock,
-  isCitationsBlockIncomplete,
   isExternalUrl,
   parseCitations,
   removeAllCitations,
+  shouldShowCitationLoading,
   syntheticCitationFromLink,
 } from "./utils";
 

frontend/src/core/citations/utils.ts

@@ -1,3 +1,10 @@
+/**
+ * Citation parsing and display helpers.
+ * Display rule: never show half-finished citations. Use shouldShowCitationLoading
+ * and show only the loading indicator until the block is complete and all
+ * [cite-N] refs are replaced.
+ */
+
 /**
  * Citation data structure representing a source reference
  */
@@ -16,8 +23,42 @@ export interface ParseCitationsResult {
   cleanContent: string;
 }
 
+/**
+ * Parse citation lines (one JSON object per line) into Citation array.
+ * Deduplicates by URL. Used for both complete and incomplete (streaming) blocks.
+ */
+function parseCitationLines(
+  blockContent: string,
+  seenUrls: Set<string>,
+): Citation[] {
+  const out: Citation[] = [];
+  const lines = blockContent.split("\n");
+  for (const line of lines) {
+    const trimmed = line.trim();
+    if (!trimmed?.startsWith("{")) continue;
+    try {
+      const citation = JSON.parse(trimmed) as Citation;
+      if (citation.id && citation.url && !seenUrls.has(citation.url)) {
+        seenUrls.add(citation.url);
+        out.push({
+          id: citation.id,
+          title: citation.title || "",
+          url: citation.url,
+          snippet: citation.snippet || "",
+        });
+      }
+    } catch {
+      // Skip invalid JSON lines - can happen during streaming
+    }
+  }
+  return out;
+}
+
 /**
  * Parse citations block from message content.
+ * Shared by all modes (Flash / Thinking / Pro / Ultra); supports incomplete
+ * <citations> blocks during SSE streaming (parses whatever complete JSON lines
+ * have arrived so far so [cite-N] can be linked progressively).
  *
  * The citations block format:
  * <citations>
@@ -33,41 +74,25 @@ export function parseCitations(content: string): ParseCitationsResult {
     return { citations: [], cleanContent: content };
   }
 
-  // Match ALL citations blocks anywhere in content (not just at the start)
-  const citationsRegex = /<citations>([\s\S]*?)<\/citations>/g;
   const citations: Citation[] = [];
-  const seenUrls = new Set<string>(); // Deduplicate by URL
+  const seenUrls = new Set<string>();
-  let cleanContent = content;
 
+  // 1) Complete blocks: <citations>...</citations>
+  const citationsRegex = /<citations>([\s\S]*?)<\/citations>/g;
   let match;
   while ((match = citationsRegex.exec(content)) !== null) {
-    const citationsBlock = match[1] ?? "";
+    citations.push(...parseCitationLines(match[1] ?? "", seenUrls));
+  }
 
-    // Parse each line as JSON
+  // 2) Incomplete block during streaming: <citations>... (no closing tag yet)
-    const lines = citationsBlock.split("\n");
+  if (content.includes("<citations>") && !content.includes("</citations>")) {
-    for (const line of lines) {
+    const openMatch = content.match(/<citations>([\s\S]*)$/);
-      const trimmed = line.trim();
+    if (openMatch?.[1] != null) {
-      if (trimmed?.startsWith("{")) {
+      citations.push(...parseCitationLines(openMatch[1], seenUrls));
-        try {
-          const citation = JSON.parse(trimmed) as Citation;
-          // Validate required fields and deduplicate
-          if (citation.id && citation.url && !seenUrls.has(citation.url)) {
-            seenUrls.add(citation.url);
-            citations.push({
-              id: citation.id,
-              title: citation.title || "",
-              url: citation.url,
-              snippet: citation.snippet || "",
-            });
-          }
-        } catch {
-          // Skip invalid JSON lines - this can happen during streaming
-        }
-      }
     }
   }
 
-  cleanContent = removeCitationsBlocks(content);
+  let cleanContent = removeCitationsBlocks(content);
 
   // Convert [cite-N] references to markdown links
   // Example: [cite-1] -> [Title](url)
@@ -95,13 +120,6 @@ export function parseCitations(content: string): ParseCitationsResult {
   return { citations, cleanContent };
 }
 
-/**
- * Return content with citations block removed and [cite-N] replaced by markdown links.
- */
-export function getCleanContent(content: string): string {
-  return parseCitations(content ?? "").cleanContent;
-}
-
 /**
  * Build a map from URL to Citation for quick lookup
  *
@@ -173,15 +191,32 @@ export function hasCitationsBlock(content: string): boolean {
   return Boolean(content?.includes("<citations>"));
 }
 
+/** Pattern for [cite-1], [cite-2], ... that should be replaced by parseCitations. */
+const UNREPLACED_CITE_REF = /\[cite-\d+\]/;
+
 /**
- * Check if content is still receiving the citations block (streaming)
+ * Whether cleanContent still contains unreplaced [cite-N] refs (half-finished citations).
- * This helps determine if we should wait before parsing
+ * When true, callers must not render this content and should show loading instead.
- *
- * @param content - The current content being streamed
- * @returns true if citations block appears to be incomplete
  */
-export function isCitationsBlockIncomplete(content: string): boolean {
+export function hasUnreplacedCitationRefs(cleanContent: string): boolean {
-  return hasCitationsBlock(content) && !content.includes("</citations>");
+  return Boolean(cleanContent && UNREPLACED_CITE_REF.test(cleanContent));
+}
+
+/**
+ * Single source of truth: true when body must not be rendered (show loading instead).
+ * Use after parseCitations: pass raw content, parsed cleanContent, and isLoading.
+ * When streaming and any citation block is present, show loading so the indicator
+ * is visible in all modes (Pro/Ultra often receive complete blocks in one chunk).
+ */
+export function shouldShowCitationLoading(
+  rawContent: string,
+  cleanContent: string,
+  isLoading: boolean,
+): boolean {
+  return (
+    (isLoading && hasCitationsBlock(rawContent)) ||
+    (hasCitationsBlock(rawContent) && hasUnreplacedCitationRefs(cleanContent))
+  );
 }
 
 /**