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