Stage 400: PR #2710 — fix: render streamed math incrementally (no flash when delta completes a KaTeX expression)

Co-authored-by: Michaelyklam <Michaelyklam@users.noreply.github.com>

static/messages.js

@@ -595,6 +595,7 @@ function attachLiveStream(activeSid, streamId, uploaded=[], options={}){
   let _smdParser=null;     // current smd parser instance (null until first content)
   let _smdWrittenLen=0;    // how many chars of displayText have been fed to smd parser
   let _smdWrittenText='';  // exact displayText snapshot used for prefix-alignment checks
+  let _streamingKatexTimer=null; // throttles live KaTeX scans while smd writes deltas
   // On reconnect, the assistantBody already has partial smd-rendered content.
   // We clear it on first new token and restart the parser from the reconnect point.
   let _smdReconnect=reconnecting;
@@ -940,6 +941,7 @@ function attachLiveStream(activeSid, streamId, uploaded=[], options={}){
   }
   // Helper: end the current smd parser (flushes remaining state) and null it out.
   function _smdEndParser(){
+    if(_streamingKatexTimer){clearTimeout(_streamingKatexTimer);_streamingKatexTimer=null;}
     if(_smdParser&&window.smd){
       try{window.smd.parser_end(_smdParser);}catch(_){}
       // parser_end may flush remaining markdown that creates new links/images —
@@ -950,6 +952,13 @@ function attachLiveStream(activeSid, streamId, uploaded=[], options={}){
     _smdWrittenLen=0;
     _smdWrittenText='';
   }
+  function _scheduleStreamingKatex(){
+    if(_streamingKatexTimer) return;
+    _streamingKatexTimer=setTimeout(()=>{
+      _streamingKatexTimer=null;
+      if(assistantBody&&typeof renderKatexBlocks==='function') renderKatexBlocks(assistantBody);
+    },150);
+  }
   // Helper: feed new displayText delta to the smd parser.
   // Only feeds chars beyond what has already been written (_smdWrittenLen).
   function _smdWrite(displayText, fade=false){
@@ -974,6 +983,7 @@ function attachLiveStream(activeSid, streamId, uploaded=[], options={}){
     // streaming-markdown does NOT sanitize URL schemes. The default live path
     // scans after writes; fade mode blocks unsafe href/src in its renderer.set_attr.
     if(assistantBody&&!fade){_sanitizeSmdLinks(assistantBody);}
+    _scheduleStreamingKatex();
   }
   // Allowed URL schemes for anchors and images rendered from agent-streamed markdown.
   // Raw file:// anchors are rewritten to /api/media before the user can click them.

static/ui.js

@@ -7228,7 +7228,10 @@ let _katexReady=false;
 
 function renderKatexBlocks(container){
   const root=container||document;
-  const blocks=root.querySelectorAll('.katex-block:not([data-rendered]),.katex-inline:not([data-rendered])');
+  const blocks=root.querySelectorAll(
+    '.katex-block:not([data-rendered]),.katex-inline:not([data-rendered]),'+
+    'equation-block:not([data-rendered]),equation-inline:not([data-rendered])'
+  );
   if(!blocks.length) return;
   if(!_katexReady){
     if(!_katexLoading){
@@ -7250,7 +7253,8 @@ function renderKatexBlocks(container){
   blocks.forEach(el=>{
     el.dataset.rendered='true';
     const src=el.textContent||'';
-    const displayMode=el.dataset.katex==='display';
+    const tagName=(el.tagName||'').toLowerCase();
+    const displayMode=el.dataset.katex==='display'||tagName==='equation-block';
     try{
       katex.render(src,el,{
         displayMode,

tests/test_streaming_katex_live_render.py

@@ -0,0 +1,36 @@
+from pathlib import Path
+
+
+REPO = Path(__file__).resolve().parents[1]
+MESSAGES_JS = (REPO / "static" / "messages.js").read_text(encoding="utf-8")
+UI_JS = (REPO / "static" / "ui.js").read_text(encoding="utf-8")
+
+
+def test_live_smd_writes_schedule_incremental_katex_rendering():
+    """Live markdown deltas should render math before the terminal done event."""
+    assert "let _streamingKatexTimer=null" in MESSAGES_JS
+    assert "function _scheduleStreamingKatex()" in MESSAGES_JS
+    assert "setTimeout(()=>{" in MESSAGES_JS
+    assert "renderKatexBlocks(assistantBody)" in MESSAGES_JS
+
+    smd_write_idx = MESSAGES_JS.index("function _smdWrite(displayText, fade=false){")
+    done_idx = MESSAGES_JS.index("source.addEventListener('done'")
+    smd_write_block = MESSAGES_JS[smd_write_idx:done_idx]
+    assert "_scheduleStreamingKatex();" in smd_write_block
+
+
+def test_streaming_katex_timer_is_cleared_when_smd_parser_ends():
+    """The final done path should not leave a stale live KaTeX timer around."""
+    end_idx = MESSAGES_JS.index("function _smdEndParser(){")
+    write_idx = MESSAGES_JS.index("function _smdWrite(displayText, fade=false){")
+    end_block = MESSAGES_JS[end_idx:write_idx]
+    assert "if(_streamingKatexTimer){clearTimeout(_streamingKatexTimer);_streamingKatexTimer=null;}" in end_block
+
+
+def test_katex_renderer_scans_live_and_settled_unrendered_nodes_under_container():
+    assert "function renderKatexBlocks(container){" in UI_JS
+    assert "const root=container||document;" in UI_JS
+    assert ".katex-block:not([data-rendered]),.katex-inline:not([data-rendered])," in UI_JS
+    assert "equation-block:not([data-rendered]),equation-inline:not([data-rendered])" in UI_JS
+    assert "const tagName=(el.tagName||'').toLowerCase();" in UI_JS
+    assert "const displayMode=el.dataset.katex==='display'||tagName==='equation-block';" in UI_JS