diff --git a/cli.py b/cli.py index 3b555a288fa..70c30f00730 100644 --- a/cli.py +++ b/cli.py @@ -890,6 +890,10 @@ def _cleanup_all_browsers(*args, **kwargs): # Guard to prevent cleanup from running multiple times on exit _cleanup_done = False +# One-shot CLI finalization runs before process cleanup so plugins can observe +# the session boundary while the agent is still attached. If a signal lands in +# that narrow window, atexit cleanup must not emit that session finalize again. +_single_query_finalize_attempted_session_ids: set[str | None] = set() # Weak reference to the active AIAgent for memory provider shutdown at exit _active_agent_ref = None _deferred_agent_startup_done = False @@ -989,11 +993,13 @@ def _run_cleanup(*, notify_session_finalize: bool = True): # Shut down memory provider (on_session_end + shutdown_all) at actual # session boundary — NOT per-turn inside run_conversation(). if notify_session_finalize: - _notify_session_finalize( - session_id=_active_agent_ref.session_id if _active_agent_ref else None, - platform="cli", - reason="shutdown", - ) + cleanup_session_id = _active_agent_ref.session_id if _active_agent_ref else None + if _should_emit_cleanup_session_finalize(cleanup_session_id): + _notify_session_finalize( + session_id=cleanup_session_id, + platform="cli", + reason="shutdown", + ) try: if _active_agent_ref and hasattr(_active_agent_ref, 'shutdown_memory_provider'): # Forward the agent's own transcript so memory providers' @@ -1011,6 +1017,14 @@ def _run_cleanup(*, notify_session_finalize: bool = True): pass +def _should_emit_cleanup_session_finalize(session_id: str | None) -> bool: + if not _single_query_finalize_attempted_session_ids: + return True + if session_id is None: + return False + return session_id not in _single_query_finalize_attempted_session_ids + + def _notify_session_finalize( *, session_id: str | None, @@ -1068,11 +1082,17 @@ def _emit_interrupted_session_end(cli, *, reason: str = "keyboard_interrupt") -> def _notify_single_query_session_finalize(cli, *, reason: str = "shutdown") -> None: agent = getattr(cli, "agent", None) session_id = getattr(agent, "session_id", None) or getattr(cli, "session_id", None) - _notify_session_finalize( - session_id=session_id, - platform=getattr(agent, "platform", None) or "cli", - reason=reason, - ) + if session_id in _single_query_finalize_attempted_session_ids: + return + + try: + _notify_session_finalize( + session_id=session_id, + platform=getattr(agent, "platform", None) or "cli", + reason=reason, + ) + finally: + _single_query_finalize_attempted_session_ids.add(session_id) def _finalize_single_query(cli) -> None: diff --git a/tests/cli/test_single_query_session_finalize.py b/tests/cli/test_single_query_session_finalize.py index 37ad56ff7a8..20a75acd37b 100644 --- a/tests/cli/test_single_query_session_finalize.py +++ b/tests/cli/test_single_query_session_finalize.py @@ -5,6 +5,12 @@ import pytest import cli +@pytest.fixture(autouse=True) +def reset_single_query_finalize_state(monkeypatch): + monkeypatch.setattr(cli, "_single_query_finalize_attempted_session_ids", set()) + monkeypatch.setattr(cli, "_cleanup_done", False) + + def test_finalize_single_query_runs_cleanup_without_reemitting_finalize_before_release(monkeypatch): calls = [] fake_cli = SimpleNamespace(_release_active_session=lambda: calls.append(("release", {}))) @@ -70,6 +76,54 @@ def test_finalize_single_query_runs_cleanup_when_finalize_hook_fails(monkeypatch assert calls == ["finalize", "cleanup", "release"] +def test_finalize_single_query_signal_window_does_not_reemit_during_atexit(monkeypatch): + calls = [] + fake_agent = SimpleNamespace(session_id="agent-session", platform="cli") + fake_cli = SimpleNamespace( + agent=fake_agent, + session_id="cli-session", + _release_active_session=lambda: calls.append(("release", {})), + ) + + def invoke_hook(name, **kwargs): + calls.append((name, kwargs)) + + def interrupted_cleanup(**_kwargs): + raise KeyboardInterrupt() + + expected_finalize = ( + "on_session_finalize", + { + "session_id": "agent-session", + "platform": "cli", + "reason": "shutdown", + }, + ) + + original_run_cleanup = cli._run_cleanup + monkeypatch.setattr("hermes_cli.plugins.invoke_hook", invoke_hook) + monkeypatch.setattr(cli, "_run_cleanup", interrupted_cleanup) + + with pytest.raises(KeyboardInterrupt): + cli._finalize_single_query(fake_cli) + + assert calls == [expected_finalize, ("release", {})] + + # Simulate later atexit cleanup after the interrupted one-shot path. The + # active agent may already be unavailable by then. + monkeypatch.setattr(cli, "_run_cleanup", original_run_cleanup) + monkeypatch.setattr(cli, "_active_agent_ref", None) + monkeypatch.setattr(cli, "_reset_terminal_input_modes_on_exit", lambda: None) + monkeypatch.setattr(cli, "_cleanup_all_terminals", lambda: None) + monkeypatch.setattr(cli, "_cleanup_all_browsers", lambda: None) + monkeypatch.setattr("tools.mcp_tool.shutdown_mcp_servers", lambda: None) + monkeypatch.setattr("agent.auxiliary_client.shutdown_cached_clients", lambda: None) + + cli._run_cleanup() + + assert calls == [expected_finalize, ("release", {})] + + def test_notify_single_query_session_finalize_uses_agent_session(monkeypatch): calls = [] fake_agent = SimpleNamespace(session_id="agent-session", platform="cli")