From 021ee345464662bc95381502b6562df19df87258 Mon Sep 17 00:00:00 2001 From: Bartok9 Date: Mon, 13 Jul 2026 03:17:40 -0400 Subject: [PATCH] fix(gateway): drain in-flight api_server runs on shutdown MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes #63529 Root cause: GatewayRunner._drain_active_agents only waited on _running_agents + cron in-flight counts. Desktop/API sessions are tracked solely inside APIServerAdapter (_inflight_agent_runs + _active_run_agents), so stop/restart logged active_at_start=0 and systemd SIGKILL'd mid-tool work. Fix: APIServerAdapter.active_agent_work_count() plus GatewayRunner._active_api_run_count() folded into the drain wait, status updates, timeout result, and shutdown logs — same pattern as _active_cron_job_count for #60432. Verification: pytest tests/gateway/test_api_server_active_work_drain.py tests/gateway/test_cron_active_work_drain.py -q → 19 passed --- gateway/platforms/api_server.py | 13 ++ gateway/run.py | 83 +++++++++- .../test_api_server_active_work_drain.py | 156 ++++++++++++++++++ 3 files changed, 245 insertions(+), 7 deletions(-) create mode 100644 tests/gateway/test_api_server_active_work_drain.py diff --git a/gateway/platforms/api_server.py b/gateway/platforms/api_server.py index b96adf9e448d..a3e4ec4088cd 100644 --- a/gateway/platforms/api_server.py +++ b/gateway/platforms/api_server.py @@ -913,6 +913,19 @@ class APIServerAdapter(BasePlatformAdapter): # _active_run_tasks). self._inflight_agent_runs: int = 0 + def active_agent_work_count(self) -> int: + """Count of in-flight agent work served by this API server adapter. + + Covers the non-streaming chat/responses paths (``_inflight_agent_runs``) + plus live ``/v1/runs`` agents (``_active_run_agents``). Used by the + gateway shutdown drain so desk/API sessions get the same drain window + as messaging sessions and cron jobs (#63529). + """ + try: + return int(self._inflight_agent_runs) + len(self._active_run_agents) + except Exception: + return 0 + def _readiness_work_counts(self) -> tuple[int, int, int]: """Return bounded work counts from each subsystem's public state.""" active_api_runs = sum( diff --git a/gateway/run.py b/gateway/run.py index 94b91d885e10..fd3705b1f3a0 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -4142,6 +4142,56 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew except Exception: return 0 + def _active_api_run_count(self) -> int: + """Count of in-flight api_server agent runs across all adapters. + + Chat/responses/desktop API sessions live entirely inside + ``APIServerAdapter`` (``_inflight_agent_runs`` + ``_active_run_agents``) + and are invisible to ``self._running_agents``. Without this fold-in the + shutdown drain can report ``active_at_start=0`` and exit while desk/API + tools are still mid-flight (#63529) — the same structural failure cron + had before ``_active_cron_job_count`` (#60432). Best-effort: skips + adapters that don't expose the helper. + """ + total = 0 + try: + from gateway.config import Platform + except Exception: + return 0 + + def _count_from(adapters) -> int: + n = 0 + if not adapters: + return 0 + try: + values = adapters.values() if hasattr(adapters, "values") else adapters + except Exception: + return 0 + for adapter in values: + try: + if getattr(adapter, "platform", None) != Platform.API_SERVER: + continue + helper = getattr(adapter, "active_agent_work_count", None) + if callable(helper): + n += int(helper()) + else: + inflight = int(getattr(adapter, "_inflight_agent_runs", 0) or 0) + active = getattr(adapter, "_active_run_agents", None) or {} + n += inflight + len(active) + except Exception: + continue + return n + + try: + total += _count_from(getattr(self, "adapters", None)) + except Exception: + pass + try: + total += _count_from(getattr(self, "_profile_adapters", None)) + except Exception: + pass + return total + # ── scale-to-zero idle detection / dormant-quiesce (Phase 0) ────────────── # The gateway-side BEHAVIOUR that consumes the relay scale-to-zero primitives # (gateway-gateway Phase 5). Pure logic lives in gateway/scale_to_zero.py; the @@ -5696,22 +5746,26 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew snapshot = self._snapshot_running_agents() last_active_count = self._running_agent_count() last_cron_count = self._active_cron_job_count() + last_api_count = self._active_api_run_count() last_status_at = 0.0 def _maybe_update_status(force: bool = False) -> None: - nonlocal last_active_count, last_cron_count, last_status_at + nonlocal last_active_count, last_cron_count, last_api_count, last_status_at now = asyncio.get_running_loop().time() active_count = self._running_agent_count() cron_count = self._active_cron_job_count() + api_count = self._active_api_run_count() if ( force or active_count != last_active_count or cron_count != last_cron_count + or api_count != last_api_count or (now - last_status_at) >= 1.0 ): self._update_runtime_status("draining") last_active_count = active_count last_cron_count = cron_count + last_api_count = api_count last_status_at = now # Cron jobs run on the scheduler's own thread pool, outside @@ -5719,7 +5773,8 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew # same wait/timeout this method already applies to chat sessions, # or a cron job's tool work gets killed with zero warning the # instant it's the only active thing running (#60432). - if not self._running_agents and last_cron_count == 0: + # API-server / desk sessions have the same structural gap (#63529). + if not self._running_agents and last_cron_count == 0 and last_api_count == 0: _maybe_update_status(force=True) return snapshot, False @@ -5729,12 +5784,20 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew deadline = asyncio.get_running_loop().time() + timeout while ( - (self._running_agents or self._active_cron_job_count()) + ( + self._running_agents + or self._active_cron_job_count() + or self._active_api_run_count() + ) and asyncio.get_running_loop().time() < deadline ): _maybe_update_status() await asyncio.sleep(0.1) - timed_out = bool(self._running_agents) or bool(self._active_cron_job_count()) + timed_out = ( + bool(self._running_agents) + or bool(self._active_cron_job_count()) + or bool(self._active_api_run_count()) + ) _maybe_update_status(force=True) return snapshot, timed_out @@ -8187,12 +8250,14 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew logger.debug("pre-drain mark_resume_pending failed for %s: %s", _sk, _e) _cron_at_start = self._active_cron_job_count() + _api_at_start = self._active_api_run_count() _drain_started_at = time.monotonic() active_agents, timed_out = await self._drain_active_agents(timeout) logger.info( "Shutdown phase: drain done at +%.2fs (drain took %.2fs, " "timed_out=%s, active_at_start=%d, active_now=%d, " - "cron_at_start=%d, cron_now=%d)", + "cron_at_start=%d, cron_now=%d, " + "api_at_start=%d, api_now=%d)", _phase_elapsed(), time.monotonic() - _drain_started_at, timed_out, @@ -8200,6 +8265,8 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew self._running_agent_count(), _cron_at_start, self._active_cron_job_count(), + _api_at_start, + self._active_api_run_count(), ) if not timed_out: @@ -8218,11 +8285,13 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew if timed_out: logger.warning( - "Gateway drain timed out after %.1fs with %d active agent(s) " - "and %d in-flight cron job(s); interrupting remaining work.", + "Gateway drain timed out after %.1fs with %d active agent(s), " + "%d in-flight cron job(s), and %d api_server run(s); " + "interrupting remaining work.", timeout, self._running_agent_count(), self._active_cron_job_count(), + self._active_api_run_count(), ) # Mark forcibly-interrupted sessions as resume_pending BEFORE # interrupting the agents. This preserves each session's diff --git a/tests/gateway/test_api_server_active_work_drain.py b/tests/gateway/test_api_server_active_work_drain.py new file mode 100644 index 000000000000..89444d18a407 --- /dev/null +++ b/tests/gateway/test_api_server_active_work_drain.py @@ -0,0 +1,156 @@ +"""Tests for #63529: the gateway shutdown drain was structurally blind to +in-flight api_server (desk/API) agent runs. + +API-server runs are tracked only inside ``APIServerAdapter`` +(``_inflight_agent_runs`` + ``_active_run_agents``) and never enter +``GatewayRunner._running_agents``. Without folding them into the drain, +stop/restart reported ``active_at_start=0`` and let systemd SIGKILL mid-tool. +Mirrors tests/gateway/test_cron_active_work_drain.py for cron (#60432). +""" + +import asyncio +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +from tests.gateway.restart_test_helpers import make_restart_runner + + +def _make_api_adapter(*, inflight: int = 0, active_ids=None): + from gateway.config import Platform + + active = {rid: MagicMock() for rid in (active_ids or [])} + adapter = SimpleNamespace( + platform=Platform.API_SERVER, + _inflight_agent_runs=inflight, + _active_run_agents=active, + ) + + def active_agent_work_count() -> int: + return int(adapter._inflight_agent_runs) + len(adapter._active_run_agents) + + adapter.active_agent_work_count = active_agent_work_count + return adapter + + +class TestActiveApiRunCount: + def test_zero_when_no_api_adapters(self): + runner, _adapter = make_restart_runner() + runner.adapters = {} + runner._profile_adapters = {} + assert runner._active_api_run_count() == 0 + + def test_sums_inflight_and_active_run_agents(self): + runner, _adapter = make_restart_runner() + adapter = _make_api_adapter(inflight=2, active_ids=["r1"]) + runner.adapters = {"api": adapter} + runner._profile_adapters = {} + assert runner._active_api_run_count() == 3 + + def test_includes_profile_adapters(self): + runner, _adapter = make_restart_runner() + runner.adapters = {"api": _make_api_adapter(inflight=1)} + runner._profile_adapters = {"p1": _make_api_adapter(active_ids=["a", "b"])} + assert runner._active_api_run_count() == 3 + + def test_ignores_non_api_platforms(self): + from gateway.config import Platform + + runner, _adapter = make_restart_runner() + other = SimpleNamespace( + platform=Platform.DISCORD, + _inflight_agent_runs=99, + _active_run_agents={"x": MagicMock()}, + active_agent_work_count=lambda: 99, + ) + runner.adapters = {"discord": other} + runner._profile_adapters = {} + assert runner._active_api_run_count() == 0 + + def test_never_raises_on_broken_adapters(self): + runner, _adapter = make_restart_runner() + + class Bad: + platform = property(lambda self: (_ for _ in ()).throw(RuntimeError("boom"))) + + runner.adapters = {"bad": Bad()} + runner._profile_adapters = None + assert runner._active_api_run_count() == 0 + + +class TestAPIServerAdapterWorkCount: + def test_active_agent_work_count_on_real_class_method(self): + from gateway.platforms.api_server import APIServerAdapter + + # Instantiate unbound helpers via object.__new__ to avoid full connect. + adapter = object.__new__(APIServerAdapter) + adapter._inflight_agent_runs = 2 + adapter._active_run_agents = {"r1": object(), "r2": object()} + assert adapter.active_agent_work_count() == 4 + + +class TestDrainWaitsForApiWork: + @pytest.mark.asyncio + async def test_drain_returns_immediately_when_nothing_active(self): + runner, _adapter = make_restart_runner() + runner.adapters = {} + runner._profile_adapters = {} + + _snapshot, timed_out = await runner._drain_active_agents(5.0) + + assert timed_out is False + + @pytest.mark.asyncio + async def test_drain_waits_for_in_flight_api_run(self): + runner, _adapter = make_restart_runner() + api = _make_api_adapter(inflight=1) + runner.adapters = {"api": api} + runner._profile_adapters = {} + + async def finish_run(): + await asyncio.sleep(0.12) + api._inflight_agent_runs = 0 + + task = asyncio.create_task(finish_run()) + _snapshot, timed_out = await runner._drain_active_agents(2.0) + await task + + assert timed_out is False, ( + "drain must wait for api_server work, not report active_at_start=0" + ) + + @pytest.mark.asyncio + async def test_drain_times_out_if_api_run_outlives_the_window(self): + runner, _adapter = make_restart_runner() + runner.adapters = {"api": _make_api_adapter(inflight=1)} + runner._profile_adapters = {} + + _snapshot, timed_out = await runner._drain_active_agents(0.1) + + assert timed_out is True + + @pytest.mark.asyncio + async def test_drain_still_waits_for_chat_and_cron(self): + import cron.scheduler as sched + + runner, _adapter = make_restart_runner() + runner._running_agents = {"session-1": MagicMock()} + sched._running_job_ids.add("job-1") + runner.adapters = {"api": _make_api_adapter(inflight=1)} + runner._profile_adapters = {} + + async def finish_all(): + await asyncio.sleep(0.12) + runner._running_agents.clear() + sched._running_job_ids.discard("job-1") + runner.adapters["api"]._inflight_agent_runs = 0 + + task = asyncio.create_task(finish_all()) + try: + _snapshot, timed_out = await runner._drain_active_agents(2.0) + finally: + await task + sched._running_job_ids.discard("job-1") + + assert timed_out is False