diff --git a/cli.py b/cli.py index ddc469098f5..3b555a288fa 100644 --- a/cli.py +++ b/cli.py @@ -952,7 +952,7 @@ def _prepare_deferred_agent_startup() -> None: exc_info=True, ) -def _run_cleanup(): +def _run_cleanup(*, notify_session_finalize: bool = True): """Run resource cleanup exactly once.""" global _cleanup_done if _cleanup_done: @@ -988,16 +988,12 @@ def _run_cleanup(): pass # Shut down memory provider (on_session_end + shutdown_all) at actual # session boundary — NOT per-turn inside run_conversation(). - try: - from hermes_cli.plugins import invoke_hook as _invoke_hook - _invoke_hook( - "on_session_finalize", + if notify_session_finalize: + _notify_session_finalize( session_id=_active_agent_ref.session_id if _active_agent_ref else None, platform="cli", reason="shutdown", ) - except Exception: - pass try: if _active_agent_ref and hasattr(_active_agent_ref, 'shutdown_memory_provider'): # Forward the agent's own transcript so memory providers' @@ -1015,6 +1011,24 @@ def _run_cleanup(): pass +def _notify_session_finalize( + *, + session_id: str | None, + platform: str = "cli", + reason: str = "shutdown", +) -> None: + try: + from hermes_cli.plugins import invoke_hook as _invoke_hook + _invoke_hook( + "on_session_finalize", + session_id=session_id, + platform=platform, + reason=reason, + ) + except Exception: + pass + + def _emit_interrupted_session_end(cli, *, reason: str = "keyboard_interrupt") -> None: """Best-effort on_session_end hook for interrupted non-interactive runs.""" agent = getattr(cli, "agent", None) @@ -1051,6 +1065,25 @@ def _emit_interrupted_session_end(cli, *, reason: str = "keyboard_interrupt") -> pass +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, + ) + + +def _finalize_single_query(cli) -> None: + """Close one-shot CLI resources before releasing the active session lease.""" + try: + _notify_single_query_session_finalize(cli) + _run_cleanup(notify_session_finalize=False) + finally: + cli._release_active_session() + + def _reset_terminal_input_modes_on_exit() -> None: """Best-effort: disable focus reporting + mouse tracking on TUI exit so they don't leak into the next shell session sharing the tab. @@ -13583,7 +13616,7 @@ def main( cli.chat(query, images=single_query_images or None) cli._print_exit_summary() finally: - cli._release_active_session() + _finalize_single_query(cli) return # Run interactive mode diff --git a/tests/cli/test_single_query_session_finalize.py b/tests/cli/test_single_query_session_finalize.py new file mode 100644 index 00000000000..37ad56ff7a8 --- /dev/null +++ b/tests/cli/test_single_query_session_finalize.py @@ -0,0 +1,216 @@ +from types import SimpleNamespace + +import pytest + +import cli + + +def test_finalize_single_query_runs_cleanup_without_reemitting_finalize_before_release(monkeypatch): + calls = [] + fake_cli = SimpleNamespace(_release_active_session=lambda: calls.append(("release", {}))) + + def cleanup(**kwargs): + calls.append(("cleanup", kwargs)) + + monkeypatch.setattr( + cli, + "_notify_single_query_session_finalize", + lambda _cli: calls.append(("finalize", {})), + ) + monkeypatch.setattr(cli, "_run_cleanup", cleanup) + + cli._finalize_single_query(fake_cli) + + assert calls == [ + ("finalize", {}), + ("cleanup", {"notify_session_finalize": False}), + ("release", {}), + ] + + +def test_finalize_single_query_releases_session_when_cleanup_fails(monkeypatch): + calls = [] + fake_cli = SimpleNamespace(_release_active_session=lambda: calls.append("release")) + + def cleanup(**kwargs): + calls.append("cleanup") + raise RuntimeError("cleanup failed") + + monkeypatch.setattr( + cli, + "_notify_single_query_session_finalize", + lambda _cli: calls.append("finalize"), + ) + monkeypatch.setattr(cli, "_run_cleanup", cleanup) + + with pytest.raises(RuntimeError, match="cleanup failed"): + cli._finalize_single_query(fake_cli) + + assert calls == ["finalize", "cleanup", "release"] + + +def test_finalize_single_query_runs_cleanup_when_finalize_hook_fails(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("finalize") + raise RuntimeError("hook failed") + + monkeypatch.setattr("hermes_cli.plugins.invoke_hook", invoke_hook) + monkeypatch.setattr(cli, "_run_cleanup", lambda **kwargs: calls.append("cleanup")) + + cli._finalize_single_query(fake_cli) + + assert calls == ["finalize", "cleanup", "release"] + + +def test_notify_single_query_session_finalize_uses_agent_session(monkeypatch): + calls = [] + fake_agent = SimpleNamespace(session_id="agent-session", platform="cli") + fake_cli = SimpleNamespace(agent=fake_agent, session_id="cli-session") + + def invoke_hook(name, **kwargs): + calls.append((name, kwargs)) + + monkeypatch.setattr("hermes_cli.plugins.invoke_hook", invoke_hook) + + cli._notify_single_query_session_finalize(fake_cli) + + assert calls == [ + ( + "on_session_finalize", + { + "session_id": "agent-session", + "platform": "cli", + "reason": "shutdown", + }, + ) + ] + + +def test_human_single_query_main_finalizes_after_query(monkeypatch): + calls = [] + + import cli as cli_mod + + class _Console: + def print(self, *_args, **_kwargs): + calls.append("query-label") + + class FakeCLI: + def __init__(self, **_kwargs): + self.console = _Console() + self.session_id = "single-query-session" + self.agent = SimpleNamespace( + session_id="single-query-session", + platform="cli", + ) + + def _claim_active_session(self, surface, *, stderr=False): + calls.append(("claim", surface, stderr)) + return True + + def _show_security_advisories(self): + calls.append("advisories") + + def chat(self, query, images=None): + calls.append(("chat", query, images)) + return "done" + + def _print_exit_summary(self): + calls.append("summary") + + monkeypatch.setattr(cli_mod, "HermesCLI", FakeCLI) + monkeypatch.setattr(cli_mod.atexit, "register", lambda *_args, **_kwargs: None) + monkeypatch.setattr( + cli_mod, + "_finalize_single_query", + lambda fake_cli: calls.append(("finalize", fake_cli.session_id)), + ) + + cli_mod.main(query="hello", quiet=False, toolsets="terminal") + + assert calls == [ + ("claim", "cli", False), + "query-label", + "advisories", + ("chat", "hello", None), + "summary", + ("finalize", "single-query-session"), + ] + + +def test_quiet_single_query_main_finalizes_while_preserving_exit_code(monkeypatch): + calls = [] + + import cli as cli_mod + + def run_conversation(*, user_message, conversation_history): + calls.append(("run", user_message, conversation_history)) + return { + "final_response": "", + "error": "provider failed", + "failed": True, + } + + class FakeCLI: + def __init__(self, **_kwargs): + self.provider = "test-provider" + self.model = "test-model" + self.session_id = "quiet-session" + self.conversation_history = [] + self._active_agent_route_signature = "same-route" + self.agent = SimpleNamespace( + session_id="quiet-session", + platform="cli", + quiet_mode=False, + suppress_status_output=False, + stream_delta_callback=object(), + tool_gen_callback=object(), + run_conversation=run_conversation, + ) + + def _claim_active_session(self, surface, *, stderr=False): + calls.append(("claim", surface, stderr)) + return True + + def _ensure_runtime_credentials(self): + calls.append("credentials") + return True + + def _resolve_turn_agent_config(self, effective_query): + calls.append(("resolve", effective_query)) + return { + "signature": "same-route", + "model": None, + "runtime": None, + "request_overrides": None, + } + + def _init_agent(self, **kwargs): + calls.append(("init", kwargs)) + return True + + monkeypatch.delenv("HERMES_KANBAN_TASK", raising=False) + monkeypatch.delenv("HERMES_KANBAN_GOAL_MODE", raising=False) + monkeypatch.setattr(cli_mod, "HermesCLI", FakeCLI) + monkeypatch.setattr(cli_mod.atexit, "register", lambda *_args, **_kwargs: None) + monkeypatch.setattr( + cli_mod, + "_finalize_single_query", + lambda fake_cli: calls.append(("finalize", fake_cli.session_id)), + ) + + with pytest.raises(SystemExit) as exc_info: + cli_mod.main(query="hello", quiet=True, toolsets="terminal") + + assert exc_info.value.code == 1 + assert ("claim", "cli", True) in calls + assert ("run", "hello", []) in calls + assert calls[-1] == ("finalize", "quiet-session")