From ecc6725855280b4da7181c12bf930d95c7ae9a8d 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] 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 6b96fb2e1be..9e645d3728f 100644 --- a/cron/scheduler.py +++ b/cron/scheduler.py @@ -298,13 +298,6 @@ _parallel_pool_max_workers: Optional[int] = None _running_job_ids: set = set() _running_lock = threading.Lock() - -def cron_jobs_in_flight() -> int: - """Return how many cron jobs are currently executing agent/tool work.""" - with _running_lock: - return len(_running_job_ids) - - # 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 ee34d72cd82..9fb03675e5b 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 test_gateway_stop_waits_for_cron_before_final_tool_kill(): 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"),