feat: add comfyui sdk

cargos/comfyui-sdk/client/comfyui_client.rs

@@ -0,0 +1,274 @@
+//! Main ComfyUI client that combines HTTP and WebSocket functionality
+
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::time::{Duration, Instant};
+use tokio::time::sleep;
+
+use crate::types::{
+    ComfyUIClientConfig, PromptRequest, TemplateExecutionResult, 
+    ExecutionOptions, ExecutionCallbacks, ExecutionError, ParameterValues
+};
+use crate::templates::{TemplateManager, WorkflowTemplate, WorkflowInstance};
+use crate::client::{HTTPClient, WebSocketClient};
+use crate::error::{ComfyUIError, Result};
+
+
+/// Main ComfyUI client
+pub struct ComfyUIClient {
+    http_client: HTTPClient,
+    ws_client: WebSocketClient,
+    template_manager: TemplateManager,
+    config: ComfyUIClientConfig,
+}
+
+impl ComfyUIClient {
+    /// Creates a new ComfyUI client
+    pub fn new(config: ComfyUIClientConfig) -> Result<Self> {
+        let http_client = HTTPClient::new(config.clone())?;
+        let ws_client = WebSocketClient::new(config.clone());
+        let template_manager = TemplateManager::new();
+
+        Ok(Self {
+            http_client,
+            ws_client,
+            template_manager,
+            config,
+        })
+    }
+
+    /// Gets the HTTP client
+    pub fn http(&self) -> &HTTPClient {
+        &self.http_client
+    }
+
+    /// Gets the WebSocket client
+    pub fn ws(&self) -> &WebSocketClient {
+        &self.ws_client
+    }
+
+    /// Gets the template manager
+    pub fn templates(&mut self) -> &mut TemplateManager {
+        &mut self.template_manager
+    }
+
+    /// Gets the template manager (read-only)
+    pub fn templates_ref(&self) -> &TemplateManager {
+        &self.template_manager
+    }
+
+    /// Gets the client configuration
+    pub fn config(&self) -> &ComfyUIClientConfig {
+        &self.config
+    }
+
+    /// Connects to ComfyUI server (both HTTP and WebSocket)
+    pub async fn connect(&mut self) -> Result<()> {
+        // Test HTTP connection first
+        self.http_client.get_queue().await
+            .map_err(|e| ComfyUIError::connection(format!("HTTP connection failed: {}", e)))?;
+
+        // Connect WebSocket
+        self.ws_client.connect().await
+            .map_err(|e| ComfyUIError::connection(format!("WebSocket connection failed: {}", e)))?;
+
+        log::info!("Successfully connected to ComfyUI server");
+        Ok(())
+    }
+
+    /// Disconnects from ComfyUI server
+    pub async fn disconnect(&mut self) -> Result<()> {
+        self.ws_client.disconnect().await?;
+        log::info!("Disconnected from ComfyUI server");
+        Ok(())
+    }
+
+    /// Checks if connected to the server
+    pub async fn is_connected(&self) -> bool {
+        self.ws_client.is_connected().await
+    }
+
+    /// Executes a workflow template
+    pub async fn execute_template(
+        &self,
+        template: &WorkflowTemplate,
+        parameters: ParameterValues,
+        options: ExecutionOptions,
+    ) -> Result<TemplateExecutionResult> {
+        let start_time = Instant::now();
+
+        // Create workflow instance
+        let instance = template.create_instance(parameters)?;
+
+        // Convert to prompt request
+        let prompt_request = self.instance_to_prompt_request(&instance)?;
+
+        // Submit prompt
+        let prompt_response = self.http_client.queue_prompt(&prompt_request).await?;
+        let prompt_id = prompt_response.prompt_id.clone();
+
+        log::info!("Submitted prompt with ID: {}", prompt_id);
+
+        // Wait for completion if WebSocket is connected
+        if self.ws_client.is_connected().await {
+            match self.wait_for_completion(&prompt_id, options.timeout).await {
+                Ok(outputs) => {
+                    let execution_time = start_time.elapsed().as_millis() as u64;
+                    Ok(TemplateExecutionResult::success(prompt_id, outputs, execution_time))
+                }
+                Err(error) => {
+                    let execution_time = start_time.elapsed().as_millis() as u64;
+                    let exec_error = ExecutionError {
+                        node_id: None,
+                        message: error.to_string(),
+                        details: None,
+                        timestamp: chrono::Utc::now(),
+                    };
+                    Ok(TemplateExecutionResult::failure(prompt_id, exec_error, execution_time))
+                }
+            }
+        } else {
+            // If no WebSocket, just return the prompt ID
+            let execution_time = start_time.elapsed().as_millis() as u64;
+            Ok(TemplateExecutionResult::success(prompt_id, HashMap::new(), execution_time))
+        }
+    }
+
+    /// Executes a template with callbacks
+    pub async fn execute_template_with_callbacks(
+        &self,
+        template: &WorkflowTemplate,
+        parameters: ParameterValues,
+        options: ExecutionOptions,
+        callbacks: Arc<dyn ExecutionCallbacks>,
+    ) -> Result<TemplateExecutionResult> {
+        // Register callbacks
+        let callback_id = self.ws_client.register_callbacks(callbacks).await;
+
+        // Execute template
+        let result = self.execute_template(template, parameters, options).await;
+
+        // Unregister callbacks
+        self.ws_client.unregister_callbacks(&callback_id).await;
+
+        result
+    }
+
+    /// Converts a workflow instance to a prompt request
+    fn instance_to_prompt_request(&self, instance: &WorkflowInstance) -> Result<PromptRequest> {
+        let workflow_json = instance.to_workflow_json()?;
+        
+        let mut prompt = HashMap::new();
+        if let serde_json::Value::Object(workflow_obj) = workflow_json {
+            for (key, value) in workflow_obj {
+                prompt.insert(key, value);
+            }
+        }
+
+        Ok(PromptRequest {
+            prompt,
+            client_id: Some(self.ws_client.client_id().to_string()),
+            extra_data: None,
+        })
+    }
+
+    /// Waits for prompt completion
+    async fn wait_for_completion(
+        &self,
+        prompt_id: &str,
+        timeout: Option<Duration>,
+    ) -> Result<HashMap<String, serde_json::Value>> {
+        let timeout_duration = timeout.unwrap_or(Duration::from_secs(300)); // 5 minutes default
+        let start_time = Instant::now();
+        let check_interval = Duration::from_millis(1000); // Check every second
+
+        loop {
+            // Check timeout
+            if start_time.elapsed() > timeout_duration {
+                return Err(ComfyUIError::timeout(format!(
+                    "Execution timed out after {:?}",
+                    timeout_duration
+                )));
+            }
+
+            // Check history for completion
+            match self.http_client.get_history_by_prompt(prompt_id).await {
+                Ok(history) => {
+                    if let Some(history_item) = history.get(prompt_id) {
+                        if history_item.status.completed {
+                            // Extract outputs
+                            let mut outputs = HashMap::new();
+                            for (node_id, node_outputs) in &history_item.outputs {
+                                let mut node_output = HashMap::new();
+                                for (output_name, output_list) in node_outputs {
+                                    node_output.insert(output_name.clone(), serde_json::to_value(output_list)?);
+                                }
+                                outputs.insert(node_id.clone(), serde_json::to_value(node_output)?);
+                            }
+                            return Ok(outputs);
+                        }
+                    }
+                }
+                Err(e) => {
+                    log::warn!("Error checking history: {}", e);
+                }
+            }
+
+            sleep(check_interval).await;
+        }
+    }
+
+    /// Converts outputs to HTTP URLs
+    pub fn outputs_to_urls(&self, outputs: &HashMap<String, serde_json::Value>) -> Vec<String> {
+        self.http_client.outputs_to_urls(outputs)
+    }
+
+    /// Interrupts the current execution
+    pub async fn interrupt(&self) -> Result<()> {
+        self.http_client.interrupt().await
+    }
+
+    /// Clears the execution queue
+    pub async fn clear_queue(&self) -> Result<()> {
+        self.http_client.clear_queue().await
+    }
+
+    /// Gets the current queue status
+    pub async fn get_queue_status(&self) -> Result<crate::types::QueueStatus> {
+        self.http_client.get_queue().await
+    }
+
+    /// Gets system statistics
+    pub async fn get_system_stats(&self) -> Result<crate::types::SystemStats> {
+        self.http_client.get_system_stats().await
+    }
+
+    /// Gets available nodes information
+    pub async fn get_object_info(&self) -> Result<crate::types::ObjectInfo> {
+        self.http_client.get_object_info().await
+    }
+
+    /// Uploads an image file
+    pub async fn upload_image<P: AsRef<std::path::Path>>(
+        &self,
+        image_path: P,
+        overwrite: bool,
+    ) -> Result<crate::types::UploadImageResponse> {
+        self.http_client.upload_image(image_path, overwrite).await
+    }
+
+    /// Gets an image by filename
+    pub async fn get_image(
+        &self,
+        filename: &str,
+        subfolder: Option<&str>,
+        image_type: Option<&str>,
+    ) -> Result<bytes::Bytes> {
+        self.http_client.get_image(filename, subfolder, image_type).await
+    }
+
+    /// Frees memory
+    pub async fn free_memory(&self) -> Result<()> {
+        self.http_client.free_memory().await
+    }
+}

cargos/comfyui-sdk/client/http_client.rs

@@ -0,0 +1,324 @@
+//! HTTP client for ComfyUI API
+
+use std::collections::HashMap;
+use std::path::Path;
+use reqwest::{Client, multipart};
+use url::Url;
+use crate::types::{
+    ComfyUIClientConfig, PromptRequest, PromptResponse, QueueStatus, 
+    HistoryItem, SystemStats, ObjectInfo, UploadImageResponse, ViewMetadata
+};
+use crate::error::{ComfyUIError, Result};
+use crate::utils::error_handling::{retry_if_retryable, RetryConfig, with_timeout};
+
+/// HTTP client for ComfyUI API
+#[derive(Debug, Clone)]
+pub struct HTTPClient {
+    client: Client,
+    base_url: Url,
+    config: ComfyUIClientConfig,
+}
+
+impl HTTPClient {
+    /// Creates a new HTTP client
+    pub fn new(config: ComfyUIClientConfig) -> Result<Self> {
+        let base_url = Url::parse(&config.base_url)?;
+        
+        let mut client_builder = Client::builder();
+        
+        if let Some(timeout) = config.timeout {
+            client_builder = client_builder.timeout(timeout);
+        }
+
+        let client = client_builder.build()?;
+
+        Ok(Self {
+            client,
+            base_url,
+            config,
+        })
+    }
+
+    /// Gets the base URL
+    pub fn base_url(&self) -> &Url {
+        &self.base_url
+    }
+
+    /// Builds a URL for an endpoint
+    fn build_url(&self, endpoint: &str) -> Result<Url> {
+        self.base_url.join(endpoint)
+            .map_err(ComfyUIError::from)
+    }
+
+    /// Executes a GET request with retry logic
+    async fn get_with_retry<T>(&self, endpoint: &str) -> Result<T>
+    where
+        T: serde::de::DeserializeOwned,
+    {
+        let url = self.build_url(endpoint)?;
+        let retry_config = RetryConfig {
+            max_attempts: self.config.retry_attempts.unwrap_or(3),
+            initial_delay: self.config.retry_delay.unwrap_or(std::time::Duration::from_millis(1000)),
+            ..Default::default()
+        };
+
+        let operation = || async {
+            let mut request = self.client.get(url.clone());
+            
+            if let Some(headers) = &self.config.headers {
+                for (key, value) in headers {
+                    request = request.header(key, value);
+                }
+            }
+
+            let response = request.send().await?;
+            
+            if !response.status().is_success() {
+                return Err(ComfyUIError::Http(reqwest::Error::from(
+                    response.error_for_status().unwrap_err()
+                )));
+            }
+
+            let data: T = response.json().await?;
+            Ok(data)
+        };
+
+        if let Some(timeout) = self.config.timeout {
+            with_timeout(retry_if_retryable(operation, retry_config), timeout).await
+        } else {
+            retry_if_retryable(operation, retry_config).await
+        }
+    }
+
+    /// Executes a POST request with retry logic
+    async fn post_with_retry<T, B>(&self, endpoint: &str, body: &B) -> Result<T>
+    where
+        T: serde::de::DeserializeOwned,
+        B: serde::Serialize,
+    {
+        let url = self.build_url(endpoint)?;
+        let retry_config = RetryConfig {
+            max_attempts: self.config.retry_attempts.unwrap_or(3),
+            initial_delay: self.config.retry_delay.unwrap_or(std::time::Duration::from_millis(1000)),
+            ..Default::default()
+        };
+
+        let operation = || async {
+            let mut request = self.client.post(url.clone()).json(body);
+            
+            if let Some(headers) = &self.config.headers {
+                for (key, value) in headers {
+                    request = request.header(key, value);
+                }
+            }
+
+            let response = request.send().await?;
+            
+            if !response.status().is_success() {
+                return Err(ComfyUIError::Http(reqwest::Error::from(
+                    response.error_for_status().unwrap_err()
+                )));
+            }
+
+            let data: T = response.json().await?;
+            Ok(data)
+        };
+
+        if let Some(timeout) = self.config.timeout {
+            with_timeout(retry_if_retryable(operation, retry_config), timeout).await
+        } else {
+            retry_if_retryable(operation, retry_config).await
+        }
+    }
+
+    /// Gets the current queue status
+    pub async fn get_queue(&self) -> Result<QueueStatus> {
+        self.get_with_retry("/queue").await
+    }
+
+    /// Submits a prompt for execution
+    pub async fn queue_prompt(&self, prompt_request: &PromptRequest) -> Result<PromptResponse> {
+        self.post_with_retry("/prompt", prompt_request).await
+    }
+
+    /// Gets execution history
+    pub async fn get_history(&self, max_items: Option<u32>) -> Result<HashMap<String, HistoryItem>> {
+        let endpoint = if let Some(max) = max_items {
+            format!("/history?max_items={}", max)
+        } else {
+            "/history".to_string()
+        };
+        
+        self.get_with_retry(&endpoint).await
+    }
+
+    /// Gets history for a specific prompt
+    pub async fn get_history_by_prompt(&self, prompt_id: &str) -> Result<HashMap<String, HistoryItem>> {
+        let endpoint = format!("/history/{}", prompt_id);
+        self.get_with_retry(&endpoint).await
+    }
+
+    /// Clears the execution queue
+    pub async fn clear_queue(&self) -> Result<()> {
+        let url = self.build_url("/queue")?;
+        let response = self.client.delete(url).send().await?;
+        
+        if !response.status().is_success() {
+            return Err(ComfyUIError::Http(reqwest::Error::from(
+                response.error_for_status().unwrap_err()
+            )));
+        }
+        
+        Ok(())
+    }
+
+    /// Cancels a specific prompt
+    pub async fn cancel_prompt(&self, prompt_id: &str) -> Result<()> {
+        let body = serde_json::json!({ "delete": [prompt_id] });
+        let url = self.build_url("/queue")?;
+        let response = self.client.post(url).json(&body).send().await?;
+
+        if !response.status().is_success() {
+            return Err(ComfyUIError::Http(reqwest::Error::from(
+                response.error_for_status().unwrap_err()
+            )));
+        }
+
+        Ok(())
+    }
+
+    /// Gets system statistics
+    pub async fn get_system_stats(&self) -> Result<SystemStats> {
+        self.get_with_retry("/system_stats").await
+    }
+
+    /// Gets object information (available nodes)
+    pub async fn get_object_info(&self) -> Result<ObjectInfo> {
+        self.get_with_retry("/object_info").await
+    }
+
+    /// Uploads an image file
+    pub async fn upload_image<P: AsRef<Path>>(&self, image_path: P, overwrite: bool) -> Result<UploadImageResponse> {
+        let path = image_path.as_ref();
+        let filename = path.file_name()
+            .and_then(|name| name.to_str())
+            .ok_or_else(|| ComfyUIError::new("Invalid filename"))?;
+
+        let file_bytes = tokio::fs::read(path).await?;
+
+        let part = multipart::Part::bytes(file_bytes)
+            .file_name(filename.to_string())
+            .mime_str("image/*")
+            .map_err(|e| ComfyUIError::new(format!("Invalid MIME type: {}", e)))?;
+
+        let form = multipart::Form::new()
+            .part("image", part)
+            .text("overwrite", overwrite.to_string());
+
+        let url = self.build_url("/upload/image")?;
+        let response = self.client.post(url).multipart(form).send().await?;
+
+        if !response.status().is_success() {
+            return Err(ComfyUIError::Http(reqwest::Error::from(
+                response.error_for_status().unwrap_err()
+            )));
+        }
+
+        let upload_response: UploadImageResponse = response.json().await?;
+        Ok(upload_response)
+    }
+
+    /// Gets an image by filename
+    pub async fn get_image(&self, filename: &str, subfolder: Option<&str>, image_type: Option<&str>) -> Result<bytes::Bytes> {
+        let mut endpoint = format!("/view?filename={}", filename);
+
+        if let Some(subfolder) = subfolder {
+            endpoint.push_str(&format!("&subfolder={}", subfolder));
+        }
+
+        if let Some(image_type) = image_type {
+            endpoint.push_str(&format!("&type={}", image_type));
+        }
+
+        let url = self.build_url(&endpoint)?;
+        let response = self.client.get(url).send().await?;
+
+        if !response.status().is_success() {
+            return Err(ComfyUIError::Http(reqwest::Error::from(
+                response.error_for_status().unwrap_err()
+            )));
+        }
+
+        let bytes = response.bytes().await?;
+        Ok(bytes)
+    }
+
+    /// Gets image metadata
+    pub async fn get_view_metadata(&self, filename: &str, subfolder: Option<&str>) -> Result<ViewMetadata> {
+        let mut endpoint = format!("/view_metadata/{}", filename);
+
+        if let Some(subfolder) = subfolder {
+            endpoint.push_str(&format!("?subfolder={}", subfolder));
+        }
+
+        self.get_with_retry(&endpoint).await
+    }
+
+    /// Converts outputs to HTTP URLs
+    pub fn outputs_to_urls(&self, outputs: &HashMap<String, serde_json::Value>) -> Vec<String> {
+        let mut urls = Vec::new();
+
+        for (_node_id, output) in outputs {
+            if let Some(output_obj) = output.as_object() {
+                if let Some(images) = output_obj.get("images").and_then(|v| v.as_array()) {
+                    for image in images {
+                        if let Some(image_obj) = image.as_object() {
+                            if let (Some(filename), Some(subfolder), Some(image_type)) = (
+                                image_obj.get("filename").and_then(|v| v.as_str()),
+                                image_obj.get("subfolder").and_then(|v| v.as_str()),
+                                image_obj.get("type").and_then(|v| v.as_str()),
+                            ) {
+                                let base_url_str = self.base_url.as_str().trim_end_matches('/');
+                                let url = format!(
+                                    "{}/view?filename={}&subfolder={}&type={}",
+                                    base_url_str, filename, subfolder, image_type
+                                );
+                                urls.push(url);
+                            }
+                        }
+                    }
+                }
+            }
+        }
+
+        urls
+    }
+
+    /// Interrupts the current execution
+    pub async fn interrupt(&self) -> Result<()> {
+        let url = self.build_url("/interrupt")?;
+        let response = self.client.post(url).send().await?;
+
+        if !response.status().is_success() {
+            return Err(ComfyUIError::Http(reqwest::Error::from(
+                response.error_for_status().unwrap_err()
+            )));
+        }
+
+        Ok(())
+    }
+
+    /// Frees memory
+    pub async fn free_memory(&self) -> Result<()> {
+        let url = self.build_url("/free")?;
+        let response = self.client.post(url).send().await?;
+
+        if !response.status().is_success() {
+            return Err(ComfyUIError::Http(reqwest::Error::from(
+                response.error_for_status().unwrap_err()
+            )));
+        }
+
+        Ok(())
+    }
+}

cargos/comfyui-sdk/client/mod.rs

@@ -0,0 +1,10 @@
+//! Client modules for ComfyUI SDK
+
+pub mod http_client;
+pub mod websocket_client;
+pub mod comfyui_client;
+
+// Re-export main types
+pub use http_client::HTTPClient;
+pub use websocket_client::WebSocketClient;
+pub use comfyui_client::ComfyUIClient;

cargos/comfyui-sdk/client/websocket_client.rs

@@ -0,0 +1,226 @@
+//! WebSocket client for ComfyUI real-time events
+
+use std::sync::Arc;
+use tokio::sync::{RwLock, mpsc};
+use tokio_tungstenite::{connect_async, tungstenite::Message};
+use futures_util::StreamExt;
+use url::Url;
+use uuid::Uuid;
+use chrono::Utc;
+
+use crate::types::{
+    ComfyUIClientConfig, WSMessage, ExecutionProgress, ExecutionResult,
+    ExecutionError, ExecutionCallbacks
+};
+use crate::error::{ComfyUIError, Result};
+use crate::utils::event_emitter::EventEmitter;
+
+/// WebSocket client for real-time ComfyUI events
+pub struct WebSocketClient {
+    config: ComfyUIClientConfig,
+    client_id: String,
+    event_emitter: EventEmitter,
+    is_connected: Arc<RwLock<bool>>,
+    shutdown_tx: Option<mpsc::UnboundedSender<()>>,
+}
+
+impl WebSocketClient {
+    /// Creates a new WebSocket client
+    pub fn new(config: ComfyUIClientConfig) -> Self {
+        Self {
+            config,
+            client_id: Uuid::new_v4().to_string(),
+            event_emitter: EventEmitter::new(),
+            is_connected: Arc::new(RwLock::new(false)),
+            shutdown_tx: None,
+        }
+    }
+
+    /// Gets the client ID
+    pub fn client_id(&self) -> &str {
+        &self.client_id
+    }
+
+    /// Checks if connected
+    pub async fn is_connected(&self) -> bool {
+        *self.is_connected.read().await
+    }
+
+    /// Connects to the WebSocket server
+    pub async fn connect(&mut self) -> Result<()> {
+        if self.is_connected().await {
+            return Ok(());
+        }
+
+        let ws_url = self.build_websocket_url()?;
+        log::info!("Connecting to WebSocket: {}", ws_url);
+
+        let (ws_stream, _) = connect_async(&ws_url).await?;
+        let (_ws_sender, mut ws_receiver) = ws_stream.split();
+
+        // Set connected state
+        *self.is_connected.write().await = true;
+
+        // Create shutdown channel
+        let (shutdown_tx, mut shutdown_rx) = mpsc::unbounded_channel();
+        self.shutdown_tx = Some(shutdown_tx);
+
+        // Clone necessary data for the message handling task
+        let event_emitter = self.event_emitter.clone();
+        let is_connected = self.is_connected.clone();
+
+        // Spawn message handling task
+        tokio::spawn(async move {
+            loop {
+                tokio::select! {
+                    // Handle incoming messages
+                    message = ws_receiver.next() => {
+                        match message {
+                            Some(Ok(Message::Text(text))) => {
+                                if let Err(e) = Self::handle_message(&text, &event_emitter).await {
+                                    log::error!("Error handling WebSocket message: {}", e);
+                                }
+                            }
+                            Some(Ok(Message::Close(_))) => {
+                                log::info!("WebSocket connection closed by server");
+                                break;
+                            }
+                            Some(Err(e)) => {
+                                log::error!("WebSocket error: {}", e);
+                                break;
+                            }
+                            None => {
+                                log::info!("WebSocket stream ended");
+                                break;
+                            }
+                            _ => {} // Ignore other message types
+                        }
+                    }
+                    // Handle shutdown signal
+                    _ = shutdown_rx.recv() => {
+                        log::info!("WebSocket shutdown requested");
+                        break;
+                    }
+                }
+            }
+
+            // Set disconnected state
+            *is_connected.write().await = false;
+            log::info!("WebSocket connection closed");
+        });
+
+        Ok(())
+    }
+
+    /// Disconnects from the WebSocket server
+    pub async fn disconnect(&mut self) -> Result<()> {
+        if !self.is_connected().await {
+            return Ok(());
+        }
+
+        if let Some(shutdown_tx) = &self.shutdown_tx {
+            let _ = shutdown_tx.send(());
+        }
+
+        // Wait for disconnection
+        let mut attempts = 0;
+        while self.is_connected().await && attempts < 50 {
+            tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+            attempts += 1;
+        }
+
+        self.shutdown_tx = None;
+        Ok(())
+    }
+
+    /// Registers execution callbacks
+    pub async fn register_callbacks(&self, callbacks: Arc<dyn ExecutionCallbacks>) -> String {
+        let callback_id = Uuid::new_v4().to_string();
+        self.event_emitter.register_callbacks(callback_id.clone(), callbacks).await;
+        callback_id
+    }
+
+    /// Unregisters execution callbacks
+    pub async fn unregister_callbacks(&self, callback_id: &str) {
+        self.event_emitter.unregister_callbacks(callback_id).await;
+    }
+
+    /// Builds the WebSocket URL
+    fn build_websocket_url(&self) -> Result<String> {
+        let base_url = Url::parse(&self.config.base_url)?;
+        let scheme = match base_url.scheme() {
+            "http" => "ws",
+            "https" => "wss",
+            _ => return Err(ComfyUIError::new("Invalid URL scheme")),
+        };
+
+        let ws_url = format!(
+            "{}://{}:{}/ws?clientId={}",
+            scheme,
+            base_url.host_str().unwrap_or("localhost"),
+            base_url.port().unwrap_or(8188),
+            self.client_id
+        );
+
+        Ok(ws_url)
+    }
+
+    /// Handles incoming WebSocket messages
+    async fn handle_message(text: &str, event_emitter: &EventEmitter) -> Result<()> {
+        let message: WSMessage = serde_json::from_str(text)
+            .map_err(|e| ComfyUIError::new(format!("Failed to parse WebSocket message: {}", e)))?;
+
+        match message {
+            WSMessage::Progress { data } => {
+                let progress = ExecutionProgress {
+                    node_id: data.node.unwrap_or_else(|| "unknown".to_string()),
+                    progress: data.value,
+                    max: data.max,
+                    timestamp: Utc::now(),
+                };
+                event_emitter.emit_progress(progress).await;
+            }
+            WSMessage::Executing { data } => {
+                if let Some(node_id) = data.node {
+                    event_emitter.emit_executing(node_id).await;
+                }
+            }
+            WSMessage::Executed { data } => {
+                let result = ExecutionResult {
+                    prompt_id: data.prompt_id,
+                    outputs: data.output,
+                    execution_time: 0, // WebSocket doesn't provide execution time
+                    timestamp: Utc::now(),
+                };
+                event_emitter.emit_executed(result).await;
+            }
+            WSMessage::ExecutionError { data } => {
+                let error = ExecutionError {
+                    node_id: data.node_id,
+                    message: data.message,
+                    details: data.details,
+                    timestamp: Utc::now(),
+                };
+                event_emitter.emit_error(error).await;
+            }
+            WSMessage::Unknown => {
+                log::debug!("Received unknown WebSocket message: {}", text);
+            }
+        }
+
+        Ok(())
+    }
+
+    /// Gets the event emitter for direct access
+    pub fn event_emitter(&self) -> &EventEmitter {
+        &self.event_emitter
+    }
+}
+
+impl Drop for WebSocketClient {
+    fn drop(&mut self) {
+        if let Some(shutdown_tx) = &self.shutdown_tx {
+            let _ = shutdown_tx.send(());
+        }
+    }
+}