| """In-memory handoff from API callers to a single ComfyUI browser page.""" |
|
|
| import hashlib |
| import json |
| import re |
| import time |
| import uuid |
| from collections.abc import Callable |
|
|
| from aiohttp import web |
|
|
|
|
| class WorkflowSubmissions: |
| def __init__(self, connected: Callable[[str], bool]) -> None: |
| self.connected = connected |
| self.clients: dict[str, dict] = {} |
| self.jobs: dict[str, dict] = {} |
|
|
| def _prune(self) -> None: |
| now = time.time() |
| self.clients = {key: value for key, value in self.clients.items() |
| if now - value["last_seen"] < 30 and self.connected(value["session_id"])} |
| for job in self.jobs.values(): |
| if job["status"] == "pending" and now - job["created_at"] > 60: |
| job.update(status="failed", error="Browser did not accept the workflow within 60 seconds") |
| job.pop("workflow", None) |
| if job["status"] == "loading" and now - job["claimed_at"] > 120: |
| job.update(status="unknown", error="Browser result unavailable; check queue/history before retrying") |
| self.jobs = {key: value for key, value in self.jobs.items() |
| if now - value["created_at"] < 3600 or value["status"] in {"loading", "unknown"}} |
|
|
| def list_clients(self) -> list[dict]: |
| self._prune() |
| return list(self.clients.values()) |
|
|
| @staticmethod |
| def _identifier(value: object, field: str) -> str: |
| if not isinstance(value, str) or not re.fullmatch(r"[A-Za-z0-9_-]{1,128}", value): |
| raise ValueError(f"{field} must be 1-128 letters, digits, underscores or hyphens") |
| return value |
|
|
| @staticmethod |
| def summary(job: dict) -> dict: |
| return {key: value for key, value in job.items() |
| if key not in {"workflow", "fingerprint", "claim_token"}} |
|
|
| def submit(self, payload: dict) -> dict: |
| self._prune() |
| workflow = payload.get("workflow") |
| if not isinstance(workflow, dict) or not isinstance(workflow.get("nodes"), list) or not isinstance(workflow.get("links"), list): |
| raise ValueError("workflow must be UI workflow JSON with nodes and links; API prompt alone cannot preserve the editor layout") |
| if any(not isinstance(node, dict) or "id" not in node or not isinstance(node.get("type"), str) for node in workflow["nodes"]): |
| raise ValueError("Every workflow node must have an id and type") |
| mode = payload.get("mode", "new_tab") |
| if mode not in ("new_tab", "replace"): |
| raise ValueError("mode must be new_tab or replace") |
| auto_queue = payload.get("auto_queue", True) |
| if type(auto_queue) is not bool: |
| raise ValueError("auto_queue must be a boolean") |
| name = payload.get("name", "Skill workflow") |
| if not isinstance(name, str) or not name.strip() or len(name) > 200 or any(c in name for c in "/\\\x00"): |
| raise ValueError("name must be a filename without directory separators, up to 200 characters") |
| request_id = self._identifier(payload.get("request_id", str(uuid.uuid4())), "request_id") |
| client_id = payload.get("client_id") |
| if client_id is not None: |
| self._identifier(client_id, "client_id") |
| fingerprint = hashlib.sha256(json.dumps( |
| [workflow, mode, auto_queue, name, client_id], sort_keys=True, allow_nan=False, |
| ).encode()).hexdigest() |
| previous = self.jobs.get(request_id) |
| if previous: |
| if previous["fingerprint"] != fingerprint: |
| raise web.HTTPConflict(text="request_id was already used with a different payload") |
| return self.summary(previous) |
| if client_id is None: |
| if len(self.clients) > 1: |
| raise web.HTTPConflict(text="Multiple browser pages available; select client_id from GET /easy-media/workflow/clients") |
| client_id = next(iter(self.clients), None) |
| if client_id not in self.clients: |
| raise web.HTTPServiceUnavailable(text="Open or refresh a ComfyUI page with the Easy Media extension first") |
| if any(job["client_id"] == client_id and job["status"] in {"pending", "loading", "unknown"} for job in self.jobs.values()): |
| raise web.HTTPConflict(text="This browser already has an unfinished submission; check its status before retrying") |
| if len(self.jobs) >= 256: |
| raise web.HTTPServiceUnavailable(text="Submission history is full; retry after old requests expire") |
| job = dict(request_id=request_id, client_id=client_id, workflow=workflow, |
| mode=mode, auto_queue=auto_queue, name=name, status="pending", |
| created_at=time.time(), fingerprint=fingerprint) |
| self.jobs[request_id] = job |
| return self.summary(job) |
|
|
| def poll(self, payload: dict) -> dict | None: |
| self._prune() |
| client_id = self._identifier(payload.get("client_id"), "client_id") |
| session_id = self._identifier(payload.get("session_id"), "session_id") |
| if not self.connected(session_id): |
| raise web.HTTPConflict(text="ComfyUI WebSocket is not connected") |
| title = payload.get("title", "ComfyUI") |
| if not isinstance(title, str) or len(title) > 300: |
| raise ValueError("title must be a string up to 300 characters") |
| self.clients[client_id] = dict(client_id=client_id, session_id=session_id, title=title, last_seen=time.time()) |
| for job in self.jobs.values(): |
| if job["client_id"] == client_id and job["status"] == "pending": |
| job.update(status="loading", claimed_at=time.time(), claim_token=str(uuid.uuid4())) |
| delivery = {**self.summary(job), "workflow": job.pop("workflow"), "claim_token": job["claim_token"]} |
| return delivery |
| return None |
|
|
| def get(self, request_id: str) -> dict: |
| self._prune() |
| if request_id not in self.jobs: |
| raise web.HTTPNotFound(text="Unknown or expired request_id") |
| return self.summary(self.jobs[request_id]) |
|
|
| def complete(self, request_id: str, payload: dict) -> dict: |
| self.get(request_id) |
| job = self.jobs[request_id] |
| if not job.get("claim_token") or payload.get("claim_token") != job["claim_token"]: |
| raise web.HTTPForbidden(text="Invalid claim token") |
| status = payload.get("status") |
| if status not in ("loaded", "queued", "failed", "unknown"): |
| raise ValueError("Invalid result status") |
| result = {"status": status} |
| if status == "queued": |
| result["prompt_id"] = self._identifier(payload.get("prompt_id"), "prompt_id") |
| if status in ("failed", "unknown"): |
| error = payload.get("error") |
| if not isinstance(error, str) or not error or len(error) > 4000: |
| raise ValueError("error must be a non-empty string up to 4000 characters") |
| result["error"] = error |
| if job["status"] not in {"loading", "unknown"}: |
| if all(job.get(key) == value for key, value in result.items()): |
| return self.summary(job) |
| raise web.HTTPConflict(text="Submission already finished") |
| job.pop("error", None) |
| job.update(result) |
| return self.summary(job) |
|
|
|
|
| def register_workflow_routes(routes: web.RouteTableDef, connected: Callable[[str], bool]) -> None: |
| submissions = WorkflowSubmissions(connected) |
|
|
| async def read_payload(request: web.Request) -> dict: |
| |
| if request.content_type != "application/json": |
| raise web.HTTPUnsupportedMediaType(text="Content-Type must be application/json") |
| origin = request.headers.get("Origin") |
| if origin and origin != f"{request.scheme}://{request.host}": |
| raise web.HTTPForbidden(text="Cross-origin workflow control is not allowed") |
| try: |
| payload = await request.json() |
| except (ValueError, UnicodeError) as error: |
| raise web.HTTPBadRequest(text="Invalid JSON") from error |
| if not isinstance(payload, dict): |
| raise web.HTTPBadRequest(text="Request must contain a JSON object") |
| return payload |
|
|
| async def respond(action: Callable[..., object], *args: object) -> web.Response: |
| try: |
| return web.json_response(action(*args)) |
| except ValueError as error: |
| return web.json_response({"error": str(error)}, status=400) |
| except web.HTTPException as error: |
| return web.json_response({"error": error.text}, status=error.status) |
|
|
| @routes.get("/easy-media/workflow/clients") |
| async def clients(_request: web.Request) -> web.Response: |
| return web.json_response({"clients": submissions.list_clients()}) |
|
|
| @routes.post("/easy-media/workflow/submit") |
| async def submit(request: web.Request) -> web.Response: |
| return await respond(submissions.submit, await read_payload(request)) |
|
|
| @routes.post("/easy-media/workflow/poll") |
| async def poll(request: web.Request) -> web.Response: |
| return await respond(lambda data: {"job": submissions.poll(data)}, await read_payload(request)) |
|
|
| @routes.get("/easy-media/workflow/submissions/{request_id}") |
| async def status(request: web.Request) -> web.Response: |
| return await respond(submissions.get, request.match_info["request_id"]) |
|
|
| @routes.post("/easy-media/workflow/submissions/{request_id}/result") |
| async def result(request: web.Request) -> web.Response: |
| return await respond(submissions.complete, request.match_info["request_id"], await read_payload(request)) |
|
|
