From e920d9000cfbcbfa009210baaa5184decca71e42 Mon Sep 17 00:00:00 2001 From: Ted Li Date: Sat, 18 Apr 2026 12:44:50 -0700 Subject: [PATCH] Fix gateway session-end memory finalization --- gateway/run.py | 135 ++++++++++++++++------ tests/gateway/test_async_memory_flush.py | 52 ++++++++- tests/gateway/test_session_model_reset.py | 27 +++++ 3 files changed, 180 insertions(+), 34 deletions(-) diff --git a/gateway/run.py b/gateway/run.py index 1525ad147766..e27834f5815e 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -875,6 +875,91 @@ async def _async_flush_memories( session_key, ) + def _get_cached_or_running_agent(self, session_key: Optional[str]) -> Any: + """Return the current agent instance for a session key, if any.""" + if not session_key: + return None + + agent = None + _cache_lock = getattr(self, "_agent_cache_lock", None) + _cache = getattr(self, "_agent_cache", None) + if _cache_lock is not None and _cache is not None: + with _cache_lock: + _cached = _cache.get(session_key) + agent = ( + _cached[0] + if isinstance(_cached, tuple) + else _cached + if _cached + else None + ) + + if agent is None: + _running_agents = getattr(self, "_running_agents", None) or {} + agent = _running_agents.get(session_key) + + if agent is _AGENT_PENDING_SENTINEL: + return None + return agent + + def _shutdown_session_memory_provider( + self, + session_id: str, + session_key: Optional[str] = None, + agent: Any = None, + ) -> Any: + """Dispatch transcript-aware session-end memory shutdown for one agent.""" + target_agent = agent or self._get_cached_or_running_agent(session_key) + if target_agent is None: + return None + + history = self.session_store.load_transcript(session_id) or [] + try: + if hasattr(target_agent, "shutdown_memory_provider"): + target_agent.shutdown_memory_provider(history) + except Exception as e: + logger.debug( + "Session-end memory shutdown failed for %s: %s", + session_id, + e, + ) + return target_agent + + async def _async_finalize_session_end( + self, + session_id: str, + session_key: Optional[str] = None, + *, + agent: Any = None, + close_agent: bool = False, + evict_cached: bool = False, + ) -> None: + """Flush built-in memory, then finalize provider memory with transcript.""" + await self._async_flush_memories(session_id, session_key) + + loop = asyncio.get_running_loop() + target_agent = await loop.run_in_executor( + None, + self._shutdown_session_memory_provider, + session_id, + session_key, + agent, + ) + + if close_agent and target_agent is not None: + try: + if hasattr(target_agent, "close"): + await loop.run_in_executor(None, target_agent.close) + except Exception as e: + logger.debug( + "Session-end agent close failed for %s: %s", + session_id, + e, + ) + + if evict_cached and session_key: + self._evict_cached_agent(session_key) + @property def should_exit_cleanly(self) -> bool: return self._exit_cleanly @@ -2100,27 +2185,12 @@ async def _session_expiry_watcher(self, interval: int = 300): for key, entry in _expired_entries: try: - await self._async_flush_memories(entry.session_id, key) - # Shut down memory provider and close tool resources - # on the cached agent. Idle agents live in - # _agent_cache (not _running_agents), so look there. - _cached_agent = None - _cache_lock = getattr(self, "_agent_cache_lock", None) - if _cache_lock is not None: - with _cache_lock: - _cached = self._agent_cache.get(key) - _cached_agent = _cached[0] if isinstance(_cached, tuple) else _cached if _cached else None - # Fall back to _running_agents in case the agent is - # still mid-turn when the expiry fires. - if _cached_agent is None: - _cached_agent = self._running_agents.get(key) - if _cached_agent and _cached_agent is not _AGENT_PENDING_SENTINEL: - self._cleanup_agent_resources(_cached_agent) - # Drop the cache entry so the AIAgent (and its LLM - # clients, tool schemas, memory provider refs) can - # be garbage-collected. Otherwise the cache grows - # unbounded across the gateway's lifetime. - self._evict_cached_agent(key) + await self._async_finalize_session_end( + entry.session_id, + key, + close_agent=True, + evict_cached=True, + ) # Mark as flushed and persist to disk so the flag # survives gateway restarts. with self.session_store._lock: @@ -4528,6 +4598,8 @@ async def _handle_reset_command(self, event: MessageEvent) -> str: # Get existing session key session_key = self._session_key_for_source(source) + _old_agent = self._get_cached_or_running_agent(session_key) + old_entry = None # Flush memories in the background (fire-and-forget) so the user # gets the "Session reset!" response immediately. @@ -4535,22 +4607,19 @@ async def _handle_reset_command(self, event: MessageEvent) -> str: old_entry = self.session_store._entries.get(session_key) if old_entry: _flush_task = asyncio.create_task( - self._async_flush_memories(old_entry.session_id, session_key) + self._async_finalize_session_end( + old_entry.session_id, + session_key=session_key, + agent=_old_agent, + close_agent=True, + ) ) self._background_tasks.add(_flush_task) _flush_task.add_done_callback(self._background_tasks.discard) except Exception as e: - logger.debug("Gateway memory flush on reset failed: %s", e) - # Close tool resources on the old agent (terminal sandboxes, browser - # daemons, background processes) before evicting from cache. - # Guard with getattr because test fixtures may skip __init__. - _cache_lock = getattr(self, "_agent_cache_lock", None) - if _cache_lock is not None: - with _cache_lock: - _cached = self._agent_cache.get(session_key) - _old_agent = _cached[0] if isinstance(_cached, tuple) else _cached if _cached else None - if _old_agent is not None: - self._cleanup_agent_resources(_old_agent) + logger.debug("Gateway session finalization on reset failed: %s", e) + if old_entry is None and _old_agent is not None: + self._cleanup_agent_resources(_old_agent) self._evict_cached_agent(session_key) try: diff --git a/tests/gateway/test_async_memory_flush.py b/tests/gateway/test_async_memory_flush.py index 0d731949048c..40d1a6620d42 100644 --- a/tests/gateway/test_async_memory_flush.py +++ b/tests/gateway/test_async_memory_flush.py @@ -7,10 +7,11 @@ 4. The background watcher can detect expired sessions """ +import threading import pytest from datetime import datetime, timedelta from pathlib import Path -from unittest.mock import patch, MagicMock +from unittest.mock import AsyncMock, patch, MagicMock from gateway.config import Platform, GatewayConfig, SessionResetPolicy from gateway.session import SessionSource, SessionStore, SessionEntry @@ -247,3 +248,52 @@ def test_legacy_entry_without_field_defaults_false(self): } entry = SessionEntry.from_dict(data) assert entry.memory_flushed is False + + +@pytest.mark.asyncio +async def test_session_expiry_watcher_finalizes_provider_memory_with_transcript(): + """Expired sessions should use the transcript-aware finalization path.""" + from gateway.run import GatewayRunner + + runner = object.__new__(GatewayRunner) + runner._running = True + runner._background_tasks = set() + runner._running_agents = {} + runner._pending_messages = {} + runner._pending_approvals = {} + runner.hooks = MagicMock() + + expired_entry = SessionEntry( + session_key="agent:main:telegram:dm:123", + session_id="sess-expired", + created_at=datetime.now() - timedelta(hours=2), + updated_at=datetime.now() - timedelta(hours=2), + platform=Platform.TELEGRAM, + chat_type="dm", + ) + lock = threading.Lock() + runner.session_store = MagicMock() + runner.session_store._entries = {expired_entry.session_key: expired_entry} + runner.session_store._lock = lock + runner.session_store._is_session_expired.return_value = True + + runner._async_finalize_session_end = AsyncMock() + + sleep_calls = {"count": 0} + + async def _fake_sleep(_delay): + sleep_calls["count"] += 1 + if sleep_calls["count"] >= 2: + runner._running = False + + with patch("gateway.run.asyncio.sleep", new=_fake_sleep): + await runner._session_expiry_watcher(interval=1) + + runner._async_finalize_session_end.assert_awaited_once_with( + "sess-expired", + "agent:main:telegram:dm:123", + close_agent=True, + evict_cached=True, + ) + assert expired_entry.memory_flushed is True + runner.session_store._save.assert_called_once() diff --git a/tests/gateway/test_session_model_reset.py b/tests/gateway/test_session_model_reset.py index 6529f3a11dd2..f7bf014a3579 100644 --- a/tests/gateway/test_session_model_reset.py +++ b/tests/gateway/test_session_model_reset.py @@ -1,4 +1,6 @@ """Tests that /new (and its /reset alias) clears the session-scoped model override.""" +import asyncio +import threading from datetime import datetime from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock @@ -124,3 +126,28 @@ async def test_new_command_only_clears_own_session(): assert session_key not in runner._session_model_overrides assert other_key in runner._session_model_overrides + + +@pytest.mark.asyncio +async def test_new_command_schedules_transcript_aware_session_finalize(): + """/new should finalize the old agent through the session-end helper.""" + runner = _make_runner() + session_key = build_session_key(_make_source()) + old_agent = MagicMock() + runner._agent_cache = {session_key: (old_agent, "sig")} + runner._agent_cache_lock = threading.Lock() + runner._async_finalize_session_end = AsyncMock() + + await runner._handle_reset_command(_make_event("/new")) + if runner._background_tasks: + await asyncio.gather(*runner._background_tasks) + + runner._async_finalize_session_end.assert_awaited_once_with( + "sess-1", + session_key=session_key, + agent=old_agent, + close_agent=True, + ) + old_agent.close.assert_not_called() + with runner._agent_cache_lock: + assert session_key not in runner._agent_cache