feat: complete Rust backend UI compatibility

2026-05-18 16:52:54 +08:00
parent f2434e725d
commit f986fe0377
41 changed files with 3585 additions and 811 deletions
@@ -0,0 +1,262 @@
+use crate::config::ComfyUIConfig;
+use serde::Serialize;
+use std::{path::Path, sync::Arc, time::Instant};
+use tokio::{
+    process::{Child, Command},
+    sync::{Mutex, RwLock},
+};
+
+#[derive(Debug, Clone, Serialize)]
+#[serde(tag = "status", rename_all = "snake_case")]
+pub enum ComfyUIStatus {
+    Disabled,
+    Starting,
+    Running { pid: Option<u32>, uptime_secs: u64 },
+    Unhealthy { last_error: String },
+    Stopped,
+}
+
+#[derive(Debug, Clone, Serialize)]
+pub struct ComfyUIStatusResponse {
+    pub status: String,
+    pub pid: Option<u32>,
+    pub uptime_secs: u64,
+    pub mode: String,
+    pub model_backend: String,
+    pub base_url: String,
+    pub version: Option<String>,
+    pub last_error: Option<String>,
+}
+
+pub struct ComfyUIProcess {
+    config: ComfyUIConfig,
+    child: Mutex<Option<Child>>,
+    status: RwLock<ComfyUIStatus>,
+    started_at: Mutex<Option<Instant>>,
+    boosapi_url: String,
+    proxy_token: String,
+}
+
+impl ComfyUIProcess {
+    pub fn new(config: ComfyUIConfig, boosapi_url: String, proxy_token: String) -> Arc<Self> {
+        let initial_status = if config.enabled {
+            ComfyUIStatus::Stopped
+        } else {
+            ComfyUIStatus::Disabled
+        };
+
+        Arc::new(Self {
+            config,
+            child: Mutex::new(None),
+            status: RwLock::new(initial_status),
+            started_at: Mutex::new(None),
+            boosapi_url,
+            proxy_token,
+        })
+    }
+
+    pub fn config(&self) -> &ComfyUIConfig {
+        &self.config
+    }
+
+    pub fn internal_base_url(&self) -> String {
+        if self.config.enabled {
+            format!("http://127.0.0.1:{}", self.config.port)
+        } else {
+            self.config.target_base_url()
+        }
+    }
+
+    pub fn model_backend_url(&self) -> String {
+        format!("{}/v1", self.boosapi_url.trim_end_matches('/'))
+    }
+
+    pub async fn status(&self) -> ComfyUIStatus {
+        let status = self.status.read().await.clone();
+        match status {
+            ComfyUIStatus::Running { pid, .. } => ComfyUIStatus::Running {
+                pid,
+                uptime_secs: self.uptime_secs().await,
+            },
+            other => other,
+        }
+    }
+
+    pub async fn status_response(&self) -> ComfyUIStatusResponse {
+        let status = self.status().await;
+        let base_url = self.internal_base_url();
+        let mode = if self.config.enabled {
+            "embedded"
+        } else {
+            "external"
+        }
+        .to_string();
+
+        match status {
+            ComfyUIStatus::Disabled => ComfyUIStatusResponse {
+                status: "disabled".to_string(),
+                pid: None,
+                uptime_secs: 0,
+                mode,
+                model_backend: self.model_backend_url(),
+                base_url,
+                version: None,
+                last_error: None,
+            },
+            ComfyUIStatus::Starting => ComfyUIStatusResponse {
+                status: "starting".to_string(),
+                pid: None,
+                uptime_secs: 0,
+                mode,
+                model_backend: self.model_backend_url(),
+                base_url,
+                version: None,
+                last_error: None,
+            },
+            ComfyUIStatus::Running { pid, uptime_secs } => ComfyUIStatusResponse {
+                status: "running".to_string(),
+                pid,
+                uptime_secs,
+                mode,
+                model_backend: self.model_backend_url(),
+                base_url,
+                version: None,
+                last_error: None,
+            },
+            ComfyUIStatus::Unhealthy { last_error } => ComfyUIStatusResponse {
+                status: "unhealthy".to_string(),
+                pid: None,
+                uptime_secs: self.uptime_secs().await,
+                mode,
+                model_backend: self.model_backend_url(),
+                base_url,
+                version: None,
+                last_error: Some(last_error),
+            },
+            ComfyUIStatus::Stopped => ComfyUIStatusResponse {
+                status: "stopped".to_string(),
+                pid: None,
+                uptime_secs: 0,
+                mode,
+                model_backend: self.model_backend_url(),
+                base_url,
+                version: None,
+                last_error: None,
+            },
+        }
+    }
+
+    pub async fn start(&self) -> anyhow::Result<()> {
+        if !self.config.enabled {
+            *self.status.write().await = ComfyUIStatus::Disabled;
+            return Ok(());
+        }
+
+        if self.child.lock().await.is_some() {
+            return Ok(());
+        }
+
+        *self.status.write().await = ComfyUIStatus::Starting;
+
+        let install_dir = self
+            .config
+            .install_dir
+            .as_ref()
+            .ok_or_else(|| anyhow::anyhow!("BOOSAPI_COMFYUI__INSTALL_DIR is required"))?;
+        let main_py = install_dir.join("main.py");
+        if !main_py.exists() {
+            let error = format!("ComfyUI main.py not found at {}", main_py.display());
+            *self.status.write().await = ComfyUIStatus::Unhealthy {
+                last_error: error.clone(),
+            };
+            anyhow::bail!(error);
+        }
+
+        write_boosapi_config(install_dir, &self.model_backend_url(), &self.proxy_token)?;
+
+        let mut command = Command::new(&self.config.python_bin);
+        command
+            .arg("main.py")
+            .arg("--listen")
+            .arg("127.0.0.1")
+            .arg("--port")
+            .arg(self.config.port.to_string())
+            .current_dir(install_dir);
+
+        for arg in &self.config.extra_args {
+            if !arg.trim().is_empty() {
+                command.arg(arg);
+            }
+        }
+
+        let child = command.spawn()?;
+        let pid = child.id();
+
+        *self.child.lock().await = Some(child);
+        *self.started_at.lock().await = Some(Instant::now());
+        *self.status.write().await = ComfyUIStatus::Running {
+            pid,
+            uptime_secs: 0,
+        };
+
+        Ok(())
+    }
+
+    pub async fn stop(&self) -> anyhow::Result<()> {
+        let mut child_slot = self.child.lock().await;
+        if let Some(child) = child_slot.as_mut() {
+            if let Err(err) = child.start_kill() {
+                *self.status.write().await = ComfyUIStatus::Unhealthy {
+                    last_error: err.to_string(),
+                };
+                return Err(err.into());
+            }
+            let _ = child.wait().await;
+        }
+
+        *child_slot = None;
+        *self.started_at.lock().await = None;
+        *self.status.write().await = if self.config.enabled {
+            ComfyUIStatus::Stopped
+        } else {
+            ComfyUIStatus::Disabled
+        };
+
+        Ok(())
+    }
+
+    pub async fn restart(&self) -> anyhow::Result<()> {
+        self.stop().await?;
+        self.start().await
+    }
+
+    pub async fn health_check(&self) -> bool {
+        let url = format!("{}/system_stats", self.internal_base_url());
+        match reqwest::Client::new().get(url).send().await {
+            Ok(response) => response.status().is_success(),
+            Err(_) => false,
+        }
+    }
+
+    async fn uptime_secs(&self) -> u64 {
+        self.started_at
+            .lock()
+            .await
+            .map(|started| started.elapsed().as_secs())
+            .unwrap_or(0)
+    }
+}
+
+fn write_boosapi_config(
+    install_dir: &Path,
+    api_base: &str,
+    proxy_token: &str,
+) -> anyhow::Result<()> {
+    let config = format!(
+        "boosapi:\n  api_base: \"{}\"\n  api_key: \"{}\"\n",
+        api_base.trim_end_matches('/'),
+        proxy_token
+    );
+    std::fs::write(install_dir.join("boosapi_config.yaml"), config)?;
+    Ok(())
+}
@@ -1,5 +1,6 @@
 pub mod announcement;
 pub mod billing;
+pub mod comfyui_process;
 pub mod cron;
 pub mod notification;
 pub mod oauth;