Merge #3191: guard session cache ownership across compression

# Conflicts:
#	CHANGELOG.md

CHANGELOG.md

@@ -3,6 +3,13 @@
 
 ## [Unreleased]
 
+## [v0.51.180] — 2026-05-30 — Release EZ (stage-batch62 — session/agent cache ownership hardening)
+
+### Fixed
+
+- Guard session and agent caches against compression/continuation id drift: `GET /api/session` now evicts a cached `Session` whose own `session_id` no longer matches the requested key (instead of trusting the LRU), the background title-update/refresh paths only adopt a cached session when its identity matches, and the compression checkpoint migration no longer re-stores a stale object under the old lineage id. Prevents a stale cached object from making `/api/session?session_id=<tip>` return an older transcript segment, which looked like a disappeared session (#3191).
+- Evicted cached agents are now torn down cleanly at the WebUI session boundary: pending session memory is committed first, and only if the lifecycle entry is clean afterward does the agent get unregistered and its memory provider shut down via `shutdown_memory_provider(messages)` (closing provider-owned clients such as Hindsight's aiohttp session) before the session DB is closed — instead of leaking those resources until garbage collection (#3166).
+
 ## [v0.51.179] — 2026-05-30 — Release EY (stage-batch61 — custom-provider reasoning efforts + clearer sidebar tooltips)
 
 ### Fixed

api/models.py

@@ -2204,6 +2204,22 @@ def get_session(sid, metadata_only=False):
         cached = SESSIONS.get(sid)
         if cached is not None:
             SESSIONS.move_to_end(sid)  # LRU: mark as recently used
+    if cached is not None:
+        # Defensive cache ownership check: compression/continuation and recovery
+        # paths can temporarily juggle Session objects across lineage ids.  A
+        # stale object stored under the wrong key makes GET /api/session return
+        # a different transcript than the requested sid, which looks exactly
+        # like a disappeared session.  Evict instead of trusting the LRU.
+        if str(getattr(cached, 'session_id', '') or '') != str(sid):
+            logger.warning(
+                "evicting mismatched cached session: requested %s but cached object is %s",
+                sid,
+                getattr(cached, 'session_id', None),
+            )
+            with LOCK:
+                if SESSIONS.get(sid) is cached:
+                    SESSIONS.pop(sid, None)
+            cached = None
     if cached is not None:
         if not metadata_only and _inactive_cache_tail_needs_disk_check(cached):
             try:

api/streaming.py

@@ -28,7 +28,7 @@ from api.config import (
     STREAM_REASONING_TEXT, STREAM_LIVE_TOOL_CALLS,
     STREAM_GOAL_RELATED, PENDING_GOAL_CONTINUATION,
     STREAM_LAST_EVENT_ID,
-    LOCK, SESSIONS, SESSION_DIR,
+    LOCK, SESSIONS, SESSIONS_MAX, SESSION_DIR,
     _get_session_agent_lock, _set_thread_env, _clear_thread_env,
     register_active_run, update_active_run, unregister_active_run,
     SESSION_AGENT_LOCKS, SESSION_AGENT_LOCKS_LOCK,
@@ -2163,7 +2163,9 @@ def _run_background_title_update(session_id: str, user_text: str, assistant_text
         if next_title:
             with _get_session_agent_lock(session_id):
                 with LOCK:
-                    s = SESSIONS.get(session_id, s)
+                    cached_session = SESSIONS.get(session_id)
+                    if cached_session is not None and getattr(cached_session, 'session_id', None) == session_id:
+                        s = cached_session
                     effective_title = str(s.title or '').strip()
                     invalid_existing_now = _looks_invalid_generated_title(s.title)
                     still_auto = (
@@ -2238,7 +2240,9 @@ def _run_background_title_refresh(session_id: str, user_text: str, assistant_tex
             return
         with _get_session_agent_lock(session_id):
             with LOCK:
-                s = SESSIONS.get(session_id, s)
+                cached_session = SESSIONS.get(session_id)
+                if cached_session is not None and getattr(cached_session, 'session_id', None) == session_id:
+                    s = cached_session
                 # Re-check: user may have renamed while we were generating
                 if str(s.title or '').strip() != current_title:
                     _put_title_status(put_event, session_id, 'skipped', 'manual_title', str(s.title or '').strip())
@@ -3752,6 +3756,33 @@ def _refresh_cached_agent_runtime(agent, agent_kwargs: dict) -> bool:
         return False
 
 
+def _cached_agent_session_identity(agent) -> str | None:
+    """Best-effort session id carried by a cached AIAgent.
+
+    The cache key is only safe when it agrees with the object's own session
+    identity. Some old/fake agents may not expose an identity; keep those
+    backwards-compatible and treat them as unverifiable rather than mismatched.
+    """
+    if agent is None:
+        return None
+    for attr in ('session_id', '_session_id'):
+        value = getattr(agent, attr, None)
+        if isinstance(value, str) and value:
+            return value
+    session_db = getattr(agent, '_session_db', None)
+    if session_db is not None:
+        for attr in ('session_id', '_session_id'):
+            value = getattr(session_db, attr, None)
+            if isinstance(value, str) and value:
+                return value
+    return None
+
+
+def _cached_agent_matches_session(agent, session_id: str) -> bool:
+    identity = _cached_agent_session_identity(agent)
+    return identity is None or identity == str(session_id)
+
+
 def _refresh_cached_agent_primary_runtime_snapshot(agent) -> None:
     """Keep AIAgent's primary-runtime snapshot aligned with refreshed creds.
 
@@ -4859,9 +4890,19 @@ def _run_agent_streaming(
                 with SESSION_AGENT_CACHE_LOCK:
                     _cached = SESSION_AGENT_CACHE.get(session_id)
                     if _cached and _cached[1] == _agent_sig:
-                        agent = _cached[0]
-                        SESSION_AGENT_CACHE.move_to_end(session_id)  # LRU: mark as recently used
-                        logger.debug('[webui] Reusing cached agent for session %s', session_id)
+                        _cached_agent = _cached[0]
+                        if _cached_agent_matches_session(_cached_agent, session_id):
+                            agent = _cached_agent
+                            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)
+                            logger.warning(
+                                '[webui] Evicted cached agent with mismatched session identity: cache_key=%s agent_session_id=%s',
+                                session_id,
+                                _cached_agent_session_identity(_cached_agent),
+                            )
+                    if agent is not None:
                         # Reopened/cache-hit sessions must register the agent
                         # so later lifecycle commits can find it.
                         try:
@@ -5498,8 +5539,22 @@ def _run_agent_streaming(
                     # parent, losing access to the recoverable history in old_sid.json.
                     s.parent_session_id = old_sid
                     with LOCK:
-                        if old_sid in SESSIONS:
-                            SESSIONS[new_sid] = SESSIONS.pop(old_sid)
+                        cached_old_session = SESSIONS.pop(old_sid, None)
+                        if cached_old_session is not None and cached_old_session is not s:
+                            cached_old_sid = str(getattr(cached_old_session, 'session_id', '') or '')
+                            if cached_old_sid == str(old_sid):
+                                SESSIONS[old_sid] = cached_old_session
+                            else:
+                                logger.warning(
+                                    "compression cache migration skipped stale object: old_sid=%s new_sid=%s cached_session_id=%s",
+                                    old_sid,
+                                    new_sid,
+                                    cached_old_sid or None,
+                                )
+                        SESSIONS[new_sid] = s
+                        SESSIONS.move_to_end(new_sid)
+                        while len(SESSIONS) > SESSIONS_MAX:
+                            SESSIONS.popitem(last=False)
                     # Migrate the per-session lock: alias new_sid to the held
                     # _agent_lock reference directly (not via old_sid lookup),
                     # then remove the old_sid entry to prevent a leak.
@@ -5512,7 +5567,16 @@ def _run_agent_streaming(
                     with SESSION_AGENT_CACHE_LOCK:
                         _cached_entry = SESSION_AGENT_CACHE.pop(old_sid, None)
                         if _cached_entry:
-                            SESSION_AGENT_CACHE[new_sid] = _cached_entry
+                            _cached_agent = _cached_entry[0]
+                            if _cached_agent_matches_session(_cached_agent, new_sid):
+                                SESSION_AGENT_CACHE[new_sid] = _cached_entry
+                            else:
+                                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),
+                                )
                     _compressed = True
                 # Also detect compression via the result dict or compressor state
                 if not _compressed:
@@ -6388,6 +6452,16 @@ def _handle_chat_steer(handler, body: dict) -> bool:
 
     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)
+                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 not cached:
         # No active agent for this session — caller falls back to interrupt
         return j(handler, {"accepted": False, "fallback": "no_cached_agent",

tests/test_issue765_streaming_persistence.py

@@ -465,14 +465,17 @@ class TestIssue765FollowupHardening:
         )
         fn_idx = src.find("def _run_background_title_update(")
         assert fn_idx != -1, "_run_background_title_update not found"
-        fn_block = src[fn_idx:fn_idx + 3200]
+        fn_block = src[fn_idx:fn_idx + 4200]
         assert "with LOCK:" in fn_block, (
             "_run_background_title_update must acquire LOCK before rebinding "
             "to canonical cached session instance"
         )
-        assert "s = SESSIONS.get(session_id, s)" in fn_block, (
-            "_run_background_title_update must rebind to canonical cached "
-            "session instance under LOCK"
+        assert "cached_session = SESSIONS.get(session_id)" in fn_block, (
+            "_run_background_title_update must read the cached session under LOCK"
+        )
+        assert "getattr(cached_session, 'session_id', None) == session_id" in fn_block, (
+            "_run_background_title_update must only rebind to a canonical cached "
+            "session instance whose id matches the requested session"
         )
 
     def test_cancel_stream_uses_agent_lock(self):

tests/test_session_cache_ownership.py

@@ -0,0 +1,122 @@
+import json
+from io import BytesIO
+from types import SimpleNamespace
+
+import api.config as config
+import api.models as models
+from api.models import Session, get_session
+
+
+def test_get_session_evicts_cached_object_with_wrong_session_id(tmp_path, monkeypatch):
+    session_dir = tmp_path / "sessions"
+    session_dir.mkdir()
+    monkeypatch.setattr(models, "SESSION_DIR", session_dir)
+    monkeypatch.setattr(models, "SESSION_INDEX_FILE", session_dir / "_index.json")
+    monkeypatch.setattr(config, "SESSION_INDEX_FILE", session_dir / "_index.json", raising=False)
+    models.SESSIONS.clear()
+
+    requested = Session(
+        session_id="requested-tip",
+        title="Requested Tip",
+        messages=[{"role": "user", "content": "correct transcript"}],
+    )
+    requested.save()
+    wrong = Session(
+        session_id="older-segment",
+        title="Older Segment",
+        messages=[{"role": "user", "content": "wrong transcript"}],
+    )
+
+    models.SESSIONS["requested-tip"] = wrong
+
+    loaded = get_session("requested-tip")
+
+    assert loaded.session_id == "requested-tip"
+    assert loaded.title == "Requested Tip"
+    assert loaded.messages == [{"role": "user", "content": "correct transcript"}]
+    assert models.SESSIONS["requested-tip"] is loaded
+
+    models.SESSIONS.clear()
+
+
+def test_get_session_metadata_only_evicts_cached_object_with_wrong_session_id(tmp_path, monkeypatch):
+    session_dir = tmp_path / "sessions"
+    session_dir.mkdir()
+    monkeypatch.setattr(models, "SESSION_DIR", session_dir)
+    monkeypatch.setattr(models, "SESSION_INDEX_FILE", session_dir / "_index.json")
+    monkeypatch.setattr(config, "SESSION_INDEX_FILE", session_dir / "_index.json", raising=False)
+    models.SESSIONS.clear()
+
+    requested = Session(
+        session_id="requested-tip",
+        title="Requested Tip",
+        messages=[{"role": "user", "content": "correct transcript"}],
+    )
+    requested.save()
+    models.SESSIONS["requested-tip"] = Session(session_id="older-segment", title="Older Segment")
+
+    loaded = get_session("requested-tip", metadata_only=True)
+
+    assert loaded.session_id == "requested-tip"
+    assert loaded.title == "Requested Tip"
+    assert models.SESSIONS.get("requested-tip") is None
+
+    models.SESSIONS.clear()
+
+
+def test_compression_cache_migration_never_moves_unverified_cached_object_to_new_sid():
+    streaming_src = open("api/streaming.py", encoding="utf-8").read()
+
+    assert "SESSIONS[new_sid] = SESSIONS.pop(old_sid)" not in streaming_src
+    assert "cached_old_session is not s" in streaming_src
+    assert "SESSIONS[new_sid] = s" in streaming_src
+
+
+def test_cached_agent_session_identity_matches_requested_sid():
+    from api.streaming import _cached_agent_matches_session, _cached_agent_session_identity
+
+    matching = SimpleNamespace(session_id="requested")
+    mismatched = SimpleNamespace(session_id="other")
+    legacy = SimpleNamespace()
+    db_owned = SimpleNamespace(_session_db=SimpleNamespace(session_id="requested"))
+
+    assert _cached_agent_session_identity(matching) == "requested"
+    assert _cached_agent_matches_session(matching, "requested") is True
+    assert _cached_agent_matches_session(mismatched, "requested") is False
+    assert _cached_agent_matches_session(db_owned, "requested") is True
+    assert _cached_agent_matches_session(legacy, "requested") is True
+
+
+def test_handle_chat_steer_evicts_mismatched_cached_agent():
+    from api.streaming import _handle_chat_steer
+
+    class Handler:
+        headers = {}
+
+        def __init__(self):
+            self.status = None
+            self.response_headers = []
+            self.wfile = BytesIO()
+
+        def send_response(self, status):
+            self.status = status
+
+        def send_header(self, key, value):
+            self.response_headers.append((key, value))
+
+        def end_headers(self):
+            pass
+
+    wrong_agent = SimpleNamespace(session_id="other-session", steer=lambda _text: True)
+    config.SESSION_AGENT_CACHE.clear()
+    config.SESSION_AGENT_CACHE["requested"] = (wrong_agent, "sig")
+    handler = Handler()
+
+    _handle_chat_steer(handler, {"session_id": "requested", "text": "please steer"})
+
+    payload = json.loads(handler.wfile.getvalue().decode("utf-8"))
+    assert handler.status == 200
+    assert payload == {"accepted": False, "fallback": "no_cached_agent", "stream_id": None}
+    assert "requested" not in config.SESSION_AGENT_CACHE
+
+    config.SESSION_AGENT_CACHE.clear()