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 保持内部使用；在文件头注释中说明展示规则。
2026-04-02 22:02:13 +08:00 · 2026-02-09 15:01:51 +08:00
parent 79c85d6410
commit a4268cb6d3
9 changed files with 309 additions and 161 deletions
--- a/frontend/src/components/ai-elements/inline-citation.tsx
+++ b/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
--- a/frontend/src/components/workspace/artifacts/artifact-file-detail.tsx
+++ b/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
-            filepath={filepath}
-            threadId={threadId}
-            content={content}
-            language={language ?? "text"}
-            cleanContent={parsed.cleanContent}
-            citationMap={parsed.citationMap}
-          />
+          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={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={{
-            a: (props: React.AnchorHTMLAttributes<HTMLAnchorElement>) => (
-              <CitationAwareLink
-                {...props}
-                citationMap={citationMap}
-                syntheticExternal
-              />
-            ),
-          }}
+          components={components}
        >
          {cleanContent ?? ""}
        </Streamdown>
--- a/frontend/src/components/workspace/messages/message-group.tsx
+++ b/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 {
-  getCleanContent,
-  hasCitationsBlock,
-  useParsedCitations,
-} from "@/core/citations";
+import { shouldShowCitationLoading, 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
-                      remarkPlugins={streamdownPlugins.remarkPlugins}
+                    <SafeCitationContent
+                      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
-                    remarkPlugins={streamdownPlugins.remarkPlugins}
+                  <SafeCitationContent
+                    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 (
      <>
--- a/frontend/src/components/workspace/messages/message-list-item.tsx
+++ b/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 {
-  isCitationsBlockIncomplete,
-  removeAllCitations,
-  useParsedCitations,
-} from "@/core/citations";
+import { removeAllCitations } 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 { contentToParse, uploadedFiles, isLoadingCitations } = useMemo(() => {
-    const reasoningContent = extractReasoningContentFromMessage(message);
-    const rawContent = extractContentFromMessage(message);
-
+  const rawContent = extractContentFromMessage(message);
+  const reasoningContent = extractReasoningContentFromMessage(message);
+  const { contentToParse, uploadedFiles } = useMemo(() => {
    if (!isLoading && reasoningContent && !rawContent) {
-      return {
-        contentToParse: reasoningContent,
-        uploadedFiles: [] as UploadedFile[],
-        isLoadingCitations: false,
-      };
+      return { contentToParse: reasoningContent, uploadedFiles: [] as UploadedFile[] };
    }
-
    if (isHuman && rawContent) {
      const { files, cleanContent: contentWithoutFiles } =
        parseUploadedFiles(rawContent);
-      return {
-        contentToParse: contentWithoutFiles,
-        uploadedFiles: files,
-        isLoadingCitations: false,
-      };
+      return { contentToParse: contentWithoutFiles, uploadedFiles: files };
    }
-
    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
-  const citationsLoadingIndicator = isLoadingCitations ? (
-    <CitationsLoadingIndicator citations={citations} className="my-3" />
-  ) : null;
-
-  // Human messages with uploaded files: render outside bubble
-  if (isHuman && uploadedFiles.length > 0) {
+  if (isHuman) {
+    const messageResponse = contentToParse ? (
+      <AIElementMessageResponse
+        remarkPlugins={humanMessagePlugins.remarkPlugins}
+        rehypePlugins={humanMessagePlugins.rehypePlugins}
+      >
+        {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}
-      {citationsLoadingIndicator}
+      <SafeCitationContent
+        content={contentToParse}
+        isLoading={isLoading}
+        rehypePlugins={[...rehypePlugins, [rehypeKatex, { output: "html" }]]}
+        className="my-3"
+        isHuman={false}
+        img={(props) => (
+          <MessageImage
+            {...props}
+            threadId={thread_id}
+            maxWidth="90%"
+          />
+        )}
+      />
    </AIElementMessageContent>
  );
 }
--- a/frontend/src/components/workspace/messages/message-list.tsx
+++ b/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}>
-                  {extractContentFromMessage(message)}
-                </MessageResponse>
+                <SafeCitationContent
+                  key={group.id}
+                  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
-                    className="mb-4"
+                  <SafeCitationContent
+                    content={extractContentFromMessage(group.messages[0])}
+                    isLoading={thread.isLoading}
                    rehypePlugins={rehypePlugins}
-                  >
-                    {extractContentFromMessage(group.messages[0])}
-                  </MessageResponse>
+                    className="mb-4"
+                  />
                )}
                <ArtifactFileList files={files} threadId={threadId} />
              </div>
--- a/frontend/src/components/workspace/messages/safe-citation-content.tsx
+++ b/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>
+  );
+}
--- a/frontend/src/components/workspace/messages/subtask-card.tsx
+++ b/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>
            </>
--- a/frontend/src/core/citations/index.ts
+++ b/frontend/src/core/citations/index.ts
@@ -1,12 +1,11 @@
 export {
  contentWithoutCitationsFromParsed,
  extractDomainFromUrl,
-  getCleanContent,
  hasCitationsBlock,
-  isCitationsBlockIncomplete,
  isExternalUrl,
  parseCitations,
  removeAllCitations,
+  shouldShowCitationLoading,
  syntheticCitationFromLink,
 } from "./utils";

--- a/frontend/src/core/citations/utils.ts
+++ b/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
-  let cleanContent = content;
+  const seenUrls = new Set<string>();

+  // 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
-    const lines = citationsBlock.split("\n");
-    for (const line of lines) {
-      const trimmed = line.trim();
-      if (trimmed?.startsWith("{")) {
-        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
-        }
-      }
+  // 2) Incomplete block during streaming: <citations>... (no closing tag yet)
+  if (content.includes("<citations>") && !content.includes("</citations>")) {
+    const openMatch = content.match(/<citations>([\s\S]*)$/);
+    if (openMatch?.[1] != null) {
+      citations.push(...parseCitationLines(openMatch[1], seenUrls));
    }
  }

-  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)
- * This helps determine if we should wait before parsing
- *
- * @param content - The current content being streamed
- * @returns true if citations block appears to be incomplete
+ * Whether cleanContent still contains unreplaced [cite-N] refs (half-finished citations).
+ * When true, callers must not render this content and should show loading instead.
 */
-export function isCitationsBlockIncomplete(content: string): boolean {
-  return hasCitationsBlock(content) && !content.includes("</citations>");
+export function hasUnreplacedCitationRefs(cleanContent: string): boolean {
+  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))
+  );
 }

 /**