From 20a900be609e3e9093b2159533bd794170bfc6b7 Mon Sep 17 00:00:00 2001 From: wlknight Date: Wed, 24 Jun 2026 06:09:44 +0800 Subject: [PATCH] fix(perf): eliminate O(DOM) streaming render freeze on long agent answers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduce _safeSmdRenderer — a renderer wrapper that hooks smd's set_attr to validate URL schemes inline as each DOM node is created, eliminating the per-token _sanitizeSmdLinks(assistantBody) full-DOM scan that caused progressive freezing on answers of 30,000+ chars (5,000+ DOM nodes). Also add parse-result caching to _scheduleRender to avoid redundant _extractInlineThinkingFromContent scans between the token event handler and the rAF render callback. Changes: - New _safeSmdRenderer(el): URL scheme safety inline via set_attr hook (identical algorithm to proven _streamFadeRenderer.set_attr, no animation) - _smdNewParser: non-fade path uses _safeSmdRenderer instead of bare default_renderer - _smdWrite: removed _sanitizeSmdLinks(assistantBody) call — safety now handled by renderer's set_attr at node-creation time - _scheduleRender(parsed): accepts optional cached parse result from token handler to avoid duplicate O(n) text scans - Token handler: passes _parseStreamState() result to _scheduleRender _sanitizeSmdLinks is preserved at _smdEndParser (stream-end safety net) and in fallback paths where smd is not available. All 166 existing tests pass. --- static/messages.js | 55 +++++++++++++++++++++++++---- tests/test_streaming_markdown.py | 59 +++++++++++++++++++++++++------- 2 files changed, 94 insertions(+), 20 deletions(-) diff --git a/static/messages.js b/static/messages.js index 3888c7627..6d6104c4c 100644 --- a/static/messages.js +++ b/static/messages.js @@ -3001,7 +3001,7 @@ function attachLiveStream(activeSid, streamId, uploaded=[], options={}){ _smdWrittenLen=0; _smdWrittenText=''; if(!window.smd){_smdParser=null;return;} - const baseRenderer=fade ? _streamFadeRenderer(el) : window.smd.default_renderer(el); + const baseRenderer=fade ? _streamFadeRenderer(el) : _safeSmdRenderer(el); const renderer=_smdRendererWithoutUnderscoreEmphasis(baseRenderer); _smdParser=window.smd.parser(renderer); } @@ -3072,9 +3072,9 @@ function attachLiveStream(activeSid, streamId, uploaded=[], options={}){ try{window.smd.parser_write(_smdParser,delta);}catch(_){} _smdWrittenLen=displayText.length; _smdWrittenText=displayText; - // 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);} + // URL scheme safety is handled by the renderer's set_attr hook + // (_safeSmdRenderer or _streamFadeRenderer), applied inline as smd + // creates each DOM node — no post-hoc full-DOM scan needed. _scheduleStreamingKatex(); } // Allowed URL schemes for anchors and images rendered from agent-streamed markdown. @@ -3245,6 +3245,36 @@ function attachLiveStream(activeSid, streamId, uploaded=[], options={}){ }; return renderer; } + // Safe renderer: wraps default_renderer with a set_attr hook that validates + // href/src URL schemes inline — no post-hoc DOM-wide querySelectorAll needed. + // Unlike _streamFadeRenderer, this does NOT wrap add_text, so smd adds new + // DOM nodes as plain text nodes (no animation spans). Used on the non-fade + // streaming path to eliminate _sanitizeSmdLinks(assistantBody) O(DOM) scans + // on every token event (#WebUI-perf). + function _safeSmdRenderer(el){ + const renderer=window.smd.default_renderer(el); + const baseSetAttr=renderer.set_attr; + renderer.set_attr=(data,attr,value)=>{ + const isHref=window.smd&&attr===window.smd.HREF; + const isSrc=window.smd&&attr===window.smd.SRC; + const safeUrl=isSrc?_SMD_SAFE_IMG_URL_RE:_SMD_SAFE_URL_RE; + if(isHref&&/^(file|workspace|session):\/\//i.test(String(value||''))){ + baseSetAttr(data,attr,_smdLinkHref(value)); + if(/^session:\/\//i.test(String(value||''))){ + const node=data&&data.nodes&&data.nodes[data.index]; + if(node&&node.classList) node.classList.add('session-link'); + } + return; + } + if((isHref||isSrc)&&!safeUrl.test(String(value||''))){ + const node=data&&data.nodes&&data.nodes[data.index]; + if(node&&node.setAttribute) node.setAttribute('data-blocked-scheme','1'); + return; + } + baseSetAttr(data,attr,value); + }; + return renderer; + } function _streamFadeWordCountOf(text){ const m=String(text||'').match(/\S+/g); return m?m.length:0; @@ -3714,7 +3744,17 @@ function attachLiveStream(activeSid, streamId, uploaded=[], options={}){ } let _lastRenderMs=0; - function _scheduleRender(){ + // Parse-result cache: _scheduleRender can accept a pre-computed _parseStreamState() + // from the token event handler, avoiding a duplicate O(n) scan inside _doRender + // when the rAF fires before the next token arrives. + let _cachedParsed=null; + let _cachedParsedText=''; + function _scheduleRender(parsed){ + // If caller provides a pre-computed parse result, cache it for _doRender. + if(parsed){ + _cachedParsed=parsed; + _cachedParsedText=assistantText; + } if(_renderPending) return; if(_streamFinalized) return; // Bug A: don't schedule new rAF after stream finalized _renderPending=true; @@ -3733,7 +3773,8 @@ function attachLiveStream(activeSid, streamId, uploaded=[], options={}){ // Guard: a pending setTimeout+rAF can outlive stream finalization. if(_streamFinalized) return; _lastRenderMs=performance.now(); - const parsed=_parseStreamState(); + const parsed=_cachedParsed&&_cachedParsedText===assistantText ? _cachedParsed : _parseStreamState(); + _cachedParsed=null; _renderLiveThinking(parsed); if(assistantBody){ const displayText = segmentStart===0 @@ -3828,7 +3869,7 @@ function attachLiveStream(activeSid, streamId, uploaded=[], options={}){ const parsed=_parseStreamState(); if(_freshSegment) appendThinking('', _liveThinkingPlacement()); if(String((parsed&&parsed.displayText)||'').trim()||assistantRow) ensureAssistantRow(); - _scheduleRender(); + _scheduleRender(parsed); }); source.addEventListener('interim_assistant',e=>{ diff --git a/tests/test_streaming_markdown.py b/tests/test_streaming_markdown.py index 2e75c3800..e5449dc6a 100644 --- a/tests/test_streaming_markdown.py +++ b/tests/test_streaming_markdown.py @@ -197,10 +197,18 @@ class TestSmdHelpers: "_smdWrittenLen=0" in fn or "_smdWrittenLen = 0" in fn ), "_smdNewParser must reset _smdWrittenLen to 0" - def test_smd_new_parser_calls_default_renderer(self): + def test_smd_new_parser_calls_safe_renderer(self): fn = extract_fn(MESSAGES_JS, "_smdNewParser") - assert fn and "default_renderer" in fn, ( - "_smdNewParser must call smd.default_renderer() to create a renderer" + assert fn and ( + "_safeSmdRenderer(" in fn or "_streamFadeRenderer(" in fn + ), ( + "_smdNewParser must use _safeSmdRenderer or _streamFadeRenderer " + "so URL scheme safety is applied inline via set_attr hook" + ) + # Verify _safeSmdRenderer itself uses smd.default_renderer internally + safefn = extract_fn(MESSAGES_JS, "_safeSmdRenderer") + assert safefn and "default_renderer" in safefn, ( + "_safeSmdRenderer must call smd.default_renderer() to create the base renderer" ) def test_smd_new_parser_calls_parser(self): @@ -706,16 +714,41 @@ class TestSmdUrlSchemeSanitization: assert "_smdFileHref" in MESSAGES_JS assert "api/media?path=" in MESSAGES_JS - def test_sanitize_called_after_smd_write(self): - # _smdWrite must invoke _sanitizeSmdLinks on assistantBody after feeding the parser, - # so anchors/images created mid-stream get their javascript:/data:/vbscript: - # hrefs/srcs stripped before the user can click them. - fn = extract_fn(MESSAGES_JS, "_smdWrite") - assert fn, "_smdWrite function not found" - assert "_sanitizeSmdLinks" in fn, ( - "_smdWrite must call _sanitizeSmdLinks(assistantBody) after parser_write " - "so unsafe URL schemes are stripped from newly-added anchors/images " - "before the user can click them" + def test_url_safety_via_renderer_set_attr(self): + # URL scheme safety is now enforced inline by the renderer's set_attr + # hook (_safeSmdRenderer or _streamFadeRenderer) as smd creates each + # DOM node, eliminating the post-hoc _sanitizeSmdLinks O(DOM) scan + # per token that caused progressive streaming freeze on long answers. + safefn = extract_fn(MESSAGES_JS, "_safeSmdRenderer") + assert safefn, "_safeSmdRenderer must exist for URL safety on the non-fade path" + assert "set_attr" in safefn, ( + "_safeSmdRenderer must override set_attr to validate href/src inline" + ) + assert "_SMD_SAFE_URL_RE" in safefn, ( + "_safeSmdRenderer set_attr must use _SMD_SAFE_URL_RE for href safety" + ) + assert "_SMD_SAFE_IMG_URL_RE" in safefn, ( + "_safeSmdRenderer set_attr must use _SMD_SAFE_IMG_URL_RE for src safety" + ) + assert "data-blocked-scheme" in safefn, ( + "_safeSmdRenderer set_attr must set data-blocked-scheme on unsafe URLs" + ) + # _safeSmdRenderer algorithm must be the same as the proven + # _streamFadeRenderer set_attr (fade path has been in production). + fadefn = extract_fn(MESSAGES_JS, "_streamFadeRenderer") + assert fadefn and "set_attr" in fadefn, "_streamFadeRenderer must also have set_attr" + # Both renderers go through _smdRendererWithoutUnderscoreEmphasis so + # the set_attr hook is part of the final parser — verify the call chain. + newparser = extract_fn(MESSAGES_JS, "_smdNewParser") + assert newparser and "fade ? _streamFadeRenderer(el) : _safeSmdRenderer(el)" in newparser, ( + "_smdNewParser must route non-fade to _safeSmdRenderer and fade to _streamFadeRenderer" + ) + # _sanitizeSmdLinks is still defined and used at parser_end as a final + # safety net — not removed, just no longer called on every token. + assert "_sanitizeSmdLinks" in MESSAGES_JS, "_sanitizeSmdLinks must still exist" + endfn = extract_fn(MESSAGES_JS, "_smdEndParser") + assert endfn and "_sanitizeSmdLinks" in endfn, ( + "_smdEndParser must still call _sanitizeSmdLinks as a final safety net" ) def test_sanitize_called_at_parser_end(self):