Nvidia/BossAI
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: complete Rust backend UI compatibility

Browse Source
This commit is contained in:
boos
2026-05-18 16:52:54 +08:00
parent f2434e725d
commit f986fe0377
41 changed files with 3585 additions and 811 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/routes/comfyui/mod.rs
+93 -34
View File
@@ -10,10 +10,11 @@ use crate::{
AppState,
};
use axum::{
body::Bytes,
extract::{Path, Query, State},
http::{HeaderMap, StatusCode},
http::{HeaderMap, Method, StatusCode, Uri},
response::{Html, IntoResponse, Response},
routing::{get, post},
routing::{any, get, post},
Json, Router,
};
use reqwest::Client;
@@ -21,9 +22,6 @@ use serde_json::Value;
use std::collections::HashMap;
use tracing::info;
/// Default ComfyUI backend address.
const DEFAULT_COMFYUI_BASE: &str = "http://127.0.0.1:8188";
/// Shared HTTP client
fn http_client() -> Client {
Client::builder()
@@ -39,7 +37,7 @@ async fn check_auth(state: &AppState, headers: &HeaderMap) -> Result<(), AppErro
.get("Authorization")
.and_then(|v| v.to_str().ok())
.and_then(|v| v.strip_prefix("Bearer "))
.and_then(|v| v.trim().split_whitespace().next())
.and_then(|v| v.split_whitespace().next())
.ok_or_else(|| AppError::Unauthorized("Missing Authorization header".into()))?;
// Admin token passes
@@ -49,9 +47,10 @@ async fn check_auth(state: &AppState, headers: &HeaderMap) -> Result<(), AppErro
// Downstream API key
if auth.starts_with("sk-") {
let result = authorize_downstream_token(auth, &state.db)
.await
.map_err(|e| AppError::Unauthorized(e))?;
let result =
authorize_downstream_token(auth, &state.db, Some(&state.config.auth.proxy_token))
.await
.map_err(AppError::Unauthorized)?;
match result {
DownstreamAuthResult::Managed { key, .. } => {
@@ -67,9 +66,9 @@ async fn check_auth(state: &AppState, headers: &HeaderMap) -> Result<(), AppErro
}
}
/// Get ComfyUI base URL from env or default
fn comfyui_base() -> String {
std::env::var("COMFYUI_BASE").unwrap_or_else(|_| DEFAULT_COMFYUI_BASE.to_string())
/// Get ComfyUI base URL from embedded/external runtime configuration.
fn comfyui_base(state: &AppState) -> String {
state.comfyui.internal_base_url()
}
// ── Route builder ────────────────────────────────────────────────────
@@ -82,10 +81,11 @@ pub fn routes() -> Router<AppState> {
.route("/api/comfyui/history/{id}", get(proxy_get_history_by_id))
.route("/api/comfyui/queue", get(proxy_get_queue))
.route("/api/comfyui/view", get(proxy_get_view))
// Static proxy (public — assets/JS/CSS)
.route("/api/comfyui/proxy", get(proxy_static_root))
.route("/api/comfyui/proxy/", get(proxy_static_root))
.route("/api/comfyui/proxy/{*path}", get(proxy_static))
.route("/api/comfyui/object_info", get(proxy_get_object_info))
// Public ComfyUI reverse proxy for the embedded iframe UI.
.route("/api/comfyui/proxy", any(proxy_static_root))
.route("/api/comfyui/proxy/", any(proxy_static_root))
.route("/api/comfyui/proxy/{*path}", any(proxy_static))
// HTML page (public)
.route("/comfyui", get(comfyui_page))
}
@@ -100,7 +100,7 @@ async fn proxy_post_prompt(
) -> AppResult<Response> {
info!("POST /api/comfyui/prompt");
check_auth(&state, &headers).await?;
proxy_to_comfyui("POST", "/prompt", None, Some(body)).await
proxy_to_comfyui(&state, "POST", "/prompt", None, Some(body)).await
}
/// GET /api/comfyui/history
@@ -110,7 +110,7 @@ async fn proxy_get_history(
) -> AppResult<Response> {
info!("GET /api/comfyui/history");
check_auth(&state, &headers).await?;
proxy_to_comfyui("GET", "/history", None, None).await
proxy_to_comfyui(&state, "GET", "/history", None, None).await
}
/// GET /api/comfyui/history/{id}
@@ -122,14 +122,24 @@ async fn proxy_get_history_by_id(
info!("GET /api/comfyui/history/{}", id);
check_auth(&state, &headers).await?;
let path = format!("/history/{}", urlencoding(&id));
proxy_to_comfyui("GET", &path, None, None).await
proxy_to_comfyui(&state, "GET", &path, None, None).await
}
/// GET /api/comfyui/queue
async fn proxy_get_queue(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Response> {
info!("GET /api/comfyui/queue");
check_auth(&state, &headers).await?;
proxy_to_comfyui("GET", "/queue", None, None).await
proxy_to_comfyui(&state, "GET", "/queue", None, None).await
}
/// GET /api/comfyui/object_info
async fn proxy_get_object_info(
State(state): State<AppState>,
headers: HeaderMap,
) -> AppResult<Response> {
info!("GET /api/comfyui/object_info");
check_auth(&state, &headers).await?;
proxy_to_comfyui(&state, "GET", "/object_info", None, None).await
}
/// GET /api/comfyui/view?filename=...&type=...&subfolder=...
@@ -143,29 +153,73 @@ async fn proxy_get_view(
let qs = build_query_string(&params);
let path = format!("/view{}", qs);
proxy_to_comfyui_binary("GET", &path).await
proxy_to_comfyui_binary(&state, "GET", &path).await
}
// ── Static proxy (public — assets/JS/CSS for iframe) ─────────────────
// ── Public ComfyUI reverse proxy for iframe UI ───────────────────────
/// GET /api/comfyui/proxy (root) — redirect to proxied root
async fn proxy_static_root() -> AppResult<Response> {
proxy_static(Path("/".to_string())).await
/// /api/comfyui/proxy (root) — proxy to ComfyUI root
async fn proxy_static_root(
State(state): State<AppState>,
method: Method,
uri: Uri,
headers: HeaderMap,
body: Bytes,
) -> AppResult<Response> {
proxy_static_request(state, method, uri, headers, "/".to_string(), body).await
}
/// GET /api/comfyui/proxy/{path} — proxy everything (HTML, assets)
async fn proxy_static(Path(path): Path<String>) -> AppResult<Response> {
info!("GET /api/comfyui/proxy/{}", path);
/// /api/comfyui/proxy/{path} — proxy ComfyUI UI assets and API calls.
async fn proxy_static(
State(state): State<AppState>,
Path(path): Path<String>,
method: Method,
uri: Uri,
headers: HeaderMap,
body: Bytes,
) -> AppResult<Response> {
proxy_static_request(state, method, uri, headers, path, body).await
}
async fn proxy_static_request(
state: AppState,
method: Method,
uri: Uri,
headers: HeaderMap,
path: String,
body: Bytes,
) -> AppResult<Response> {
let target_path = if path.is_empty() || path == "/" {
"/".to_string()
} else {
format!("/{}", path)
};
let target_path = match uri.query() {
Some(query) => format!("{target_path}?{query}"),
None => target_path,
};
let client = http_client();
let url = format!("{}{}", comfyui_base(), target_path);
let resp = client
.get(&url)
let url = format!("{}{}", comfyui_base(&state), target_path);
info!("{} /api/comfyui/proxy{} -> {}", method, uri.path(), url);
let reqwest_method = method
.as_str()
.parse()
.map_err(|e| AppError::BadRequest(format!("Invalid HTTP method: {e}")))?;
let mut req = client.request(reqwest_method, &url);
if let Some(content_type) = headers.get(axum::http::header::CONTENT_TYPE) {
req = req.header(reqwest::header::CONTENT_TYPE, content_type.as_bytes());
}
if let Some(accept) = headers.get(axum::http::header::ACCEPT) {
req = req.header(reqwest::header::ACCEPT, accept.as_bytes());
}
if !body.is_empty() {
req = req.body(body);
}
let resp = req
.send()
.await
.map_err(|e| AppError::UpstreamError(format!("ComfyUI unreachable: {e}")))?;
@@ -306,13 +360,14 @@ fn urlencoding(s: &str) -> String {
/// Proxy to ComfyUI — returns JSON/text response
async fn proxy_to_comfyui(
state: &AppState,
method: &str,
path: &str,
_query: Option<&HashMap<String, String>>,
body: Option<Value>,
) -> AppResult<Response> {
let client = http_client();
let url = format!("{}{}", comfyui_base(), path);
let url = format!("{}{}", comfyui_base(state), path);
let req = match method {
"POST" => {
@@ -347,9 +402,13 @@ async fn proxy_to_comfyui(
}
/// Proxy to ComfyUI — returns binary response (for images)
async fn proxy_to_comfyui_binary(_method: &str, path: &str) -> AppResult<Response> {
async fn proxy_to_comfyui_binary(
state: &AppState,
_method: &str,
path: &str,
) -> AppResult<Response> {
let client = http_client();
let url = format!("{}{}", comfyui_base(), path);
let url = format!("{}{}", comfyui_base(state), path);
let req = client.get(&url);
let resp = req