Merge pull request #27 from LofiSu/experimental

Experimental

.gitignore

@@ -35,5 +35,8 @@ coverage/
 skills/custom/*
 logs/
 
+# Local git hooks (keep only on this machine, do not push)
+.githooks/
+
 # pnpm
 .pnpm-store

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>

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">;
+const InlineCitationCardBody = ({
-
-export const InlineCitationCardBody = ({
   className,
   ...props
-}: InlineCitationCardBodyProps) => (
+}: ComponentProps<"div">) => (
   <HoverCardContent className={cn("relative w-80 p-0", className)} {...props} />
 );
 
-export type InlineCitationSourceProps = ComponentProps<"div"> & {
+const InlineCitationSource = ({
-  title?: string;
-  url?: string;
-  description?: string;
-};
-
-export 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

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" && (
+        {previewable &&
-          <ArtifactFilePreview
+          viewMode === "preview" &&
-            filepath={filepath}
+          language === "markdown" &&
-            threadId={threadId}
+          content && (
-            content={content}
+            <SafeCitationContent
-            language={language ?? "text"}
+              content={content}
-            cleanContent={parsed.cleanContent}
+              isLoading={thread.isLoading}
-            citationMap={parsed.citationMap}
+              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={{
+          components={components}
-            a: (props: React.AnchorHTMLAttributes<HTMLAnchorElement>) => (
-              <CitationAwareLink
-                {...props}
-                citationMap={citationMap}
-                syntheticExternal
-              />
-            ),
-          }}
         >
           {cleanContent ?? ""}
         </Streamdown>

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
+                    <SafeCitationContent
-                      remarkPlugins={streamdownPlugins.remarkPlugins}
+                      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
+                  <SafeCitationContent
-                    remarkPlugins={streamdownPlugins.remarkPlugins}
+                    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 && (
+        {isMarkdown && (
-          <div className="mt-2 ml-8">
+          <SafeCitationContent
-            <CitationsLoadingIndicator citations={citations} />
+            content={fileContent}
-          </div>
+            isLoading={threadIsLoading && isLast}
+            rehypePlugins={rehypePlugins}
+            loadingOnly
+            className="mt-2 ml-8"
+          />
         )}
       </>
     );

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,
@@ -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}>
+                <SafeCitationContent
-                  {extractContentFromMessage(message)}
+                  key={group.id}
-                </MessageResponse>
+                  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
+                  <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,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>
+  );
+}

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,10 @@
 export {
   contentWithoutCitationsFromParsed,
   extractDomainFromUrl,
-  getCleanContent,
-  hasCitationsBlock,
-  isCitationsBlockIncomplete,
   isExternalUrl,
   parseCitations,
   removeAllCitations,
+  shouldShowCitationLoading,
   syntheticCitationFromLink,
 } from "./utils";
 

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,

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,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)
+ * 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.
+ * 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);
 }
 
 /**