281 lines
9.4 KiB
Rust
281 lines
9.4 KiB
Rust
//! WebSocket 消息处理器
|
|
//! 专门处理 ComfyUI WebSocket 消息的解析和事件转换
|
|
|
|
use anyhow::{Result, anyhow};
|
|
use std::collections::HashMap;
|
|
use std::sync::Arc;
|
|
use tokio::sync::{RwLock, broadcast};
|
|
use tracing::{info, warn, error, debug};
|
|
use serde_json::Value;
|
|
|
|
use crate::business::services::realtime_monitor_v2::RealtimeEventV2;
|
|
use crate::data::repositories::comfyui_repository::ComfyUIRepository;
|
|
|
|
/// WebSocket 消息处理器
|
|
pub struct WebSocketHandler {
|
|
/// 事件发送器
|
|
event_sender: broadcast::Sender<RealtimeEventV2>,
|
|
/// 提示 ID 到执行 ID 的映射
|
|
prompt_execution_map: Arc<RwLock<HashMap<String, String>>>,
|
|
/// 数据仓库
|
|
repository: Arc<ComfyUIRepository>,
|
|
}
|
|
|
|
impl WebSocketHandler {
|
|
/// 创建新的消息处理器
|
|
pub fn new(
|
|
event_sender: broadcast::Sender<RealtimeEventV2>,
|
|
prompt_execution_map: Arc<RwLock<HashMap<String, String>>>,
|
|
repository: Arc<ComfyUIRepository>,
|
|
) -> Self {
|
|
Self {
|
|
event_sender,
|
|
prompt_execution_map,
|
|
repository,
|
|
}
|
|
}
|
|
|
|
/// 处理 WebSocket 消息
|
|
pub async fn handle_message(&self, message: Value) -> Result<()> {
|
|
debug!("处理 WebSocket 消息: {:?}", message);
|
|
|
|
// 解析消息类型
|
|
let message_type = message.get("type")
|
|
.and_then(|t| t.as_str())
|
|
.unwrap_or("unknown");
|
|
|
|
let timestamp = chrono::Utc::now().to_rfc3339();
|
|
|
|
match message_type {
|
|
"executing" => self.handle_executing_message(message, timestamp).await?,
|
|
"progress" => self.handle_progress_message(message, timestamp).await?,
|
|
"executed" => self.handle_executed_message(message, timestamp).await?,
|
|
"execution_error" => self.handle_execution_error_message(message, timestamp).await?,
|
|
"status" => self.handle_status_message(message, timestamp).await?,
|
|
_ => {
|
|
debug!("未处理的消息类型: {}", message_type);
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// 处理执行消息
|
|
async fn handle_executing_message(&self, message: Value, timestamp: String) -> Result<()> {
|
|
if let Some(data) = message.get("data") {
|
|
let prompt_id = data.get("prompt_id")
|
|
.and_then(|p| p.as_str())
|
|
.unwrap_or("")
|
|
.to_string();
|
|
|
|
let node_id = data.get("node")
|
|
.and_then(|n| n.as_str())
|
|
.map(|s| s.to_string());
|
|
|
|
let execution_id = {
|
|
let map = self.prompt_execution_map.read().await;
|
|
map.get(&prompt_id).cloned()
|
|
};
|
|
|
|
if let Some(node_id) = node_id {
|
|
let _ = self.event_sender.send(RealtimeEventV2::NodeExecuting {
|
|
prompt_id: prompt_id.clone(),
|
|
execution_id: execution_id.clone(),
|
|
node_id: node_id.clone(),
|
|
node_title: None,
|
|
timestamp: timestamp.clone(),
|
|
});
|
|
} else {
|
|
let _ = self.event_sender.send(RealtimeEventV2::ExecutionStarted {
|
|
prompt_id,
|
|
execution_id,
|
|
node_id: None,
|
|
timestamp,
|
|
});
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
/// 处理进度消息
|
|
async fn handle_progress_message(&self, message: Value, timestamp: String) -> Result<()> {
|
|
if let Some(data) = message.get("data") {
|
|
let prompt_id = data.get("prompt_id")
|
|
.and_then(|p| p.as_str())
|
|
.unwrap_or("")
|
|
.to_string();
|
|
|
|
let node_id = data.get("node")
|
|
.and_then(|n| n.as_str())
|
|
.unwrap_or("")
|
|
.to_string();
|
|
|
|
let progress = data.get("value")
|
|
.and_then(|v| v.as_u64())
|
|
.unwrap_or(0) as u32;
|
|
|
|
let max_progress = data.get("max")
|
|
.and_then(|m| m.as_u64())
|
|
.unwrap_or(100) as u32;
|
|
|
|
let percentage = if max_progress > 0 {
|
|
(progress as f32 / max_progress as f32) * 100.0
|
|
} else {
|
|
0.0
|
|
};
|
|
|
|
let execution_id = {
|
|
let map = self.prompt_execution_map.read().await;
|
|
map.get(&prompt_id).cloned()
|
|
};
|
|
|
|
let _ = self.event_sender.send(RealtimeEventV2::ExecutionProgress {
|
|
prompt_id,
|
|
execution_id,
|
|
node_id,
|
|
progress,
|
|
max_progress,
|
|
percentage,
|
|
timestamp,
|
|
});
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
/// 处理执行完成消息
|
|
async fn handle_executed_message(&self, message: Value, timestamp: String) -> Result<()> {
|
|
if let Some(data) = message.get("data") {
|
|
let prompt_id = data.get("prompt_id")
|
|
.and_then(|p| p.as_str())
|
|
.unwrap_or("")
|
|
.to_string();
|
|
|
|
let outputs = data.get("output")
|
|
.and_then(|o| o.as_object())
|
|
.map(|obj| {
|
|
obj.iter()
|
|
.map(|(k, v)| (k.clone(), v.clone()))
|
|
.collect::<HashMap<String, serde_json::Value>>()
|
|
})
|
|
.unwrap_or_default();
|
|
|
|
let output_urls = self.extract_output_urls(&outputs);
|
|
|
|
let execution_id = {
|
|
let map = self.prompt_execution_map.read().await;
|
|
map.get(&prompt_id).cloned()
|
|
};
|
|
|
|
// 更新数据库中的执行状态
|
|
if let Some(exec_id) = &execution_id {
|
|
if let Ok(Some(mut execution)) = self.repository.get_execution(exec_id).await {
|
|
execution.set_results(outputs.clone(), output_urls.clone());
|
|
let _ = self.repository.update_execution(&execution).await;
|
|
}
|
|
}
|
|
|
|
let _ = self.event_sender.send(RealtimeEventV2::ExecutionCompleted {
|
|
prompt_id: prompt_id.clone(),
|
|
execution_id: execution_id.clone(),
|
|
outputs,
|
|
output_urls,
|
|
execution_time: 0, // TODO: 计算实际执行时间
|
|
timestamp,
|
|
});
|
|
|
|
// 清理映射
|
|
{
|
|
let mut map = self.prompt_execution_map.write().await;
|
|
map.remove(&prompt_id);
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
/// 处理执行错误消息
|
|
async fn handle_execution_error_message(&self, message: Value, timestamp: String) -> Result<()> {
|
|
if let Some(data) = message.get("data") {
|
|
let prompt_id = data.get("prompt_id")
|
|
.and_then(|p| p.as_str())
|
|
.unwrap_or("")
|
|
.to_string();
|
|
|
|
let node_id = data.get("node_id")
|
|
.and_then(|n| n.as_str())
|
|
.map(|s| s.to_string());
|
|
|
|
let error_msg = data.get("exception_message")
|
|
.and_then(|e| e.as_str())
|
|
.unwrap_or("未知错误")
|
|
.to_string();
|
|
|
|
let details = data.get("traceback").cloned();
|
|
|
|
let execution_id = {
|
|
let map = self.prompt_execution_map.read().await;
|
|
map.get(&prompt_id).cloned()
|
|
};
|
|
|
|
// 更新数据库中的执行状态
|
|
if let Some(exec_id) = &execution_id {
|
|
if let Ok(Some(mut execution)) = self.repository.get_execution(exec_id).await {
|
|
execution.set_error(error_msg.clone());
|
|
let _ = self.repository.update_execution(&execution).await;
|
|
}
|
|
}
|
|
|
|
let _ = self.event_sender.send(RealtimeEventV2::ExecutionFailed {
|
|
prompt_id: prompt_id.clone(),
|
|
execution_id: execution_id.clone(),
|
|
node_id,
|
|
error: error_msg,
|
|
details,
|
|
timestamp,
|
|
});
|
|
|
|
// 清理映射
|
|
{
|
|
let mut map = self.prompt_execution_map.write().await;
|
|
map.remove(&prompt_id);
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
/// 处理状态消息
|
|
async fn handle_status_message(&self, message: Value, timestamp: String) -> Result<()> {
|
|
if let Some(data) = message.get("data") {
|
|
let queue_running = data.get("status")
|
|
.and_then(|s| s.get("exec_info"))
|
|
.and_then(|e| e.get("queue_remaining"))
|
|
.and_then(|q| q.as_u64())
|
|
.unwrap_or(0) as u32;
|
|
|
|
let _ = self.event_sender.send(RealtimeEventV2::QueueUpdated {
|
|
running_count: if queue_running > 0 { 1 } else { 0 },
|
|
pending_count: queue_running,
|
|
timestamp,
|
|
});
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
/// 提取输出 URLs
|
|
fn extract_output_urls(&self, outputs: &HashMap<String, serde_json::Value>) -> Vec<String> {
|
|
let mut urls = Vec::new();
|
|
|
|
for (_, output) in outputs {
|
|
if let Some(images) = output.get("images").and_then(|v| v.as_array()) {
|
|
for image in images {
|
|
if let Some(filename) = image.get("filename").and_then(|v| v.as_str()) {
|
|
let url = format!("/view?filename={}", filename);
|
|
urls.push(url);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
urls
|
|
}
|
|
}
|