fix: audit follow-ups for the next-free-GPU dispatcher

Lifecycle:
- Queue clears and deletes leave a run whose master prompt is currently
  executing alone, as ComfyUI leaves the running prompt alone; previously an
  executing orchestrator-only head run could be cancelled out from under its
  waiting Collector.
- The sweep now also drops worker-assigned runs whose queued delegate prompt
  left the queue without the Collector finishing, releasing their result
  queues instead of leaking them and polling forever.
- Completed runs release their result queues, and orchestration no longer
  pre-creates result queues when load balance picks the master alone.

Transport:
- A websocket timeout during connect, before anything was sent, is a plain
  failure the dispatcher may retry elsewhere; only timeouts after the prompt
  was sent remain AmbiguousDispatchError.

Efficiency and cleanup:
- One WorkerSnapshot per worker per cycle answers every run's "is my job
  there" question (queue first, history only when needed).
- Sweep and cancellation snapshot queued/running prompt ids once per pass.
- probe_worker_queues reuses dispatch._probe_worker_queue; the dead
  aggregate idle/busy probe helper is gone.

Rebased on "never replay load-balanced runs on missing history": an
unresolved submission stays on hold until delivery or positive evidence,
so there is no rejection path to fix and no bounded hold.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

api/orchestration/dispatch.py

@@ -98,6 +98,12 @@ async def _dispatch_via_websocket(worker_url, payload, client_id, timeout=60.0):
     try:
         async with asyncio.timeout_at(deadline):
             ws = await session.ws_connect(ws_url, heartbeat=20)
+    except asyncio.TimeoutError as exc:
+        # Nothing was sent: a definite failure the caller may retry elsewhere.
+        raise RuntimeError("Worker websocket connect timed out before dispatch; nothing was sent") from exc
+
+    try:
+        async with asyncio.timeout_at(deadline):
             await ws.send_json(ws_payload)
             while True:
                 msg = await ws.receive()

api/orchestration/load_balance_queue.py

@@ -14,9 +14,9 @@ A delegate prompt blocks ComfyUI's single executor while it waits for its
 worker, which would stall every full workflow queued behind it on a
 participating master. So on a participating master the delegate prompt is
 queued only once the worker's final result has arrived (or the worker is
-seen idle or unreachable without delivering, so the Collector can report
-it); an orchestrator-only master has nothing else to run and queues it at
-once.
+seen to have finished or lost the run without delivering, so the Collector
+can report it); an orchestrator-only master has nothing else to run and
+queues it at once.
 
 The claim is atomic through ComfyUI's queue mutex: a run can be removed from
 the queue by the dispatcher or popped by the master executor, never both.
@@ -37,7 +37,8 @@ Two lifecycle rules keep a run from executing twice or vanishing:
   run itself. Work that was not yet submitted is dropped along with its
   result queues; work a worker already accepted stays tracked until its
   result arrives or the worker is seen done, and is then discarded instead
-  of reappearing.
+  of reappearing. A run whose master prompt is currently executing is left
+  alone by clears and deletes, exactly as ComfyUI leaves the running prompt.
 """
 import asyncio
 import threading
@@ -47,8 +48,8 @@ from collections import OrderedDict
 import server
 
 from ...utils.logging import debug_log, log
-from ...utils.network import build_worker_url, probe_worker
-from .dispatch import _extract_queue_remaining
+from ...utils.network import build_worker_url
+from .dispatch import _probe_worker_queue
 from .dispatch import _select_idle_round_robin
 
 try:
@@ -74,7 +75,6 @@ REQUEUE_RETRY_MIN_SECONDS = 1.0
 REQUEUE_RETRY_MAX_SECONDS = 30.0
 # Consecutive failed probes before a dispatched worker counts as gone.
 WORKER_UNREACHABLE_PROBES = 3
-
 # How far back the worker's history is searched for an unresolved run.
 HISTORY_LOOKBACK_ITEMS = 200
 # Job ids whose result queues were discarded; a late worker must not recreate them.
@@ -90,19 +90,9 @@ def _worker_key(worker):
 async def probe_worker_queues(workers, probe_concurrency=8, probe_timeout=3.0):
     """Return ``[{"worker", "queue_remaining"}]`` for every worker that answered."""
     semaphore = asyncio.Semaphore(max(1, int(probe_concurrency or 1)))
-
-    async def _probe(worker):
-        async with semaphore:
-            if worker.get("type") == "serverless":
-                from .serverless_runtime import get_manager
-                status = await get_manager().health(worker)
-                return {"worker": worker, "queue_remaining": status["queueRemaining"]}
-            payload = await probe_worker(build_worker_url(worker), timeout=probe_timeout)
-            if payload is None:
-                return None
-            return {"worker": worker, "queue_remaining": _extract_queue_remaining(payload)}
-
-    statuses = await asyncio.gather(*[_probe(worker) for worker in workers])
+    statuses = await asyncio.gather(
+        *[_probe_worker_queue(worker, semaphore, probe_timeout) for worker in workers]
+    )
     return [status for status in statuses if status is not None]
 
 
@@ -118,51 +108,87 @@ def _graph_has_job(graph, job_ids):
     return False
 
 
-async def fetch_worker_job_state(worker, job_ids, timeout=5.0):
-    """Ask the worker itself whether a specific run is there.
+class WorkerSnapshot:
+    """What one worker says about the runs it holds, fetched at most once.
 
-    Returns ``"running"`` when a prompt carrying one of ``job_ids`` is in the
-    worker's running or pending queue, ``"finished"`` when it is in the
-    worker's recent history, ``"absent"`` when neither snapshot contains it,
-    and ``None`` when the worker could not be asked. Absence is not proof that
-    a submission was rejected: history can be truncated, cleared or lost on
-    restart. Direct workers only; serverless pools never reach this path.
+    The queue (running and pending prompts) is fetched on first use; the
+    history, which is larger, only when a run is not found in the queue.
+    One snapshot serves every run on that worker within a dispatcher cycle.
+    Direct workers only; serverless pools never reach this path.
     """
-    import aiohttp
-    from ...utils.network import get_client_session
 
-    job_ids = {str(job_id) for job_id in job_ids}
-    base_url = build_worker_url(worker)
-    session = await get_client_session()
-    try:
-        async with session.get(f"{base_url}/queue", timeout=aiohttp.ClientTimeout(total=timeout)) as response:
-            if response.status != 200:
-                return None
-            queue_info = await response.json()
-        if not isinstance(queue_info, dict):
-            return None
-        for key in ("queue_running", "queue_pending"):
-            for item in queue_info.get(key) or []:
-                if isinstance(item, (list, tuple)) and len(item) > 2 and _graph_has_job(item[2], job_ids):
-                    return "running"
+    def __init__(self, worker, timeout=5.0):
+        self.worker = worker
+        self.timeout = timeout
+        self._queue_graphs = None
+        self._history_graphs = None
+        self.failed = False
+
+    async def _fetch(self, endpoint, params=None):
+        import aiohttp
+        from ...utils.network import get_client_session
+
+        session = await get_client_session()
         async with session.get(
-            f"{base_url}/history",
-            params={"max_items": str(HISTORY_LOOKBACK_ITEMS)},
-            timeout=aiohttp.ClientTimeout(total=timeout),
+            f"{build_worker_url(self.worker)}{endpoint}",
+            params=params,
+            timeout=aiohttp.ClientTimeout(total=self.timeout),
         ) as response:
             if response.status != 200:
-                return None
-            history = await response.json()
-        if not isinstance(history, dict):
+                raise RuntimeError(f"{endpoint} returned HTTP {response.status}")
+            payload = await response.json()
+        if not isinstance(payload, dict):
+            raise RuntimeError(f"{endpoint} returned a non-object body")
+        return payload
+
+    async def _queue(self):
+        if self._queue_graphs is None:
+            queue_info = await self._fetch("/queue")
+            graphs = []
+            for key in ("queue_running", "queue_pending"):
+                for item in queue_info.get(key) or []:
+                    if isinstance(item, (list, tuple)) and len(item) > 2:
+                        graphs.append(item[2])
+            self._queue_graphs = graphs
+        return self._queue_graphs
+
+    async def _history(self):
+        if self._history_graphs is None:
+            history = await self._fetch("/history", params={"max_items": str(HISTORY_LOOKBACK_ITEMS)})
+            graphs = []
+            for entry in history.values():
+                item = entry.get("prompt") if isinstance(entry, dict) else None
+                if isinstance(item, (list, tuple)) and len(item) > 2:
+                    graphs.append(item[2])
+            self._history_graphs = graphs
+        return self._history_graphs
+
+    async def job_state(self, job_ids):
+        """``"running"`` (in the worker's queue), ``"finished"`` (in its recent
+        history), ``"absent"`` (in neither), or ``None`` when the worker could
+        not be asked.
+
+        Absence is not proof that a submission was rejected: history can be
+        truncated, cleared or lost on restart.
+        """
+        if self.failed:
             return None
-        for entry in history.values():
-            item = entry.get("prompt") if isinstance(entry, dict) else None
-            if isinstance(item, (list, tuple)) and len(item) > 2 and _graph_has_job(item[2], job_ids):
+        job_ids = {str(job_id) for job_id in job_ids}
+        try:
+            if any(_graph_has_job(graph, job_ids) for graph in await self._queue()):
+                return "running"
+            if any(_graph_has_job(graph, job_ids) for graph in await self._history()):
                 return "finished"
-        return "absent"
-    except Exception as exc:
-        debug_log(f"Load balance: could not check worker {_worker_key(worker)} for a run: {exc}")
-        return None
+            return "absent"
+        except Exception as exc:
+            self.failed = True
+            debug_log(f"Load balance: could not check worker {_worker_key(self.worker)} for a run: {exc}")
+            return None
+
+
+async def fetch_worker_job_state(worker, job_ids, timeout=5.0):
+    """Ask the worker itself whether a specific run is there (see ``WorkerSnapshot``)."""
+    return await WorkerSnapshot(worker, timeout).job_state(job_ids)
 
 
 async def select_idle_worker(workers, trace_execution_id=None, probe_concurrency=8, probe_timeout=3.0):
@@ -260,6 +286,11 @@ class PendingRun:
         """Out of the master queue for a hand-off: cancellation must be recorded here."""
         return self.claiming or self.on_worker
 
+    @property
+    def collector_unqueued(self):
+        """Out of the master queue on purpose: being handed off, held unresolved, or awaiting its delegate."""
+        return self.claiming or self.unresolved is not None or self.requeue_pending
+
     def accepts(self, worker):
         key = _worker_key(worker)
         return any(_worker_key(candidate) == key for candidate in self.candidates)
@@ -275,6 +306,7 @@ class LoadBalanceQueue:
         self._by_job = {}  # job_id -> PendingRun
         self._recent_dispatch = {}  # worker id -> monotonic time
         self._cancelled_jobs = OrderedDict()
+        self._snapshots = {}  # worker id -> WorkerSnapshot, valid for one cycle
         self._wake = asyncio.Event()
         self._cycle_lock = asyncio.Lock()
         self._task = None
@@ -377,13 +409,17 @@ class LoadBalanceQueue:
             if self._by_job.get(job_id) is record:
                 self._by_job.pop(job_id, None)
 
+    def _release_result_queues(self, record):
+        pending_jobs = self._pending_jobs_getter()
+        if isinstance(pending_jobs, dict):
+            for job_id in record.job_ids:
+                pending_jobs.pop(job_id, None)
+
     def _discard(self, record):
         """Drop the run and the result queues nobody will consume any more."""
         self._remove(record)
-        pending_jobs = self._pending_jobs_getter()
+        self._release_result_queues(record)
         for job_id in record.job_ids:
-            if isinstance(pending_jobs, dict):
-                pending_jobs.pop(job_id, None)
             self._cancelled_jobs[job_id] = True
             self._cancelled_jobs.move_to_end(job_id)
         while len(self._cancelled_jobs) > CANCELLED_JOB_MEMORY:
@@ -430,10 +466,15 @@ class LoadBalanceQueue:
         self._hooks_installed = True
 
     def cancel_matching(self, predicate):
-        """Cancel in-flight runs whose queue item ``predicate`` would have deleted."""
+        """Cancel in-flight runs whose queue item ``predicate`` would have deleted.
+
+        A run whose master prompt is currently executing is left alone, as
+        ComfyUI's clear and delete leave the running prompt alone.
+        """
         cancelled = []
+        running = self._running_ids()
         for record in list(self._records.values()):
-            if record.cancelled or not record.in_flight:
+            if record.cancelled or not record.in_flight or record.prompt_id in running:
                 continue
             item = record.unresolved_item or (0, record.prompt_id, {}, {}, [])
             try:
@@ -503,6 +544,8 @@ class LoadBalanceQueue:
         On failure or interruption the run is cancelled: a hand-off still in
         flight is stopped before submission, and work a worker already
         accepted is discarded when it ends, so no result queue is left behind.
+        On success the run is forgotten once every job finished, including the
+        result queues a master-run never needed.
         """
         if failed:
             if record.prompt_id in self._records:
@@ -511,20 +554,27 @@ class LoadBalanceQueue:
         record.finished_jobs.add(str(job_id))
         if record.finished_jobs >= record.job_ids:
             self._remove(record)
+            self._release_result_queues(record)
 
     # ----- queue inspection ---------------------------------------------
     def _prompt_queue(self):
         return self._prompt_queue_getter()
 
-    def _is_running(self, prompt_id):
+    def _running_ids(self):
         prompt_queue = self._prompt_queue()
         with prompt_queue.mutex:
-            return any(item[1] == prompt_id for item in prompt_queue.currently_running.values())
+            return {item[1] for item in prompt_queue.currently_running.values()}
+
+    def _queued_ids(self):
+        prompt_queue = self._prompt_queue()
+        with prompt_queue.mutex:
+            return {item[1] for item in prompt_queue.queue}
+
+    def _is_running(self, prompt_id):
+        return prompt_id in self._running_ids()
 
     def _is_queued(self, prompt_id):
-        prompt_queue = self._prompt_queue()
-        with prompt_queue.mutex:
-            return any(item[1] == prompt_id for item in prompt_queue.queue)
+        return prompt_id in self._queued_ids()
 
     def _take_from_queue(self, prompt_id):
         """Remove the prompt from the master queue; return the item or None."""
@@ -579,13 +629,20 @@ class LoadBalanceQueue:
             self._task = None
 
     def sweep(self):
-        """Drop runs that left the master queue without being collected."""
+        """Drop runs that left the master queue without being collected.
+
+        Runs that are out of the queue on purpose (being handed off, held
+        unresolved, or waiting for their delegate prompt to be queued) and
+        cancelled runs are left to the dispatcher's own bookkeeping.
+        """
+        queued = self._queued_ids()
+        running = self._running_ids()
         for record in list(self._records.values()):
-            if not record.queued or record.in_flight or record.cancelled:
+            if not record.queued or record.collector_unqueued or record.cancelled:
                 continue
-            if self._is_queued(record.prompt_id):
+            if record.prompt_id in queued:
                 continue
-            if self._is_running(record.prompt_id):
+            if record.prompt_id in running:
                 if record.assigned is None and record.master_may_run:
                     record._assign(MASTER)
                     debug_log(f"Load balance: master is running run {record.prompt_id}")
@@ -605,7 +662,11 @@ class LoadBalanceQueue:
     async def run_cycle(self, probe_concurrency=8, probe_timeout=3.0):
         """One dispatcher pass: sweep, resolve, retry, then hand runs to idle workers."""
         async with self._cycle_lock:
-            return await self._run_cycle(probe_concurrency, probe_timeout)
+            self._snapshots = {}
+            try:
+                return await self._run_cycle(probe_concurrency, probe_timeout)
+            finally:
+                self._snapshots = {}
 
     async def _run_cycle(self, probe_concurrency, probe_timeout):
         self.sweep()
@@ -662,14 +723,14 @@ class LoadBalanceQueue:
         record.delegate_needed = False
         return True
 
-    async def _probe_one(self, record, worker):
-        """Return 'unreachable', 'idle' or 'busy' for the worker running ``record``."""
-        statuses = await probe_worker_queues([worker])
-        if not statuses:
-            record.missed_probes += 1
-            return "unreachable"
-        record.missed_probes = 0
-        return "idle" if statuses[0]["queue_remaining"] == 0 else "busy"
+    async def _job_state(self, record, worker):
+        """Ask ``worker`` about ``record``'s jobs, sharing one snapshot per worker per cycle."""
+        key = _worker_key(worker)
+        snapshot = self._snapshots.get(key)
+        if snapshot is None:
+            snapshot = WorkerSnapshot(worker)
+            self._snapshots[key] = snapshot
+        return await snapshot.job_state(record.job_ids)
 
     async def _worker_done_without_result(self, record):
         """True when the worker no longer has this run and never delivered.
@@ -680,7 +741,7 @@ class LoadBalanceQueue:
         """
         if time.monotonic() - record.dispatched_at < RECENT_DISPATCH_TTL_SECONDS:
             return False
-        state = await fetch_worker_job_state(record.assigned, record.job_ids)
+        state = await self._job_state(record, record.assigned)
         if state is None:
             record.missed_probes += 1
             return record.missed_probes >= WORKER_UNREACHABLE_PROBES
@@ -695,7 +756,9 @@ class LoadBalanceQueue:
         an idle worker may have run the job and lost the delivery, a busy one
         may be on unrelated work. Missing history is not rejection evidence,
         even on repeated checks: accepted work may have aged out or been cleared.
+        A worker that cannot be asked keeps the run on hold as well.
         """
+        now = time.monotonic()
         for record in list(self._records.values()):
             worker = record.unresolved
             if worker is None:
@@ -703,9 +766,9 @@ class LoadBalanceQueue:
             if record.delivered:
                 self._accept_unresolved(record, "delivered its result")
                 continue
-            if time.monotonic() - record.unresolved_since < RECENT_DISPATCH_TTL_SECONDS:
+            if now - record.unresolved_since < RECENT_DISPATCH_TTL_SECONDS:
                 continue
-            state = await fetch_worker_job_state(worker, record.job_ids)
+            state = await self._job_state(record, worker)
             if state == "running":
                 self._accept_unresolved(record, "has it queued or running")
             elif state == "finished":
@@ -764,8 +827,8 @@ class LoadBalanceQueue:
                 if not await self._worker_done_without_result(record):
                     continue
                 log(
-                    f"[Distributed] Load balance: worker {record.assigned_id} is idle or unreachable without "
-                    f"delivering run {record.prompt_id[:8]}; queueing its collector so the run can report."
+                    f"[Distributed] Load balance: worker {record.assigned_id} no longer has run "
+                    f"{record.prompt_id[:8]} and never delivered it; queueing its collector so the run can report."
                 )
                 record.wait_for_result = False
             if record.requeue_retry_at > now:

api/queue_orchestration.py

@@ -566,6 +566,9 @@ async def orchestrate_distributed_execution(
     # collect-only prompt once the worker delivers, not now, so the master
     # executor stays free for the next full workflow.
     hold_delegate_for_worker = None
+    # Load balance picked the master alone: no worker will ever deliver, so
+    # no result queues are needed for this run.
+    master_runs_alone = False
     if load_balance_requested and not image_list_mode:
         candidate_workers = list(active_workers)
         if not delegate_master:
@@ -617,6 +620,7 @@ async def orchestrate_distributed_execution(
             # Master selected as least busy; run master workload only.
             active_workers = []
             delegate_master = False
+            master_runs_alone = True
             trace_debug(
                 execution_trace_id,
                 "Load-balance selected master for execution (workers skipped).",
@@ -711,7 +715,7 @@ async def orchestrate_distributed_execution(
         manager.ensure_capacity(len(serverless_jobs))
 
     # Do not leave collector queues behind when offline compilation/media validation fails.
-    if not image_list_mode:
+    if not image_list_mode and not master_runs_alone:
         for job_id in job_id_map.values():
             await _ensure_distributed_queue(job_id)
 

tests/test_dispatch_selection.py

@@ -196,6 +196,19 @@ class ExplicitTargetTransportTests(unittest.IsolatedAsyncioTestCase):
         self.assertEqual(calls[0]['partial_execution_targets'], ['9'])
         self.assertNotIn('partial_execution_targets', calls[1])
 
+    async def test_websocket_connect_timeout_is_a_definite_failure(self):
+        """A timeout before anything was sent must not read as ambiguous acceptance."""
+        never = asyncio.Event()
+
+        async def hang(*_args, **_kwargs):
+            await never.wait()
+
+        session = types.SimpleNamespace(ws_connect=hang)
+        with patch.object(dispatch, "get_client_session", AsyncMock(return_value=session)):
+            with self.assertRaisesRegex(RuntimeError, "nothing was sent") as caught:
+                await dispatch._dispatch_via_websocket("http://worker", {"prompt": {}}, "client", timeout=0.05)
+        self.assertNotIsInstance(caught.exception, dispatch.AmbiguousDispatchError)
+
     async def test_websocket_sends_targets_all_the_way_to_wire(self):
         import json
         sent = []

tests/test_load_balance_orchestration.py

@@ -188,6 +188,14 @@ class DeferredLoadBalanceTests(unittest.IsolatedAsyncioTestCase):
         self.assertEqual(self.lb_queue.records[0].unresolved["id"], "w1")
         self.q.queue_prompt_payload.assert_not_awaited()
 
+    async def test_master_selected_alone_creates_no_result_queues(self):
+        self.select_idle_worker.return_value = ({"id": "master", "name": "Master", "type": "local"}, 2)
+        _, _, worker_count, _ = await self._run()
+        self.assertEqual(worker_count, 0)
+        self.q.queue_prompt_payload.assert_awaited_once()
+        self.assertEqual(self.ps.distributed_pending_jobs, {})
+        self.assertEqual(self.lb_queue.records, [])
+
     async def test_orchestrator_only_master_queues_delegate_immediately(self):
         self.settings["master_delegate_only"] = True
         self.select_idle_worker.return_value = (self.worker, 1)
@@ -223,18 +231,17 @@ class DeferredLoadBalanceTests(unittest.IsolatedAsyncioTestCase):
         network.build_worker_url = lambda worker, *a: f"http://{worker['id']}"
         network.probe_worker = AsyncMock(return_value=None)
         dispatch = sys.modules[PKG + ".api.orchestration.dispatch"]
-        dispatch._extract_queue_remaining = lambda payload: int(payload["exec_info"]["queue_remaining"])
+
+        async def always_busy_probe(worker, semaphore, probe_timeout):
+            return {"worker": worker, "queue_remaining": 1}
+
+        dispatch._probe_worker_queue = always_busy_probe
         dispatch._select_idle_round_robin = lambda statuses: statuses[0]
         spec = importlib.util.spec_from_file_location(
             PKG + ".api.orchestration.load_balance_queue", ROOT / "api" / "orchestration" / "load_balance_queue.py")
         lb = importlib.util.module_from_spec(spec)
         spec.loader.exec_module(lb)
         lb.select_idle_worker = AsyncMock(return_value=(None, 1))
-
-        async def always_busy(url, timeout=3.0):
-            return {"exec_info": {"queue_remaining": 1}}
-
-        lb.probe_worker = always_busy
         sys.modules[PKG + ".api.orchestration.load_balance_queue"] = lb
         queue = lb.get_load_balance_queue(self.ps)
         queue.poll_interval = 0.01

tests/test_load_balance_queue.py

@@ -12,6 +12,8 @@ from unittest.mock import AsyncMock, patch
 
 ROOT = Path(__file__).resolve().parents[1]
 PKG = "dist_lb_queue_testpkg"
+# The running test case; the dispatch stub's probe reads its ``remaining`` table.
+_ACTIVE = {"case": None}
 
 
 def _module(name, **attrs):
@@ -84,9 +86,15 @@ def _load_module(prompt_queue):
     sys.modules[PKG + ".utils.network"] = _module(
         "network", build_worker_url=lambda worker, *a: f"http://{worker['id']}", probe_worker=AsyncMock(return_value=None)
     )
+    async def _probe_worker_queue(worker, semaphore, probe_timeout):
+        async with semaphore:
+            case = _ACTIVE["case"]
+            value = case.remaining.get(worker["id"]) if case is not None else None
+            return None if value is None else {"worker": worker, "queue_remaining": value}
+
     sys.modules[PKG + ".api.orchestration.dispatch"] = _module(
         "dispatch",
-        _extract_queue_remaining=lambda payload: int(payload.get("exec_info", {}).get("queue_remaining", 0)),
+        _probe_worker_queue=_probe_worker_queue,
         _select_idle_round_robin=lambda statuses: statuses[0],
     )
     spec = importlib.util.spec_from_file_location(
@@ -111,35 +119,37 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase):
         self.dispatched = []
         self.prepared = []
         self.requeued = []
-
-        async def fake_probe(url, timeout=3.0):
-            worker_id = url.rsplit("/", 1)[-1]
-            value = self.remaining.get(worker_id)
-            return None if value is None else {"exec_info": {"queue_remaining": value}}
-
-        self.lb.probe_worker = fake_probe
+        _ACTIVE["case"] = self
         # What a worker answers when asked about a specific run (its job ids):
         # "running" | "finished" | "absent" | None (could not be asked).
         self.job_state = {}
         self.job_state_queries = []
+        self.snapshots_taken = []
+        self.real_snapshot_class = self.lb.WorkerSnapshot
+        case = self
 
-        async def fake_job_state(worker, job_ids, timeout=5.0):
-            self.job_state_queries.append((worker["id"], set(job_ids)))
-            return self.job_state.get(worker["id"], "absent")
+        class FakeSnapshot:
+            def __init__(self, worker, timeout=5.0):
+                self.worker = worker
+                case.snapshots_taken.append(worker["id"])
 
-        self.fetch_worker_job_state = self.lb.fetch_worker_job_state
-        self.lb.fetch_worker_job_state = fake_job_state
+            async def job_state(self, job_ids):
+                case.job_state_queries.append((self.worker["id"], set(job_ids)))
+                return case.job_state.get(self.worker["id"], "absent")
+
+        self.lb.WorkerSnapshot = FakeSnapshot
 
     def _queue_run(self, prompt_id, master_may_run=True, dispatch_error=None, requeue_errors=0, queued=True,
                    prepare_error=None, prepare_gate=None, ambiguous=False, requeue_gate=None):
         """Register a run. ``prepare_gate`` (an Event) delays preparation until set;
-        ``ambiguous`` makes the send raise AmbiguousDispatchError after the worker
-        has in fact accepted the job (its acknowledgement was lost)."""
+        ``ambiguous`` makes the first send raise AmbiguousDispatchError after the
+        worker has in fact accepted the job (its acknowledgement was lost)."""
         if queued:
             number = len(self.prompt_queue.queue) + len(self.prompt_queue.currently_running)
             self.prompt_queue.put((number, prompt_id, {"full": prompt_id}, {}, [], {}))
         self.pending_jobs[f"job-{prompt_id}"] = asyncio.Queue()
         remaining_requeue_errors = [requeue_errors]
+        ambiguous_sends_left = [1 if ambiguous else 0]
 
         async def prepare_for_worker(worker):
             if prepare_gate is not None:
@@ -154,7 +164,8 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase):
             if dispatch_error is not None:
                 raise dispatch_error
             self.dispatched.append((prompt_id, worker["id"]))
-            if ambiguous:
+            if ambiguous_sends_left[0] > 0:
+                ambiguous_sends_left[0] -= 1
                 raise self.lb.AmbiguousDispatchError("acknowledgement lost")
 
         async def requeue_as_delegate(worker):
@@ -193,7 +204,7 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase):
         self.assertEqual(record.unresolved["id"], "w1")
         self.assertIsNone(record.assigned)
         self.assertEqual(self.prompt_queue.prompt_ids(), [])  # not restored: the master cannot pick it up
-        record.unresolved_since = 0  # past the grace period
+        record.unresolved_since = time.monotonic() - self.lb.RECENT_DISPATCH_TTL_SECONDS - 1  # past the grace period
         self.queue._recent_dispatch.clear()
         return record
 
@@ -275,7 +286,7 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase):
                 raise AssertionError(f"Unexpected request: {url}")
 
         network = sys.modules[PKG + ".utils.network"]
-        self.lb.fetch_worker_job_state = self.fetch_worker_job_state
+        self.lb.WorkerSnapshot = self.real_snapshot_class
         with patch.object(network, "get_client_session", AsyncMock(return_value=Session()), create=True):
             for _ in range(5):
                 self.queue._recent_dispatch.clear()
@@ -523,6 +534,113 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase):
         self.assertNotIn("job-p1", self.pending_jobs)
         self.assertEqual(self.prompt_queue.prompt_ids(), [])
 
+    # ----- audit regressions ---------------------------------------------------
+    async def test_clear_queue_leaves_an_executing_head_run_alone(self):
+        """ComfyUI's clear removes pending prompts only; the running one, whose Collector
+        is waiting for a worker, keeps its hand-off."""
+        gate = asyncio.Event()
+        record = self._queue_run("p1", master_may_run=False, prepare_gate=gate)
+        self.prompt_queue.start_next()
+        waiter = asyncio.ensure_future(self.queue.wait_for_assignment(record))
+        await asyncio.sleep(0)
+        self.remaining["w1"] = 0
+        cycle = asyncio.ensure_future(self.queue.run_cycle())
+        for _ in range(200):
+            if record.claiming:
+                break
+            await asyncio.sleep(0.005)
+        self.assertTrue(record.claiming)
+        self.prompt_queue.wipe_queue()
+        self.prompt_queue.delete_queue_item(lambda item: item[1] == "p1")
+        self.assertFalse(record.cancelled)
+        gate.set()
+        await cycle
+        self.assertEqual(self.dispatched, [("p1", "w1")])
+        assignment = await asyncio.wait_for(waiter, timeout=2)
+        self.assertEqual(assignment["id"], "w1")
+        self.assertIs(self.queue.lookup_job("job-p1"), record)
+
+    async def test_sweep_drops_assigned_run_whose_delegate_left_the_queue(self):
+        record = self._queue_run("p1", master_may_run=False)
+        self.remaining["w1"] = 0
+        await self.queue.run_cycle()
+        self.assertEqual(self.prompt_queue.prompt_ids(), ["p1"])  # delegate queued
+        self.assertEqual(record.assigned_id, "w1")
+        # The delegate prompt is gone without its Collector ever finishing
+        # (an execution error before the Collector node, for example).
+        self.prompt_queue.queue = []
+        self.queue.sweep()
+        self.assertIsNone(self.queue.lookup_job("job-p1"))
+        self.assertNotIn("job-p1", self.pending_jobs)
+        self.assertEqual(len(self.queue), 0)
+
+    async def test_sweep_keeps_runs_whose_delegate_is_executing_or_still_held(self):
+        running = self._queue_run("p1", master_may_run=False)
+        held = self._queue_run("p2")
+        self.remaining = {"w1": 0, "w2": 0}
+        await self.queue.run_cycle()
+        self.assertEqual(sorted(self.dispatched), [("p1", "w1"), ("p2", "w2")])
+        self.prompt_queue.start_next()  # p1's delegate executes
+        self.queue.sweep()
+        self.assertIs(self.queue.lookup_job("job-p1"), running)
+        self.assertIs(self.queue.lookup_job("job-p2"), held)  # delegate held until delivery
+
+    async def test_master_run_completion_releases_its_result_queues(self):
+        record = self._queue_run("p1")
+        self.prompt_queue.start_next()
+        self.assertEqual(await self.queue.wait_for_assignment(record), self.lb.MASTER)
+        self.assertIn("job-p1", self.pending_jobs)
+        self.queue.finish_job(record, "job-p1")
+        self.assertIsNone(self.queue.lookup_job("job-p1"))
+        self.assertNotIn("job-p1", self.pending_jobs)
+        self.assertFalse(self.queue.is_cancelled_job("job-p1"))  # a normal completion, not a cancel
+
+    async def test_one_worker_snapshot_serves_every_run_in_a_cycle(self):
+        first = self._queue_run("p1")
+        second = self._queue_run("p2")
+        self.remaining["w1"] = 0
+        await self.queue.run_cycle()
+        self.queue._recent_dispatch.clear()
+        await self.queue.run_cycle()
+        self.assertEqual(sorted(self.dispatched), [("p1", "w1"), ("p2", "w1")])
+        first.dispatched_at = second.dispatched_at = 0
+        self.job_state["w1"] = "running"
+        self.snapshots_taken.clear()
+        self.job_state_queries.clear()
+        await self.queue.run_cycle()
+        self.assertEqual(self.snapshots_taken, ["w1"])  # one fetch, two questions
+        self.assertEqual(sorted(query[1] for query in self.job_state_queries), [{"job-p1"}, {"job-p2"}])
+        self.assertEqual(self.requeued, [])
+
+    async def test_worker_snapshot_reads_queue_then_history_lazily(self):
+        calls = []
+
+        class Snapshot(self.real_snapshot_class):
+            async def _fetch(self, endpoint, params=None):
+                calls.append(endpoint)
+                if endpoint == "/queue":
+                    return {"queue_running": [[1, "x", {"n": {"inputs": {"multi_job_id": "job-a"}}}]],
+                            "queue_pending": []}
+                return {"h1": {"prompt": [0, "y", {"n": {"inputs": {"multi_job_id": "job-b"}}}]}}
+
+        snapshot = Snapshot({"id": "w1"})
+        self.assertEqual(await snapshot.job_state({"job-a"}), "running")
+        self.assertEqual(calls, ["/queue"])
+        self.assertEqual(await snapshot.job_state({"job-b"}), "finished")
+        self.assertEqual(calls, ["/queue", "/history"])
+        self.assertEqual(await snapshot.job_state({"job-c"}), "absent")
+        self.assertEqual(calls, ["/queue", "/history"])  # both cached
+
+        class Broken(self.real_snapshot_class):
+            async def _fetch(self, endpoint, params=None):
+                calls.append(endpoint)
+                raise RuntimeError("connection refused")
+
+        broken = Broken({"id": "w2"})
+        self.assertIsNone(await broken.job_state({"job-a"}))
+        self.assertIsNone(await broken.job_state({"job-a"}))
+        self.assertEqual(calls.count("/queue"), 2)  # a failed worker is not asked again this cycle
+
     # ----- interrupted assignment on an orchestrator-only master ---------------
     async def test_interrupt_during_assignment_cleans_result_queue_and_stops_dispatch(self):
         gate = asyncio.Event()