fix(gateway): finalize session-end memory providers with transcript

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,29 +4598,28 @@ 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.
         try:
             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:

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

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