fix(cli): prevent duplicate one-shot finalize on interrupted cleanup (#43320)

Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
This commit is contained in:
mnajafian-nv 2026-06-09 22:41:04 -07:00 committed by GitHub
parent 1967c590ed
commit f8fd30942c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 84 additions and 10 deletions

40
cli.py
View file

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