From 104ffeae23a7dc08265f1b76ae7d5d68c3eaf2e4 Mon Sep 17 00:00:00 2001 From: kshitijk4poor <82637225+kshitijk4poor@users.noreply.github.com> Date: Mon, 13 Jul 2026 16:27:18 +0530 Subject: [PATCH] fix(gateway): complete API-server shutdown drain --- gateway/platforms/api_server.py | 59 +++++- gateway/run.py | 50 +---- .../test_api_server_active_work_drain.py | 194 ++++++++++++------ tests/gateway/test_shutdown_cache_cleanup.py | 5 + 4 files changed, 198 insertions(+), 110 deletions(-) diff --git a/gateway/platforms/api_server.py b/gateway/platforms/api_server.py index a3e4ec4088cd..cff62197fa8f 100644 --- a/gateway/platforms/api_server.py +++ b/gateway/platforms/api_server.py @@ -914,18 +914,50 @@ class APIServerAdapter(BasePlatformAdapter): 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. + """Return all live agent work owned by this API 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). + ``/v1/runs`` registers an asyncio task before it constructs and stores + its agent, so ``_active_run_agents`` has a real queued-before-agent gap. + Reuse the task-based accounting used by the concurrent-run limit: it + covers that gap and excludes completed tasks retained until cleanup. """ try: - return int(self._inflight_agent_runs) + len(self._active_run_agents) + return int(self._inflight_agent_runs) + sum( + not task.done() for task in self._active_run_tasks.values() + ) except Exception: return 0 + @staticmethod + def _gateway_is_draining() -> bool: + """Whether the owning gateway currently refuses new agent turns.""" + try: + from gateway.run import _gateway_runner_ref + + runner = _gateway_runner_ref() + return bool( + runner + and ( + getattr(runner, "_draining", False) + or getattr(runner, "_external_drain_active", False) + ) + ) + except Exception: + return False + + def _draining_response(self) -> Optional["web.Response"]: + """Return a retryable response while the gateway drains existing work.""" + if not self._gateway_is_draining(): + return None + return web.json_response( + _openai_error( + "Gateway is draining existing work; retry shortly.", + code="gateway_draining", + ), + status=503, + headers={"Retry-After": "1"}, + ) + def _readiness_work_counts(self) -> tuple[int, int, int]: """Return bounded work counts from each subsystem's public state.""" active_api_runs = sum( @@ -1937,6 +1969,9 @@ class APIServerAdapter(BasePlatformAdapter): auth_err = self._check_auth(request) if auth_err: return auth_err + draining = self._draining_response() + if draining is not None: + return draining gateway_session_key, key_err = self._parse_session_key_header(request) if key_err is not None: return key_err @@ -1981,6 +2016,9 @@ class APIServerAdapter(BasePlatformAdapter): auth_err = self._check_auth(request) if auth_err: return auth_err + draining = self._draining_response() + if draining is not None: + return draining gateway_session_key, key_err = self._parse_session_key_header(request) if key_err is not None: return key_err @@ -2122,6 +2160,9 @@ class APIServerAdapter(BasePlatformAdapter): auth_err = self._check_auth(request) if auth_err: return auth_err + draining = self._draining_response() + if draining is not None: + return draining # Bound total in-flight agent runs (configurable; #7483). limited = self._concurrency_limited_response() @@ -3258,6 +3299,9 @@ class APIServerAdapter(BasePlatformAdapter): auth_err = self._check_auth(request) if auth_err: return auth_err + draining = self._draining_response() + if draining is not None: + return draining # Bound total in-flight agent runs (configurable; #7483). limited = self._concurrency_limited_response() @@ -4233,6 +4277,9 @@ class APIServerAdapter(BasePlatformAdapter): auth_err = self._check_auth(request) if auth_err: return auth_err + draining = self._draining_response() + if draining is not None: + return draining # Long-term memory scope header (see chat_completions for details). gateway_session_key, key_err = self._parse_session_key_header(request) diff --git a/gateway/run.py b/gateway/run.py index fd3705b1f3a0..3382d5323c2a 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -4143,55 +4143,19 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew return 0 def _active_api_run_count(self) -> int: - """Count of in-flight api_server agent runs across all adapters. + """Count API-server work that is outside ``_running_agents``. - 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. + The primary API server owns the sole HTTP listener. Secondary multiplex + profiles cannot create an ``api_server`` adapter because it binds a port, + so only the primary registry is a supported source of this work. """ - total = 0 try: - from gateway.config import Platform + adapter = getattr(self, "adapters", {}).get(Platform.API_SERVER) + helper = getattr(adapter, "active_agent_work_count", None) + return max(0, int(helper())) if callable(helper) else 0 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 diff --git a/tests/gateway/test_api_server_active_work_drain.py b/tests/gateway/test_api_server_active_work_drain.py index 89444d18a407..780eb88d8906 100644 --- a/tests/gateway/test_api_server_active_work_drain.py +++ b/tests/gateway/test_api_server_active_work_drain.py @@ -1,11 +1,10 @@ -"""Tests for #63529: the gateway shutdown drain was structurally blind to -in-flight api_server (desk/API) agent runs. +"""Regression coverage for #63529 API-server shutdown draining. -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). +API-server work is adapter-owned rather than tracked by +``GatewayRunner._running_agents``. The shutdown drain must account for the +same live state as the API concurrency limiter, including a ``/v1/runs`` task +that exists before its agent has been constructed, and it must refuse new API +turns once the gateway starts draining. """ import asyncio @@ -13,81 +12,106 @@ from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest +from aiohttp import web +from aiohttp.test_utils import TestClient, TestServer +from gateway.config import Platform, PlatformConfig +from gateway.platforms.api_server import APIServerAdapter 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 +class _RunTask: + def __init__(self, done: bool = False): + self._done = done - active = {rid: MagicMock() for rid in (active_ids or [])} + def done(self) -> bool: + return self._done + + +def _make_api_adapter(*, inflight: int = 0, queued_ids=()): + tasks = {run_id: _RunTask() for run_id in queued_ids} adapter = SimpleNamespace( platform=Platform.API_SERVER, _inflight_agent_runs=inflight, - _active_run_agents=active, + _active_run_tasks=tasks, ) def active_agent_work_count() -> int: - return int(adapter._inflight_agent_runs) + len(adapter._active_run_agents) + return int(adapter._inflight_agent_runs) + sum( + not task.done() for task in adapter._active_run_tasks.values() + ) adapter.active_agent_work_count = active_agent_work_count return adapter +def _make_admission_app(adapter: APIServerAdapter) -> web.Application: + app = web.Application() + app.router.add_post("/api/sessions/{session_id}/chat", adapter._handle_session_chat) + app.router.add_post( + "/api/sessions/{session_id}/chat/stream", adapter._handle_session_chat_stream + ) + app.router.add_post("/v1/chat/completions", adapter._handle_chat_completions) + app.router.add_post("/v1/responses", adapter._handle_responses) + app.router.add_post("/v1/runs", adapter._handle_runs) + return app + + class TestActiveApiRunCount: - def test_zero_when_no_api_adapters(self): + def test_zero_when_no_api_adapter(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): + def test_delegates_to_primary_api_adapter(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"])} + runner.adapters = { + Platform.API_SERVER: _make_api_adapter(inflight=2, queued_ids=["r1"]) + } 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 = {} + runner.adapters = {Platform.DISCORD: other} assert runner._active_api_run_count() == 0 - def test_never_raises_on_broken_adapters(self): + def test_never_raises_on_broken_adapter(self): runner, _adapter = make_restart_runner() class Bad: - platform = property(lambda self: (_ for _ in ()).throw(RuntimeError("boom"))) + platform = Platform.API_SERVER - runner.adapters = {"bad": Bad()} - runner._profile_adapters = None + @staticmethod + def active_agent_work_count() -> int: + raise RuntimeError("boom") + + runner.adapters = {Platform.API_SERVER: Bad()} 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) + def test_counts_live_run_task_before_agent_creation(self): + adapter = APIServerAdapter(PlatformConfig(enabled=True)) adapter._inflight_agent_runs = 2 - adapter._active_run_agents = {"r1": object(), "r2": object()} - assert adapter.active_agent_work_count() == 4 + adapter._active_run_tasks = { + "queued": _RunTask(), + "finished": _RunTask(done=True), + } + adapter._active_run_agents = {} + + 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 class TestDrainWaitsForApiWork: @@ -95,56 +119,79 @@ class TestDrainWaitsForApiWork: 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): + async def test_drain_waits_for_real_queued_run_before_agent_creation(self): + """A live /v1/runs task must block drain before it has an agent.""" runner, _adapter = make_restart_runner() - api = _make_api_adapter(inflight=1) - runner.adapters = {"api": api} - runner._profile_adapters = {} + api = APIServerAdapter(PlatformConfig(enabled=True)) + runner.adapters = {Platform.API_SERVER: api} + app = _make_admission_app(api) + original_create_task = asyncio.create_task + task_started = asyncio.Event() + allow_task = asyncio.Event() - async def finish_run(): - await asyncio.sleep(0.12) - api._inflight_agent_runs = 0 + def delayed_create_task(coro): + async def delayed(): + task_started.set() + await allow_task.wait() + return await coro - task = asyncio.create_task(finish_run()) - _snapshot, timed_out = await runner._drain_active_agents(2.0) - await task + return original_create_task(delayed()) - assert timed_out is False, ( - "drain must wait for api_server work, not report active_at_start=0" - ) + mock_agent = MagicMock() + mock_agent.run_conversation.return_value = {"final_response": "done"} + mock_agent.session_prompt_tokens = 0 + mock_agent.session_completion_tokens = 0 + mock_agent.session_total_tokens = 0 + + with patch( + "gateway.platforms.api_server.asyncio.create_task", + side_effect=delayed_create_task, + ), patch.object(api, "_create_agent", return_value=mock_agent): + async with TestClient(TestServer(app)) as client: + response = await client.post("/v1/runs", json={"input": "hello"}) + assert response.status == 202 + await task_started.wait() + + assert api._active_run_agents == {} + assert runner._active_api_run_count() == 1 + drain_task = original_create_task(runner._drain_active_agents(2.0)) + await asyncio.sleep(0.1) + assert not drain_task.done() + + allow_task.set() + _snapshot, timed_out = await drain_task + + 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 = {"api": _make_api_adapter(inflight=1)} - runner._profile_adapters = {} + 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 @pytest.mark.asyncio - async def test_drain_still_waits_for_chat_and_cron(self): + 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 = {"api": _make_api_adapter(inflight=1)} - runner._profile_adapters = {} + 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["api"]._inflight_agent_runs = 0 + runner.adapters[Platform.API_SERVER]._active_run_tasks.clear() task = asyncio.create_task(finish_all()) try: @@ -154,3 +201,28 @@ class TestDrainWaitsForApiWork: sched._running_job_ids.discard("job-1") assert timed_out is False + + +class TestDrainAdmission: + @pytest.mark.asyncio + async def test_drain_refuses_every_agent_start_endpoint(self): + adapter = APIServerAdapter(PlatformConfig(enabled=True)) + runner = SimpleNamespace(_draining=True, _external_drain_active=False) + app = _make_admission_app(adapter) + paths = ( + "/api/sessions/missing/chat", + "/api/sessions/missing/chat/stream", + "/v1/chat/completions", + "/v1/responses", + "/v1/runs", + ) + + with patch("gateway.run._gateway_runner_ref", lambda: runner): + async with TestClient(TestServer(app)) as client: + for path in paths: + response = await client.post(path, json={}) + payload = await response.json() + + assert response.status == 503 + assert response.headers["Retry-After"] == "1" + assert payload["error"]["code"] == "gateway_draining" diff --git a/tests/gateway/test_shutdown_cache_cleanup.py b/tests/gateway/test_shutdown_cache_cleanup.py index e03f651c9628..45a3461d04ee 100644 --- a/tests/gateway/test_shutdown_cache_cleanup.py +++ b/tests/gateway/test_shutdown_cache_cleanup.py @@ -60,6 +60,11 @@ class _FakeGateway: # there's never in-flight cron work to report. return 0 + def _active_api_run_count(self): + # The shutdown log also reports adapter-owned API work (#63529). + # This fake has no API server adapter, so it is always idle. + return 0 + def _update_runtime_status(self, *_a, **_kw): pass