From ab21fbfd89f4f168afcc024c3cf329140671ea98 Mon Sep 17 00:00:00 2001 From: Teknium Date: Wed, 8 Apr 2026 04:22:55 -0700 Subject: [PATCH] fix: add gateway coverage for session boundary hooks, move test to tests/cli/ - Fire on_session_finalize and on_session_reset in gateway _handle_reset_command() - Fire on_session_finalize during gateway stop() for each active agent - Move CLI test from tests/ root to tests/cli/ (matches recent restructure) - Add 5 gateway tests covering reset hooks, ordering, shutdown, and error handling - Place on_session_reset after new session is guaranteed to exist (covers the get_or_create_session fallback path) --- gateway/run.py | 30 +++- .../{ => cli}/test_session_boundary_hooks.py | 0 tests/gateway/test_session_boundary_hooks.py | 158 ++++++++++++++++++ 3 files changed, 186 insertions(+), 2 deletions(-) rename tests/{ => cli}/test_session_boundary_hooks.py (100%) create mode 100644 tests/gateway/test_session_boundary_hooks.py diff --git a/gateway/run.py b/gateway/run.py index 149b1f59d..7a551be16 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -1481,6 +1481,14 @@ class GatewayRunner: logger.debug("Interrupted running agent for session %s during shutdown", session_key[:20]) except Exception as e: logger.debug("Failed interrupting agent during shutdown: %s", e) + # Fire plugin on_session_finalize hook before memory shutdown + try: + from hermes_cli.plugins import invoke_hook as _invoke_hook + _invoke_hook("on_session_finalize", + session_id=getattr(agent, 'session_id', None), + platform="gateway") + except Exception: + pass # Shut down memory provider at actual session boundary try: if hasattr(agent, 'shutdown_memory_provider'): @@ -3274,6 +3282,15 @@ class GatewayRunner: # the configured default instead of the previously switched model. self._session_model_overrides.pop(session_key, None) + # Fire plugin on_session_finalize hook (session boundary) + try: + from hermes_cli.plugins import invoke_hook as _invoke_hook + _old_sid = old_entry.session_id if old_entry else None + _invoke_hook("on_session_finalize", session_id=_old_sid, + platform=source.platform.value if source.platform else "") + except Exception: + pass + # Emit session:end hook (session is ending) await self.hooks.emit("session:end", { "platform": source.platform.value if source.platform else "", @@ -3287,7 +3304,7 @@ class GatewayRunner: "user_id": source.user_id, "session_key": session_key, }) - + # Resolve session config info to surface to the user try: session_info = self._format_session_info() @@ -3298,9 +3315,18 @@ class GatewayRunner: header = "✨ Session reset! Starting fresh." else: # No existing session, just create one - self.session_store.get_or_create_session(source, force_new=True) + new_entry = self.session_store.get_or_create_session(source, force_new=True) header = "✨ New session started!" + # Fire plugin on_session_reset hook (new session guaranteed to exist) + try: + from hermes_cli.plugins import invoke_hook as _invoke_hook + _new_sid = new_entry.session_id if new_entry else None + _invoke_hook("on_session_reset", session_id=_new_sid, + platform=source.platform.value if source.platform else "") + except Exception: + pass + if session_info: return f"{header}\n\n{session_info}" return header diff --git a/tests/test_session_boundary_hooks.py b/tests/cli/test_session_boundary_hooks.py similarity index 100% rename from tests/test_session_boundary_hooks.py rename to tests/cli/test_session_boundary_hooks.py diff --git a/tests/gateway/test_session_boundary_hooks.py b/tests/gateway/test_session_boundary_hooks.py new file mode 100644 index 000000000..31e02980a --- /dev/null +++ b/tests/gateway/test_session_boundary_hooks.py @@ -0,0 +1,158 @@ +"""Tests that on_session_finalize and on_session_reset plugin hooks fire in the gateway.""" +from datetime import datetime +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from gateway.config import GatewayConfig, Platform, PlatformConfig +from gateway.platforms.base import MessageEvent +from gateway.session import SessionEntry, SessionSource, build_session_key + + +def _make_source() -> SessionSource: + return SessionSource( + platform=Platform.TELEGRAM, + user_id="u1", + chat_id="c1", + user_name="tester", + chat_type="dm", + ) + + +def _make_event(text: str) -> MessageEvent: + return MessageEvent(text=text, source=_make_source(), message_id="m1") + + +def _make_runner(): + from gateway.run import GatewayRunner + + runner = object.__new__(GatewayRunner) + runner.config = GatewayConfig( + platforms={Platform.TELEGRAM: PlatformConfig(enabled=True, token="***")} + ) + adapter = MagicMock() + adapter.send = AsyncMock() + runner.adapters = {Platform.TELEGRAM: adapter} + runner._voice_mode = {} + runner.hooks = SimpleNamespace(emit=AsyncMock(), loaded_hooks=False) + runner._session_model_overrides = {} + runner._pending_model_notes = {} + runner._background_tasks = set() + + session_key = build_session_key(_make_source()) + session_entry = SessionEntry( + session_key=session_key, + session_id="sess-old", + created_at=datetime.now(), + updated_at=datetime.now(), + platform=Platform.TELEGRAM, + chat_type="dm", + ) + new_session_entry = SessionEntry( + session_key=session_key, + session_id="sess-new", + created_at=datetime.now(), + updated_at=datetime.now(), + platform=Platform.TELEGRAM, + chat_type="dm", + ) + runner.session_store = MagicMock() + runner.session_store.get_or_create_session.return_value = new_session_entry + runner.session_store.reset_session.return_value = new_session_entry + runner.session_store._entries = {session_key: session_entry} + runner.session_store._generate_session_key.return_value = session_key + runner._running_agents = {} + runner._pending_messages = {} + runner._pending_approvals = {} + runner._session_db = None + runner._agent_cache_lock = None + runner._is_user_authorized = lambda _source: True + runner._format_session_info = lambda: "" + + return runner + + +@pytest.mark.asyncio +@patch("hermes_cli.plugins.invoke_hook") +async def test_reset_fires_finalize_hook(mock_invoke_hook): + """/new must fire on_session_finalize with the OLD session id.""" + runner = _make_runner() + + await runner._handle_reset_command(_make_event("/new")) + + mock_invoke_hook.assert_any_call( + "on_session_finalize", session_id="sess-old", platform="telegram" + ) + + +@pytest.mark.asyncio +@patch("hermes_cli.plugins.invoke_hook") +async def test_reset_fires_reset_hook(mock_invoke_hook): + """/new must fire on_session_reset with the NEW session id.""" + runner = _make_runner() + + await runner._handle_reset_command(_make_event("/new")) + + mock_invoke_hook.assert_any_call( + "on_session_reset", session_id="sess-new", platform="telegram" + ) + + +@pytest.mark.asyncio +@patch("hermes_cli.plugins.invoke_hook") +async def test_finalize_before_reset(mock_invoke_hook): + """on_session_finalize must fire before on_session_reset.""" + runner = _make_runner() + + await runner._handle_reset_command(_make_event("/new")) + + calls = [c for c in mock_invoke_hook.call_args_list + if c[0][0] in ("on_session_finalize", "on_session_reset")] + hook_names = [c[0][0] for c in calls] + assert hook_names == ["on_session_finalize", "on_session_reset"] + + +@pytest.mark.asyncio +@patch("hermes_cli.plugins.invoke_hook") +async def test_shutdown_fires_finalize_for_active_agents(mock_invoke_hook): + """Gateway stop() must fire on_session_finalize for each active agent.""" + from gateway.run import GatewayRunner + + runner = object.__new__(GatewayRunner) + runner._running = True + runner._background_tasks = set() + runner._pending_messages = {} + runner._pending_approvals = {} + runner._shutdown_event = MagicMock() + runner.adapters = {} + runner._exit_reason = "test" + + agent1 = MagicMock() + agent1.session_id = "sess-a" + agent2 = MagicMock() + agent2.session_id = "sess-b" + runner._running_agents = {"key-a": agent1, "key-b": agent2} + + with patch("gateway.status.remove_pid_file"), \ + patch("gateway.status.write_runtime_status"): + await runner.stop() + + finalize_calls = [ + c for c in mock_invoke_hook.call_args_list + if c[0][0] == "on_session_finalize" + ] + session_ids = {c[1]["session_id"] for c in finalize_calls} + assert session_ids == {"sess-a", "sess-b"} + + +@pytest.mark.asyncio +@patch("hermes_cli.plugins.invoke_hook", side_effect=Exception("boom")) +async def test_hook_error_does_not_break_reset(mock_invoke_hook): + """Plugin hook errors must not prevent /new from completing.""" + runner = _make_runner() + + result = await runner._handle_reset_command(_make_event("/new")) + + # Should still return a success message despite hook errors + assert "Session reset" in result or "New session" in result