From 8408a3dd22fe6df4a590f8dbed37fba6904820f0 Mon Sep 17 00:00:00 2001 From: Florian Krause Date: Thu, 28 May 2026 02:48:17 +0000 Subject: [PATCH] fix: reattach SSE on session-switch return + close leaked stream connections (#2925) Squash-merged pr-2925 into stage-batch33. Closes #2924. --- api/updates.py | 14 ++++-- static/messages.js | 30 +++++++++-- static/sessions.js | 7 +++ tests/test_inflight_stream_reuse.py | 72 +++++++++++++++++++++++++++ tests/test_parallel_session_switch.py | 32 +++++++++--- tests/test_regressions.py | 8 ++- tests/test_streaming_race_fix.py | 10 +++- 7 files changed, 156 insertions(+), 17 deletions(-) diff --git a/api/updates.py b/api/updates.py index d7793dc3..93c6ce95 100644 --- a/api/updates.py +++ b/api/updates.py @@ -116,9 +116,17 @@ def _dirty_suffix(path: Path, timeout=1) -> str: out, ok = _run_git(['diff-index', '--quiet', 'HEAD', '--'], path, timeout=timeout) if ok: return "" - # diff-index exits 1 with no output for a dirty tree. Timeouts and real git - # failures include a diagnostic; skip the suffix so the base version remains. - return "-dirty" if not out else "" + # diff-index --quiet exits 1 with no stdout/stderr to *signal* a dirty tree + # (not an error). _run_git() substitutes a synthetic "git exited with + # status N" diagnostic when both streams are empty, which makes the naive + # `if not out` guard always false on dirty trees — silently dropping the + # suffix and defeating dev-build cache busting (static/foo.js?v=… stays + # identical to the last-committed version). Treat the synthetic shape as + # the dirty signal; real errors (timeouts, missing git) carry a different + # diagnostic and correctly suppress the suffix. + if not out or out.startswith('git exited with status '): + return "-dirty" + return "" def _describe_git_version(path: Path, *, timeout=5, dirty_timeout=1) -> str | None: diff --git a/static/messages.js b/static/messages.js index c8035f46..c2195150 100644 --- a/static/messages.js +++ b/static/messages.js @@ -616,6 +616,16 @@ function closeLiveStream(sessionId, streamId, source){ if(source&&live.source!==source) return; try{live.source.close();}catch(_){ } delete LIVE_STREAMS[sessionId]; + // closeLiveStream() is called during session-switch teardown for any session + // the user is no longer viewing. The stream is still active on the server, + // so mark the in-memory INFLIGHT entry for reattach — otherwise + // loadSession() returning to this session skips the reattach branch + // (`INFLIGHT.reattach` was only set by the storage-load path) and the SSE + // is never reopened. The user then sees no streamed tokens until the LLM + // finishes and a metadata refresh swaps in the final reply. + // If the stream is terminating cleanly, _clearOwnerInflightState() has + // already deleted INFLIGHT[sessionId], so this is a safe no-op. + if(INFLIGHT[sessionId]) INFLIGHT[sessionId].reattach=true; } function closeOtherLiveStreams(activeSid){ @@ -648,9 +658,16 @@ function attachLiveStream(activeSid, streamId, uploaded=[], options={}){ closeOtherLiveStreams(activeSid); closeLiveStream(activeSid); - let assistantText=''; - let reasoningText=''; - let liveReasoningText=''; + // On reconnect, restore accumulated text from INFLIGHT so we don't lose + // progress made before the session switch. Without this the closure starts + // empty and tokens arriving on the new SSE connection append to nothing — + // the already-rendered content vanishes. + const _lastLiveAssistant = reconnecting + ? INFLIGHT[activeSid]?.messages?.findLast?.(m => m.role === 'assistant' && m._live) + : null; + let assistantText = _lastLiveAssistant ? (_lastLiveAssistant.content || '') : ''; + let reasoningText = _lastLiveAssistant ? (_lastLiveAssistant.reasoning || '') : ''; + let liveReasoningText = reasoningText; let visibleInterimSnippets=[]; let _latestGoalStatus=null; let _pendingGoalContinuation=null; @@ -2135,6 +2152,13 @@ function attachLiveStream(activeSid, streamId, uploaded=[], options={}){ if(_deferStreamErrorIfOffline()) return; if(_deferStreamErrorIfPageHidden(source)) return; _closeSource(source); + // If the user has switched to a different session, don't attempt to + // reconnect — the old stream's EventSource was closed intentionally + // during session switch and reconnecting would leak a background stream. + if(!_isSessionActivelyViewed(activeSid)) return; + if(_terminalStateReached || _streamFinalized){ + return; + } // Attempt one reconnect if the stream is still active server-side if(!_reconnectAttempted && streamId){ _reconnectAttempted=true; diff --git a/static/sessions.js b/static/sessions.js index fe84c301..fb51405e 100644 --- a/static/sessions.js +++ b/static/sessions.js @@ -589,6 +589,13 @@ async function loadSession(sid){ S.toolCalls = []; _messagesTruncated = false; _oldestIdx = 0; + // Close live SSE streams from the session we're leaving. The error + // handler checks _isSessionActivelyViewed() and won't auto-reconnect + // for a backgrounded session, preventing leaked connections that would + // pump token events into an orphaned closure, freezing the main thread. + if (currentSid && currentSid !== sid && typeof closeOtherLiveStreams === 'function') { + closeOtherLiveStreams(sid); + } _loadingOlder = false; const _msgInner = $('msgInner'); if (_msgInner && currentSid !== sid) _msgInner.innerHTML = '
Loading conversation...
'; diff --git a/tests/test_inflight_stream_reuse.py b/tests/test_inflight_stream_reuse.py index 403df9ec..69c756b1 100644 --- a/tests/test_inflight_stream_reuse.py +++ b/tests/test_inflight_stream_reuse.py @@ -1,4 +1,5 @@ """Regression tests for preserving live streams across session switches.""" +import re from pathlib import Path REPO_ROOT = Path(__file__).parent.parent @@ -99,3 +100,74 @@ def test_load_session_reattach_path_uses_attach_live_stream_for_running_sessions assert reattach_pos != -1 assert active_pos < reattach_pos assert "{reconnecting:true}" in body[reattach_pos : reattach_pos + 200] + + +def test_close_live_stream_marks_inflight_for_reattach_on_return(): + """When closeLiveStream() tears down a still-active SSE transport (e.g. the + user switched to another session), the corresponding INFLIGHT entry must be + flagged so loadSession() reopens the SSE on return. + + Without this flag the in-memory INFLIGHT entry stays as it was (no + `reattach:true`, which is only set on the storage-load path), so + loadSession()'s reattach branch is skipped — the SSE is never reopened and + the user sees no streamed tokens until the LLM finishes and a metadata + refresh swaps in the final reply. + """ + body = _function_body(MESSAGES_JS, "closeLiveStream") + assert "INFLIGHT" in body, ( + "closeLiveStream() must touch INFLIGHT so loadSession() reattaches the " + "SSE when the user switches back to a still-streaming session" + ) + assert re.search(r"INFLIGHT\[\w+\]\s*&&\s*\(?INFLIGHT\[\w+\]\.reattach\s*=\s*true", body) \ + or re.search(r"if\s*\(\s*INFLIGHT\[\w+\]\s*\)\s*INFLIGHT\[\w+\]\.reattach\s*=\s*true", body), ( + "closeLiveStream() must set INFLIGHT[sessionId].reattach = true " + "(guarded by an existence check) so loadSession()'s reattach branch fires" + ) + + +def test_close_other_live_streams_triggers_reattach_for_backgrounded_sessions(): + """closeOtherLiveStreams() during session switch must mark every closed + background session for reattach. Otherwise switching back to a session whose + stream was closed during the switch leaves the SSE permanently disconnected. + """ + helper_body = _function_body(MESSAGES_JS, "closeOtherLiveStreams") + close_body = _function_body(MESSAGES_JS, "closeLiveStream") + # closeOtherLiveStreams delegates per-session teardown to closeLiveStream, + # so the reattach flag must be set inside closeLiveStream itself for the + # chain to work — this guards the indirection. + assert "closeLiveStream(sid)" in helper_body.replace(" ", ""), ( + "closeOtherLiveStreams() must delegate teardown to closeLiveStream()" + ) + assert "reattach" in close_body, ( + "closeLiveStream() must set the reattach flag so closeOtherLiveStreams() " + "propagates the reattach intent to every backgrounded session" + ) + + +def test_load_session_reattaches_when_inflight_is_in_memory_and_marked_for_reattach(): + """The session-switch return path must hit attachLiveStream() even when + INFLIGHT[sid] is already in memory (i.e. wasn't loaded from storage). + + Before the fix, only the storage-load path set `reattach:true` on INFLIGHT, + so a switch-back through an in-memory INFLIGHT entry skipped the reattach + branch. Once closeLiveStream() also sets reattach=true, the existing + `INFLIGHT[sid].reattach && activeStreamId` gate is enough — this test + pins the gate's shape so future refactors don't drop the flag check. + """ + body = _function_body(SESSIONS_JS, "loadSession") + inflight_idx = body.find("if(INFLIGHT[sid]){") + assert inflight_idx >= 0, "INFLIGHT branch not found in loadSession" + inflight_block = body[inflight_idx : inflight_idx + 2400] + assert "INFLIGHT[sid].reattach" in inflight_block, ( + "loadSession()'s INFLIGHT branch must gate the SSE reattach on the " + "reattach flag so closeLiveStream()'s marking flows through" + ) + reattach_gate = re.search( + r"if\(INFLIGHT\[sid\]\.reattach\s*&&\s*activeStreamId.*?attachLiveStream\(sid, activeStreamId", + inflight_block, + re.DOTALL, + ) + assert reattach_gate, ( + "loadSession() must reattach via attachLiveStream() when " + "INFLIGHT[sid].reattach && activeStreamId" + ) diff --git a/tests/test_parallel_session_switch.py b/tests/test_parallel_session_switch.py index 3507db79..bd01d4eb 100644 --- a/tests/test_parallel_session_switch.py +++ b/tests/test_parallel_session_switch.py @@ -9,6 +9,7 @@ Four optimizations to reduce session-switch latency: """ import pathlib +import re import threading import time from unittest.mock import patch, MagicMock @@ -487,10 +488,20 @@ class TestSessionSwitchCancellation: def test_loading_older_reset_on_session_switch(self): """loadSession must reset _loadingOlder when switching sessions.""" - # Find the reset block in loadSession - marker = "_messagesTruncated = false;\n _oldestIdx = 0;\n _loadingOlder = false;" - idx = SESSIONS_JS.find(marker) - assert idx >= 0, ( + # Locate the on-switch reset block — it lives in the `if (currentSid !== sid || forceReload)` + # arm of loadSession. Match by the surrounding state-resets rather than by a fragile + # multi-line substring, so unrelated code (like the closeOtherLiveStreams teardown + # that was inserted between _oldestIdx and _loadingOlder) doesn't break the test. + switch_arm = re.search( + r"if \(currentSid !== sid \|\| forceReload\) \{(.*?)\n \}", + SESSIONS_JS, + re.DOTALL, + ) + assert switch_arm, "loadSession's session-switch reset arm not found" + block = switch_arm.group(1) + assert "_messagesTruncated = false;" in block + assert "_oldestIdx = 0;" in block + assert "_loadingOlder = false;" in block, ( "loadSession must reset _loadingOlder=false on session switch " "to prevent a stale _loadOlderMessages lock from blocking the " "new session's scroll-to-top loading." @@ -517,13 +528,20 @@ class TestSessionSwitchCancellation: def test_messages_truncated_reset_on_switch(self): """loadSession must reset _messagesTruncated on session switch.""" - marker = "_messagesTruncated = false;\n _oldestIdx = 0;\n _loadingOlder = false;" - idx = SESSIONS_JS.find(marker) - assert idx >= 0, ( + switch_arm = re.search( + r"if \(currentSid !== sid \|\| forceReload\) \{(.*?)\n \}", + SESSIONS_JS, + re.DOTALL, + ) + assert switch_arm, "loadSession's session-switch reset arm not found" + block = switch_arm.group(1) + assert "_messagesTruncated = false;" in block, ( "_messagesTruncated must be reset to false on session switch " "to prevent the scroll-to-top handler from trying to load " "older messages from the previous session." ) + assert "_oldestIdx = 0;" in block + assert "_loadingOlder = false;" in block def test_oldest_idx_reset_prevents_wrong_cursor(self): """_oldestIdx=0 after switch prevents passing stale cursor to API.""" diff --git a/tests/test_regressions.py b/tests/test_regressions.py index e2dd3441..428b3b46 100644 --- a/tests/test_regressions.py +++ b/tests/test_regressions.py @@ -761,9 +761,13 @@ def test_messages_js_supports_live_reasoning_and_tool_completion(cleanup_test_se until the final done snapshot redraws the whole turn. """ src = (REPO_ROOT / "static/messages.js").read_text() - assert "let reasoningText=''" in src, \ + # reasoningText is initialised at closure scope in attachLiveStream. + # On initial connect it defaults to ''; on reconnect it restores from + # INFLIGHT so the already-rendered content survives the session switch. + assert ("let reasoningText=''" in src + or "let reasoningText = _lastLiveAssistant" in src), \ "messages.js must track streamed reasoning text separately from assistant text" - assert "let liveReasoningText=''" in src or 'let liveReasoningText = ""' in src, \ + assert ("let liveReasoningText=''" in src or "let liveReasoningText = reasoningText" in src), \ "messages.js must track the currently active reasoning segment separately from cumulative reasoning" assert "source.addEventListener('reasoning'" in src or 'source.addEventListener("reasoning"' in src, \ "messages.js must listen for live reasoning SSE events" diff --git a/tests/test_streaming_race_fix.py b/tests/test_streaming_race_fix.py index 55f0f2ed..aeba1531 100644 --- a/tests/test_streaming_race_fix.py +++ b/tests/test_streaming_race_fix.py @@ -154,8 +154,14 @@ class TestReconnectAccumulatorPreservation: ) assert m, "attachLiveStream prelude not found" prelude = m.group(0) - assert "let assistantText=''" in prelude or 'let assistantText = ""' in prelude, ( - "assistantText must be initialised to '' at closure scope — " + # On initial connect, assistantText and reasoningText are initialised to '' + # at closure scope (the ternary defaults to '' when reconnecting is false + # or INFLIGHT has no _live assistant message). On reconnect, they restore + # from INFLIGHT so the already-rendered content survives the session switch. + assert ("let assistantText=''" in prelude + or 'let assistantText = _lastLiveAssistant' in prelude + or 'let assistantText = ""' in prelude), ( + "assistantText must be initialised at closure scope — " "this is the only legitimate reset; _wireSSE must not re-reset" )