Close evicted cached agents on identity mismatch

This commit is contained in:
AJV20
2026-05-30 19:22:56 -04:00
parent 82f92eb78d
commit 6ecc32be00
4 changed files with 99 additions and 10 deletions

View File

@@ -3,6 +3,9 @@
## [Unreleased]
### 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 (#3215).
## [v0.51.180] — 2026-05-30 — Release EZ (stage-batch62 — session/agent cache ownership hardening)
### Fixed

View File

@@ -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",

View File

@@ -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()

View File

@@ -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,35 @@ 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)",
]
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)",
]
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"