refactor(workflow_ui): 抽出 native_proxy 子模块

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
boos
2026-06-18 22:22:19 +08:00
parent f43daed3ee
commit a37546cb2d
2 changed files with 665 additions and 637 deletions
+19 -637
View File
@@ -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,
+646
View File
@@ -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
}