Files
cbpark84-ComfyBridgeIO/client/websocket_client.py
T
cbpark84andClaude Sonnet 4.6 9a2bcdd4fa feat: initial project structure for ComfyBridgeIO
- ComfyUI custom nodes: BridgeInput (11 types) + BridgeOutput
- AnyType wildcard for universal output capture
- Custom server route: GET /comfybridge/results/{prompt_id}
- Python async client (ComfyBridge) + sync wrapper (ComfyBridgeSync)
- Workflow JSON parser and patcher (deepcopy, original unchanged)
- WebSocket completion detection
- BridgeResult container with PIL/bytes helpers
- Examples: basic_text_to_image.py, agent_loop.py
- Unit tests for workflow parser and client

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-07-16 10:45:26 +09:00

54 lines
1.6 KiB
Python

"""WebSocket 통신 및 워크플로우 완료 대기."""
import json
import logging
import uuid
from collections.abc import Callable
from typing import Any
logger = logging.getLogger(__name__)
async def wait_for_completion(
ws,
target_prompt_id: str,
progress_callback: Callable[[float], None] | None = None,
) -> None:
"""
WebSocket 메시지를 수신하며 워크플로우 완료를 기다린다.
완료 신호: { "type": "executing", "data": { "node": null } }
에러 신호: { "type": "execution_error" }
"""
async for raw in ws:
try:
if isinstance(raw, bytes):
continue # 바이너리 프리뷰 메시지 무시
data = json.loads(raw)
except json.JSONDecodeError:
continue
msg_type = data.get("type")
msg_data = data.get("data", {})
# 다른 prompt_id 메시지 무시
pid = msg_data.get("prompt_id")
if pid and pid != target_prompt_id:
continue
if msg_type == "executing":
node = msg_data.get("node")
if node is None:
logger.debug(f"[ComfyBridgeIO] Execution complete: {target_prompt_id}")
return # 완료
elif msg_type == "progress" and progress_callback:
value = msg_data.get("value", 0)
maximum = msg_data.get("max", 1)
if maximum > 0:
progress_callback(value / maximum)
elif msg_type == "execution_error":
err = msg_data.get("exception_message", "Unknown error")
raise RuntimeError(f"[ComfyBridgeIO] Execution error: {err}")