tool: extract helpers

1 files changed, 282 insertions, 204 deletions

diff --git a/comfier.py b/comfier.py
index c32b0d7..dc37173 100644
--- a/comfier.py
+++ b/comfier.py
@@ -55,23 +55,6 @@ class Tools:
     def __init__(self):
         self.valves = self.Valves()
 
-    async def comfier_generate_seed(
-        self, bits: Annotated[int, Field(ge=8, le=63)] = 63
-    ) -> str:
-        """
-        Generate a new random seed for image generation.
-
-        Calling this is optional: comfier_generate_image picks a random seed itself and reports it in its result.
-        Use this tool only when a seed is needed before generating, for example to create several images with distinct, known seeds.
-        Returns the seed as a decimal string, which can be passed as the seed of comfier_generate_image.
-
-        :param bits: Maximum bit-width of the seed, between 8 and 63. Larger values allow larger seeds. Defaults to 63.
-        :return: The generated seed as a decimal string.
-        """
-
-        clamped = max(8, min(int(bits), 63))
-        return f"{random.getrandbits(clamped)}"
-
     async def comfier_generate_image(
         self,
         positive_prompt: str,
@@ -199,6 +182,23 @@ class Tools:
             "images": image_urls,
         }
 
+    async def comfier_generate_seed(
+        self, bits: Annotated[int, Field(ge=8, le=63)] = 63
+    ) -> str:
+        """
+        Generate a new random seed for image generation.
+
+        Calling this is optional: comfier_generate_image picks a random seed itself and reports it in its result.
+        Use this tool only when a seed is needed before generating, for example to create several images with distinct, known seeds.
+        Returns the seed as a decimal string, which can be passed as the seed of comfier_generate_image.
+
+        :param bits: Maximum bit-width of the seed, between 8 and 63. Larger values allow larger seeds. Defaults to 63.
+        :return: The generated seed as a decimal string.
+        """
+
+        clamped = max(8, min(int(bits), 63))
+        return f"{random.getrandbits(clamped)}"
+
     def _apply_workflow_parameters(
         self,
         workflow: dict[str, Any],
@@ -269,51 +269,110 @@ class Tools:
         except Exception:
             pass
 
-    async def _verify_prompt_alive(
+    async def _collect_outputs(
         self,
+        ws: aiohttp.ClientWebSocketResponse,
         session: aiohttp.ClientSession,
         http_base: str,
+        workflow: dict[str, Any],
         prompt_id: str,
-    ) -> bool:
+        event_emitter: Optional[EventEmitter],
+    ) -> list[dict[str, Any]]:
         """
-        Verify with ComfyUI that the backend is reachable and still tracking prompt_id
-        in either /queue (running/pending) or /history.
+        Follow the ComfyUI WebSocket until the prompt finishes, emitting progress.
+
+        Returns the image descriptors reported by the workflow's output nodes.
+        Raises ConnectionError if the backend disappears, RuntimeError if the
+        job is interrupted or fails.
         """
 
-        timeout = aiohttp.ClientTimeout(total=3.0)
-        try:
-            async with session.get(f"{http_base}/queue", timeout=timeout) as q_resp:
-                q_resp.raise_for_status()
-                q_data = await q_resp.json()
+        output_images: list[dict[str, Any]] = []
+        last_emit_time = 0.0
+        idle_ticks = 0
 
-            running = q_data.get("queue_running") or []
-            pending = q_data.get("queue_pending") or []
-            in_queue = any(
-                isinstance(item, list) and len(item) > 1 and item[1] == prompt_id
-                for item in (running + pending)
-            )
-            if in_queue:
-                return True
+        while True:
+            try:
+                msg = await asyncio.wait_for(ws.receive(), timeout=1.0)
+            except asyncio.TimeoutError:
+                idle_ticks += 1
+                if idle_ticks >= 5:
+                    idle_ticks = 0
+                    await self._ensure_prompt_alive(session, http_base, prompt_id)
+                continue
 
-            async with session.get(
-                f"{http_base}/history/{prompt_id}", timeout=timeout
-            ) as h_resp:
-                h_resp.raise_for_status()
-                h_data = await h_resp.json() or {}
-                if prompt_id not in h_data:
-                    return False
+            idle_ticks = 0
 
-                status_info = (h_data[prompt_id] or {}).get("status") or {}
-                if status_info.get("status_str") == "error":
-                    raise RuntimeError(
-                        "The job was terminated with an error or interruption on the backend."
-                    )
-                return True
+            if msg.type != aiohttp.WSMsgType.TEXT:
+                self._raise_if_ws_closed(ws, msg.type)
+                continue
 
-        except aiohttp.ClientError as err:
-            raise ConnectionError(
-                f"Lost connection to ComfyUI backend ({err})."
-            ) from err
+            payload = json.loads(msg.data) or {}
+            msg_type = payload.get("type")
+            data = payload.get("data") or {}
+
+            if data.get("prompt_id") and data.get("prompt_id") != prompt_id:
+                continue
+
+            if msg_type == "executing":
+                if data.get("node") is None:
+                    return output_images
+                await self._emit_node_status(event_emitter, workflow, str(data["node"]))
+
+            elif msg_type == "progress":
+                last_emit_time = await self._emit_progress(
+                    event_emitter, data, last_emit_time
+                )
+
+            elif msg_type == "executed":
+                images = (data.get("output") or {}).get("images") or []
+                if isinstance(images, list):
+                    output_images.extend(img for img in images if isinstance(img, dict))
+
+            elif msg_type == "execution_interrupted":
+                raise RuntimeError("Generation was interrupted on the backend.")
+
+            elif msg_type == "execution_error":
+                error_message = data.get("exception_message", "Unknown ComfyUI error")
+                raise RuntimeError(f"ComfyUI execution error: {error_message}")
+
+    async def _emit_node_status(
+        self,
+        event_emitter: Optional[EventEmitter],
+        workflow: dict[str, Any],
+        node_id: str,
+    ) -> None:
+        """Announce the workflow node ComfyUI has started executing."""
+
+        node_info = workflow.get(node_id) or {}
+        node_title = (
+            (node_info.get("_meta") or {}).get("title")
+            or node_info.get("class_type")
+            or f"Node {node_id}"
+        )
+        await self._emit_status(event_emitter, f"Running: {node_title}...", done=False)
+
+    async def _emit_progress(
+        self,
+        event_emitter: Optional[EventEmitter],
+        data: dict[str, Any],
+        last_emit_time: float,
+    ) -> float:
+        """Emit a throttled sampling progress update. Returns the new last-emit time."""
+
+        val = data.get("value", 0)
+        maximum = data.get("max", 1)
+        now = time.monotonic()
+
+        if (now - last_emit_time) < 0.25 and val != maximum:
+            return last_emit_time
+
+        percentage = int((val / maximum) * 100) if maximum else 0
+        await self._emit_status(
+            event_emitter,
+            f"Sampling: step {val}/{maximum} ({percentage}%)",
+            done=False,
+        )
+        return now
 
     async def _emit_status(
         self,
@@ -337,6 +396,16 @@ class Tools:
             }
         )
 
+    async def _ensure_prompt_alive(
+        self, session: aiohttp.ClientSession, http_base: str, prompt_id: str
+    ) -> None:
+        """Raise ConnectionError if ComfyUI no longer tracks the prompt."""
+
+        if not await self._verify_prompt_alive(session, http_base, prompt_id):
+            raise ConnectionError(
+                f"ComfyUI backend lost prompt {prompt_id} (backend may have restarted or crashed)."
+            )
+
     async def _execute(
         self,
         workflow: dict[str, Any],
@@ -352,11 +421,8 @@ class Tools:
         headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
         client_id = str(uuid.uuid4())
 
-        output_images: list[dict[str, Any]] = []
         stored_urls: list[str] = []
-        last_emit_time = 0.0
         prompt_id: Optional[str] = None
-        idle_ticks = 0
 
         try:
             async with aiohttp.ClientSession(headers=headers) as session:
@@ -367,113 +433,15 @@ class Tools:
                     await self._emit_status(
                         event_emitter, "Submitting workflow to ComfyUI...", done=False
                     )
-
-                    async with session.post(
-                        f"{http_base}/prompt",
-                        json={"prompt": workflow, "client_id": client_id},
-                    ) as response:
-                        if response.status >= 400:
-                            error_body = await response.text()
-                            raise RuntimeError(
-                                f"ComfyUI rejected workflow (HTTP {response.status}): {error_body}"
-                            )
-                        queued = await response.json()
-                        prompt_id = queued["prompt_id"]
-
+                    prompt_id = await self._submit_workflow(
+                        session, http_base, workflow, client_id
+                    )
                     await self._emit_status(
                         event_emitter, "Waiting for ComfyUI queue...", done=False
                     )
-
-                    while True:
-                        try:
-                            msg = await asyncio.wait_for(ws.receive(), timeout=1.0)
-                            idle_ticks = 0
-                        except asyncio.TimeoutError:
-                            idle_ticks += 1
-                            if idle_ticks >= 5 and prompt_id:
-                                idle_ticks = 0
-                                is_alive = await self._verify_prompt_alive(
-                                    session, http_base, prompt_id
-                                )
-                                if not is_alive:
-                                    raise ConnectionError(
-                                        f"ComfyUI backend lost prompt {prompt_id} (backend may have restarted or crashed)."
-                                    ) from None
-                            continue
-
-                        if msg.type != aiohttp.WSMsgType.TEXT:
-                            if msg.type in (
-                                aiohttp.WSMsgType.CLOSE,
-                                aiohttp.WSMsgType.CLOSING,
-                                aiohttp.WSMsgType.CLOSED,
-                                aiohttp.WSMsgType.ERROR,
-                            ):
-                                ws_err = ws.exception()
-                                detail = (
-                                    f": {ws_err}"
-                                    if ws_err
-                                    else f" (code {ws.close_code})"
-                                )
-                                raise ConnectionError(
-                                    f"ComfyUI WebSocket disconnected unexpectedly before completion{detail}."
-                                )
-                            continue
-
-                        payload = json.loads(msg.data) or {}
-                        msg_type = payload.get("type")
-                        data = payload.get("data") or {}
-
-                        if data.get("prompt_id") and data.get("prompt_id") != prompt_id:
-                            continue
-
-                        if msg_type == "executing":
-                            node_id = data.get("node")
-                            if node_id is None:
-                                break
-                            node_info = workflow.get(str(node_id)) or {}
-                            node_meta = node_info.get("_meta") or {}
-                            node_title = (
-                                node_meta.get("title")
-                                or node_info.get("class_type")
-                                or f"Node {node_id}"
-                            )
-                            await self._emit_status(
-                                event_emitter, f"Running: {node_title}...", done=False
-                            )
-
-                        elif msg_type == "progress":
-                            val = data.get("value", 0)
-                            maximum = data.get("max", 1)
-                            percentage = int((val / maximum) * 100) if maximum else 0
-                            now = time.monotonic()
-                            if (now - last_emit_time) >= 0.25 or val == maximum:
-                                last_emit_time = now
-                                await self._emit_status(
-                                    event_emitter,
-                                    f"Sampling: step {val}/{maximum} ({percentage}%)",
-                                    done=False,
-                                )
-
-                        elif msg_type == "executed":
-                            output_data = data.get("output") or {}
-                            images = output_data.get("images") or []
-                            if isinstance(images, list):
-                                output_images.extend(
-                                    img for img in images if isinstance(img, dict)
-                                )
-
-                        elif msg_type == "execution_interrupted":
-                            raise RuntimeError(
-                                "Generation was interrupted on the backend."
-                            )
-
-                        elif msg_type == "execution_error":
-                            error_message = data.get(
-                                "exception_message", "Unknown ComfyUI error"
-                            )
-                            raise RuntimeError(
-                                f"ComfyUI execution error: {error_message}"
-                            )
+                    output_images = await self._collect_outputs(
+                        ws, session, http_base, workflow, prompt_id, event_emitter
+                    )
 
                 if not output_images:
                     raise RuntimeError(
@@ -485,53 +453,11 @@ class Tools:
                 )
 
                 for img in output_images:
-                    filename = img.get("filename")
-                    if not filename:
-                        continue
-
-                    params = {
-                        "filename": filename,
-                        "subfolder": img.get("subfolder", ""),
-                        "type": img.get("type", "output"),
-                    }
-
-                    async with session.get(
-                        f"{http_base}/view", params=params
-                    ) as img_response:
-                        img_response.raise_for_status()
-                        image_bytes = await img_response.read()
-                        content_type = (
-                            img_response.headers.get("Content-Type")
-                            or mimetypes.guess_type(filename)[0]
-                            or "image/png"
-                        )
-
-                    uploaded: Any = upload_image(
-                        request=request,
-                        image_data=image_bytes,
-                        content_type=content_type,
-                        metadata=metadata,
-                        user=user,
+                    url = await self._store_image(
+                        session, http_base, img, request, user, metadata
                     )
-
-                    if inspect.isawaitable(uploaded):
-                        uploaded = await uploaded
-
-                    if isinstance(uploaded, tuple):
-                        uploaded = uploaded[-1]
-
-                    if isinstance(uploaded, str):
-                        url = uploaded
-                    elif isinstance(uploaded, dict) and "id" in uploaded:
-                        url = f"/api/v1/files/{uploaded['id']}/content"
-                    elif hasattr(uploaded, "id"):
-                        url = f"/api/v1/files/{uploaded.id}/content"  # pyright: ignore[reportAttributeAccessIssue]
-                    else:
-                        raise RuntimeError(
-                            f"Unexpected return type from upload_image: {type(uploaded)!r}"
-                        )
-
-                    stored_urls.append(url)
+                    if url:
+                        stored_urls.append(url)
 
         except asyncio.CancelledError:
             await asyncio.shield(self._cancel(http_base, headers, prompt_id))
@@ -635,6 +561,25 @@ class Tools:
 
         return parsed & ((1 << 63) - 1)
 
+    def _raise_if_ws_closed(
+        self, ws: aiohttp.ClientWebSocketResponse, msg_type: aiohttp.WSMsgType
+    ) -> None:
+        """Raise ConnectionError if a non-text WebSocket message signals a closed connection."""
+
+        if msg_type not in (
+            aiohttp.WSMsgType.CLOSE,
+            aiohttp.WSMsgType.CLOSING,
+            aiohttp.WSMsgType.CLOSED,
+            aiohttp.WSMsgType.ERROR,
+        ):
+            return
+
+        ws_err = ws.exception()
+        detail = f": {ws_err}" if ws_err else f" (code {ws.close_code})"
+        raise ConnectionError(
+            f"ComfyUI WebSocket disconnected unexpectedly before completion{detail}."
+        )
+
     async def _resolve_dimensions(self, size: Optional[str]) -> tuple[int, int]:
         """Parse 'WIDTHxHEIGHT' from the size parameter or fall back to global IMAGE_SIZE."""
 
@@ -719,3 +664,136 @@ class Tools:
             raw_steps = await Config.get("image_generation.steps", 25)
 
         return int(raw_steps)
+
+    async def _store_image(
+        self,
+        session: aiohttp.ClientSession,
+        http_base: str,
+        img: dict[str, Any],
+        request: Request,
+        user: Any,
+        metadata: dict[str, Any],
+    ) -> Optional[str]:
+        """
+        Download one generated image from ComfyUI and store it in Open WebUI.
+
+        Returns the stored image's URL, or None if the descriptor has no filename.
+        """
+
+        filename = img.get("filename")
+        if not filename:
+            return None
+
+        params = {
+            "filename": filename,
+            "subfolder": img.get("subfolder", ""),
+            "type": img.get("type", "output"),
+        }
+
+        async with session.get(f"{http_base}/view", params=params) as img_response:
+            img_response.raise_for_status()
+            image_bytes = await img_response.read()
+            content_type = (
+                img_response.headers.get("Content-Type")
+                or mimetypes.guess_type(filename)[0]
+                or "image/png"
+            )
+
+        uploaded: Any = upload_image(
+            request=request,
+            image_data=image_bytes,
+            content_type=content_type,
+            metadata=metadata,
+            user=user,
+        )
+
+        if inspect.isawaitable(uploaded):
+            uploaded = await uploaded
+
+        return self._upload_result_to_url(uploaded)
+
+    async def _submit_workflow(
+        self,
+        session: aiohttp.ClientSession,
+        http_base: str,
+        workflow: dict[str, Any],
+        client_id: str,
+    ) -> str:
+        """Queue the workflow on ComfyUI and return the new prompt ID."""
+
+        async with session.post(
+            f"{http_base}/prompt",
+            json={"prompt": workflow, "client_id": client_id},
+        ) as response:
+            if response.status >= 400:
+                error_body = await response.text()
+                raise RuntimeError(
+                    f"ComfyUI rejected workflow (HTTP {response.status}): {error_body}"
+                )
+            queued = await response.json()
+            return queued["prompt_id"]
+
+    def _upload_result_to_url(self, uploaded: Any) -> str:
+        """Turn the return value of Open WebUI's upload_image into a file content URL."""
+
+        if isinstance(uploaded, tuple):
+            uploaded = uploaded[-1]
+
+        if isinstance(uploaded, str):
+            return uploaded
+
+        if isinstance(uploaded, dict) and "id" in uploaded:
+            return f"/api/v1/files/{uploaded['id']}/content"
+
+        if hasattr(uploaded, "id"):
+            return f"/api/v1/files/{uploaded.id}/content"  # pyright: ignore[reportAttributeAccessIssue]
+
+        raise RuntimeError(
+            f"Unexpected return type from upload_image: {type(uploaded)!r}"
+        )
+
+    async def _verify_prompt_alive(
+        self,
+        session: aiohttp.ClientSession,
+        http_base: str,
+        prompt_id: str,
+    ) -> bool:
+        """
+        Verify with ComfyUI that the backend is reachable and still tracking prompt_id
+        in either /queue (running/pending) or /history.
+        """
+
+        timeout = aiohttp.ClientTimeout(total=3.0)
+        try:
+            async with session.get(f"{http_base}/queue", timeout=timeout) as q_resp:
+                q_resp.raise_for_status()
+                q_data = await q_resp.json()
+
+            running = q_data.get("queue_running") or []
+            pending = q_data.get("queue_pending") or []
+            in_queue = any(
+                isinstance(item, list) and len(item) > 1 and item[1] == prompt_id
+                for item in (running + pending)
+            )
+            if in_queue:
+                return True
+
+            async with session.get(
+                f"{http_base}/history/{prompt_id}", timeout=timeout
+            ) as h_resp:
+                h_resp.raise_for_status()
+                h_data = await h_resp.json() or {}
+                if prompt_id not in h_data:
+                    return False
+
+                status_info = (h_data[prompt_id] or {}).get("status") or {}
+                if status_info.get("status_str") == "error":
+                    raise RuntimeError(
+                        "The job was terminated with an error or interruption on the backend."
+                    )
+                return True
+
+        except aiohttp.ClientError as err:
+            raise ConnectionError(
+                f"Lost connection to ComfyUI backend ({err})."
+            ) from err