feat: polish workflow studio and workspace assets

src/routes/comfyui/mod.rs

@@ -11,7 +11,7 @@ use crate::{
 };
 use axum::{
     body::Bytes,
-    extract::{Path, Query, State},
+    extract::{ws::Message, Path, Query, State, WebSocketUpgrade},
     http::{HeaderMap, Method, StatusCode, Uri},
     response::{Html, IntoResponse, Response},
     routing::{any, get, post},
@@ -20,7 +20,9 @@ use axum::{
 use reqwest::Client;
 use serde_json::Value;
 use std::collections::HashMap;
+use tokio_tungstenite::{connect_async, tungstenite};
 use tracing::info;
+use tungstenite::client::IntoClientRequest;
 
 /// Shared HTTP client
 fn http_client() -> Client {
@@ -85,11 +87,186 @@ pub fn routes() -> Router<AppState> {
         // Public ComfyUI reverse proxy for the embedded iframe UI.
         .route("/api/comfyui/proxy", any(proxy_static_root))
         .route("/api/comfyui/proxy/", any(proxy_static_root))
+        .route("/api/comfyui/proxy/ws", get(proxy_ws))
+        .route(
+            "/api/comfyui/proxy/api/experiment/models",
+            get(proxy_model_folders),
+        )
+        .route(
+            "/api/comfyui/proxy/api/experiment/models/{folder}",
+            get(proxy_models_in_folder),
+        )
+        .route(
+            "/api/comfyui/proxy/api/view_metadata/{folder}",
+            get(proxy_model_metadata),
+        )
         .route("/api/comfyui/proxy/{*path}", any(proxy_static))
         // HTML page (public)
         .route("/comfyui", get(comfyui_page))
 }
 
+async fn proxy_model_folders(State(state): State<AppState>) -> AppResult<Response> {
+    let mut folders = proxy_to_comfyui_json(&state, "GET", "/api/experiment/models", None).await?;
+    if let Some(items) = folders.as_array_mut() {
+        if !items
+            .iter()
+            .any(|item| item.get("name").and_then(Value::as_str) == Some("boosapi_models"))
+        {
+            items.insert(0, serde_json::json!({ "name": "boosapi_models" }));
+        }
+    }
+    Ok((
+        StatusCode::OK,
+        [(axum::http::header::CONTENT_TYPE, "application/json")],
+        folders.to_string(),
+    )
+        .into_response())
+}
+
+async fn proxy_models_in_folder(
+    State(state): State<AppState>,
+    Path(folder): Path<String>,
+) -> AppResult<Response> {
+    if folder == "boosapi_models" {
+        let models = list_boosapi_models(&state).await?;
+        let items: Vec<Value> = models
+            .into_iter()
+            .map(|name| serde_json::json!({ "name": name, "pathIndex": 0 }))
+            .collect();
+        return Ok((
+            StatusCode::OK,
+            [(axum::http::header::CONTENT_TYPE, "application/json")],
+            Value::Array(items).to_string(),
+        )
+            .into_response());
+    }
+    let path = format!("/api/experiment/models/{}", urlencoding(&folder));
+    proxy_to_comfyui_response(&state, "GET", &path, None).await
+}
+
+async fn proxy_model_metadata(
+    State(state): State<AppState>,
+    Path(folder): Path<String>,
+    uri: Uri,
+) -> AppResult<Response> {
+    if folder == "boosapi_models" {
+        let filename = uri
+            .query()
+            .and_then(|query| {
+                query.split('&').find_map(|part| {
+                    let (key, value) = part.split_once('=')?;
+                    (key == "filename").then_some(value.replace("%2F", "/").replace("%20", " "))
+                })
+            })
+            .unwrap_or_else(|| "BoosAPI model".to_string());
+        let body = serde_json::json!({
+            "title": filename,
+            "author": "BoosAPI",
+            "description": "BoosAPI 后台模型路由。双击可创建 BoosAPI 节点，实际调用本地 /v1 模型网关。",
+            "modelspec.tags": "boosapi,api,model",
+        });
+        return Ok((
+            StatusCode::OK,
+            [(axum::http::header::CONTENT_TYPE, "application/json")],
+            body.to_string(),
+        )
+            .into_response());
+    }
+    let path = match uri.query() {
+        Some(query) => format!("/api/view_metadata/{}?{}", urlencoding(&folder), query),
+        None => format!("/api/view_metadata/{}", urlencoding(&folder)),
+    };
+    proxy_to_comfyui_response(&state, "GET", &path, None).await
+}
+
+async fn proxy_ws(
+    State(state): State<AppState>,
+    ws: WebSocketUpgrade,
+    uri: Uri,
+) -> AppResult<Response> {
+    let mut target = comfyui_base(&state)
+        .replacen("https://", "wss://", 1)
+        .replacen("http://", "ws://", 1);
+    target.push_str("/ws");
+    if let Some(query) = uri.query() {
+        target.push('?');
+        target.push_str(query);
+    }
+    Ok(ws.on_upgrade(move |socket| proxy_ws_bridge(socket, target)))
+}
+
+async fn proxy_ws_bridge(mut client_socket: axum::extract::ws::WebSocket, target: String) {
+    let request = match target.as_str().into_client_request() {
+        Ok(request) => request,
+        Err(err) => {
+            let _ = client_socket
+                .send(Message::Text(format!("invalid ComfyUI websocket URL: {err}").into()))
+                .await;
+            return;
+        }
+    };
+    let Ok((upstream_socket, _)) = connect_async(request).await else {
+        let _ = client_socket
+            .send(Message::Text("ComfyUI websocket connect failed".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,
+                }
+            }
+        }
+    }
+}
+
 // ── Business route handlers (require auth) ───────────────────────────
 
 /// POST /api/comfyui/prompt
@@ -243,6 +420,7 @@ async fn proxy_static_request(
             .replace("comfyui", "BoosAPI");
         // Rewrite URLs
         let rewritten = rewrite_proxy_urls(&rebranded);
+        let rewritten = inject_boosapi_model_library_patch(&rewritten);
         Ok((
             StatusCode::from_u16(status.as_u16()).unwrap(),
             [(axum::http::header::CONTENT_TYPE, "text/html; charset=utf-8")],
@@ -274,6 +452,96 @@ async fn proxy_static_request(
     }
 }
 
+async fn proxy_to_comfyui_json(
+    state: &AppState,
+    method: &str,
+    path: &str,
+    body: Option<Value>,
+) -> AppResult<Value> {
+    let client = http_client();
+    let url = format!("{}{}", comfyui_base(state), path);
+    let req = match method {
+        "POST" => client.post(&url).json(&body.unwrap_or(Value::Null)),
+        _ => client.get(&url),
+    };
+    let resp = req
+        .send()
+        .await
+        .map_err(|e| AppError::UpstreamError(format!("ComfyUI unreachable: {e}")))?;
+    let status = resp.status();
+    let text = resp
+        .text()
+        .await
+        .map_err(|e| AppError::UpstreamError(e.to_string()))?;
+    if !status.is_success() {
+        return Err(AppError::UpstreamError(text));
+    }
+    serde_json::from_str(&text).map_err(|e| AppError::UpstreamError(e.to_string()))
+}
+
+async fn proxy_to_comfyui_response(
+    state: &AppState,
+    method: &str,
+    path: &str,
+    body: Option<Value>,
+) -> AppResult<Response> {
+    let client = http_client();
+    let url = format!("{}{}", comfyui_base(state), path);
+    let req = match method {
+        "POST" => client.post(&url).json(&body.unwrap_or(Value::Null)),
+        _ => client.get(&url),
+    };
+    let resp = req
+        .send()
+        .await
+        .map_err(|e| AppError::UpstreamError(format!("ComfyUI unreachable: {e}")))?;
+    let status = resp.status();
+    let content_type = resp
+        .headers()
+        .get("content-type")
+        .and_then(|v| v.to_str().ok())
+        .unwrap_or("application/json")
+        .to_string();
+    let bytes = resp
+        .bytes()
+        .await
+        .map_err(|e| AppError::UpstreamError(e.to_string()))?;
+    Ok((
+        StatusCode::from_u16(status.as_u16()).unwrap_or(StatusCode::BAD_GATEWAY),
+        [(axum::http::header::CONTENT_TYPE, content_type)],
+        bytes,
+    )
+        .into_response())
+}
+
+async fn list_boosapi_models(state: &AppState) -> AppResult<Vec<String>> {
+    let rows: Vec<String> = sqlx::query_scalar(
+        "SELECT DISTINCT COALESCE(NULLIF(TRIM(model_mapping), ''), model_pattern)
+         FROM token_routes
+         WHERE COALESCE(enabled, 1) = 1
+           AND TRIM(COALESCE(model_pattern, '')) <> ''
+         ORDER BY 1 COLLATE NOCASE",
+    )
+    .fetch_all(&*state.db)
+    .await?;
+
+    let mut seen = std::collections::HashSet::new();
+    let mut models: Vec<String> = rows
+        .into_iter()
+        .map(|value| value.trim().to_string())
+        .filter(|value| !value.is_empty())
+        .filter(|value| seen.insert(value.clone()))
+        .collect();
+    if models.is_empty() {
+        models.extend([
+            "gpt-5.5".to_string(),
+            "gpt-5.4".to_string(),
+            "gpt-image-2".to_string(),
+        ]);
+    }
+    Ok(models)
+}
+
 /// GET /comfyui — HTML page with iframe
 async fn comfyui_page() -> Html<&'static str> {
     Html(
@@ -335,6 +603,61 @@ fn rewrite_proxy_urls(html: &str) -> String {
     result
 }
 
+fn inject_boosapi_model_library_patch(html: &str) -> String {
+    let patch = r#"<script>
+(() => {
+  const isBoosapiModelText = (text) => {
+    const value = String(text || '').replace(/\s+/g, ' ').trim();
+    return value && !value.includes('boosapi_models') && !value.endsWith('_models');
+  };
+  const isImageModel = (name) => /image|dall|flux|sd|midjourney|stable|wan|seedream/i.test(name);
+  const addBoosapiNode = (model) => {
+    try {
+      const graph = window.app?.graph || window.LiteGraph?.LGraphCanvas?.active_canvas?.graph;
+      const canvas = window.app?.canvas || window.LiteGraph?.LGraphCanvas?.active_canvas;
+      if (!graph || !window.LiteGraph?.createNode) return false;
+      const nodeType = isImageModel(model) ? 'BoosAPITextToImage' : 'BoosAPILLM';
+      const node = window.LiteGraph.createNode(nodeType, nodeType);
+      if (!node) return false;
+      node.pos = canvas?.convertEventToCanvasOffset
+        ? canvas.convertEventToCanvasOffset({ clientX: window.innerWidth / 2, clientY: window.innerHeight / 2 })
+        : [120, 120];
+      graph.add(node);
+      const modelWidget = node.widgets?.find((item) => item.name === 'model');
+      if (modelWidget) modelWidget.value = model;
+      graph.setDirtyCanvas?.(true, true);
+      canvas?.selectNode?.(node);
+      return true;
+    } catch (err) {
+      console.warn('[BoosAPI] failed to add model node', err);
+      return false;
+    }
+  };
+  const patchModelClicks = () => {
+    document.querySelectorAll('#model-library .p-tree-node-content, #model-library [role="treeitem"], .model-lib-tree-explorer .p-tree-node-content').forEach((el) => {
+      if (el.__boosapiModelPatch) return;
+      el.__boosapiModelPatch = true;
+      el.addEventListener('dblclick', (event) => {
+        const text = String(el.textContent || '').replace(/\s+/g, ' ').trim();
+        if (!isBoosapiModelText(text)) return;
+        const hasBoosapiAncestor = Array.from(el.parentElement?.querySelectorAll?.('*') || [])
+          .some((child) => String(child.textContent || '').includes('boosapi_models'));
+        if (!hasBoosapiAncestor && !text.match(/gpt|claude|gemini|image|dall|flux|seed|deepseek|qwen|glm|doubao|wan/i)) return;
+        if (addBoosapiNode(text)) {
+          event.preventDefault();
+          event.stopPropagation();
+        }
+      }, true);
+    });
+  };
+  document.addEventListener('click', () => setTimeout(patchModelClicks, 100), true);
+  setInterval(patchModelClicks, 1200);
+  setTimeout(patchModelClicks, 500);
+})();
+</script>"#;
+    html.replace("</head>", &format!("{patch}\n</head>"))
+}
+
 /// Build query string from params
 fn build_query_string(params: &HashMap<String, String>) -> String {
     if params.is_empty() {