diff --git a/api/config.py b/api/config.py index 8230f2f79..8efafbb4d 100644 --- a/api/config.py +++ b/api/config.py @@ -7603,6 +7603,25 @@ def _evict_session_agent(session_id: str) -> None: agent = entry[0] if isinstance(entry, tuple) else None if agent is None: return + # A live run for this session may still hold this agent's _session_db (the + # worker assigns agent._session_db at run start). Never close it out from + # under an in-flight turn — ACTIVE_RUNS is the authoritative liveness signal + # (mirrors the worker's own LRU-eviction guard in streaming.py). When a run + # is live we still drop the cache handle above (harmless — the worker holds + # a local ref), but skip the lifecycle commit + _session_db.close() so the + # running turn can finish persisting. Hardens /clear + model-switch eviction + # too, not just truncate (#5096 Bug D). + _run_active = False + try: + with ACTIVE_RUNS_LOCK: + for _entry in (ACTIVE_RUNS or {}).values(): + if (_entry or {}).get("session_id") == session_id: + _run_active = True + break + except Exception: + _run_active = False + if _run_active: + return should_close = True try: from api.session_lifecycle import commit_session_memory, discard_session, has_uncommitted_work, unregister_agent diff --git a/tests/test_evict_session_agent_active_run_guard.py b/tests/test_evict_session_agent_active_run_guard.py new file mode 100644 index 000000000..c1033fb0b --- /dev/null +++ b/tests/test_evict_session_agent_active_run_guard.py @@ -0,0 +1,80 @@ +"""Regression: _evict_session_agent must not close a live session's _session_db. + +Bug-D follow-up (#5096): truncate/clear/model-switch all call +api.config._evict_session_agent(). The worker assigns agent._session_db at run +start, so eviction must consult ACTIVE_RUNS (the authoritative liveness signal, +same as the worker's own LRU-eviction guard) and skip the lifecycle commit + +_session_db.close() while a run is in flight on that session. Otherwise a +truncate racing an in-flight turn on the same session (reachable via a second +client / direct API; the UI gates it behind S.busy) closes the SessionDB the +running worker is still persisting through. +""" + +import api.config as config + + +class _FakeSessionDB: + def __init__(self): + self.closed = False + + def close(self): + self.closed = True + + +class _FakeAgent: + def __init__(self): + self._session_db = _FakeSessionDB() + + +def _seed_cache(session_id, agent): + with config.SESSION_AGENT_CACHE_LOCK: + config.SESSION_AGENT_CACHE[session_id] = (agent, "sig") + + +def _clear_active_runs(): + with config.ACTIVE_RUNS_LOCK: + config.ACTIVE_RUNS.clear() + + +def test_evict_skips_session_db_close_when_run_active(monkeypatch): + """A live run on the session => the cache handle drops but the DB stays open.""" + sid = "live-session-evict-guard" + agent = _FakeAgent() + _seed_cache(sid, agent) + _clear_active_runs() + config.register_active_run("stream-xyz", session_id=sid) + try: + config._evict_session_agent(sid) + # The agent handle is removed from the cache (harmless — the worker + # holds its own local ref) ... + with config.SESSION_AGENT_CACHE_LOCK: + assert sid not in config.SESSION_AGENT_CACHE + # ... but the live worker's SessionDB must NOT be closed. + assert agent._session_db.closed is False + finally: + _clear_active_runs() + with config.SESSION_AGENT_CACHE_LOCK: + config.SESSION_AGENT_CACHE.pop(sid, None) + + +def test_evict_closes_session_db_when_no_run_active(monkeypatch): + """No live run => normal eviction closes the SessionDB (idle path unchanged).""" + sid = "idle-session-evict-guard" + agent = _FakeAgent() + _seed_cache(sid, agent) + _clear_active_runs() + + # Neutralize the lifecycle commit machinery so the test isolates the + # ACTIVE_RUNS guard + close decision (no uncommitted work => should_close). + monkeypatch.setattr("api.session_lifecycle.has_uncommitted_work", lambda *_a, **_k: False) + monkeypatch.setattr("api.session_lifecycle.unregister_agent", lambda *_a, **_k: None) + monkeypatch.setattr("api.session_lifecycle.discard_session", lambda *_a, **_k: None) + try: + config._evict_session_agent(sid) + with config.SESSION_AGENT_CACHE_LOCK: + assert sid not in config.SESSION_AGENT_CACHE + assert agent._session_db.closed is True + finally: + _clear_active_runs() + with config.SESSION_AGENT_CACHE_LOCK: + config.SESSION_AGENT_CACHE.pop(sid, None)