Chat widget: markdown styling, KaTeX math rendering, SSE streaming fix

- Add comprehensive CSS for markdown elements in chat bubbles (headers, lists, tables, blockquotes, links, code blocks, horizontal rules) - Integrate KaTeX via marked-katex-extension for inline ($) and display ($$) math rendering with DOMPurify whitelist for MathML/SVG elements - Add normalizeDisplayMath() pre-processor to handle LLM output where $$ delimiters appear inline rather than on their own lines - Add dark theme KaTeX styles matching the SpiceBook color palette - Fix asyncio.to_thread usage in chat SSE endpoint for streaming
2026-02-23 14:25:42 -07:00 · 2026-02-23 14:25:42 -07:00 · 6f239c185e
commit 6f239c185e
parent 99d1ca28d2
5 changed files with 260 additions and 25 deletions
--- a/backend/src/spicebook/routers/chat.py
+++ b/backend/src/spicebook/routers/chat.py
@ -23,12 +23,12 @@ def _sse_event(event: str, data: dict | list) -> str:
    return f"event: {event}\ndata: {payload}\n\n"


-def _build_notebook_context(req: ChatStreamRequest) -> str:
+async def _build_notebook_context(req: ChatStreamRequest) -> str:
    """Extract SPICE cells and markdown notes from the notebook for LLM context."""
    if not req.notebook or not req.notebook.notebook_id:
        return ""

-    nb = load_notebook(settings.notebook_dir, req.notebook.notebook_id)
+    nb = await asyncio.to_thread(load_notebook, settings.notebook_dir, req.notebook.notebook_id)
    if nb is None:
        return ""

@ -64,7 +64,7 @@ async def chat_stream(req: ChatStreamRequest):

    async def generate():
        # Build context from the notebook if available
-        context = _build_notebook_context(req)
+        context = await _build_notebook_context(req)

        question = req.question
        if req.notebook and req.notebook.title:
--- a/frontend/package-lock.json
+++ b/frontend/package-lock.json
@ -28,7 +28,10 @@
        "astro-seo-meta": "^5.2.0",
        "clsx": "^2.1.0",
        "dompurify": "^3.3.1",
+        "katex": "^0.16.33",
        "lucide-react": "^0.468.0",
+        "marked": "^17.0.3",
+        "marked-katex-extension": "^5.1.7",
        "react": "^19.0.0",
        "react-dom": "^19.0.0",
        "satori": "^0.19.2",
@ -6123,6 +6126,31 @@
      "integrity": "sha512-H8jvkz1O50L3dMZCsLqiuB2tA7muqbSg1AtGEkN0leAqGjsUzDJir3Zwr02BhqdcITPg3ei3mZ+HjMocAknhhg==",
      "license": "MIT"
    },
+    "node_modules/katex": {
+      "version": "0.16.33",
+      "resolved": "https://registry.npmjs.org/katex/-/katex-0.16.33.tgz",
+      "integrity": "sha512-q3N5u+1sY9Bu7T4nlXoiRBXWfwSefNGoKeOwekV+gw0cAXQlz2Ww6BLcmBxVDeXBMUDQv6fK5bcNaJLxob3ZQA==",
+      "funding": [
+        "https://opencollective.com/katex",
+        "https://github.com/sponsors/katex"
+      ],
+      "license": "MIT",
+      "dependencies": {
+        "commander": "^8.3.0"
+      },
+      "bin": {
+        "katex": "cli.js"
+      }
+    },
+    "node_modules/katex/node_modules/commander": {
+      "version": "8.3.0",
+      "resolved": "https://registry.npmjs.org/commander/-/commander-8.3.0.tgz",
+      "integrity": "sha512-OkTL9umf+He2DZkUq8f8J9of7yL6RJKI24dVITBmNfZBmri9zYZQrKkuXiKhyfPSu8tUhnVBB1iKXevvnlR4Ww==",
+      "license": "MIT",
+      "engines": {
+        "node": ">= 12"
+      }
+    },
    "node_modules/kleur": {
      "version": "3.0.3",
      "resolved": "https://registry.npmjs.org/kleur/-/kleur-3.0.3.tgz",
@ -6479,6 +6507,28 @@
        "url": "https://github.com/sponsors/wooorm"
      }
    },
+    "node_modules/marked": {
+      "version": "17.0.3",
+      "resolved": "https://registry.npmjs.org/marked/-/marked-17.0.3.tgz",
+      "integrity": "sha512-jt1v2ObpyOKR8p4XaUJVk3YWRJ5n+i4+rjQopxvV32rSndTJXvIzuUdWWIy/1pFQMkQmvTXawzDNqOH/CUmx6A==",
+      "license": "MIT",
+      "bin": {
+        "marked": "bin/marked.js"
+      },
+      "engines": {
+        "node": ">= 20"
+      }
+    },
+    "node_modules/marked-katex-extension": {
+      "version": "5.1.7",
+      "resolved": "https://registry.npmjs.org/marked-katex-extension/-/marked-katex-extension-5.1.7.tgz",
+      "integrity": "sha512-CVFzrqwpXGVaHByqcVvO/JfzW/OMWrAF3pEfNYNIruzBzM64moANSHapCg1qbzEN+NGf5unHwkMfwJIXHzyDAw==",
+      "license": "MIT",
+      "peerDependencies": {
+        "katex": ">=0.16 <0.17",
+        "marked": ">=4 <18"
+      }
+    },
    "node_modules/mdast-util-definitions": {
      "version": "6.0.0",
      "resolved": "https://registry.npmjs.org/mdast-util-definitions/-/mdast-util-definitions-6.0.0.tgz",
--- a/frontend/package.json
+++ b/frontend/package.json
@ -29,7 +29,10 @@
    "astro-seo-meta": "^5.2.0",
    "clsx": "^2.1.0",
    "dompurify": "^3.3.1",
+    "katex": "^0.16.33",
    "lucide-react": "^0.468.0",
+    "marked": "^17.0.3",
+    "marked-katex-extension": "^5.1.7",
    "react": "^19.0.0",
    "react-dom": "^19.0.0",
    "satori": "^0.19.2",
--- a/frontend/src/components/chat/ChatWidget.tsx
+++ b/frontend/src/components/chat/ChatWidget.tsx
@ -12,33 +12,63 @@ import { useChatStore } from '../../lib/chat-store';
 import type { ChatMessage } from '../../lib/chat-store';
 import { streamChat } from '../../lib/chat-api';
 import { useNotebookStore } from '../../lib/notebook-store';
+import { marked } from 'marked';
+import markedKatex from 'marked-katex-extension';
+import DOMPurify from 'dompurify';
+import 'katex/dist/katex.min.css';
 import '../../styles/chat-widget.css';

-// ── Lightweight markdown ────────────────────────────────
+// ── Markdown rendering (marked + KaTeX + DOMPurify) ──────

-function escapeHtml(text: string): string {
-  return text
-    .replace(/&/g, '&amp;')
-    .replace(/</g, '&lt;')
-    .replace(/>/g, '&gt;')
-    .replace(/"/g, '&quot;');
+marked.setOptions({
+  breaks: true,
+  gfm: true,
+});
+
+marked.use(markedKatex({
+  throwOnError: false,
+  nonStandard: true,
+}));
+
+// KaTeX generates SVG + spans with specific classes/attributes.
+// DOMPurify must allow these through.
+const KATEX_TAGS = [
+  'math', 'semantics', 'mrow', 'mi', 'mo', 'mn', 'msup', 'msub',
+  'mfrac', 'munderover', 'mover', 'munder', 'msqrt', 'mroot',
+  'mtable', 'mtr', 'mtd', 'mspace', 'mtext', 'menclose',
+  'annotation', 'annotation-xml',
+];
+
+const KATEX_ATTRS = [
+  'xmlns', 'encoding', 'mathvariant', 'displaystyle', 'scriptlevel',
+  'fence', 'stretchy', 'symmetric', 'lspace', 'rspace',
+  'linethickness', 'columnalign', 'rowalign', 'columnspacing',
+  'rowspacing', 'columnlines', 'rowlines', 'frame', 'framespacing',
+  'width', 'height', 'voffset', 'accent', 'accentunder',
+  'notation', 'minsize', 'maxsize', 'movablelimits',
+  'aria-hidden', 'focusable', 'role', 'tabindex',
+  'viewBox', 'preserveAspectRatio', 'd', 'fill', 'stroke',
+  'stroke-width', 'stroke-linecap', 'stroke-linejoin',
+  'transform', 'clip-path',
+];
+
+// LLMs often emit $$...$$ inline (e.g. "is:$$\frac{1}{s}$$\nNext")
+// but marked-katex-extension requires $$ on its own line for display mode.
+// This normalizer ensures $$ delimiters get their own lines.
+function normalizeDisplayMath(text: string): string {
+  return text.replace(/\$\$([\s\S]*?)\$\$/g, (_match, inner: string) => {
+    const trimmed = inner.trim();
+    return `\n$$\n${trimmed}\n$$\n`;
+  });
 }

 function renderMarkdown(text: string): string {
-  let html = escapeHtml(text);
-  // Code blocks (``` ... ```)
-  html = html.replace(/```(\w*)\n([\s\S]*?)```/g, (_m, _lang, code) => {
-    return `<pre><code>${code.trim()}</code></pre>`;
+  const normalized = normalizeDisplayMath(text);
+  const raw = marked.parse(normalized, { async: false }) as string;
+  return DOMPurify.sanitize(raw, {
+    ADD_TAGS: KATEX_TAGS,
+    ADD_ATTR: [...KATEX_ATTRS, 'target', 'class', 'style'],
  });
-  html = html.replace(/\*\*(.+?)\*\*/g, '<strong>$1</strong>');
-  html = html.replace(/(?<!\*)\*([^*]+?)\*(?!\*)/g, '<em>$1</em>');
-  html = html.replace(/`([^`]+?)`/g, '<code>$1</code>');
-  html = html.replace(
-    /\[([^\]]+?)\]\((https?:\/\/[^)]+?)\)/g,
-    '<a href="$2" target="_blank" rel="noopener noreferrer">$1</a>',
-  );
-  html = html.replace(/\n/g, '<br>');
-  return html;
 }

 // ── Components ──────────────────────────────────────────
@ -158,6 +188,12 @@ export default function ChatWidget() {
  const inputRef = useRef<HTMLInputElement>(null);
  const abortRef = useRef<AbortController | null>(null);

+  // Token batching: accumulate tokens in a ref, flush to state once per animation frame.
+  // Without this, React 19 batches all rapid set() calls from the SSE loop into one
+  // render at stream end — so the user sees nothing until the LLM finishes.
+  const pendingTokensRef = useRef('');
+  const flushRafRef = useRef(0);
+
  const activeConv = getActiveConversation();
  const messages = activeConv?.messages ?? [];

@ -225,8 +261,18 @@ export default function ChatWidget() {
            setStatusText((evt.data as { text: string }).text);
            break;
          case 'token':
-            appendToLastAssistant((evt.data as { text: string }).text);
-            setStatusText('');
+            // Accumulate in ref; flush to Zustand once per animation frame
+            pendingTokensRef.current += (evt.data as { text: string }).text;
+            if (!flushRafRef.current) {
+              flushRafRef.current = requestAnimationFrame(() => {
+                if (pendingTokensRef.current) {
+                  appendToLastAssistant(pendingTokensRef.current);
+                  pendingTokensRef.current = '';
+                }
+                flushRafRef.current = 0;
+              });
+              setStatusText('');
+            }
            break;
          case 'reasoning':
            setReasoningText((prev) => prev + (evt.data as { text: string }).text);
@ -250,6 +296,15 @@ export default function ChatWidget() {
        );
      }
    } finally {
+      // Flush any remaining buffered tokens
+      if (flushRafRef.current) {
+        cancelAnimationFrame(flushRafRef.current);
+        flushRafRef.current = 0;
+      }
+      if (pendingTokensRef.current) {
+        appendToLastAssistant(pendingTokensRef.current);
+        pendingTokensRef.current = '';
+      }
      setStreaming(false);
      setStatusText('');
      abortRef.current = null;
--- a/frontend/src/styles/chat-widget.css
+++ b/frontend/src/styles/chat-widget.css
@ -174,6 +174,133 @@
  padding: 0;
 }

+/* Markdown headings */
+.chat-bubble.assistant h1,
+.chat-bubble.assistant h2,
+.chat-bubble.assistant h3,
+.chat-bubble.assistant h4 {
+  color: var(--color-sb-text-bright);
+  font-weight: 600;
+  margin: 0.75rem 0 0.25rem;
+  line-height: 1.3;
+}
+
+.chat-bubble.assistant h1 { font-size: 1rem; }
+.chat-bubble.assistant h2 { font-size: 0.9375rem; }
+.chat-bubble.assistant h3 { font-size: 0.875rem; }
+.chat-bubble.assistant h4 { font-size: 0.8125rem; }
+
+.chat-bubble.assistant h1:first-child,
+.chat-bubble.assistant h2:first-child,
+.chat-bubble.assistant h3:first-child,
+.chat-bubble.assistant h4:first-child {
+  margin-top: 0;
+}
+
+/* Paragraphs */
+.chat-bubble.assistant p {
+  margin: 0.375rem 0;
+}
+
+.chat-bubble.assistant p:first-child {
+  margin-top: 0;
+}
+
+.chat-bubble.assistant p:last-child {
+  margin-bottom: 0;
+}
+
+/* Lists */
+.chat-bubble.assistant ul,
+.chat-bubble.assistant ol {
+  margin: 0.375rem 0;
+  padding-left: 1.25rem;
+}
+
+.chat-bubble.assistant li {
+  margin: 0.125rem 0;
+}
+
+.chat-bubble.assistant li > ul,
+.chat-bubble.assistant li > ol {
+  margin: 0.125rem 0;
+}
+
+/* Horizontal rules */
+.chat-bubble.assistant hr {
+  border: none;
+  border-top: 1px solid var(--color-sb-border);
+  margin: 0.625rem 0;
+}
+
+/* Blockquotes */
+.chat-bubble.assistant blockquote {
+  border-left: 3px solid var(--color-sb-accent);
+  margin: 0.5rem 0;
+  padding: 0.375rem 0.625rem;
+  background: rgba(59, 130, 246, 0.08);
+  border-radius: 0 0.25rem 0.25rem 0;
+  color: var(--color-sb-text);
+}
+
+.chat-bubble.assistant blockquote p {
+  margin: 0.25rem 0;
+}
+
+/* Tables — wrapped in a scrollable container via overflow on the bubble */
+.chat-bubble.assistant table {
+  width: 100%;
+  border-collapse: collapse;
+  margin: 0.5rem 0;
+  font-size: 0.75rem;
+  display: block;
+  overflow-x: auto;
+}
+
+.chat-bubble.assistant th,
+.chat-bubble.assistant td {
+  padding: 0.3rem 0.5rem;
+  border: 1px solid var(--color-sb-border);
+  text-align: left;
+}
+
+.chat-bubble.assistant th {
+  background: var(--color-sb-surface);
+  color: var(--color-sb-text-bright);
+  font-weight: 600;
+}
+
+.chat-bubble.assistant tr:nth-child(even) {
+  background: rgba(255, 255, 255, 0.02);
+}
+
+/* Links */
+.chat-bubble.assistant a {
+  color: var(--color-sb-accent);
+  text-decoration: none;
+}
+
+.chat-bubble.assistant a:hover {
+  text-decoration: underline;
+}
+
+/* KaTeX math — dark theme overrides */
+.chat-bubble.assistant .katex {
+  font-size: 0.9em;
+  color: var(--color-sb-text-bright);
+}
+
+.chat-bubble.assistant .katex-display {
+  margin: 0.5rem 0;
+  overflow-x: auto;
+  overflow-y: hidden;
+  padding: 0.375rem 0;
+}
+
+.chat-bubble.assistant .katex-display > .katex {
+  white-space: normal;
+}
+
 /* Status messages */
 .chat-status {
  text-align: center;