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
This commit is contained in:
joaomarcos 2026-07-07 23:05:26 -03:00 committed by Teknium
parent e6077af279
commit 24e9ed73c2
5 changed files with 538 additions and 13 deletions

View file

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

View file

@ -4080,6 +4080,26 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
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 @@ 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_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 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
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 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
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 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
)
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 @@ 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()
_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 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
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 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
"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

View file

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

View file

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

View file

@ -54,6 +54,12 @@ class _FakeGateway:
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