feat: 添加ComfyUI服务器持久化功能，支持服务器状态自动保存与恢复，增强系统的可靠性和可维护性

README.md

@@ -7,6 +7,7 @@
 - 工作流上传
 - 工作流版本控制
 - 工作流加载
+- ComfyUI 服务器管理（支持持久化）
 - 需配置工作流服务器 (详见**WAAS(工作流即服务) Demo API 服务器**)
 
 # WAAS(工作流即服务) Demo API 服务器
@@ -42,6 +43,21 @@
       version TEXT NOT NULL,
       workflow_json TEXT NOT NULL
   )
+  
+  CREATE TABLE IF NOT EXISTS comfyui_servers (
+      name TEXT PRIMARY KEY,
+      http_url TEXT NOT NULL,
+      ws_url TEXT NOT NULL,
+      status TEXT NOT NULL DEFAULT 'offline',
+      last_heartbeat DATETIME,
+      last_health_check DATETIME,
+      current_tasks INTEGER DEFAULT 0,
+      max_concurrent_tasks INTEGER DEFAULT 1,
+      capabilities TEXT,
+      server_metadata TEXT,
+      created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
+      updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
+  )
   ```
 - 路由
 
@@ -100,6 +116,58 @@
   }
   ```
 
+# ComfyUI 服务器持久化功能
+
+## 概述
+
+ComfyUI 服务器管理器现在支持持久化功能，这意味着即使管理服务重启，也能从数据库中恢复已注册的服务器及其状态信息。
+
+## 主要特性
+
+- **自动持久化**: 服务器注册、状态变化、注销时自动保存到数据库
+- **自动恢复**: 服务启动时自动从数据库加载已注册的服务器
+- **健康检查**: 定期检查服务器健康状态并更新数据库
+- **心跳机制**: 支持服务器主动发送心跳，自动检测超时
+
+## 使用方法
+
+### 1. 数据库迁移
+
+首次使用需要运行迁移脚本：
+
+```bash
+python migrate_database.py
+```
+
+### 2. 测试功能
+
+运行测试脚本验证功能：
+
+```bash
+python test_persistence.py
+```
+
+### 3. API 使用
+
+```python
+from workflow_service.comfy.comfy_server import server_manager
+
+# 注册服务器
+await server_manager.register_server(
+    name="my_server",
+    http_url="http://localhost:8188",
+    ws_url="ws://localhost:8188",
+    max_concurrent_tasks=2
+)
+
+# 获取服务器状态
+servers = await server_manager.get_all_servers()
+```
+
+详细文档请参考 [PERSISTENCE_README.md](PERSISTENCE_README.md)
+
+---
+
 # JS 脚本使用指南
 
 ### MacOS

workflow_service/comfy/comfy_server.py

@@ -9,12 +9,19 @@ from enum import Enum
 
 import aiohttp
 from aiohttp import ClientTimeout
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy import select, update, delete
+from sqlalchemy.orm import selectinload
+
+from workflow_service.database.connection import AsyncSessionLocal
+from workflow_service.database.models import ComfyUIServer as ComfyUIServerModel
 
 logger = logging.getLogger(__name__)
 
 
 class ServerStatus(Enum):
     """服务器状态枚举"""
+
     ONLINE = "online"
     OFFLINE = "offline"
     BUSY = "busy"
@@ -24,6 +31,7 @@ class ServerStatus(Enum):
 @dataclass
 class ComfyUIServerInfo:
     """ComfyUI服务器信息"""
+
     name: str  # 服务器名称，由注册方自行拟定
     http_url: str
     ws_url: str
@@ -34,26 +42,141 @@ class ComfyUIServerInfo:
     max_concurrent_tasks: int = 1
     capabilities: Dict[str, any] = None  # 服务器能力信息，如支持的模型等
     metadata: Dict[str, any] = None  # 其他元数据
-    
+
     def __post_init__(self):
         if self.capabilities is None:
             self.capabilities = {}
         if self.metadata is None:
             self.metadata = {}
 
+    @classmethod
+    def from_model(cls, model: ComfyUIServerModel) -> "ComfyUIServerInfo":
+        """从数据库模型创建实例"""
+        capabilities = json.loads(model.capabilities) if model.capabilities else {}
+        metadata = json.loads(model.server_metadata) if model.server_metadata else {}
+
+        return cls(
+            name=model.name,
+            http_url=model.http_url,
+            ws_url=model.ws_url,
+            status=ServerStatus(model.status),
+            last_heartbeat=model.last_heartbeat,
+            last_health_check=model.last_health_check,
+            current_tasks=model.current_tasks,
+            max_concurrent_tasks=model.max_concurrent_tasks,
+            capabilities=capabilities,
+            metadata=metadata,
+        )
+
+    def to_model(self) -> ComfyUIServerModel:
+        """转换为数据库模型"""
+        return ComfyUIServerModel(
+            name=self.name,
+            http_url=self.http_url,
+            ws_url=self.ws_url,
+            status=self.status.value,
+            last_heartbeat=self.last_heartbeat,
+            last_health_check=self.last_health_check,
+            current_tasks=self.current_tasks,
+            max_concurrent_tasks=self.max_concurrent_tasks,
+            capabilities=json.dumps(self.capabilities),
+            server_metadata=json.dumps(self.metadata),
+        )
+
 
 class ComfyUIServerManager:
     """ComfyUI服务器管理器"""
-    
+
     def __init__(self):
         self.servers: Dict[str, ComfyUIServerInfo] = {}  # name -> server_info
         self.lock = asyncio.Lock()
         self.health_check_interval = 30  # 健康检查间隔（秒）
         self.heartbeat_timeout = 60  # 心跳超时时间（秒）
         self._health_check_task: Optional[asyncio.Task] = None
-        # 延迟启动健康检查，避免在模块导入时启动异步任务
-        # self._start_health_check()
-    
+        self._initialized = False
+
+    async def initialize(self):
+        """初始化管理器，从数据库加载服务器信息"""
+        if self._initialized:
+            return
+
+        async with self.lock:
+            try:
+                await self._load_servers_from_db()
+                self._initialized = True
+                logger.info(f"从数据库加载了 {len(self.servers)} 个服务器")
+            except Exception as e:
+                logger.error(f"从数据库加载服务器失败: {e}")
+                self._initialized = True  # 即使失败也标记为已初始化，避免重复尝试
+
+    async def _load_servers_from_db(self):
+        """从数据库加载服务器信息"""
+        async with AsyncSessionLocal() as session:
+            async with session.begin():
+                result = await session.execute(select(ComfyUIServerModel))
+                models = result.scalars().all()
+
+                for model in models:
+                    try:
+                        server_info = ComfyUIServerInfo.from_model(model)
+                        self.servers[server_info.name] = server_info
+                    except Exception as e:
+                        logger.error(f"加载服务器 {model.name} 失败: {e}")
+
+    async def _save_server_to_db(self, server: ComfyUIServerInfo):
+        """保存服务器信息到数据库"""
+        try:
+            async with AsyncSessionLocal() as session:
+                async with session.begin():
+                    # 检查是否已存在
+                    existing = await session.execute(
+                        select(ComfyUIServerModel).where(
+                            ComfyUIServerModel.name == server.name
+                        )
+                    )
+                    existing_model = existing.scalar_one_or_none()
+
+                    if existing_model:
+                        # 更新现有记录
+                        await session.execute(
+                            update(ComfyUIServerModel)
+                            .where(ComfyUIServerModel.name == server.name)
+                            .values(
+                                http_url=server.http_url,
+                                ws_url=server.ws_url,
+                                status=server.status.value,
+                                last_heartbeat=server.last_heartbeat,
+                                last_health_check=server.last_health_check,
+                                current_tasks=server.current_tasks,
+                                max_concurrent_tasks=server.max_concurrent_tasks,
+                                capabilities=json.dumps(server.capabilities),
+                                server_metadata=json.dumps(server.metadata),
+                                updated_at=datetime.utcnow(),
+                            )
+                        )
+                    else:
+                        # 创建新记录
+                        new_model = server.to_model()
+                        session.add(new_model)
+
+                    await session.commit()
+        except Exception as e:
+            logger.error(f"保存服务器 {server.name} 到数据库失败: {e}")
+
+    async def _delete_server_from_db(self, name: str):
+        """从数据库删除服务器"""
+        try:
+            async with AsyncSessionLocal() as session:
+                async with session.begin():
+                    await session.execute(
+                        delete(ComfyUIServerModel).where(
+                            ComfyUIServerModel.name == name
+                        )
+                    )
+                    await session.commit()
+        except Exception as e:
+            logger.error(f"从数据库删除服务器 {name} 失败: {e}")
+
     async def register_server(
         self,
         name: str,
@@ -61,92 +184,116 @@ class ComfyUIServerManager:
         ws_url: str,
         max_concurrent_tasks: int = 1,
         capabilities: Optional[Dict[str, any]] = None,
-        metadata: Optional[Dict[str, any]] = None
+        metadata: Optional[Dict[str, any]] = None,
     ) -> bool:
         """注册新的ComfyUI服务器"""
+        # 确保已初始化
+        await self.initialize()
+
         # 确保健康检查任务已启动
         await self._ensure_health_check_started()
-        
+
         async with self.lock:
             if name in self.servers:
                 logger.warning(f"服务器名称 {name} 已存在，将更新配置")
-            
+
             server_info = ComfyUIServerInfo(
                 name=name,
-                http_url=http_url.rstrip('/'),
-                ws_url=ws_url.rstrip('/'),
-                status=ServerStatus.OFFLINE, # 注册时状态为离线，等待健康检查
+                http_url=http_url.rstrip("/"),
+                ws_url=ws_url.rstrip("/"),
+                status=ServerStatus.OFFLINE,  # 注册时状态为离线，等待健康检查
                 max_concurrent_tasks=max_concurrent_tasks,
                 capabilities=capabilities or {},
-                metadata=metadata or {}
+                metadata=metadata or {},
             )
-            
+
             # 立即进行健康检查
             if await self._check_server_health(server_info):
                 server_info.status = ServerStatus.ONLINE
                 server_info.last_health_check = datetime.now()
                 server_info.last_heartbeat = datetime.now()
-            
+
             self.servers[name] = server_info
+
+            # 保存到数据库
+            await self._save_server_to_db(server_info)
+
             logger.info(f"服务器 {name} 注册成功: {http_url}")
             return True
-    
+
     async def unregister_server(self, name: str) -> bool:
         """注销ComfyUI服务器"""
         async with self.lock:
             if name in self.servers:
                 del self.servers[name]
+                # 从数据库删除
+                await self._delete_server_from_db(name)
                 logger.info(f"服务器 {name} 已注销")
                 return True
             return False
-    
+
     async def update_server_heartbeat(self, name: str) -> bool:
         """更新服务器心跳"""
         async with self.lock:
             if name in self.servers:
                 self.servers[name].last_heartbeat = datetime.now()
                 self.servers[name].status = ServerStatus.ONLINE
+                # 保存到数据库
+                await self._save_server_to_db(self.servers[name])
                 return True
             return False
-    
-    async def get_available_server(self, required_capabilities: Optional[Dict[str, any]] = None) -> Optional[ComfyUIServerInfo]:
+
+    async def get_available_server(
+        self, required_capabilities: Optional[Dict[str, any]] = None
+    ) -> Optional[ComfyUIServerInfo]:
         """获取可用的服务器（负载均衡）"""
+        # 确保已初始化
+        await self.initialize()
+
         # 确保健康检查任务已启动
         await self._ensure_health_check_started()
-        
+
         async with self.lock:
             available_servers = []
-            
+
             for server in self.servers.values():
-                if (server.status == ServerStatus.ONLINE and 
-                    server.current_tasks < server.max_concurrent_tasks):
-                    
+                if (
+                    server.status == ServerStatus.ONLINE
+                    and server.current_tasks < server.max_concurrent_tasks
+                ):
+
                     # 检查能力要求
                     if required_capabilities:
-                        if not self._check_capabilities(server.capabilities, required_capabilities):
+                        if not self._check_capabilities(
+                            server.capabilities, required_capabilities
+                        ):
                             continue
-                    
+
                     available_servers.append(server)
-            
+
             if not available_servers:
                 return None
-            
+
             # 简单的负载均衡：选择当前任务数最少的服务器
             return min(available_servers, key=lambda s: s.current_tasks)
-    
+
     async def allocate_server(self, name: str) -> bool:
         """分配服务器资源"""
         async with self.lock:
             if name in self.servers:
                 server = self.servers[name]
-                if (server.status == ServerStatus.ONLINE and 
-                    server.current_tasks < server.max_concurrent_tasks):
+                if (
+                    server.status == ServerStatus.ONLINE
+                    and server.current_tasks < server.max_concurrent_tasks
+                ):
                     server.current_tasks += 1
                     if server.current_tasks >= server.max_concurrent_tasks:
                         server.status = ServerStatus.BUSY
+                    # 保存到数据库
+                    await self._save_server_to_db(server)
                     return True
             return False
-    
+
     async def release_server(self, name: str) -> bool:
         """释放服务器资源"""
         async with self.lock:
@@ -154,22 +301,35 @@ class ComfyUIServerManager:
                 server = self.servers[name]
                 if server.current_tasks > 0:
                     server.current_tasks -= 1
-                    if server.status == ServerStatus.BUSY and server.current_tasks < server.max_concurrent_tasks:
+                    if (
+                        server.status == ServerStatus.BUSY
+                        and server.current_tasks < server.max_concurrent_tasks
+                    ):
                         server.status = ServerStatus.ONLINE
+                    # 保存到数据库
+                    await self._save_server_to_db(server)
                     return True
             return False
-    
+
     async def get_server_status(self, name: str) -> Optional[ComfyUIServerInfo]:
         """获取服务器状态"""
+        # 确保已初始化
+        await self.initialize()
+
         async with self.lock:
             return self.servers.get(name)
-    
+
     async def get_all_servers(self) -> List[ComfyUIServerInfo]:
         """获取所有服务器信息"""
+        # 确保已初始化
+        await self.initialize()
+
         async with self.lock:
             return list(self.servers.values())
-    
-    def _check_capabilities(self, server_capabilities: Dict[str, any], required_capabilities: Dict[str, any]) -> bool:
+
+    def _check_capabilities(
+        self, server_capabilities: Dict[str, any], required_capabilities: Dict[str, any]
+    ) -> bool:
         """检查服务器是否满足能力要求"""
         for key, value in required_capabilities.items():
             if key not in server_capabilities:
@@ -180,7 +340,7 @@ class ComfyUIServerManager:
             elif server_capabilities[key] != value:
                 return False
         return True
-    
+
     async def _check_server_health(self, server: ComfyUIServerInfo) -> bool:
         """检查服务器健康状态"""
         try:
@@ -191,9 +351,9 @@ class ComfyUIServerManager:
                         return True
         except Exception as e:
             logger.debug(f"服务器 {server.name} 健康检查失败: {e}")
-        
+
         return False
-    
+
     async def _health_check_loop(self):
         """健康检查循环"""
         while True:
@@ -201,35 +361,46 @@ class ComfyUIServerManager:
                 await self._perform_health_checks()
             except Exception as e:
                 logger.error(f"健康检查过程中发生错误: {e}")
-            
+
             await asyncio.sleep(self.health_check_interval)
-    
+
     async def _perform_health_checks(self):
         """执行健康检查"""
         current_time = datetime.now()
-        
+
         async with self.lock:
             for server in self.servers.values():
                 # 检查心跳超时
-                if (server.last_heartbeat and 
-                    current_time - server.last_heartbeat > timedelta(seconds=self.heartbeat_timeout)):
+                if (
+                    server.last_heartbeat
+                    and current_time - server.last_heartbeat
+                    > timedelta(seconds=self.heartbeat_timeout)
+                ):
                     server.status = ServerStatus.OFFLINE
                     logger.warning(f"服务器 {server.name} 心跳超时，标记为离线")
-                
+                    # 保存到数据库
+                    await self._save_server_to_db(server)
+
                 # 定期健康检查
-                if (not server.last_health_check or 
-                    current_time - server.last_health_check > timedelta(seconds=self.health_check_interval)):
-                    
+                if (
+                    not server.last_health_check
+                    or current_time - server.last_health_check
+                    > timedelta(seconds=self.health_check_interval)
+                ):
+
                     is_healthy = await self._check_server_health(server)
                     server.last_health_check = current_time
-                    
+
                     if is_healthy and server.status != ServerStatus.ONLINE:
                         server.status = ServerStatus.ONLINE
                         logger.info(f"服务器 {server.name} 恢复在线")
                     elif not is_healthy and server.status == ServerStatus.ONLINE:
                         server.status = ServerStatus.ERROR
                         logger.warning(f"服务器 {server.name} 健康检查失败")
-    
+
+                    # 保存到数据库
+                    await self._save_server_to_db(server)
+
     def _start_health_check(self):
         """启动健康检查任务"""
         try:
@@ -240,12 +411,12 @@ class ComfyUIServerManager:
         except RuntimeError:
             # 没有运行中的事件循环，延迟启动
             logger.debug("没有运行中的事件循环，健康检查将在事件循环启动后自动开始")
-    
+
     async def _ensure_health_check_started(self):
         """确保健康检查任务已启动"""
         if self._health_check_task is None or self._health_check_task.done():
             self._start_health_check()
-    
+
     async def shutdown(self):
         """关闭管理器"""
         if self._health_check_task and not self._health_check_task.done():
@@ -266,6 +437,8 @@ async def get_available_servers() -> List[ComfyUIServerInfo]:
     return await server_manager.get_all_servers()
 
 
-async def get_server_for_task(required_capabilities: Optional[Dict[str, any]] = None) -> Optional[ComfyUIServerInfo]:
+async def get_server_for_task(
+    required_capabilities: Optional[Dict[str, any]] = None
+) -> Optional[ComfyUIServerInfo]:
     """为任务获取合适的服务器（兼容性函数）"""
     return await server_manager.get_available_server(required_capabilities)

workflow_service/database/__init__.py

@@ -3,7 +3,7 @@
 提供工作流和工作流运行的数据库操作接口
 """
 
-from .models import Workflow, WorkflowRun, WorkflowRunNode
+from .models import Workflow, WorkflowRun, WorkflowRunNode, ComfyUIServer
 from .api import (
     init_db,
     save_workflow,
@@ -27,6 +27,7 @@ __all__ = [
     "Workflow",
     "WorkflowRun", 
     "WorkflowRunNode",
+    "ComfyUIServer",
     
     # API函数
     "init_db",

workflow_service/database/api.py

@@ -11,7 +11,7 @@ import uuid
 from sqlalchemy.future import select
 from sqlalchemy import delete
 
-from .models import Base, Workflow, WorkflowRun, WorkflowRunNode
+from .models import Base, Workflow, WorkflowRun, WorkflowRunNode, ComfyUIServer
 from .connection import async_engine, AsyncSessionLocal
 
 

workflow_service/database/models.py

@@ -4,7 +4,7 @@
 """
 
 from datetime import datetime
-from sqlalchemy import Column, String, Text, DateTime, ForeignKey
+from sqlalchemy import Column, String, Text, DateTime, ForeignKey, Integer, Boolean
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import relationship
 
@@ -109,3 +109,38 @@ class WorkflowRunNode(Base):
             "output_data": self.output_data,
             "error_message": self.error_message,
         }
+
+
+class ComfyUIServer(Base):
+    """ComfyUI 服务器模型"""
+
+    __tablename__ = "comfyui_servers"
+
+    name = Column(String, primary_key=True)
+    http_url = Column(String, nullable=False)
+    ws_url = Column(String, nullable=False)
+    status = Column(String, nullable=False, default="offline")
+    last_heartbeat = Column(DateTime)
+    last_health_check = Column(DateTime)
+    current_tasks = Column(Integer, default=0)
+    max_concurrent_tasks = Column(Integer, default=1)
+    capabilities = Column(Text)  # JSON 字符串
+    server_metadata = Column(Text)  # JSON 字符串
+    created_at = Column(DateTime, default=datetime.utcnow)
+    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
+
+    def to_dict(self) -> dict:
+        return {
+            "name": self.name,
+            "http_url": self.http_url,
+            "ws_url": self.ws_url,
+            "status": self.status,
+            "last_heartbeat": self.last_heartbeat.isoformat() if self.last_heartbeat else None,
+            "last_health_check": self.last_health_check.isoformat() if self.last_health_check else None,
+            "current_tasks": self.current_tasks,
+            "max_concurrent_tasks": self.max_concurrent_tasks,
+            "capabilities": self.capabilities,
+            "server_metadata": self.server_metadata,
+            "created_at": self.created_at.isoformat() if self.created_at else None,
+            "updated_at": self.updated_at.isoformat() if self.updated_at else None,
+        }