diff --git a/gateway/run_shutdown_session_end.py b/gateway/run_shutdown_session_end.py index cb268f7564..1d8e9ab71a 100644 --- a/gateway/run_shutdown_session_end.py +++ b/gateway/run_shutdown_session_end.py @@ -35,6 +35,24 @@ class GatewaySessionEndMixin: adapter, str(source.chat_id), message, source.platform.value if source.platform else "", "Failed to send session-end plugin message to %s:%s: %s", metadata=metadata) + async def _reset_session_off_loop(self, *, session_key: Optional[str] = None, **kwargs: Any) -> None: + """Plugin ``on_session_reset`` off-loop under the profile owning ``session_key``, bounded like + finalize: a viewer's ``/new`` resets its conversation on the owner (gateway/session_viewer_close).""" + def _call() -> None: + from hermes_cli.lifecycle import invoke_hook + invoke_hook("on_session_reset", **kwargs) + + try: + await asyncio.wait_for( + self._run_housekeeping_in_executor(self._run_release_in_profile_scope, _call, (), session_key), + timeout=self._FINALIZE_TIMEOUT_S, + ) + except TimeoutError: + logger.warning("Session reset hooks (%s) exceeded %ss; proceeding without blocking the event loop.", + kwargs.get("old_session_id"), self._FINALIZE_TIMEOUT_S) + except Exception: + logger.debug("Session reset hooks (%s) failed", kwargs.get("old_session_id"), exc_info=True) + async def _finalize_session_off_loop( self, *, session_id: Any, platform: str, reason: str, session_key: Optional[str] = None, **extra: Any, ) -> list[str]: diff --git a/gateway/session_viewer_close.py b/gateway/session_viewer_close.py index 4e6708cd10..c60dbdb1e6 100644 --- a/gateway/session_viewer_close.py +++ b/gateway/session_viewer_close.py @@ -33,6 +33,11 @@ def _busy(authority, session_id, target): return conn.execute(_BUSY_SQL, (session_id, session_id, target)).fetchone() is not None +# The viewer's /new: the classic CLI's ``new_session``; Ink's ``tui_close`` (/new or a resume switch, +# both of which main's TUI followed with on_session_reset for the session it opened next). +_RESET_REASONS = frozenset({'new_session', 'tui_close'}) + + class _PendingClose: """One viewer close; identity, not the reason, tells a settler its close is still the pending one.""" __slots__ = ('reason',) @@ -69,8 +74,13 @@ async def settle_viewer_close(authority, session_id): del authority.pending_closes[session_id] from gateway.session_acp_lifecycle import end_idle_viewer_session, reopen_if_reattached from gateway.session_runtime_workers import tracked_write - await tracked_write(authority, partial(end_idle_viewer_session, authority, session_id, close.reason), - then=partial(reopen_if_reattached, authority, session_id)) + ended = await tracked_write(authority, partial(end_idle_viewer_session, authority, session_id, close.reason), + then=partial(reopen_if_reattached, authority, session_id)) + if not ended: + await _requeue_if_busy(authority, session_id, live, target, close) + return + if live.subscribers: + return # a viewer reattached during the stamp; reopen_if_reattached reopened the row platform = (row or {}).get('source') or (live.source.platform.value if live.source else '') await _run_close_lifecycle(authority, live, target, platform, close.reason) @@ -79,6 +89,18 @@ def _still_pending(authority, session_id, live, close): return not live.subscribers and authority.pending_closes.get(session_id) is close +async def _requeue_if_busy(authority, session_id, live, target, close): + """The stamp was refused. If input was admitted after the idle check, the close goes back to the + drain's next idle point; any other refusal (row already ended) leaves nothing to finish.""" + if not await asyncio.to_thread(_busy, authority, session_id, target): + return + if live.subscribers or session_id in authority.pending_closes: + return # a reattach cancelled it, or a newer close replaced it + authority.pending_closes[session_id] = close + # The drain may have passed its idle point while the close was out of the map. + await settle_viewer_close(authority, session_id) + + async def _run_close_lifecycle(authority, live, target, platform, reason): """What main's client ran at exit and ``/new``, in the process that owns the session: the memory providers' end-of-session commit (the cached agent is evicted at a boundary, as a @@ -94,3 +116,9 @@ async def _run_close_lifecycle(authority, live, target, platform, reason): if messages: # No viewer is left to show them (session.detach carries no messages yet). logger.info('Session-end plugin message(s) for %s dropped: no viewer', target) + reset = getattr(runner, '_reset_session_off_loop', None) + if reason in _RESET_REASONS and reset is not None: + # The gateway's own /new arguments. The viewer opens its next session only after this + # detach, so the new id is not known here. + await reset(session_key=live.route, session_id=None, reason='new_session', platform=platform, + old_session_id=target, new_session_id=None) diff --git a/hermes_cli/cli_session_mixin.py b/hermes_cli/cli_session_mixin.py index 7888dc3aba..5d4dc6b93c 100644 --- a/hermes_cli/cli_session_mixin.py +++ b/hermes_cli/cli_session_mixin.py @@ -574,7 +574,9 @@ class CLISessionMixin: _mm.on_session_switch( self.session_id, parent_session_id=old_session_id or "", reset=True, reason="new_session") - self._notify_session_boundary("on_session_reset") + from hermes_cli.cli_shutdown import _gateway_owns_session_end + if not _gateway_owns_session_end(self.agent): # the gateway resets when end_session closes it + self._notify_session_boundary("on_session_reset") if not silent: if title: diff --git a/tests/gateway/test_viewer_close_lifecycle.py b/tests/gateway/test_viewer_close_lifecycle.py index 121f9d02a6..9e7e30408d 100644 --- a/tests/gateway/test_viewer_close_lifecycle.py +++ b/tests/gateway/test_viewer_close_lifecycle.py @@ -7,7 +7,7 @@ nothing ran until idle eviction; and a close sent while a turn was live left ``e """ import asyncio import threading -from types import SimpleNamespace +from types import MethodType, SimpleNamespace import pytest @@ -18,6 +18,7 @@ async def _owner(tmp_path, monkeypatch): from gateway import run from gateway.config import GatewayConfig from gateway.session import SessionStore + from gateway.run_shutdown_session_end import GatewaySessionEndMixin from gateway.session_authority import initialize_session_authority monkeypatch.setenv('HERMES_HOME', str(tmp_path)) @@ -29,11 +30,18 @@ async def _owner(tmp_path, monkeypatch): calls.append(('finalize', kwargs)) return [] + async def housekeeping(func, *args): + return func(*args) + + monkeypatch.setattr('hermes_cli.lifecycle.invoke_hook', lambda name, **kwargs: calls.append((name, kwargs))) runner = SimpleNamespace(session_store=store, _session_db=store._db, adapters={}, _draining=False, _cached_agent_for=lambda route: None, _evict_cached_agent_at_boundary=lambda key: calls.append(('commit', key)), - _finalize_session_off_loop=finalize) + _finalize_session_off_loop=finalize, _FINALIZE_TIMEOUT_S=10.0, + _run_housekeeping_in_executor=housekeeping, + _run_release_in_profile_scope=lambda target, args, key: target(*args)) runner._adapter_for_source = lambda source: runner.adapters.get(source.platform) + runner._reset_session_off_loop = MethodType(GatewaySessionEndMixin._reset_session_off_loop, runner) authority = await initialize_session_authority(runner, profile_id=str(tmp_path), instance_id='owner') return store._db, authority, calls @@ -83,6 +91,22 @@ async def test_last_viewer_close_commits_memory_and_finalizes_on_the_owner(tmp_p 'session_id': 'chat', 'platform': 'cli', 'reason': 'cli_close', 'session_key': route})] +@pytest.mark.asyncio +@pytest.mark.parametrize('reason', ['new_session', 'tui_close']) # classic CLI /new; Ink /new or resume switch +async def test_a_viewer_new_resets_the_session_on_the_owner_after_finalize(tmp_path, monkeypatch, reason): + db, authority, calls = await _owner(tmp_path, monkeypatch) + viewer, subscription = await _open(db, authority, tmp_path) + route = authority.sessions['chat'].route + + await _call(viewer, 'session.detach', session_id='chat', subscription_id=subscription, end_reason=reason) + + # Once, after finalize, with the gateway's own /new arguments (the next session is not open yet). + assert calls == [('commit', route), + ('finalize', {'session_id': 'chat', 'platform': 'cli', 'reason': reason, 'session_key': route}), + ('on_session_reset', {'session_id': None, 'reason': 'new_session', 'platform': 'cli', + 'old_session_id': 'chat', 'new_session_id': None})] + + @pytest.mark.asyncio async def test_close_while_a_turn_is_queued_completes_when_the_fifo_goes_idle(tmp_path, monkeypatch): from gateway.session_contract import SessionRef @@ -105,6 +129,44 @@ async def test_close_while_a_turn_is_queued_completes_when_the_fifo_goes_idle(tm +@pytest.mark.asyncio +async def test_a_turn_queued_during_a_slow_finalize_starts_after_it(tmp_path, monkeypatch): + from gateway import session_finite + from gateway.session_contract import SessionRef + from hermes_state_runtime import admit_session_input + + db, authority, _calls = await _owner(tmp_path, monkeypatch) + viewer, subscription = await _open(db, authority, tmp_path) + executed, finalizing, release = [], asyncio.Event(), asyncio.Event() + + async def execute(authority, ref, row): + executed.append(row['payload']['text']) + return 'done' + + async def slow_finalize(**kwargs): + finalizing.set() + await release.wait() # a plugin's on_session_finalize, up to the finalize timeout + return [] + + monkeypatch.setattr(session_finite, 'execute_finite_admission', execute) + authority.runner._finalize_session_off_loop = slow_finalize + ref = SessionRef(authority.profile_id, 'chat') + user = authority.sessions['chat'].source.user_id # the session's own operator may run turns + admit_session_input(db, epoch=authority.epoch, principal_id=user, session_id='chat', + request_id='r1', payload={'text': 'mid-turn'}) + await _call(viewer, 'session.detach', session_id='chat', subscription_id=subscription, end_reason='cli_close') + authority._schedule(ref) + # The turn ends, the drain's idle point completes the close and blocks in finalize. + await asyncio.wait_for(finalizing.wait(), 10) + + admit_session_input(db, epoch=authority.epoch, principal_id=user, session_id='chat', + request_id='r2', payload={'text': 'automation'}) + authority._schedule(ref) + assert executed == ['mid-turn'] # it waits for the plugins, as the review notes + release.set() + await async_wait_until(lambda: executed == ['mid-turn', 'automation']) + + def _gate_busy_check(monkeypatch): """Hold every settler inside its off-loop idle check until the test releases it.""" import gateway.session_viewer_close as viewer_close @@ -158,6 +220,67 @@ async def test_a_close_stamps_the_row_only_once_it_is_claimed(tmp_path, monkeypa assert [name for name, _ in calls] == ['commit', 'finalize'] +def _gate_stamp(monkeypatch): + """Hold the stamp's write (past its viewer check) until the test releases it.""" + import gateway.session_acp_lifecycle as lifecycle + stamp, entered, release = lifecycle.end_idle_local_session, [], threading.Event() + + def gated(*args, **kwargs): + entered.append(kwargs) + release.wait(timeout=30) + return stamp(*args, **kwargs) + + monkeypatch.setattr(lifecycle, 'end_idle_local_session', gated) + return entered, release + + +@pytest.mark.asyncio +async def test_a_viewer_reattaching_during_the_stamp_keeps_the_session_live(tmp_path, monkeypatch): + db, authority, calls = await _owner(tmp_path, monkeypatch) + viewer, subscription = await _open(db, authority, tmp_path) + entered, release = _gate_stamp(monkeypatch) + + closing = asyncio.create_task(_call(viewer, 'session.detach', session_id='chat', + subscription_id=subscription, end_reason='cli_close')) + await async_wait_until(lambda: len(entered) == 1) + await _call(_viewer(authority, tmp_path), 'session.resume', session_id='chat') + release.set() + await closing + + # The row is reopened under the new viewer, and its agent and Relay scope are left alone. + assert len(authority.sessions['chat'].subscribers) == 1 + assert calls == [] and authority.pending_closes == {} + assert db.get_session('chat')['ended_at'] is None + + +@pytest.mark.asyncio +async def test_input_admitted_during_the_stamp_defers_the_close_to_the_idle_point(tmp_path, monkeypatch): + from gateway.session_contract import SessionRef + from hermes_state_runtime import admit_session_input, cancel_session_input + + db, authority, calls = await _owner(tmp_path, monkeypatch) + viewer, subscription = await _open(db, authority, tmp_path) + entered, release = _gate_stamp(monkeypatch) + + closing = asyncio.create_task(_call(viewer, 'session.detach', session_id='chat', + subscription_id=subscription, end_reason='cli_close')) + await async_wait_until(lambda: len(entered) == 1) + row = admit_session_input(db, epoch=authority.epoch, principal_id='native', session_id='chat', + request_id='r1', payload={'text': 'automation'}) + release.set() + await closing + + # The busy FIFO refused the stamp: nothing ran yet and the close waits for the idle point. + assert calls == [] and 'chat' in authority.pending_closes + assert db.get_session('chat')['ended_at'] is None + + cancel_session_input(db, epoch=authority.epoch, admission_id=row['admission_id']) + authority._schedule(SessionRef(authority.profile_id, 'chat')) + await async_wait_until(lambda: len(calls) == 2) + assert db.get_session('chat')['end_reason'] == 'cli_close' + assert [name for name, _ in calls] == ['commit', 'finalize'] + + @pytest.mark.asyncio async def test_concurrent_settlers_run_the_close_lifecycle_once(tmp_path, monkeypatch): from gateway.session_viewer_close import settle_viewer_close diff --git a/tests/hermes_cli/test_gateway_session_end_ownership.py b/tests/hermes_cli/test_gateway_session_end_ownership.py index 63a474d210..96fa8a7451 100644 --- a/tests/hermes_cli/test_gateway_session_end_ownership.py +++ b/tests/hermes_cli/test_gateway_session_end_ownership.py @@ -45,6 +45,16 @@ def test_new_does_not_finalize_a_gateway_session_in_the_cli(mock_finalize, mock_ mock_finalize.assert_not_called() +@patch("hermes_cli.lifecycle.invoke_hook") +@patch("hermes_cli.lifecycle.finalize_session") +def test_new_leaves_on_session_reset_to_the_gateway(mock_finalize, mock_invoke_hook, cli_mod): + cli = cli_mod.HermesCLI() + cli.agent = _gateway_agent("gw-session") + cli.new_session(silent=True) + # The detach's ``new_session`` reason makes the owner fire it (tests/gateway/test_viewer_close_lifecycle.py). + assert not [c for c in mock_invoke_hook.call_args_list if c.args == ("on_session_reset",)] + + @patch("hermes_cli.lifecycle.invoke_hook") def test_exit_mid_turn_leaves_on_session_end_to_the_gateway_turn(mock_invoke_hook, cli_mod): cli = cli_mod.HermesCLI()