fix: preserve image-list worker selection across queued runs

tests/test_image_list_master_selection.py

@@ -1,8 +1,9 @@
 """Image-list participation uses the existing master orchestrator setting."""
 import json
 import sys
+import types
 import unittest
-from unittest.mock import AsyncMock, patch
+from unittest.mock import AsyncMock, Mock, patch
 
 from test_serverless_integration import PKG, load_modules
 
@@ -76,7 +77,7 @@ class MasterSelectionTests(unittest.IsolatedAsyncioTestCase):
         self.config["settings"]["master_delegate_only"] = True
         await self.queue()
         self.assertIs(self.prepared_inputs()["image_list_master_participates"], False)
-        self.assertTrue(self.preflight.call_args.kwargs["require_idle"])
+        self.assertIs(self.preflight.call_args.kwargs["require_idle"], False)
         # Upstream source stays on master even in orchestrator-only mode.
         self.assertIn("1", self.q.queue_prompt_payload.call_args.args[0])
 
@@ -107,10 +108,154 @@ class MasterSelectionTests(unittest.IsolatedAsyncioTestCase):
         self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), [])
         self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
 
-    async def test_busy_workers_do_not_prevent_enabled_master(self):
-        self.preflight.return_value = {}
+    def use_real_preflight(self, queues):
+        """Exercise capability/queue parsing with read-only fake HTTP responses."""
+        from test_dispatch_selection import dispatch, _FakeSession, _valid_capability
+
+        responses = {(worker_id, "capabilities"): (200, _valid_capability())
+                     for worker_id in queues}
+        responses.update({(worker_id, "queue"): (200, queue)
+                          for worker_id, queue in queues.items()})
+        session = _FakeSession(responses)
+        for target, value in (
+            ("get_client_session", AsyncMock(return_value=session)),
+            ("build_worker_url", lambda worker, path: f"http://{worker['id']}{path}"),
+        ):
+            patched = patch.object(dispatch, target, value)
+            patched.start()
+            self.addCleanup(patched.stop)
+        self.preflight = AsyncMock(wraps=dispatch.preflight_image_list_workers)
+        sys.modules[PKG + ".api.orchestration.dispatch"].preflight_image_list_workers = self.preflight
+        return responses
+
+    async def run_prepared_selection(self):
+        """Run real runtime selection up to the scheduler, without any item I/O."""
+        rt = sys.modules[PKG + ".image_list.runtime"]
+        inputs = self.prepared_inputs()
+        image = object()
+        context = rt.create_master_job(
+            "2", "3", self.q.queue_prompt_payload.call_args.args[0],
+            [{"image": image, "index": 0}],
+            json.loads(inputs["image_list_enabled_worker_ids"]),
+            inputs["image_list_master_callback"],
+            master_participates=inputs["image_list_master_participates"],
+            explicit_targets=True,
+        )
+        self.addCleanup(rt.remove_master_job, context)
+        self.scheduler = AsyncMock(return_value=[{"image": image}])
+        self.local_slot = Mock(return_value=types.SimpleNamespace(close=AsyncMock()))
+        self.runtime_session = AsyncMock(return_value=object())
+        config = sys.modules[PKG + ".utils.config"]
+        network = sys.modules[PKG + ".utils.network"]
+        with patch.object(config, "load_config", lambda: self.config), \
+             patch.object(config, "get_worker_timeout_seconds", lambda: 60, create=True), \
+             patch.object(config, "get_image_list_job_timeout_seconds", lambda: 0, create=True), \
+             patch.object(network, "get_client_session", self.runtime_session, create=True), \
+             patch.object(network, "build_worker_url", Mock(), create=True), \
+             patch.object(rt, "dynamic_schedule", self.scheduler), \
+             patch.dict(sys.modules, {PKG + ".image_list.local_executor": types.SimpleNamespace(
+                 LocalExecutionSlot=self.local_slot)}):
+            await rt.run_image_list_job(context)
+        self.assertIsNone(rt.get_master_job(context))
+        return self.scheduler.call_args.args[1]
+
+    async def test_busy_worker_at_enqueue_is_available_to_scheduler_when_idle_at_start(self):
+        responses = self.use_real_preflight({
+            "worker": {"queue_running": [[1, "prompt-a"]], "queue_pending": []},
+        })
+        result = await self.queue()
+        self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), ["worker"])
+        self.assertEqual(result[2], 1)  # Candidates, not a promise of execution-time availability.
+        self.q.dispatch_worker_prompt.assert_not_awaited()
+
+        responses[("worker", "queue")] = (200, {"queue_running": [], "queue_pending": []})
+        # A newly enabled worker is not implicitly added to this queued selection.
+        self.config["workers"].append({"id": "later", "enabled": True})
+        slots = await self.run_prepared_selection()
+        self.assertEqual([slot for slot in slots if isinstance(slot, str)], ["worker"])
+        self.assertEqual(len(slots), 2)
+        self.assertEqual(self.preflight.await_count, 2)
+        self.assertEqual(self.preflight.await_args_list[0].args, self.preflight.await_args_list[1].args)
+
+    async def test_worker_still_busy_at_start_is_excluded_from_scheduler(self):
+        self.use_real_preflight({
+            "worker": {"queue_running": [], "queue_pending": [[1, "prompt-a"]]},
+        })
         await self.queue()
-        self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), [])
+        self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), ["worker"])
+        slots = await self.run_prepared_selection()
+        self.assertEqual(len(slots), 1)
+        self.assertNotIn("worker", slots)
+        self.assertEqual(self.preflight.await_count, 2)
+        self.runtime_session.assert_not_awaited()
+
+    async def test_explicit_empty_selection_stays_master_only_at_start(self):
+        self.use_real_preflight({"worker": {"queue_running": [], "queue_pending": []}})
+        await self.queue(enabled_worker_ids=[])
+        slots = await self.run_prepared_selection()
+        self.assertEqual(len(slots), 1)
+        self.assertNotIn("worker", slots)
+        self.preflight.assert_not_awaited()
+        self.runtime_session.assert_not_awaited()
+
+    async def test_busy_incompatible_worker_fails_before_master_queue(self):
+        responses = self.use_real_preflight({
+            "worker": {"queue_running": [[1, "prompt-a"]], "queue_pending": []},
+        })
+        capability = responses[("worker", "capabilities")][1]
+        for field, invalid, message in (
+            ("protocol_version", 1, "incompatible"),
+            ("execution_targets", False, "execution_targets"),
+            ("value_results", False, "value_results"),
+        ):
+            with self.subTest(field=field), patch.dict(capability, {field: invalid}):
+                if field == "value_results":
+                    self.prompt["3"]["inputs"]["value_10"] = ["2", 2]
+                with self.assertRaisesRegex(RuntimeError, message):
+                    await self.queue()
+                self.q.queue_prompt_payload.assert_not_awaited()
+                self.q.dispatch_worker_prompt.assert_not_awaited()
+
+    async def test_orchestrator_with_all_workers_busy_queues_for_idle_start(self):
+        responses = self.use_real_preflight({
+            "worker": {"queue_running": [[1, "prompt-a"]], "queue_pending": []},
+        })
+        result = await self.queue(delegate_master=True)
+        self.q.queue_prompt_payload.assert_awaited_once()
+        self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), ["worker"])
+        self.assertIs(self.prepared_inputs()["image_list_master_participates"], False)
+        self.assertEqual(result[2], 1)
+        self.q.dispatch_worker_prompt.assert_not_awaited()
+
+        responses[("worker", "queue")] = (200, {"queue_running": [], "queue_pending": []})
+        self.assertEqual(await self.run_prepared_selection(), ["worker"])
+        self.local_slot.assert_not_called()
+        self.assertEqual(self.preflight.await_count, 2)
+        self.assertIs(self.preflight.await_args_list[0].kwargs["require_idle"], False)
+        self.assertIs(self.preflight.await_args_list[1].kwargs["require_idle"], True)
+
+    async def test_orchestrator_rechecks_availability_without_master_fallback(self):
+        responses = self.use_real_preflight({"worker": {"queue_running": [], "queue_pending": []}})
+        await self.queue(delegate_master=True)
+        responses[("worker", "queue")] = (200, {"queue_running": [[1, "other"]], "queue_pending": []})
+        with self.assertRaisesRegex(RuntimeError, "No selected feature-compatible direct worker is idle"):
+            await self.run_prepared_selection()
+        self.assertEqual(self.preflight.await_count, 2)
+        self.scheduler.assert_not_awaited()
+        self.local_slot.assert_not_called()
+
+    async def test_orchestrator_retains_busy_candidates_when_another_worker_is_idle(self):
+        self.config["workers"].append({"id": "idle", "enabled": True})
+        responses = self.use_real_preflight({
+            "worker": {"queue_running": [[1, "prompt-a"]], "queue_pending": []},
+            "idle": {"queue_running": [], "queue_pending": []},
+        })
+        await self.queue(delegate_master=True)
+        self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), ["worker", "idle"])
+        responses[("worker", "queue")] = (200, {"queue_running": [], "queue_pending": []})
+        responses[("idle", "queue")] = (200, {"queue_running": [[1, "other"]], "queue_pending": []})
+        self.assertEqual(await self.run_prepared_selection(), ["worker"])
+        self.local_slot.assert_not_called()
 
     async def test_orchestrator_without_workers_fails_before_queue(self):
         self.config["settings"]["master_delegate_only"] = True