From b3483d943c75d538aa2208df387dc651628059dd Mon Sep 17 00:00:00 2001 From: ai-ag2026 <261867348+ai-ag2026@users.noreply.github.com> Date: Sat, 30 May 2026 06:07:06 +0200 Subject: [PATCH] fix: close evicted agents from WebUI cache --- api/streaming.py | 82 +++++++++++++++---- ...est_memory_session_lifecycle_generation.py | 8 +- tests/test_streaming_agent_cache_lifecycle.py | 73 +++++++++++++++++ tests/test_v050259_sessiondb_fd_leak.py | 14 +++- 4 files changed, 154 insertions(+), 23 deletions(-) create mode 100644 tests/test_streaming_agent_cache_lifecycle.py diff --git a/api/streaming.py b/api/streaming.py index da77fef02..60c851948 100644 --- a/api/streaming.py +++ b/api/streaming.py @@ -3609,6 +3609,69 @@ def _agent_cache_api_key_sig(resolved_api_key, credential_pool) -> str: return _hashlib.sha256((resolved_api_key or '').encode()).hexdigest()[:16] +def _lifecycle_commit_session_memory(session_id: str, *, agent=None, wait: bool = False) -> bool: + from api.session_lifecycle import commit_session_memory + + return commit_session_memory(session_id, agent=agent, wait=wait) + + +def _lifecycle_has_uncommitted_work(session_id: str) -> bool: + from api.session_lifecycle import has_uncommitted_work + + return has_uncommitted_work(session_id) + + +def _lifecycle_unregister_agent(session_id: str) -> None: + from api.session_lifecycle import unregister_agent + + unregister_agent(session_id) + + +def _close_evicted_agent_at_session_boundary(session_id: str, agent) -> bool: + """Commit and tear down an evicted cached agent at a WebUI session boundary. + + WebUI keeps AIAgent instances in an LRU cache so memory providers can carry + state across turns. When an agent is evicted, commit pending memory first; + if the lifecycle entry is clean afterwards, unregister and call + shutdown_memory_provider(messages) so provider-owned clients such as + Hindsight's aiohttp session are closed instead of being garbage-collected + later. Passing the cached transcript mirrors gateway cleanup semantics for + providers that use on_session_end(messages) during shutdown. + """ + if agent is None: + return True + + should_close_evicted_agent = True + try: + _lifecycle_commit_session_memory(session_id, agent=agent, wait=True) + if not _lifecycle_has_uncommitted_work(session_id): + _lifecycle_unregister_agent(session_id) + else: + should_close_evicted_agent = False + except Exception: + should_close_evicted_agent = False + logger.debug("Lifecycle commit on eviction failed for %s", session_id, exc_info=True) + + if not should_close_evicted_agent: + return False + + try: + shutdown_memory_provider = getattr(agent, 'shutdown_memory_provider', None) + if callable(shutdown_memory_provider): + session_messages = vars(agent).get('_session_messages', []) + shutdown_memory_provider(session_messages) + except Exception: + logger.debug("Failed to shut down evicted agent memory provider for session %s", session_id, exc_info=True) + + try: + session_db = getattr(agent, '_session_db', None) + if session_db is not None: + session_db.close() + except Exception: + logger.debug("Failed to close evicted agent session DB for session %s", session_id, exc_info=True) + return True + + def _refresh_cached_agent_runtime(agent, agent_kwargs: dict) -> bool: """Refresh volatile runtime credentials on a reused cached AIAgent. @@ -4880,24 +4943,7 @@ def _run_agent_streaming( for _evicted_sid, _evicted_entry in _evicted_items: try: _evicted_agent = _evicted_entry[0] if isinstance(_evicted_entry, tuple) else None - _should_close_evicted_agent = True - if _evicted_agent is not None: - try: - from api.session_lifecycle import ( - commit_session_memory as _lifecycle_commit, - has_uncommitted_work as _lifecycle_has_uncommitted_work, - unregister_agent as _lifecycle_unregister_agent, - ) - _lifecycle_commit(_evicted_sid, agent=_evicted_agent, wait=True) - if not _lifecycle_has_uncommitted_work(_evicted_sid): - _lifecycle_unregister_agent(_evicted_sid) - else: - _should_close_evicted_agent = False - except Exception: - _should_close_evicted_agent = False - logger.debug("Lifecycle commit on eviction failed for %s", _evicted_sid, exc_info=True) - if _should_close_evicted_agent and _evicted_agent is not None and getattr(_evicted_agent, '_session_db', None) is not None: - _evicted_agent._session_db.close() + _close_evicted_agent_at_session_boundary(_evicted_sid, _evicted_agent) except Exception: logger.debug("Failed to close evicted agent for session %s", _evicted_sid, exc_info=True) logger.debug('[webui] Evicted LRU agent from cache: %s', _evicted_sid) diff --git a/tests/test_memory_session_lifecycle_generation.py b/tests/test_memory_session_lifecycle_generation.py index 8b5c40227..2132c1056 100644 --- a/tests/test_memory_session_lifecycle_generation.py +++ b/tests/test_memory_session_lifecycle_generation.py @@ -329,8 +329,12 @@ def test_lru_eviction_commits_outside_cache_lock(): assert "commit_session_memory" not in locked_section assert "_lifecycle_commit" not in locked_section assert "SESSION_AGENT_CACHE.popitem" in locked_section - assert "_lifecycle_commit" in outside_section - assert "wait=True" in outside_section + assert "_close_evicted_agent_at_session_boundary" in outside_section + helper_start = src.index("def _close_evicted_agent_at_session_boundary") + helper_end = src.index("\ndef _refresh_cached_agent_runtime", helper_start) + helper_section = src[helper_start:helper_end] + assert "_lifecycle_commit_session_memory" in helper_section + assert "wait=True" in helper_section assert "outside the cache lock" in outside_section diff --git a/tests/test_streaming_agent_cache_lifecycle.py b/tests/test_streaming_agent_cache_lifecycle.py new file mode 100644 index 000000000..d57aebd92 --- /dev/null +++ b/tests/test_streaming_agent_cache_lifecycle.py @@ -0,0 +1,73 @@ +from unittest.mock import MagicMock + + +def test_evicted_agent_lifecycle_commits_unregisters_and_shutdowns(monkeypatch): + import api.streaming as streaming + + events = [] + + def fake_commit(session_id, *, agent=None, wait=False): + events.append(("commit", session_id, agent, wait)) + return True + + def fake_has_uncommitted_work(session_id): + events.append(("has_uncommitted", session_id)) + return False + + def fake_unregister(session_id): + events.append(("unregister", session_id)) + + monkeypatch.setattr(streaming, "_lifecycle_commit_session_memory", fake_commit) + monkeypatch.setattr(streaming, "_lifecycle_has_uncommitted_work", fake_has_uncommitted_work) + monkeypatch.setattr(streaming, "_lifecycle_unregister_agent", fake_unregister) + + session_db = MagicMock() + agent = MagicMock() + agent._session_db = session_db + agent._session_messages = [{"role": "user", "content": "hello"}] + + streaming._close_evicted_agent_at_session_boundary("old-session", agent) + + assert ("commit", "old-session", agent, True) in events + assert ("has_uncommitted", "old-session") in events + assert ("unregister", "old-session") in events + agent.shutdown_memory_provider.assert_called_once_with(agent._session_messages) + session_db.close.assert_called_once() + + +def test_evicted_agent_lifecycle_shutdown_uses_empty_messages_when_missing(monkeypatch): + import api.streaming as streaming + + monkeypatch.setattr(streaming, "_lifecycle_commit_session_memory", lambda *a, **kw: True) + monkeypatch.setattr(streaming, "_lifecycle_has_uncommitted_work", lambda session_id: False) + monkeypatch.setattr(streaming, "_lifecycle_unregister_agent", MagicMock()) + + agent = MagicMock() + agent._session_db = MagicMock() + + streaming._close_evicted_agent_at_session_boundary("old-session", agent) + + agent.shutdown_memory_provider.assert_called_once_with([]) + agent._session_db.close.assert_called_once() + + +def test_evicted_agent_lifecycle_keeps_provider_alive_when_commit_still_dirty(monkeypatch): + import api.streaming as streaming + + def fake_commit(session_id, *, agent=None, wait=False): + return True + + def fake_has_uncommitted_work(session_id): + return True + + monkeypatch.setattr(streaming, "_lifecycle_commit_session_memory", fake_commit) + monkeypatch.setattr(streaming, "_lifecycle_has_uncommitted_work", fake_has_uncommitted_work) + monkeypatch.setattr(streaming, "_lifecycle_unregister_agent", MagicMock()) + + agent = MagicMock() + agent._session_db = MagicMock() + + streaming._close_evicted_agent_at_session_boundary("dirty-session", agent) + + agent.shutdown_memory_provider.assert_not_called() + agent._session_db.close.assert_not_called() diff --git a/tests/test_v050259_sessiondb_fd_leak.py b/tests/test_v050259_sessiondb_fd_leak.py index 80290fc7a..6d3b748ee 100644 --- a/tests/test_v050259_sessiondb_fd_leak.py +++ b/tests/test_v050259_sessiondb_fd_leak.py @@ -85,9 +85,17 @@ def test_lru_eviction_closes_evicted_agent_session_db(): "is the original bug shape." ) - # Positive pattern: must close on the evicted agent. - assert "_evicted_agent._session_db.close()" in block, ( - "LRU eviction must close the evicted agent\'s _session_db. " + # Positive pattern: eviction must call the lifecycle helper, and the helper + # must close the evicted agent's SessionDB after provider teardown. + assert "_close_evicted_agent_at_session_boundary(_evicted_sid, _evicted_agent)" in block, ( + "LRU eviction must route the evicted agent through the session-boundary " + "close helper." + ) + helper_start = src.index("def _close_evicted_agent_at_session_boundary") + helper_end = src.index("\ndef _refresh_cached_agent_runtime", helper_start) + helper_block = src[helper_start:helper_end] + assert "session_db.close()" in helper_block, ( + "LRU eviction helper must close the evicted agent's _session_db. " "(Opus pre-release follow-up to PR #1421.)" )