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.
This commit is contained in:
teknium1 2026-07-07 21:34:03 -07:00 committed by Teknium
parent 8a573bb6e7
commit ecc6725855
2 changed files with 5 additions and 12 deletions

View file

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

View file

@ -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"),