fix(gateway): close cross-turn reap race and cover API-server disconnect
Addresses the hermes-sweeper review on #76188: 1. task_id is session-scoped (task_id == session_id), not turn-scoped, and the reap runs on a detached thread. A replacement turn could claim the same session and spawn a legitimate process before the previous turn's reaper thread actually enumerates its targets, killing that new process by mistake. Fixed by gating the reap on the existing run_generation mechanism (_is_session_run_current) instead of inventing a new ownership token: the timeout path captures its own run_generation at turn start, the interrupt path captures the generation immediately after invalidating it. If a newer turn has since claimed the session, the reap is skipped — that newer turn owns its own baseline, so nothing is left permanently unreaped. 2. gateway/platforms/api_server.py's SSE handlers for chat-completions and the /api/sessions responses endpoint run their own agent lifecycle via _run_agent() and never passed through TurnRunner, so client-disconnect abandonment there had no baseline and no reap — contradicting the PR's stated disconnect coverage. Both disconnect handlers now snapshot/reap through the same tools.process_registry primitives, via a small _reap_disconnected_agent_processes() helper shared by both call sites.
This commit is contained in:
parent
7c543bf409
commit
dbbb10d396
|
|
@ -617,6 +617,32 @@ def _multimodal_validation_error(exc: ValueError, *, param: str) -> "web.Respons
|
|||
)
|
||||
|
||||
|
||||
def _reap_disconnected_agent_processes(agent: Any) -> None:
|
||||
"""Reap background processes an SSE-disconnected agent's turn created.
|
||||
|
||||
Mirrors the gateway-turn cleanup in ``gateway/run.py`` (#76115) for this
|
||||
API-server surface, which runs its own agent lifecycle via ``_run_agent``
|
||||
and never passes through ``TurnRunner`` — so it needs its own trigger for
|
||||
the same baseline-diff reap. Fire-and-forget on a daemon thread so the
|
||||
SSE handler's own cleanup isn't blocked on process-tree teardown.
|
||||
"""
|
||||
process_task_id = getattr(agent, "_gateway_turn_process_task_id", "")
|
||||
process_baseline = getattr(agent, "_gateway_turn_process_baseline", None)
|
||||
if not process_task_id or process_baseline is None:
|
||||
return
|
||||
import threading
|
||||
|
||||
from gateway.run import _reap_gateway_turn_processes
|
||||
|
||||
threading.Thread(
|
||||
target=_reap_gateway_turn_processes,
|
||||
args=(process_task_id, process_baseline),
|
||||
kwargs={"source": "api_server_sse_disconnect"},
|
||||
name=f"api-turn-reaper-{process_task_id[:12]}",
|
||||
daemon=True,
|
||||
).start()
|
||||
|
||||
|
||||
def _session_chat_user_message(body: Dict[str, Any], *, param: str = "message") -> tuple[Any, Optional["web.Response"]]:
|
||||
"""Parse and normalize session chat ``message`` / ``input`` like chat completions."""
|
||||
user_message = body.get("message") or body.get("input")
|
||||
|
|
@ -4234,6 +4260,7 @@ class APIServerAdapter(BasePlatformAdapter):
|
|||
agent.interrupt("SSE client disconnected")
|
||||
except Exception:
|
||||
pass
|
||||
_reap_disconnected_agent_processes(agent)
|
||||
if not agent_task.done():
|
||||
agent_task.cancel()
|
||||
try:
|
||||
|
|
@ -4813,6 +4840,7 @@ class APIServerAdapter(BasePlatformAdapter):
|
|||
agent.interrupt("SSE client disconnected")
|
||||
except Exception:
|
||||
pass
|
||||
_reap_disconnected_agent_processes(agent)
|
||||
if not agent_task.done():
|
||||
agent_task.cancel()
|
||||
try:
|
||||
|
|
@ -5815,6 +5843,7 @@ class APIServerAdapter(BasePlatformAdapter):
|
|||
session_key=gateway_session_key or session_id or "",
|
||||
session_id=session_id or "",
|
||||
)
|
||||
agent = None
|
||||
try:
|
||||
agent = self._create_agent(
|
||||
ephemeral_system_prompt=ephemeral_system_prompt,
|
||||
|
|
@ -5834,6 +5863,16 @@ class APIServerAdapter(BasePlatformAdapter):
|
|||
if agent_ref is not None:
|
||||
agent_ref[0] = agent
|
||||
effective_task_id = session_id or str(uuid.uuid4())
|
||||
# Baseline for selective background-process reaping on
|
||||
# SSE client disconnect — mirrors gateway/run.py's
|
||||
# gateway-turn cleanup (#76115); this API-server surface
|
||||
# runs its own agent lifecycle and doesn't go through
|
||||
# TurnRunner, so it needs its own baseline.
|
||||
from tools.process_registry import process_registry
|
||||
agent._gateway_turn_process_task_id = effective_task_id
|
||||
agent._gateway_turn_process_baseline = (
|
||||
process_registry.snapshot_running_ids(effective_task_id)
|
||||
)
|
||||
result = agent.run_conversation(
|
||||
user_message=user_message,
|
||||
conversation_history=conversation_history,
|
||||
|
|
@ -5956,6 +5995,14 @@ class APIServerAdapter(BasePlatformAdapter):
|
|||
{"input_tokens": 0, "output_tokens": 0, "total_tokens": 0},
|
||||
)
|
||||
finally:
|
||||
# Turn finished (success, auth failure, or crash) — clear
|
||||
# ownership markers so a disconnect landing after this
|
||||
# point can't reap background work this turn left
|
||||
# running on purpose. Mirrors the same race-window guard
|
||||
# in gateway/run.py's _run_sync_with_timeout_lifecycle.
|
||||
if agent is not None:
|
||||
agent._gateway_turn_process_task_id = ""
|
||||
agent._gateway_turn_process_baseline = frozenset()
|
||||
clear_session_vars(tokens)
|
||||
|
||||
self._activate_admitted_request()
|
||||
|
|
|
|||
|
|
@ -2694,8 +2694,37 @@ def _reap_gateway_turn_processes(
|
|||
process_baseline,
|
||||
*,
|
||||
source: str,
|
||||
is_still_current: Optional[Callable[[], bool]] = None,
|
||||
) -> int:
|
||||
"""Reap only background processes created by one abandoned turn."""
|
||||
"""Reap only background processes created by one abandoned turn.
|
||||
|
||||
``task_id`` is session-scoped (task_id == session_id), not turn-scoped,
|
||||
so a *replacement* turn on the same session can start and spawn its own
|
||||
legitimate process while this reap is still in flight. ``is_still_current``
|
||||
— a closure over the run_generation captured when the reaping turn began
|
||||
or was interrupted — lets the caller detect that a newer turn has since
|
||||
claimed the session and bail out instead of killing that newer turn's
|
||||
process. The newer turn snapshots its own baseline independently, so
|
||||
skipping here does not leave anything permanently unreaped.
|
||||
"""
|
||||
if is_still_current is not None:
|
||||
try:
|
||||
if not is_still_current():
|
||||
logger.debug(
|
||||
"Skipping reap for turn %s (%s): a newer turn already "
|
||||
"claimed this session; it owns its own baseline.",
|
||||
task_id,
|
||||
source,
|
||||
)
|
||||
return 0
|
||||
except Exception:
|
||||
logger.debug(
|
||||
"is_still_current check failed for turn %s (%s); reaping anyway",
|
||||
task_id,
|
||||
source,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
from tools.process_registry import process_registry
|
||||
|
||||
try:
|
||||
|
|
@ -2734,6 +2763,7 @@ def _abandon_timed_out_gateway_turn(
|
|||
worker_done: threading.Event,
|
||||
timeout_fired: threading.Event,
|
||||
cleanup_lock: threading.Lock,
|
||||
is_still_current: Optional[Callable[[], bool]] = None,
|
||||
) -> bool:
|
||||
"""Interrupt one timed-out turn and reap only processes it created."""
|
||||
with cleanup_lock:
|
||||
|
|
@ -2753,6 +2783,7 @@ def _abandon_timed_out_gateway_turn(
|
|||
task_id,
|
||||
process_baseline,
|
||||
source="gateway_turn_timeout",
|
||||
is_still_current=is_still_current,
|
||||
)
|
||||
except Exception:
|
||||
logger.warning(
|
||||
|
|
@ -2773,6 +2804,7 @@ def _watch_gateway_turn_inactivity(
|
|||
timeout_fired: threading.Event,
|
||||
cleanup_lock: threading.Lock,
|
||||
poll_interval: float = 5.0,
|
||||
is_still_current: Optional[Callable[[], bool]] = None,
|
||||
) -> None:
|
||||
"""Thread watchdog that remains runnable when gateway asyncio is starved."""
|
||||
while not worker_done.wait(max(0.01, poll_interval)):
|
||||
|
|
@ -2794,6 +2826,7 @@ def _watch_gateway_turn_inactivity(
|
|||
worker_done=worker_done,
|
||||
timeout_fired=timeout_fired,
|
||||
cleanup_lock=cleanup_lock,
|
||||
is_still_current=is_still_current,
|
||||
)
|
||||
return
|
||||
|
||||
|
|
@ -22190,6 +22223,8 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
|
|||
return
|
||||
_iac_state = self._peek_session_state(session_key)
|
||||
running_agent = _iac_state.turn.agent if _iac_state else None
|
||||
_process_task_id = ""
|
||||
_process_baseline = None
|
||||
if running_agent and running_agent is not _AGENT_PENDING_SENTINEL:
|
||||
running_agent.interrupt(interrupt_reason)
|
||||
_process_task_id = getattr(
|
||||
|
|
@ -22198,15 +22233,29 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
|
|||
_process_baseline = getattr(
|
||||
running_agent, "_gateway_turn_process_baseline", None
|
||||
)
|
||||
if _process_task_id and _process_baseline is not None:
|
||||
threading.Thread(
|
||||
target=_reap_gateway_turn_processes,
|
||||
args=(_process_task_id, _process_baseline),
|
||||
kwargs={"source": "gateway_turn_interrupt"},
|
||||
name=f"gateway-turn-reaper-{_process_task_id[:12]}",
|
||||
daemon=True,
|
||||
).start()
|
||||
self._invalidate_session_run_generation(session_key, reason=invalidation_reason)
|
||||
# Bump the generation *before* scheduling the reap thread and capture
|
||||
# the post-bump value: task_id is session-scoped (task_id ==
|
||||
# session_id), so if a replacement turn claims this session and
|
||||
# spawns its own process before the reap thread actually runs, that
|
||||
# claim bumps the generation again. The closure below then sees a
|
||||
# stale generation and skips — the replacement turn's own baseline
|
||||
# covers its own cleanup, so nothing is left permanently unreaped.
|
||||
_generation_at_interrupt = self._invalidate_session_run_generation(
|
||||
session_key, reason=invalidation_reason
|
||||
)
|
||||
if _process_task_id and _process_baseline is not None:
|
||||
threading.Thread(
|
||||
target=_reap_gateway_turn_processes,
|
||||
args=(_process_task_id, _process_baseline),
|
||||
kwargs={
|
||||
"source": "gateway_turn_interrupt",
|
||||
"is_still_current": lambda: self._is_session_run_current(
|
||||
session_key, _generation_at_interrupt
|
||||
),
|
||||
},
|
||||
name=f"gateway-turn-reaper-{_process_task_id[:12]}",
|
||||
daemon=True,
|
||||
).start()
|
||||
adapter = self._adapter_for_source(source)
|
||||
interrupt_session_activity = getattr(
|
||||
type(adapter), "interrupt_session_activity", None
|
||||
|
|
@ -24169,6 +24218,17 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
|
|||
_turn_worker_done = threading.Event()
|
||||
_turn_timeout_fired = threading.Event()
|
||||
_turn_cleanup_lock = threading.Lock()
|
||||
# task_id above is session-scoped, not turn-scoped (#76115
|
||||
# review): gate the eventual reap on this exact claim still
|
||||
# being current, so a replacement turn that starts on the same
|
||||
# session before the watchdog fires doesn't get its own fresh
|
||||
# process killed by this turn's stale baseline.
|
||||
_turn_run_generation = run_generation
|
||||
_turn_is_current = (
|
||||
(lambda: self._is_session_run_current(session_key, _turn_run_generation))
|
||||
if _turn_run_generation is not None
|
||||
else (lambda: True)
|
||||
)
|
||||
|
||||
def _run_sync_with_timeout_lifecycle():
|
||||
try:
|
||||
|
|
@ -24203,6 +24263,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
|
|||
"timeout_fired": _turn_timeout_fired,
|
||||
"cleanup_lock": _turn_cleanup_lock,
|
||||
"poll_interval": 5.0,
|
||||
"is_still_current": _turn_is_current,
|
||||
},
|
||||
name=f"gateway-turn-watchdog-{_turn_task_id[:12]}",
|
||||
daemon=True,
|
||||
|
|
@ -24315,6 +24376,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
|
|||
"worker_done": _turn_worker_done,
|
||||
"timeout_fired": _turn_timeout_fired,
|
||||
"cleanup_lock": _turn_cleanup_lock,
|
||||
"is_still_current": _turn_is_current,
|
||||
},
|
||||
name=f"gateway-turn-reaper-{_turn_task_id[:12]}",
|
||||
daemon=True,
|
||||
|
|
|
|||
Loading…
Reference in New Issue