refactor(workflow_ui): 抽出 ws 子模块
Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
boos
@@ -1,4 +1,3 @@
|
|||||||
use crate::workflow::executor::{WorkflowProgressCallback, WorkflowProgressEvent};
|
|
||||||
use crate::{
|
use crate::{
|
||||||
error::{AppError, AppResult},
|
error::{AppError, AppResult},
|
||||||
workflow::{
|
workflow::{
|
||||||
@@ -8,7 +7,6 @@ use crate::{
|
|||||||
};
|
};
|
||||||
use axum::{
|
use axum::{
|
||||||
extract::{
|
extract::{
|
||||||
ws::{Message, WebSocketUpgrade},
|
|
||||||
DefaultBodyLimit, Path, Query, State,
|
DefaultBodyLimit, Path, Query, State,
|
||||||
},
|
},
|
||||||
http::{header, HeaderMap, Method, StatusCode},
|
http::{header, HeaderMap, Method, StatusCode},
|
||||||
@@ -23,15 +21,9 @@ use serde_json::{json, Value};
|
|||||||
use std::{
|
use std::{
|
||||||
collections::{HashMap, HashSet},
|
collections::{HashMap, HashSet},
|
||||||
path::{Path as StdPath, PathBuf},
|
path::{Path as StdPath, PathBuf},
|
||||||
sync::{mpsc, Arc},
|
|
||||||
time::{Duration, Instant},
|
time::{Duration, Instant},
|
||||||
};
|
};
|
||||||
use tokio::sync::broadcast;
|
|
||||||
use tokio_tungstenite::{connect_async, tungstenite};
|
|
||||||
use tungstenite::client::IntoClientRequest;
|
|
||||||
|
|
||||||
static WORKFLOW_WS_SESSIONS: Lazy<parking_lot::Mutex<HashMap<String, broadcast::Sender<Value>>>> =
|
|
||||||
Lazy::new(|| parking_lot::Mutex::new(HashMap::new()));
|
|
||||||
pub(super) 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()));
|
||||||
pub(super) const NATIVE_WORKFLOW_CACHE_MAX_AGE: Duration = Duration::from_secs(21600);
|
pub(super) const NATIVE_WORKFLOW_CACHE_MAX_AGE: Duration = Duration::from_secs(21600);
|
||||||
@@ -166,6 +158,10 @@ use run_convert::{
|
|||||||
node_output_for_comfy, run_outputs_for_comfy, workflow_graph_to_comfy_workflow,
|
node_output_for_comfy, run_outputs_for_comfy, workflow_graph_to_comfy_workflow,
|
||||||
workflow_run_to_history, workflow_run_to_history_task, workflow_run_to_queue_task,
|
workflow_run_to_history, workflow_run_to_history_task, workflow_run_to_queue_task,
|
||||||
};
|
};
|
||||||
|
mod ws;
|
||||||
|
use ws::{
|
||||||
|
publish_run_events, workflow_progress_callback_with_run_sender, workflow_ws,
|
||||||
|
};
|
||||||
|
|
||||||
pub fn routes() -> Router<AppState> {
|
pub fn routes() -> Router<AppState> {
|
||||||
Router::new()
|
Router::new()
|
||||||
@@ -660,27 +656,6 @@ pub(super) fn simple_urlencode(value: &str) -> String {
|
|||||||
.replace('/', "%2F")
|
.replace('/', "%2F")
|
||||||
}
|
}
|
||||||
|
|
||||||
fn workflow_session_sender(session_id: &str) -> broadcast::Sender<Value> {
|
|
||||||
let mut sessions = WORKFLOW_WS_SESSIONS.lock();
|
|
||||||
if let Some(sender) = sessions.get(session_id) {
|
|
||||||
return sender.clone();
|
|
||||||
}
|
|
||||||
let (sender, _) = broadcast::channel(128);
|
|
||||||
sessions.insert(session_id.to_string(), sender.clone());
|
|
||||||
sender
|
|
||||||
}
|
|
||||||
|
|
||||||
fn publish_workflow_ws_event(session_id: Option<&str>, event_type: &str, data: Value) {
|
|
||||||
let Some(session_id) = session_id.map(str::trim).filter(|value| !value.is_empty()) else {
|
|
||||||
return;
|
|
||||||
};
|
|
||||||
let sender = workflow_session_sender(session_id);
|
|
||||||
let _ = sender.send(json!({
|
|
||||||
"type": event_type,
|
|
||||||
"data": data,
|
|
||||||
}));
|
|
||||||
}
|
|
||||||
|
|
||||||
pub(super) fn workflow_node_label(node_type: &str, node_id: &str) -> String {
|
pub(super) fn workflow_node_label(node_type: &str, node_id: &str) -> String {
|
||||||
if node_type.trim().is_empty() {
|
if node_type.trim().is_empty() {
|
||||||
node_id.to_string()
|
node_id.to_string()
|
||||||
@@ -689,244 +664,6 @@ pub(super) fn workflow_node_label(node_type: &str, node_id: &str) -> String {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(super) fn workflow_progress_callback_with_run_sender(
|
|
||||||
session_id: Option<String>,
|
|
||||||
run_sender: Option<mpsc::SyncSender<String>>,
|
|
||||||
) -> Option<WorkflowProgressCallback> {
|
|
||||||
let session_id = session_id.map(|value| value.trim().to_string())?;
|
|
||||||
if session_id.is_empty() {
|
|
||||||
return None;
|
|
||||||
}
|
|
||||||
let run_sender = Arc::new(parking_lot::Mutex::new(run_sender));
|
|
||||||
Some(Arc::new(move |event| match event {
|
|
||||||
WorkflowProgressEvent::Started { prompt_id, total } => {
|
|
||||||
if let Some(sender) = run_sender.lock().take() {
|
|
||||||
let _ = sender.send(prompt_id.clone());
|
|
||||||
}
|
|
||||||
publish_workflow_ws_event(
|
|
||||||
Some(&session_id),
|
|
||||||
"status",
|
|
||||||
json!({ "status": { "exec_info": { "queue_remaining": 1 } } }),
|
|
||||||
);
|
|
||||||
publish_workflow_ws_event(
|
|
||||||
Some(&session_id),
|
|
||||||
"execution_start",
|
|
||||||
json!({
|
|
||||||
"prompt_id": prompt_id,
|
|
||||||
"timestamp": timestamp_ms(),
|
|
||||||
}),
|
|
||||||
);
|
|
||||||
publish_workflow_ws_event(
|
|
||||||
Some(&session_id),
|
|
||||||
"progress",
|
|
||||||
json!({ "value": 0, "max": total.max(1) }),
|
|
||||||
);
|
|
||||||
}
|
|
||||||
WorkflowProgressEvent::NodeStarted {
|
|
||||||
prompt_id,
|
|
||||||
node_id,
|
|
||||||
node_type,
|
|
||||||
index,
|
|
||||||
total,
|
|
||||||
} => {
|
|
||||||
let label = workflow_node_label(&node_type, &node_id);
|
|
||||||
publish_workflow_ws_event(
|
|
||||||
Some(&session_id),
|
|
||||||
"executing",
|
|
||||||
json!({
|
|
||||||
"prompt_id": prompt_id,
|
|
||||||
"node": node_id,
|
|
||||||
"node_type": node_type,
|
|
||||||
"display_node": label,
|
|
||||||
"label": label,
|
|
||||||
}),
|
|
||||||
);
|
|
||||||
publish_workflow_ws_event(
|
|
||||||
Some(&session_id),
|
|
||||||
"progress",
|
|
||||||
json!({ "value": index, "max": total.max(1) }),
|
|
||||||
);
|
|
||||||
}
|
|
||||||
WorkflowProgressEvent::NodeCompleted {
|
|
||||||
prompt_id,
|
|
||||||
node_id,
|
|
||||||
node_type,
|
|
||||||
index,
|
|
||||||
total,
|
|
||||||
output,
|
|
||||||
} => {
|
|
||||||
let label = workflow_node_label(&node_type, &node_id);
|
|
||||||
let output = node_output_for_comfy(&output).unwrap_or(output);
|
|
||||||
publish_workflow_ws_event(
|
|
||||||
Some(&session_id),
|
|
||||||
"executed",
|
|
||||||
json!({
|
|
||||||
"prompt_id": prompt_id,
|
|
||||||
"node": node_id,
|
|
||||||
"node_type": node_type,
|
|
||||||
"display_node": label,
|
|
||||||
"label": label,
|
|
||||||
"output": output,
|
|
||||||
}),
|
|
||||||
);
|
|
||||||
publish_workflow_ws_event(
|
|
||||||
Some(&session_id),
|
|
||||||
"progress",
|
|
||||||
json!({ "value": index, "max": total.max(1) }),
|
|
||||||
);
|
|
||||||
}
|
|
||||||
WorkflowProgressEvent::NodeProgress {
|
|
||||||
prompt_id,
|
|
||||||
node_id,
|
|
||||||
node_type,
|
|
||||||
value,
|
|
||||||
max,
|
|
||||||
message,
|
|
||||||
} => {
|
|
||||||
let label = workflow_node_label(&node_type, &node_id);
|
|
||||||
publish_workflow_ws_event(
|
|
||||||
Some(&session_id),
|
|
||||||
"executing",
|
|
||||||
json!({
|
|
||||||
"prompt_id": prompt_id,
|
|
||||||
"node": node_id,
|
|
||||||
"node_type": node_type,
|
|
||||||
"display_node": label,
|
|
||||||
"label": label,
|
|
||||||
"message": message,
|
|
||||||
}),
|
|
||||||
);
|
|
||||||
publish_workflow_ws_event(
|
|
||||||
Some(&session_id),
|
|
||||||
"progress",
|
|
||||||
json!({ "value": value, "max": max.max(1), "message": message }),
|
|
||||||
);
|
|
||||||
}
|
|
||||||
WorkflowProgressEvent::Finished {
|
|
||||||
prompt_id,
|
|
||||||
success,
|
|
||||||
error,
|
|
||||||
} => {
|
|
||||||
publish_workflow_ws_event(
|
|
||||||
Some(&session_id),
|
|
||||||
"executing",
|
|
||||||
json!({
|
|
||||||
"prompt_id": prompt_id,
|
|
||||||
"node": Value::Null,
|
|
||||||
"display_node": Value::Null,
|
|
||||||
}),
|
|
||||||
);
|
|
||||||
publish_workflow_ws_event(
|
|
||||||
Some(&session_id),
|
|
||||||
"status",
|
|
||||||
json!({ "status": { "exec_info": { "queue_remaining": 0 } } }),
|
|
||||||
);
|
|
||||||
publish_workflow_ws_event(
|
|
||||||
Some(&session_id),
|
|
||||||
if success {
|
|
||||||
"execution_success"
|
|
||||||
} else {
|
|
||||||
"execution_error"
|
|
||||||
},
|
|
||||||
if success {
|
|
||||||
json!({ "prompt_id": prompt_id, "timestamp": timestamp_ms() })
|
|
||||||
} else {
|
|
||||||
json!({
|
|
||||||
"prompt_id": prompt_id,
|
|
||||||
"timestamp": timestamp_ms(),
|
|
||||||
"node_id": "workflow",
|
|
||||||
"node_type": "Workflow",
|
|
||||||
"exception_message": error.unwrap_or_else(|| "Workflow execution failed".to_string()),
|
|
||||||
"exception_type": "BoosAPIWorkflowError",
|
|
||||||
"traceback": [],
|
|
||||||
})
|
|
||||||
},
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}))
|
|
||||||
}
|
|
||||||
|
|
||||||
pub(super) fn publish_run_events(
|
|
||||||
session_id: Option<&str>,
|
|
||||||
run: &crate::workflow::models::WorkflowRunSummary,
|
|
||||||
outputs: &Value,
|
|
||||||
) {
|
|
||||||
let Some(session_id) = session_id.map(str::trim).filter(|value| !value.is_empty()) else {
|
|
||||||
return;
|
|
||||||
};
|
|
||||||
publish_workflow_ws_event(
|
|
||||||
Some(session_id),
|
|
||||||
"execution_start",
|
|
||||||
json!({
|
|
||||||
"prompt_id": run.public_id,
|
|
||||||
"timestamp": timestamp_ms(),
|
|
||||||
}),
|
|
||||||
);
|
|
||||||
|
|
||||||
for node in &run.node_runs {
|
|
||||||
let label = workflow_node_label(&node.node_type, &node.node_id);
|
|
||||||
publish_workflow_ws_event(
|
|
||||||
Some(session_id),
|
|
||||||
"executing",
|
|
||||||
json!({
|
|
||||||
"prompt_id": run.public_id,
|
|
||||||
"node": node.node_id,
|
|
||||||
"node_type": node.node_type,
|
|
||||||
"display_node": label,
|
|
||||||
"label": label,
|
|
||||||
}),
|
|
||||||
);
|
|
||||||
if let Some(output) = outputs.get(&node.node_id) {
|
|
||||||
publish_workflow_ws_event(
|
|
||||||
Some(session_id),
|
|
||||||
"executed",
|
|
||||||
json!({
|
|
||||||
"prompt_id": run.public_id,
|
|
||||||
"node": node.node_id,
|
|
||||||
"node_type": node.node_type,
|
|
||||||
"display_node": label,
|
|
||||||
"label": label,
|
|
||||||
"output": output,
|
|
||||||
}),
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
publish_workflow_ws_event(
|
|
||||||
Some(session_id),
|
|
||||||
"executing",
|
|
||||||
json!({
|
|
||||||
"prompt_id": run.public_id,
|
|
||||||
"node": Value::Null,
|
|
||||||
"display_node": Value::Null,
|
|
||||||
}),
|
|
||||||
);
|
|
||||||
publish_workflow_ws_event(
|
|
||||||
Some(session_id),
|
|
||||||
if run.status == "failed" {
|
|
||||||
"execution_error"
|
|
||||||
} else {
|
|
||||||
"execution_success"
|
|
||||||
},
|
|
||||||
if run.status == "failed" {
|
|
||||||
json!({
|
|
||||||
"prompt_id": run.public_id,
|
|
||||||
"timestamp": timestamp_ms(),
|
|
||||||
"node_id": run.node_runs.last().map(|node| node.node_id.as_str()).unwrap_or("workflow"),
|
|
||||||
"node_type": run.node_runs.last().map(|node| node.node_type.as_str()).unwrap_or("Workflow"),
|
|
||||||
"exception_message": run.error_message.clone().unwrap_or_else(|| "Workflow execution failed".to_string()),
|
|
||||||
"exception_type": "BoosAPIWorkflowError",
|
|
||||||
"traceback": [],
|
|
||||||
})
|
|
||||||
} else {
|
|
||||||
json!({
|
|
||||||
"prompt_id": run.public_id,
|
|
||||||
"timestamp": timestamp_ms(),
|
|
||||||
})
|
|
||||||
},
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn experiment_models(
|
async fn experiment_models(
|
||||||
State(state): State<AppState>,
|
State(state): State<AppState>,
|
||||||
headers: HeaderMap,
|
headers: HeaderMap,
|
||||||
@@ -1158,206 +895,6 @@ async fn view_video(
|
|||||||
.into_response())
|
.into_response())
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn workflow_ws(
|
|
||||||
State(state): State<AppState>,
|
|
||||||
headers: HeaderMap,
|
|
||||||
ws: WebSocketUpgrade,
|
|
||||||
Query(query): Query<HashMap<String, String>>,
|
|
||||||
) -> AppResult<Response> {
|
|
||||||
let _ = auth_caller(&state, &headers).await?;
|
|
||||||
let sid = query
|
|
||||||
.get("clientId")
|
|
||||||
.or_else(|| query.get("client_id"))
|
|
||||||
.cloned()
|
|
||||||
.unwrap_or_else(|| "boosapi-workflow".to_string());
|
|
||||||
if native_workflow_enabled(&state) && !state.config.comfyui.uses_runninghub_backend() {
|
|
||||||
return Ok(
|
|
||||||
ws.on_upgrade(move |socket| native_workflow_ws_bridge(state, socket, query, sid))
|
|
||||||
);
|
|
||||||
}
|
|
||||||
Ok(ws.on_upgrade(move |socket| async move {
|
|
||||||
use tokio::time::{interval, Duration};
|
|
||||||
let mut socket = socket;
|
|
||||||
let mut receiver = workflow_session_sender(&sid).subscribe();
|
|
||||||
let mut ticker = interval(Duration::from_secs(20));
|
|
||||||
let _ = socket
|
|
||||||
.send(Message::Text(
|
|
||||||
json!({
|
|
||||||
"type": "status",
|
|
||||||
"data": {
|
|
||||||
"status": { "exec_info": { "queue_remaining": 0 } },
|
|
||||||
"sid": sid
|
|
||||||
}
|
|
||||||
})
|
|
||||||
.to_string()
|
|
||||||
.into(),
|
|
||||||
))
|
|
||||||
.await;
|
|
||||||
loop {
|
|
||||||
tokio::select! {
|
|
||||||
event = receiver.recv() => {
|
|
||||||
match event {
|
|
||||||
Ok(event) => {
|
|
||||||
if socket.send(Message::Text(event.to_string().into())).await.is_err() {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
Err(broadcast::error::RecvError::Lagged(_)) => {}
|
|
||||||
Err(broadcast::error::RecvError::Closed) => break,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
_ = ticker.tick() => {
|
|
||||||
if socket.send(Message::Text(json!({"type":"status","data":{"status":{"exec_info":{"queue_remaining":0}},"sid":sid}}).to_string().into())).await.is_err() {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
next = socket.recv() => {
|
|
||||||
match next {
|
|
||||||
Some(Ok(Message::Close(_))) | None => break,
|
|
||||||
Some(Ok(_)) => {}
|
|
||||||
Some(Err(_)) => break,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}))
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn native_workflow_ws_bridge(
|
|
||||||
state: AppState,
|
|
||||||
mut client_socket: axum::extract::ws::WebSocket,
|
|
||||||
mut query: HashMap<String, String>,
|
|
||||||
sid: String,
|
|
||||||
) {
|
|
||||||
let mut ws_url = native_workflow_base(&state)
|
|
||||||
.replacen("https://", "wss://", 1)
|
|
||||||
.replacen("http://", "ws://", 1);
|
|
||||||
query.entry("clientId".to_string()).or_insert(sid.clone());
|
|
||||||
if let Some(api_key) = state
|
|
||||||
.config
|
|
||||||
.comfyui
|
|
||||||
.api_key
|
|
||||||
.as_deref()
|
|
||||||
.map(str::trim)
|
|
||||||
.filter(|value| !value.is_empty())
|
|
||||||
{
|
|
||||||
query
|
|
||||||
.entry("token".to_string())
|
|
||||||
.or_insert_with(|| api_key.to_string());
|
|
||||||
}
|
|
||||||
ws_url.push_str("/ws");
|
|
||||||
ws_url.push_str(&build_query_string(&query));
|
|
||||||
|
|
||||||
let request = match ws_url.as_str().into_client_request() {
|
|
||||||
Ok(mut request) => {
|
|
||||||
if let Some(api_key) = state
|
|
||||||
.config
|
|
||||||
.comfyui
|
|
||||||
.api_key
|
|
||||||
.as_deref()
|
|
||||||
.map(str::trim)
|
|
||||||
.filter(|value| !value.is_empty())
|
|
||||||
{
|
|
||||||
if let Ok(value) = tungstenite::http::HeaderValue::from_str(api_key) {
|
|
||||||
request.headers_mut().insert("X-API-Key", value);
|
|
||||||
}
|
|
||||||
if let Ok(value) =
|
|
||||||
tungstenite::http::HeaderValue::from_str(&format!("Bearer {api_key}"))
|
|
||||||
{
|
|
||||||
request
|
|
||||||
.headers_mut()
|
|
||||||
.insert(tungstenite::http::header::AUTHORIZATION, value);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
request
|
|
||||||
}
|
|
||||||
Err(err) => {
|
|
||||||
let _ = client_socket
|
|
||||||
.send(Message::Text(
|
|
||||||
json!({"type":"execution_error","data":{"exception_message":format!("invalid ComfyUI websocket URL: {err}")}})
|
|
||||||
.to_string()
|
|
||||||
.into(),
|
|
||||||
))
|
|
||||||
.await;
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
let upstream_target = ws_url
|
|
||||||
.split('?')
|
|
||||||
.next()
|
|
||||||
.unwrap_or(ws_url.as_str())
|
|
||||||
.to_string();
|
|
||||||
let (upstream_socket, _) = match connect_async(request).await {
|
|
||||||
Ok(value) => value,
|
|
||||||
Err(err) => {
|
|
||||||
tracing::warn!(target = %upstream_target, error = %err, "ComfyUI native websocket connect failed");
|
|
||||||
let _ = client_socket
|
|
||||||
.send(Message::Text(
|
|
||||||
json!({"type":"execution_error","data":{"exception_message":format!("ComfyUI websocket connect failed: {err}")}})
|
|
||||||
.to_string()
|
|
||||||
.into(),
|
|
||||||
))
|
|
||||||
.await;
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
use futures::{SinkExt, StreamExt};
|
|
||||||
let (mut upstream_write, mut upstream_read) = upstream_socket.split();
|
|
||||||
loop {
|
|
||||||
tokio::select! {
|
|
||||||
from_upstream = upstream_read.next() => {
|
|
||||||
match from_upstream {
|
|
||||||
Some(Ok(message)) => {
|
|
||||||
let outgoing = match message {
|
|
||||||
tungstenite::Message::Text(text) => Message::Text(text.to_string().into()),
|
|
||||||
tungstenite::Message::Binary(bytes) => Message::Binary(bytes.to_vec().into()),
|
|
||||||
tungstenite::Message::Ping(bytes) => Message::Ping(bytes.to_vec().into()),
|
|
||||||
tungstenite::Message::Pong(bytes) => Message::Pong(bytes.to_vec().into()),
|
|
||||||
tungstenite::Message::Close(frame) => {
|
|
||||||
let _ = client_socket.send(Message::Close(frame.map(|frame| axum::extract::ws::CloseFrame {
|
|
||||||
code: frame.code.into(),
|
|
||||||
reason: frame.reason.to_string().into(),
|
|
||||||
}))).await;
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
tungstenite::Message::Frame(_) => continue,
|
|
||||||
};
|
|
||||||
if client_socket.send(outgoing).await.is_err() {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
Some(Err(_)) | None => break,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
from_client = client_socket.recv() => {
|
|
||||||
match from_client {
|
|
||||||
Some(Ok(message)) => {
|
|
||||||
let outgoing = match message {
|
|
||||||
Message::Text(text) => tungstenite::Message::Text(text.to_string().into()),
|
|
||||||
Message::Binary(bytes) => tungstenite::Message::Binary(bytes.to_vec().into()),
|
|
||||||
Message::Ping(bytes) => tungstenite::Message::Ping(bytes.to_vec().into()),
|
|
||||||
Message::Pong(bytes) => tungstenite::Message::Pong(bytes.to_vec().into()),
|
|
||||||
Message::Close(frame) => {
|
|
||||||
let _ = upstream_write.send(tungstenite::Message::Close(frame.map(|frame| tungstenite::protocol::CloseFrame {
|
|
||||||
code: frame.code.into(),
|
|
||||||
reason: frame.reason.to_string().into(),
|
|
||||||
}))).await;
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
};
|
|
||||||
if upstream_write.send(outgoing).await.is_err() {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
Some(Err(_)) | None => break,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
pub(super) async fn workflow_context(state: &AppState, caller: &WorkflowCaller) -> AppResult<RunContext> {
|
pub(super) async fn workflow_context(state: &AppState, caller: &WorkflowCaller) -> AppResult<RunContext> {
|
||||||
let bearer = match caller {
|
let bearer = match caller {
|
||||||
WorkflowCaller::Admin => state.builtin_keys.comfyui_agent.1.clone(),
|
WorkflowCaller::Admin => state.builtin_keys.comfyui_agent.1.clone(),
|
||||||
|
|||||||
@@ -0,0 +1,494 @@
|
|||||||
|
use crate::workflow::executor::{WorkflowProgressCallback, WorkflowProgressEvent};
|
||||||
|
use crate::{error::AppResult, AppState};
|
||||||
|
use axum::{
|
||||||
|
extract::{
|
||||||
|
ws::{Message, WebSocketUpgrade},
|
||||||
|
Query, State,
|
||||||
|
},
|
||||||
|
http::HeaderMap,
|
||||||
|
response::Response,
|
||||||
|
};
|
||||||
|
use once_cell::sync::Lazy;
|
||||||
|
use serde_json::{json, Value};
|
||||||
|
use std::{
|
||||||
|
collections::HashMap,
|
||||||
|
sync::{mpsc, Arc},
|
||||||
|
};
|
||||||
|
use tokio::sync::broadcast;
|
||||||
|
use tokio_tungstenite::{connect_async, tungstenite};
|
||||||
|
use tungstenite::client::IntoClientRequest;
|
||||||
|
|
||||||
|
use super::workflow_node_label;
|
||||||
|
use super::{
|
||||||
|
build_query_string, native_workflow_base, native_workflow_enabled, node_output_for_comfy,
|
||||||
|
timestamp_ms,
|
||||||
|
};
|
||||||
|
use super::auth::auth_caller;
|
||||||
|
|
||||||
|
static WORKFLOW_WS_SESSIONS: Lazy<parking_lot::Mutex<HashMap<String, broadcast::Sender<Value>>>> =
|
||||||
|
Lazy::new(|| parking_lot::Mutex::new(HashMap::new()));
|
||||||
|
|
||||||
|
fn workflow_session_sender(session_id: &str) -> broadcast::Sender<Value> {
|
||||||
|
let mut sessions = WORKFLOW_WS_SESSIONS.lock();
|
||||||
|
if let Some(sender) = sessions.get(session_id) {
|
||||||
|
return sender.clone();
|
||||||
|
}
|
||||||
|
let (sender, _) = broadcast::channel(128);
|
||||||
|
sessions.insert(session_id.to_string(), sender.clone());
|
||||||
|
sender
|
||||||
|
}
|
||||||
|
|
||||||
|
fn publish_workflow_ws_event(session_id: Option<&str>, event_type: &str, data: Value) {
|
||||||
|
let Some(session_id) = session_id.map(str::trim).filter(|value| !value.is_empty()) else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
let sender = workflow_session_sender(session_id);
|
||||||
|
let _ = sender.send(json!({
|
||||||
|
"type": event_type,
|
||||||
|
"data": data,
|
||||||
|
}));
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn workflow_progress_callback_with_run_sender(
|
||||||
|
session_id: Option<String>,
|
||||||
|
run_sender: Option<mpsc::SyncSender<String>>,
|
||||||
|
) -> Option<WorkflowProgressCallback> {
|
||||||
|
let session_id = session_id.map(|value| value.trim().to_string())?;
|
||||||
|
if session_id.is_empty() {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
let run_sender = Arc::new(parking_lot::Mutex::new(run_sender));
|
||||||
|
Some(Arc::new(move |event| match event {
|
||||||
|
WorkflowProgressEvent::Started { prompt_id, total } => {
|
||||||
|
if let Some(sender) = run_sender.lock().take() {
|
||||||
|
let _ = sender.send(prompt_id.clone());
|
||||||
|
}
|
||||||
|
publish_workflow_ws_event(
|
||||||
|
Some(&session_id),
|
||||||
|
"status",
|
||||||
|
json!({ "status": { "exec_info": { "queue_remaining": 1 } } }),
|
||||||
|
);
|
||||||
|
publish_workflow_ws_event(
|
||||||
|
Some(&session_id),
|
||||||
|
"execution_start",
|
||||||
|
json!({
|
||||||
|
"prompt_id": prompt_id,
|
||||||
|
"timestamp": timestamp_ms(),
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
publish_workflow_ws_event(
|
||||||
|
Some(&session_id),
|
||||||
|
"progress",
|
||||||
|
json!({ "value": 0, "max": total.max(1) }),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
WorkflowProgressEvent::NodeStarted {
|
||||||
|
prompt_id,
|
||||||
|
node_id,
|
||||||
|
node_type,
|
||||||
|
index,
|
||||||
|
total,
|
||||||
|
} => {
|
||||||
|
let label = workflow_node_label(&node_type, &node_id);
|
||||||
|
publish_workflow_ws_event(
|
||||||
|
Some(&session_id),
|
||||||
|
"executing",
|
||||||
|
json!({
|
||||||
|
"prompt_id": prompt_id,
|
||||||
|
"node": node_id,
|
||||||
|
"node_type": node_type,
|
||||||
|
"display_node": label,
|
||||||
|
"label": label,
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
publish_workflow_ws_event(
|
||||||
|
Some(&session_id),
|
||||||
|
"progress",
|
||||||
|
json!({ "value": index, "max": total.max(1) }),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
WorkflowProgressEvent::NodeCompleted {
|
||||||
|
prompt_id,
|
||||||
|
node_id,
|
||||||
|
node_type,
|
||||||
|
index,
|
||||||
|
total,
|
||||||
|
output,
|
||||||
|
} => {
|
||||||
|
let label = workflow_node_label(&node_type, &node_id);
|
||||||
|
let output = node_output_for_comfy(&output).unwrap_or(output);
|
||||||
|
publish_workflow_ws_event(
|
||||||
|
Some(&session_id),
|
||||||
|
"executed",
|
||||||
|
json!({
|
||||||
|
"prompt_id": prompt_id,
|
||||||
|
"node": node_id,
|
||||||
|
"node_type": node_type,
|
||||||
|
"display_node": label,
|
||||||
|
"label": label,
|
||||||
|
"output": output,
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
publish_workflow_ws_event(
|
||||||
|
Some(&session_id),
|
||||||
|
"progress",
|
||||||
|
json!({ "value": index, "max": total.max(1) }),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
WorkflowProgressEvent::NodeProgress {
|
||||||
|
prompt_id,
|
||||||
|
node_id,
|
||||||
|
node_type,
|
||||||
|
value,
|
||||||
|
max,
|
||||||
|
message,
|
||||||
|
} => {
|
||||||
|
let label = workflow_node_label(&node_type, &node_id);
|
||||||
|
publish_workflow_ws_event(
|
||||||
|
Some(&session_id),
|
||||||
|
"executing",
|
||||||
|
json!({
|
||||||
|
"prompt_id": prompt_id,
|
||||||
|
"node": node_id,
|
||||||
|
"node_type": node_type,
|
||||||
|
"display_node": label,
|
||||||
|
"label": label,
|
||||||
|
"message": message,
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
publish_workflow_ws_event(
|
||||||
|
Some(&session_id),
|
||||||
|
"progress",
|
||||||
|
json!({ "value": value, "max": max.max(1), "message": message }),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
WorkflowProgressEvent::Finished {
|
||||||
|
prompt_id,
|
||||||
|
success,
|
||||||
|
error,
|
||||||
|
} => {
|
||||||
|
publish_workflow_ws_event(
|
||||||
|
Some(&session_id),
|
||||||
|
"executing",
|
||||||
|
json!({
|
||||||
|
"prompt_id": prompt_id,
|
||||||
|
"node": Value::Null,
|
||||||
|
"display_node": Value::Null,
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
publish_workflow_ws_event(
|
||||||
|
Some(&session_id),
|
||||||
|
"status",
|
||||||
|
json!({ "status": { "exec_info": { "queue_remaining": 0 } } }),
|
||||||
|
);
|
||||||
|
publish_workflow_ws_event(
|
||||||
|
Some(&session_id),
|
||||||
|
if success {
|
||||||
|
"execution_success"
|
||||||
|
} else {
|
||||||
|
"execution_error"
|
||||||
|
},
|
||||||
|
if success {
|
||||||
|
json!({ "prompt_id": prompt_id, "timestamp": timestamp_ms() })
|
||||||
|
} else {
|
||||||
|
json!({
|
||||||
|
"prompt_id": prompt_id,
|
||||||
|
"timestamp": timestamp_ms(),
|
||||||
|
"node_id": "workflow",
|
||||||
|
"node_type": "Workflow",
|
||||||
|
"exception_message": error.unwrap_or_else(|| "Workflow execution failed".to_string()),
|
||||||
|
"exception_type": "BoosAPIWorkflowError",
|
||||||
|
"traceback": [],
|
||||||
|
})
|
||||||
|
},
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
|
||||||
|
// APPEND_MARKER_PUBLISH_RUN_EVENTS
|
||||||
|
|
||||||
|
pub(super) fn publish_run_events(
|
||||||
|
session_id: Option<&str>,
|
||||||
|
run: &crate::workflow::models::WorkflowRunSummary,
|
||||||
|
outputs: &Value,
|
||||||
|
) {
|
||||||
|
let Some(session_id) = session_id.map(str::trim).filter(|value| !value.is_empty()) else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
publish_workflow_ws_event(
|
||||||
|
Some(session_id),
|
||||||
|
"execution_start",
|
||||||
|
json!({
|
||||||
|
"prompt_id": run.public_id,
|
||||||
|
"timestamp": timestamp_ms(),
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
|
||||||
|
for node in &run.node_runs {
|
||||||
|
let label = workflow_node_label(&node.node_type, &node.node_id);
|
||||||
|
publish_workflow_ws_event(
|
||||||
|
Some(session_id),
|
||||||
|
"executing",
|
||||||
|
json!({
|
||||||
|
"prompt_id": run.public_id,
|
||||||
|
"node": node.node_id,
|
||||||
|
"node_type": node.node_type,
|
||||||
|
"display_node": label,
|
||||||
|
"label": label,
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
if let Some(output) = outputs.get(&node.node_id) {
|
||||||
|
publish_workflow_ws_event(
|
||||||
|
Some(session_id),
|
||||||
|
"executed",
|
||||||
|
json!({
|
||||||
|
"prompt_id": run.public_id,
|
||||||
|
"node": node.node_id,
|
||||||
|
"node_type": node.node_type,
|
||||||
|
"display_node": label,
|
||||||
|
"label": label,
|
||||||
|
"output": output,
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
publish_workflow_ws_event(
|
||||||
|
Some(session_id),
|
||||||
|
"executing",
|
||||||
|
json!({
|
||||||
|
"prompt_id": run.public_id,
|
||||||
|
"node": Value::Null,
|
||||||
|
"display_node": Value::Null,
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
publish_workflow_ws_event(
|
||||||
|
Some(session_id),
|
||||||
|
if run.status == "failed" {
|
||||||
|
"execution_error"
|
||||||
|
} else {
|
||||||
|
"execution_success"
|
||||||
|
},
|
||||||
|
if run.status == "failed" {
|
||||||
|
json!({
|
||||||
|
"prompt_id": run.public_id,
|
||||||
|
"timestamp": timestamp_ms(),
|
||||||
|
"node_id": run.node_runs.last().map(|node| node.node_id.as_str()).unwrap_or("workflow"),
|
||||||
|
"node_type": run.node_runs.last().map(|node| node.node_type.as_str()).unwrap_or("Workflow"),
|
||||||
|
"exception_message": run.error_message.clone().unwrap_or_else(|| "Workflow execution failed".to_string()),
|
||||||
|
"exception_type": "BoosAPIWorkflowError",
|
||||||
|
"traceback": [],
|
||||||
|
})
|
||||||
|
} else {
|
||||||
|
json!({
|
||||||
|
"prompt_id": run.public_id,
|
||||||
|
"timestamp": timestamp_ms(),
|
||||||
|
})
|
||||||
|
},
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
// APPEND_MARKER_WS_HANDLERS
|
||||||
|
|
||||||
|
pub(super) async fn workflow_ws(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
ws: WebSocketUpgrade,
|
||||||
|
Query(query): Query<HashMap<String, String>>,
|
||||||
|
) -> AppResult<Response> {
|
||||||
|
let _ = auth_caller(&state, &headers).await?;
|
||||||
|
let sid = query
|
||||||
|
.get("clientId")
|
||||||
|
.or_else(|| query.get("client_id"))
|
||||||
|
.cloned()
|
||||||
|
.unwrap_or_else(|| "boosapi-workflow".to_string());
|
||||||
|
if native_workflow_enabled(&state) && !state.config.comfyui.uses_runninghub_backend() {
|
||||||
|
return Ok(
|
||||||
|
ws.on_upgrade(move |socket| native_workflow_ws_bridge(state, socket, query, sid))
|
||||||
|
);
|
||||||
|
}
|
||||||
|
Ok(ws.on_upgrade(move |socket| async move {
|
||||||
|
use tokio::time::{interval, Duration};
|
||||||
|
let mut socket = socket;
|
||||||
|
let mut receiver = workflow_session_sender(&sid).subscribe();
|
||||||
|
let mut ticker = interval(Duration::from_secs(20));
|
||||||
|
let _ = socket
|
||||||
|
.send(Message::Text(
|
||||||
|
json!({
|
||||||
|
"type": "status",
|
||||||
|
"data": {
|
||||||
|
"status": { "exec_info": { "queue_remaining": 0 } },
|
||||||
|
"sid": sid
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.to_string()
|
||||||
|
.into(),
|
||||||
|
))
|
||||||
|
.await;
|
||||||
|
loop {
|
||||||
|
tokio::select! {
|
||||||
|
event = receiver.recv() => {
|
||||||
|
match event {
|
||||||
|
Ok(event) => {
|
||||||
|
if socket.send(Message::Text(event.to_string().into())).await.is_err() {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(broadcast::error::RecvError::Lagged(_)) => {}
|
||||||
|
Err(broadcast::error::RecvError::Closed) => break,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ = ticker.tick() => {
|
||||||
|
if socket.send(Message::Text(json!({"type":"status","data":{"status":{"exec_info":{"queue_remaining":0}},"sid":sid}}).to_string().into())).await.is_err() {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
next = socket.recv() => {
|
||||||
|
match next {
|
||||||
|
Some(Ok(Message::Close(_))) | None => break,
|
||||||
|
Some(Ok(_)) => {}
|
||||||
|
Some(Err(_)) => break,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
|
||||||
|
// APPEND_MARKER_BRIDGE
|
||||||
|
|
||||||
|
async fn native_workflow_ws_bridge(
|
||||||
|
state: AppState,
|
||||||
|
mut client_socket: axum::extract::ws::WebSocket,
|
||||||
|
mut query: HashMap<String, String>,
|
||||||
|
sid: String,
|
||||||
|
) {
|
||||||
|
let mut ws_url = native_workflow_base(&state)
|
||||||
|
.replacen("https://", "wss://", 1)
|
||||||
|
.replacen("http://", "ws://", 1);
|
||||||
|
query.entry("clientId".to_string()).or_insert(sid.clone());
|
||||||
|
if let Some(api_key) = state
|
||||||
|
.config
|
||||||
|
.comfyui
|
||||||
|
.api_key
|
||||||
|
.as_deref()
|
||||||
|
.map(str::trim)
|
||||||
|
.filter(|value| !value.is_empty())
|
||||||
|
{
|
||||||
|
query
|
||||||
|
.entry("token".to_string())
|
||||||
|
.or_insert_with(|| api_key.to_string());
|
||||||
|
}
|
||||||
|
ws_url.push_str("/ws");
|
||||||
|
ws_url.push_str(&build_query_string(&query));
|
||||||
|
|
||||||
|
let request = match ws_url.as_str().into_client_request() {
|
||||||
|
Ok(mut request) => {
|
||||||
|
if let Some(api_key) = state
|
||||||
|
.config
|
||||||
|
.comfyui
|
||||||
|
.api_key
|
||||||
|
.as_deref()
|
||||||
|
.map(str::trim)
|
||||||
|
.filter(|value| !value.is_empty())
|
||||||
|
{
|
||||||
|
if let Ok(value) = tungstenite::http::HeaderValue::from_str(api_key) {
|
||||||
|
request.headers_mut().insert("X-API-Key", value);
|
||||||
|
}
|
||||||
|
if let Ok(value) =
|
||||||
|
tungstenite::http::HeaderValue::from_str(&format!("Bearer {api_key}"))
|
||||||
|
{
|
||||||
|
request
|
||||||
|
.headers_mut()
|
||||||
|
.insert(tungstenite::http::header::AUTHORIZATION, value);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
request
|
||||||
|
}
|
||||||
|
Err(err) => {
|
||||||
|
let _ = client_socket
|
||||||
|
.send(Message::Text(
|
||||||
|
json!({"type":"execution_error","data":{"exception_message":format!("invalid ComfyUI websocket URL: {err}")}})
|
||||||
|
.to_string()
|
||||||
|
.into(),
|
||||||
|
))
|
||||||
|
.await;
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let upstream_target = ws_url
|
||||||
|
.split('?')
|
||||||
|
.next()
|
||||||
|
.unwrap_or(ws_url.as_str())
|
||||||
|
.to_string();
|
||||||
|
let (upstream_socket, _) = match connect_async(request).await {
|
||||||
|
Ok(value) => value,
|
||||||
|
Err(err) => {
|
||||||
|
tracing::warn!(target = %upstream_target, error = %err, "ComfyUI native websocket connect failed");
|
||||||
|
let _ = client_socket
|
||||||
|
.send(Message::Text(
|
||||||
|
json!({"type":"execution_error","data":{"exception_message":format!("ComfyUI websocket connect failed: {err}")}})
|
||||||
|
.to_string()
|
||||||
|
.into(),
|
||||||
|
))
|
||||||
|
.await;
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
use futures::{SinkExt, StreamExt};
|
||||||
|
let (mut upstream_write, mut upstream_read) = upstream_socket.split();
|
||||||
|
loop {
|
||||||
|
tokio::select! {
|
||||||
|
from_upstream = upstream_read.next() => {
|
||||||
|
match from_upstream {
|
||||||
|
Some(Ok(message)) => {
|
||||||
|
let outgoing = match message {
|
||||||
|
tungstenite::Message::Text(text) => Message::Text(text.to_string().into()),
|
||||||
|
tungstenite::Message::Binary(bytes) => Message::Binary(bytes.to_vec().into()),
|
||||||
|
tungstenite::Message::Ping(bytes) => Message::Ping(bytes.to_vec().into()),
|
||||||
|
tungstenite::Message::Pong(bytes) => Message::Pong(bytes.to_vec().into()),
|
||||||
|
tungstenite::Message::Close(frame) => {
|
||||||
|
let _ = client_socket.send(Message::Close(frame.map(|frame| axum::extract::ws::CloseFrame {
|
||||||
|
code: frame.code.into(),
|
||||||
|
reason: frame.reason.to_string().into(),
|
||||||
|
}))).await;
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
tungstenite::Message::Frame(_) => continue,
|
||||||
|
};
|
||||||
|
if client_socket.send(outgoing).await.is_err() {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Some(Err(_)) | None => break,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
from_client = client_socket.recv() => {
|
||||||
|
match from_client {
|
||||||
|
Some(Ok(message)) => {
|
||||||
|
let outgoing = match message {
|
||||||
|
Message::Text(text) => tungstenite::Message::Text(text.to_string().into()),
|
||||||
|
Message::Binary(bytes) => tungstenite::Message::Binary(bytes.to_vec().into()),
|
||||||
|
Message::Ping(bytes) => tungstenite::Message::Ping(bytes.to_vec().into()),
|
||||||
|
Message::Pong(bytes) => tungstenite::Message::Pong(bytes.to_vec().into()),
|
||||||
|
Message::Close(frame) => {
|
||||||
|
let _ = upstream_write.send(tungstenite::Message::Close(frame.map(|frame| tungstenite::protocol::CloseFrame {
|
||||||
|
code: frame.code.into(),
|
||||||
|
reason: frame.reason.to_string().into(),
|
||||||
|
}))).await;
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
if upstream_write.send(outgoing).await.is_err() {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Some(Err(_)) | None => break,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user