fix: settle unacknowledged submissions by asking the worker about the run itself

Second round of review follow-ups for the next-free-GPU dispatcher.

Ambiguous acceptance is no longer inferred from aggregate queue depth. The
dispatcher asks the worker about the specific run: a prompt carrying one
of its job ids in the worker's running or pending queue, or in its recent
history, means the run was accepted and its result is collected (or
reported) there; only after consecutive checks find no trace of it is the
original prompt restored. A worker that cannot be asked keeps the run on
hold. The same question decides when a held delegate prompt must be queued
because the worker finished (or lost) the run without delivering.

An eagerly dispatched run whose acknowledgement was lost is now collected
after the worker delivers: the record remembers that no master prompt was
ever queued for it, so acceptance queues its collector instead of treating
it like a collector that is already executing.

Cancellation now reaches the queue-insertion boundary: queue_prompt_payload
takes an insert_if guard that is checked under the queue mutex right before
the item is put, after validation has yielded. The delegate requeue passes
the run's cancellation flag, so a clear or delete that lands during
validation refuses the insertion instead of being discovered afterwards.

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

api/orchestration/load_balance_queue.py

@@ -25,9 +25,12 @@ Two lifecycle rules keep a run from executing twice or vanishing:
 
 * A submission whose acceptance is unknown (the transport raised
   ``AmbiguousDispatchError``) is never replayed. The run is held as
-  *unresolved* until the worker delivers, is seen busy (accepted), or is
-  seen idle without delivering (never accepted), and only then is it either
-  collected or put back in the master queue.
+  *unresolved* until the worker delivers, or until the worker's own queue
+  and history are checked for this run's job ids: found there means it was
+  accepted and the result is collected; confirmed absent on consecutive
+  checks means it never arrived and the original prompt goes back. Aggregate
+  queue depth is never used for this, since an idle worker may have run the
+  job and lost the delivery, and a busy worker may be on something else.
 * While a run is out of the master queue for a hand-off, cancellation
   (queue clear, queue delete, interrupting its Collector) is recorded on the
   run itself. Work that was not yet submitted is dropped along with its
@@ -70,6 +73,11 @@ 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
+# An unresolved submission is restored only after this many consecutive
+# worker queue+history checks found no trace of its job ids.
+ABSENT_CONFIRMATIONS = 2
+# 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.
 CANCELLED_JOB_MEMORY = 512
 
@@ -99,6 +107,64 @@ async def probe_worker_queues(workers, probe_concurrency=8, probe_timeout=3.0):
     return [status for status in statuses if status is not None]
 
 
+def _graph_has_job(graph, job_ids):
+    if not isinstance(graph, dict):
+        return False
+    for node in graph.values():
+        if not isinstance(node, dict):
+            continue
+        inputs = node.get("inputs")
+        if isinstance(inputs, dict) and str(inputs.get("multi_job_id", "")) in job_ids:
+            return True
+    return False
+
+
+async def fetch_worker_job_state(worker, job_ids, timeout=5.0):
+    """Ask the worker itself whether a specific run is there.
+
+    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 the worker answered both
+    questions and has no trace of it, and ``None`` when the worker could not
+    be asked. 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"
+        async with session.get(
+            f"{base_url}/history",
+            params={"max_items": str(HISTORY_LOOKBACK_ITEMS)},
+            timeout=aiohttp.ClientTimeout(total=timeout),
+        ) as response:
+            if response.status != 200:
+                return None
+            history = await response.json()
+        if not isinstance(history, dict):
+            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):
+                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
+
+
 async def select_idle_worker(workers, trace_execution_id=None, probe_concurrency=8, probe_timeout=3.0):
     """Pick one idle participant, round-robin; ``(None, probed)`` when all are busy.
 
@@ -149,6 +215,10 @@ class PendingRun:
         self.unresolved = None
         self.unresolved_item = None
         self.unresolved_since = 0.0
+        self.absent_checks = 0
+        # No master prompt exists for this run yet (sent straight to a worker
+        # by a participating master): its collector still has to be queued.
+        self.delegate_needed = False
         # A worker runs the job but the master's delegate prompt is not queued
         # yet; the dispatcher keeps retrying so the result is still collected.
         self.requeue_pending = False
@@ -236,6 +306,7 @@ class LoadBalanceQueue:
         """
         self.register(record)
         record.queued = True
+        record.delegate_needed = True
         self._hold_delegate(record, worker)
         self._ensure_task()
         self.wake()
@@ -253,6 +324,7 @@ class LoadBalanceQueue:
         record.unresolved_item = item
         record.unresolved_since = time.monotonic()
         record.missed_probes = 0
+        record.absent_checks = 0
         self.mark_worker_busy(_worker_key(worker))
         self._ensure_task()
         self.wake()
@@ -571,6 +643,11 @@ class LoadBalanceQueue:
         try:
             await record.requeue_as_delegate(worker)
         except Exception as exc:
+            if record.cancelled or getattr(exc, "cancelled_insertion", False):
+                # Cancelled while the delegate prompt was being validated: the
+                # insertion was refused on purpose, nothing to retry.
+                debug_log(f"Load balance: delegate for run {record.prompt_id} not queued, run was cancelled")
+                return False
             record.requeue_failures += 1
             delay = min(REQUEUE_RETRY_MIN_SECONDS * (2 ** (record.requeue_failures - 1)), REQUEUE_RETRY_MAX_SECONDS)
             record.requeue_retry_at = time.monotonic() + delay
@@ -582,6 +659,7 @@ class LoadBalanceQueue:
             return False
         record.requeue_pending = False
         record.requeue_failures = 0
+        record.delegate_needed = False
         return True
 
     async def _probe_one(self, record, worker):
@@ -594,16 +672,29 @@ class LoadBalanceQueue:
         return "idle" if statuses[0]["queue_remaining"] == 0 else "busy"
 
     async def _worker_done_without_result(self, record):
-        """True when the dispatched worker is idle or gone and never delivered."""
+        """True when the worker no longer has this run and never delivered.
+
+        Asks the worker about the run itself: still queued or running means
+        wait; finished (delivery lost) or no trace of it means the collector
+        should be queued so the run can report; unreachable counts up.
+        """
         if time.monotonic() - record.dispatched_at < RECENT_DISPATCH_TTL_SECONDS:
             return False
-        state = await self._probe_one(record, record.assigned)
-        if state == "unreachable":
+        state = await fetch_worker_job_state(record.assigned, record.job_ids)
+        if state is None:
+            record.missed_probes += 1
             return record.missed_probes >= WORKER_UNREACHABLE_PROBES
-        return state == "idle"
+        record.missed_probes = 0
+        return state != "running"
 
     async def _resolve_unresolved(self):
-        """Settle submissions whose acceptance was unknown; never replay them."""
+        """Settle submissions whose acceptance was unknown; never replay them.
+
+        The worker is asked about this run specifically (its job ids in the
+        worker's queue or history). Aggregate idle/busy state is not evidence:
+        an idle worker may have run the job and lost the delivery, a busy one
+        may be on unrelated work.
+        """
         for record in list(self._records.values()):
             worker = record.unresolved
             if worker is None:
@@ -613,16 +704,24 @@ class LoadBalanceQueue:
                 continue
             if time.monotonic() - record.unresolved_since < RECENT_DISPATCH_TTL_SECONDS:
                 continue
-            state = await self._probe_one(record, worker)
-            if state == "busy":
-                self._accept_unresolved(record, "is busy with it")
-            elif state == "idle":
-                self._reject_unresolved(record)
-            elif record.missed_probes in (1, WORKER_UNREACHABLE_PROBES):
-                log(
-                    f"[Distributed] Load balance: worker {_worker_key(worker)} is unreachable; run "
-                    f"{record.prompt_id[:8]} stays on hold until its acceptance can be confirmed."
-                )
+            state = await fetch_worker_job_state(worker, record.job_ids)
+            if state == "running":
+                self._accept_unresolved(record, "has it queued or running")
+            elif state == "finished":
+                self._accept_unresolved(record, "ran it (result not received yet)")
+            elif state == "absent":
+                record.absent_checks += 1
+                record.missed_probes = 0
+                if record.absent_checks >= ABSENT_CONFIRMATIONS:
+                    self._reject_unresolved(record)
+            else:
+                record.absent_checks = 0
+                record.missed_probes += 1
+                if record.missed_probes in (1, WORKER_UNREACHABLE_PROBES):
+                    log(
+                        f"[Distributed] Load balance: worker {_worker_key(worker)} cannot be asked about run "
+                        f"{record.prompt_id[:8]}; it stays on hold until its acceptance can be confirmed."
+                    )
 
     def _accept_unresolved(self, record, how):
         worker, item = record.unresolved, record.unresolved_item
@@ -631,16 +730,22 @@ class LoadBalanceQueue:
         log(f"[Distributed] Load balance: worker {_worker_key(worker)} {how}; run {record.prompt_id[:8]} is confirmed there")
         if record.cancelled:
             record._assign(worker)
-            record.dispatched_at = record.unresolved_since
+            record.dispatched_at = time.monotonic()
             return
-        if item is not None and record.master_may_run:
+        if record.delegate_needed or (item is not None and record.master_may_run):
+            # Participating master: its collector is queued once the worker
+            # delivers (an eager run never had a master prompt at all).
             self._hold_delegate(record, worker)
-        else:
+        elif item is not None:
+            # Orchestrator-only: queue the collector now.
             record._assign(worker)
-            record.dispatched_at = record.unresolved_since
-            if item is not None:
-                record.requeue_pending = True
-                record.requeue_retry_at = 0.0
+            record.dispatched_at = time.monotonic()
+            record.requeue_pending = True
+            record.requeue_retry_at = 0.0
+        else:
+            # The collector is already executing on the master and waiting.
+            record._assign(worker)
+            record.dispatched_at = time.monotonic()
 
     def _reject_unresolved(self, record):
         worker, item = record.unresolved, record.unresolved_item
@@ -735,7 +840,7 @@ class LoadBalanceQueue:
                 # delegate prompt is queued once the worker delivers.
                 self._hold_delegate(record, worker)
             else:
-                if item is not None and not await self._requeue(record, worker):
+                if item is not None and not await self._requeue(record, worker) and not record.cancelled:
                     # The worker is running the job; keep the run tracked and
                     # keep trying to put its collector back on the master queue.
                     record.requeue_pending = True

api/queue_orchestration.py

@@ -380,7 +380,12 @@ def _register_waiting_run(prompt_index, job_id_map, candidates, *, master_may_ru
         delegate_prompt = _build_master_prompt(
             prompt_index, [worker["id"]], job_id_map, master_url, True, workflow_meta
         )
-        await queue_prompt_payload(delegate_prompt, workflow_meta, client_id, prompt_id=prompt_id)
+        # Validation yields to the loop; a clear or delete that lands meanwhile
+        # must stop the insertion itself, not just the attempt.
+        await queue_prompt_payload(
+            delegate_prompt, workflow_meta, client_id, prompt_id=prompt_id,
+            insert_if=lambda: not record.cancelled,
+        )
 
     record = PendingRun(
         prompt_id,

tests/test_async_helpers.py

@@ -1,3 +1,4 @@
+import asyncio
 import importlib.util
 import sys
 import types
@@ -63,6 +64,39 @@ async_helpers, prompt_server = _load_async_helpers_module()
 
 
 class QueuePromptPayloadTests(unittest.IsolatedAsyncioTestCase):
+    async def test_insert_if_is_checked_at_insertion_after_validation_yields(self):
+        module, prompt_server = _load_async_helpers_module()
+        validating = asyncio.Event()
+        resume = asyncio.Event()
+        state = {"cancelled": False}
+
+        async def slow_validate(prompt_id, prompt, partial_execution_targets):
+            validating.set()
+            await resume.wait()
+            return (True, None, ["9"], {})
+
+        module.execution.validate_prompt = slow_validate
+        task = asyncio.ensure_future(module.queue_prompt_payload(
+            {"9": {"class_type": "SaveImage", "inputs": {}}}, None, "client", prompt_id="run-1",
+            insert_if=lambda: not state["cancelled"],
+        ))
+        await validating.wait()
+        state["cancelled"] = True  # the queue is cleared while validation is in progress
+        resume.set()
+        with self.assertRaises(module.QueueInsertionCancelled):
+            await task
+        self.assertEqual(prompt_server.prompt_queue.items, [])
+        self.assertEqual(prompt_server.number, 12)  # no queue number consumed either
+
+        state["cancelled"] = False
+        resume.set()
+        prompt_id = await module.queue_prompt_payload(
+            {"9": {"class_type": "SaveImage", "inputs": {}}}, None, "client", prompt_id="run-1",
+            insert_if=lambda: not state["cancelled"],
+        )
+        self.assertEqual(prompt_id, "run-1")
+        self.assertEqual([item[1] for item in prompt_server.prompt_queue.items], ["run-1"])
+
     async def test_queue_prompt_payload_includes_create_time_and_client_metadata(self):
         result = await async_helpers.queue_prompt_payload(
             {"1": {"class_type": "Node"}},

tests/test_load_balance_orchestration.py

@@ -52,6 +52,7 @@ class FakePendingRun:
         self.send_to_worker = send_to_worker
         self.requeue_as_delegate = requeue_as_delegate
         self.unresolved = None
+        self.cancelled = False
 
 
 class AmbiguousDispatchError(RuntimeError):
@@ -146,6 +147,11 @@ class DeferredLoadBalanceTests(unittest.IsolatedAsyncioTestCase):
         self.assertNotIn("1", delegate_prompt)
         self.assertEqual(json.loads(delegate_prompt["2"]["inputs"]["enabled_worker_ids"]), ["w1"])
         self.assertTrue(delegate_prompt["2"]["inputs"]["delegate_only"])
+        # The insertion guard follows the run's cancellation flag.
+        insert_if = self.q.queue_prompt_payload.call_args.kwargs["insert_if"]
+        self.assertTrue(insert_if())
+        record.cancelled = True
+        self.assertFalse(insert_if())
 
     async def test_idle_worker_is_dispatched_right_away_and_master_stays_free(self):
         self.select_idle_worker.return_value = (self.worker, 2)

tests/test_load_balance_queue.py

@@ -20,6 +20,12 @@ def _module(name, **attrs):
     return value
 
 
+class InsertionRefused(RuntimeError):
+    """What queue_prompt_payload raises when insert_if refuses the insertion."""
+
+    cancelled_insertion = True
+
+
 class FakePromptQueue:
     """Just enough of ComfyUI's PromptQueue: a heap, a mutex and currently_running."""
 
@@ -112,9 +118,19 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase):
             return None if value is None else {"exec_info": {"queue_remaining": value}}
 
         self.lb.probe_worker = fake_probe
+        # 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 = []
+
+        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")
+
+        self.lb.fetch_worker_job_state = fake_job_state
 
     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):
+                   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)."""
@@ -141,9 +157,14 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase):
                 raise self.lb.AmbiguousDispatchError("acknowledgement lost")
 
         async def requeue_as_delegate(worker):
+            if requeue_gate is not None:
+                await requeue_gate.wait()  # prompt validation yielding to the loop
             if remaining_requeue_errors[0] > 0:
                 remaining_requeue_errors[0] -= 1
                 raise RuntimeError("delegate prompt failed validation")
+            # Mirrors queue_prompt_payload(insert_if=lambda: not record.cancelled).
+            if record.cancelled:
+                raise InsertionRefused("cancelled before insertion")
             self.requeued.append((prompt_id, worker["id"]))
             self.prompt_queue.put((100 + len(self.requeued), prompt_id, {"delegate": worker["id"]}, {}, [], {}))
 
@@ -163,69 +184,135 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase):
         self.fail("hand-off did not start")
 
     # ----- ambiguous acceptance ------------------------------------------------
-    async def test_lost_acknowledgement_never_replays_the_run(self):
-        """Acceptance unknown: the original prompt must not run on the master or elsewhere."""
-        record = self._queue_run("p1", ambiguous=True)
+    async def _ambiguous_run(self, prompt_id="p1", **kwargs):
+        record = self._queue_run(prompt_id, ambiguous=True, **kwargs)
         self.remaining["w1"] = 0
         await self.queue.run_cycle()
-        self.assertEqual(self.dispatched, [("p1", "w1")])
-        self.assertEqual(self.prompt_queue.prompt_ids(), [])  # not restored: the master cannot pick it up
-        self.assertIsNone(record.assigned)
+        self.assertEqual(self.dispatched, [(prompt_id, "w1")])
         self.assertEqual(record.unresolved["id"], "w1")
-        self.assertIs(self.queue.lookup_job("job-p1"), record)
-        # Another free worker must not be handed the same run, and sweeps keep it.
-        self.remaining["w2"] = 0
+        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
         self.queue._recent_dispatch.clear()
+        return record
+
+    async def test_lost_acknowledgement_never_replays_the_run(self):
+        """The worker took the job, lost the acknowledgement, ran it and lost the
+        delivery too: it now looks idle. The original prompt must not run again."""
+        record = await self._ambiguous_run()
+        self.remaining["w1"] = 0  # idle again; nothing delivered
+        self.job_state["w1"] = "finished"  # but the worker's history has the run
+        await self.queue.run_cycle()
+        self.assertEqual(len(self.dispatched), 1)  # no second submission
+        self.assertEqual(self.prompt_queue.queue, [])  # original prompt never restored
+        self.assertIsNone(record.unresolved)
+        self.assertEqual(record.assigned_id, "w1")
+        self.assertEqual(self.job_state_queries, [("w1", {"job-p1"})])
+        # Nothing else may claim it either.
+        self.remaining["w2"] = 0
         await self.queue.run_cycle()
-        self.queue.sweep()
         self.assertEqual(len(self.dispatched), 1)
-        self.assertIs(self.queue.lookup_job("job-p1"), record)
-        # The worker had accepted it after all: its result is collected normally.
+        # The lost result is reported through the collector, never by re-running.
+        record.dispatched_at = 0
+        await self.queue.run_cycle()
+        self.assertEqual(self.requeued, [("p1", "w1")])
+        self.assertEqual([item[2] for item in self.prompt_queue.queue], [{"delegate": "w1"}])
+        self.assertEqual(len(self.dispatched), 1)
+
+    async def test_lost_acknowledgement_with_job_still_running_is_collected(self):
+        record = await self._ambiguous_run()
+        self.remaining["w1"] = 0  # aggregate probe says idle (e.g. a race); the job itself is there
+        self.job_state["w1"] = "running"
+        await self.queue.run_cycle()
+        self.assertEqual(record.assigned_id, "w1")
+        self.assertTrue(record.requeue_pending and record.wait_for_result)
+        self.assertEqual(len(self.dispatched), 1)
+        self.queue.result_delivered("job-p1")
+        await self.queue.run_cycle()
+        self.assertEqual(self.requeued, [("p1", "w1")])
+        self.queue.finish_job(record, "job-p1")
+        self.assertIsNone(self.queue.lookup_job("job-p1"))
+
+    async def test_lost_acknowledgement_delivery_resolves_it_without_asking(self):
+        record = await self._ambiguous_run()
         self.queue.result_delivered("job-p1")
         await self.queue.run_cycle()
-        self.assertIsNone(record.unresolved)
         self.assertEqual(record.assigned_id, "w1")
         self.assertEqual(self.requeued, [("p1", "w1")])
-        self.assertEqual(self.prompt_queue.prompt_ids(), ["p1"])
+        self.assertEqual(self.job_state_queries, [])
         self.assertEqual(len(self.dispatched), 1)
 
-    async def test_lost_acknowledgement_is_resolved_by_seeing_the_worker_busy(self):
-        record = self._queue_run("p1", master_may_run=False, ambiguous=True)
-        self.remaining["w1"] = 0
-        await self.queue.run_cycle()
-        record.unresolved_since = 0  # past the grace period
-        self.remaining["w1"] = 1  # the worker is running it
-        await self.queue.run_cycle()
-        self.assertIsNone(record.unresolved)
-        self.assertEqual(record.assigned_id, "w1")
-        self.assertEqual(self.requeued, [("p1", "w1")])  # orchestrator-only: collector queued at once
-
-    async def test_lost_acknowledgement_restores_run_when_worker_never_took_it(self):
-        record = self._queue_run("p1", ambiguous=True)
-        self.remaining["w1"] = 0
+    async def test_busy_worker_on_unrelated_work_is_not_taken_as_acceptance(self):
+        record = await self._ambiguous_run(master_may_run=False)
+        self.remaining["w1"] = 3  # busy, but with other prompts
+        self.job_state["w1"] = "absent"
         await self.queue.run_cycle()
+        self.assertEqual(record.unresolved["id"], "w1")  # one absent check is not enough
         self.assertEqual(self.prompt_queue.prompt_ids(), [])
-        await self.queue.run_cycle()  # within the grace period: still on hold
-        self.assertEqual(self.prompt_queue.prompt_ids(), [])
-        record.unresolved_since = 0
-        await self.queue.run_cycle()  # worker idle, nothing delivered: it never ran the job
-        self.assertIsNone(record.unresolved)
+        await self.queue.run_cycle()
+        self.assertIsNone(record.unresolved)  # confirmed absent: never accepted
         self.assertIsNone(record.assigned)
         self.assertEqual(self.prompt_queue.prompt_ids(), ["p1"])  # original prompt back, same position
         self.assertEqual(self.prompt_queue.queue[0][2], {"full": "p1"})
         self.assertIs(self.queue.lookup_job("job-p1"), record)
+        self.assertEqual(self.requeued, [])
 
-    async def test_lost_acknowledgement_waits_while_worker_is_unreachable(self):
-        record = self._queue_run("p1", ambiguous=True)
-        self.remaining["w1"] = 0
-        await self.queue.run_cycle()
-        record.unresolved_since = 0
-        del self.remaining["w1"]
+    async def test_lost_acknowledgement_stays_on_hold_while_worker_cannot_be_asked(self):
+        record = await self._ambiguous_run()
+        self.job_state["w1"] = None
         for _ in range(self.lb.WORKER_UNREACHABLE_PROBES + 2):
             await self.queue.run_cycle()
         self.assertEqual(record.unresolved["id"], "w1")  # never replayed, never dropped
         self.assertEqual(self.prompt_queue.prompt_ids(), [])
         self.assertEqual(len(self.dispatched), 1)
+        # A single absent answer between failures does not count as confirmation.
+        self.job_state["w1"] = "absent"
+        await self.queue.run_cycle()
+        self.job_state["w1"] = None
+        await self.queue.run_cycle()
+        self.assertEqual(record.absent_checks, 0)
+        self.assertEqual(record.unresolved["id"], "w1")
+
+    async def test_eager_run_with_lost_acknowledgement_is_collected_after_delivery(self):
+        """Sent straight to an idle worker by a participating master, no master prompt exists yet."""
+        record = self._queue_run("p1", master_may_run=False, queued=False)
+        self.queue.register_dispatched(record, self.workers[0])
+        self.queue.mark_unresolved(record, self.workers[0])  # the eager dispatch raised ambiguity
+        self.assertIsNone(record.assigned)
+        await self.queue.run_cycle()
+        self.assertEqual(self.requeued, [])
+        self.queue.result_delivered("job-p1")
+        await self.queue.run_cycle()
+        self.assertEqual(record.assigned_id, "w1")
+        self.assertEqual(self.requeued, [("p1", "w1")])  # a collector is queued for the result
+        self.assertEqual(self.prompt_queue.prompt_ids(), ["p1"])
+        self.assertFalse(record.requeue_pending)
+        self.queue.finish_job(record, "job-p1")
+        self.assertIsNone(self.queue.lookup_job("job-p1"))
+
+    async def test_eager_run_with_lost_acknowledgement_confirmed_running_is_collected(self):
+        record = self._queue_run("p1", master_may_run=False, queued=False)
+        self.queue.register_dispatched(record, self.workers[0])
+        self.queue.mark_unresolved(record, self.workers[0])
+        record.unresolved_since = 0
+        self.job_state["w1"] = "running"
+        await self.queue.run_cycle()
+        self.assertEqual(record.assigned_id, "w1")
+        self.assertTrue(record.requeue_pending)
+        self.queue.result_delivered("job-p1")
+        await self.queue.run_cycle()
+        self.assertEqual(self.requeued, [("p1", "w1")])
+
+    async def test_eager_run_with_lost_acknowledgement_never_accepted_is_dropped(self):
+        record = self._queue_run("p1", master_may_run=False, queued=False)
+        self.queue.register_dispatched(record, self.workers[0])
+        self.queue.mark_unresolved(record, self.workers[0])
+        record.unresolved_since = 0
+        for _ in range(self.lb.ABSENT_CONFIRMATIONS):
+            await self.queue.run_cycle()
+        self.assertIsNone(self.queue.lookup_job("job-p1"))  # nothing to restore: it never existed on the master
+        self.assertNotIn("job-p1", self.pending_jobs)
+        self.assertEqual(self.prompt_queue.prompt_ids(), [])
 
     async def test_definite_failure_before_submission_still_restores_the_run(self):
         record = self._queue_run("p1", dispatch_error=RuntimeError("connection refused"))
@@ -329,11 +416,38 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase):
         self.prompt_queue.wipe_queue()
         self.assertTrue(record.cancelled)
         record.unresolved_since = 0
-        await self.queue.run_cycle()  # worker idle: never took it; the user cancelled, so do not restore
+        for _ in range(self.lb.ABSENT_CONFIRMATIONS):
+            await self.queue.run_cycle()  # confirmed never taken; the user cancelled, so do not restore
         self.assertEqual(self.prompt_queue.prompt_ids(), [])
         self.assertIsNone(self.queue.lookup_job("job-p1"))
         self.assertNotIn("job-p1", self.pending_jobs)
 
+    async def test_clear_queue_during_delegate_validation_does_not_insert_it(self):
+        """Cancellation that lands while the delegate prompt is being validated."""
+        gate = asyncio.Event()
+        record = self._queue_run("p1", master_may_run=False, requeue_gate=gate)
+        self.remaining["w1"] = 0
+        cycle = asyncio.ensure_future(self.queue.run_cycle())
+        for _ in range(200):
+            if self.dispatched:
+                break
+            await asyncio.sleep(0.005)
+        self.assertEqual(self.dispatched, [("p1", "w1")])  # accepted; delegate now validating
+        self.prompt_queue.wipe_queue()
+        self.assertTrue(record.cancelled)
+        gate.set()
+        await cycle
+        self.assertEqual(self.requeued, [])
+        self.assertEqual(self.prompt_queue.prompt_ids(), [])  # never inserted
+        self.assertEqual(record.requeue_failures, 0)  # refused insertion is not a retryable failure
+        self.assertFalse(record.requeue_pending)
+        self.assertIs(self.queue.lookup_job("job-p1"), record)  # accepted work stays owned
+        self.queue.result_delivered("job-p1")
+        await self.queue.run_cycle()
+        self.assertIsNone(self.queue.lookup_job("job-p1"))
+        self.assertNotIn("job-p1", self.pending_jobs)
+        self.assertEqual(self.prompt_queue.prompt_ids(), [])
+
     # ----- interrupted assignment on an orchestrator-only master ---------------
     async def test_interrupt_during_assignment_cleans_result_queue_and_stops_dispatch(self):
         gate = asyncio.Event()
@@ -485,9 +599,9 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase):
         self.assertTrue(record.requeue_pending and record.wait_for_result)
         self.queue.sweep()
         self.assertIs(self.queue.lookup_job("job-p1"), record)
-        # Still nothing while the worker is busy, even after the dispatch grace.
+        # Still nothing while the worker is running it, even after the dispatch grace.
         record.dispatched_at = 0
-        self.remaining["w2"] = 1
+        self.job_state["w2"] = "running"
         await self.queue.run_cycle()
         self.assertEqual(self.requeued, [])
         # The worker delivers: the collector is queued behind the real work.
@@ -499,13 +613,17 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase):
         self.queue.finish_job(record, "job-p1")
         self.assertIsNone(self.queue.lookup_job("job-p1"))
 
-    async def test_delegate_is_queued_when_worker_goes_idle_without_delivering(self):
+    async def test_delegate_is_queued_when_worker_finished_without_delivering(self):
         record = self._queue_run("p1")
         self.remaining["w1"] = 0
         await self.queue.run_cycle()
         self.assertEqual(self.requeued, [])
         record.dispatched_at = 0  # past the dispatch grace
-        await self.queue.run_cycle()  # worker idle, nothing delivered: let the collector report it
+        self.job_state["w1"] = "running"
+        await self.queue.run_cycle()  # still running there: keep waiting
+        self.assertEqual(self.requeued, [])
+        self.job_state["w1"] = "finished"
+        await self.queue.run_cycle()  # ran it, nothing delivered: let the collector report it
         self.assertEqual(self.requeued, [("p1", "w1")])
         self.assertFalse(record.wait_for_result)
 
@@ -514,7 +632,7 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase):
         self.remaining["w1"] = 0
         await self.queue.run_cycle()
         record.dispatched_at = 0
-        del self.remaining["w1"]  # probes fail from now on
+        self.job_state["w1"] = None  # the worker cannot be asked from now on
         for _ in range(self.lb.WORKER_UNREACHABLE_PROBES - 1):
             await self.queue.run_cycle()
             self.assertEqual(self.requeued, [])

utils/async_helpers.py

@@ -2,6 +2,7 @@
 Async helper utilities for ComfyUI-Distributed.
 """
 import asyncio
+import contextlib
 import threading
 import time
 import uuid
@@ -79,6 +80,12 @@ def _summarize_node_errors(node_errors: dict) -> str:
     return " | ".join(parts)
 
 
+class QueueInsertionCancelled(RuntimeError):
+    """The prompt was valid but its insertion was refused by ``insert_if``."""
+
+    cancelled_insertion = True
+
+
 class PromptValidationError(RuntimeError):
     """Raised when a prompt fails ComfyUI validation with structured context."""
 
@@ -127,12 +134,17 @@ async def queue_prompt_payload(
     *,
     partial_execution_targets=None,
     prompt_id=None,
+    insert_if=None,
 ):
     """Validate and queue a prompt, preserving any explicitly requested roots.
 
     ``prompt_id`` lets a caller re-queue a prompt under an id it already
     handed out (the load-balance dispatcher swaps a waiting prompt for its
     delegate form without changing the id clients track).
+
+    ``insert_if`` is checked under the queue mutex right before insertion,
+    after validation has yielded to the event loop; when it returns false
+    the prompt is not queued and ``QueueInsertionCancelled`` is raised.
     """
     targets = (validate_execution_targets(prompt_obj, partial_execution_targets)
                if partial_execution_targets is not None else None)
@@ -179,10 +191,15 @@ async def queue_prompt_payload(
         if key in extra_data:
             sensitive[key] = extra_data.pop(key)
 
-    number = getattr(prompt_server, "number", 0)
-    prompt_server.number = number + 1
-    prompt_queue_item = (number, prompt_id, prompt, extra_data, valid[2], sensitive)
-    prompt_server.prompt_queue.put(prompt_queue_item)
+    prompt_queue = prompt_server.prompt_queue
+    mutex = getattr(prompt_queue, "mutex", None)
+    with (mutex if mutex is not None else contextlib.nullcontext()):
+        if insert_if is not None and not insert_if():
+            raise QueueInsertionCancelled(f"Prompt {prompt_id} was cancelled before it could be queued")
+        number = getattr(prompt_server, "number", 0)
+        prompt_server.number = number + 1
+        prompt_queue_item = (number, prompt_id, prompt, extra_data, valid[2], sensitive)
+        prompt_queue.put(prompt_queue_item)
 
     if include_queue_metadata:
         return {