Merge pull request #27 from LofiSu/experimental

Experimental
2026-04-02 22:02:13 +08:00 · 2026-02-09 15:59:41 +08:00
parent cbe0f3b32f cef8d389fd
commit 0f0346f73e
12 changed files with 358 additions and 206 deletions
--- a/.gitignore
+++ b/.gitignore
@@ -35,5 +35,8 @@ coverage/
 skills/custom/*
 logs/

+# Local git hooks (keep only on this machine, do not push)
+.githooks/
+
 # pnpm
 .pnpm-store
--- a/frontend/src/app/workspace/chats/[thread_id]/page.tsx
+++ b/frontend/src/app/workspace/chats/[thread_id]/page.tsx
@@ -28,8 +28,9 @@ import { Welcome } from "@/components/workspace/welcome";
 import { useI18n } from "@/core/i18n/hooks";
 import { useNotification } from "@/core/notification/hooks";
 import { useLocalSettings } from "@/core/settings";
-import { type AgentThread } from "@/core/threads";
+import { type AgentThread, type AgentThreadState } from "@/core/threads";
 import { useSubmitThread, useThreadStream } from "@/core/threads/hooks";
+import type { Message } from "@langchain/langgraph-sdk";
 import {
  pathOfThread,
  textOfMessage,
@@ -88,10 +89,12 @@ export default function ChatPage() {
  }, [threadIdFromPath]);

  const { showNotification } = useNotification();
+  const [finalState, setFinalState] = useState<AgentThreadState | null>(null);
  const thread = useThreadStream({
    isNewThread,
    threadId,
    onFinish: (state) => {
+      setFinalState(state);
      if (document.hidden || !document.hasFocus()) {
        let body = "Conversation finished";
        const lastMessage = state.messages[state.messages.length - 1];
@@ -111,6 +114,9 @@ export default function ChatPage() {
      }
    },
  });
+  useEffect(() => {
+    if (thread.isLoading) setFinalState(null);
+  }, [thread.isLoading]);

  const title = useMemo(() => {
    let result = isNewThread
@@ -239,6 +245,11 @@ export default function ChatPage() {
                  className={cn("size-full", !isNewThread && "pt-10")}
                  threadId={threadId}
                  thread={thread}
+                  messagesOverride={
+                    !thread.isLoading && finalState?.messages
+                      ? (finalState.messages as Message[])
+                      : undefined
+                  }
                  paddingBottom={todoListCollapsed ? 160 : 280}
                />
              </div>
--- 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,
@@ -22,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>
@@ -202,6 +205,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,
@@ -48,6 +48,9 @@ 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";

 export function ArtifactFileDetail({
@@ -89,6 +92,7 @@ export function ArtifactFileDetail({
  const previewable = useMemo(() => {
    return (language === "html" && !isWriteFile) || language === "markdown";
  }, [isWriteFile, language]);
+  const { thread } = useThread();
  const { content } = useArtifactContent({
    threadId,
    filepath: filepathFromProps,
@@ -247,16 +251,27 @@ export function ArtifactFileDetail({
        </div>
      </ArtifactHeader>
      <ArtifactContent className="p-0">
-        {previewable && viewMode === "preview" && (
-          <ArtifactFilePreview
-            filepath={filepath}
-            threadId={threadId}
-            content={content}
-            language={language ?? "text"}
-            cleanContent={parsed.cleanContent}
-            citationMap={parsed.citationMap}
-          />
-        )}
+        {previewable &&
+          viewMode === "preview" &&
+          language === "markdown" &&
+          content && (
+            <SafeCitationContent
+              content={content}
+              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"
@@ -291,20 +306,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
@@ -22,14 +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 {
-  getCleanContent,
-  hasCitationsBlock,
-  useParsedCitations,
-} from "@/core/citations";
 import { useI18n } from "@/core/i18n/hooks";
 import {
  extractReasoningContentFromMessage,
@@ -47,6 +41,8 @@ import { Tooltip } from "../tooltip";

 import { useThread } from "./context";

+import { SafeCitationContent } from "./safe-citation-content";
+
 export function MessageGroup({
  className,
  messages,
@@ -124,16 +120,20 @@ 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>
              ) : (
-                <ToolCall key={step.id} {...step} isLoading={isLoading} />
+                <ToolCall
+                  key={step.id}
+                  {...step}
+                  isLoading={isLoading}
+                  rehypePlugins={rehypePlugins}
+                />
              ),
            )}
          {lastToolCallStep && (
@@ -143,6 +143,7 @@ export function MessageGroup({
                {...lastToolCallStep}
                isLast={true}
                isLoading={isLoading}
+                rehypePlugins={rehypePlugins}
              />
            </FlipDisplay>
          )}
@@ -177,12 +178,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>
@@ -201,6 +201,7 @@ function ToolCall({
  result,
  isLast = false,
  isLoading = false,
+  rehypePlugins,
 }: {
  id?: string;
  messageId?: string;
@@ -209,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 } =
@@ -217,7 +219,6 @@ function ToolCall({
  const threadIsLoading = thread.isLoading;

  const fileContent = typeof args.content === "string" ? args.content : "";
-  const { citations } = useParsedCitations(fileContent);

  if (name === "web_search") {
    let label: React.ReactNode = t.toolCalls.searchForRelatedInfo;
@@ -363,12 +364,9 @@ 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;

    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"
+          />
        )}
      </>
    );
--- 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,
@@ -19,6 +18,7 @@ import { useRehypeSplitWordsIntoSpans } from "@/core/rehype";
 import type { Subtask } from "@/core/tasks";
 import { useUpdateSubtask } from "@/core/tasks/context";
 import type { AgentThreadState } from "@/core/threads";
+import type { Message } from "@langchain/langgraph-sdk";
 import { cn } from "@/lib/utils";

 import { ArtifactFileList } from "../artifacts/artifact-file-list";
@@ -26,6 +26,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";

@@ -33,16 +34,20 @@ export function MessageList({
  className,
  threadId,
  thread,
+  messagesOverride,
  paddingBottom = 160,
 }: {
  className?: string;
  threadId: string;
  thread: UseStream<AgentThreadState>;
+  /** When set (e.g. from onFinish), use instead of thread.messages so SSE end shows complete state. */
+  messagesOverride?: Message[];
  paddingBottom?: number;
 }) {
  const { t } = useI18n();
  const rehypePlugins = useRehypeSplitWordsIntoSpans(thread.isLoading);
  const updateSubtask = useUpdateSubtask();
+  const messages = messagesOverride ?? thread.messages;
  if (thread.isThreadLoading) {
    return <MessageListSkeleton />;
  }
@@ -51,7 +56,7 @@ export function MessageList({
      className={cn("flex size-full flex-col justify-center", className)}
    >
      <ConversationContent className="mx-auto w-full max-w-(--container-width-md) gap-8 pt-12">
-        {groupMessages(thread.messages, (group) => {
+        {groupMessages(messages, (group) => {
          if (group.type === "human" || group.type === "assistant") {
            return (
              <MessageListItem
@@ -64,9 +69,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 +89,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,85 @@
+"use client";
+
+import type { ImgHTMLAttributes } from "react";
+import type { ReactNode } 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,
+  type UseParsedCitationsResult,
+} 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 }) => 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. */
+export function SafeCitationContent({
+  content,
+  isLoading,
+  rehypePlugins,
+  className,
+  remarkPlugins = streamdownPlugins.remarkPlugins,
+  isHuman = false,
+  img,
+  loadingOnly = false,
+  renderBody,
+}: SafeCitationContentProps) {
+  const parsed = useParsedCitations(content);
+  const { citations, cleanContent, citationMap } = parsed;
+  const showLoading = shouldShowCitationLoading(content, cleanContent, isLoading);
+  const components = useMemo(
+    () =>
+      createCitationMarkdownComponents({
+        citationMap,
+        isHuman,
+        isLoadingCitations: false,
+        img,
+      }),
+    [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}
+      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,10 @@
 export {
  contentWithoutCitationsFromParsed,
  extractDomainFromUrl,
-  getCleanContent,
-  hasCitationsBlock,
-  isCitationsBlockIncomplete,
  isExternalUrl,
  parseCitations,
  removeAllCitations,
+  shouldShowCitationLoading,
  syntheticCitationFromLink,
 } from "./utils";

--- 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,
--- 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,29 +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
- *
- * @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).
 */
@@ -173,15 +175,30 @@ 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.
+ * Never show body when cleanContent still has [cite-N] (e.g. refs arrived before
+ * <citations> block in stream); also show loading while streaming with citation block.
+ */
+export function shouldShowCitationLoading(
+  rawContent: string,
+  cleanContent: string,
+  isLoading: boolean,
+): boolean {
+  if (hasUnreplacedCitationRefs(cleanContent)) return true;
+  return isLoading && hasCitationsBlock(rawContent);
 }

 /**