Import AITURK IDE 1.0.0-beta.1 from Hermes 63279301; preserve MIT license

tests/gateway/test_shutdown_cache_cleanup.py

@@ -0,0 +1,178 @@
+"""Regression tests for gateway shutdown cleaning up cached agent memory providers (issue #11205).
+
+When the gateway shuts down, ``stop()`` called ``_finalize_shutdown_agents()``
+which only drained agents in ``_running_agents``.  Idle agents sitting in
+``_agent_cache`` (LRU cache) were never cleaned up, so their
+``MemoryProvider.on_session_end()`` hooks never fired.
+
+The fix adds an explicit sweep of ``_agent_cache`` after
+``_finalize_shutdown_agents`` in the ``_stop_impl`` coroutine.
+"""
+
+import asyncio
+import threading
+from collections import OrderedDict
+from unittest.mock import MagicMock
+
+import pytest
+
+# Import the module (not the class) to reach stop() and helpers
+import gateway.run as gw_mod
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+class _FakeGateway:
+    """Minimal stand-in with just enough state for ``stop()`` to run."""
+
+    def __init__(self):
+        self._running = True
+        self._draining = False
+        self._restart_requested = False
+        self._restart_detached = False
+        self._restart_via_service = False
+        self._stop_task = None
+        self._exit_cleanly = False
+        self._exit_with_failure = False
+        self._exit_reason = None
+        self._exit_code = None
+        self._restart_drain_timeout = 0.01
+        self._running_agents = {}
+        self._running_agents_ts = {}
+        self._agent_cache = OrderedDict()
+        self._agent_cache_lock = threading.Lock()
+        self.adapters = {}
+        self._background_tasks = set()
+        self._failed_platforms = []
+        self._shutdown_event = asyncio.Event()
+        self._pending_messages = {}
+        self._pending_approvals = {}
+        self._busy_ack_ts = {}
+
+    def _running_agent_count(self):
+        return len(self._running_agents)
+
+    def _active_cron_job_count(self):
+        # stop() reads this alongside _running_agent_count when logging the
+        # drain snapshot (#60432) -- this fake has no cron scheduler, so
+        # there's never in-flight cron work to report.
+        return 0
+
+    def _active_api_run_count(self):
+        # The shutdown log also reports adapter-owned API work (#63529).
+        # This fake has no API server adapter, so it is always idle.
+        return 0
+
+    def _update_runtime_status(self, *_a, **_kw):
+        pass
+
+    def _clear_plugin_message_injector(self):
+        pass
+
+    async def _run_in_executor_with_context(self, func, *args):
+        # stop() offloads agent-resource cleanup off the loop (#53175); run
+        # inline in tests so the bounded-cleanup path is exercised.
+        return func(*args)
+
+    async def _cleanup_agent_resources_off_loop(self, agent, *, context=""):
+        # Mirror the real bounded helper, inline (no executor/timeout) so the
+        # fake exercises the same call shape stop() now uses.
+        self._cleanup_agent_resources(agent)
+
+    async def _notify_active_sessions_of_shutdown(self):
+        pass
+
+    async def _cancel_secondary_profile_reconnect_tasks(self):
+        pass
+
+    async def _drain_active_agents(self, timeout, cron_timeout=None):
+        return {}, False
+
+    async def _finalize_shutdown_agents(self, agents):
+        for agent in agents.values():
+            self._cleanup_agent_resources(agent)
+
+    def _cleanup_agent_resources(self, agent):
+        if agent is None:
+            return
+        try:
+            if hasattr(agent, "shutdown_memory_provider"):
+                agent.shutdown_memory_provider()
+        except Exception:
+            pass
+        try:
+            if hasattr(agent, "close"):
+                agent.close()
+        except Exception:
+            pass
+
+    def _evict_cached_agent(self, key):
+        pass
+
+    def _release_running_agent_state(self, session_key, **_kwargs):
+        agent = self._running_agents.pop(session_key, None)
+        self._running_agents_ts.pop(session_key, None)
+        self._cleanup_agent_resources(agent)
+        return agent is not None
+
+
+def _make_mock_agent():
+    a = MagicMock()
+    a.shutdown_memory_provider = MagicMock()
+    a.close = MagicMock()
+    return a
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+class TestCachedAgentCleanupOnShutdown:
+    """Verify that ``stop()`` calls ``_cleanup_agent_resources`` on idle
+    cached agents, triggering ``shutdown_memory_provider()`` (which calls
+    ``on_session_end``)."""
+
+    @pytest.mark.asyncio
+    async def test_cached_agent_memory_provider_shut_down(self):
+        """A cached agent's shutdown_memory_provider is called during gateway stop."""
+        gw = _FakeGateway()
+        agent = _make_mock_agent()
+        gw._agent_cache["session-1"] = (agent, "sig-123")
+
+        # Call the real stop() from GatewayRunner
+        await gw_mod.GatewayRunner.stop(gw)
+
+        agent.shutdown_memory_provider.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_cache_cleared_after_shutdown(self):
+        """The _agent_cache dict is cleared after stop."""
+        gw = _FakeGateway()
+        agent = _make_mock_agent()
+        gw._agent_cache["s1"] = (agent, "sig1")
+
+        await gw_mod.GatewayRunner.stop(gw)
+
+        assert len(gw._agent_cache) == 0
+
+
+class TestRunningAgentsNotDoubleCleaned:
+    """Verify behavior when agents appear in both _running_agents and _agent_cache."""
+
+    @pytest.mark.asyncio
+    async def test_running_and_cached_agent_cleaned_at_least_once(self):
+        """An agent in both _running_agents and _agent_cache gets
+        shutdown_memory_provider called at least once."""
+        gw = _FakeGateway()
+        shared = _make_mock_agent()
+
+        gw._running_agents["s1"] = shared
+        gw._agent_cache["s1"] = (shared, "sig1")
+
+        await gw_mod.GatewayRunner.stop(gw)
+
+        # Called at least once — either from _finalize_shutdown_agents
+        # or from the cache sweep (or both)
+        assert shared.shutdown_memory_provider.call_count >= 1