From c2ecfe0a729ef442d5d28a1e6c833ee4d885689f Mon Sep 17 00:00:00 2001 From: HexLab98 Date: Wed, 8 Jul 2026 08:24:34 +0700 Subject: [PATCH 1/5] 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 21baab05482f4..0a9a9813180f5 100644 --- a/cron/scheduler.py +++ b/cron/scheduler.py @@ -298,6 +298,13 @@ def _is_token(line: str) -> bool: _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 11eff2fd9df4a..a3752d5beb013 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -5548,20 +5548,30 @@ async def _handle_active_session_busy_message(self, event: MessageEvent, session 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 @@ def _maybe_update_status(force: bool = False) -> None: 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 @@ def _phase_elapsed() -> float: 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 @@ def _phase_elapsed() -> float: 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 From a2f3aafdf65b541a067426c035043fc56edf6567 Mon Sep 17 00:00:00 2001 From: HexLab98 Date: Wed, 8 Jul 2026 08:24:34 +0700 Subject: [PATCH 2/5] test(gateway): cover cron drain during gateway shutdown (#60432) --- tests/gateway/test_update_cron_drain.py | 86 +++++++++++++++++++++++++ 1 file changed, 86 insertions(+) create mode 100644 tests/gateway/test_update_cron_drain.py diff --git a/tests/gateway/test_update_cron_drain.py b/tests/gateway/test_update_cron_drain.py new file mode 100644 index 0000000000000..ee34d72cd82ee --- /dev/null +++ b/tests/gateway/test_update_cron_drain.py @@ -0,0 +1,86 @@ +"""Regression tests for #60432. + +``/update`` (and other gateway shutdown paths) must drain in-flight cron jobs +before ``process_registry.kill_all()`` runs in final-cleanup. Cron work runs on +a thread-pool worker and is tracked in ``cron.scheduler._running_job_ids``, not +in ``GatewayRunner._running_agents`` — so a zero-agent drain must still wait +for cron to finish (or time out and take the interrupt/kill path). +""" +import asyncio +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from tests.gateway.restart_test_helpers import make_restart_runner + + +@pytest.mark.asyncio +async def test_drain_active_agents_waits_for_in_flight_cron_jobs(): + runner, _adapter = make_restart_runner() + runner._running_agents = {} + + cron_count = [1] + + def _cron_in_flight(): + return cron_count[0] + + async def finish_cron(): + await asyncio.sleep(0.15) + cron_count[0] = 0 + + with patch("cron.scheduler.cron_jobs_in_flight", side_effect=_cron_in_flight): + task = asyncio.create_task(finish_cron()) + _snapshot, timed_out = await runner._drain_active_agents(1.0) + await task + + assert timed_out is False + assert _snapshot == {} + + +@pytest.mark.asyncio +async def test_drain_active_agents_times_out_when_cron_still_running(): + runner, _adapter = make_restart_runner() + runner._running_agents = {} + + with patch("cron.scheduler.cron_jobs_in_flight", return_value=1): + _snapshot, timed_out = await runner._drain_active_agents(0.05) + + assert timed_out is True + assert _snapshot == {} + + +@pytest.mark.asyncio +async def test_gateway_stop_waits_for_cron_before_final_tool_kill(): + """Graceful cron completion must finish before final-cleanup kill_all.""" + runner, adapter = make_restart_runner() + runner._restart_drain_timeout = 1.0 + + cron_count = [1] + call_order: list[str] = [] + + def _cron_in_flight(): + return cron_count[0] + + def _fake_kill_all(task_id=None): + call_order.append("kill_all") + return 0 + + async def finish_cron(): + await asyncio.sleep(0.12) + cron_count[0] = 0 + + with ( + patch("cron.scheduler.cron_jobs_in_flight", side_effect=_cron_in_flight), + patch("gateway.status.remove_pid_file"), + patch("gateway.status.write_runtime_status"), + patch("agent.auxiliary_client.shutdown_cached_clients"), + patch("tools.process_registry.process_registry") as registry_mock, + ): + registry_mock.kill_all.side_effect = _fake_kill_all + adapter.disconnect = AsyncMock() + + cron_task = asyncio.create_task(finish_cron()) + await runner.stop() + await cron_task + + assert call_order == ["kill_all"] From c28b5d5220854fae7d1c7ec687a8b47778a9dafd Mon Sep 17 00:00:00 2001 From: joaomarcos Date: Tue, 7 Jul 2026 23:05:26 -0300 Subject: [PATCH 3/5] fix(gateway,cron): make shutdown drain visible to in-flight cron work Cron jobs run through cron/scheduler.py's own ThreadPoolExecutor via a standalone AIAgent (run_job/run_one_job), entirely outside GatewayRunner._running_agents -- the dict _drain_active_agents() and every other active-work check on that class reads. A gateway shutdown (/update, /restart, and SIGUSR1 all funnel through the same stop()) could log active_at_start=0 and immediately kill tool subprocesses while a cron job's terminal command was still running, with no wait and no indication anything was interrupted. Real-world impact (from the issue): a scheduled daily briefing cron job was in flight during /update, its tool subprocess got killed by the unconditional shutdown cleanup, and the job was never marked failed -- it simply never completed or delivered, with no error surfaced anywhere. A repro with a 30-minute `sleep` cron job in flight during /update reproduced the same pattern: subprocess killed at +0.22s of drain (active_at_start=0), the job's agent thread continued in-process and produced a plausible-looking final response from the truncated tool output, and the scheduler marked the run successful. Root cause is layered, not a single line: 1. GatewayRunner._drain_active_agents() only waits on _running_agents. Cron work was invisible to it, so drain returned instantly whenever the only active work was a cron job. 2. Even with visibility, the shutdown's final tool-subprocess kill (process_registry.kill_all()) is a global, unconditional sweep with no per-job targeting -- a long-running cron job that outlives the drain timeout still gets its subprocess killed. 3. cron/scheduler.py had no way to detect that a job's tool subprocess was killed out from under it mid-run; the agent thread kept going and its eventual (often degraded but plausible-looking) response got reported as a normal successful completion. Fix, three parts: - cron/scheduler.py: expose get_running_job_ids() (thread-safe snapshot of the existing _running_job_ids set, already used to prevent double-dispatch) so the gateway can read cron's in-flight state without reaching into private module internals. - gateway/run.py: GatewayRunner._active_cron_job_count() reads that snapshot. _drain_active_agents() now waits on (_running_agents OR active cron jobs), so a cron-only workload gets the same bounded wait chat sessions already get instead of an instant active_at_start=0. Shutdown drain logging gains cron_active_at_start/cron_active_now fields alongside the existing ones (unchanged, for compat). - cron/scheduler.py: mark_running_jobs_interrupted(reason), called by gateway/run.py's _kill_tool_subprocesses() right after process_registry.kill_all(), marks every job still in _running_job_ids at that instant as failed/interrupted via the existing mark_job_run() -- and records the job IDs in _interrupted_job_ids BEFORE writing, so run_one_job()'s own eventual completion for the same run (racing in its own thread) checks that flag and skips its normal write instead of clobbering the interrupted status with a false "ok" produced from the now-truncated tool output. This does not attempt to correlate a killed PID to a specific job ID (process_registry tracks PIDs, not job IDs) -- any job still dispatched at the moment of a forced kill is treated as interrupted, matching the existing coarser precedent set by _interrupt_running_agents(), which interrupts every entry in _running_agents on a drain timeout without per-agent correlation either. Deliberately out of scope (flagged in the issue as a separate, lower-priority concern): startup-time reconciliation of cron runs that started but never reached a terminal status. Testing: - tests/cron/test_shutdown_interrupt.py (12 tests): get_running_job_ids snapshot semantics, mark_running_jobs_interrupted marking/no-op/ partial-failure behavior, and -- the core race guard -- run_one_job skipping its own last_status write (both the success path and the exception path) when the shutdown path already marked the run interrupted, with a control test proving ordinary un-interrupted completions are unaffected. - tests/gateway/test_cron_active_work_drain.py (9 tests): _active_cron_job_count reading cron state and failing closed (0) if the cron module is unavailable; _drain_active_agents waiting for an in-flight cron job the same way it waits for chat sessions, timing out if the job outruns the window, and leaving existing chat-session drain behavior unchanged; a full runner.stop() integration test (drain-timeout path) proving mark_running_jobs_interrupted actually fires with the right job ID when a tool subprocess is force-killed, plus a no-op control when nothing cron-related is in flight. - tests/gateway/test_shutdown_cache_cleanup.py: added _active_cron_job_count() to that file's hand-rolled _FakeGateway test double, which stop() now calls -- without it those 8 pre-existing tests AttributeError (caught by fail-then-pass below, not a production bug). Fail-then-pass: reverted gateway/run.py + cron/scheduler.py, all 21 new tests fail (fixture/attribute errors -- the feature doesn't exist yet); restored, all 21 pass. Regression check: ran the full plausibly-affected surface -- tests/gateway/{test_gateway_shutdown,test_restart_drain, test_restart_notification,test_restart_redelivery_dedup, test_restart_resume_pending,test_restart_service_detection, test_shutdown_cache_cleanup,test_stuck_loop,test_clean_shutdown_marker, test_external_drain_control,test_session_state_cleanup, test_update_command,test_update_streaming}.py plus tests/cron/ (944 tests) -- against a clean upstream/main checkout and against this branch. Diffed the two FAILED lists: identical, 20 pre-existing failures on both sides (Windows-locale/cp1252 file-encoding issues and Unix-permission-bit assertions that don't apply on this Windows dev box), zero new failures, zero fixed-by-accident. The 8 test_shutdown_cache_cleanup.py failures found mid-development were from the _FakeGateway gap above, fixed in the same commit and confirmed clean on the final rerun (diff against baseline: exit 0). Fixes #60432 --- cron/scheduler.py | 87 +++++++- gateway/run.py | 64 +++++- tests/cron/test_shutdown_interrupt.py | 209 +++++++++++++++++++ tests/gateway/test_cron_active_work_drain.py | 185 ++++++++++++++++ tests/gateway/test_shutdown_cache_cleanup.py | 6 + 5 files changed, 538 insertions(+), 13 deletions(-) create mode 100644 tests/cron/test_shutdown_interrupt.py create mode 100644 tests/gateway/test_cron_active_work_drain.py diff --git a/cron/scheduler.py b/cron/scheduler.py index 0a9a9813180f5..1994a857fede6 100644 --- a/cron/scheduler.py +++ b/cron/scheduler.py @@ -305,6 +305,87 @@ def cron_jobs_in_flight() -> int: return len(_running_job_ids) +# Job IDs the gateway shutdown path force-killed the tool subprocess of +# while still in ``_running_job_ids`` (see ``mark_running_jobs_interrupted`` +# below). ``run_one_job``'s own completion path checks this set before +# writing its own ``last_status`` so a cron agent thread that keeps running +# in-process after its tool was killed out from under it — and produces a +# plausible-looking final response from truncated output — can never +# overwrite the interrupted status with a false "ok" (#60432). +_interrupted_job_ids: set = set() + + +def get_running_job_ids() -> "frozenset[str]": + """Thread-safe snapshot of cron job IDs currently executing. + + A job ID is a member from the moment ``_submit_with_guard`` dispatches + it onto the parallel/sequential pool until ``_process_job`` returns — + i.e. for the job's *entire* run, tool calls included, not just the + ticker's dispatch instant. + + The gateway shutdown path (``gateway/run.py::GatewayRunner. + _drain_active_agents``) reads this to treat in-flight cron work as + active the same way it already treats in-flight chat sessions via + ``_running_agents`` — cron jobs run through their own thread pool here, + entirely outside that dict, so without this the drain is structurally + blind to them (#60432). + """ + with _running_lock: + return frozenset(_running_job_ids) + + +def mark_running_jobs_interrupted(reason: str) -> list: + """Best-effort: mark every currently in-flight cron job interrupted. + + Called by the gateway shutdown path immediately after it force-kills + tool subprocesses (``process_registry.kill_all()``). A job whose tool + subprocess was just killed out from under it must never be allowed to + report success — even though its agent thread is still alive in this + same process and may go on to produce a plausible-looking final + response from the now-truncated tool output. + + Records the job IDs in ``_interrupted_job_ids`` BEFORE writing + ``last_status`` so ``run_one_job``'s own eventual completion for the + same job (racing in its own thread) sees the flag and skips its normal + write instead of clobbering this one — see the check near the end of + ``run_one_job``. This does not attempt to correlate the killed + subprocess PID to a specific job ID (the process registry tracks PIDs, + not cron job IDs); any job still dispatched at the moment of a forced + kill is treated as interrupted, matching the coarser precedent already + set by ``GatewayRunner._interrupt_running_agents``, which interrupts + every entry in ``_running_agents`` on a drain timeout without + per-agent correlation either. + + Returns the list of job IDs marked, for the caller to log. + """ + with _running_lock: + job_ids = list(_running_job_ids) + _interrupted_job_ids.update(job_ids) + marked = [] + for job_id in job_ids: + try: + mark_job_run(job_id, False, reason) + marked.append(job_id) + except Exception as e: + logger.warning("Failed to mark job %s interrupted: %s", job_id, e) + return marked + + +def _consume_interrupted_flag(job_id: str) -> bool: + """Return True and clear the flag if the shutdown path already marked + ``job_id`` interrupted (see ``mark_running_jobs_interrupted``). + + Called by ``run_one_job`` right before it would otherwise write its own + ``last_status``. Consuming (discarding) rather than just checking keeps + the flag from leaking across a later, unrelated run of the same job ID + (recurring jobs reuse their ID every fire).""" + with _running_lock: + if job_id in _interrupted_job_ids: + _interrupted_job_ids.discard(job_id) + return True + return False + + # 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 @@ -3377,12 +3458,14 @@ def run_one_job(job: dict, *, adapters=None, loop=None, verbose: bool = False) - success = False error = "Agent completed but produced empty response (model error, timeout, or misconfiguration)" - mark_job_run(job["id"], success, error, delivery_error=delivery_error) + if not _consume_interrupted_flag(job["id"]): + mark_job_run(job["id"], success, error, delivery_error=delivery_error) return True except Exception as e: logger.error("Error processing job %s: %s", job['id'], e) - mark_job_run(job["id"], False, str(e)) + if not _consume_interrupted_flag(job["id"]): + mark_job_run(job["id"], False, str(e)) return False diff --git a/gateway/run.py b/gateway/run.py index a3752d5beb013..527da4ff5132d 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -4080,6 +4080,26 @@ def _request_clean_exit(self, reason: str) -> None: def _running_agent_count(self) -> int: return len(self._running_agents) + def _active_cron_job_count(self) -> int: + """Count of cron jobs currently executing, from the cron scheduler's + own in-flight tracking (``cron.scheduler._running_job_ids``). + + Cron jobs run through a standalone ``AIAgent`` on the scheduler's own + thread pool (``cron/scheduler.py::run_job``), entirely outside + ``self._running_agents`` — the dict every OTHER active-work check on + this class (``_running_agent_count``, ``_drain_active_agents``) reads. + Without this, the shutdown drain is structurally blind to in-flight + cron work: it can report ``active_at_start=0`` and proceed straight + to killing tool subprocesses while a cron job's terminal command is + still running (#60432). Best-effort: returns 0 if the cron module + can't be imported (e.g. a minimal test double for this class). + """ + try: + from cron.scheduler import get_running_job_ids + return len(get_running_job_ids()) + except Exception: + return 0 + # ── scale-to-zero idle detection / dormant-quiesce (Phase 0) ────────────── # The gateway-side BEHAVIOUR that consumes the relay scale-to-zero primitives # (gateway-gateway Phase 5). Pure logic lives in gateway/scale_to_zero.py; the @@ -5548,18 +5568,16 @@ async def _handle_active_session_busy_message(self, event: MessageEvent, session 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_cron_count = self._active_cron_job_count() last_status_at = 0.0 def _maybe_update_status(force: bool = False) -> None: nonlocal last_active_count, last_cron_count, last_status_at now = asyncio.get_running_loop().time() active_count = self._running_agent_count() - cron_count = cron_jobs_in_flight() + cron_count = self._active_cron_job_count() if ( force or active_count != last_active_count @@ -5571,6 +5589,11 @@ def _maybe_update_status(force: bool = False) -> None: last_cron_count = cron_count last_status_at = now + # Cron jobs run on the scheduler's own thread pool, outside + # ``self._running_agents`` — fold their in-flight count into the + # same wait/timeout this method already applies to chat sessions, + # or a cron job's tool work gets killed with zero warning the + # instant it's the only active thing running (#60432). if not self._running_agents and last_cron_count == 0: _maybe_update_status(force=True) return snapshot, False @@ -5581,12 +5604,12 @@ def _maybe_update_status(force: bool = False) -> None: deadline = asyncio.get_running_loop().time() + timeout while ( - (self._running_agents or cron_jobs_in_flight()) + (self._running_agents or self._active_cron_job_count()) and asyncio.get_running_loop().time() < deadline ): _maybe_update_status() await asyncio.sleep(0.1) - timed_out = bool(self._running_agents or cron_jobs_in_flight()) + timed_out = bool(self._running_agents) or bool(self._active_cron_job_count()) _maybe_update_status(force=True) return snapshot, timed_out @@ -7957,6 +7980,27 @@ def _kill_tool_subprocesses(phase: str) -> None: ) except Exception as _e: logger.debug("process_registry.kill_all (%s) error: %s", phase, _e) + try: + # Any cron job still dispatched at this instant just had + # its tool subprocess killed above (kill_all() has no + # per-job-ID targeting — it's a global sweep). Its agent + # thread is still alive in this process and may go on to + # produce a plausible-looking final response from the + # now-truncated tool output; mark the run interrupted so + # the scheduler can never report that as success (#60432). + # No-op when no cron job is in flight. + from cron.scheduler import mark_running_jobs_interrupted + _interrupted = mark_running_jobs_interrupted( + f"Gateway shutdown ({phase}) killed the job's tool " + "subprocess before the run finished." + ) + if _interrupted: + logger.warning( + "Shutdown (%s): marked %d in-flight cron job(s) interrupted: %s", + phase, len(_interrupted), ", ".join(_interrupted), + ) + except Exception as _e: + logger.debug("mark_running_jobs_interrupted (%s) error: %s", phase, _e) try: from tools.async_delegation import interrupt_all as _interrupt_async _async_n = _interrupt_async(reason=f"gateway shutdown ({phase})") @@ -8017,9 +8061,7 @@ def _phase_elapsed() -> float: 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() + _cron_at_start = self._active_cron_job_count() _drain_started_at = time.monotonic() active_agents, timed_out = await self._drain_active_agents(timeout) logger.info( @@ -8032,7 +8074,7 @@ def _phase_elapsed() -> float: len(active_agents), self._running_agent_count(), _cron_at_start, - cron_jobs_in_flight(), + self._active_cron_job_count(), ) if not timed_out: @@ -8055,7 +8097,7 @@ def _phase_elapsed() -> float: "and %d in-flight cron job(s); interrupting remaining work.", timeout, self._running_agent_count(), - cron_jobs_in_flight(), + self._active_cron_job_count(), ) # Mark forcibly-interrupted sessions as resume_pending BEFORE # interrupting the agents. This preserves each session's diff --git a/tests/cron/test_shutdown_interrupt.py b/tests/cron/test_shutdown_interrupt.py new file mode 100644 index 0000000000000..2b5077fc9d277 --- /dev/null +++ b/tests/cron/test_shutdown_interrupt.py @@ -0,0 +1,209 @@ +"""Tests for #60432: cron jobs must not be silently invisible to gateway +shutdown, and a job whose tool subprocess got killed by shutdown must +never be reported as a successful run. + +Covers the cron/scheduler.py primitives directly: + - get_running_job_ids() -- thread-safe snapshot the gateway drain reads + - mark_running_jobs_interrupted() -- called by the gateway right after + it force-kills tool subprocesses + - the interrupted-flag race guard in run_one_job(), which must win over + the job's own thread finishing normally with a plausible-looking + result AFTER its tool was already killed out from under it +""" + +from unittest.mock import patch + +import pytest + + +@pytest.fixture(autouse=True) +def _reset_scheduler_state(): + """Every test starts from a clean slate and leaves one behind, since + these sets are module-level globals shared across the test process.""" + import cron.scheduler as sched + + sched._running_job_ids.clear() + sched._interrupted_job_ids.clear() + yield + sched._running_job_ids.clear() + sched._interrupted_job_ids.clear() + + +class TestGetRunningJobIds: + def test_empty_when_nothing_running(self): + import cron.scheduler as sched + + assert sched.get_running_job_ids() == frozenset() + + def test_reflects_in_flight_jobs(self): + import cron.scheduler as sched + + sched._running_job_ids.add("job-1") + sched._running_job_ids.add("job-2") + + result = sched.get_running_job_ids() + + assert result == frozenset({"job-1", "job-2"}) + + def test_snapshot_is_immutable_and_independent(self): + """Mutating _running_job_ids after the call must not change the + already-returned snapshot -- callers (the gateway drain loop) rely + on this to safely count in a tight polling loop.""" + import cron.scheduler as sched + + sched._running_job_ids.add("job-1") + snapshot = sched.get_running_job_ids() + sched._running_job_ids.add("job-2") + + assert snapshot == frozenset({"job-1"}) + + +class TestMarkRunningJobsInterrupted: + def test_no_op_when_nothing_running(self): + import cron.scheduler as sched + + with patch("cron.scheduler.mark_job_run") as mock_mark: + marked = sched.mark_running_jobs_interrupted("shutdown") + + assert marked == [] + mock_mark.assert_not_called() + + def test_marks_every_in_flight_job(self): + import cron.scheduler as sched + + sched._running_job_ids.update({"job-1", "job-2"}) + + with patch("cron.scheduler.mark_job_run") as mock_mark: + marked = sched.mark_running_jobs_interrupted("gateway shutdown (final-cleanup)") + + assert sorted(marked) == ["job-1", "job-2"] + assert mock_mark.call_count == 2 + called_ids = {c.args[0] for c in mock_mark.call_args_list} + assert called_ids == {"job-1", "job-2"} + for c in mock_mark.call_args_list: + # success must be False -- an interrupted run is never "ok". + assert c.args[1] is False + assert "gateway shutdown" in c.args[2] + + def test_sets_interrupted_flag_for_consumption_by_run_one_job(self): + import cron.scheduler as sched + + sched._running_job_ids.add("job-1") + + with patch("cron.scheduler.mark_job_run"): + sched.mark_running_jobs_interrupted("shutdown") + + assert "job-1" in sched._interrupted_job_ids + + def test_one_job_marking_failure_does_not_block_the_others(self): + """mark_job_run raising for one job (e.g. a jobs.json write race) + must not prevent the rest from being marked -- this runs during + shutdown, there's no retry window.""" + import cron.scheduler as sched + + sched._running_job_ids.update({"job-1", "job-2"}) + + def _side_effect(job_id, success, reason, **kwargs): + if job_id == "job-1": + raise OSError("disk full") + + with patch("cron.scheduler.mark_job_run", side_effect=_side_effect): + marked = sched.mark_running_jobs_interrupted("shutdown") + + assert marked == ["job-2"] + + +class TestConsumeInterruptedFlag: + def test_false_when_not_marked(self): + import cron.scheduler as sched + + assert sched._consume_interrupted_flag("job-1") is False + + def test_true_and_clears_when_marked(self): + import cron.scheduler as sched + + sched._interrupted_job_ids.add("job-1") + + assert sched._consume_interrupted_flag("job-1") is True + # Consumed -- a second check (e.g. a later, unrelated fire of the + # same recurring job ID) must not still read as interrupted. + assert sched._consume_interrupted_flag("job-1") is False + + +class TestRunOneJobHonoursInterruptedFlag: + """run_one_job() must not let a job's own completion overwrite a + status the shutdown path already wrote for the same run.""" + + def _make_job(self, job_id="job-1"): + return {"id": job_id, "name": "test job", "prompt": "do work"} + + def test_success_path_skipped_when_interrupted(self): + import cron.scheduler as sched + + job = self._make_job() + sched._interrupted_job_ids.add(job["id"]) + + with patch("cron.scheduler.claim_dispatch", return_value=True), \ + patch("agent.secret_scope.set_secret_scope", return_value=None), \ + patch("agent.secret_scope.build_profile_secret_scope", return_value=None), \ + patch("agent.secret_scope.reset_secret_scope"), \ + patch( + "cron.scheduler.run_job", + return_value=(True, "full output", "final response", None), + ), \ + patch("cron.scheduler.save_job_output", return_value="/tmp/out.md"), \ + patch("cron.scheduler._is_cron_silence_response", return_value=False), \ + patch("cron.scheduler._deliver_result", return_value=None), \ + patch("cron.scheduler.mark_job_run") as mock_mark: + result = sched.run_one_job(job) + + assert result is True + # The would-be "success" write must NOT happen -- the shutdown + # path already wrote the authoritative interrupted status. + mock_mark.assert_not_called() + # Flag is consumed so a later, unrelated fire of the same job ID + # isn't permanently silenced. + assert job["id"] not in sched._interrupted_job_ids + + def test_success_path_writes_normally_when_not_interrupted(self): + """Control case: the guard must not swallow ordinary, un-interrupted + completions -- only ones the shutdown path explicitly flagged.""" + import cron.scheduler as sched + + job = self._make_job() + + with patch("cron.scheduler.claim_dispatch", return_value=True), \ + patch("agent.secret_scope.set_secret_scope", return_value=None), \ + patch("agent.secret_scope.build_profile_secret_scope", return_value=None), \ + patch("agent.secret_scope.reset_secret_scope"), \ + patch( + "cron.scheduler.run_job", + return_value=(True, "full output", "final response", None), + ), \ + patch("cron.scheduler.save_job_output", return_value="/tmp/out.md"), \ + patch("cron.scheduler._is_cron_silence_response", return_value=False), \ + patch("cron.scheduler._deliver_result", return_value=None), \ + patch("cron.scheduler.mark_job_run") as mock_mark: + result = sched.run_one_job(job) + + assert result is True + mock_mark.assert_called_once() + assert mock_mark.call_args.args[0] == job["id"] + assert mock_mark.call_args.args[1] is True # success + + def test_exception_path_also_honours_interrupted_flag(self): + import cron.scheduler as sched + + job = self._make_job() + sched._interrupted_job_ids.add(job["id"]) + + with patch("cron.scheduler.claim_dispatch", return_value=True), \ + patch("agent.secret_scope.set_secret_scope", return_value=None), \ + patch("agent.secret_scope.build_profile_secret_scope", return_value=None), \ + patch("agent.secret_scope.reset_secret_scope"), \ + patch("cron.scheduler.run_job", side_effect=RuntimeError("boom")), \ + patch("cron.scheduler.mark_job_run") as mock_mark: + result = sched.run_one_job(job) + + assert result is False + mock_mark.assert_not_called() diff --git a/tests/gateway/test_cron_active_work_drain.py b/tests/gateway/test_cron_active_work_drain.py new file mode 100644 index 0000000000000..f4bf1ca0af7d6 --- /dev/null +++ b/tests/gateway/test_cron_active_work_drain.py @@ -0,0 +1,185 @@ +"""Tests for #60432: the gateway shutdown drain was structurally blind to +in-flight cron work. Cron jobs run through cron/scheduler.py's own thread +pool, entirely outside ``GatewayRunner._running_agents`` -- the dict every +other active-work check on this class reads. A shutdown (``/update``, +``/restart``, SIGUSR1 -- they all funnel through the same ``stop()``) could +report ``active_at_start=0`` and immediately kill tool subprocesses while a +cron job's terminal command was still running. + +These tests cover the gateway side of the fix: + - _active_cron_job_count() reads cron.scheduler's in-flight job set + - _drain_active_agents() waits for cron work the same way it already + waits for chat sessions + - the final tool-subprocess kill marks any still-in-flight cron job + interrupted + +See tests/cron/test_shutdown_interrupt.py for the cron-side primitives +this relies on (get_running_job_ids, mark_running_jobs_interrupted). +""" + +import asyncio +from unittest.mock import MagicMock, patch + +import pytest + +from tests.gateway.restart_test_helpers import make_restart_runner + + +@pytest.fixture(autouse=True) +def _reset_cron_running_set(): + import cron.scheduler as sched + + sched._running_job_ids.clear() + sched._interrupted_job_ids.clear() + yield + sched._running_job_ids.clear() + sched._interrupted_job_ids.clear() + + +def _make_async_noop(): + async def _noop(*args, **kwargs): + return None + + return _noop + + +class TestActiveCronJobCount: + def test_zero_when_no_cron_jobs_running(self): + runner, _adapter = make_restart_runner() + assert runner._active_cron_job_count() == 0 + + def test_reflects_cron_scheduler_state(self): + import cron.scheduler as sched + + runner, _adapter = make_restart_runner() + sched._running_job_ids.add("job-1") + + assert runner._active_cron_job_count() == 1 + + def test_never_raises_if_cron_module_unavailable(self): + """Best-effort: a broken/absent import must not take shutdown + counting down with it.""" + runner, _adapter = make_restart_runner() + + with patch( + "cron.scheduler.get_running_job_ids", side_effect=ImportError("boom") + ): + assert runner._active_cron_job_count() == 0 + + +class TestDrainWaitsForCronWork: + @pytest.mark.asyncio + async def test_drain_returns_immediately_when_nothing_active(self): + runner, _adapter = make_restart_runner() + + _snapshot, timed_out = await runner._drain_active_agents(5.0) + + assert timed_out is False + + @pytest.mark.asyncio + async def test_drain_waits_for_in_flight_cron_job(self): + """Before this fix, a cron-only workload made active_at_start=0 + and the drain returned instantly -- this is the exact repro from + the issue (a `sleep 1800` cron job in flight during /update).""" + import cron.scheduler as sched + + runner, _adapter = make_restart_runner() + sched._running_job_ids.add("job-1") + + async def finish_job(): + await asyncio.sleep(0.12) + sched._running_job_ids.discard("job-1") + + task = asyncio.create_task(finish_job()) + _snapshot, timed_out = await runner._drain_active_agents(2.0) + await task + + assert timed_out is False, ( + "drain must wait for the cron job to finish, not report " + "active_at_start=0 and return instantly" + ) + + @pytest.mark.asyncio + async def test_drain_times_out_if_cron_job_outlives_the_window(self): + import cron.scheduler as sched + + runner, _adapter = make_restart_runner() + sched._running_job_ids.add("job-1") # never removed within the window + + _snapshot, timed_out = await runner._drain_active_agents(0.1) + + assert timed_out is True + + @pytest.mark.asyncio + async def test_drain_still_waits_for_chat_sessions_unchanged(self): + """Regression guard: folding cron into the check must not break + the pre-existing chat-session drain behavior.""" + runner, _adapter = make_restart_runner() + runner._running_agents = {"session-1": MagicMock()} + + async def finish_agent(): + await asyncio.sleep(0.12) + runner._running_agents.clear() + + task = asyncio.create_task(finish_agent()) + _snapshot, timed_out = await runner._drain_active_agents(2.0) + await task + + assert timed_out is False + + +class TestKillToolSubprocessesMarksCronInterrupted: + @pytest.mark.asyncio + async def test_in_flight_cron_job_marked_interrupted_on_forced_kill(self, monkeypatch): + import cron.scheduler as sched + import tools.process_registry as _pr + import tools.terminal_tool as _tt + import tools.browser_tool as _bt + + runner, adapter = make_restart_runner() + runner._restart_drain_timeout = 0.01 # force the timeout path + adapter.disconnect = _make_async_noop() + + sched._running_job_ids.add("job-1") + + monkeypatch.setattr(_pr.process_registry, "kill_all", lambda task_id=None: 1) + monkeypatch.setattr(_tt, "cleanup_all_environments", lambda: None) + monkeypatch.setattr(_bt, "cleanup_all_browsers", lambda: None) + + marked_calls = [] + real_mark = sched.mark_running_jobs_interrupted + + def _spy(reason): + result = real_mark(reason) + marked_calls.append((reason, result)) + return result + + monkeypatch.setattr(sched, "mark_running_jobs_interrupted", _spy) + + with patch("gateway.status.remove_pid_file"), patch("gateway.status.write_runtime_status"), \ + patch("cron.scheduler.mark_job_run"): + await runner.stop() + + assert marked_calls, "mark_running_jobs_interrupted was never called during shutdown" + assert any(result == ["job-1"] for _reason, result in marked_calls) + + @pytest.mark.asyncio + async def test_no_cron_jobs_running_is_a_silent_no_op(self, monkeypatch): + """Graceful shutdown with nothing in flight must not spuriously + mark or log anything cron-related.""" + import tools.process_registry as _pr + import tools.terminal_tool as _tt + import tools.browser_tool as _bt + + runner, adapter = make_restart_runner() + adapter.disconnect = _make_async_noop() + + monkeypatch.setattr(_pr.process_registry, "kill_all", lambda task_id=None: 0) + monkeypatch.setattr(_tt, "cleanup_all_environments", lambda: None) + monkeypatch.setattr(_bt, "cleanup_all_browsers", lambda: None) + + with patch("gateway.status.remove_pid_file"), patch("gateway.status.write_runtime_status"), \ + patch("cron.scheduler.mark_job_run") as mock_mark: + await runner.stop() + + mock_mark.assert_not_called() diff --git a/tests/gateway/test_shutdown_cache_cleanup.py b/tests/gateway/test_shutdown_cache_cleanup.py index 1b122a0d105a5..e03f651c9628b 100644 --- a/tests/gateway/test_shutdown_cache_cleanup.py +++ b/tests/gateway/test_shutdown_cache_cleanup.py @@ -54,6 +54,12 @@ def __init__(self): def _running_agent_count(self): return len(self._running_agents) + def _active_cron_job_count(self): + # stop() reads this alongside _running_agent_count when logging the + # drain snapshot (#60432) -- this fake has no cron scheduler, so + # there's never in-flight cron work to report. + return 0 + def _update_runtime_status(self, *_a, **_kw): pass From 0f98c5192b31ee4a10f2b53fb6ae7c5a24ca46f4 Mon Sep 17 00:00:00 2001 From: joaomarcos Date: Tue, 7 Jul 2026 23:22:57 -0300 Subject: [PATCH 4/5] fix(cron): stop interrupted jobs from delivering their pre-kill output Follow-up to the previous commit on #60432. The status-write guard (_consume_interrupted_flag, checked right before mark_job_run) closes the false-success bookkeeping gap, but run_one_job delivers its result BEFORE that check: delivery happens right after run_job() returns, mark_job_run happens at the very end. A job whose tool subprocess was killed mid-flight can still produce a plausible-looking final_response from the truncated output, and that response would reach the user via _deliver_result before the interrupted flag was ever consulted -- correct status in jobs.json, wrong message already sent. Adds _is_interrupted(), a non-destructive peek at the same _interrupted_job_ids set (_consume_interrupted_flag stays as the consuming, authoritative check right before the status write -- this needed a peek instead since the flag has to still be visible there). Checked right after save_job_output, before the deliver_content decision: if the run looked successful but was flagged interrupted, force success=False with an explicit interruption message. This routes delivery through the existing _summarize_cron_failure_for_delivery path (the same one a real failure already uses) instead of the raw final_response, so the user gets an honest "this run was interrupted" instead of a truncated/misleading result. Testing: 4 new tests in tests/cron/test_shutdown_interrupt.py -- _is_interrupted peek semantics (false/true/does-not-clear, as opposed to the consuming _consume_interrupted_flag), and the delivery-gate test itself, which mocks run_job to return a normal-looking success with a "plausible final response" while the job is pre-marked interrupted, and asserts _deliver_result receives the failure summary ("This run was interrupted.") instead, with the summarizer's error argument confirmed to mention the interruption. Fail-then-pass: reverted cron/scheduler.py only, the 4 new tests fail (3 on the missing _is_interrupted attribute, 1 -- the delivery-gate test -- on _summarize_cron_failure_for_delivery never being called, i.e. the raw response would have gone out); restored, all 16 tests in the file pass. Regression: tests/cron/ (683 tests) + test_cron_active_work_drain.py + test_gateway_shutdown.py + test_shutdown_cache_cleanup.py -- 11 pre-existing failures (Unix file-permission-bit and path-tilde assertions that don't apply on this Windows dev box), matching the same set already established as pre-existing in the prior commit's regression check. Zero new failures. Continues #60432 --- cron/scheduler.py | 29 +++++++++++ tests/cron/test_shutdown_interrupt.py | 69 +++++++++++++++++++++++++++ 2 files changed, 98 insertions(+) diff --git a/cron/scheduler.py b/cron/scheduler.py index 1994a857fede6..6b96fb2e1bea8 100644 --- a/cron/scheduler.py +++ b/cron/scheduler.py @@ -371,6 +371,21 @@ def mark_running_jobs_interrupted(reason: str) -> list: return marked +def _is_interrupted(job_id: str) -> bool: + """Non-destructive peek at whether the shutdown path has marked + ``job_id`` interrupted (see ``mark_running_jobs_interrupted``). + + Called by ``run_one_job`` BEFORE it decides what to deliver — a job + whose tool subprocess was killed mid-flight may still produce a + plausible-looking ``final_response`` from the truncated output, and + that must not go out to the user as if it were a normal result. + Unlike ``_consume_interrupted_flag`` below, this does not clear the + flag: the later, authoritative check (right before ``last_status`` is + written) still needs to see it.""" + with _running_lock: + return job_id in _interrupted_job_ids + + def _consume_interrupted_flag(job_id: str) -> bool: """Return True and clear the flag if the shutdown path already marked ``job_id`` interrupted (see ``mark_running_jobs_interrupted``). @@ -3420,6 +3435,20 @@ def run_one_job(job: dict, *, adapters=None, loop=None, verbose: bool = False) - if verbose: logger.info("Output saved to: %s", output_file) + # If the gateway shutdown killed this job's tool subprocess + # mid-flight (#60432), the agent may still have produced a + # plausible-looking final_response from the truncated output -- + # force the failure path so the delivered message is an honest + # "this run was interrupted" summary instead of that response. + # Peek-only: the flag stays set for the authoritative check + # right before mark_job_run below. + if success and _is_interrupted(job["id"]): + success = False + error = ( + "Interrupted by gateway shutdown before the run finished " + "(tool subprocess was killed mid-flight)." + ) + # Deliver the final response to the origin/target chat. # If the agent responded with [SILENT], skip delivery (but # output is already saved above). Failed jobs always deliver. diff --git a/tests/cron/test_shutdown_interrupt.py b/tests/cron/test_shutdown_interrupt.py index 2b5077fc9d277..a95567ff88b49 100644 --- a/tests/cron/test_shutdown_interrupt.py +++ b/tests/cron/test_shutdown_interrupt.py @@ -113,6 +113,35 @@ def _side_effect(job_id, success, reason, **kwargs): assert marked == ["job-2"] +class TestIsInterrupted: + """Peek-only check used at the delivery gate -- must NOT clear the + flag, unlike _consume_interrupted_flag.""" + + def test_false_when_not_marked(self): + import cron.scheduler as sched + + assert sched._is_interrupted("job-1") is False + + def test_true_when_marked(self): + import cron.scheduler as sched + + sched._interrupted_job_ids.add("job-1") + + assert sched._is_interrupted("job-1") is True + + def test_does_not_clear_the_flag(self): + import cron.scheduler as sched + + sched._interrupted_job_ids.add("job-1") + + sched._is_interrupted("job-1") + + # Still set -- the later, authoritative check before mark_job_run + # must still see it. + assert "job-1" in sched._interrupted_job_ids + assert sched._is_interrupted("job-1") is True + + class TestConsumeInterruptedFlag: def test_false_when_not_marked(self): import cron.scheduler as sched @@ -165,6 +194,46 @@ def test_success_path_skipped_when_interrupted(self): # isn't permanently silenced. assert job["id"] not in sched._interrupted_job_ids + def test_interrupted_job_delivers_failure_summary_not_raw_response(self): + """The status-write guard alone isn't enough: delivery happens + BEFORE mark_job_run in run_one_job's own flow, so a job that kept + running post-kill and produced a plausible-looking final_response + must not have that response sent to the user just because the + eventual status write gets suppressed. Interrupted jobs must route + through the same failure-summary delivery path a real failure + would.""" + import cron.scheduler as sched + + job = self._make_job() + sched._interrupted_job_ids.add(job["id"]) + + with patch("cron.scheduler.claim_dispatch", return_value=True), \ + patch("agent.secret_scope.set_secret_scope", return_value=None), \ + patch("agent.secret_scope.build_profile_secret_scope", return_value=None), \ + patch("agent.secret_scope.reset_secret_scope"), \ + patch( + "cron.scheduler.run_job", + return_value=(True, "full output", "a plausible final response", None), + ), \ + patch("cron.scheduler.save_job_output", return_value="/tmp/out.md"), \ + patch( + "cron.scheduler._summarize_cron_failure_for_delivery", + return_value="This run was interrupted.", + ) as mock_summarize, \ + patch("cron.scheduler._is_cron_silence_response", return_value=False), \ + patch("cron.scheduler._deliver_result", return_value=None) as mock_deliver, \ + patch("cron.scheduler.mark_job_run"): + result = sched.run_one_job(job) + + assert result is True + mock_summarize.assert_called_once() + # The summarizer's error argument must mention the interruption, + # not be silently None / the agent's own (possibly absent) error. + assert "interrupt" in mock_summarize.call_args.args[1].lower() + delivered_content = mock_deliver.call_args.args[1] + assert delivered_content == "This run was interrupted." + assert "plausible final response" not in delivered_content + def test_success_path_writes_normally_when_not_interrupted(self): """Control case: the guard must not swallow ordinary, un-interrupted completions -- only ones the shutdown path explicitly flagged.""" From 5fb1ea2e5b2fe36c16f1740170cde8e3e249b68c Mon Sep 17 00:00:00 2001 From: teknium1 <127238744+teknium1@users.noreply.github.com> Date: Tue, 7 Jul 2026 21:34:03 -0700 Subject: [PATCH 5/5] fix(gateway,cron): reconcile #60612 + #60631 onto one drain surface Keep #60631's get_running_job_ids() snapshot + _active_cron_job_count() (import-guarded for minimal test doubles) as the single read path, and retarget #60612's drain tests at it. Drops the redundant cron_jobs_in_flight() helper so there is one surface, not two. --- cron/scheduler.py | 7 ------- tests/gateway/test_update_cron_drain.py | 10 +++++----- 2 files changed, 5 insertions(+), 12 deletions(-) diff --git a/cron/scheduler.py b/cron/scheduler.py index 6b96fb2e1bea8..9e645d3728f3d 100644 --- a/cron/scheduler.py +++ b/cron/scheduler.py @@ -298,13 +298,6 @@ def _is_token(line: str) -> bool: _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) - - # Job IDs the gateway shutdown path force-killed the tool subprocess of # while still in ``_running_job_ids`` (see ``mark_running_jobs_interrupted`` # below). ``run_one_job``'s own completion path checks this set before diff --git a/tests/gateway/test_update_cron_drain.py b/tests/gateway/test_update_cron_drain.py index ee34d72cd82ee..9fb03675e5b86 100644 --- a/tests/gateway/test_update_cron_drain.py +++ b/tests/gateway/test_update_cron_drain.py @@ -22,13 +22,13 @@ async def test_drain_active_agents_waits_for_in_flight_cron_jobs(): cron_count = [1] def _cron_in_flight(): - return cron_count[0] + return frozenset(f"job-{i}" for i in range(cron_count[0])) async def finish_cron(): await asyncio.sleep(0.15) cron_count[0] = 0 - with patch("cron.scheduler.cron_jobs_in_flight", side_effect=_cron_in_flight): + with patch("cron.scheduler.get_running_job_ids", side_effect=_cron_in_flight): task = asyncio.create_task(finish_cron()) _snapshot, timed_out = await runner._drain_active_agents(1.0) await task @@ -42,7 +42,7 @@ async def test_drain_active_agents_times_out_when_cron_still_running(): runner, _adapter = make_restart_runner() runner._running_agents = {} - with patch("cron.scheduler.cron_jobs_in_flight", return_value=1): + with patch("cron.scheduler.get_running_job_ids", return_value=frozenset({"job-1"})): _snapshot, timed_out = await runner._drain_active_agents(0.05) assert timed_out is True @@ -59,7 +59,7 @@ async def test_gateway_stop_waits_for_cron_before_final_tool_kill(): call_order: list[str] = [] def _cron_in_flight(): - return cron_count[0] + return frozenset(f"job-{i}" for i in range(cron_count[0])) def _fake_kill_all(task_id=None): call_order.append("kill_all") @@ -70,7 +70,7 @@ async def finish_cron(): cron_count[0] = 0 with ( - patch("cron.scheduler.cron_jobs_in_flight", side_effect=_cron_in_flight), + patch("cron.scheduler.get_running_job_ids", side_effect=_cron_in_flight), patch("gateway.status.remove_pid_file"), patch("gateway.status.write_runtime_status"), patch("agent.auxiliary_client.shutdown_cached_clients"),