fix(cli): run one-shot query cleanup before lease release (#43036)

* fix(cli): run one-shot query cleanup before lease release

Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>

* test(cli): cover quiet one-shot cleanup finalization

Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>

---------

Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>

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

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")