fix(image-list): bound acknowledgements and configure run deadline

api/orchestration/dispatch.py

@@ -77,27 +77,44 @@ async def _dispatch_via_websocket(worker_url, payload, client_id, timeout=60.0):
     ws_url = worker_url.replace("http://", "ws://").replace("https://", "wss://")
     ws_url = f"{ws_url}/distributed/worker_ws"
     session = await get_client_session()
+    loop = asyncio.get_running_loop()
+    deadline = loop.time() + timeout
+    ws = None
 
-    async with session.ws_connect(ws_url, heartbeat=20, timeout=timeout) as ws:
-        await ws.send_json(ws_payload)
-        async for msg in ws:
-            if msg.type == aiohttp.WSMsgType.TEXT:
-                data = json.loads(msg.data or "{}")
-                if data.get("type") == "dispatch_ack" and data.get("request_id") == request_id:
-                    if data.get("ok"):
-                        return data.get("prompt_id")
-                    error_text = data.get("error") or "Worker rejected websocket dispatch."
-                    validation_error = data.get("validation_error")
-                    node_errors = data.get("node_errors")
-                    if validation_error:
-                        error_text = f"{error_text} | validation_error={validation_error}"
-                    if node_errors:
-                        error_text = f"{error_text} | node_errors={node_errors}"
-                    raise RuntimeError(error_text)
-            elif msg.type in (aiohttp.WSMsgType.ERROR, aiohttp.WSMsgType.CLOSED):
-                raise RuntimeError(f"Worker websocket closed unexpectedly: {msg.type}")
-
-    raise RuntimeError("Worker websocket closed before dispatch_ack was received.")
+    try:
+        async with asyncio.timeout_at(deadline):
+            ws = await session.ws_connect(ws_url, heartbeat=20)
+            await ws.send_json(ws_payload)
+            while True:
+                msg = await ws.receive()
+                if msg.type == aiohttp.WSMsgType.TEXT:
+                    data = json.loads(msg.data or "{}")
+                    if data.get("type") == "dispatch_ack" and data.get("request_id") == request_id:
+                        if data.get("ok"):
+                            return data.get("prompt_id")
+                        error_text = data.get("error") or "Worker rejected websocket dispatch."
+                        validation_error = data.get("validation_error")
+                        node_errors = data.get("node_errors")
+                        if validation_error:
+                            error_text = f"{error_text} | validation_error={validation_error}"
+                        if node_errors:
+                            error_text = f"{error_text} | node_errors={node_errors}"
+                        raise RuntimeError(error_text)
+                elif msg.type in (aiohttp.WSMsgType.ERROR, aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.CLOSED):
+                    raise RuntimeError(f"Worker websocket closed before acknowledgement: {msg.type}; submission status is ambiguous")
+    except asyncio.TimeoutError as exc:
+        raise RuntimeError(
+            "Worker websocket dispatch timed out before dispatch_ack; submission status is ambiguous."
+        ) from exc
+    finally:
+        if ws is not None:
+            # Closing must not discard an already acknowledged prompt ID or
+            # replace a definite rejection with an ambiguous timeout.
+            try:
+                async with asyncio.timeout_at(min(deadline, loop.time() + 1.0)):
+                    await ws.close()
+            except (asyncio.TimeoutError, aiohttp.ClientError):
+                pass
 
 
 async def dispatch_worker_prompt(

tests/test_image_list_timeout_config.py

@@ -0,0 +1,106 @@
+import importlib.util
+import sys
+import types
+from pathlib import Path
+
+import pytest
+
+
+def _load_config_module():
+    module_path = Path(__file__).resolve().parents[1] / "utils" / "config.py"
+    package_name = "image_list_timeout_config_testpkg"
+
+    for mod_name in list(sys.modules):
+        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
+            del sys.modules[mod_name]
+
+    root_pkg = types.ModuleType(package_name)
+    root_pkg.__path__ = [str(module_path.parent)]
+    sys.modules[package_name] = root_pkg
+
+    logging_module = types.ModuleType(f"{package_name}.logging")
+    logging_module.log = lambda *_args, **_kwargs: None
+    logging_module.debug_log = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.logging"] = logging_module
+
+    constants_module = types.ModuleType(f"{package_name}.constants")
+    constants_module.HEARTBEAT_TIMEOUT = 30
+    sys.modules[f"{package_name}.constants"] = constants_module
+
+    spec = importlib.util.spec_from_file_location(f"{package_name}.config", module_path)
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    spec.loader.exec_module(module)
+    return module
+
+
+config_module = _load_config_module()
+
+
+SETTING_KEY = "image_list_job_timeout_seconds"
+
+
+def _config_with_timeout(value):
+    return {"settings": {SETTING_KEY: value}}
+
+
+def test_default_config_sets_image_list_job_timeout_to_zero():
+    assert config_module.get_default_config()["settings"][SETTING_KEY] == 0.0
+
+
+def test_missing_value_defaults_to_disabled_timeout():
+    assert config_module.get_image_list_job_timeout_seconds({}) == 0.0
+    assert config_module.get_image_list_job_timeout_seconds({"settings": {}}) == 0.0
+
+
+def test_none_config_loads_current_config(monkeypatch):
+    cfg = _config_with_timeout("12.5")
+    monkeypatch.setattr(config_module, "load_config", lambda: cfg)
+
+    assert config_module.get_image_list_job_timeout_seconds() == 12.5
+
+
+@pytest.mark.parametrize(
+    ("value", "expected"),
+    [
+        (0, 0.0),
+        (0.0, 0.0),
+        ("0", 0.0),
+        (" 0.0 ", 0.0),
+        (1, 1.0),
+        (2.5, 2.5),
+        ("3", 3.0),
+        (" 4.25 ", 4.25),
+    ],
+)
+def test_accepts_zero_and_positive_numeric_values(value, expected):
+    result = config_module.get_image_list_job_timeout_seconds(_config_with_timeout(value))
+
+    assert result == expected
+    assert isinstance(result, float)
+
+
+@pytest.mark.parametrize(
+    "value",
+    [
+        True,
+        False,
+        -1,
+        -0.01,
+        float("nan"),
+        float("inf"),
+        float("-inf"),
+        "not-a-number",
+        "",
+        "   ",
+        "-1",
+        "nan",
+        "inf",
+        None,
+        [],
+        {},
+    ],
+)
+def test_rejects_invalid_values_with_setting_name(value):
+    with pytest.raises(ValueError, match=SETTING_KEY):
+        config_module.get_image_list_job_timeout_seconds(_config_with_timeout(value))

tests/test_image_list_websocket_deadline.py

@@ -0,0 +1,212 @@
+import asyncio
+import importlib.util
+import sys
+import types
+import unittest
+import uuid
+from pathlib import Path
+from typing import Any, cast
+from unittest.mock import patch
+
+import aiohttp
+from aiohttp import web
+
+
+def _load_dispatch_module():
+    module_path = Path(__file__).resolve().parents[1] / "api" / "orchestration" / "dispatch.py"
+    package_name = f"image_list_ws_deadline_testpkg_{uuid.uuid4().hex}"
+    package_modules = {}
+
+    for suffix in ("", ".api", ".api.orchestration", ".utils"):
+        package = types.ModuleType(package_name + suffix)
+        package.__path__ = []
+        package_modules[package_name + suffix] = package
+
+    logging_module = types.ModuleType(f"{package_name}.utils.logging")
+    setattr(logging_module, "debug_log", lambda *_args, **_kwargs: None)
+    setattr(logging_module, "log", lambda *_args, **_kwargs: None)
+    package_modules[logging_module.__name__] = logging_module
+
+    network_module = types.ModuleType(f"{package_name}.utils.network")
+    setattr(network_module, "build_worker_url", lambda *_args, **_kwargs: "http://example.invalid")
+
+    async def _unused_probe_worker(*_args, **_kwargs):
+        return None
+
+    async def _unused_get_client_session(*_args, **_kwargs):
+        raise RuntimeError("get_client_session must be patched by this test")
+
+    setattr(network_module, "probe_worker", _unused_probe_worker)
+    setattr(network_module, "get_client_session", _unused_get_client_session)
+    package_modules[network_module.__name__] = network_module
+
+    module_name = f"{package_name}.api.orchestration.dispatch"
+    spec = importlib.util.spec_from_file_location(module_name, module_path)
+    assert spec is not None and spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    package_modules[module_name] = module
+
+    previous_modules = {name: sys.modules.get(name) for name in package_modules}
+    try:
+        sys.modules.update(package_modules)
+        spec.loader.exec_module(module)
+    finally:
+        for name, previous in previous_modules.items():
+            if previous is None:
+                sys.modules.pop(name, None)
+            else:
+                sys.modules[name] = previous
+    return module
+
+
+dispatch = _load_dispatch_module()
+
+
+class WebsocketDispatchDeadlineTests(unittest.IsolatedAsyncioTestCase):
+    async def asyncSetUp(self):
+        self.mode = None
+        self.received_payload = None
+        self.dispatch_received = asyncio.Event()
+        self.peer_finished = asyncio.Event()
+        self.irrelevant_messages_sent = 0
+
+        app = web.Application()
+        app.router.add_get("/distributed/worker_ws", self._handle_websocket)
+        self.runner = web.AppRunner(app)
+        await self.runner.setup()
+        self.site = web.TCPSite(self.runner, "127.0.0.1", 0)
+        await self.site.start()
+        server = cast(Any, self.site._server)
+        port = server.sockets[0].getsockname()[1]
+        self.worker_url = f"http://127.0.0.1:{port}"
+        self.session = aiohttp.ClientSession()
+
+    async def asyncTearDown(self):
+        await self.session.close()
+        await self.runner.cleanup()
+
+    async def _get_client_session(self):
+        return self.session
+
+    async def _handle_websocket(self, request):
+        ws = web.WebSocketResponse()
+        await ws.prepare(request)
+        try:
+            payload = await ws.receive_json()
+            self.received_payload = payload
+            self.dispatch_received.set()
+
+            if self.mode in ("success", "slow-close"):
+                await ws.send_json({
+                    "type": "dispatch_ack",
+                    "request_id": "unrelated-request",
+                    "ok": True,
+                    "prompt_id": "wrong-prompt",
+                })
+                await ws.send_json({"type": "heartbeat"})
+                await ws.send_json({
+                    "type": "dispatch_ack",
+                    "request_id": payload["request_id"],
+                    "ok": True,
+                    "prompt_id": "prompt-123",
+                })
+                if self.mode == "slow-close":
+                    await asyncio.sleep(0.15)
+            elif self.mode == "rejection":
+                await ws.send_json({
+                    "type": "dispatch_ack",
+                    "request_id": payload["request_id"],
+                    "ok": False,
+                    "error": "Invalid prompt",
+                    "validation_error": {"node": "bad"},
+                    "node_errors": {"bad": "missing input"},
+                })
+            elif self.mode in ("no-ack", "wait-for-cancel"):
+                async def send_irrelevant_messages():
+                    while True:
+                        try:
+                            await ws.send_json({"type": "heartbeat"})
+                            self.irrelevant_messages_sent += 1
+                            await ws.send_json({
+                                "type": "dispatch_ack",
+                                "request_id": "unrelated-request",
+                                "ok": True,
+                                "prompt_id": "not-this-request",
+                            })
+                            self.irrelevant_messages_sent += 1
+                        except (ConnectionError, RuntimeError):
+                            return
+                        await asyncio.sleep(0.01)
+
+                sender = asyncio.create_task(send_irrelevant_messages())
+                try:
+                    async for _message in ws:
+                        pass
+                finally:
+                    sender.cancel()
+                    try:
+                        await sender
+                    except asyncio.CancelledError:
+                        pass
+        finally:
+            self.peer_finished.set()
+        return ws
+
+    async def _dispatch(self, timeout):
+        with patch.object(dispatch, "get_client_session", self._get_client_session):
+            return await dispatch._dispatch_via_websocket(
+                self.worker_url,
+                {"prompt": {"node": {"class_type": "Example"}}, "workflow": {"name": "test"}},
+                "client-456",
+                timeout=timeout,
+            )
+
+    async def test_returns_matching_acknowledgement_prompt_id(self):
+        self.mode = "success"
+
+        prompt_id = await self._dispatch(timeout=0.5)
+
+        self.assertEqual(prompt_id, "prompt-123")
+        received_payload = self.received_payload
+        assert received_payload is not None
+        self.assertEqual(received_payload["type"], "dispatch_prompt")
+        self.assertTrue(received_payload["request_id"])
+        self.assertEqual(received_payload["client_id"], "client-456")
+        self.assertEqual(received_payload["prompt"], {"node": {"class_type": "Example"}})
+        self.assertEqual(received_payload["workflow"], {"name": "test"})
+
+    async def test_acknowledged_prompt_id_survives_slow_close(self):
+        self.mode = "slow-close"
+        self.assertEqual(await self._dispatch(timeout=0.03), "prompt-123")
+
+    async def test_matching_rejection_preserves_worker_error_details(self):
+        self.mode = "rejection"
+
+        with self.assertRaisesRegex(RuntimeError, "Invalid prompt.*validation_error=.*node_errors="):
+            await self._dispatch(timeout=0.5)
+
+    async def test_repeated_irrelevant_messages_do_not_reset_ack_deadline(self):
+        self.mode = "no-ack"
+        started = asyncio.get_running_loop().time()
+
+        with self.assertRaisesRegex(RuntimeError, "submission status is ambiguous"):
+            await asyncio.wait_for(self._dispatch(timeout=0.12), timeout=0.5)
+
+        elapsed = asyncio.get_running_loop().time() - started
+        self.assertLess(elapsed, 0.5)
+        self.assertGreaterEqual(self.irrelevant_messages_sent, 4)
+        await asyncio.wait_for(self.peer_finished.wait(), timeout=0.5)
+
+    async def test_outer_cancellation_propagates_and_closes_peer(self):
+        self.mode = "wait-for-cancel"
+        task = asyncio.create_task(self._dispatch(timeout=2.0))
+        await asyncio.wait_for(self.dispatch_received.wait(), timeout=0.5)
+
+        task.cancel()
+        with self.assertRaises(asyncio.CancelledError):
+            await asyncio.wait_for(task, timeout=0.5)
+        await asyncio.wait_for(self.peer_finished.wait(), timeout=0.5)
+
+
+if __name__ == "__main__":
+    unittest.main()

utils/config.py

@@ -5,6 +5,7 @@ import asyncio
 import copy
 import os
 import json
+import math
 from contextlib import asynccontextmanager
 from .logging import log
 from .secret_fields import is_secret_field_name as _is_secret_field_name
@@ -81,7 +82,8 @@ def get_default_config():
             "worker_probe_concurrency": 8,
             "worker_prep_concurrency": 4,
             "media_sync_concurrency": 2,
-            "media_sync_timeout_seconds": 120
+            "media_sync_timeout_seconds": 120,
+            "image_list_job_timeout_seconds": 0.0
         },
         "tunnel": {
             "status": "stopped",
@@ -215,6 +217,37 @@ def get_worker_timeout_seconds(default: int = HEARTBEAT_TIMEOUT) -> int:
         return max(1, int(default))
 
 
+def get_image_list_job_timeout_seconds(config=None) -> float:
+    """Return the validated image-list total-job timeout in seconds.
+
+    A value of zero disables the total-job deadline. Missing settings default
+    to zero; configured values must be finite, non-negative numbers.
+    """
+    setting_key = "settings.image_list_job_timeout_seconds"
+    cfg = load_config() if config is None else config
+    if not isinstance(cfg, dict):
+        raise ValueError(f"{setting_key} must be a finite non-negative number")
+
+    settings = cfg.get("settings", {})
+    if not isinstance(settings, dict):
+        raise ValueError(f"{setting_key} must be a finite non-negative number")
+
+    value = settings.get("image_list_job_timeout_seconds", 0.0)
+    if isinstance(value, bool) or not isinstance(value, (int, float, str)):
+        raise ValueError(f"{setting_key} must be a finite non-negative number")
+    if isinstance(value, str) and not value.strip():
+        raise ValueError(f"{setting_key} must be a finite non-negative number")
+
+    try:
+        timeout = float(value)
+    except (OverflowError, ValueError):
+        raise ValueError(f"{setting_key} must be a finite non-negative number") from None
+
+    if not math.isfinite(timeout) or timeout < 0:
+        raise ValueError(f"{setting_key} must be a finite non-negative number")
+    return timeout
+
+
 def is_master_delegate_only() -> bool:
     """Returns True when master should skip local workload and act as orchestrator only."""
     try: