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>
This commit is contained in:
mnajafian-nv 2026-06-09 21:52:13 -07:00 committed by GitHub
parent 96af61b6ef
commit d03cdd63eb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 257 additions and 8 deletions

49
cli.py
View file

@ -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