diff --git a/gateway/run_shutdown.py b/gateway/run_shutdown.py index ef0ffe2ef6..e0dd54c01a 100644 --- a/gateway/run_shutdown.py +++ b/gateway/run_shutdown.py @@ -826,12 +826,9 @@ class GatewayShutdownMixin: _maybe_update_status(force=True) if not self._running_agents and not (_cron0 or _api0 or _deferred0): return snapshot, False - # Cron has its own deadline: a chat turn is announced+resumable; a killed cron run is a permanent failure. - # ``timeout`` (``restart_drain_timeout``) defaults to 0 because interrupting a chat turn is - # announced and resumable; a cron run killed mid-flight is recorded in jobs.json as a permanent - # failure nobody is waiting on. Sharing one budget meant the default config could report - # ``timed_out=True`` after 0.00s with a cron job in flight and kill it — the drain never even - # entered this loop (#82161). + # Cron and api_server runs ride the cron floor: a chat turn is announced+resumable, but a killed + # cron run is a permanent failure and a killed /v1 run fails a caller blocked on its result. + # On ``restart_drain_timeout``'s 0 default they were killed after 0.00s (#82161, #132989). started = loop.time() deadline = started + timeout cron_deadline = started + (timeout if cron_timeout is None else cron_timeout) @@ -839,7 +836,7 @@ class GatewayShutdownMixin: def _still_draining() -> bool: now = loop.time() agents, cron, api, deferred = self._drain_work_counts() - return bool(((agents or api or deferred) and now < deadline) or (cron and now < cron_deadline)) + return bool(((agents or deferred) and now < deadline) or ((cron or api) and now < cron_deadline)) # Both budgets at 0 = an expired deadline (loop unentered), so timed_out still comes from real state. while _still_draining(): diff --git a/tests/gateway/test_cron_drain_floor.py b/tests/gateway/test_cron_drain_floor.py index bb826df2dc..8ec20c5a92 100644 --- a/tests/gateway/test_cron_drain_floor.py +++ b/tests/gateway/test_cron_drain_floor.py @@ -17,9 +17,11 @@ post-drain cleanup window. """ import asyncio +from types import SimpleNamespace import pytest +from gateway.config import Platform from gateway.restart import ( CRON_DRAIN_CLEANUP_RESERVE_S, DEFAULT_GATEWAY_CRON_DRAIN_TIMEOUT, @@ -66,6 +68,23 @@ class TestDrainWaitsForCronOnDefaultConfig: ) assert runner._active_cron_job_count() == 0 + @pytest.mark.asyncio + async def test_zero_drain_timeout_still_waits_for_api_server_run(self): + """#132989: a /v1 run's caller is blocked on its result, so it rides the cron floor too.""" + runner, _adapter = make_restart_runner() + live = [1] + runner.adapters = {Platform.API_SERVER: SimpleNamespace(active_agent_work_count=lambda: live[0])} + + async def finish_run(): + await asyncio.sleep(0.12) + live[0] = 0 + + task = asyncio.create_task(finish_run()) + _snapshot, timed_out = await runner._drain_active_agents(0.0, 2.0) + await task + + assert timed_out is False, "api_server run was interrupted on the 0s chat budget (#132989)" + @pytest.mark.asyncio async def test_cron_floor_is_bounded_not_indefinite(self): """A job that never finishes must still lose, or a cron-triggered