From 175c1d2e3b6acaee897fbe85356591641d1cf4f5 Mon Sep 17 00:00:00 2001 From: ruitanglin Date: Mon, 9 Feb 2026 12:13:06 +0800 Subject: [PATCH] refactor(frontend): simplify and deduplicate Citation-related code - Extract removeCitationsBlocks in utils, reuse in parseCitations and removeAllCitations - Add hasCitationsBlock; isCitationsBlockIncomplete now uses it - Add useParsedCitations hook (parseCitations + buildCitationMap) for message/artifact - Add CitationAwareLink to unify link rendering (message-list-item + artifact-file-detail) - Add getCleanContent helper; message-group uses it and useParsedCitations - ArtifactFileDetail: single useParsedCitations, pass cleanContent/citationMap to Preview - Stop exporting buildCitationMap and removeCitationsBlocks from citations index - Remove duplicate MessageLink and inline link logic in artifact preview Co-authored-by: Cursor --- .../ai-elements/inline-citation.tsx | 78 ++++++++++- .../artifacts/artifact-file-detail.tsx | 88 ++++--------- .../workspace/messages/message-group.tsx | 22 ++-- .../workspace/messages/message-list-item.tsx | 121 +++++------------- frontend/src/core/citations/index.ts | 5 +- .../core/citations/use-parsed-citations.ts | 28 ++++ frontend/src/core/citations/utils.ts | 51 +++++--- 7 files changed, 202 insertions(+), 191 deletions(-) create mode 100644 frontend/src/core/citations/use-parsed-citations.ts diff --git a/frontend/src/components/ai-elements/inline-citation.tsx b/frontend/src/components/ai-elements/inline-citation.tsx index dde6e31..ef9f893 100644 --- a/frontend/src/components/ai-elements/inline-citation.tsx +++ b/frontend/src/components/ai-elements/inline-citation.tsx @@ -12,10 +12,15 @@ import { HoverCardContent, HoverCardTrigger, } from "@/components/ui/hover-card"; -import { cn } from "@/lib/utils"; +import { + cn, + externalLinkClass, + externalLinkClassNoUnderline, +} from "@/lib/utils"; import { ExternalLinkIcon, ArrowLeftIcon, ArrowRightIcon } from "lucide-react"; import { type ComponentProps, + Children, createContext, useCallback, useContext, @@ -23,7 +28,11 @@ import { useState, } from "react"; import type { Citation } from "@/core/citations"; -import { extractDomainFromUrl } from "@/core/citations"; +import { + extractDomainFromUrl, + isExternalUrl, + syntheticCitationFromLink, +} from "@/core/citations"; import { Shimmer } from "./shimmer"; import { useI18n } from "@/core/i18n/hooks"; @@ -360,6 +369,71 @@ export const CitationLink = ({ ); }; +/** + * Renders a link with optional citation badge. Use in markdown components (message + artifact). + * - citationMap: URL -> Citation; links in map render as CitationLink. + * - isHuman: when true, never render as CitationLink (plain link). + * - isLoadingCitations: when true and not human, non-citation links use no-underline style. + * - syntheticExternal: when true, external URLs not in citationMap render as CitationLink with synthetic citation. + */ +export type CitationAwareLinkProps = ComponentProps<"a"> & { + citationMap: Map; + isHuman?: boolean; + isLoadingCitations?: boolean; + syntheticExternal?: boolean; +}; + +export const CitationAwareLink = ({ + href, + children, + citationMap, + isHuman = false, + isLoadingCitations = false, + syntheticExternal = false, + className, + ...rest +}: CitationAwareLinkProps) => { + if (!href) return {children}; + + const citation = citationMap.get(href); + + if (citation && !isHuman) { + return ( + + {children} + + ); + } + + if (syntheticExternal && isExternalUrl(href)) { + const linkText = + typeof children === "string" + ? children + : String(Children.toArray(children).join("")).trim() || href; + return ( + + {children} + + ); + } + + const noUnderline = !isHuman && isLoadingCitations; + return ( + + {children} + + ); +}; + /** * Shared CitationsLoadingIndicator component * Used across message-list-item and message-group to show loading citations diff --git a/frontend/src/components/workspace/artifacts/artifact-file-detail.tsx b/frontend/src/components/workspace/artifacts/artifact-file-detail.tsx index 1178899..53bfcb1 100644 --- 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 { CitationLink } from "@/components/ai-elements/inline-citation"; +import { CitationAwareLink } from "@/components/ai-elements/inline-citation"; import { Select, SelectItem } from "@/components/ui/select"; import { SelectContent, @@ -33,19 +33,14 @@ import { ToggleGroup, ToggleGroupItem } from "@/components/ui/toggle-group"; import { CodeEditor } from "@/components/workspace/code-editor"; import { useArtifactContent } from "@/core/artifacts/hooks"; import { urlOfArtifact } from "@/core/artifacts/utils"; -import { - buildCitationMap, - isExternalUrl, - parseCitations, - removeAllCitations, - syntheticCitationFromLink, -} from "@/core/citations"; +import type { Citation } from "@/core/citations"; +import { removeAllCitations, useParsedCitations } from "@/core/citations"; import { useI18n } from "@/core/i18n/hooks"; import { installSkill } from "@/core/skills/api"; import { streamdownPlugins } from "@/core/streamdown"; import { checkCodeFile, getFileName } from "@/core/utils/files"; import { env } from "@/env"; -import { cn, externalLinkClass } from "@/lib/utils"; +import { cn } from "@/lib/utils"; import { Tooltip } from "../tooltip"; @@ -96,15 +91,11 @@ export function ArtifactFileDetail({ enabled: isCodeFile && !isWriteFile, }); - // Parse citations and get clean content for code editor - const cleanContent = useMemo(() => { - if (language === "markdown" && content) { - return parseCitations(content).cleanContent; - } - return content; - }, [content, language]); - - // Get content without ANY citations for copy/download + const parsed = useParsedCitations( + language === "markdown" ? (content ?? "") : "", + ); + const cleanContent = + language === "markdown" && content ? parsed.cleanContent : (content ?? ""); const contentWithoutCitations = useMemo(() => { if (language === "markdown" && content) { return removeAllCitations(content); @@ -260,6 +251,8 @@ export function ArtifactFileDetail({ threadId={threadId} content={content} language={language ?? "text"} + cleanContent={parsed.cleanContent} + citationMap={parsed.citationMap} /> )} {isCodeFile && viewMode === "code" && ( @@ -285,21 +278,16 @@ export function ArtifactFilePreview({ threadId, content, language, + cleanContent, + citationMap, }: { filepath: string; threadId: string; content: string; language: string; + cleanContent: string; + citationMap: Map; }) { - const { cleanContent, citationMap } = React.useMemo(() => { - const parsed = parseCitations(content ?? ""); - const map = buildCitationMap(parsed.citations); - return { - cleanContent: parsed.cleanContent, - citationMap: map, - }; - }, [content]); - if (language === "markdown") { return (
@@ -307,45 +295,13 @@ export function ArtifactFilePreview({ className="size-full" {...streamdownPlugins} components={{ - a: ({ - href, - children, - }: React.AnchorHTMLAttributes) => { - if (!href) return {children}; - const citation = citationMap.get(href); - if (citation) { - return ( - - {children} - - ); - } - if (isExternalUrl(href)) { - const linkText = - typeof children === "string" - ? children - : String(React.Children.toArray(children).join("")).trim() || - href; - return ( - - {children} - - ); - } - return ( - - {children} - - ); - }, + a: (props: React.AnchorHTMLAttributes) => ( + + ), }} > {cleanContent ?? ""} diff --git a/frontend/src/components/workspace/messages/message-group.tsx b/frontend/src/components/workspace/messages/message-group.tsx index f1b7a71..524d16f 100644 --- a/frontend/src/components/workspace/messages/message-group.tsx +++ b/frontend/src/components/workspace/messages/message-group.tsx @@ -25,7 +25,11 @@ 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 { parseCitations } from "@/core/citations"; +import { + getCleanContent, + hasCitationsBlock, + useParsedCitations, +} from "@/core/citations"; import { useI18n } from "@/core/i18n/hooks"; import { extractReasoningContentFromMessage, @@ -124,7 +128,7 @@ export function MessageGroup({ remarkPlugins={streamdownPlugins.remarkPlugins} rehypePlugins={rehypePlugins} > - {parseCitations(step.reasoning ?? "").cleanContent} + {getCleanContent(step.reasoning ?? "")} } > @@ -177,10 +181,7 @@ export function MessageGroup({ remarkPlugins={streamdownPlugins.remarkPlugins} rehypePlugins={rehypePlugins} > - { - parseCitations(lastReasoningStep.reasoning ?? "") - .cleanContent - } + {getCleanContent(lastReasoningStep.reasoning ?? "")} } > @@ -215,12 +216,8 @@ function ToolCall({ const { thread } = useThread(); const threadIsLoading = thread.isLoading; - // Move useMemo to top level to comply with React Hooks rules const fileContent = typeof args.content === "string" ? args.content : ""; - const { citations } = useMemo( - () => parseCitations(fileContent), - [fileContent], - ); + const { citations } = useParsedCitations(fileContent); if (name === "web_search") { let label: React.ReactNode = t.toolCalls.searchForRelatedInfo; @@ -370,9 +367,8 @@ function ToolCall({ const isMarkdown = path?.toLowerCase().endsWith(".md") || path?.toLowerCase().endsWith(".markdown"); - const hasCitationsBlock = fileContent.includes(""); const showCitationsLoading = - isMarkdown && threadIsLoading && hasCitationsBlock && isLast; + isMarkdown && threadIsLoading && hasCitationsBlock(fileContent) && isLast; return ( <> diff --git a/frontend/src/components/workspace/messages/message-list-item.tsx b/frontend/src/components/workspace/messages/message-list-item.tsx index 7858951..787d921 100644 --- a/frontend/src/components/workspace/messages/message-list-item.tsx +++ b/frontend/src/components/workspace/messages/message-list-item.tsx @@ -5,7 +5,7 @@ import { memo, useMemo } from "react"; import rehypeKatex from "rehype-katex"; import { - CitationLink, + CitationAwareLink, CitationsLoadingIndicator, } from "@/components/ai-elements/inline-citation"; import { @@ -17,11 +17,9 @@ import { import { Badge } from "@/components/ui/badge"; import { resolveArtifactURL } from "@/core/artifacts/utils"; import { - type Citation, - buildCitationMap, isCitationsBlockIncomplete, - parseCitations, removeAllCitations, + useParsedCitations, } from "@/core/citations"; import { extractContentFromMessage, @@ -31,11 +29,7 @@ import { } from "@/core/messages/utils"; import { useRehypeSplitWordsIntoSpans } from "@/core/rehype"; import { humanMessagePlugins, streamdownPlugins } from "@/core/streamdown"; -import { - cn, - externalLinkClass, - externalLinkClassNoUnderline, -} from "@/lib/utils"; +import { cn } from "@/lib/utils"; import { CopyButton } from "../copy-button"; @@ -79,49 +73,6 @@ export function MessageListItem({ ); } -/** - * Custom link component that handles citations and external links - * Only links in citationMap are rendered as CitationLink badges - * Other links (project URLs, regular links) are rendered as plain links - * During citation loading (streaming), non-citation links are rendered without underline so they match final citation style (p3) - */ -function MessageLink({ - href, - children, - citationMap, - isHuman, - isLoadingCitations, -}: React.AnchorHTMLAttributes & { - citationMap: Map; - isHuman: boolean; - isLoadingCitations?: boolean; -}) { - if (!href) return {children}; - - const citation = citationMap.get(href); - - // Only render as CitationLink badge if it's a citation (in citationMap) and not human message - if (citation && !isHuman) { - return ( - - {children} - - ); - } - - const noUnderline = !isHuman && isLoadingCitations; - return ( - - {children} - - ); -} - /** * Custom image component that handles artifact URLs */ @@ -165,50 +116,44 @@ function MessageContent_({ const isHuman = message.type === "human"; const { thread_id } = useParams<{ thread_id: string }>(); - // Extract and parse citations and uploaded files from message content - const { citations, cleanContent, uploadedFiles, isLoadingCitations } = - useMemo(() => { - const reasoningContent = extractReasoningContentFromMessage(message); - const rawContent = extractContentFromMessage(message); + // Content to parse for citations (and optionally uploaded files) + const { contentToParse, uploadedFiles, isLoadingCitations } = useMemo(() => { + const reasoningContent = extractReasoningContentFromMessage(message); + const rawContent = extractContentFromMessage(message); - // When only reasoning content exists (no main content), also parse citations - if (!isLoading && reasoningContent && !rawContent) { - const { citations, cleanContent } = parseCitations(reasoningContent); - return { - citations, - cleanContent, - uploadedFiles: [], - isLoadingCitations: false, - }; - } + if (!isLoading && reasoningContent && !rawContent) { + return { + contentToParse: reasoningContent, + uploadedFiles: [] as UploadedFile[], + isLoadingCitations: false, + }; + } - // For human messages, parse uploaded files first - if (isHuman && rawContent) { - const { files, cleanContent: contentWithoutFiles } = - parseUploadedFiles(rawContent); - const { citations, cleanContent: finalContent } = - parseCitations(contentWithoutFiles); - return { - citations, - cleanContent: finalContent, - uploadedFiles: files, - isLoadingCitations: false, - }; - } + if (isHuman && rawContent) { + const { files, cleanContent: contentWithoutFiles } = + parseUploadedFiles(rawContent); + return { + contentToParse: contentWithoutFiles, + uploadedFiles: files, + isLoadingCitations: false, + }; + } - const { citations, cleanContent } = parseCitations(rawContent ?? ""); - const isLoadingCitations = - isLoading && isCitationsBlockIncomplete(rawContent ?? ""); + return { + contentToParse: rawContent ?? "", + uploadedFiles: [] as UploadedFile[], + isLoadingCitations: + isLoading && isCitationsBlockIncomplete(rawContent ?? ""), + }; + }, [isLoading, message, isHuman]); - return { citations, cleanContent, uploadedFiles: [], isLoadingCitations }; - }, [isLoading, message, isHuman]); - - const citationMap = useMemo(() => buildCitationMap(citations), [citations]); + const { citations, cleanContent, citationMap } = + useParsedCitations(contentToParse); // Shared markdown components const markdownComponents = useMemo(() => ({ a: (props: React.AnchorHTMLAttributes) => ( - ; +} + +/** + * 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); + return { + citations: parsed.citations, + cleanContent: parsed.cleanContent, + citationMap, + }; + }, [content]); +} diff --git a/frontend/src/core/citations/utils.ts b/frontend/src/core/citations/utils.ts index 1937af1..d2384a4 100644 --- a/frontend/src/core/citations/utils.ts +++ b/frontend/src/core/citations/utils.ts @@ -67,14 +67,7 @@ export function parseCitations(content: string): ParseCitationsResult { } } - // Remove ALL citations blocks from content (both complete and incomplete) - cleanContent = content.replace(/[\s\S]*?<\/citations>/g, "").trim(); - - // Also remove incomplete citations blocks (during streaming) - // Match without closing tag or followed by anything until end of string - if (cleanContent.includes("")) { - cleanContent = cleanContent.replace(/[\s\S]*$/g, "").trim(); - } + cleanContent = removeCitationsBlocks(content); // Convert [cite-N] references to markdown links // Example: [cite-1] -> [Title](url) @@ -102,6 +95,13 @@ 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 * @@ -153,6 +153,26 @@ export function extractDomainFromUrl(url: string): string { } } +/** + * Remove all blocks from content (complete and incomplete). + * Does not remove [cite-N] or markdown links; use removeAllCitations for that. + */ +export function removeCitationsBlocks(content: string): string { + if (!content) return content; + let result = content.replace(/[\s\S]*?<\/citations>/g, "").trim(); + if (result.includes("")) { + result = result.replace(/[\s\S]*$/g, "").trim(); + } + return result; +} + +/** + * Whether content contains a block (open tag). + */ +export function hasCitationsBlock(content: string): boolean { + return Boolean(content?.includes("")); +} + /** * Check if content is still receiving the citations block (streaming) * This helps determine if we should wait before parsing @@ -161,15 +181,7 @@ export function extractDomainFromUrl(url: string): string { * @returns true if citations block appears to be incomplete */ export function isCitationsBlockIncomplete(content: string): boolean { - if (!content) { - return false; - } - - // Check if we have an opening tag but no closing tag - const hasOpenTag = content.includes(""); - const hasCloseTag = content.includes(""); - - return hasOpenTag && !hasCloseTag; + return hasCitationsBlock(content) && !content.includes(""); } /** @@ -188,11 +200,8 @@ export function removeAllCitations(content: string): string { return content; } - let result = content; - // Step 1: Remove all blocks (complete and incomplete) - result = result.replace(/[\s\S]*?<\/citations>/g, ""); - result = result.replace(/[\s\S]*$/g, ""); + let result = removeCitationsBlocks(content); // Step 2: Remove all [cite-N] references result = result.replace(/\[cite-\d+\]/g, "");