fix(gateway): drain in-flight cron jobs before shutdown tool kill

/update and other shutdown paths only waited on gateway session agents,
so active cron tool work was killed immediately in final-cleanup while
the scheduler could still mark the job successful (#60432).
This commit is contained in:
HexLab98 2026-07-08 08:24:34 +07:00 committed by Teknium
parent a208b7eeb4
commit 862aee4956
2 changed files with 35 additions and 7 deletions

View file

@ -298,6 +298,13 @@ _parallel_pool_max_workers: Optional[int] = None
_running_job_ids: set = set()
_running_lock = threading.Lock()
def cron_jobs_in_flight() -> int:
"""Return how many cron jobs are currently executing agent/tool work."""
with _running_lock:
return len(_running_job_ids)
# Sequential (env-mutating) cron jobs — workdir jobs that touch
# process-global runtime state — must run one at a time, but must NOT block the
# ticker thread. A persistent single-thread executor preserves ordering across

View file

@ -5548,20 +5548,30 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
return True
async def _drain_active_agents(self, timeout: float) -> tuple[Dict[str, Any], bool]:
from cron.scheduler import cron_jobs_in_flight
snapshot = self._snapshot_running_agents()
last_active_count = self._running_agent_count()
last_cron_count = cron_jobs_in_flight()
last_status_at = 0.0
def _maybe_update_status(force: bool = False) -> None:
nonlocal last_active_count, last_status_at
nonlocal last_active_count, last_cron_count, last_status_at
now = asyncio.get_running_loop().time()
active_count = self._running_agent_count()
if force or active_count != last_active_count or (now - last_status_at) >= 1.0:
cron_count = cron_jobs_in_flight()
if (
force
or active_count != last_active_count
or cron_count != last_cron_count
or (now - last_status_at) >= 1.0
):
self._update_runtime_status("draining")
last_active_count = active_count
last_cron_count = cron_count
last_status_at = now
if not self._running_agents:
if not self._running_agents and last_cron_count == 0:
_maybe_update_status(force=True)
return snapshot, False
@ -5570,10 +5580,13 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
return snapshot, True
deadline = asyncio.get_running_loop().time() + timeout
while self._running_agents and asyncio.get_running_loop().time() < deadline:
while (
(self._running_agents or cron_jobs_in_flight())
and asyncio.get_running_loop().time() < deadline
):
_maybe_update_status()
await asyncio.sleep(0.1)
timed_out = bool(self._running_agents)
timed_out = bool(self._running_agents or cron_jobs_in_flight())
_maybe_update_status(force=True)
return snapshot, timed_out
@ -8004,16 +8017,22 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
except Exception as _e:
logger.debug("pre-drain mark_resume_pending failed for %s: %s", _sk, _e)
from cron.scheduler import cron_jobs_in_flight
_cron_at_start = cron_jobs_in_flight()
_drain_started_at = time.monotonic()
active_agents, timed_out = await self._drain_active_agents(timeout)
logger.info(
"Shutdown phase: drain done at +%.2fs (drain took %.2fs, "
"timed_out=%s, active_at_start=%d, active_now=%d)",
"timed_out=%s, active_at_start=%d, active_now=%d, "
"cron_at_start=%d, cron_now=%d)",
_phase_elapsed(),
time.monotonic() - _drain_started_at,
timed_out,
len(active_agents),
self._running_agent_count(),
_cron_at_start,
cron_jobs_in_flight(),
)
if not timed_out:
@ -8032,9 +8051,11 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
if timed_out:
logger.warning(
"Gateway drain timed out after %.1fs with %d active agent(s); interrupting remaining work.",
"Gateway drain timed out after %.1fs with %d active agent(s) "
"and %d in-flight cron job(s); interrupting remaining work.",
timeout,
self._running_agent_count(),
cron_jobs_in_flight(),
)
# Mark forcibly-interrupted sessions as resume_pending BEFORE
# interrupting the agents. This preserves each session's