From 9932899b884f60904993faf232feab3b43e029d9 Mon Sep 17 00:00:00 2001 From: ai-ag2026 <261867348+ai-ag2026@users.noreply.github.com> Date: Sat, 30 May 2026 18:13:46 +0200 Subject: [PATCH] fix: guard session cache ownership across compression --- CHANGELOG.md | 4 + api/models.py | 16 +++ api/streaming.py | 92 ++++++++++++-- tests/test_issue765_streaming_persistence.py | 11 +- tests/test_session_cache_ownership.py | 122 +++++++++++++++++++ 5 files changed, 232 insertions(+), 13 deletions(-) create mode 100644 tests/test_session_cache_ownership.py diff --git a/CHANGELOG.md b/CHANGELOG.md index c8ec9465..580675af 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,10 @@ ## [Unreleased] +### Fixed + +- Guard session and agent caches against compression/continuation id drift so a stale cached object cannot make `/api/session?session_id=` return an older transcript segment. + ## [v0.51.171] — 2026-05-30 — Release EQ (stage-batch53 — tool-output card badge + Neon opt-in skin) ### Added diff --git a/api/models.py b/api/models.py index aa80bfcf..0203c5ba 100644 --- a/api/models.py +++ b/api/models.py @@ -2199,6 +2199,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: diff --git a/api/streaming.py b/api/streaming.py index da77fef0..daf9a062 100644 --- a/api/streaming.py +++ b/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()) @@ -3685,6 +3689,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. @@ -4792,9 +4823,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: @@ -5448,8 +5489,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. @@ -5462,7 +5517,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: @@ -6338,6 +6402,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", diff --git a/tests/test_issue765_streaming_persistence.py b/tests/test_issue765_streaming_persistence.py index 53a9d34a..7a74bf78 100644 --- a/tests/test_issue765_streaming_persistence.py +++ b/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): diff --git a/tests/test_session_cache_ownership.py b/tests/test_session_cache_ownership.py new file mode 100644 index 00000000..3e3ee336 --- /dev/null +++ b/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()