From 7fb344f92f40c00a35e13cd2e2d386cd777723db Mon Sep 17 00:00:00 2001
From: kshitijk4poor <82637225+kshitijk4poor@users.noreply.github.com>
Date: Mon, 5 Oct 2026 13:48:07 +0530
Subject: [PATCH] fix(gateway): give in-flight api_server runs the cron drain
 floor

stop()'s drain grouped api_server (/v1) runs with chat turns on the
restart_drain_timeout budget, which defaults to 0. A default-config stop
therefore never entered the drain loop and hard-interrupted every
in-flight /v1 run after 0.00s, even though a caller is blocked on its
result and nothing resumes it. #82161 split only cron off onto its own
watchdog-clamped floor (agent.cron_drain_timeout).

Let api_server runs ride that same clamped floor (no new config key);
chat turns and deferred post-turn workers stay on the chat budget, and
cron_drain_timeout: 0 still opts out. Reuses the existing
_active_api_run_count counter via _drain_work_counts.

Fixes #132989. Diagnosis and the 'ride the cron deadline' approach by
the reporter, @ItoDonno.
---
 gateway/run_shutdown.py                | 11 ++++-------
 tests/gateway/test_cron_drain_floor.py | 19 +++++++++++++++++++
 2 files changed, 23 insertions(+), 7 deletions(-)

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