mirror of
https://github.com/rookiestar28/ComfyUI-OpenClaw.git
synced 2026-08-14 00:48:07 +00:00
260 lines
9.3 KiB
Python
260 lines
9.3 KiB
Python
"""
|
|
Queue Submit Service (F5 + R33).
|
|
Submits prompt workflows to ComfyUI execution queue with execution budgets.
|
|
|
|
- Uses internal HTTP call to POST /prompt
|
|
- Handles client_id and extra metadata
|
|
- R33: Applies concurrency caps and render size budgets
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
import uuid
|
|
from typing import Any, Dict, Optional
|
|
|
|
try:
|
|
from api.errors import APIError, ErrorCode, create_error_response
|
|
except ImportError:
|
|
# Fallback if api module not found (e.g. some test environments)
|
|
# Define minimal mocks to avoid crash
|
|
class ErrorCode:
|
|
DEPENDENCY_UNAVAILABLE = "dependency_unavailable"
|
|
QUEUE_SUBMIT_FAILED = "queue_submit_failed"
|
|
INTERNAL_ERROR = "internal_error"
|
|
|
|
class APIError(Exception):
|
|
def __init__(self, message, code="internal_error", status=500, detail=None):
|
|
super().__init__(message)
|
|
self.code = code
|
|
self.status = status
|
|
self.detail = detail or {}
|
|
|
|
|
|
logger = logging.getLogger("ComfyUI-OpenClaw.services.queue")
|
|
try:
|
|
from .structured_logging import (
|
|
configure_logger_for_structured_output,
|
|
emit_structured_log,
|
|
)
|
|
except ImportError:
|
|
from services.structured_logging import ( # type: ignore
|
|
configure_logger_for_structured_output,
|
|
emit_structured_log,
|
|
)
|
|
|
|
configure_logger_for_structured_output(logger)
|
|
|
|
import os
|
|
|
|
try:
|
|
from .tenant_context import get_current_tenant_id
|
|
except ImportError:
|
|
from services.tenant_context import get_current_tenant_id # type: ignore
|
|
|
|
# ComfyUI internal server URL fallback
|
|
COMFYUI_URL = (
|
|
os.environ.get("OPENCLAW_COMFYUI_URL")
|
|
or os.environ.get("MOLTBOT_COMFYUI_URL")
|
|
or "http://127.0.0.1:8188"
|
|
)
|
|
# IMPORTANT: keep this fixed and prompt-free; ComfyUI forwards it into API-node hidden inputs.
|
|
COMFY_USAGE_SOURCE = "comfyui-openclaw"
|
|
|
|
|
|
def _build_queue_extra_data(
|
|
extra_data: Optional[Dict[str, Any]] = None,
|
|
*,
|
|
tenant_id: Optional[str] = None,
|
|
) -> Dict[str, Any]:
|
|
payload_extra = dict(extra_data or {})
|
|
payload_extra.setdefault("comfy_usage_source", COMFY_USAGE_SOURCE)
|
|
|
|
openclaw_value = payload_extra.get("openclaw")
|
|
openclaw_extra = dict(openclaw_value) if isinstance(openclaw_value, dict) else {}
|
|
openclaw_extra.setdefault("tenant_id", tenant_id or get_current_tenant_id())
|
|
payload_extra["openclaw"] = openclaw_extra
|
|
|
|
return payload_extra
|
|
|
|
|
|
async def submit_prompt(
|
|
prompt_workflow: Dict[str, Any],
|
|
client_id: Optional[str] = None,
|
|
extra_data: Optional[Dict[str, Any]] = None,
|
|
source: str = "unknown", # R33: Source tracking
|
|
trace_id: Optional[str] = None, # R33: Trace ID for logging
|
|
tenant_id: Optional[str] = None, # S49: tenant context for budget + audit metadata
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
Submit a prompt workflow to ComfyUI with execution budgets (R33).
|
|
|
|
Args:
|
|
prompt_workflow: The full workflow JSON (API format)
|
|
client_id: Optional client ID for WebSocket mapping
|
|
extra_data: Extra metadata to attach (logging, etc)
|
|
source: Source type ("webhook" | "trigger" | "scheduler" | "bridge" | "unknown")
|
|
trace_id: Optional trace ID for logging/correlation
|
|
|
|
Returns:
|
|
Dict containing 'prompt_id' and 'number' (queue position) or error info.
|
|
|
|
Raises:
|
|
BudgetExceededError: If concurrency or size budgets are exceeded
|
|
"""
|
|
emit_structured_log(
|
|
logger,
|
|
level=logging.INFO,
|
|
event="queue.submit.start",
|
|
fields={
|
|
"source": source,
|
|
"trace_id": trace_id,
|
|
"has_extra_data": bool(extra_data),
|
|
"tenant_id": tenant_id or get_current_tenant_id(),
|
|
},
|
|
)
|
|
# NOTE: Must try relative import first. In ComfyUI runtime, `services` is not a top-level module.
|
|
# Keeping this order prevents "No module named 'services.execution_budgets'" during queue submit.
|
|
try:
|
|
from .execution_budgets import check_render_size, get_limiter
|
|
except ImportError:
|
|
from services.execution_budgets import ( # type: ignore
|
|
check_render_size,
|
|
get_limiter,
|
|
)
|
|
|
|
# R33: Check render size budget
|
|
check_render_size(prompt_workflow, trace_id=trace_id)
|
|
|
|
if client_id is None:
|
|
client_id = str(uuid.uuid4())
|
|
|
|
payload = {"prompt": prompt_workflow, "client_id": client_id}
|
|
payload["extra_data"] = _build_queue_extra_data(extra_data, tenant_id=tenant_id)
|
|
|
|
# NOTE: Debug-only full payload logging for troubleshooting mismatched outputs.
|
|
# Enable with OPENCLAW_DEBUG_PROMPT_PAYLOAD=1. This may include sensitive prompt content.
|
|
if os.environ.get("OPENCLAW_DEBUG_PROMPT_PAYLOAD", "").strip().lower() in (
|
|
"1",
|
|
"true",
|
|
"yes",
|
|
"on",
|
|
):
|
|
try:
|
|
logger.warning(
|
|
"DEBUG prompt payload (trace=%s source=%s): %s",
|
|
trace_id,
|
|
source,
|
|
json.dumps(payload, ensure_ascii=False),
|
|
)
|
|
except Exception:
|
|
logger.warning(
|
|
"DEBUG prompt payload (trace=%s source=%s): <failed to serialize>",
|
|
trace_id,
|
|
source,
|
|
)
|
|
|
|
# R33: Acquire concurrency budget
|
|
limiter = get_limiter()
|
|
async with limiter.acquire(
|
|
source=source,
|
|
trace_id=trace_id,
|
|
tenant_id=tenant_id or get_current_tenant_id(),
|
|
):
|
|
# Use aiohttp to post to local ComfyUI instance
|
|
# We assume we are running INSIDE ComfyUI process, but for HTTP access we use localhost
|
|
# unless we can hook internal server entry point.
|
|
# For MVP, HTTP loopback is safest and standard.
|
|
|
|
# R62: Lazy import aiohttp to avoid hard dependency crash at startup
|
|
try:
|
|
import aiohttp
|
|
except ImportError:
|
|
msg = "aiohttp is required for queue submission but not installed."
|
|
logger.error(msg)
|
|
raise APIError(
|
|
message=msg,
|
|
code=ErrorCode.DEPENDENCY_UNAVAILABLE,
|
|
status=503,
|
|
detail={"package": "aiohttp"},
|
|
)
|
|
|
|
url = f"{COMFYUI_URL}/prompt"
|
|
|
|
try:
|
|
async with aiohttp.ClientSession() as session:
|
|
async with session.post(url, json=payload) as resp:
|
|
if resp.status == 200:
|
|
data = await resp.json()
|
|
|
|
# R102 Hook
|
|
try:
|
|
q_size = data.get("number", 0)
|
|
from .security_telemetry import get_security_telemetry
|
|
|
|
get_security_telemetry().record_queue_saturation(q_size)
|
|
except:
|
|
pass
|
|
|
|
logger.info(
|
|
f"Queued prompt: {data.get('prompt_id')} (source={source}, trace_id={trace_id})"
|
|
)
|
|
emit_structured_log(
|
|
logger,
|
|
level=logging.INFO,
|
|
event="queue.submit.success",
|
|
fields={
|
|
"source": source,
|
|
"trace_id": trace_id,
|
|
"prompt_id": data.get("prompt_id"),
|
|
"queue_number": data.get("number"),
|
|
},
|
|
)
|
|
return data
|
|
else:
|
|
text = await resp.text()
|
|
logger.error(
|
|
f"Failed to queue prompt: {resp.status} - {text} (source={source}, trace_id={trace_id})"
|
|
)
|
|
emit_structured_log(
|
|
logger,
|
|
level=logging.ERROR,
|
|
event="queue.submit.upstream_error",
|
|
fields={
|
|
"source": source,
|
|
"trace_id": trace_id,
|
|
"upstream_status": resp.status,
|
|
},
|
|
)
|
|
# R61: Use APIError for queue failure
|
|
raise APIError(
|
|
message=f"Queue submission failed: {resp.status}",
|
|
code=ErrorCode.QUEUE_SUBMIT_FAILED,
|
|
status=502,
|
|
detail={
|
|
"upstream_status": resp.status,
|
|
"upstream_response": text[:200],
|
|
},
|
|
)
|
|
except APIError:
|
|
raise
|
|
except Exception as e:
|
|
logger.error(
|
|
f"Error submitting to queue: {e} (source={source}, trace_id={trace_id})"
|
|
)
|
|
emit_structured_log(
|
|
logger,
|
|
level=logging.ERROR,
|
|
event="queue.submit.error",
|
|
fields={
|
|
"source": source,
|
|
"trace_id": trace_id,
|
|
"error_type": type(e).__name__,
|
|
},
|
|
)
|
|
# R61: Wrap generic exceptions too
|
|
raise APIError(
|
|
message=f"Queue submission error: {str(e)}",
|
|
code=ErrorCode.INTERNAL_ERROR,
|
|
status=500,
|
|
)
|