fix(chat): recover viewport anchor via sessionIdx when content-derived key goes stale

Clean rebase of allenliang2022's #5469 (rebase-first).

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

static/ui.js

@@ -868,13 +868,27 @@ function _restoreMessageViewportAnchor(anchor, rawIdxDelta){
   const container=$('messages');
   if(!container||!anchor) return false;
   const anchorKey=String(anchor.key||'');
-  let row=anchorKey?Array.from(container.querySelectorAll('[data-message-anchor-key]')).find(el=>el&&el.dataset&&el.dataset.messageAnchorKey===anchorKey):null;
-  if(row&&row.getClientRects&&row.getClientRects().length===0) row=null;
-  if(!row&&anchorKey) return false;
   const sessionIdx=Number(anchor.sessionIdx);
   const hasSessionIdx=Number.isFinite(sessionIdx);
+  let row=anchorKey?Array.from(container.querySelectorAll('[data-message-anchor-key]')).find(el=>el&&el.dataset&&el.dataset.messageAnchorKey===anchorKey):null;
+  if(row&&row.getClientRects&&row.getClientRects().length===0) row=null;
+  // The anchor key is content-derived (role|ts|attachments|first-160-chars, built by
+  // _messageViewportAnchorKeyForMessage) so it goes STALE while a live assistant
+  // message is still streaming: every chunk that changes the first 160 chars
+  // recomputes that row's data-message-anchor-key, so a snapshot captured mid-stream
+  // no longer matches by key. We used to concede the moment the keyed lookup missed
+  // (`if(!row&&anchorKey) return false`), and the caller then fell back to an ABSOLUTE
+  // scrollTop=snapshot.top that does NOT compensate the above-viewport height growth
+  // from that same streaming chunk — the residual DESKTOP scroll jump-back. (Desktop
+  // rests at overflow-anchor:none, so #5392's mobile overflow-anchor guard is a no-op
+  // here; this is a distinct code path.) The anchored row is still in the DOM under
+  // its STABLE session-relative index, so recover it via sessionIdx before conceding.
+  // A genuinely removed anchor (message compressed/deleted away) misses key AND
+  // sessionIdx and still returns false. A missing sessionIdx is NOT degraded to the
+  // window-relative rawIdx (which could resolve to a different message), preserving
+  // the original per-tier guard.
   if(!row&&hasSessionIdx) row=container.querySelector(`[data-session-msg-idx="${sessionIdx}"]`);
-  if(!row&&hasSessionIdx) return false;
+  if(!row&&(anchorKey||hasSessionIdx)) return false;
   const targetIdx=Number(anchor.rawIdx)+Number(rawIdxDelta||0);
   if(!row&&Number.isFinite(targetIdx)) row=container.querySelector(`[data-msg-idx="${targetIdx}"]`);
   if(!row) return false;

tests/test_issue4295_midstream_scroll_anchor.py

@@ -467,3 +467,99 @@ def test_same_frame_restore_uses_the_shared_anchor_remount_fallback():
 
     assert "if(!restoredViaAnchor&&typeof_remountMessageViewportAnchor==='function'&&_remountMessageViewportAnchor(snapshot.anchor))" in compact
     assert "_messageVirtualScrollTopForVisibleIdx" not in body
+
+
+def test_stale_content_key_recovers_via_session_index_and_compensates_height():
+    """Desktop scroll jump-back root fix (residual, distinct from #5392 mobile path).
+
+    The viewport anchor key is content-derived (role|ts|attachments|first-160-chars),
+    so while a live assistant message is streaming, each chunk that changes the first
+    160 chars recomputes that row's data-message-anchor-key. A scroll snapshot captured
+    a chunk earlier now has a STALE key that no longer matches by key lookup.
+
+    The OLD code conceded immediately on a stale key (`if(!row&&anchorKey) return false`),
+    and the caller then fell back to an ABSOLUTE scrollTop that does not compensate the
+    above-viewport height change from that same chunk -> desktop jump-back. The row is
+    still in the DOM under its stable data-session-msg-idx, so the fix recovers it via
+    sessionIdx and applies the relative-offset realign (which compensates the height).
+
+    Behavior test (not a source-string check): the anchor row's on-screen top must be
+    held at its captured topOffset. With buggy code the keyed lookup misses, sessionIdx
+    is never consulted, _restoreMessageViewportAnchor returns false, and this asserts FAIL.
+    """
+
+    script = f"""
+const assert = require('assert');
+let _programmaticScroll = false;
+let _programmaticScrollSetAt = 0;
+// Scenario: reader is unpinned mid-transcript. A streaming chunk (a) changed the
+// anchor row's content-derived key (stale key -> keyed lookup misses) and (b) grew
+// content above the viewport. We assert the realign still runs (via sessionIdx) and
+// writes the exact relative-offset compensation instead of conceding to an absolute
+// jump. Geometry is chosen so the write is unambiguous:
+//   containerRect.top = 100, captured topOffset = 300 (row was 300px below container top)
+//   after the chunk, the row's live rect.top = 250 => (250-100) = 150 from container top
+//   realign write: scrollTop += (150) - 300 = -150  => 1200 - 150 = 1050
+// A buggy body dead-ends on the stale key (returns false, no scrollTop write at all),
+// so both asserts below FAIL on the old code.
+const anchorRow = {{
+  getBoundingClientRect() {{ return {{ top: 250, bottom: 400 }}; }}
+}};
+const container = {{
+  scrollTop: 1200,
+  scrollHeight: 4600,
+  clientHeight: 600,
+  getBoundingClientRect() {{ return {{ top: 100, bottom: 700 }}; }},
+  querySelectorAll(selector) {{
+    // Stale key: NO node matches the captured (old) content key.
+    if (selector === '[data-message-anchor-key]') return [];
+    return [];
+  }},
+  querySelector(selector) {{
+    // The row is still present under its stable session index (and raw index).
+    if (selector === '[data-session-msg-idx="3"]') return anchorRow;
+    if (selector === '[data-msg-idx="3"]') return anchorRow;
+    return null;
+  }}
+}};
+function $(id) {{ return id === 'messages' ? container : null; }}
+function requestAnimationFrame(fn) {{ fn(); }}
+function setTimeout(fn) {{ fn(); }}
+const performance = {{ now() {{ return 0; }} }};
+function _suppressBrowserOverflowAnchor() {{ return null; }}  // desktop no-op
+function _deferClearProgrammaticScroll() {{ _programmaticScroll = false; }}
+{_function_body(UI_JS, "_restoreMessageViewportAnchor")}
+// Captured snapshot: anchor row was at topOffset=300, content key now STALE, but the
+// row still exists at sessionIdx=3 / rawIdx=3.
+const ok = _restoreMessageViewportAnchor({{
+  key: 'assistant|200|0|STALE_first_160_chars_from_an_earlier_chunk',
+  sessionIdx: 3,
+  rawIdx: 3,
+  topOffset: 300,
+}}, 0);
+// Must NOT concede on the stale key: sessionIdx recovers the row and the realign runs.
+assert.strictEqual(ok, true, 'stale content key must recover via sessionIdx, not return false');
+// Relative realign: scrollTop += (rect.top - containerRect.top) - targetTop
+//                              = 1200 + ((250 - 100) - 300) = 1200 + (150 - 300) = 1050.
+// This compensates the above-viewport height change; the buggy path would have left
+// an absolute jump (or no write at all), so this exact value pins the correct behavior.
+assert.strictEqual(container.scrollTop, 1050, 'realign must compensate the above-height change, not leave an absolute jump');
+"""
+    subprocess.run(["node", "-e", script], check=True, capture_output=True, text=True)
+
+
+def test_stale_key_does_not_dead_end_before_session_index_fallback():
+    """Guard the exact dead-code regression: a present content key must not short-circuit
+    the sessionIdx fallback. The old body had `if(!row&&anchorKey) return false;` BEFORE
+    sessionIdx was consulted, so a stale key could never reach the session-index recovery.
+    This asserts that early dead-end is gone and sessionIdx is resolved before conceding.
+    """
+
+    compact = _compact(_function_body(UI_JS, "_restoreMessageViewportAnchor"))
+
+    # The old dead-end (concede on stale key before trying sessionIdx) must be gone.
+    assert "if(!row&&anchorKey)returnfalse;" not in compact
+    # sessionIdx must be computed and the row recovered by it before any concede.
+    assert 'container.querySelector(`[data-session-msg-idx="${{sessionIdx}}"]`)'.replace("{{", "{").replace("}}", "}") in compact
+    # The concede now requires BOTH key AND sessionIdx to have failed.
+    assert "if(!row&&(anchorKey||hasSessionIdx))returnfalse;" in compact

tests/test_tars_scroll_reset_regressions.py

@@ -163,8 +163,13 @@ def test_preserve_scroll_restores_unpinned_viewport_after_dom_rebuild():
     assert "row.dataset.messageAnchorKey=_messageViewportAnchorKeyForMessage(m);" in UI_JS
     assert "seg.dataset.messageAnchorKey=_messageViewportAnchorKeyForMessage(m);" in UI_JS
     assert "container.querySelector(`[data-session-msg-idx=\"${sessionIdx}\"]`)" in UI_JS
-    assert "if(!row&&anchorKey) return false;" in UI_JS
+    # Desktop scroll jump-back root fix: a stale content-derived anchor key must NOT
-    assert "if(!row&&hasSessionIdx) return false;" in UI_JS
+    # dead-end before the sessionIdx fallback (the old `if(!row&&anchorKey) return false;`
+    # sat before sessionIdx was consulted, so a live-stream stale key could never reach
+    # the session-index recovery and the caller fell back to an absolute jump). The
+    # concede now requires BOTH key and sessionIdx lookups to have failed.
+    assert "if(!row&&anchorKey) return false;" not in UI_JS
+    assert "if(!row&&(anchorKey||hasSessionIdx)) return false;" in UI_JS
     assert "_restoreMessageViewportAnchor(snapshot.anchor,0)" in restore
     assert "if(!restoredViaAnchor){" in restore
     assert "el.scrollTop=Math.max(0,Math.min(Number(snapshot.top)||0,maxTop))" in restore