diff --git a/gateway/platforms/api_server.py b/gateway/platforms/api_server.py index 62dff2878fa59..26bcd6c4aee7e 100644 --- a/gateway/platforms/api_server.py +++ b/gateway/platforms/api_server.py @@ -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() diff --git a/gateway/run.py b/gateway/run.py index 67ccf2620d2a1..72a43b99c9678 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -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,