diff --git a/.gitignore b/.gitignore index 1bfd7f6..559f548 100644 --- 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 diff --git a/frontend/src/app/workspace/chats/[thread_id]/page.tsx b/frontend/src/app/workspace/chats/[thread_id]/page.tsx index d521b7d..5d39207 100644 --- 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(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} /> diff --git a/frontend/src/components/ai-elements/inline-citation.tsx b/frontend/src/components/ai-elements/inline-citation.tsx index 49ba280..506e428 100644 --- 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; +type InlineCitationCardProps = ComponentProps; -export const InlineCitationCard = (props: InlineCitationCardProps) => ( +const InlineCitationCard = (props: InlineCitationCardProps) => ( ); -export type InlineCitationCardBodyProps = ComponentProps<"div">; - -export const InlineCitationCardBody = ({ +const InlineCitationCardBody = ({ className, ...props -}: InlineCitationCardBodyProps) => ( +}: ComponentProps<"div">) => ( ); -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; +}) => (
{title && (

{title}

@@ -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; + isHuman?: boolean; + isLoadingCitations?: boolean; + syntheticExternal?: boolean; + /** Optional custom img component (e.g. MessageImage with threadId). Omit for artifact. */ + img?: (props: ImgHTMLAttributes & { 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) => ReactElement; + img?: (props: ImgHTMLAttributes & { threadId?: string; maxWidth?: string }) => ReactNode; +} { + const { + citationMap, + isHuman = false, + isLoadingCitations = false, + syntheticExternal = false, + img, + } = options; + const a = (props: AnchorHTMLAttributes) => ( + + ); + return img ? { a, img } : { a }; +} + /** * 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 4c88aa9..2bfae75 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 { 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({
- {previewable && viewMode === "preview" && ( - - )} + {previewable && + viewMode === "preview" && + language === "markdown" && + content && ( + ( + + )} + /> + )} {isCodeFile && viewMode === "code" && ( ; }) { if (language === "markdown") { + const components = createCitationMarkdownComponents({ + citationMap, + syntheticExternal: true, + }); return (
) => ( - - ), - }} + components={components} > {cleanContent ?? ""} diff --git a/frontend/src/components/workspace/messages/message-group.tsx b/frontend/src/components/workspace/messages/message-group.tsx index 524d16f..3a2bb94 100644 --- 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({ - {getCleanContent(step.reasoning ?? "")} - + /> } > ) : ( - + ), )} {lastToolCallStep && ( @@ -143,6 +143,7 @@ export function MessageGroup({ {...lastToolCallStep} isLast={true} isLoading={isLoading} + rehypePlugins={rehypePlugins} /> )} @@ -177,12 +178,11 @@ export function MessageGroup({ - {getCleanContent(lastReasoningStep.reasoning ?? "")} - + /> } > @@ -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; isLast?: boolean; isLoading?: boolean; + rehypePlugins: ReturnType; }) { 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({ )} - {showCitationsLoading && ( -
- -
+ {isMarkdown && ( + )} ); diff --git a/frontend/src/components/workspace/messages/message-list-item.tsx b/frontend/src/components/workspace/messages/message-list-item.tsx index 787d921..fced6fe 100644 --- 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) => ( - - ), - img: (props: React.ImgHTMLAttributes) => ( - - ), - }), [citationMap, thread_id, isHuman, isLoadingCitations]); - - // Render message response - // Human messages use humanMessagePlugins (no autolink) to prevent URL bleeding into adjacent text - const messageResponse = cleanContent ? ( - - {cleanContent} - - ) : null; - - // Uploaded files list const filesList = uploadedFiles.length > 0 && thread_id ? ( ) : null; - // Citations loading indicator - const citationsLoadingIndicator = isLoadingCitations ? ( - - ) : null; - - // Human messages with uploaded files: render outside bubble - if (isHuman && uploadedFiles.length > 0) { + if (isHuman) { + const messageResponse = contentToParse ? ( + + {contentToParse} + + ) : null; return (
{filesList} @@ -201,12 +151,23 @@ function MessageContent_({ ); } - // Default rendering return ( {filesList} - {messageResponse} - {citationsLoadingIndicator} + ( + + )} + /> ); } diff --git a/frontend/src/components/workspace/messages/message-list.tsx b/frontend/src/components/workspace/messages/message-list.tsx index a615515..430d94e 100644 --- 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; + /** 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 ; } @@ -51,7 +56,7 @@ export function MessageList({ className={cn("flex size-full flex-col justify-center", className)} > - {groupMessages(thread.messages, (group) => { + {groupMessages(messages, (group) => { if (group.type === "human" || group.type === "assistant") { return ( - {extractContentFromMessage(message)} - + ); } return null; @@ -81,12 +89,12 @@ export function MessageList({ return (
{group.messages[0] && hasContent(group.messages[0]) && ( - - {extractContentFromMessage(group.messages[0])} - + className="mb-4" + /> )}
diff --git a/frontend/src/components/workspace/messages/safe-citation-content.tsx b/frontend/src/components/workspace/messages/safe-citation-content.tsx new file mode 100644 index 0000000..e453b89 --- /dev/null +++ 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 & { 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 ( + + ); + } + if (loadingOnly) return null; + if (renderBody) return renderBody(parsed); + if (!cleanContent) return null; + + return ( + + {cleanContent} + + ); +} diff --git a/frontend/src/components/workspace/messages/subtask-card.tsx b/frontend/src/components/workspace/messages/subtask-card.tsx index 924ffc1..a8d0e63 100644 --- 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({ > {task.result} + task.result ? ( + + ) : null } > diff --git a/frontend/src/core/citations/index.ts b/frontend/src/core/citations/index.ts index c56d293..ac41316 100644 --- 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"; diff --git a/frontend/src/core/citations/use-parsed-citations.ts b/frontend/src/core/citations/use-parsed-citations.ts index 1dc8463..7be76de 100644 --- 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(); + for (const c of parsed.citations) citationMap.set(c.url, c); return { citations: parsed.citations, cleanContent: parsed.cleanContent, diff --git a/frontend/src/core/citations/utils.ts b/frontend/src/core/citations/utils.ts index 0567b02..fb5cded 100644 --- 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, +): 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 + * blocks during SSE streaming (parses whatever complete JSON lines + * have arrived so far so [cite-N] can be linked progressively). * * The citations block format: * @@ -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 = /([\s\S]*?)<\/citations>/g; const citations: Citation[] = []; - const seenUrls = new Set(); // Deduplicate by URL - let cleanContent = content; + const seenUrls = new Set(); + // 1) Complete blocks: ... + const citationsRegex = /([\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: ... (no closing tag yet) + if (content.includes("") && !content.includes("")) { + const openMatch = content.match(/([\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 { - const map = new Map(); - 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("")); } +/** 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(""); +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 + * 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); } /**