Close evicted cached agents on identity mismatch
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user