use crate::data::models::comfyui::*; use anyhow::{anyhow, Result}; use reqwest::{Client, Response}; use serde_json::Value; use std::time::Duration; use tracing::{debug, error, info, warn}; /// ComfyUI Infrastructure API 服务 /// /// 提供对 ComfyUI Workflow Service & Management API 的完整封装 /// 基于 OpenAPI 3.1.0 规范实现所有端点 /// 使用 /api/run/ 端点进行工作流执行 #[derive(Debug, Clone)] pub struct ComfyuiInfrastructureService { /// HTTP 客户端 client: Client, /// 服务配置 config: ComfyuiConfig, } impl ComfyuiInfrastructureService { /// 创建新的 ComfyUI Infrastructure 服务实例 pub fn new(config: ComfyuiConfig) -> Result { let timeout = Duration::from_secs(config.timeout.unwrap_or(30)); let client = Client::builder() .timeout(timeout) .build() .map_err(|e| anyhow!("Failed to create HTTP client: {}", e))?; info!("ComfyUI Infrastructure service initialized with base URL: {}", config.base_url); Ok(Self { client, config }) } /// 获取所有工作流 /// /// GET /api/workflow /// 返回所有已发布的工作流列表 pub async fn get_all_workflows(&self) -> Result> { let url = format!("{}/api/workflow", self.config.base_url.trim_end_matches('/')); debug!("Getting all workflows from: {}", url); let response = self.execute_get_request(&url).await?; let workflows: Vec = response.json().await .map_err(|e| anyhow!("Failed to parse workflows response: {}", e))?; info!("Retrieved {} workflows", workflows.len()); Ok(workflows) } /// 发布工作流 /// /// POST /api/workflow /// 发布新的工作流或更新现有工作流 pub async fn publish_workflow(&self, request: PublishWorkflowRequest) -> Result { let url = format!("{}/api/workflow", self.config.base_url.trim_end_matches('/')); debug!("Publishing workflow '{}' to: {}", request.name, url); let response = self.execute_post_request(&url, &request).await?; // 根据 OpenAPI 规范,成功响应状态码是 201 if response.status().as_u16() == 201 { // 响应体可能为空,创建成功响应 let workflow_response = PublishWorkflowResponse { success: true, message: Some("Workflow published successfully".to_string()), workflow: None, // API 规范中响应体为空 }; info!("Workflow '{}' published successfully", request.name); Ok(workflow_response) } else { let status = response.status(); let error_text = response.text().await.unwrap_or_default(); Err(anyhow!("Failed to publish workflow: HTTP {} - {}", status, error_text)) } } /// 删除工作流 /// /// DELETE /api/workflow/{workflow_name} /// 删除指定的工作流 pub async fn delete_workflow(&self, workflow_name: &str) -> Result { let encoded_name = urlencoding::encode(workflow_name); let url = format!("{}/api/workflow/{}", self.config.base_url.trim_end_matches('/'), encoded_name); debug!("Deleting workflow '{}' from: {}", workflow_name, url); let response = self.execute_delete_request(&url).await?; if response.status().is_success() { let delete_response = DeleteWorkflowResponse { success: true, message: Some("Workflow deleted successfully".to_string()), deleted_workflow: Some(workflow_name.to_string()), }; info!("Workflow '{}' deleted successfully", workflow_name); Ok(delete_response) } else { let status = response.status(); let error_text = response.text().await.unwrap_or_default(); Err(anyhow!("Failed to delete workflow: HTTP {} - {}", status, error_text)) } } /// 执行工作流 /// /// POST /api/run/ /// 执行指定的工作流 pub async fn execute_workflow(&self, request: ExecuteWorkflowRequest) -> Result { let mut url = format!("{}/api/run/", self.config.base_url.trim_end_matches('/')); // 添加查询参数 url.push_str(&format!("?base_name={}", urlencoding::encode(&request.base_name))); if let Some(version) = &request.version { url.push_str(&format!("&version={}", urlencoding::encode(version))); } debug!("Executing workflow '{}' at: {}", request.base_name, url); let response = self.execute_post_request(&url, &request.request_data).await?; if response.status().is_success() { let response_url = response.url().to_string(); info!("Response URL: {}", response_url); info!("Workflow '{}' execution started", request.base_name); // 总是尝试用 GET 请求获取最终结果 info!("Using GET request to fetch final result from: {}", response_url); match self.client.get(&response_url).send().await { Ok(final_response) => { if final_response.status().is_success() { match final_response.json::().await { Ok(final_result) => { info!("Successfully retrieved final result via GET"); let execute_response = ExecuteWorkflowResponse { task_id: None, // GET 请求通常不返回 task_id status: Some("completed".to_string()), message: Some("Workflow execution completed successfully".to_string()), result: Some(final_result), error: None, }; return Ok(execute_response); } Err(e) => { warn!("Failed to parse final result JSON: {}", e); } } } else { warn!("Failed to fetch final result: HTTP {}", final_response.status()); } } Err(e) => { warn!("Failed to request final result: {}", e); } } // 如果无法获取最终结果,解析初始响应 let initial_result: Value = response.json().await .unwrap_or_else(|_| serde_json::json!({})); let execute_response = ExecuteWorkflowResponse { task_id: initial_result.get("task_id").and_then(|v| v.as_str()).map(String::from), status: initial_result.get("status").and_then(|v| v.as_str()).map(String::from), message: Some("Workflow execution started".to_string()), result: Some(initial_result), error: None, }; Ok(execute_response) } else { let status = response.status(); let error_text = response.text().await.unwrap_or_default(); // 解析错误详情,如果是服务器端代理错误,提供更清晰的错误信息 let error_message = if error_text.contains("comfyui-for-waas-ui") { format!("服务器端 ComfyUI 服务错误: {} (状态码: {})", error_text, status) } else { format!("HTTP {} - {}", status, error_text) }; let execute_response = ExecuteWorkflowResponse { task_id: None, status: Some("failed".to_string()), message: None, result: None, error: Some(error_message.clone()), }; warn!("Workflow '{}' execution failed: {}", request.base_name, error_message); Ok(execute_response) } } /// 获取工作流规范 /// /// GET /api/spec/ /// 获取指定工作流的规范信息 pub async fn get_workflow_spec(&self, request: GetWorkflowSpecRequest) -> Result { let mut url = format!("{}/api/spec/", self.config.base_url.trim_end_matches('/')); // 添加查询参数 url.push_str(&format!("?base_name={}", urlencoding::encode(&request.base_name))); if let Some(version) = &request.version { url.push_str(&format!("&version={}", urlencoding::encode(version))); } debug!("Getting workflow spec for '{}' from: {}", request.base_name, url); let response = self.execute_get_request(&url).await?; let spec: Value = response.json().await .map_err(|e| anyhow!("Failed to parse workflow spec response: {}", e))?; let spec_response = GetWorkflowSpecResponse { spec: spec.clone(), name: Some(request.base_name.clone()), version: request.version.clone(), description: spec.get("description").and_then(|v| v.as_str()).map(String::from), }; info!("Retrieved workflow spec for '{}'", request.base_name); Ok(spec_response) } /// 获取服务器状态 /// /// GET /api/servers/status /// 获取所有已配置的ComfyUI服务器的配置信息和实时状态 pub async fn get_servers_status(&self) -> Result> { let url = format!("{}/api/servers/status", self.config.base_url.trim_end_matches('/')); debug!("Getting servers status from: {}", url); let response = self.execute_get_request(&url).await?; let servers: Vec = response.json().await .map_err(|e| anyhow!("Failed to parse servers status response: {}", e))?; info!("Retrieved status for {} servers", servers.len()); Ok(servers) } /// 获取服务器文件列表 /// /// GET /api/servers/{server_index}/files /// 获取指定ComfyUI服务器的输入和输出文件夹中的文件列表 pub async fn list_server_files(&self, server_index: i32) -> Result { let url = format!("{}/api/servers/{}/files", self.config.base_url.trim_end_matches('/'), server_index); debug!("Getting server files for server {} from: {}", server_index, url); let response = self.execute_get_request(&url).await?; let files: ServerFiles = response.json().await .map_err(|e| anyhow!("Failed to parse server files response: {}", e))?; info!("Retrieved file list for server {}: {} input files, {} output files", server_index, files.input_files.len(), files.output_files.len()); Ok(files) } /// 获取 API 根信息 /// /// GET / /// 提供一个API的快速使用指南 pub async fn get_api_root(&self) -> Result { let url = self.config.base_url.trim_end_matches('/').to_string(); debug!("Getting API root info from: {}", url); let response = self.execute_get_request(&url).await?; // 尝试解析为结构化响应,如果失败则返回基本信息 let result: Value = response.json().await .unwrap_or_else(|_| serde_json::json!({"guide": "ComfyUI API"})); let root_response = ApiRootResponse { guide: result.get("guide").and_then(|v| v.as_str()) .unwrap_or("ComfyUI Workflow Service & Management API").to_string(), version: result.get("version").and_then(|v| v.as_str()).map(String::from), endpoints: result.get("endpoints").and_then(|v| v.as_array()) .map(|arr| arr.iter().filter_map(|v| v.as_str().map(String::from)).collect()), }; info!("Retrieved API root information"); Ok(root_response) } // ========== 私有辅助方法 ========== /// 执行 GET 请求 async fn execute_get_request(&self, url: &str) -> Result { let response = self.client.get(url).send().await .map_err(|e| anyhow!("GET request failed: {}", e))?; Ok(response) } /// 执行 POST 请求 async fn execute_post_request(&self, url: &str, body: &T) -> Result { info!("Infrastructure service making POST request to: {}", url); info!("Request body: {:?}", serde_json::to_value(body).unwrap_or_default()); let response = self.client.post(url).json(body).send().await .map_err(|e| { error!("POST request failed to {}: {}", url, e); anyhow!("POST request failed: {}", e) })?; info!("Response status: {}", response.status()); info!("Response URL: {}", response.url()); Ok(response) } /// 执行 DELETE 请求 async fn execute_delete_request(&self, url: &str) -> Result { let response = self.client.delete(url).send().await .map_err(|e| anyhow!("DELETE request failed: {}", e))?; Ok(response) } /// 验证配置 pub fn validate_config(&self) -> Result<()> { if self.config.base_url.is_empty() { return Err(anyhow!("Base URL cannot be empty")); } if !self.config.base_url.starts_with("http://") && !self.config.base_url.starts_with("https://") { return Err(anyhow!("Base URL must start with http:// or https://")); } if let Some(timeout) = self.config.timeout { if timeout == 0 { return Err(anyhow!("Timeout must be greater than 0")); } } if let Some(retry_attempts) = self.config.retry_attempts { if retry_attempts > 10 { warn!("Retry attempts is very high ({}), consider reducing it", retry_attempts); } } Ok(()) } /// 获取服务配置 pub fn get_config(&self) -> &ComfyuiConfig { &self.config } /// 更新服务配置 pub fn update_config(&mut self, config: ComfyuiConfig) -> Result<()> { // 验证新配置 let temp_service = Self::new(config.clone())?; temp_service.validate_config()?; // 更新配置和客户端 self.config = config; let timeout = Duration::from_secs(self.config.timeout.unwrap_or(30)); self.client = Client::builder() .timeout(timeout) .build() .map_err(|e| anyhow!("Failed to create HTTP client: {}", e))?; info!("ComfyUI service configuration updated"); Ok(()) } /// 测试连接 pub async fn test_connection(&self) -> Result { match self.get_api_root().await { Ok(_) => { info!("ComfyUI service connection test successful"); Ok(true) } Err(e) => { error!("ComfyUI service connection test failed: {}", e); Ok(false) } } } }