diff --git a/gateway/platforms/api_server.py b/gateway/platforms/api_server.py index 24cc8f1bca55d..d24cffc05fc6b 100644 --- a/gateway/platforms/api_server.py +++ b/gateway/platforms/api_server.py @@ -1474,6 +1474,25 @@ class APIServerAdapter(BasePlatformAdapter): except Exception: return 0 + def interrupt_active_runs(self, reason: str) -> int: + """Interrupt active API-run agents during gateway shutdown. + + The gateway drain accounts for API-server work through + ``active_agent_work_count()``, but those agents are owned by this + adapter rather than ``GatewayRunner._running_agents``. Expose the same + cooperative interrupt used by ``POST /v1/runs/{run_id}/stop`` so a + shutdown timeout can stop long-running API work before process teardown. + """ + interrupted = 0 + for run_id, agent in list(self._active_run_agents.items()): + try: + agent.interrupt(reason) + interrupted += 1 + logger.debug("[api_server] interrupted active run %s during shutdown", run_id) + except Exception as exc: + logger.debug("[api_server] failed interrupting active run %s: %s", run_id, exc) + return interrupted + @staticmethod def _gateway_is_draining() -> bool: """Whether the owning gateway currently refuses new agent turns.""" diff --git a/gateway/run.py b/gateway/run.py index 40be1f3af9658..2d2934e48d4de 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -7420,6 +7420,16 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew except Exception: return 0 + def _interrupt_api_server_runs(self, reason: str) -> int: + """Interrupt API-server agents that are not in ``_running_agents``.""" + try: + adapter = getattr(self, "adapters", {}).get(Platform.API_SERVER) + helper = getattr(adapter, "interrupt_active_runs", None) + return max(0, int(helper(reason))) if callable(helper) else 0 + except Exception as exc: + logger.debug("Failed interrupting api_server runs during shutdown: %s", exc) + return 0 + # ── 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 @@ -9249,6 +9259,9 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew logger.debug("Interrupted running agent for session %s during shutdown", session_key) except Exception as e: logger.debug("Failed interrupting agent during shutdown: %s", e) + interrupted_api = self._interrupt_api_server_runs(reason) + if interrupted_api: + logger.debug("Interrupted %d api_server run(s) during shutdown", interrupted_api) async def _notify_active_sessions_of_shutdown(self) -> None: """Send shutdown/restart notifications to active chats and home channels. @@ -12916,7 +12929,10 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew _INTERRUPT_REASON_GATEWAY_RESTART if self._restart_requested else _INTERRUPT_REASON_GATEWAY_SHUTDOWN ) interrupt_deadline = asyncio.get_running_loop().time() + 5.0 - while self._running_agents and asyncio.get_running_loop().time() < interrupt_deadline: + while ( + self._running_agents + or self._active_api_run_count() + ) and asyncio.get_running_loop().time() < interrupt_deadline: self._update_runtime_status("draining") await asyncio.sleep(0.1) diff --git a/tests/gateway/test_api_server_active_work_drain.py b/tests/gateway/test_api_server_active_work_drain.py index 2dedee6c72433..45a4a467980ba 100644 --- a/tests/gateway/test_api_server_active_work_drain.py +++ b/tests/gateway/test_api_server_active_work_drain.py @@ -93,6 +93,23 @@ class TestAPIServerAdapterWorkCount: assert adapter.active_agent_work_count() == 3 + def test_does_not_double_count_started_run_agent(self): + adapter = APIServerAdapter(PlatformConfig(enabled=True)) + adapter._inflight_agent_runs = 0 + adapter._active_run_tasks = {"run-1": _RunTask()} + adapter._active_run_agents = {"run-1": object()} + + assert adapter.active_agent_work_count() == 1 + + def test_interrupt_active_runs_interrupts_adapter_owned_agents(self): + adapter = APIServerAdapter(PlatformConfig(enabled=True)) + agent = MagicMock() + adapter._active_run_agents = {"run-1": agent} + + assert adapter.interrupt_active_runs("gateway shutdown") == 1 + + agent.interrupt.assert_called_once_with("gateway shutdown") + class TestDrainWaitsForApiWork: @@ -141,6 +158,50 @@ class TestDrainWaitsForApiWork: assert timed_out is False + @pytest.mark.asyncio + async def test_drain_times_out_if_api_run_outlives_the_window(self): + runner, _adapter = make_restart_runner() + runner.adapters = {Platform.API_SERVER: _make_api_adapter(queued_ids=["run-1"])} + + _snapshot, timed_out = await runner._drain_active_agents(0.1) + + assert timed_out is True + + def test_shutdown_interrupt_reaches_api_server_runs(self): + runner, _adapter = make_restart_runner() + api = APIServerAdapter(PlatformConfig(enabled=True)) + agent = MagicMock() + api._active_run_agents = {"run-1": agent} + runner.adapters = {Platform.API_SERVER: api} + + runner._interrupt_running_agents("gateway shutdown") + + agent.interrupt.assert_called_once_with("gateway shutdown") + + @pytest.mark.asyncio + async def test_drain_still_waits_for_chat_cron_and_api_work(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 = {Platform.API_SERVER: _make_api_adapter(queued_ids=["run-1"])} + + async def finish_all(): + await asyncio.sleep(0.12) + runner._running_agents.clear() + sched._running_job_ids.discard("job-1") + runner.adapters[Platform.API_SERVER]._active_run_tasks.clear() + + 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 + class TestDrainAdmission: @pytest.mark.asyncio