mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-07-08 13:12:08 +00:00
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:
parent
e6077af279
commit
24e9ed73c2
5 changed files with 538 additions and 13 deletions
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
209
tests/cron/test_shutdown_interrupt.py
Normal file
209
tests/cron/test_shutdown_interrupt.py
Normal 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()
|
||||
185
tests/gateway/test_cron_active_work_drain.py
Normal file
185
tests/gateway/test_cron_active_work_drain.py
Normal 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()
|
||||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue