refactor: 重构WorkflowQueueManager和相关API，移除冗余代码，简化任务处理逻辑，增强服务器状态检查功能

workflow_service/comfy/comfy_queue.py

@@ -1,27 +1,13 @@
 import asyncio
-import json
 import logging
-import os
+from typing import Any, Optional
-import random
-import uuid
-import websockets
-from datetime import datetime
-from typing import Dict, Any, Optional
 
 import aiohttp
-from aiohttp import ClientTimeout
 
 from workflow_service.comfy.comfy_workflow import ComfyWorkflow
+from workflow_service.comfy.comfy_run import ComfyRun
 from workflow_service.config import Settings
 from workflow_service.comfy.comfy_server import server_manager, ComfyUIServerInfo
-from workflow_service.database.api import (
-    create_workflow_run,
-    update_workflow_run_status,
-    create_workflow_run_nodes,
-    update_workflow_run_node_status,
-    get_workflow_run,
-    get_workflow_run_nodes,
-)
 
 settings = Settings()
 
@@ -29,22 +15,6 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 
-API_INPUT_PREFIX = "INPUT_"
-API_OUTPUT_PREFIX = "OUTPUT_"
-
-
-class ComfyUIExecutionError(Exception):
-    def __init__(self, error_data: dict):
-        self.error_data = error_data
-        # 创建一个对开发者友好的异常消息
-        message = (
-            f"ComfyUI节点执行失败。节点ID: {error_data.get('node_id')}, "
-            f"节点类型: {error_data.get('node_type')}. "
-            f"错误: {error_data.get('exception_message', 'N/A')}"
-        )
-        super().__init__(message)
-
-
 # 全局任务队列管理器
 class WorkflowQueueManager:
     def __init__(self, monitor_interval: int = 5):
@@ -71,33 +41,60 @@ class WorkflowQueueManager:
                 pass
             logger.info("队列监控任务已停止")
 
-    async def set_monitor_interval(self, interval: int):
+    async def add_task(
-        """设置监控间隔"""
+        self,
-        if interval < 1:
+        workflow: ComfyWorkflow,
-            raise ValueError("监控间隔不能小于1秒")
+        request_data: dict,
-        self._monitor_interval = interval
+    ):
-        logger.info(f"队列监控间隔已设置为 {interval} 秒")
+        """添加新任务到队列"""
-
-        # 如果监控任务正在运行，重启它以应用新间隔
-        if self._queue_monitor_task and not self._queue_monitor_task.done():
-            await self.stop_queue_monitor()
-            await self.start_queue_monitor()
-
-    async def trigger_queue_processing(self):
-        """手动触发队列处理（用于测试或紧急情况）"""
-        logger.info("手动触发队列处理")
-        asyncio.create_task(self._process_queue())
-
-    async def get_queue_status(self) -> dict:
-        """获取队列状态信息"""
         async with self.lock:
-            return {
+            # 创建ComfyRun实例并保存到数据库
-                "pending_tasks_count": len(self.pending_tasks),
+            comfy_run = await ComfyRun.create(workflow, request_data)
-                "running_tasks_count": len(self.running_tasks),
+
-                "monitor_active": self._queue_monitor_task is not None
+            # 添加到待处理队列
-                and not self._queue_monitor_task.done(),
+            self.pending_tasks.append(comfy_run.run_id)
-                "monitor_interval": self._monitor_interval,
+            logger.info(
-            }
+                f"任务 {comfy_run.run_id} 已添加到队列，当前队列长度: {len(self.pending_tasks)}"
+            )
+
+            return comfy_run.run_id
+
+    async def get_server_status(
+        self, server: ComfyUIServerInfo, session: aiohttp.ClientSession
+    ) -> dict[str, Any]:
+        """
+        检查单个ComfyUI服务器的详细状态。
+        返回一个包含可达性、队列状态和详细队列内容的字典。
+        """
+        # [BUG修复] 确保初始字典结构与成功时的结构一致，以满足Pydantic模型
+        status_info = {
+            "is_reachable": False,
+            "is_free": False,
+            "queue_details": {"running_count": 0, "pending_count": 0},
+        }
+        try:
+            queue_url = f"{server.http_url}/queue"
+            async with session.get(queue_url, timeout=60) as response:
+                response.raise_for_status()
+                queue_data = await response.json()
+
+                status_info["is_reachable"] = True
+
+                running_count = len(queue_data.get("queue_running", []))
+                pending_count = len(queue_data.get("queue_pending", []))
+
+                status_info["queue_details"] = {
+                    "running_count": running_count,
+                    "pending_count": pending_count,
+                }
+
+                status_info["is_free"] = running_count == 0 and pending_count == 0
+
+        except Exception as e:
+            # 当请求失败时，将返回上面定义的、结构正确的初始 status_info
+            logger.warning(f"无法检查服务器 {server.http_url} 的队列状态: {e}")
+
+        return status_info
 
     async def _monitor_queue(self):
         """定期监控队列，当有可用服务器时自动处理"""
@@ -107,11 +104,12 @@ class WorkflowQueueManager:
 
                 # 每次定时检测触发时，打印当前状态
                 async with self.lock:
+                    from datetime import datetime
+
                     pending_count = len(self.pending_tasks)
                     running_count = len(self.running_tasks)
                     current_time = datetime.now()
                     timestamp = current_time.strftime("%Y-%m-%d %H:%M:%S")
-                    time_only = current_time.strftime("%H:%M:%S")
 
                     # 计算距离上次处理队列的时间
                     time_since_last_processing = ""
@@ -161,44 +159,12 @@ class WorkflowQueueManager:
                 # 出错时稍微延长等待时间，避免频繁重试
                 await asyncio.sleep(self._monitor_interval * 2)
 
-    async def add_task(
-        self,
-        workflow: ComfyWorkflow,
-        request_data: dict,
-    ):
-        """添加新任务到队列"""
-        workflow_run_id = str(uuid.uuid4())
-        async with self.lock:
-            # 创建任务记录
-            await create_workflow_run(
-                workflow_run_id=workflow_run_id,
-                workflow_name=workflow.workflow_name,
-                workflow_json=json.dumps(workflow.workflow_data.model_dump()),
-                api_spec=json.dumps(workflow.get_api_spec().model_dump()),
-                request_data=json.dumps(request_data),
-            )
-
-            # 创建工作流节点记录
-            nodes_data = []
-            for node in workflow.workflow_data.nodes:
-                nodes_data.append({"id": str(node.id), "type": node.type})
-            await create_workflow_run_nodes(workflow_run_id, nodes_data)
-
-            # 添加到待处理队列
-            self.pending_tasks.append(workflow_run_id)
-            logger.info(
-                f"任务 {workflow_run_id} 已添加到队列，当前队列长度: {len(self.pending_tasks)}"
-            )
-
-            # 注意：不再手动触发队列处理，由定时监控自动处理
-            # 这样可以避免在没有可用服务器时的不必要尝试
-
-            return workflow_run_id
-
     async def _process_queue(self):
         """处理队列中的任务"""
         async with self.lock:
             # 更新上次处理队列的时间
+            from datetime import datetime
+
             self._last_processing_time = datetime.now()
 
             if not self.pending_tasks:
@@ -221,13 +187,7 @@ class WorkflowQueueManager:
                 f"开始处理任务 {workflow_run_id}，使用服务器 {server.name} ({server.http_url})"
             )
 
-            # 分配服务器资源
+            # 记录运行中任务
-            await server_manager.allocate_server(server.name)
-
-            # 标记任务为运行中
-            await update_workflow_run_status(
-                workflow_run_id, "running", server.http_url
-            )
             self.running_tasks[server.http_url] = {
                 "workflow_run_id": workflow_run_id,
                 "started_at": datetime.now(),
@@ -251,307 +211,25 @@ class WorkflowQueueManager:
 
     async def _execute_task(self, workflow_run_id: str, server: ComfyUIServerInfo):
         """执行任务"""
-        cleanup_paths = []
         try:
-            # 获取工作流数据
+            # 从run_id创建ComfyRun实例并执行
-            workflow_run = await get_workflow_run(workflow_run_id)
+            comfy_run = await ComfyRun.from_run_id(workflow_run_id)
-            if not workflow_run:
+            if not comfy_run:
                 raise Exception(f"找不到工作流运行记录: {workflow_run_id}")
 
-            workflow_data = json.loads(workflow_run.workflow_json)
+            # 执行工作流（ComfyRun内部处理所有逻辑包括资源管理）
-            request_data = json.loads(workflow_run.request_data)
+            await comfy_run.execute(server)
-            workflow = ComfyWorkflow(workflow_run.workflow_name, workflow_data)
-            # 执行工作流
-            result = await _execute_prompt_on_server(
-                workflow, request_data, server, workflow_run_id
-            )
-
-            # 保存处理后的结果到数据库
-            await update_workflow_run_status(
-                workflow_run_id,
-                "completed",
-                result=json.dumps(result, ensure_ascii=False),
-            )
 
         except Exception as e:
             logger.error(f"执行任务 {workflow_run_id} 时出错: {e}")
-            await update_workflow_run_status(
-                workflow_run_id, "failed", error_message=str(e)
-            )
-        finally:
-            # 清理临时文件
-            if cleanup_paths:
-                logger.info(f"正在清理 {len(cleanup_paths)} 个临时文件...")
-                for path in cleanup_paths:
-                    try:
-                        if os.path.exists(path):
-                            os.remove(path)
-                            logger.info(f"  - 已删除: {path}")
-                    except Exception as e:
-                        logger.warning(f"  - 删除 {path} 时出错: {e}")
 
+        finally:
             # 清理运行状态
             async with self.lock:
                 if server.http_url in self.running_tasks:
                     del self.running_tasks[server.http_url]
 
-                # 释放服务器资源
-                await server_manager.release_server(server.name)
-
-    async def get_task_status(self, workflow_run_id: str) -> dict:
-        """获取任务状态"""
-        workflow_run = await get_workflow_run(workflow_run_id)
-        if not workflow_run:
-            return {"error": "任务不存在"}
-        nodes = await get_workflow_run_nodes(workflow_run_id)
-
-        result = {
-            "id": workflow_run_id,
-            "status": workflow_run.status,
-            "created_at": workflow_run.created_at,
-            "started_at": workflow_run.started_at,
-            "completed_at": workflow_run.completed_at,
-            "server_url": workflow_run.server_url,
-            "error_message": workflow_run.error_message,
-            "nodes": nodes,
-        }
-
-        # 如果任务完成，从数据库获取结果
-        if workflow_run.status == "completed" and workflow_run.result:
-            try:
-                result["result"] = json.loads(workflow_run.result)
-            except (json.JSONDecodeError, TypeError):
-                result["result"] = None
-
-        return result
-
-    async def get_server_status(
-        self, server: ComfyUIServerInfo, session: aiohttp.ClientSession
-    ) -> dict[str, Any]:
-        """
-        检查单个ComfyUI服务器的详细状态。
-        返回一个包含可达性、队列状态和详细队列内容的字典。
-        """
-        # [BUG修复] 确保初始字典结构与成功时的结构一致，以满足Pydantic模型
-        status_info = {
-            "is_reachable": False,
-            "is_free": False,
-            "queue_details": {"running_count": 0, "pending_count": 0},
-        }
-        try:
-            queue_url = f"{server.http_url}/queue"
-            async with session.get(queue_url, timeout=60) as response:
-                response.raise_for_status()
-                queue_data = await response.json()
-
-                status_info["is_reachable"] = True
-
-                running_count = len(queue_data.get("queue_running", []))
-                pending_count = len(queue_data.get("queue_pending", []))
-
-                status_info["queue_details"] = {
-                    "running_count": running_count,
-                    "pending_count": pending_count,
-                }
-
-                status_info["is_free"] = running_count == 0 and pending_count == 0
-
-        except Exception as e:
-            # 当请求失败时，将返回上面定义的、结构正确的初始 status_info
-            logger.warning(f"无法检查服务器 {server.http_url} 的队列状态: {e}")
-
-        return status_info
-
 
 # 全局队列管理器实例
 # 从配置中读取监控间隔，默认为5秒
 queue_manager = WorkflowQueueManager(monitor_interval=5)
-
-
-async def _execute_prompt_on_server(
-    workflow: ComfyWorkflow,
-    request_data: dict,
-    server: ComfyUIServerInfo,
-    workflow_run_id: str,
-) -> dict:
-    """
-    在指定的服务器上执行一个准备好的prompt。
-    现在支持节点级别的状态跟踪。
-    """
-    client_id = str(uuid.uuid4())
-
-    # 构建prompt
-    prompt = await workflow.build_prompt(server, request_data)
-
-    # 更新工作流运行状态，记录prompt_id和client_id
-    await update_workflow_run_status(
-        workflow_run_id,
-        "running",
-        server.http_url,
-        None,  # prompt_id将在_queue_prompt后更新
-        client_id,
-    )
-
-    # 提交到ComfyUI
-    prompt_id = await _queue_prompt(workflow, prompt, client_id, server.http_url)
-
-    # 更新prompt_id
-    await update_workflow_run_status(
-        workflow_run_id, "running", server.http_url, prompt_id, client_id
-    )
-
-    logger.info(
-        f"工作流 {workflow_run_id} 已在 {server.http_url} 上入队，Prompt ID: {prompt_id}"
-    )
-
-    # 获取执行结果，现在支持节点级别的状态跟踪
-    results = await _get_execution_results(
-        workflow, prompt_id, client_id, server.ws_url, workflow_run_id
-    )
-    return results
-
-
-async def _queue_prompt(
-    workflow: ComfyWorkflow, prompt: dict, client_id: str, http_url: str
-) -> str:
-    """通过HTTP POST将工作流任务提交到指定的ComfyUI服务器。"""
-    for node_id in prompt:
-        prompt[node_id]["inputs"][f"cache_buster_{uuid.uuid4().hex}"] = random.random()
-    payload = {
-        "prompt": prompt,
-        "client_id": client_id,
-        "extra_data": {
-            "api_key_comfy_org": "",
-            "extra_pnginfo": {"workflow": workflow.workflow_data.model_dump()},
-        },
-    }
-    async with aiohttp.ClientSession(timeout=ClientTimeout(total=90)) as session:
-        prompt_url = f"{http_url}/prompt"
-        try:
-            async with session.post(prompt_url, json=payload) as response:
-                logger.info(f"ComfyUI /prompt 端点返回的响应: {response}")
-                response.raise_for_status()
-                result = await response.json()
-                logger.info(f"ComfyUI /prompt 端点返回的响应: {result}")
-                if "prompt_id" not in result:
-                    raise Exception(f"从 ComfyUI /prompt 端点返回的响应无效: {result}")
-                return result["prompt_id"]
-        except Exception as e:
-            logger.error(f"提交到 ComfyUI /prompt 端点时发生错误: {e}")
-            raise e
-
-
-async def _get_execution_results(
-    workflow: ComfyWorkflow,
-    prompt_id: str,
-    client_id: str,
-    ws_url: str,
-    workflow_run_id: str,
-) -> dict:
-    """
-    通过WebSocket连接到指定的ComfyUI服务器，聚合执行结果。
-    现在支持节点级别的状态跟踪。
-    """
-    full_ws_url = f"{ws_url}?clientId={client_id}"
-    aggregated_outputs = {}
-
-    try:
-        async with websockets.connect(full_ws_url) as websocket:
-            while True:
-                try:
-                    out = await websocket.recv()
-                    if not isinstance(out, str):
-                        continue
-
-                    message = json.loads(out)
-                    msg_type = message.get("type")
-                    data = message.get("data")
-
-                    if not (data and data.get("prompt_id") == prompt_id):
-                        continue
-
-                    # 捕获并处理执行错误
-                    if msg_type == "execution_error":
-                        error_data = data
-                        logger.error(
-                            f"ComfyUI执行错误 (Prompt ID: {prompt_id}): {error_data}"
-                        )
-
-                        # 更新节点状态为失败
-                        node_id = error_data.get("node_id")
-                        if node_id:
-                            await update_workflow_run_node_status(
-                                workflow_run_id,
-                                node_id,
-                                "failed",
-                                error_message=error_data.get(
-                                    "exception_message", "Unknown error"
-                                ),
-                            )
-
-                        # 抛出自定义异常，将错误详情传递出去
-                        raise ComfyUIExecutionError(error_data)
-
-                    # 处理节点开始执行
-                    if msg_type == "executing" and data.get("node"):
-                        node_id = data.get("node")
-                        logger.info(f"节点 {node_id} 开始执行 (Prompt ID: {prompt_id})")
-
-                        # 更新节点状态为运行中
-                        await update_workflow_run_node_status(
-                            workflow_run_id, node_id, "running"
-                        )
-
-                    # 处理节点执行完成
-                    if msg_type == "executed":
-                        node_id = data.get("node")
-                        output_data = data.get("output")
-                        if node_id and output_data:
-                            node = next(
-                                (
-                                    x
-                                    for x in workflow.workflow_data.nodes
-                                    if str(x.id) == node_id
-                                ),
-                                None,
-                            )
-                            if (
-                                node
-                                and node.title
-                                and node.title.startswith(API_OUTPUT_PREFIX)
-                            ):
-                                title = node.title.replace(API_OUTPUT_PREFIX, "")
-                                aggregated_outputs[title] = output_data
-                            logger.info(
-                                f"收到节点 {node_id} 的输出 (Prompt ID: {prompt_id})"
-                            )
-
-                            # 更新节点状态为完成
-                            await update_workflow_run_node_status(
-                                workflow_run_id,
-                                node_id,
-                                "completed",
-                                output_data=json.dumps(output_data),
-                            )
-
-                    # 处理整个工作流执行完成
-                    elif msg_type == "executing" and data.get("node") is None:
-                        logger.info(f"Prompt ID: {prompt_id} 执行完成。")
-                        return aggregated_outputs
-
-                except websockets.exceptions.ConnectionClosed as e:
-                    logger.warning(
-                        f"WebSocket 连接已关闭 (Prompt ID: {prompt_id})。错误: {e}"
-                    )
-                    return aggregated_outputs
-                except Exception as e:
-                    # 重新抛出我们自己的异常，或者处理其他意外错误
-                    if not isinstance(e, ComfyUIExecutionError):
-                        logger.error(f"处理 prompt {prompt_id} 时发生意外错误: {e}")
-                    raise e
-
-    except websockets.exceptions.InvalidURI as e:
-        logger.error(
-            f"错误: 尝试连接的WebSocket URI无效: '{full_ws_url}'. 原始URL: '{ws_url}'. 错误: {e}"
-        )
-        raise e

workflow_service/comfy/comfy_run.py

@@ -0,0 +1,345 @@
+import json
+import logging
+import random
+import uuid
+from typing import Optional
+
+import aiohttp
+import websockets
+from aiohttp import ClientTimeout
+
+from workflow_service.comfy.comfy_workflow import ComfyWorkflow, API_OUTPUT_PREFIX
+from workflow_service.comfy.comfy_server import ComfyUIServerInfo, server_manager
+from workflow_service.database.api import (
+    create_workflow_run,
+    create_workflow_run_nodes,
+    get_workflow_run,
+    update_workflow_run_status,
+    update_workflow_run_node_status,
+)
+
+
+class ComfyUIExecutionError(Exception):
+    """ComfyUI执行错误"""
+
+    def __init__(self, error_data):
+        self.error_data = error_data
+        super().__init__(f"ComfyUI execution error: {error_data}")
+
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class ComfyRun:
+    """
+    ComfyUI工作流运行实例，封装单个工作流运行的所有操作
+    """
+
+    def __init__(self, workflow: ComfyWorkflow, run_id: str, request_data: dict):
+        """
+        初始化工作流运行实例
+
+        Args:
+            workflow: ComfyWorkflow实例
+            run_id: 运行ID
+            request_data: 请求数据
+        """
+        self.workflow = workflow
+        self.run_id = run_id
+        self.request_data = request_data
+        self.client_id = str(uuid.uuid4())
+        self.prompt_id: Optional[str] = None
+        self.server: Optional[ComfyUIServerInfo] = None
+
+    @classmethod
+    async def create(cls, workflow: ComfyWorkflow, request_data: dict) -> "ComfyRun":
+        """
+        创建新的ComfyRun实例并保存到数据库
+
+        Args:
+            workflow: ComfyWorkflow实例
+            request_data: 请求数据
+
+        Returns:
+            ComfyRun实例
+        """
+        run_id = str(uuid.uuid4())
+
+        # 创建任务记录
+        await create_workflow_run(
+            workflow_run_id=run_id,
+            workflow_name=workflow.workflow_name,
+            workflow_json=json.dumps(workflow.workflow_data.model_dump()),
+            api_spec=json.dumps(workflow.get_api_spec().model_dump()),
+            request_data=json.dumps(request_data),
+        )
+
+        # 创建工作流节点记录
+        nodes_data = []
+        for node in workflow.workflow_data.nodes:
+            nodes_data.append({"id": str(node.id), "type": node.type})
+        await create_workflow_run_nodes(run_id, nodes_data)
+
+        return cls(workflow, run_id, request_data)
+
+    @classmethod
+    async def from_run_id(cls, run_id: str) -> Optional["ComfyRun"]:
+        """
+        从run_id创建ComfyRun实例
+
+        Args:
+            run_id: 运行ID
+
+        Returns:
+            ComfyRun实例，如果找不到则返回None
+        """
+        workflow_run = await get_workflow_run(run_id)
+        if not workflow_run:
+            return None
+
+        workflow_data = json.loads(workflow_run.workflow_json)
+        workflow = ComfyWorkflow(workflow_run.workflow_name, workflow_data)
+        request_data = json.loads(workflow_run.request_data)
+
+        return cls(workflow, run_id, request_data)
+
+    async def execute(self, server: ComfyUIServerInfo) -> dict:
+        """
+        在指定的服务器上执行工作流
+
+        Args:
+            server: ComfyUI服务器信息
+
+        Returns:
+            执行结果
+        """
+        self.server = server
+
+        try:
+            # 分配服务器资源
+            await server_manager.allocate_server(server.name)
+
+            # 构建prompt
+            prompt = await self.workflow.build_prompt(server, self.request_data)
+
+            # 更新运行状态为running
+            await self._update_status(
+                "running", server.http_url, client_id=self.client_id
+            )
+
+            # 提交到ComfyUI
+            self.prompt_id = await self._queue_prompt(prompt, server.http_url)
+
+            # 更新prompt_id
+            await self._update_status(
+                "running", server.http_url, self.prompt_id, self.client_id
+            )
+
+            logger.info(
+                f"工作流 {self.run_id} 已在 {server.http_url} 上入队，Prompt ID: {self.prompt_id}"
+            )
+
+            # 获取执行结果
+            results = await self._get_execution_results(server.ws_url)
+
+            # 标记完成
+            await self._update_status(
+                "completed", result=json.dumps(results, ensure_ascii=False)
+            )
+
+            return results
+
+        except Exception as e:
+            # 标记失败
+            await self._update_status("failed", error_message=str(e))
+            raise
+
+        finally:
+            # 释放服务器资源
+            try:
+                await server_manager.release_server(server.name)
+            except Exception as e:
+                logger.error(f"释放服务器资源时出错: {e}")
+
+    async def _update_status(
+        self,
+        status: str,
+        server_url: Optional[str] = None,
+        prompt_id: Optional[str] = None,
+        client_id: Optional[str] = None,
+        result: Optional[str] = None,
+        error_message: Optional[str] = None,
+    ):
+        """更新工作流运行状态"""
+        await update_workflow_run_status(
+            self.run_id, status, server_url, prompt_id, client_id, error_message, result
+        )
+
+    async def _queue_prompt(self, prompt: dict, http_url: str) -> str:
+        """提交工作流到ComfyUI服务器"""
+        # 添加随机缓存破坏器避免缓存
+        for node_id in prompt:
+            prompt[node_id]["inputs"][
+                f"cache_buster_{uuid.uuid4().hex}"
+            ] = random.random()
+
+        payload = {
+            "prompt": prompt,
+            "client_id": self.client_id,
+            "extra_data": {
+                "api_key_comfy_org": "",
+                "extra_pnginfo": {"workflow": self.workflow.workflow_data.model_dump()},
+            },
+        }
+
+        async with aiohttp.ClientSession(timeout=ClientTimeout(total=90)) as session:
+            prompt_url = f"{http_url}/prompt"
+            try:
+                async with session.post(prompt_url, json=payload) as response:
+                    logger.info(f"ComfyUI /prompt 端点返回的响应: {response}")
+                    response.raise_for_status()
+                    result = await response.json()
+                    logger.info(f"ComfyUI /prompt 端点返回的响应: {result}")
+
+                    if "prompt_id" not in result:
+                        raise Exception(
+                            f"从 ComfyUI /prompt 端点返回的响应无效: {result}"
+                        )
+
+                    return result["prompt_id"]
+
+            except Exception as e:
+                logger.error(f"提交到 ComfyUI /prompt 端点时发生错误: {e}")
+                raise
+
+    async def _get_execution_results(self, ws_url: str) -> dict:
+        """
+        通过WebSocket连接获取执行结果，支持节点级别的状态跟踪
+        """
+        aggregated_outputs = {}
+        full_ws_url = f"{ws_url}?clientId={self.client_id}"
+
+        try:
+            async with websockets.connect(full_ws_url) as websocket:
+                logger.info(f"已连接到WebSocket: {full_ws_url}")
+
+                while True:
+                    try:
+                        out = await websocket.recv()
+                        if not isinstance(out, str):
+                            continue
+
+                        message = json.loads(out)
+                        msg_type = message.get("type")
+                        data = message.get("data")
+
+                        if not (data and data.get("prompt_id") == self.prompt_id):
+                            continue
+
+                        # 捕获并处理执行错误
+                        if msg_type == "execution_error":
+                            error_data = data
+                            logger.error(
+                                f"ComfyUI执行错误 (Prompt ID: {self.prompt_id}): {error_data}"
+                            )
+
+                            # 更新节点状态为失败
+                            node_id = error_data.get("node_id")
+                            if node_id:
+                                await update_workflow_run_node_status(
+                                    self.run_id,
+                                    node_id,
+                                    "failed",
+                                    error_message=error_data.get(
+                                        "exception_message", "Unknown error"
+                                    ),
+                                )
+
+                            # 抛出自定义异常，将错误详情传递出去
+                            raise ComfyUIExecutionError(error_data)
+
+                        # 处理节点开始执行
+                        if msg_type == "executing" and data.get("node"):
+                            node_id = data.get("node")
+                            logger.info(
+                                f"节点 {node_id} 开始执行 (Prompt ID: {self.prompt_id})"
+                            )
+
+                            # 更新节点状态为运行中
+                            await update_workflow_run_node_status(
+                                self.run_id, node_id, "running"
+                            )
+
+                        # 处理节点执行完成
+                        elif msg_type == "executed":
+                            await self._handle_node_executed(data, aggregated_outputs)
+
+                        # 处理整个工作流执行完成
+                        elif msg_type == "executing" and data.get("node") is None:
+                            logger.info(f"Prompt ID: {self.prompt_id} 执行完成。")
+                            return aggregated_outputs
+
+                    except websockets.exceptions.ConnectionClosed as e:
+                        logger.warning(
+                            f"WebSocket 连接已关闭 (Prompt ID: {self.prompt_id})。错误: {e}"
+                        )
+                        return aggregated_outputs
+                    except Exception as e:
+                        # 重新抛出我们自己的异常，或者处理其他意外错误
+                        if not isinstance(e, ComfyUIExecutionError):
+                            logger.error(
+                                f"处理 prompt {self.prompt_id} 时发生意外错误: {e}"
+                            )
+                        raise e
+
+        except websockets.exceptions.InvalidURI as e:
+            logger.error(
+                f"错误: 尝试连接的WebSocket URI无效: '{full_ws_url}'. 原始URL: '{ws_url}'. 错误: {e}"
+            )
+            raise e
+        except Exception as e:
+            logger.error(f"WebSocket连接出错: {e}")
+            raise
+
+    async def _handle_node_executed(self, data: dict, aggregated_outputs: dict):
+        """处理节点执行完成事件"""
+        node_id = data.get("node")
+        output_data = data.get("output")
+
+        if not node_id or not output_data:
+            return
+
+        # 查找对应的节点
+        node = next(
+            (x for x in self.workflow.workflow_data.nodes if str(x.id) == node_id),
+            None,
+        )
+
+        if not node:
+            return
+
+        # 如果是输出节点，收集结果
+        if node.title and node.title.startswith(API_OUTPUT_PREFIX):
+            title = node.title.replace(API_OUTPUT_PREFIX, "")
+            aggregated_outputs[title] = output_data
+
+        logger.info(f"收到节点 {node_id} 的输出 (Prompt ID: {self.prompt_id})")
+
+        # 更新节点状态为完成
+        await update_workflow_run_node_status(
+            self.run_id,
+            node_id,
+            "completed",
+            output_data=json.dumps(output_data),
+        )
+
+    def get_status_info(self) -> dict:
+        """获取运行状态信息"""
+        return {
+            "run_id": self.run_id,
+            "workflow_name": self.workflow.workflow_name,
+            "client_id": self.client_id,
+            "prompt_id": self.prompt_id,
+            "server_url": self.server.http_url if self.server else None,
+        }

workflow_service/routes/run.py

@@ -7,7 +7,12 @@ from fastapi.responses import JSONResponse
 
 from workflow_service.comfy.comfy_queue import queue_manager
 from workflow_service.comfy.comfy_workflow import ComfyWorkflow
-from workflow_service.database.api import get_workflow, get_workflow_runs_recent
+from workflow_service.database.api import (
+    get_workflow,
+    get_workflow_run,
+    get_workflow_run_nodes,
+    get_workflow_runs_recent,
+)
 
 run_router = APIRouter(
     prefix="/api/run",
@@ -133,7 +138,30 @@ async def get_run_status(workflow_run_id: str):
     获取工作流执行状态。
     """
     try:
-        status = await queue_manager.get_task_status(workflow_run_id)
+        """获取任务状态"""
-        return status
+        workflow_run = await get_workflow_run(workflow_run_id)
+        if not workflow_run:
+            return {"error": "任务不存在"}
+        nodes = await get_workflow_run_nodes(workflow_run_id)
+
+        result = {
+            "id": workflow_run_id,
+            "status": workflow_run.status,
+            "created_at": workflow_run.created_at,
+            "started_at": workflow_run.started_at,
+            "completed_at": workflow_run.completed_at,
+            "server_url": workflow_run.server_url,
+            "error_message": workflow_run.error_message,
+            "nodes": nodes,
+        }
+
+        # 如果任务完成，从数据库获取结果
+        if workflow_run.status == "completed" and workflow_run.result:
+            try:
+                result["result"] = json.loads(workflow_run.result)
+            except (json.JSONDecodeError, TypeError):
+                result["result"] = None
+
+        return result
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"获取状态失败: {str(e)}")