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.
This commit is contained in:
kshitijk4poor
2026-10-05 13:48:07 +05:30
committed by kshitij
parent a9cd8de4ee
commit 7fb344f92f
2 changed files with 23 additions and 7 deletions

View File

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

View File

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