Merge pull request #27 from LofiSu/experimental

Experimental

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