Harden citation token and mermaid loose chart parsing

Author: owenisas

frontend/src/components/ui/RichText.tsx

@@ -67,13 +67,13 @@ function parseMermaidXyChart(code: string): ParsedGraph | null {
 
   const xAxisMatch = code.match(/^\s*x-axis\s+(.+)$/im)
   const yAxisMatch = code.match(/^\s*y-axis\s+"([^"]+)"/im)
-  const barMatch = code.match(/^\s*bar\s+\[([^\]]+)\]/im)
-  const lineMatch = code.match(/^\s*line\s+\[([^\]]+)\]/im)
+  const barMatch = code.match(/^\s*bar(?:\s+"([^"]+)")?\s+\[([^\]]+)\]/im)
+  const lineMatch = code.match(/^\s*line(?:\s+"([^"]+)")?\s+\[([^\]]+)\]/im)
 
   const values = barMatch
-    ? parseMermaidNumbers(barMatch[1])
+    ? parseMermaidNumbers(barMatch[2])
     : lineMatch
-      ? parseMermaidNumbers(lineMatch[1])
+      ? parseMermaidNumbers(lineMatch[2])
       : []
   if (values.length === 0) return null
 
@@ -86,7 +86,7 @@ function parseMermaidXyChart(code: string): ParsedGraph | null {
     type: barMatch ? "bar" : "line",
     labels: labels.slice(0, count),
     series: [{
-      name: yAxisMatch?.[1]?.trim() || "Value",
+      name: barMatch?.[1]?.trim() || lineMatch?.[1]?.trim() || yAxisMatch?.[1]?.trim() || "Value",
       values: values.slice(0, count),
     }],
   }
@@ -123,11 +123,21 @@ function parseMermaidGraph(code: string): ParsedGraph | null {
 
 const OPENAI_CITATION_OPEN = "\uE200cite\uE202"
 const OPENAI_CITATION_CLOSE = "\uE201"
+const OPENAI_CITATION_OPEN_ALT = "\u{F0000}cite\u{F0002}"
+const OPENAI_CITATION_CLOSE_ALT = "\u{F0001}"
 
 function parseOpenAICitationToken(text: string, from: number): { length: number; ids: number[] } | null {
-  if (!text.startsWith(OPENAI_CITATION_OPEN, from)) return null
-  const payloadStart = from + OPENAI_CITATION_OPEN.length
-  const end = text.indexOf(OPENAI_CITATION_CLOSE, payloadStart)
+  let payloadStart = -1
+  let end = -1
+  if (text.startsWith(OPENAI_CITATION_OPEN, from)) {
+    payloadStart = from + OPENAI_CITATION_OPEN.length
+    end = text.indexOf(OPENAI_CITATION_CLOSE, payloadStart)
+  } else if (text.startsWith(OPENAI_CITATION_OPEN_ALT, from)) {
+    payloadStart = from + OPENAI_CITATION_OPEN_ALT.length
+    end = text.indexOf(OPENAI_CITATION_CLOSE_ALT, payloadStart)
+  } else {
+    return null
+  }
   if (end === -1) return null
 
   const payload = text.slice(payloadStart, end)
@@ -148,20 +158,34 @@ function parseOpenAICitationToken(text: string, from: number): { length: number;
 
 function parseLooseMermaidChartBlock(lines: string[], startLine: number): { graph: ParsedGraph; endLine: number; snippet: string } | null {
   const start = lines[startLine]?.trim() ?? ""
-  const maybeMermaidStart = /^xychart-beta\b/i.test(start) || /^pie(?:\s+showData)?$/i.test(start)
+  const isXyStart = /^xychart-beta\b/i.test(start)
+  const isPieStart = /^pie(?:\s+showData)?$/i.test(start)
+  const maybeMermaidStart = isXyStart || isPieStart
   if (!maybeMermaidStart) return null
+  const mode: "xy" | "pie" = isXyStart ? "xy" : "pie"
+
+  const isChartLine = (t: string) => {
+    if (!t) return true
+    if (mode === "xy") return /^(title|x-axis|y-axis|line|bar)\b/i.test(t)
+    return /^(title\b|(?:"[^"]+"|'[^']+'|[^:]+)\s*:\s*-?\d+(?:\.\d+)?)$/i.test(t)
+  }
 
   const collected: string[] = [lines[startLine]]
   let endLine = startLine
   for (let i = startLine + 1; i < lines.length; i++) {
     const t = lines[i].trim()
-    if (!t) break
+    if (!t) {
+      collected.push(lines[i])
+      endLine = i
+      continue
+    }
     if (/^#{1,6}\s/.test(t)) break
     if (/^```/.test(t)) break
     if (/^[-*]\s+/.test(t)) break
     if (/^\d+[.)]\s+/.test(t)) break
     if (parseGraphTag(t)) break
     if (parseTableLine(lines[i])) break
+    if (!isChartLine(t)) break
     collected.push(lines[i])
     endLine = i
   }
@@ -324,7 +348,10 @@ function renderInline(text: string, variant: Variant, COLORS: ThemeColors, VS: R
   let key = 0
 
   while (cursor < text.length) {
-    const nextOpenAICitation = findNextIndex(text, cursor, OPENAI_CITATION_OPEN)
+    const nextOpenAICitation = Math.min(
+      findNextIndex(text, cursor, OPENAI_CITATION_OPEN),
+      findNextIndex(text, cursor, OPENAI_CITATION_OPEN_ALT),
+    )
     const nextCode = findNextIndex(text, cursor, "`")
     const nextBoldItalic = allowBold ? findNextIndex(text, cursor, "***") : Number.POSITIVE_INFINITY
     const nextBold = allowBold ? findNextIndex(text, cursor, "**") : Number.POSITIVE_INFINITY