diff --git a/gateway/session_viewer_close.py b/gateway/session_viewer_close.py index 8133cc6255..9e1fdf3f3e 100644 --- a/gateway/session_viewer_close.py +++ b/gateway/session_viewer_close.py @@ -33,9 +33,17 @@ def _busy(authority, session_id, target): return conn.execute(_BUSY_SQL, (session_id, session_id, target)).fetchone() is not None +class _PendingClose: + """One viewer close; identity, not the reason, tells a settler its close is still the pending one.""" + __slots__ = ('reason',) + + def __init__(self, reason): + self.reason = reason + + async def close_viewer_session(authority, session_id, reason): """``session.detach {end_reason}`` left *session_id* without viewers: close it now, or once idle.""" - authority.pending_closes[session_id] = reason + authority.pending_closes[session_id] = _PendingClose(reason) await settle_viewer_close(authority, session_id) @@ -43,25 +51,36 @@ async def settle_viewer_close(authority, session_id): """Complete a pending viewer close once *session_id* has no viewer and an idle FIFO: stamp the row with the viewer's reason, commit the cached agent's memory and fire ``on_session_finalize``. The drain calls this at every idle point; a viewer attaching meanwhile cancels the close.""" - reason = authority.pending_closes.get(session_id) + close = authority.pending_closes.get(session_id) live = authority.sessions.get(session_id) - if reason is None or live is None or live.subscribers: + if close is None or live is None or live.subscribers: authority.pending_closes.pop(session_id, None) return 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, reason), + await tracked_write(authority, partial(end_idle_viewer_session, authority, session_id, close.reason), then=partial(reopen_if_reattached, authority, session_id)) target = _physical_target(authority, session_id) - if live.subscribers or authority.pending_closes.get(session_id) != reason: + # Every await below can let a viewer reattach (attach drops the pending close) or another + # settler run, so the close is re-checked after each one. + if not _still_pending(authority, session_id, live, close): return if await asyncio.to_thread(_busy, authority, session_id, target): return # the drain's next idle point completes it - authority.pending_closes.pop(session_id, None) - await _run_close_lifecycle(authority, live, target, reason) + row = await asyncio.to_thread(authority.db.get_session, target) + if not _still_pending(authority, session_id, live, close): + return + # Claimed with no await since the check: exactly one settler runs the lifecycle. + del authority.pending_closes[session_id] + platform = (row or {}).get('source') or (live.source.platform.value if live.source else '') + await _run_close_lifecycle(authority, live, target, platform, close.reason) -async def _run_close_lifecycle(authority, live, target, reason): +def _still_pending(authority, session_id, live, close): + return not live.subscribers and authority.pending_closes.get(session_id) is close + + +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 gateway ``/new`` does) and the finalize hooks under the owning profile's scope.""" @@ -72,8 +91,6 @@ async def _run_close_lifecycle(authority, live, target, reason): finalize = getattr(runner, '_finalize_session_off_loop', None) if finalize is None: return - row = await asyncio.to_thread(authority.db.get_session, target) - platform = (row or {}).get('source') or (live.source.platform.value if live.source else '') messages = await finalize(session_id=target, platform=platform, reason=reason, session_key=live.route) if messages: # No viewer is left to show them (session.detach carries no messages yet). diff --git a/hermes_cli/cli_shutdown.py b/hermes_cli/cli_shutdown.py index d8ed31451b..efb37915ec 100644 --- a/hermes_cli/cli_shutdown.py +++ b/hermes_cli/cli_shutdown.py @@ -213,7 +213,8 @@ def _emit_interrupted_session_end(cli, *, reason: str = "keyboard_interrupt") -> with suppress(Exception): agent.interrupt(reason.replace("_", " ")) - if session_id in _cli()._handed_off_session_ids: # gateway owns the lifecycle now + # The gateway owns the lifecycle now; a gateway turn fires on_session_end itself when interrupted. + if session_id in _cli()._handed_off_session_ids or _gateway_owns_session_end(agent): return if session_id: with suppress(Exception): @@ -232,7 +233,9 @@ def _notify_single_query_session_finalize(cli, *, reason: str = "shutdown") -> N agent, session_id = _oneshot_agent_and_session(cli) if session_id in _cli()._single_query_finalize_attempted_session_ids: return - if session_id in _cli()._handed_off_session_ids: # gateway owns the lifecycle now + # The gateway owns the lifecycle now; for a gateway session the flush's end_session detached with + # an end_reason, so the gateway finalizes it. + if session_id in _cli()._handed_off_session_ids or _gateway_owns_session_end(agent): return try: diff --git a/tests/gateway/test_viewer_close_lifecycle.py b/tests/gateway/test_viewer_close_lifecycle.py index 2a6a77c28f..f99a33097e 100644 --- a/tests/gateway/test_viewer_close_lifecycle.py +++ b/tests/gateway/test_viewer_close_lifecycle.py @@ -5,6 +5,8 @@ Regression for #134889 LC-1/LC-2/LC-6. Main's classic CLI and TUI ran the memory process at exit and ``/new``. Under the gateway the client is a viewer with no memory manager, so nothing ran until idle eviction; and a close sent while a turn was live left ``ended_at`` unset for good. """ +import asyncio +import threading from types import SimpleNamespace import pytest @@ -101,3 +103,55 @@ async def test_close_while_a_turn_is_queued_completes_when_the_fifo_goes_idle(tm assert db.get_session('chat')['end_reason'] == 'cli_close' assert [name for name, _ in calls] == ['commit', 'finalize'] + + +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 + entered, release = [], threading.Event() + + def busy(*args): + entered.append(args) + release.wait(timeout=30) + return False + + monkeypatch.setattr(viewer_close, '_busy', busy) + return entered, release + + +@pytest.mark.asyncio +async def test_a_viewer_reattaching_during_the_idle_check_cancels_the_close(tmp_path, monkeypatch): + db, authority, calls = await _owner(tmp_path, monkeypatch) + viewer, subscription = await _open(db, authority, tmp_path) + entered, release = _gate_busy_check(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 conversation is being watched again: its agent stays cached and nothing is finalized. + assert calls == [] + + +@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 + + db, authority, calls = await _owner(tmp_path, monkeypatch) + viewer, subscription = await _open(db, authority, tmp_path) + entered, release = _gate_busy_check(monkeypatch) + + # The detach's own settle and a drain idle point both reach the idle check. + 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) + drain_idle = asyncio.create_task(settle_viewer_close(authority, 'chat')) + await async_wait_until(lambda: len(entered) == 2) + release.set() + await asyncio.gather(closing, drain_idle) + + assert [name for name, _ in calls] == ['commit', 'finalize'] + assert db.get_session('chat')['end_reason'] == 'cli_close' diff --git a/tests/hermes_cli/test_gateway_session_end_ownership.py b/tests/hermes_cli/test_gateway_session_end_ownership.py index f5420348d7..63a474d210 100644 --- a/tests/hermes_cli/test_gateway_session_end_ownership.py +++ b/tests/hermes_cli/test_gateway_session_end_ownership.py @@ -6,6 +6,7 @@ and the gateway, which owns the agent, the memory manager and the Relay scope, r well finalized every session twice, once in a process that holds none of that state; and the CLI's mid-turn ``on_session_end`` duplicated the one the gateway's interrupted turn fires. """ +from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest @@ -53,3 +54,28 @@ def test_exit_mid_turn_leaves_on_session_end_to_the_gateway_turn(mock_invoke_hoo with patch.object(cli_mod, "_run_cleanup"), patch.object(cli, "_print_exit_summary"): cli._tui_shutdown() assert not [c for c in mock_invoke_hook.call_args_list if c.args == ("on_session_end",)] + + +@patch("hermes_cli.lifecycle.finalize_session") +def test_one_shot_ends_a_gateway_session_through_the_gateway_only(mock_finalize, cli_mod, monkeypatch): + store = MagicMock() + agent = GatewayAgent(connection=MagicMock(), session_db=store, session_id="gw-session") + released = [] + one_shot = SimpleNamespace(agent=agent, session_id="gw-session", conversation_history=[], _session_db=store, + _release_active_session=lambda: released.append(True)) + monkeypatch.setattr(cli_mod, "_wait_for_oneshot_background_completions", lambda _cli: None) + monkeypatch.setattr(cli_mod, "_run_cleanup", lambda **_kwargs: None) + + cli_mod._finalize_single_query(one_shot) + + # The detach carries the end reason, so the gateway runs the lifecycle; the CLI never does. + store.end_session.assert_called_once_with("gw-session", "cli_close") + mock_finalize.assert_not_called() + assert released == [True] + + +@patch("hermes_cli.lifecycle.invoke_hook") +def test_interrupted_one_shot_leaves_on_session_end_to_the_gateway_turn(mock_invoke_hook, cli_mod): + one_shot = SimpleNamespace(agent=_gateway_agent("gw-session"), session_id="gw-session") + cli_mod._emit_interrupted_session_end(one_shot) + assert not [c for c in mock_invoke_hook.call_args_list if c.args == ("on_session_end",)]