fix(gateway): drain in-flight api_server runs on shutdown

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

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

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

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