feat: 实时通信与高级功能

2025-08-08 15:10:06 +08:00
parent 8223061aea
commit e1a8cd4d9d
12 changed files with 3736 additions and 1 deletions
--- a/apps/desktop/src-tauri/src/business/services/websocket_handler.rs
+++ b/apps/desktop/src-tauri/src/business/services/websocket_handler.rs
@@ -0,0 +1,280 @@
+//! 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
+    }
+}