refactor(workflow_ui): 抽出 native_proxy 子模块
Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
+19
-637
@@ -1,20 +1,17 @@
|
|||||||
use crate::workflow::executor::{WorkflowProgressCallback, WorkflowProgressEvent};
|
use crate::workflow::executor::{WorkflowProgressCallback, WorkflowProgressEvent};
|
||||||
use crate::{
|
use crate::{
|
||||||
error::{AppError, AppResult},
|
error::{AppError, AppResult},
|
||||||
services::workspace::{self, WorkspaceAssetInput},
|
|
||||||
workflow::{
|
workflow::{
|
||||||
models::{RunContext, WorkflowGraph},
|
models::{RunContext, WorkflowGraph},
|
||||||
},
|
},
|
||||||
AppState,
|
AppState,
|
||||||
};
|
};
|
||||||
use axum::{
|
use axum::{
|
||||||
body::Bytes,
|
|
||||||
extract::{
|
extract::{
|
||||||
multipart::Multipart,
|
|
||||||
ws::{Message, WebSocketUpgrade},
|
ws::{Message, WebSocketUpgrade},
|
||||||
DefaultBodyLimit, FromRequest, Path, Query, State,
|
DefaultBodyLimit, Path, Query, State,
|
||||||
},
|
},
|
||||||
http::{header, HeaderMap, HeaderValue, Method, StatusCode, Uri},
|
http::{header, HeaderMap, Method, StatusCode},
|
||||||
response::{IntoResponse, Response},
|
response::{IntoResponse, Response},
|
||||||
routing::{any, get, post},
|
routing::{any, get, post},
|
||||||
Json, Router,
|
Json, Router,
|
||||||
@@ -27,7 +24,7 @@ use std::{
|
|||||||
collections::{HashMap, HashSet},
|
collections::{HashMap, HashSet},
|
||||||
path::{Path as StdPath, PathBuf},
|
path::{Path as StdPath, PathBuf},
|
||||||
sync::{mpsc, Arc},
|
sync::{mpsc, Arc},
|
||||||
time::{Duration, Instant, SystemTime, UNIX_EPOCH},
|
time::{Duration, Instant},
|
||||||
};
|
};
|
||||||
use tokio::sync::broadcast;
|
use tokio::sync::broadcast;
|
||||||
use tokio_tungstenite::{connect_async, tungstenite};
|
use tokio_tungstenite::{connect_async, tungstenite};
|
||||||
@@ -35,9 +32,9 @@ use tungstenite::client::IntoClientRequest;
|
|||||||
|
|
||||||
static WORKFLOW_WS_SESSIONS: Lazy<parking_lot::Mutex<HashMap<String, broadcast::Sender<Value>>>> =
|
static WORKFLOW_WS_SESSIONS: Lazy<parking_lot::Mutex<HashMap<String, broadcast::Sender<Value>>>> =
|
||||||
Lazy::new(|| parking_lot::Mutex::new(HashMap::new()));
|
Lazy::new(|| parking_lot::Mutex::new(HashMap::new()));
|
||||||
static NATIVE_WORKFLOW_JSON_CACHE: Lazy<parking_lot::Mutex<HashMap<String, CachedJson>>> =
|
pub(super) static NATIVE_WORKFLOW_JSON_CACHE: Lazy<parking_lot::Mutex<HashMap<String, CachedJson>>> =
|
||||||
Lazy::new(|| parking_lot::Mutex::new(HashMap::new()));
|
Lazy::new(|| parking_lot::Mutex::new(HashMap::new()));
|
||||||
const NATIVE_WORKFLOW_CACHE_MAX_AGE: Duration = Duration::from_secs(21600);
|
pub(super) const NATIVE_WORKFLOW_CACHE_MAX_AGE: Duration = Duration::from_secs(21600);
|
||||||
const WORKFLOW_UPLOAD_BODY_LIMIT_BYTES: usize = 256 * 1024 * 1024;
|
const WORKFLOW_UPLOAD_BODY_LIMIT_BYTES: usize = 256 * 1024 * 1024;
|
||||||
const OFFICIAL_WORKFLOW_TEMPLATES_BASE: &str =
|
const OFFICIAL_WORKFLOW_TEMPLATES_BASE: &str =
|
||||||
"https://raw.githubusercontent.com/Comfy-Org/workflow_templates/main/templates";
|
"https://raw.githubusercontent.com/Comfy-Org/workflow_templates/main/templates";
|
||||||
@@ -55,9 +52,9 @@ pub(super) enum WorkflowCaller {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
struct CachedJson {
|
pub(super) struct CachedJson {
|
||||||
inserted_at: Instant,
|
pub(super) inserted_at: Instant,
|
||||||
value: Value,
|
pub(super) value: Value,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Deserialize, Serialize)]
|
#[derive(Debug, Deserialize, Serialize)]
|
||||||
@@ -97,10 +94,10 @@ pub(super) struct UserDataQuery {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Deserialize)]
|
#[derive(Debug, Deserialize)]
|
||||||
struct UploadQuery {
|
pub(super) struct UploadQuery {
|
||||||
subfolder: Option<String>,
|
pub(super) subfolder: Option<String>,
|
||||||
#[serde(rename = "type")]
|
#[serde(rename = "type")]
|
||||||
file_type: Option<String>,
|
pub(super) file_type: Option<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Deserialize, Serialize)]
|
#[derive(Debug, Deserialize, Serialize)]
|
||||||
@@ -144,7 +141,7 @@ use validate::validate_workflow;
|
|||||||
mod queue;
|
mod queue;
|
||||||
use queue::{
|
use queue::{
|
||||||
history, history_by_id, history_mutate, interrupt, internal_logs, internal_logs_raw, queue,
|
history, history_by_id, history_mutate, interrupt, internal_logs, internal_logs_raw, queue,
|
||||||
queue_mutate, queue_prompt, queue_status, repair_native_prompt_payload, view,
|
queue_mutate, queue_prompt, queue_status, view,
|
||||||
};
|
};
|
||||||
mod model_inject;
|
mod model_inject;
|
||||||
use model_inject::{
|
use model_inject::{
|
||||||
@@ -157,6 +154,13 @@ use repair::{
|
|||||||
inject_native_prompt_cloud_api_key, repair_dynamic_combo_nodes, repair_resize_image_mask_nodes,
|
inject_native_prompt_cloud_api_key, repair_dynamic_combo_nodes, repair_resize_image_mask_nodes,
|
||||||
repair_text_generate_ltx2_prompt_nodes, repair_wan2_api_nodes,
|
repair_text_generate_ltx2_prompt_nodes, repair_wan2_api_nodes,
|
||||||
};
|
};
|
||||||
|
mod native_proxy;
|
||||||
|
use native_proxy::{
|
||||||
|
native_workflow_api_proxy, native_workflow_base, native_workflow_enabled, native_workflow_json,
|
||||||
|
native_workflow_json_cached, native_workflow_metadata_ttl, native_workflow_request,
|
||||||
|
native_workflow_response, native_workflow_status_ttl, native_workflow_video_response,
|
||||||
|
native_workflow_view_response, timestamp_ms, upload_image, upload_mask, workflow_cache_status,
|
||||||
|
};
|
||||||
|
|
||||||
pub fn routes() -> Router<AppState> {
|
pub fn routes() -> Router<AppState> {
|
||||||
Router::new()
|
Router::new()
|
||||||
@@ -493,46 +497,6 @@ pub(super) fn auth_caller_optional(state: &AppState, headers: &HeaderMap) -> Opt
|
|||||||
caller_from_headers(state, headers).ok()
|
caller_from_headers(state, headers).ok()
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn workflow_cache_status(
|
|
||||||
State(state): State<AppState>,
|
|
||||||
headers: HeaderMap,
|
|
||||||
) -> AppResult<Json<Value>> {
|
|
||||||
let _ = auth_caller_optional(&state, &headers);
|
|
||||||
let base = native_workflow_base(&state);
|
|
||||||
let now = Instant::now();
|
|
||||||
let mut items = NATIVE_WORKFLOW_JSON_CACHE
|
|
||||||
.lock()
|
|
||||||
.iter()
|
|
||||||
.filter(|(key, _)| key.starts_with(&base))
|
|
||||||
.map(|(key, item)| {
|
|
||||||
json!({
|
|
||||||
"key": key,
|
|
||||||
"ageSecs": now.duration_since(item.inserted_at).as_secs(),
|
|
||||||
"kind": workflow_cache_kind_from_key(key),
|
|
||||||
})
|
|
||||||
})
|
|
||||||
.collect::<Vec<_>>();
|
|
||||||
items.sort_by(|a, b| {
|
|
||||||
a.get("key")
|
|
||||||
.and_then(Value::as_str)
|
|
||||||
.cmp(&b.get("key").and_then(Value::as_str))
|
|
||||||
});
|
|
||||||
Ok(Json(json!({
|
|
||||||
"success": true,
|
|
||||||
"backend": if state.config.comfyui.uses_runninghub_backend() {
|
|
||||||
"runninghub"
|
|
||||||
} else if state.config.comfyui.uses_cloud_backend() {
|
|
||||||
"cloud"
|
|
||||||
} else if state.config.comfyui.uses_native_backend() {
|
|
||||||
"external"
|
|
||||||
} else {
|
|
||||||
"boosapi"
|
|
||||||
},
|
|
||||||
"baseUrl": base,
|
|
||||||
"items": items,
|
|
||||||
})))
|
|
||||||
}
|
|
||||||
|
|
||||||
pub(super) fn workflow_linked_node_ids(links: &[Value]) -> HashSet<i64> {
|
pub(super) fn workflow_linked_node_ids(links: &[Value]) -> HashSet<i64> {
|
||||||
let mut ids = HashSet::new();
|
let mut ids = HashSet::new();
|
||||||
for link in links {
|
for link in links {
|
||||||
@@ -711,571 +675,6 @@ async fn internal_folder_paths(
|
|||||||
Ok(Json(json!({ "input": [], "output": [], "temp": [] })))
|
Ok(Json(json!({ "input": [], "output": [], "temp": [] })))
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn upload_image(
|
|
||||||
State(state): State<AppState>,
|
|
||||||
headers: HeaderMap,
|
|
||||||
Query(query): Query<UploadQuery>,
|
|
||||||
body: Bytes,
|
|
||||||
) -> AppResult<Response> {
|
|
||||||
let caller = auth_caller(&state, &headers).await?;
|
|
||||||
if state.config.comfyui.uses_native_backend() {
|
|
||||||
return upload_image_to_native_backend(&state, &headers, body).await;
|
|
||||||
}
|
|
||||||
let user_id = caller_user_id(&caller);
|
|
||||||
let multipart = Multipart::from_request(
|
|
||||||
axum::http::Request::builder()
|
|
||||||
.method(Method::POST)
|
|
||||||
.header(
|
|
||||||
header::CONTENT_TYPE,
|
|
||||||
headers
|
|
||||||
.get(header::CONTENT_TYPE)
|
|
||||||
.and_then(|value| value.to_str().ok())
|
|
||||||
.unwrap_or_default(),
|
|
||||||
)
|
|
||||||
.body(axum::body::Body::from(body))
|
|
||||||
.map_err(|err| AppError::BadRequest(err.to_string()))?,
|
|
||||||
&state,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.map_err(|err| AppError::BadRequest(err.to_string()))?;
|
|
||||||
upload_image_to_local_backend(user_id, query, multipart).await
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn upload_image_to_local_backend(
|
|
||||||
user_id: Option<i64>,
|
|
||||||
query: UploadQuery,
|
|
||||||
mut multipart: Multipart,
|
|
||||||
) -> AppResult<Response> {
|
|
||||||
let mut image_name: Option<String> = None;
|
|
||||||
let mut image_bytes: Option<Vec<u8>> = None;
|
|
||||||
let mut subfolder = query.subfolder.unwrap_or_default();
|
|
||||||
let mut file_type = query.file_type.unwrap_or_else(|| "input".to_string());
|
|
||||||
|
|
||||||
while let Some(field) = multipart
|
|
||||||
.next_field()
|
|
||||||
.await
|
|
||||||
.map_err(|err| AppError::BadRequest(err.to_string()))?
|
|
||||||
{
|
|
||||||
let name = field.name().unwrap_or_default().to_string();
|
|
||||||
if name == "image" {
|
|
||||||
image_name = field.file_name().map(str::to_string);
|
|
||||||
image_bytes = Some(
|
|
||||||
field
|
|
||||||
.bytes()
|
|
||||||
.await
|
|
||||||
.map_err(|err| AppError::BadRequest(err.to_string()))?
|
|
||||||
.to_vec(),
|
|
||||||
);
|
|
||||||
} else if name == "subfolder" {
|
|
||||||
subfolder = field.text().await.unwrap_or_default();
|
|
||||||
} else if name == "type" {
|
|
||||||
file_type = field.text().await.unwrap_or_else(|_| "input".to_string());
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
let Some(bytes) = image_bytes else {
|
|
||||||
return Err(AppError::BadRequest("image is required".into()));
|
|
||||||
};
|
|
||||||
let file_name =
|
|
||||||
image_name.unwrap_or_else(|| format!("upload-{}.png", uuid::Uuid::new_v4().simple()));
|
|
||||||
let root = user_home(user_id).join("userdata").join(&file_type);
|
|
||||||
let dir = if subfolder.trim().is_empty() {
|
|
||||||
root.clone()
|
|
||||||
} else {
|
|
||||||
root.join(safe_view_path(&subfolder))
|
|
||||||
};
|
|
||||||
tokio::fs::create_dir_all(&dir)
|
|
||||||
.await
|
|
||||||
.map_err(|err| AppError::Internal(anyhow::anyhow!(err)))?;
|
|
||||||
let path = dir.join(&file_name);
|
|
||||||
tokio::fs::write(&path, bytes)
|
|
||||||
.await
|
|
||||||
.map_err(|err| AppError::Internal(anyhow::anyhow!(err)))?;
|
|
||||||
Ok(Json(json!({
|
|
||||||
"name": file_name,
|
|
||||||
"subfolder": subfolder,
|
|
||||||
"type": file_type,
|
|
||||||
}))
|
|
||||||
.into_response())
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn upload_image_to_native_backend(
|
|
||||||
state: &AppState,
|
|
||||||
headers: &HeaderMap,
|
|
||||||
body: Bytes,
|
|
||||||
) -> AppResult<Response> {
|
|
||||||
if body.is_empty() {
|
|
||||||
return Err(AppError::BadRequest("empty upload request body".into()));
|
|
||||||
}
|
|
||||||
let content_type = headers
|
|
||||||
.get(header::CONTENT_TYPE)
|
|
||||||
.and_then(|value| value.to_str().ok())
|
|
||||||
.filter(|value| {
|
|
||||||
value
|
|
||||||
.to_ascii_lowercase()
|
|
||||||
.starts_with("multipart/form-data")
|
|
||||||
})
|
|
||||||
.ok_or_else(|| AppError::BadRequest("multipart/form-data content-type is required".into()))?
|
|
||||||
.to_string();
|
|
||||||
let response = native_workflow_request(state, Method::POST, "/upload/image", None)
|
|
||||||
.header(reqwest::header::CONTENT_TYPE, content_type)
|
|
||||||
.body(body)
|
|
||||||
.send()
|
|
||||||
.await
|
|
||||||
.map_err(|err| AppError::UpstreamError(format!("ComfyUI upload failed: {err}")))?;
|
|
||||||
let status = response.status();
|
|
||||||
let content_type = response
|
|
||||||
.headers()
|
|
||||||
.get(reqwest::header::CONTENT_TYPE)
|
|
||||||
.and_then(|value| value.to_str().ok())
|
|
||||||
.unwrap_or("application/json")
|
|
||||||
.to_string();
|
|
||||||
let bytes = response
|
|
||||||
.bytes()
|
|
||||||
.await
|
|
||||||
.map_err(|err| AppError::UpstreamError(err.to_string()))?;
|
|
||||||
if !status.is_success() {
|
|
||||||
let body = String::from_utf8_lossy(&bytes);
|
|
||||||
tracing::warn!(status = status.as_u16(), body = %body, "Workflow native upload failed");
|
|
||||||
}
|
|
||||||
Ok((
|
|
||||||
StatusCode::from_u16(status.as_u16()).unwrap_or(StatusCode::BAD_GATEWAY),
|
|
||||||
[(header::CONTENT_TYPE, content_type)],
|
|
||||||
bytes,
|
|
||||||
)
|
|
||||||
.into_response())
|
|
||||||
}
|
|
||||||
|
|
||||||
fn timestamp_ms() -> i64 {
|
|
||||||
SystemTime::now()
|
|
||||||
.duration_since(UNIX_EPOCH)
|
|
||||||
.map(|duration| duration.as_millis() as i64)
|
|
||||||
.unwrap_or_default()
|
|
||||||
}
|
|
||||||
|
|
||||||
fn native_workflow_enabled(state: &AppState) -> bool {
|
|
||||||
state.config.comfyui.uses_native_backend()
|
|
||||||
}
|
|
||||||
|
|
||||||
pub(super) fn native_workflow_metadata_ttl(state: &AppState) -> Duration {
|
|
||||||
if state.config.comfyui.uses_runninghub_backend() {
|
|
||||||
Duration::from_secs(21600)
|
|
||||||
} else {
|
|
||||||
Duration::from_secs(300)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
pub(super) fn native_workflow_status_ttl(state: &AppState) -> Duration {
|
|
||||||
if state.config.comfyui.uses_runninghub_backend() {
|
|
||||||
Duration::from_secs(120)
|
|
||||||
} else {
|
|
||||||
Duration::from_secs(30)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn native_workflow_base(state: &AppState) -> String {
|
|
||||||
state.config.comfyui.target_base_url()
|
|
||||||
}
|
|
||||||
|
|
||||||
fn native_workflow_http_path(state: &AppState, path: &str) -> String {
|
|
||||||
if state.config.comfyui.uses_cloud_backend() {
|
|
||||||
format!(
|
|
||||||
"/api{}",
|
|
||||||
path.strip_prefix('/')
|
|
||||||
.map(|value| format!("/{value}"))
|
|
||||||
.unwrap_or_else(|| path.to_string())
|
|
||||||
)
|
|
||||||
} else {
|
|
||||||
path.to_string()
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
pub(super) fn native_workflow_request(
|
|
||||||
state: &AppState,
|
|
||||||
method: Method,
|
|
||||||
path: &str,
|
|
||||||
body: Option<Value>,
|
|
||||||
) -> reqwest::RequestBuilder {
|
|
||||||
let client = reqwest::Client::new();
|
|
||||||
let url = format!(
|
|
||||||
"{}{}",
|
|
||||||
native_workflow_base(state),
|
|
||||||
native_workflow_http_path(state, path)
|
|
||||||
);
|
|
||||||
let reqwest_method =
|
|
||||||
reqwest::Method::from_bytes(method.as_str().as_bytes()).unwrap_or(reqwest::Method::GET);
|
|
||||||
let mut request = client.request(reqwest_method, url);
|
|
||||||
if let Some(api_key) = state
|
|
||||||
.config
|
|
||||||
.comfyui
|
|
||||||
.api_key
|
|
||||||
.as_deref()
|
|
||||||
.map(str::trim)
|
|
||||||
.filter(|value| !value.is_empty())
|
|
||||||
{
|
|
||||||
request = request
|
|
||||||
.header("X-API-Key", api_key)
|
|
||||||
.header("Authorization", format!("Bearer {api_key}"));
|
|
||||||
}
|
|
||||||
if let Some(body) = body {
|
|
||||||
request = request.json(&body);
|
|
||||||
}
|
|
||||||
request
|
|
||||||
}
|
|
||||||
|
|
||||||
pub(super) async fn native_workflow_json(
|
|
||||||
state: &AppState,
|
|
||||||
method: Method,
|
|
||||||
path: &str,
|
|
||||||
body: Option<Value>,
|
|
||||||
) -> AppResult<Json<Value>> {
|
|
||||||
let response = native_workflow_request(state, method, path, body)
|
|
||||||
.send()
|
|
||||||
.await
|
|
||||||
.map_err(|err| AppError::UpstreamError(format!("ComfyUI backend unreachable: {err}")))?;
|
|
||||||
let status = response.status();
|
|
||||||
let text = response
|
|
||||||
.text()
|
|
||||||
.await
|
|
||||||
.map_err(|err| AppError::UpstreamError(err.to_string()))?;
|
|
||||||
let value = serde_json::from_str::<Value>(&text).unwrap_or_else(|_| json!({ "body": text }));
|
|
||||||
if !status.is_success() {
|
|
||||||
return Err(AppError::UpstreamError(value.to_string()));
|
|
||||||
}
|
|
||||||
Ok(Json(value))
|
|
||||||
}
|
|
||||||
|
|
||||||
pub(super) async fn native_workflow_json_cached(
|
|
||||||
state: &AppState,
|
|
||||||
method: Method,
|
|
||||||
path: &str,
|
|
||||||
body: Option<Value>,
|
|
||||||
ttl: Duration,
|
|
||||||
) -> AppResult<Json<Value>> {
|
|
||||||
if method != Method::GET || body.is_some() || ttl.is_zero() {
|
|
||||||
return native_workflow_json(state, method, path, body).await;
|
|
||||||
}
|
|
||||||
|
|
||||||
let cache_key = format!(
|
|
||||||
"{}|{}|{}",
|
|
||||||
native_workflow_base(state),
|
|
||||||
method.as_str(),
|
|
||||||
native_workflow_http_path(state, path)
|
|
||||||
);
|
|
||||||
let now = Instant::now();
|
|
||||||
if let Some(cached) = NATIVE_WORKFLOW_JSON_CACHE.lock().get(&cache_key).cloned() {
|
|
||||||
if now.duration_since(cached.inserted_at) <= ttl {
|
|
||||||
return Ok(Json(cached.value));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
let value = match native_workflow_json(state, method, path, None).await {
|
|
||||||
Ok(value) => value.0,
|
|
||||||
Err(err) => {
|
|
||||||
if let Some(cached) = NATIVE_WORKFLOW_JSON_CACHE.lock().get(&cache_key).cloned() {
|
|
||||||
tracing::warn!(
|
|
||||||
path,
|
|
||||||
error = %err,
|
|
||||||
"using stale ComfyUI workflow metadata cache after refresh failure"
|
|
||||||
);
|
|
||||||
return Ok(Json(cached.value));
|
|
||||||
}
|
|
||||||
return Err(err);
|
|
||||||
}
|
|
||||||
};
|
|
||||||
let mut cache = NATIVE_WORKFLOW_JSON_CACHE.lock();
|
|
||||||
if cache.len() > 64 {
|
|
||||||
cache.retain(|_, item| {
|
|
||||||
now.duration_since(item.inserted_at) <= NATIVE_WORKFLOW_CACHE_MAX_AGE
|
|
||||||
});
|
|
||||||
}
|
|
||||||
cache.insert(
|
|
||||||
cache_key,
|
|
||||||
CachedJson {
|
|
||||||
inserted_at: now,
|
|
||||||
value: value.clone(),
|
|
||||||
},
|
|
||||||
);
|
|
||||||
Ok(Json(value))
|
|
||||||
}
|
|
||||||
|
|
||||||
fn workflow_cache_kind_from_key(key: &str) -> &'static str {
|
|
||||||
if key.contains("/object_info") {
|
|
||||||
"object_info"
|
|
||||||
} else if key.contains("/experiment/models") {
|
|
||||||
"models"
|
|
||||||
} else if key.contains("/system_stats") {
|
|
||||||
"status"
|
|
||||||
} else {
|
|
||||||
"json"
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
pub(super) async fn native_workflow_response(
|
|
||||||
state: &AppState,
|
|
||||||
method: Method,
|
|
||||||
path: &str,
|
|
||||||
body: Option<Value>,
|
|
||||||
) -> AppResult<Response> {
|
|
||||||
let response = native_workflow_request(state, method, path, body)
|
|
||||||
.send()
|
|
||||||
.await
|
|
||||||
.map_err(|err| AppError::UpstreamError(format!("ComfyUI backend unreachable: {err}")))?;
|
|
||||||
let status = response.status();
|
|
||||||
let content_type = response
|
|
||||||
.headers()
|
|
||||||
.get(reqwest::header::CONTENT_TYPE)
|
|
||||||
.and_then(|value| value.to_str().ok())
|
|
||||||
.unwrap_or("application/json")
|
|
||||||
.to_string();
|
|
||||||
let bytes = response
|
|
||||||
.bytes()
|
|
||||||
.await
|
|
||||||
.map_err(|err| AppError::UpstreamError(err.to_string()))?;
|
|
||||||
Ok((
|
|
||||||
StatusCode::from_u16(status.as_u16()).unwrap_or(StatusCode::BAD_GATEWAY),
|
|
||||||
[(header::CONTENT_TYPE, content_type)],
|
|
||||||
bytes,
|
|
||||||
)
|
|
||||||
.into_response())
|
|
||||||
}
|
|
||||||
|
|
||||||
pub(super) async fn native_workflow_view_response(
|
|
||||||
state: &AppState,
|
|
||||||
caller: &WorkflowCaller,
|
|
||||||
params: &HashMap<String, String>,
|
|
||||||
) -> AppResult<Response> {
|
|
||||||
native_workflow_file_response(state, caller, "/view", params).await
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn native_workflow_video_response(
|
|
||||||
state: &AppState,
|
|
||||||
caller: &WorkflowCaller,
|
|
||||||
params: &HashMap<String, String>,
|
|
||||||
) -> AppResult<Response> {
|
|
||||||
native_workflow_file_response(state, caller, "/viewvideo", params).await
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn native_workflow_file_response(
|
|
||||||
state: &AppState,
|
|
||||||
caller: &WorkflowCaller,
|
|
||||||
path: &str,
|
|
||||||
params: &HashMap<String, String>,
|
|
||||||
) -> AppResult<Response> {
|
|
||||||
let query = build_query_string(params);
|
|
||||||
let response = native_workflow_request(state, Method::GET, &format!("{path}{query}"), None)
|
|
||||||
.send()
|
|
||||||
.await
|
|
||||||
.map_err(|err| AppError::UpstreamError(format!("ComfyUI backend unreachable: {err}")))?;
|
|
||||||
let status = response.status();
|
|
||||||
let content_type = response
|
|
||||||
.headers()
|
|
||||||
.get(reqwest::header::CONTENT_TYPE)
|
|
||||||
.and_then(|value| value.to_str().ok())
|
|
||||||
.unwrap_or("application/octet-stream")
|
|
||||||
.to_string();
|
|
||||||
let bytes = response
|
|
||||||
.bytes()
|
|
||||||
.await
|
|
||||||
.map_err(|err| AppError::UpstreamError(err.to_string()))?;
|
|
||||||
if status.is_success() && matches!(path, "/view" | "/viewvideo") {
|
|
||||||
archive_native_workflow_asset(state, caller, path, params, &content_type, bytes.as_ref())
|
|
||||||
.await;
|
|
||||||
}
|
|
||||||
Ok((
|
|
||||||
StatusCode::from_u16(status.as_u16()).unwrap_or(StatusCode::BAD_GATEWAY),
|
|
||||||
[(header::CONTENT_TYPE, content_type)],
|
|
||||||
bytes,
|
|
||||||
)
|
|
||||||
.into_response())
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn archive_native_workflow_asset(
|
|
||||||
state: &AppState,
|
|
||||||
caller: &WorkflowCaller,
|
|
||||||
path: &str,
|
|
||||||
params: &HashMap<String, String>,
|
|
||||||
content_type: &str,
|
|
||||||
bytes: &[u8],
|
|
||||||
) {
|
|
||||||
let file_type = params.get("type").map(String::as_str).unwrap_or("").trim();
|
|
||||||
if file_type != "output" {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
let Some(kind) = workspace_kind_from_mime(content_type) else {
|
|
||||||
return;
|
|
||||||
};
|
|
||||||
let filename = params
|
|
||||||
.get("filename")
|
|
||||||
.map(String::as_str)
|
|
||||||
.unwrap_or("")
|
|
||||||
.trim();
|
|
||||||
if filename.is_empty() || bytes.is_empty() {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
let user_id = caller_user_id(caller);
|
|
||||||
let source = "workflow-studio";
|
|
||||||
let already_exists = sqlx::query_scalar::<_, i64>(
|
|
||||||
"SELECT COUNT(*)
|
|
||||||
FROM workspace_assets
|
|
||||||
WHERE deleted_at IS NULL
|
|
||||||
AND (
|
|
||||||
(? = 1 AND user_id IS NULL)
|
|
||||||
OR (? = 0 AND user_id = ?)
|
|
||||||
)
|
|
||||||
AND source = ?
|
|
||||||
AND json_extract(metadata, '$.workflowUi.filename') = ?
|
|
||||||
AND json_extract(metadata, '$.workflowUi.type') = ?
|
|
||||||
AND json_extract(metadata, '$.workflowUi.subfolder') = ?",
|
|
||||||
)
|
|
||||||
.bind(user_id.is_none())
|
|
||||||
.bind(user_id.is_none())
|
|
||||||
.bind(user_id)
|
|
||||||
.bind(source)
|
|
||||||
.bind(filename)
|
|
||||||
.bind(file_type)
|
|
||||||
.bind(params.get("subfolder").map(String::as_str).unwrap_or(""))
|
|
||||||
.fetch_one(&*state.db)
|
|
||||||
.await
|
|
||||||
.unwrap_or_default();
|
|
||||||
if already_exists > 0 {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
|
|
||||||
let input = WorkspaceAssetInput {
|
|
||||||
user_id,
|
|
||||||
source: source.to_string(),
|
|
||||||
kind: kind.to_string(),
|
|
||||||
title: Some(filename.chars().take(80).collect()),
|
|
||||||
prompt: None,
|
|
||||||
model: None,
|
|
||||||
remote_url: None,
|
|
||||||
metadata: json!({
|
|
||||||
"workspaceAsset": { "role": "output" },
|
|
||||||
"workflowUi": {
|
|
||||||
"filename": filename,
|
|
||||||
"type": file_type,
|
|
||||||
"subfolder": params.get("subfolder").map(String::as_str).unwrap_or(""),
|
|
||||||
"endpoint": path,
|
|
||||||
}
|
|
||||||
}),
|
|
||||||
};
|
|
||||||
if let Err(err) = workspace::record_bytes_asset(state, input, bytes, content_type).await {
|
|
||||||
tracing::warn!(%err, filename, "failed to archive workflow-ui asset into workspace");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn native_workflow_api_proxy(
|
|
||||||
State(state): State<AppState>,
|
|
||||||
headers: HeaderMap,
|
|
||||||
Path(path): Path<String>,
|
|
||||||
method: Method,
|
|
||||||
uri: Uri,
|
|
||||||
body: Bytes,
|
|
||||||
) -> AppResult<Response> {
|
|
||||||
let _ = auth_caller_optional(&state, &headers);
|
|
||||||
if !native_workflow_enabled(&state) {
|
|
||||||
return Err(AppError::NotFound("Workflow API route not found".into()));
|
|
||||||
}
|
|
||||||
let target = match uri.query() {
|
|
||||||
Some(query) => format!("/{path}?{query}"),
|
|
||||||
None => format!("/{path}"),
|
|
||||||
};
|
|
||||||
if state.config.comfyui.uses_cloud_backend() {
|
|
||||||
return native_workflow_raw_response(&state, method, &target, &headers, body).await;
|
|
||||||
}
|
|
||||||
let body = if body.is_empty() {
|
|
||||||
None
|
|
||||||
} else {
|
|
||||||
Some(
|
|
||||||
serde_json::from_slice::<Value>(&body)
|
|
||||||
.unwrap_or_else(|_| json!({ "body": String::from_utf8_lossy(&body).to_string() })),
|
|
||||||
)
|
|
||||||
};
|
|
||||||
native_workflow_response(&state, method, &target, body).await
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn native_workflow_raw_response(
|
|
||||||
state: &AppState,
|
|
||||||
method: Method,
|
|
||||||
path: &str,
|
|
||||||
headers: &HeaderMap,
|
|
||||||
body: Bytes,
|
|
||||||
) -> AppResult<Response> {
|
|
||||||
let is_prompt_post = method == Method::POST && path.split('?').next() == Some("/prompt");
|
|
||||||
let mut request = native_workflow_request(state, method, path, None);
|
|
||||||
if let Some(content_type) = headers
|
|
||||||
.get(header::CONTENT_TYPE)
|
|
||||||
.and_then(|value| value.to_str().ok())
|
|
||||||
{
|
|
||||||
request = request.header(reqwest::header::CONTENT_TYPE, content_type);
|
|
||||||
}
|
|
||||||
if let Some(accept) = headers
|
|
||||||
.get(header::ACCEPT)
|
|
||||||
.and_then(|value| value.to_str().ok())
|
|
||||||
{
|
|
||||||
request = request.header(reqwest::header::ACCEPT, accept);
|
|
||||||
}
|
|
||||||
if !body.is_empty() {
|
|
||||||
if is_prompt_post {
|
|
||||||
if let Ok(mut value) = serde_json::from_slice::<Value>(&body) {
|
|
||||||
repair_native_prompt_body_value(state, &mut value).await;
|
|
||||||
request = request.json(&value);
|
|
||||||
} else {
|
|
||||||
request = request.body(body);
|
|
||||||
}
|
|
||||||
} else {
|
|
||||||
request = request.body(body);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
let response = request
|
|
||||||
.send()
|
|
||||||
.await
|
|
||||||
.map_err(|err| AppError::UpstreamError(format!("ComfyUI backend unreachable: {err}")))?;
|
|
||||||
let status = response.status();
|
|
||||||
let response_headers = response.headers().clone();
|
|
||||||
let content_type = response_headers
|
|
||||||
.get(reqwest::header::CONTENT_TYPE)
|
|
||||||
.and_then(|value| value.to_str().ok())
|
|
||||||
.unwrap_or("application/json")
|
|
||||||
.to_string();
|
|
||||||
let bytes = response
|
|
||||||
.bytes()
|
|
||||||
.await
|
|
||||||
.map_err(|err| AppError::UpstreamError(err.to_string()))?;
|
|
||||||
let mut builder = Response::builder()
|
|
||||||
.status(StatusCode::from_u16(status.as_u16()).unwrap_or(StatusCode::BAD_GATEWAY))
|
|
||||||
.header(header::CONTENT_TYPE, content_type);
|
|
||||||
for name in [
|
|
||||||
header::CACHE_CONTROL,
|
|
||||||
header::ETAG,
|
|
||||||
header::LAST_MODIFIED,
|
|
||||||
header::LOCATION,
|
|
||||||
header::SET_COOKIE,
|
|
||||||
header::WWW_AUTHENTICATE,
|
|
||||||
] {
|
|
||||||
if let Some(value) = response_headers
|
|
||||||
.get(name.as_str())
|
|
||||||
.and_then(|value| value.to_str().ok())
|
|
||||||
.and_then(|value| HeaderValue::from_str(value).ok())
|
|
||||||
{
|
|
||||||
builder = builder.header(name, value);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
builder
|
|
||||||
.body(axum::body::Body::from(bytes))
|
|
||||||
.map_err(|err| AppError::Internal(anyhow::anyhow!(err)))
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn repair_native_prompt_body_value(state: &AppState, value: &mut Value) {
|
|
||||||
if let Ok(mut payload) = serde_json::from_value::<PromptPayload>(value.clone()) {
|
|
||||||
repair_native_prompt_payload(state, &mut payload).await;
|
|
||||||
if let Ok(repaired) = serde_json::to_value(payload) {
|
|
||||||
*value = repaired;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
pub(super) fn build_query_string(params: &HashMap<String, String>) -> String {
|
pub(super) fn build_query_string(params: &HashMap<String, String>) -> String {
|
||||||
if params.is_empty() {
|
if params.is_empty() {
|
||||||
return String::new();
|
return String::new();
|
||||||
@@ -1805,23 +1204,6 @@ pub(super) fn publish_run_events(
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn upload_mask(
|
|
||||||
State(state): State<AppState>,
|
|
||||||
headers: HeaderMap,
|
|
||||||
body: Bytes,
|
|
||||||
) -> AppResult<Response> {
|
|
||||||
upload_image(
|
|
||||||
State(state),
|
|
||||||
headers,
|
|
||||||
Query(UploadQuery {
|
|
||||||
subfolder: None,
|
|
||||||
file_type: Some("input".into()),
|
|
||||||
}),
|
|
||||||
body,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn experiment_models(
|
async fn experiment_models(
|
||||||
State(state): State<AppState>,
|
State(state): State<AppState>,
|
||||||
headers: HeaderMap,
|
headers: HeaderMap,
|
||||||
|
|||||||
@@ -0,0 +1,646 @@
|
|||||||
|
use super::auth::{auth_caller, caller_user_id};
|
||||||
|
use super::queue::repair_native_prompt_payload;
|
||||||
|
use super::{
|
||||||
|
auth_caller_optional, build_query_string, safe_view_path, user_home, workspace_kind_from_mime,
|
||||||
|
CachedJson, PromptPayload, UploadQuery, WorkflowCaller, NATIVE_WORKFLOW_CACHE_MAX_AGE,
|
||||||
|
NATIVE_WORKFLOW_JSON_CACHE,
|
||||||
|
};
|
||||||
|
use crate::{
|
||||||
|
error::{AppError, AppResult},
|
||||||
|
services::workspace::{self, WorkspaceAssetInput},
|
||||||
|
AppState,
|
||||||
|
};
|
||||||
|
use axum::{
|
||||||
|
body::Bytes,
|
||||||
|
extract::{multipart::Multipart, FromRequest, Path, Query, State},
|
||||||
|
http::{header, HeaderMap, HeaderValue, Method, StatusCode, Uri},
|
||||||
|
response::{IntoResponse, Response},
|
||||||
|
Json,
|
||||||
|
};
|
||||||
|
use serde_json::{json, Value};
|
||||||
|
use std::{
|
||||||
|
collections::HashMap,
|
||||||
|
time::{Duration, Instant, SystemTime, UNIX_EPOCH},
|
||||||
|
};
|
||||||
|
|
||||||
|
pub(super) async fn workflow_cache_status(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
let _ = auth_caller_optional(&state, &headers);
|
||||||
|
let base = native_workflow_base(&state);
|
||||||
|
let now = Instant::now();
|
||||||
|
let mut items = NATIVE_WORKFLOW_JSON_CACHE
|
||||||
|
.lock()
|
||||||
|
.iter()
|
||||||
|
.filter(|(key, _)| key.starts_with(&base))
|
||||||
|
.map(|(key, item)| {
|
||||||
|
json!({
|
||||||
|
"key": key,
|
||||||
|
"ageSecs": now.duration_since(item.inserted_at).as_secs(),
|
||||||
|
"kind": workflow_cache_kind_from_key(key),
|
||||||
|
})
|
||||||
|
})
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
items.sort_by(|a, b| {
|
||||||
|
a.get("key")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.cmp(&b.get("key").and_then(Value::as_str))
|
||||||
|
});
|
||||||
|
Ok(Json(json!({
|
||||||
|
"success": true,
|
||||||
|
"backend": if state.config.comfyui.uses_runninghub_backend() {
|
||||||
|
"runninghub"
|
||||||
|
} else if state.config.comfyui.uses_cloud_backend() {
|
||||||
|
"cloud"
|
||||||
|
} else if state.config.comfyui.uses_native_backend() {
|
||||||
|
"external"
|
||||||
|
} else {
|
||||||
|
"boosapi"
|
||||||
|
},
|
||||||
|
"baseUrl": base,
|
||||||
|
"items": items,
|
||||||
|
})))
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) async fn upload_image(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
Query(query): Query<UploadQuery>,
|
||||||
|
body: Bytes,
|
||||||
|
) -> AppResult<Response> {
|
||||||
|
let caller = auth_caller(&state, &headers).await?;
|
||||||
|
if state.config.comfyui.uses_native_backend() {
|
||||||
|
return upload_image_to_native_backend(&state, &headers, body).await;
|
||||||
|
}
|
||||||
|
let user_id = caller_user_id(&caller);
|
||||||
|
let multipart = Multipart::from_request(
|
||||||
|
axum::http::Request::builder()
|
||||||
|
.method(Method::POST)
|
||||||
|
.header(
|
||||||
|
header::CONTENT_TYPE,
|
||||||
|
headers
|
||||||
|
.get(header::CONTENT_TYPE)
|
||||||
|
.and_then(|value| value.to_str().ok())
|
||||||
|
.unwrap_or_default(),
|
||||||
|
)
|
||||||
|
.body(axum::body::Body::from(body))
|
||||||
|
.map_err(|err| AppError::BadRequest(err.to_string()))?,
|
||||||
|
&state,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.map_err(|err| AppError::BadRequest(err.to_string()))?;
|
||||||
|
upload_image_to_local_backend(user_id, query, multipart).await
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn upload_image_to_local_backend(
|
||||||
|
user_id: Option<i64>,
|
||||||
|
query: UploadQuery,
|
||||||
|
mut multipart: Multipart,
|
||||||
|
) -> AppResult<Response> {
|
||||||
|
let mut image_name: Option<String> = None;
|
||||||
|
let mut image_bytes: Option<Vec<u8>> = None;
|
||||||
|
let mut subfolder = query.subfolder.unwrap_or_default();
|
||||||
|
let mut file_type = query.file_type.unwrap_or_else(|| "input".to_string());
|
||||||
|
|
||||||
|
while let Some(field) = multipart
|
||||||
|
.next_field()
|
||||||
|
.await
|
||||||
|
.map_err(|err| AppError::BadRequest(err.to_string()))?
|
||||||
|
{
|
||||||
|
let name = field.name().unwrap_or_default().to_string();
|
||||||
|
if name == "image" {
|
||||||
|
image_name = field.file_name().map(str::to_string);
|
||||||
|
image_bytes = Some(
|
||||||
|
field
|
||||||
|
.bytes()
|
||||||
|
.await
|
||||||
|
.map_err(|err| AppError::BadRequest(err.to_string()))?
|
||||||
|
.to_vec(),
|
||||||
|
);
|
||||||
|
} else if name == "subfolder" {
|
||||||
|
subfolder = field.text().await.unwrap_or_default();
|
||||||
|
} else if name == "type" {
|
||||||
|
file_type = field.text().await.unwrap_or_else(|_| "input".to_string());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let Some(bytes) = image_bytes else {
|
||||||
|
return Err(AppError::BadRequest("image is required".into()));
|
||||||
|
};
|
||||||
|
let file_name =
|
||||||
|
image_name.unwrap_or_else(|| format!("upload-{}.png", uuid::Uuid::new_v4().simple()));
|
||||||
|
let root = user_home(user_id).join("userdata").join(&file_type);
|
||||||
|
let dir = if subfolder.trim().is_empty() {
|
||||||
|
root.clone()
|
||||||
|
} else {
|
||||||
|
root.join(safe_view_path(&subfolder))
|
||||||
|
};
|
||||||
|
tokio::fs::create_dir_all(&dir)
|
||||||
|
.await
|
||||||
|
.map_err(|err| AppError::Internal(anyhow::anyhow!(err)))?;
|
||||||
|
let path = dir.join(&file_name);
|
||||||
|
tokio::fs::write(&path, bytes)
|
||||||
|
.await
|
||||||
|
.map_err(|err| AppError::Internal(anyhow::anyhow!(err)))?;
|
||||||
|
Ok(Json(json!({
|
||||||
|
"name": file_name,
|
||||||
|
"subfolder": subfolder,
|
||||||
|
"type": file_type,
|
||||||
|
}))
|
||||||
|
.into_response())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn upload_image_to_native_backend(
|
||||||
|
state: &AppState,
|
||||||
|
headers: &HeaderMap,
|
||||||
|
body: Bytes,
|
||||||
|
) -> AppResult<Response> {
|
||||||
|
if body.is_empty() {
|
||||||
|
return Err(AppError::BadRequest("empty upload request body".into()));
|
||||||
|
}
|
||||||
|
let content_type = headers
|
||||||
|
.get(header::CONTENT_TYPE)
|
||||||
|
.and_then(|value| value.to_str().ok())
|
||||||
|
.filter(|value| {
|
||||||
|
value
|
||||||
|
.to_ascii_lowercase()
|
||||||
|
.starts_with("multipart/form-data")
|
||||||
|
})
|
||||||
|
.ok_or_else(|| AppError::BadRequest("multipart/form-data content-type is required".into()))?
|
||||||
|
.to_string();
|
||||||
|
let response = native_workflow_request(state, Method::POST, "/upload/image", None)
|
||||||
|
.header(reqwest::header::CONTENT_TYPE, content_type)
|
||||||
|
.body(body)
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.map_err(|err| AppError::UpstreamError(format!("ComfyUI upload failed: {err}")))?;
|
||||||
|
let status = response.status();
|
||||||
|
let content_type = response
|
||||||
|
.headers()
|
||||||
|
.get(reqwest::header::CONTENT_TYPE)
|
||||||
|
.and_then(|value| value.to_str().ok())
|
||||||
|
.unwrap_or("application/json")
|
||||||
|
.to_string();
|
||||||
|
let bytes = response
|
||||||
|
.bytes()
|
||||||
|
.await
|
||||||
|
.map_err(|err| AppError::UpstreamError(err.to_string()))?;
|
||||||
|
if !status.is_success() {
|
||||||
|
let body = String::from_utf8_lossy(&bytes);
|
||||||
|
tracing::warn!(status = status.as_u16(), body = %body, "Workflow native upload failed");
|
||||||
|
}
|
||||||
|
Ok((
|
||||||
|
StatusCode::from_u16(status.as_u16()).unwrap_or(StatusCode::BAD_GATEWAY),
|
||||||
|
[(header::CONTENT_TYPE, content_type)],
|
||||||
|
bytes,
|
||||||
|
)
|
||||||
|
.into_response())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn timestamp_ms() -> i64 {
|
||||||
|
SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.map(|duration| duration.as_millis() as i64)
|
||||||
|
.unwrap_or_default()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn native_workflow_enabled(state: &AppState) -> bool {
|
||||||
|
state.config.comfyui.uses_native_backend()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn native_workflow_metadata_ttl(state: &AppState) -> Duration {
|
||||||
|
if state.config.comfyui.uses_runninghub_backend() {
|
||||||
|
Duration::from_secs(21600)
|
||||||
|
} else {
|
||||||
|
Duration::from_secs(300)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn native_workflow_status_ttl(state: &AppState) -> Duration {
|
||||||
|
if state.config.comfyui.uses_runninghub_backend() {
|
||||||
|
Duration::from_secs(120)
|
||||||
|
} else {
|
||||||
|
Duration::from_secs(30)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn native_workflow_base(state: &AppState) -> String {
|
||||||
|
state.config.comfyui.target_base_url()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn native_workflow_http_path(state: &AppState, path: &str) -> String {
|
||||||
|
if state.config.comfyui.uses_cloud_backend() {
|
||||||
|
format!(
|
||||||
|
"/api{}",
|
||||||
|
path.strip_prefix('/')
|
||||||
|
.map(|value| format!("/{value}"))
|
||||||
|
.unwrap_or_else(|| path.to_string())
|
||||||
|
)
|
||||||
|
} else {
|
||||||
|
path.to_string()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn native_workflow_request(
|
||||||
|
state: &AppState,
|
||||||
|
method: Method,
|
||||||
|
path: &str,
|
||||||
|
body: Option<Value>,
|
||||||
|
) -> reqwest::RequestBuilder {
|
||||||
|
let client = reqwest::Client::new();
|
||||||
|
let url = format!(
|
||||||
|
"{}{}",
|
||||||
|
native_workflow_base(state),
|
||||||
|
native_workflow_http_path(state, path)
|
||||||
|
);
|
||||||
|
let reqwest_method =
|
||||||
|
reqwest::Method::from_bytes(method.as_str().as_bytes()).unwrap_or(reqwest::Method::GET);
|
||||||
|
let mut request = client.request(reqwest_method, url);
|
||||||
|
if let Some(api_key) = state
|
||||||
|
.config
|
||||||
|
.comfyui
|
||||||
|
.api_key
|
||||||
|
.as_deref()
|
||||||
|
.map(str::trim)
|
||||||
|
.filter(|value| !value.is_empty())
|
||||||
|
{
|
||||||
|
request = request
|
||||||
|
.header("X-API-Key", api_key)
|
||||||
|
.header("Authorization", format!("Bearer {api_key}"));
|
||||||
|
}
|
||||||
|
if let Some(body) = body {
|
||||||
|
request = request.json(&body);
|
||||||
|
}
|
||||||
|
request
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) async fn native_workflow_json(
|
||||||
|
state: &AppState,
|
||||||
|
method: Method,
|
||||||
|
path: &str,
|
||||||
|
body: Option<Value>,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
let response = native_workflow_request(state, method, path, body)
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.map_err(|err| AppError::UpstreamError(format!("ComfyUI backend unreachable: {err}")))?;
|
||||||
|
let status = response.status();
|
||||||
|
let text = response
|
||||||
|
.text()
|
||||||
|
.await
|
||||||
|
.map_err(|err| AppError::UpstreamError(err.to_string()))?;
|
||||||
|
let value = serde_json::from_str::<Value>(&text).unwrap_or_else(|_| json!({ "body": text }));
|
||||||
|
if !status.is_success() {
|
||||||
|
return Err(AppError::UpstreamError(value.to_string()));
|
||||||
|
}
|
||||||
|
Ok(Json(value))
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) async fn native_workflow_json_cached(
|
||||||
|
state: &AppState,
|
||||||
|
method: Method,
|
||||||
|
path: &str,
|
||||||
|
body: Option<Value>,
|
||||||
|
ttl: Duration,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
if method != Method::GET || body.is_some() || ttl.is_zero() {
|
||||||
|
return native_workflow_json(state, method, path, body).await;
|
||||||
|
}
|
||||||
|
|
||||||
|
let cache_key = format!(
|
||||||
|
"{}|{}|{}",
|
||||||
|
native_workflow_base(state),
|
||||||
|
method.as_str(),
|
||||||
|
native_workflow_http_path(state, path)
|
||||||
|
);
|
||||||
|
let now = Instant::now();
|
||||||
|
if let Some(cached) = NATIVE_WORKFLOW_JSON_CACHE.lock().get(&cache_key).cloned() {
|
||||||
|
if now.duration_since(cached.inserted_at) <= ttl {
|
||||||
|
return Ok(Json(cached.value));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let value = match native_workflow_json(state, method, path, None).await {
|
||||||
|
Ok(value) => value.0,
|
||||||
|
Err(err) => {
|
||||||
|
if let Some(cached) = NATIVE_WORKFLOW_JSON_CACHE.lock().get(&cache_key).cloned() {
|
||||||
|
tracing::warn!(
|
||||||
|
path,
|
||||||
|
error = %err,
|
||||||
|
"using stale ComfyUI workflow metadata cache after refresh failure"
|
||||||
|
);
|
||||||
|
return Ok(Json(cached.value));
|
||||||
|
}
|
||||||
|
return Err(err);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let mut cache = NATIVE_WORKFLOW_JSON_CACHE.lock();
|
||||||
|
if cache.len() > 64 {
|
||||||
|
cache.retain(|_, item| {
|
||||||
|
now.duration_since(item.inserted_at) <= NATIVE_WORKFLOW_CACHE_MAX_AGE
|
||||||
|
});
|
||||||
|
}
|
||||||
|
cache.insert(
|
||||||
|
cache_key,
|
||||||
|
CachedJson {
|
||||||
|
inserted_at: now,
|
||||||
|
value: value.clone(),
|
||||||
|
},
|
||||||
|
);
|
||||||
|
Ok(Json(value))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn workflow_cache_kind_from_key(key: &str) -> &'static str {
|
||||||
|
if key.contains("/object_info") {
|
||||||
|
"object_info"
|
||||||
|
} else if key.contains("/experiment/models") {
|
||||||
|
"models"
|
||||||
|
} else if key.contains("/system_stats") {
|
||||||
|
"status"
|
||||||
|
} else {
|
||||||
|
"json"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) async fn native_workflow_response(
|
||||||
|
state: &AppState,
|
||||||
|
method: Method,
|
||||||
|
path: &str,
|
||||||
|
body: Option<Value>,
|
||||||
|
) -> AppResult<Response> {
|
||||||
|
let response = native_workflow_request(state, method, path, body)
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.map_err(|err| AppError::UpstreamError(format!("ComfyUI backend unreachable: {err}")))?;
|
||||||
|
let status = response.status();
|
||||||
|
let content_type = response
|
||||||
|
.headers()
|
||||||
|
.get(reqwest::header::CONTENT_TYPE)
|
||||||
|
.and_then(|value| value.to_str().ok())
|
||||||
|
.unwrap_or("application/json")
|
||||||
|
.to_string();
|
||||||
|
let bytes = response
|
||||||
|
.bytes()
|
||||||
|
.await
|
||||||
|
.map_err(|err| AppError::UpstreamError(err.to_string()))?;
|
||||||
|
Ok((
|
||||||
|
StatusCode::from_u16(status.as_u16()).unwrap_or(StatusCode::BAD_GATEWAY),
|
||||||
|
[(header::CONTENT_TYPE, content_type)],
|
||||||
|
bytes,
|
||||||
|
)
|
||||||
|
.into_response())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) async fn native_workflow_view_response(
|
||||||
|
state: &AppState,
|
||||||
|
caller: &WorkflowCaller,
|
||||||
|
params: &HashMap<String, String>,
|
||||||
|
) -> AppResult<Response> {
|
||||||
|
native_workflow_file_response(state, caller, "/view", params).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) async fn native_workflow_video_response(
|
||||||
|
state: &AppState,
|
||||||
|
caller: &WorkflowCaller,
|
||||||
|
params: &HashMap<String, String>,
|
||||||
|
) -> AppResult<Response> {
|
||||||
|
native_workflow_file_response(state, caller, "/viewvideo", params).await
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn native_workflow_file_response(
|
||||||
|
state: &AppState,
|
||||||
|
caller: &WorkflowCaller,
|
||||||
|
path: &str,
|
||||||
|
params: &HashMap<String, String>,
|
||||||
|
) -> AppResult<Response> {
|
||||||
|
let query = build_query_string(params);
|
||||||
|
let response = native_workflow_request(state, Method::GET, &format!("{path}{query}"), None)
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.map_err(|err| AppError::UpstreamError(format!("ComfyUI backend unreachable: {err}")))?;
|
||||||
|
let status = response.status();
|
||||||
|
let content_type = response
|
||||||
|
.headers()
|
||||||
|
.get(reqwest::header::CONTENT_TYPE)
|
||||||
|
.and_then(|value| value.to_str().ok())
|
||||||
|
.unwrap_or("application/octet-stream")
|
||||||
|
.to_string();
|
||||||
|
let bytes = response
|
||||||
|
.bytes()
|
||||||
|
.await
|
||||||
|
.map_err(|err| AppError::UpstreamError(err.to_string()))?;
|
||||||
|
if status.is_success() && matches!(path, "/view" | "/viewvideo") {
|
||||||
|
archive_native_workflow_asset(state, caller, path, params, &content_type, bytes.as_ref())
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
Ok((
|
||||||
|
StatusCode::from_u16(status.as_u16()).unwrap_or(StatusCode::BAD_GATEWAY),
|
||||||
|
[(header::CONTENT_TYPE, content_type)],
|
||||||
|
bytes,
|
||||||
|
)
|
||||||
|
.into_response())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn archive_native_workflow_asset(
|
||||||
|
state: &AppState,
|
||||||
|
caller: &WorkflowCaller,
|
||||||
|
path: &str,
|
||||||
|
params: &HashMap<String, String>,
|
||||||
|
content_type: &str,
|
||||||
|
bytes: &[u8],
|
||||||
|
) {
|
||||||
|
let file_type = params.get("type").map(String::as_str).unwrap_or("").trim();
|
||||||
|
if file_type != "output" {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let Some(kind) = workspace_kind_from_mime(content_type) else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
let filename = params
|
||||||
|
.get("filename")
|
||||||
|
.map(String::as_str)
|
||||||
|
.unwrap_or("")
|
||||||
|
.trim();
|
||||||
|
if filename.is_empty() || bytes.is_empty() {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let user_id = caller_user_id(caller);
|
||||||
|
let source = "workflow-studio";
|
||||||
|
let already_exists = sqlx::query_scalar::<_, i64>(
|
||||||
|
"SELECT COUNT(*)
|
||||||
|
FROM workspace_assets
|
||||||
|
WHERE deleted_at IS NULL
|
||||||
|
AND (
|
||||||
|
(? = 1 AND user_id IS NULL)
|
||||||
|
OR (? = 0 AND user_id = ?)
|
||||||
|
)
|
||||||
|
AND source = ?
|
||||||
|
AND json_extract(metadata, '$.workflowUi.filename') = ?
|
||||||
|
AND json_extract(metadata, '$.workflowUi.type') = ?
|
||||||
|
AND json_extract(metadata, '$.workflowUi.subfolder') = ?",
|
||||||
|
)
|
||||||
|
.bind(user_id.is_none())
|
||||||
|
.bind(user_id.is_none())
|
||||||
|
.bind(user_id)
|
||||||
|
.bind(source)
|
||||||
|
.bind(filename)
|
||||||
|
.bind(file_type)
|
||||||
|
.bind(params.get("subfolder").map(String::as_str).unwrap_or(""))
|
||||||
|
.fetch_one(&*state.db)
|
||||||
|
.await
|
||||||
|
.unwrap_or_default();
|
||||||
|
if already_exists > 0 {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
let input = WorkspaceAssetInput {
|
||||||
|
user_id,
|
||||||
|
source: source.to_string(),
|
||||||
|
kind: kind.to_string(),
|
||||||
|
title: Some(filename.chars().take(80).collect()),
|
||||||
|
prompt: None,
|
||||||
|
model: None,
|
||||||
|
remote_url: None,
|
||||||
|
metadata: json!({
|
||||||
|
"workspaceAsset": { "role": "output" },
|
||||||
|
"workflowUi": {
|
||||||
|
"filename": filename,
|
||||||
|
"type": file_type,
|
||||||
|
"subfolder": params.get("subfolder").map(String::as_str).unwrap_or(""),
|
||||||
|
"endpoint": path,
|
||||||
|
}
|
||||||
|
}),
|
||||||
|
};
|
||||||
|
if let Err(err) = workspace::record_bytes_asset(state, input, bytes, content_type).await {
|
||||||
|
tracing::warn!(%err, filename, "failed to archive workflow-ui asset into workspace");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) async fn native_workflow_api_proxy(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
Path(path): Path<String>,
|
||||||
|
method: Method,
|
||||||
|
uri: Uri,
|
||||||
|
body: Bytes,
|
||||||
|
) -> AppResult<Response> {
|
||||||
|
let _ = auth_caller_optional(&state, &headers);
|
||||||
|
if !native_workflow_enabled(&state) {
|
||||||
|
return Err(AppError::NotFound("Workflow API route not found".into()));
|
||||||
|
}
|
||||||
|
let target = match uri.query() {
|
||||||
|
Some(query) => format!("/{path}?{query}"),
|
||||||
|
None => format!("/{path}"),
|
||||||
|
};
|
||||||
|
if state.config.comfyui.uses_cloud_backend() {
|
||||||
|
return native_workflow_raw_response(&state, method, &target, &headers, body).await;
|
||||||
|
}
|
||||||
|
let body = if body.is_empty() {
|
||||||
|
None
|
||||||
|
} else {
|
||||||
|
Some(
|
||||||
|
serde_json::from_slice::<Value>(&body)
|
||||||
|
.unwrap_or_else(|_| json!({ "body": String::from_utf8_lossy(&body).to_string() })),
|
||||||
|
)
|
||||||
|
};
|
||||||
|
native_workflow_response(&state, method, &target, body).await
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn native_workflow_raw_response(
|
||||||
|
state: &AppState,
|
||||||
|
method: Method,
|
||||||
|
path: &str,
|
||||||
|
headers: &HeaderMap,
|
||||||
|
body: Bytes,
|
||||||
|
) -> AppResult<Response> {
|
||||||
|
let is_prompt_post = method == Method::POST && path.split('?').next() == Some("/prompt");
|
||||||
|
let mut request = native_workflow_request(state, method, path, None);
|
||||||
|
if let Some(content_type) = headers
|
||||||
|
.get(header::CONTENT_TYPE)
|
||||||
|
.and_then(|value| value.to_str().ok())
|
||||||
|
{
|
||||||
|
request = request.header(reqwest::header::CONTENT_TYPE, content_type);
|
||||||
|
}
|
||||||
|
if let Some(accept) = headers
|
||||||
|
.get(header::ACCEPT)
|
||||||
|
.and_then(|value| value.to_str().ok())
|
||||||
|
{
|
||||||
|
request = request.header(reqwest::header::ACCEPT, accept);
|
||||||
|
}
|
||||||
|
if !body.is_empty() {
|
||||||
|
if is_prompt_post {
|
||||||
|
if let Ok(mut value) = serde_json::from_slice::<Value>(&body) {
|
||||||
|
repair_native_prompt_body_value(state, &mut value).await;
|
||||||
|
request = request.json(&value);
|
||||||
|
} else {
|
||||||
|
request = request.body(body);
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
request = request.body(body);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let response = request
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.map_err(|err| AppError::UpstreamError(format!("ComfyUI backend unreachable: {err}")))?;
|
||||||
|
let status = response.status();
|
||||||
|
let response_headers = response.headers().clone();
|
||||||
|
let content_type = response_headers
|
||||||
|
.get(reqwest::header::CONTENT_TYPE)
|
||||||
|
.and_then(|value| value.to_str().ok())
|
||||||
|
.unwrap_or("application/json")
|
||||||
|
.to_string();
|
||||||
|
let bytes = response
|
||||||
|
.bytes()
|
||||||
|
.await
|
||||||
|
.map_err(|err| AppError::UpstreamError(err.to_string()))?;
|
||||||
|
let mut builder = Response::builder()
|
||||||
|
.status(StatusCode::from_u16(status.as_u16()).unwrap_or(StatusCode::BAD_GATEWAY))
|
||||||
|
.header(header::CONTENT_TYPE, content_type);
|
||||||
|
for name in [
|
||||||
|
header::CACHE_CONTROL,
|
||||||
|
header::ETAG,
|
||||||
|
header::LAST_MODIFIED,
|
||||||
|
header::LOCATION,
|
||||||
|
header::SET_COOKIE,
|
||||||
|
header::WWW_AUTHENTICATE,
|
||||||
|
] {
|
||||||
|
if let Some(value) = response_headers
|
||||||
|
.get(name.as_str())
|
||||||
|
.and_then(|value| value.to_str().ok())
|
||||||
|
.and_then(|value| HeaderValue::from_str(value).ok())
|
||||||
|
{
|
||||||
|
builder = builder.header(name, value);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
builder
|
||||||
|
.body(axum::body::Body::from(bytes))
|
||||||
|
.map_err(|err| AppError::Internal(anyhow::anyhow!(err)))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn repair_native_prompt_body_value(state: &AppState, value: &mut Value) {
|
||||||
|
if let Ok(mut payload) = serde_json::from_value::<PromptPayload>(value.clone()) {
|
||||||
|
repair_native_prompt_payload(state, &mut payload).await;
|
||||||
|
if let Ok(repaired) = serde_json::to_value(payload) {
|
||||||
|
*value = repaired;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) async fn upload_mask(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
body: Bytes,
|
||||||
|
) -> AppResult<Response> {
|
||||||
|
upload_image(
|
||||||
|
State(state),
|
||||||
|
headers,
|
||||||
|
Query(UploadQuery {
|
||||||
|
subfolder: None,
|
||||||
|
file_type: Some("input".into()),
|
||||||
|
}),
|
||||||
|
body,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user