fix(image-list): isolate unavailable workers during preflight

api/orchestration/dispatch.py

@@ -173,9 +173,10 @@ async def preflight_image_list_workers(workers, require_idle=True, timeout=3.0):
 
     async def inspect(worker):
         worker_id = str(worker.get("id") or "unknown")
-        async with semaphore:
-            cap_url = build_worker_url(worker, "/distributed/image-list/capabilities")
-            try:
+        check = "capability"
+        try:
+            async with semaphore:
+                cap_url = build_worker_url(worker, "/distributed/image-list/capabilities")
                 async with session.get(cap_url, timeout=aiohttp.ClientTimeout(total=timeout)) as response:
                     if response.status != 200:
                         raise RuntimeError(
@@ -183,37 +184,58 @@ async def preflight_image_list_workers(workers, require_idle=True, timeout=3.0):
                             f"(HTTP {response.status}); upgrade that worker before running this feature"
                         )
                     capability = await response.json()
-            except asyncio.TimeoutError as exc:
-                raise RuntimeError(f"Image-list worker {worker_id} capability check timed out") from exc
-            except aiohttp.ClientError as exc:
-                raise RuntimeError(f"Image-list worker {worker_id} capability check failed: {exc}") from exc
-            nodes = set(capability.get("node_ids", [])) if isinstance(capability, dict) else set()
-            if (not isinstance(capability, dict) or capability.get("supported") is not True
-                    or capability.get("protocol_version") != 1 or not required_nodes.issubset(nodes)):
-                raise RuntimeError(
-                    f"Image-list worker {worker_id} is incompatible: requires protocol v1 and nodes "
-                    f"{sorted(required_nodes)}; no work was submitted"
-                )
-            queue_url = build_worker_url(worker, "/queue")
-            try:
+                if not isinstance(capability, dict):
+                    raise RuntimeError(f"Image-list worker {worker_id} returned malformed capability data")
+                supported = capability.get("supported")
+                protocol_version = capability.get("protocol_version")
+                nodes = capability.get("node_ids")
+                if (
+                    not isinstance(supported, bool)
+                    or type(protocol_version) is not int
+                    or not isinstance(nodes, list)
+                    or any(not isinstance(node_id, str) for node_id in nodes)
+                ):
+                    raise RuntimeError(f"Image-list worker {worker_id} returned malformed capability data")
+                if (
+                    supported is not True
+                    or protocol_version != 1
+                    or not required_nodes.issubset(nodes)
+                ):
+                    raise RuntimeError(
+                        f"Image-list worker {worker_id} is incompatible: requires protocol v1 and nodes "
+                        f"{sorted(required_nodes)}; no work was submitted"
+                    )
+
+                check = "queue"
+                queue_url = build_worker_url(worker, "/queue")
                 async with session.get(queue_url, timeout=aiohttp.ClientTimeout(total=timeout)) as response:
                     if response.status != 200:
                         raise RuntimeError(f"Image-list worker {worker_id} queue preflight returned HTTP {response.status}")
                     queue = await response.json()
-            except asyncio.TimeoutError as exc:
-                raise RuntimeError(f"Image-list worker {worker_id} queue preflight timed out") from exc
-            except aiohttp.ClientError as exc:
-                raise RuntimeError(f"Image-list worker {worker_id} queue preflight failed: {exc}") from exc
-            if not isinstance(queue, dict) or not isinstance(queue.get("queue_running"), list) or not isinstance(queue.get("queue_pending"), list):
-                raise RuntimeError(f"Image-list worker {worker_id} returned malformed /queue state")
-            idle = not queue["queue_running"] and not queue["queue_pending"]
-            return worker_id, idle
+                if (
+                    not isinstance(queue, dict)
+                    or not isinstance(queue.get("queue_running"), list)
+                    or not isinstance(queue.get("queue_pending"), list)
+                ):
+                    raise RuntimeError(f"Image-list worker {worker_id} returned malformed /queue state")
+                idle = not queue["queue_running"] and not queue["queue_pending"]
+                return worker_id, idle
+        except (asyncio.TimeoutError, aiohttp.ClientConnectionError) as exc:
+            log(
+                f"[Distributed] Image-list worker {worker_id} unavailable during {check} check; "
+                f"skipping: {exc}"
+            )
+            return worker_id, None
 
     statuses = await asyncio.gather(*(inspect(worker) for worker in workers))
-    result = {worker_id: idle for worker_id, idle in statuses}
-    if require_idle and not any(result.values()):
+    idle_workers = {
+        worker_id: True
+        for worker_id, idle in statuses
+        if idle is True
+    }
+    if require_idle and not idle_workers:
         raise RuntimeError("No selected feature-compatible direct worker is idle; unrelated worker prompts were left untouched")
-    return {worker_id: True for worker_id, idle in result.items() if idle}
+    return idle_workers
 
 
 async def select_active_workers(

api/orchestration/prompt_transform.py

@@ -442,14 +442,30 @@ def prepare_delegate_master_prompt(prompt_obj, collector_ids):
 
 
 def _resolve_direct_workers_for_image_list(config, requested_ids):
-    """Resolve only enabled direct workers already selected by orchestration."""
+    """Resolve enabled workers by default, or exact explicitly selected direct workers."""
     requested = None if requested_ids is None else {str(worker_id) for worker_id in requested_ids}
+    configured_workers = (config or {}).get("workers", [])
+    configured_ids = {
+        str(worker.get("id") or "").strip()
+        for worker in configured_workers
+        if str(worker.get("id") or "").strip()
+    }
+    if requested is not None:
+        unknown_ids = requested - configured_ids
+        if unknown_ids:
+            raise RuntimeError(
+                "Unknown explicitly selected image-list worker ID(s): "
+                + ", ".join(sorted(unknown_ids))
+            )
     direct = []
-    for worker in (config or {}).get("workers", []):
+    for worker in configured_workers:
         worker_id = str(worker.get("id") or "").strip()
-        if not worker_id or not worker.get("enabled", False):
+        if not worker_id:
             continue
-        if requested is not None and worker_id not in requested:
+        if requested is None:
+            if not worker.get("enabled", False):
+                continue
+        elif worker_id not in requested:
             continue
         if str(worker.get("type") or "local").strip().lower() == "serverless":
             continue

tests/test_dispatch_selection.py

@@ -55,7 +55,13 @@ def _load_dispatch_module():
             def __init__(self, total=None):
                 self.total = total
 
-        class _ClientConnectorError(Exception):
+        class _ClientError(Exception):
+            pass
+
+        class _ClientConnectorError(_ClientError):
+            pass
+
+        class _ClientConnectionError(_ClientError):
             pass
 
         class _WSMsgType:
@@ -76,7 +82,9 @@ def _load_dispatch_module():
                 self.closed = True
 
         aiohttp_module.ClientTimeout = _ClientTimeout
+        setattr(aiohttp_module, "ClientError", _ClientError)
         aiohttp_module.ClientConnectorError = _ClientConnectorError
+        setattr(aiohttp_module, "ClientConnectionError", _ClientConnectionError)
         aiohttp_module.WSMsgType = _WSMsgType
         aiohttp_module.TCPConnector = _TCPConnector
         aiohttp_module.ClientSession = _ClientSession
@@ -100,6 +108,75 @@ def _load_dispatch_module():
 dispatch = _load_dispatch_module()
 
 
+def _load_prompt_transform_module():
+    module_path = Path(__file__).resolve().parents[1] / "api" / "orchestration" / "prompt_transform.py"
+    package_name = "prompt_transform_worker_selection_testpkg"
+    for suffix in ("", ".api", ".api.orchestration", ".utils"):
+        package = types.ModuleType(package_name + suffix)
+        package.__path__ = []
+        sys.modules[package_name + suffix] = package
+    logging_module = types.ModuleType(f"{package_name}.utils.logging")
+    setattr(logging_module, "debug_log", lambda *_args, **_kwargs: None)
+    sys.modules[f"{package_name}.utils.logging"] = logging_module
+
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.api.orchestration.prompt_transform",
+        module_path,
+    )
+    assert spec is not None and spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+prompt_transform = _load_prompt_transform_module()
+
+
+class _FakeResponse:
+    def __init__(self, status=200, payload=None, error=None):
+        self.status = status
+        self.payload = payload
+        self.error = error
+
+    async def __aenter__(self):
+        if self.error is not None:
+            raise self.error
+        return self
+
+    async def __aexit__(self, *_args):
+        return None
+
+    async def json(self):
+        return self.payload
+
+
+class _FakeSession:
+    def __init__(self, responses):
+        self.responses = responses
+
+    def get(self, url, timeout=None):
+        del timeout
+        worker_id = url.removeprefix("http://").split("/", 1)[0]
+        endpoint = "capabilities" if url.endswith("/capabilities") else "queue"
+        response = self.responses[(worker_id, endpoint)]
+        if isinstance(response, BaseException):
+            return _FakeResponse(error=response)
+        status, payload = response
+        return _FakeResponse(status=status, payload=payload)
+
+
+def _valid_capability():
+    return {
+        "supported": True,
+        "protocol_version": 1,
+        "node_ids": ["DistributedImageListStart", "DistributedImageListEnd"],
+    }
+
+
+def _valid_queue(running=None, pending=None):
+    return {"queue_running": running or [], "queue_pending": pending or []}
+
+
 class DispatchSelectionTests(unittest.IsolatedAsyncioTestCase):
     async def test_select_active_workers_filters_offline(self):
         workers = [
@@ -249,6 +326,194 @@ class DispatchSelectionTests(unittest.IsolatedAsyncioTestCase):
 
         self.assertIsNone(selected)
 
+    async def test_image_list_preflight_skips_unreachable_worker_and_keeps_healthy_worker(self):
+        workers = [{"id": "healthy"}, {"id": "offline"}]
+        responses = {
+            (worker_id, "capabilities"): (200, _valid_capability())
+            for worker_id in ("healthy", "offline")
+        }
+        responses.update({
+            (worker_id, "queue"): (200, _valid_queue())
+            for worker_id in ("healthy", "offline")
+        })
+        responses[("offline", "capabilities")] = dispatch.aiohttp.ClientConnectionError("connection refused")
+        session = _FakeSession(responses)
+
+        async def fake_session():
+            return session
+
+        with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
+            dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
+        ), patch.object(dispatch, "log") as log_mock:
+            result = await dispatch.preflight_image_list_workers(workers)
+
+        self.assertEqual(result, {"healthy": True})
+        self.assertTrue(any("offline" in str(call) for call in log_mock.call_args_list))
+
+    async def test_image_list_preflight_skips_timeout_and_keeps_healthy_worker(self):
+        workers = [{"id": "healthy"}, {"id": "slow"}]
+        responses: dict = {
+            (worker_id, "capabilities"): (200, _valid_capability())
+            for worker_id in ("healthy", "slow")
+        }
+        responses.update({
+            (worker_id, "queue"): (200, _valid_queue())
+            for worker_id in ("healthy", "slow")
+        })
+        responses[("slow", "queue")] = asyncio.TimeoutError()
+        session = _FakeSession(responses)
+
+        async def fake_session():
+            return session
+
+        with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
+            dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
+        ), patch.object(dispatch, "log") as log_mock:
+            result = await dispatch.preflight_image_list_workers(workers)
+
+        self.assertEqual(result, {"healthy": True})
+        self.assertTrue(any("slow" in str(call) for call in log_mock.call_args_list))
+
+    async def test_image_list_preflight_excludes_busy_workers(self):
+        workers = [{"id": "idle"}, {"id": "busy"}]
+        responses = {
+            (worker_id, "capabilities"): (200, _valid_capability())
+            for worker_id in ("idle", "busy")
+        }
+        responses.update({
+            ("idle", "queue"): (200, _valid_queue()),
+            ("busy", "queue"): (200, _valid_queue(running=[{"prompt": "existing"}])),
+        })
+        session = _FakeSession(responses)
+
+        async def fake_session():
+            return session
+
+        with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
+            dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
+        ):
+            result = await dispatch.preflight_image_list_workers(workers)
+
+        self.assertEqual(result, {"idle": True})
+
+    async def test_image_list_preflight_reports_no_idle_workers_when_all_are_unavailable(self):
+        workers = [{"id": "offline-a"}, {"id": "offline-b"}]
+        responses = {
+            (worker_id, "capabilities"): dispatch.aiohttp.ClientConnectionError("unreachable")
+            for worker_id in ("offline-a", "offline-b")
+        }
+        session = _FakeSession(responses)
+
+        async def fake_session():
+            return session
+
+        with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
+            dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
+        ):
+            self.assertEqual(
+                await dispatch.preflight_image_list_workers(workers, require_idle=False),
+                {},
+            )
+            with self.assertRaisesRegex(RuntimeError, "No selected feature-compatible direct worker is idle"):
+                await dispatch.preflight_image_list_workers(workers, require_idle=True)
+
+    async def test_image_list_preflight_fails_closed_for_malformed_capabilities(self):
+        malformed = {"supported": True, "protocol_version": 1, "node_ids": None}
+        session = _FakeSession({("worker", "capabilities"): (200, malformed)})
+
+        async def fake_session():
+            return session
+
+        with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
+            dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
+        ):
+            with self.assertRaisesRegex(RuntimeError, "malformed.*capabilit"):
+                await dispatch.preflight_image_list_workers([{"id": "worker"}])
+
+    async def test_image_list_preflight_fails_closed_for_incompatible_protocol(self):
+        incompatible = _valid_capability()
+        incompatible["protocol_version"] = 2
+        session = _FakeSession({("worker", "capabilities"): (200, incompatible)})
+
+        async def fake_session():
+            return session
+
+        with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
+            dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
+        ):
+            with self.assertRaisesRegex(RuntimeError, "incompatible"):
+                await dispatch.preflight_image_list_workers([{"id": "worker"}])
+
+    async def test_image_list_preflight_fails_closed_for_missing_required_node(self):
+        incompatible = _valid_capability()
+        incompatible["node_ids"] = ["DistributedImageListStart"]
+        session = _FakeSession({("worker", "capabilities"): (200, incompatible)})
+
+        async def fake_session():
+            return session
+
+        with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
+            dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
+        ):
+            with self.assertRaisesRegex(RuntimeError, "incompatible"):
+                await dispatch.preflight_image_list_workers([{"id": "worker"}])
+
+    async def test_image_list_preflight_fails_closed_for_malformed_queue(self):
+        responses = {
+            ("worker", "capabilities"): (200, _valid_capability()),
+            ("worker", "queue"): (200, {"queue_running": [], "queue_pending": None}),
+        }
+        session = _FakeSession(responses)
+
+        async def fake_session():
+            return session
+
+        with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
+            dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
+        ):
+            with self.assertRaisesRegex(RuntimeError, "malformed /queue"):
+                await dispatch.preflight_image_list_workers([{"id": "worker"}])
+
+    async def test_image_list_preflight_does_not_hide_unexpected_errors(self):
+        session = _FakeSession({("worker", "capabilities"): ValueError("programming error")})
+
+        async def fake_session():
+            return session
+
+        with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
+            dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
+        ):
+            with self.assertRaisesRegex(ValueError, "programming error"):
+                await dispatch.preflight_image_list_workers([{"id": "worker"}])
+
+
+class ImageListDirectWorkerResolutionTests(unittest.TestCase):
+    def setUp(self):
+        self.config = {"workers": [
+            {"id": "enabled", "enabled": True, "host": "enabled-host", "port": 8188},
+            {"id": "disabled", "enabled": False, "host": "disabled-host", "port": 8189},
+            {"id": "serverless", "enabled": True, "type": "serverless"},
+            {"id": "other", "enabled": True, "host": "other-host", "port": 8190},
+        ]}
+
+    def test_implicit_selection_uses_only_enabled_direct_workers(self):
+        workers = prompt_transform._resolve_direct_workers_for_image_list(self.config, None)
+        self.assertEqual([worker["id"] for worker in workers], ["enabled", "other"])
+
+    def test_explicit_selection_includes_exact_disabled_direct_worker(self):
+        workers = prompt_transform._resolve_direct_workers_for_image_list(self.config, ["disabled"])
+        self.assertEqual([worker["id"] for worker in workers], ["disabled"])
+
+    def test_explicit_empty_selection_returns_no_workers(self):
+        self.assertEqual(prompt_transform._resolve_direct_workers_for_image_list(self.config, []), [])
+
+    def test_unknown_explicit_id_raises_instead_of_falling_back(self):
+        with self.assertRaisesRegex(RuntimeError, "missing"):
+            prompt_transform._resolve_direct_workers_for_image_list(self.config, ["enabled", "missing"])
+
+    def test_explicit_serverless_selection_is_excluded(self):
+        self.assertEqual(prompt_transform._resolve_direct_workers_for_image_list(self.config, ["serverless"]), [])
+
 
 if __name__ == "__main__":
     unittest.main()