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::{
|
||||
error::{AppError, AppResult},
|
||||
workflow::{
|
||||
@@ -8,7 +7,6 @@ use crate::{
|
||||
};
|
||||
use axum::{
|
||||
extract::{
|
||||
ws::{Message, WebSocketUpgrade},
|
||||
DefaultBodyLimit, Path, Query, State,
|
||||
},
|
||||
http::{header, HeaderMap, Method, StatusCode},
|
||||
@@ -23,15 +21,9 @@ use serde_json::{json, Value};
|
||||
use std::{
|
||||
collections::{HashMap, HashSet},
|
||||
path::{Path as StdPath, PathBuf},
|
||||
sync::{mpsc, Arc},
|
||||
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>>> =
|
||||
Lazy::new(|| parking_lot::Mutex::new(HashMap::new()));
|
||||
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,
|
||||
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> {
|
||||
Router::new()
|
||||
@@ -660,27 +656,6 @@ pub(super) fn simple_urlencode(value: &str) -> String {
|
||||
.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 {
|
||||
if node_type.trim().is_empty() {
|
||||
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(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
@@ -1158,206 +895,6 @@ async fn view_video(
|
||||
.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> {
|
||||
let bearer = match caller {
|
||||
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