Merge pull request #3221 from nesquena/release/stage-batchA

Release stage-batchA → v0.51.181 (agent-cache eviction teardown #3218 + streaming finalize race #3201)

CHANGELOG.md

@@ -3,6 +3,12 @@
 
 ## [Unreleased]
 
+## [v0.51.181] — 2026-05-30 — Release FA (stage-batchA — agent-cache eviction teardown + streaming finalize race)
+
+### Fixed
+- Cached agents evicted after session-identity mismatches, unsafe runtime refresh, credential self-heal, or skipped compression migration now go through the normal session-boundary teardown path, committing pending memory and closing provider/session resources instead of silently dropping the cache entry (#3218, closes #3215).
+- Assistant streaming text is no longer lost when a stream completes while you have switched to a different session tab: the SSE `done` handler now sets the stream-finalized flag immediately (before the fade window), so a `stream_end` event arriving mid-fade can no longer trigger `_restoreSettledSession()` and overwrite the live messages with a stale server snapshot (#3201, closes #3195).
+
 ## [v0.51.180] — 2026-05-30 — Release EZ (stage-batch62 — session/agent cache ownership hardening)
 
 ### Fixed

api/streaming.py

@@ -3577,8 +3577,11 @@ def _attempt_credential_self_heal(
             return None
 
         # 2. Evict the cached agent for this session
+        _evicted_entry = None
         with SESSION_AGENT_CACHE_LOCK:
-            SESSION_AGENT_CACHE.pop(session_id, None)
+            _evicted_entry = SESSION_AGENT_CACHE.pop(session_id, None)
+        if _evicted_entry is not None:
+            _close_cached_agent_entry_at_session_boundary(session_id, _evicted_entry)
 
         # 3. Invalidate the credential pool for this provider
         invalidate_credential_pool_cache(provider_id)
@@ -3680,6 +3683,12 @@ def _close_evicted_agent_at_session_boundary(session_id: str, agent) -> bool:
     return True
 
 
+def _close_cached_agent_entry_at_session_boundary(session_id: str, cache_entry) -> bool:
+    """Commit and tear down a popped SESSION_AGENT_CACHE entry outside the cache lock."""
+    agent = cache_entry[0] if isinstance(cache_entry, tuple) else None
+    return _close_evicted_agent_at_session_boundary(session_id, agent)
+
+
 def _refresh_cached_agent_runtime(agent, agent_kwargs: dict) -> bool:
     """Refresh volatile runtime credentials on a reused cached AIAgent.
 
@@ -4887,6 +4896,7 @@ def _run_agent_streaming(
                 _agent_sig = _hashlib.sha256(_sig_blob.encode()).hexdigest()[:16]
 
                 agent = None
+                _identity_mismatch_entry = None
                 with SESSION_AGENT_CACHE_LOCK:
                     _cached = SESSION_AGENT_CACHE.get(session_id)
                     if _cached and _cached[1] == _agent_sig:
@@ -4896,7 +4906,7 @@ def _run_agent_streaming(
                             SESSION_AGENT_CACHE.move_to_end(session_id)  # LRU: mark as recently used
                             logger.debug('[webui] Reusing cached agent for session %s', session_id)
                         else:
-                            SESSION_AGENT_CACHE.pop(session_id, None)
+                            _identity_mismatch_entry = SESSION_AGENT_CACHE.pop(session_id, None)
                             logger.warning(
                                 '[webui] Evicted cached agent with mismatched session identity: cache_key=%s agent_session_id=%s',
                                 session_id,
@@ -4911,6 +4921,12 @@ def _run_agent_streaming(
                         except Exception:
                             logger.debug("Lifecycle register_agent failed for cached session %s", session_id, exc_info=True)
 
+                if _identity_mismatch_entry is not None:
+                    try:
+                        _close_cached_agent_entry_at_session_boundary(session_id, _identity_mismatch_entry)
+                    except Exception:
+                        logger.debug("Failed to close identity-mismatched cached agent for session %s", session_id, exc_info=True)
+
                 if agent is not None:
                     # Refresh volatile runtime credentials selected from provider
                     # pools without discarding cross-turn agent/provider state.
@@ -4919,13 +4935,14 @@ def _run_agent_streaming(
                             '[webui] Cached agent runtime could not be safely refreshed; rebuilding agent for session %s',
                             session_id,
                         )
-                        try:
-                            if getattr(agent, '_session_db', None) is not None:
-                                agent._session_db.close()
-                        except Exception:
-                            pass
+                        _stale_runtime_entry = None
                         with SESSION_AGENT_CACHE_LOCK:
-                            SESSION_AGENT_CACHE.pop(session_id, None)
+                            _stale_runtime_entry = SESSION_AGENT_CACHE.pop(session_id, None)
+                        if _stale_runtime_entry is not None:
+                            try:
+                                _close_cached_agent_entry_at_session_boundary(session_id, _stale_runtime_entry)
+                            except Exception:
+                                logger.debug("Failed to close stale-runtime cached agent for session %s", session_id, exc_info=True)
                         agent = None
 
                 if agent is not None:
@@ -5564,6 +5581,7 @@ def _run_agent_streaming(
                     # Migrate cached agent to the new session ID so the turn
                     # count survives context compression.
                     from api.config import SESSION_AGENT_CACHE, SESSION_AGENT_CACHE_LOCK
+                    _skipped_agent_migration_entry = None
                     with SESSION_AGENT_CACHE_LOCK:
                         _cached_entry = SESSION_AGENT_CACHE.pop(old_sid, None)
                         if _cached_entry:
@@ -5571,12 +5589,18 @@ def _run_agent_streaming(
                             if _cached_agent_matches_session(_cached_agent, new_sid):
                                 SESSION_AGENT_CACHE[new_sid] = _cached_entry
                             else:
+                                _skipped_agent_migration_entry = _cached_entry
                                 logger.warning(
                                     '[webui] Skipped cached agent migration with mismatched session identity: old_sid=%s new_sid=%s agent_session_id=%s',
                                     old_sid,
                                     new_sid,
                                     _cached_agent_session_identity(_cached_agent),
                                 )
+                    if _skipped_agent_migration_entry is not None:
+                        try:
+                            _close_cached_agent_entry_at_session_boundary(old_sid, _skipped_agent_migration_entry)
+                        except Exception:
+                            logger.debug("Failed to close skipped compression-migration cached agent for session %s", old_sid, exc_info=True)
                     _compressed = True
                 # Also detect compression via the result dict or compressor state
                 if not _compressed:
@@ -6450,18 +6474,24 @@ def _handle_chat_steer(handler, body: dict) -> bool:
     if not text:
         return bad(handler, "text required")
 
+    evicted_cached_entry = None
     with _cfg.SESSION_AGENT_CACHE_LOCK:
         cached = _cfg.SESSION_AGENT_CACHE.get(sid)
         if cached:
             agent = cached[0]
             if not _cached_agent_matches_session(agent, sid):
-                _cfg.SESSION_AGENT_CACHE.pop(sid, None)
+                evicted_cached_entry = _cfg.SESSION_AGENT_CACHE.pop(sid, None)
                 logger.warning(
                     '[webui] Evicted cached agent before steer due to mismatched session identity: cache_key=%s agent_session_id=%s',
                     sid,
                     _cached_agent_session_identity(agent),
                 )
                 cached = None
+    if evicted_cached_entry is not None:
+        try:
+            _close_cached_agent_entry_at_session_boundary(sid, evicted_cached_entry)
+        except Exception:
+            logger.debug("Failed to close steer identity-mismatched cached agent for session %s", sid, exc_info=True)
     if not cached:
         # No active agent for this session — caller falls back to interrupt
         return j(handler, {"accepted": False, "fallback": "no_cached_agent",

static/messages.js

@@ -1821,6 +1821,11 @@ function attachLiveStream(activeSid, streamId, uploaded=[], options={}){
 
     source.addEventListener('done',e=>{
       if(_streamFinalized) return;
+      // Set _streamFinalized IMMEDIATELY — before any fade delay. Without this,
+      // a stream_end event arriving during the fade window sees
+      // _streamFinalized=false, calls _restoreSettledSession(), and overwrites
+      // S.messages with stale server data (issue #3195).
+      _streamFinalized=true;
       _terminalStateReached=true;
       if(_persistTimer){clearTimeout(_persistTimer);_persistTimer=null;}
       const _doneData=JSON.parse(e.data);

tests/test_session_cache_ownership.py

@@ -87,7 +87,8 @@ def test_cached_agent_session_identity_matches_requested_sid():
     assert _cached_agent_matches_session(legacy, "requested") is True
 
 
-def test_handle_chat_steer_evicts_mismatched_cached_agent():
+def test_handle_chat_steer_evicts_mismatched_cached_agent(monkeypatch):
+    import api.streaming as streaming
     from api.streaming import _handle_chat_steer
 
     class Handler:
@@ -108,6 +109,12 @@ def test_handle_chat_steer_evicts_mismatched_cached_agent():
             pass
 
     wrong_agent = SimpleNamespace(session_id="other-session", steer=lambda _text: True)
+    closed_entries = []
+    monkeypatch.setattr(
+        streaming,
+        "_close_cached_agent_entry_at_session_boundary",
+        lambda session_id, entry: closed_entries.append((session_id, entry)),
+    )
     config.SESSION_AGENT_CACHE.clear()
     config.SESSION_AGENT_CACHE["requested"] = (wrong_agent, "sig")
     handler = Handler()
@@ -118,5 +125,6 @@ def test_handle_chat_steer_evicts_mismatched_cached_agent():
     assert handler.status == 200
     assert payload == {"accepted": False, "fallback": "no_cached_agent", "stream_id": None}
     assert "requested" not in config.SESSION_AGENT_CACHE
+    assert closed_entries == [("requested", (wrong_agent, "sig"))]
 
     config.SESSION_AGENT_CACHE.clear()

tests/test_streaming_agent_cache_lifecycle.py

@@ -51,6 +51,22 @@ def test_evicted_agent_lifecycle_shutdown_uses_empty_messages_when_missing(monke
     agent._session_db.close.assert_called_once()
 
 
+def test_cached_agent_entry_lifecycle_extracts_agent_from_cache_tuple(monkeypatch):
+    import api.streaming as streaming
+
+    closed = []
+    monkeypatch.setattr(
+        streaming,
+        "_close_evicted_agent_at_session_boundary",
+        lambda session_id, agent: closed.append((session_id, agent)) or True,
+    )
+
+    agent = MagicMock()
+
+    assert streaming._close_cached_agent_entry_at_session_boundary("old-session", (agent, "sig")) is True
+    assert closed == [("old-session", agent)]
+
+
 def test_evicted_agent_lifecycle_keeps_provider_alive_when_commit_still_dirty(monkeypatch):
     import api.streaming as streaming
 
@@ -71,3 +87,37 @@ def test_evicted_agent_lifecycle_keeps_provider_alive_when_commit_still_dirty(mo
 
     agent.shutdown_memory_provider.assert_not_called()
     agent._session_db.close.assert_not_called()
+
+
+def test_identity_mismatch_cache_evictions_close_entries_outside_cache_lock():
+    src = open("api/streaming.py", encoding="utf-8").read()
+
+    expected_markers = [
+        "_identity_mismatch_entry = SESSION_AGENT_CACHE.pop(session_id, None)",
+        "_stale_runtime_entry = SESSION_AGENT_CACHE.pop(session_id, None)",
+        "_skipped_agent_migration_entry = _cached_entry",
+        "evicted_cached_entry = _cfg.SESSION_AGENT_CACHE.pop(sid, None)",
+        "_evicted_entry = SESSION_AGENT_CACHE.pop(session_id, None)",
+    ]
+    for marker in expected_markers:
+        assert marker in src
+
+    close_markers = [
+        "_close_cached_agent_entry_at_session_boundary(session_id, _identity_mismatch_entry)",
+        "_close_cached_agent_entry_at_session_boundary(session_id, _stale_runtime_entry)",
+        "_close_cached_agent_entry_at_session_boundary(old_sid, _skipped_agent_migration_entry)",
+        "_close_cached_agent_entry_at_session_boundary(sid, evicted_cached_entry)",
+        "_close_cached_agent_entry_at_session_boundary(session_id, _evicted_entry)",
+    ]
+    lines = src.splitlines()
+    for marker in close_markers:
+        close_idx = next(i for i, line in enumerate(lines) if marker in line)
+        lock_idx = max(i for i, line in enumerate(lines[:close_idx]) if "with SESSION_AGENT_CACHE_LOCK:" in line)
+        lock_indent = len(lines[lock_idx]) - len(lines[lock_idx].lstrip())
+        between = lines[lock_idx + 1:close_idx]
+        assert any(
+            line.strip()
+            and not line.lstrip().startswith("#")
+            and len(line) - len(line.lstrip()) <= lock_indent
+            for line in between
+        ), f"{marker} still appears inside the SESSION_AGENT_CACHE_LOCK block"

tests/test_streaming_race_fix.py

@@ -71,6 +71,29 @@ class TestStreamFinalized:
             "'done' handler must call finalizeThinkingCard() to close thinking card"
         )
 
+    def test_done_sets_stream_finalized_before_fade_window(self):
+        """#3195 regression: the 'done' handler must set _streamFinalized=true
+        IMMEDIATELY after the early-return guard — before the fade machinery /
+        _finishDone() closure runs. Otherwise a stream_end event arriving during
+        the fade window sees _streamFinalized=false, calls _restoreSettledSession(),
+        and overwrites S.messages with stale server data (assistant text between
+        tool-call blocks vanishes on switching back to a settled session).
+        """
+        src = read('static/messages.js')
+        m = re.search(r"source\.addEventListener\('done'.*?\}\);", src, re.DOTALL)
+        assert m, "'done' handler not found"
+        fn = m.group(0)
+        guard_idx = fn.find('if(_streamFinalized) return;')
+        assert guard_idx != -1, "'done' handler must early-return on _streamFinalized"
+        finalize_idx = fn.find('_streamFinalized=true', guard_idx)
+        terminal_idx = fn.find('_terminalStateReached=true', guard_idx)
+        assert finalize_idx != -1, "'done' handler must set _streamFinalized=true"
+        assert terminal_idx != -1, "'done' handler must set _terminalStateReached"
+        assert finalize_idx < terminal_idx, (
+            "_streamFinalized=true must be set immediately after the guard "
+            "(before _terminalStateReached / fade machinery) — #3195"
+        )
+
     def test_apperror_sets_stream_finalized(self):
         src = read('static/messages.js')
         m = re.search(r"source\.addEventListener\('apperror'.*?\}\);", src, re.DOTALL)