"""Queue management route functions for ComfyUI API."""
import httpx
from src.auth.base import ComfyAuth
from src.client.get_data import get_data
from src.client.response import ResponseGetData
from src.utils.logging import log_call
class QueueError(Exception):
"""Raised when queue operations fail."""
def __init__(self, message: str, response: ResponseGetData | None = None):
super().__init__(message)
self.response = response
@log_call(action_name="get_queue", level_name="route")
async def get_queue(
auth: ComfyAuth,
*,
session: httpx.AsyncClient | None = None,
) -> ResponseGetData:
"""Get current ComfyUI execution queue status.
Returns information about queued and running workflows, including
prompt IDs and queue positions.
Args:
auth: ComfyAuth instance for base URL and authentication
session: Optional HTTPX client for connection pooling
Returns:
ResponseGetData with:
- response: {"queue_running": [...], "queue_pending": [...]}
- status: 200 on success
- is_success: True on success
Raises:
QueueError: If request fails
Example:
>>> auth = NoAuth("http://127.0.0.1:8188")
>>> res = await get_queue(auth=auth)
>>> print(len(res.response["queue_pending"]))
3
"""
url = f"{auth.base_url}/queue"
res = await get_data(
auth=auth,
method="GET",
url=url,
session=session,
)
if not res.is_success:
raise QueueError(
f"Failed to get queue: HTTP {res.status}",
response=res,
)
return res
@log_call(action_name="cancel_prompt", level_name="route")
async def cancel_prompt(
auth: ComfyAuth,
prompt_id: str,
*,
session: httpx.AsyncClient | None = None,
) -> ResponseGetData:
"""Cancel a queued or running prompt in ComfyUI.
Removes the specified prompt from the execution queue. If the prompt
is currently running, it will be interrupted.
Args:
auth: ComfyAuth instance for base URL and authentication
prompt_id: Prompt ID to cancel
session: Optional HTTPX client for connection pooling
Returns:
ResponseGetData with:
- response: Confirmation message or empty
- status: 200 on success
- is_success: True on success
Raises:
QueueError: If cancellation fails
Example:
>>> auth = NoAuth("http://127.0.0.1:8188")
>>> res = await cancel_prompt(auth=auth, prompt_id="abc-123")
>>> print(res.is_success)
True
"""
url = f"{auth.base_url}/queue"
res = await get_data(
auth=auth,
method="POST",
url=url,
body={"delete": [prompt_id]},
session=session,
)
if not res.is_success:
raise QueueError(
f"Failed to cancel prompt {prompt_id}: HTTP {res.status}",
response=res,
)
return res
