From 862aee49564cacbbe9e397adf40f7f87cb575f16 Mon Sep 17 00:00:00 2001
From: HexLab98 <liruixinch@outlook.com>
Date: Wed, 8 Jul 2026 08:24:34 +0700
Subject: [PATCH] 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).
---
 cron/scheduler.py |  7 +++++++
 gateway/run.py    | 35 ++++++++++++++++++++++++++++-------
 2 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/cron/scheduler.py b/cron/scheduler.py
index 21baab05482..0a9a9813180 100644
--- a/cron/scheduler.py
+++ b/cron/scheduler.py
@@ -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
diff --git a/gateway/run.py b/gateway/run.py
index 11eff2fd9df..a3752d5beb0 100644
--- a/gateway/run.py
+++ b/gateway/run.py
@@ -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