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:
joaomarcos 2026-08-01 11:43:53 -03:00
parent 7c543bf409
commit dbbb10d396
2 changed files with 119 additions and 10 deletions

View File

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

View File

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