Nvidia/BossAI
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add workflow studio and postgres support

Browse Source
This commit is contained in:
boos
2026-05-25 16:03:07 +08:00
parent 7cdb02ef1d
commit cb35d5fc6e
96 changed files with 12414 additions and 15548 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/routes/admin/compat.rs
+93 -6
View File
@@ -100,6 +100,14 @@ struct IdPath {
id: i64,
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
struct DatabaseConnectionRequest {
dialect: Option<String>,
connection_string: Option<String>,
ssl: Option<bool>,
}
#[derive(Debug, Deserialize)]
struct RouteChannelPath {
route_id: i64,
@@ -599,23 +607,102 @@ async fn test_system_proxy() -> Json<Value> {
}
async fn database_runtime(State(state): State<AppState>) -> Json<Value> {
let dialect = crate::db::DbKind::from_url(&state.config.db.url)
.map(|kind| kind.as_str().to_string())
.unwrap_or_else(|_| "unknown".to_string());
Json(json!({
"dialect": "sqlite",
"connectionString": state.config.db.url,
"dialect": dialect,
"connectionString": mask_connection_string(&state.config.db.url),
"ssl": state.config.db.ssl,
"active": {
"dialect": dialect,
"connection": mask_connection_string(&state.config.db.url),
"ssl": state.config.db.ssl,
},
"saved": state.config.db.runtime_db_url.as_ref().map(|url| json!({
"dialect": state.config.db.runtime_db_dialect.as_deref().unwrap_or("postgres"),
"connection": mask_connection_string(url),
"ssl": state.config.db.runtime_db_ssl.unwrap_or(false),
})),
"restartRequired": false,
}))
}
async fn update_database_runtime() -> Json<Value> {
Json(json!({ "success": true, "message": "database config is managed by config.toml" }))
Json(json!({
"success": false,
"message": "数据库运行配置由 config.toml / BOOSAPI_DB__URL 管理；修改后需要重启服务",
"restartRequired": true,
}))
}
async fn test_database_connection() -> Json<Value> {
Json(json!({ "success": true }))
async fn test_database_connection(
Json(req): Json<DatabaseConnectionRequest>,
) -> AppResult<Json<Value>> {
let connection = req
.connection_string
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.ok_or_else(|| AppError::BadRequest("connectionString is required".into()))?;
let detected = crate::db::test_connection(connection, 2)
.await
.map_err(|err| AppError::BadRequest(err.to_string()))?;
if let Some(expected) = req
.dialect
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
{
if expected != detected.as_str() {
return Err(AppError::BadRequest(format!(
"dialect mismatch: selected {}, detected {}",
expected,
detected.as_str()
)));
}
}
Ok(Json(json!({
"success": true,
"dialect": detected.as_str(),
"connection": mask_connection_string(connection),
"ssl": req.ssl.unwrap_or(false),
})))
}
async fn database_migrate() -> Json<Value> {
Json(json!({ "success": true }))
Json(json!({
"success": false,
"message": "Web 迁移入口已禁用。请在服务器执行 boosapi migrate-sqlite-to-postgres --sqlite-url <sqlite-url> --postgres-url <postgres-url>，需要覆盖目标库时追加 --force-truncate。",
"cliOnly": true,
}))
}
fn mask_connection_string(raw: &str) -> String {
let trimmed = raw.trim();
if trimmed.starts_with("sqlite:") {
return trimmed.to_string();
}
let Some(scheme_end) = trimmed.find("://") else {
return trimmed.to_string();
};
let after_scheme = scheme_end + 3;
let Some(at_rel) = trimmed[after_scheme..].find('@') else {
return trimmed.to_string();
};
let at = after_scheme + at_rel;
let credentials = &trimmed[after_scheme..at];
let masked_credentials = if let Some(colon) = credentials.find(':') {
format!("{}:****", &credentials[..colon])
} else {
"****".to_string()
};
format!(
"{}{}{}",
&trimmed[..after_scheme],
masked_credentials,
&trimmed[at..]
)
}
async fn export_backup(
src/routes/admin/mod.rs
+2 -1
View File
@@ -19,7 +19,7 @@ mod accounts;
mod compat;
mod downstream_keys;
mod model_candidates;
mod models_marketplace;
pub(crate) mod models_marketplace;
mod sites;
mod tokens;
mod users;
@@ -63,6 +63,7 @@ pub fn routes(state: AppState) -> Router<AppState> {
"/api/sites",
get(sites::list_sites).post(sites::create_site),
)
.route("/api/sites/detect", post(sites::detect_site))
.route(
"/api/sites/{id}",
get(sites::get_site)
src/routes/admin/sites.rs
+113
View File
@@ -12,6 +12,11 @@ use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use tracing::info;
#[derive(Debug, Deserialize)]
pub struct DetectSiteRequest {
pub url: String,
}
#[derive(Debug, Deserialize)]
pub struct CreateSiteRequest {
pub name: String,
@@ -90,6 +95,18 @@ struct SiteApiEndpointRow {
last_failure_reason: Option<String>,
}
/// POST /api/sites/detect
pub async fn detect_site(Json(payload): Json<DetectSiteRequest>) -> AppResult<Json<Value>> {
let url = normalize_site_url(&payload.url)?;
let (platform, initialization_preset_id) = detect_site_platform(&url);
Ok(Json(json!({
"url": url,
"platform": platform,
"initializationPresetId": initialization_preset_id,
})))
}
/// GET /api/sites
pub async fn list_sites(State(state): State<AppState>) -> AppResult<Json<Value>> {
info!("Listing all sites");
@@ -435,6 +452,102 @@ pub async fn update_site_disabled_models(
Ok(Json(json!({ "models": normalised })))
}
fn normalize_site_url(raw: &str) -> AppResult<String> {
let trimmed = raw.trim().trim_end_matches('/').to_string();
if trimmed.is_empty() {
return Err(AppError::BadRequest("url is required".into()));
}
let lower = trimmed.to_ascii_lowercase();
if !(lower.starts_with("http://") || lower.starts_with("https://")) {
return Err(AppError::BadRequest(
"url must start with http:// or https://".into(),
));
}
Ok(trimmed)
}
fn detect_site_platform(url: &str) -> (&'static str, Option<&'static str>) {
let lower = url.to_ascii_lowercase();
if lower.starts_with("https://coding.dashscope.aliyuncs.com/apps/anthropic") {
return ("claude", Some("codingplan-claude"));
}
if lower.starts_with("https://coding.dashscope.aliyuncs.com/v1") {
return ("openai", Some("codingplan-openai"));
}
if lower.starts_with("https://open.bigmodel.cn/api/anthropic") {
return ("claude", Some("zhipu-coding-plan-claude"));
}
if lower.starts_with("https://open.bigmodel.cn/api/coding/paas/v4") {
return ("openai", Some("zhipu-coding-plan-openai"));
}
if lower.starts_with("https://api.deepseek.com/anthropic") {
return ("claude", Some("deepseek-claude"));
}
if lower.starts_with("https://api.deepseek.com/v1") || lower == "https://api.deepseek.com" {
return ("openai", Some("deepseek-openai"));
}
if lower.starts_with("https://api.moonshot.cn/anthropic") {
return ("claude", Some("moonshot-claude"));
}
if lower.starts_with("https://api.moonshot.cn/v1") {
return ("openai", Some("moonshot-openai"));
}
if lower.starts_with("https://api.minimaxi.com/anthropic") {
return ("claude", Some("minimax-claude"));
}
if lower.starts_with("https://api.minimaxi.com/v1") {
return ("openai", Some("minimax-openai"));
}
if lower.starts_with("https://api-inference.modelscope.cn/v1") {
return ("openai", Some("modelscope-openai"));
}
if lower.starts_with("https://api-inference.modelscope.cn") {
return ("claude", Some("modelscope-claude"));
}
if lower.starts_with("https://ark.cn-beijing.volces.com/api/coding/v3") {
return ("openai", Some("doubao-coding-openai"));
}
if lower.starts_with("https://ark.cn-beijing.volces.com/api/v3")
|| lower.starts_with("https://ark.cn-beijing.volces.com")
{
return ("openai", Some("volcengine-ark-openai"));
}
if lower.contains("anthropic") || lower.contains("claude") {
return ("claude", None);
}
if lower.contains("gemini") || lower.contains("googleapis.com") {
return ("gemini", None);
}
if lower.contains("new-api") {
return ("new-api", None);
}
if lower.contains("one-api") {
return ("one-api", None);
}
if lower.contains("anyrouter") {
return ("anyrouter", None);
}
if lower.contains("veloera") {
return ("veloera", None);
}
if lower.contains("one-hub") {
return ("one-hub", None);
}
if lower.contains("done-hub") {
return ("done-hub", None);
}
if lower.contains("sub2api") {
return ("sub2api", None);
}
if lower.contains("cliproxy") || lower.contains("cpa") {
return ("cliproxyapi", None);
}
("openai", None)
}
async fn replace_site_api_endpoints(
tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
site_id: i64,
src/routes/auth.rs
+40 -4
View File
@@ -8,13 +8,14 @@ use crate::{
AppState,
};
use axum::{
extract::{Path, State},
extract::{ConnectInfo, Path, State},
http::HeaderMap,
routing::{get, post},
Json, Router,
};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use std::net::SocketAddr;
use tracing::info;
// ── Request / Response types ─────────────────────────────────────────
@@ -75,6 +76,7 @@ pub fn routes() -> Router<AppState> {
/// POST /api/auth/register
async fn register(
State(state): State<AppState>,
ConnectInfo(addr): ConnectInfo<SocketAddr>,
Json(payload): Json<RegisterRequest>,
) -> AppResult<Json<Value>> {
info!("Registering user: {}", payload.email);
@@ -90,6 +92,23 @@ async fn register(
));
}
let client_ip = Some(addr.ip());
if let Err(limit) = state
.registration_limiter
.check_and_record_attempt(client_ip, &payload.email)
{
tracing::warn!(
email = %payload.email,
client_ip = ?client_ip,
retry_after_secs = limit.retry_after_secs,
"registration rate limited",
);
return Err(AppError::RateLimitedWithRetry {
message: limit.message.to_string(),
retry_after_secs: limit.retry_after_secs,
});
}
// Check if email already exists
let existing: Option<User> = sqlx::query_as::<_, User>("SELECT * FROM users WHERE email = ?")
.bind(&payload.email)
@@ -103,17 +122,33 @@ async fn register(
let password_hash = crate::auth::password::hash_password(&payload.password)
.map_err(|e| AppError::Internal(anyhow::anyhow!("Failed to hash password: {e}")))?;
let mut tx = state.db.begin().await?;
let result = sqlx::query(
"INSERT INTO users (username, email, password_hash, role) VALUES (?, ?, ?, 'user')",
)
.bind(&payload.username)
.bind(&payload.email)
.bind(&password_hash)
.execute(&*state.db)
.execute(&mut *tx)
.await?;
let user_id = result.last_insert_rowid();
let default_key = format!("sk-user-{}", uuid::Uuid::new_v4().simple());
let default_key_name = format!("{} 的默认密钥", payload.username.trim());
sqlx::query(
"INSERT INTO downstream_api_keys (name, key, user_id, enabled) VALUES (?, ?, ?, 1)",
)
.bind(&default_key_name)
.bind(&default_key)
.bind(user_id)
.execute(&mut *tx)
.await?;
tx.commit().await?;
state.registration_limiter.record_success(client_ip);
let user: User = sqlx::query_as::<_, User>("SELECT * FROM users WHERE id = ?")
.bind(result.last_insert_rowid())
.bind(user_id)
.fetch_one(&*state.db)
.await?;
@@ -144,7 +179,8 @@ async fn register(
"role": user.role,
"status": user.status,
"createdAt": user.created_at,
}
},
"defaultApiKey": default_key
}
})))
}
src/routes/comfy_proxy.rs
-457
View File
@@ -1,457 +0,0 @@
//! `/comfy/*` reverse proxy.
//!
//! Authenticates the caller via JWT/admin token (Authorization header, cookie,
//! or `?token=` query param where browsers can't set headers),
//! looks up the per-user container in [`ComfyUIPool`], and forwards
//! both HTTP and WS traffic to that container's loopback port.
//!
//! The proxy strips `/comfy` on requests so ComfyUI can keep serving from
//! its normal root while browsers see a scoped iframe URL.
use std::sync::Arc;
use std::time::{Duration, Instant};
use axum::{
body::Body,
extract::{
ws::{Message, WebSocket, WebSocketUpgrade},
OriginalUri, RawQuery, State,
},
http::{
header::{
CONNECTION, HOST, PROXY_AUTHENTICATE, PROXY_AUTHORIZATION, TE, TRAILER,
TRANSFER_ENCODING, UPGRADE,
},
HeaderMap, HeaderName, HeaderValue, Method, Request, Response, StatusCode, Uri,
},
response::IntoResponse,
routing::{any, get},
Router,
};
use futures::{SinkExt, StreamExt};
use tokio_tungstenite::tungstenite::Message as TMessage;
use tracing::{debug, warn};
use crate::services::comfyui_pool::{ComfyUIPool, ContainerState, PoolError};
use crate::AppState;
const COLD_START_HOLD: Duration = Duration::from_secs(30);
const COLD_START_POLL_INTERVAL: Duration = Duration::from_millis(500);
const COMFY_SESSION_COOKIE: &str = "boosapi_comfy_token";
const ADMIN_WORKSPACE_USER_ID: i64 = 0;
pub fn routes() -> Router<AppState> {
Router::new()
// WS upgrade endpoints — must be matched before the catch-all.
.route("/comfy/ws", get(ws_handler))
.route("/comfy/ws/", get(ws_handler))
// Catch-all for HTTP. `{*path}` is axum 0.8 wildcard syntax.
.route("/comfy", any(http_handler))
.route("/comfy/", any(http_handler))
.route("/comfy/{*path}", any(http_handler))
}
// ── Auth ───────────────────────────────────────────────────────────────
fn user_id_from_token(state: &AppState, token: &str) -> Option<i64> {
let token = token.trim();
if token.is_empty() {
return None;
}
if token == state.config.auth.auth_token || token == state.config.auth.proxy_token {
return Some(ADMIN_WORKSPACE_USER_ID);
}
let claims = crate::auth::jwt::validate_jwt(&state.config.auth.jwt_secret, token).ok()?;
claims.sub.parse::<i64>().ok()
}
/// Resolve the calling user from `Authorization: Bearer <token>`.
fn user_id_from_headers(state: &AppState, headers: &HeaderMap) -> Option<i64> {
let raw = headers
.get(axum::http::header::AUTHORIZATION)
.and_then(|v| v.to_str().ok())?;
let token = raw.strip_prefix("Bearer ").unwrap_or(raw).trim();
user_id_from_token(state, token)
}
fn user_id_from_cookie(state: &AppState, headers: &HeaderMap) -> Option<i64> {
let raw = headers
.get(axum::http::header::COOKIE)
.and_then(|v| v.to_str().ok())?;
for part in raw.split(';') {
let mut it = part.trim().splitn(2, '=');
let name = it.next()?.trim();
let value = it.next().unwrap_or("").trim();
if name == COMFY_SESSION_COOKIE {
let token = urlencoding::decode(value).ok()?;
return user_id_from_token(state, token.as_ref());
}
}
None
}
fn token_from_query(raw_query: Option<&str>) -> Option<String> {
let q = raw_query?;
for pair in q.split('&') {
let mut it = pair.splitn(2, '=');
let k = it.next()?;
let v = it.next().unwrap_or("");
if k == "token" {
return Some(urlencoding::decode(v).ok()?.into_owned());
}
}
None
}
fn query_without_token(raw_query: Option<&str>) -> Option<String> {
let query = raw_query?;
let kept: Vec<&str> = query
.split('&')
.filter(|pair| pair.splitn(2, '=').next() != Some("token"))
.filter(|pair| !pair.is_empty())
.collect();
if kept.is_empty() {
None
} else {
Some(kept.join("&"))
}
}
fn comfy_session_cookie(token: &str) -> String {
format!(
"{}={}; Path=/comfy; HttpOnly; SameSite=Lax",
COMFY_SESSION_COOKIE,
urlencoding::encode(token)
)
}
// ── HTTP forwarder ─────────────────────────────────────────────────────
async fn http_handler(State(state): State<AppState>, req: Request<Body>) -> Response<Body> {
let query_token = token_from_query(req.uri().query());
let user_id = user_id_from_headers(&state, req.headers())
.or_else(|| user_id_from_cookie(&state, req.headers()))
.or_else(|| {
query_token
.as_deref()
.and_then(|token| user_id_from_token(&state, token))
});
let Some(user_id) = user_id else {
return unauthorized();
};
let entry = match state.comfyui_pool.get_or_start(user_id).await {
Ok(e) => e,
Err(PoolError::AtCapacity(_)) => return at_capacity(),
Err(e) => return upstream_error(format!("pool: {}", e)),
};
if entry.state != ContainerState::Ready {
if let Err(e) = wait_for_ready(&state.comfyui_pool, user_id).await {
return upstream_error(e);
}
}
state.comfyui_pool.touch(user_id).await;
let path = req.uri().path();
let upstream_path = path.strip_prefix("/comfy").unwrap_or(path);
let upstream_path = if upstream_path.is_empty() {
"/"
} else {
upstream_path
};
let query = query_without_token(req.uri().query());
let port = match state.comfyui_pool.entry(user_id).await {
Some(e) => e.host_port,
None => return upstream_error("container vanished".into()),
};
let upstream = match build_upstream_uri("http", port, upstream_path, query.as_deref()) {
Ok(u) => u,
Err(e) => return upstream_error(e),
};
let mut response = forward_http(state.comfyui_pool.clone(), user_id, req, upstream).await;
if let Some(token) = query_token {
if let Ok(value) = HeaderValue::from_str(&comfy_session_cookie(&token)) {
response
.headers_mut()
.append(axum::http::header::SET_COOKIE, value);
}
}
response
}
fn build_upstream_uri(
scheme: &str,
port: u16,
path: &str,
query: Option<&str>,
) -> Result<Uri, String> {
let mut s = format!("{}://127.0.0.1:{}{}", scheme, port, path);
if let Some(q) = query {
s.push('?');
s.push_str(q);
}
Uri::try_from(s).map_err(|e| format!("bad upstream uri: {}", e))
}
async fn forward_http(
pool: Arc<ComfyUIPool>,
user_id: i64,
req: Request<Body>,
upstream: Uri,
) -> Response<Body> {
let (parts, body) = req.into_parts();
let client = match reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.timeout(Duration::from_secs(120))
.build()
{
Ok(c) => c,
Err(e) => return upstream_error(format!("client: {}", e)),
};
let body_bytes = match axum::body::to_bytes(body, 64 * 1024 * 1024).await {
Ok(b) => b,
Err(e) => return upstream_error(format!("read body: {}", e)),
};
let mut rb = client.request(reqwest_method(&parts.method), upstream.to_string());
for (name, value) in parts.headers.iter() {
if is_hop_by_hop(name) || is_browser_context_header(name) {
continue;
}
if let Ok(v) = value.to_str() {
rb = rb.header(name.as_str(), v);
}
}
rb = rb.body(body_bytes);
let resp = match rb.send().await {
Ok(r) => r,
Err(e) => return upstream_error(format!("upstream send: {}", e)),
};
pool.touch(user_id).await;
let status = resp.status();
let mut builder = Response::builder().status(status);
for (name, value) in resp.headers().iter() {
if is_hop_by_hop(name) {
continue;
}
if let Ok(v) = HeaderValue::from_bytes(value.as_bytes()) {
builder = builder.header(name.as_str(), v);
}
}
let bytes = match resp.bytes().await {
Ok(b) => b,
Err(e) => return upstream_error(format!("upstream read: {}", e)),
};
builder.body(Body::from(bytes)).unwrap_or_else(|_| {
Response::builder()
.status(StatusCode::BAD_GATEWAY)
.body(Body::empty())
.unwrap()
})
}
fn reqwest_method(m: &Method) -> reqwest::Method {
reqwest::Method::from_bytes(m.as_str().as_bytes()).unwrap_or(reqwest::Method::GET)
}
fn is_hop_by_hop(name: &HeaderName) -> bool {
name == &CONNECTION
|| name == &PROXY_AUTHENTICATE
|| name == &PROXY_AUTHORIZATION
|| name == &TE
|| name == &TRAILER
|| name == &TRANSFER_ENCODING
|| name == &UPGRADE
|| name == &HOST
|| name.as_str().eq_ignore_ascii_case("keep-alive")
}
fn is_browser_context_header(name: &HeaderName) -> bool {
let name = name.as_str();
name.eq_ignore_ascii_case("origin")
|| name.eq_ignore_ascii_case("referer")
|| name.eq_ignore_ascii_case("sec-fetch-site")
|| name.eq_ignore_ascii_case("sec-fetch-mode")
|| name.eq_ignore_ascii_case("sec-fetch-dest")
|| name.eq_ignore_ascii_case("sec-fetch-user")
}
// ── WebSocket forwarder ────────────────────────────────────────────────
async fn ws_handler(
State(state): State<AppState>,
OriginalUri(orig_uri): OriginalUri,
RawQuery(raw_query): RawQuery,
headers: HeaderMap,
ws: WebSocketUpgrade,
) -> Response<Body> {
// Resolve user from header (preferred) or `?token=` (browser fallback).
let user_id = user_id_from_headers(&state, &headers)
.or_else(|| user_id_from_cookie(&state, &headers))
.or_else(|| {
token_from_query(raw_query.as_deref()).and_then(|t| user_id_from_token(&state, &t))
});
let Some(user_id) = user_id else {
return unauthorized();
};
let entry = match state.comfyui_pool.get_or_start(user_id).await {
Ok(e) => e,
Err(PoolError::AtCapacity(_)) => return at_capacity(),
Err(e) => return upstream_error(format!("pool: {}", e)),
};
if entry.state != ContainerState::Ready {
if let Err(e) = wait_for_ready(&state.comfyui_pool, user_id).await {
return upstream_error(e);
}
}
let port = match state.comfyui_pool.entry(user_id).await {
Some(e) => e.host_port,
None => return upstream_error("container vanished".into()),
};
// Forward whatever query the browser sent (clientId etc).
let upstream = match build_upstream_uri("ws", port, "/ws", orig_uri.query()) {
Ok(u) => u,
Err(e) => return upstream_error(e),
};
let pool = state.comfyui_pool.clone();
ws.on_upgrade(move |socket| async move {
pool.ws_open(user_id).await;
if let Err(e) = pump_ws(socket, upstream.to_string(), pool.clone(), user_id).await {
warn!(user_id, error = %e, "ws proxy ended with error");
}
pool.ws_close(user_id).await;
})
.into_response()
}
async fn pump_ws(
client: WebSocket,
upstream_url: String,
pool: Arc<ComfyUIPool>,
user_id: i64,
) -> Result<(), String> {
let (upstream_ws, _) = tokio_tungstenite::connect_async(&upstream_url)
.await
.map_err(|e| format!("ws connect: {}", e))?;
let (mut up_tx, mut up_rx) = upstream_ws.split();
let (mut cl_tx, mut cl_rx) = client.split();
let pool_a = pool.clone();
let pool_b = pool;
let to_upstream = async move {
while let Some(frame) = cl_rx.next().await {
let frame = match frame {
Ok(f) => f,
Err(e) => return Err::<(), String>(format!("client recv: {}", e)),
};
pool_a.touch(user_id).await;
let out = match frame {
Message::Text(t) => TMessage::Text(t.to_string().into()),
Message::Binary(b) => TMessage::Binary(b.to_vec()),
Message::Ping(p) => TMessage::Ping(p.to_vec()),
Message::Pong(p) => TMessage::Pong(p.to_vec()),
Message::Close(_) => TMessage::Close(None),
};
if up_tx.send(out).await.is_err() {
break;
}
}
Ok(())
};
let to_client = async move {
while let Some(frame) = up_rx.next().await {
let frame = match frame {
Ok(f) => f,
Err(e) => return Err::<(), String>(format!("upstream recv: {}", e)),
};
pool_b.touch(user_id).await;
let out = match frame {
TMessage::Text(t) => Message::Text(t.to_string().into()),
TMessage::Binary(b) => Message::Binary(b.into()),
TMessage::Ping(p) => Message::Ping(p.into()),
TMessage::Pong(p) => Message::Pong(p.into()),
TMessage::Close(_) => Message::Close(None),
TMessage::Frame(_) => continue,
};
if cl_tx.send(out).await.is_err() {
break;
}
}
Ok(())
};
tokio::select! {
r = to_upstream => r,
r = to_client => r,
}
}
// ── Cold-start hold ────────────────────────────────────────────────────
async fn wait_for_ready(pool: &ComfyUIPool, user_id: i64) -> Result<(), String> {
let deadline = Instant::now() + COLD_START_HOLD;
loop {
match pool.state(user_id).await {
ContainerState::Ready => return Ok(()),
ContainerState::Failed | ContainerState::Stopped => {
return Err("container failed to start".into());
}
ContainerState::Starting => {
if Instant::now() >= deadline {
return Err("cold start timeout".into());
}
tokio::time::sleep(COLD_START_POLL_INTERVAL).await;
}
}
}
}
// ── Error responses ────────────────────────────────────────────────────
fn unauthorized() -> Response<Body> {
Response::builder()
.status(StatusCode::UNAUTHORIZED)
.header("WWW-Authenticate", "Bearer")
.body(Body::from(r#"{"error":"unauthorized"}"#))
.unwrap()
}
fn at_capacity() -> Response<Body> {
Response::builder()
.status(StatusCode::SERVICE_UNAVAILABLE)
.header("Retry-After", "60")
.header("Content-Type", "application/json")
.body(Body::from(
r#"{"error":"comfyui pool at capacity","retry_after":60}"#,
))
.unwrap()
}
fn upstream_error(msg: String) -> Response<Body> {
debug!(%msg, "comfy proxy upstream error");
Response::builder()
.status(StatusCode::BAD_GATEWAY)
.header("Content-Type", "application/json")
.body(Body::from(format!(
r#"{{"error":"comfy proxy: {}"}}"#,
msg.replace('"', "\\\"")
)))
.unwrap()
}
src/routes/comfyui_agent.rs
+152 -8
View File
@@ -63,6 +63,7 @@ pub fn routes() -> Router<AppState> {
Router::new()
.route("/api/comfyui-agent/status", get(status_handler))
.route("/api/comfyui-agent/models", get(models_handler))
.route("/api/comfyui-agent/image-models", get(image_models_handler))
.route("/api/comfyui-agent/chat", post(chat_handler))
}
@@ -70,6 +71,8 @@ pub fn routes() -> Router<AppState> {
struct AgentChatRequest {
messages: Vec<AgentMessage>,
model: Option<String>,
#[serde(rename = "imageModel")]
image_model: Option<String>,
}
#[derive(Debug, Deserialize)]
@@ -85,6 +88,10 @@ struct AgentMessage {
struct AgentAttachment {
name: String,
mime_type: Option<String>,
#[serde(default, rename = "mediaKind")]
media_kind: Option<String>,
#[serde(default, rename = "referenceLabel")]
reference_label: Option<String>,
text: Option<String>,
data_url: Option<String>,
}
@@ -290,6 +297,52 @@ fn is_exact_model_pattern(value: &str) -> bool {
&& !normalized.contains('?')
}
fn is_image_model_pattern(value: &str) -> bool {
if !is_exact_model_pattern(value) {
return false;
}
let normalized = value.trim().to_ascii_lowercase();
let image_keywords = [
"gpt-image",
"dall-e",
"dalle",
"midjourney",
"recraft",
"kolors",
"cogview",
"seedream",
"sdxl",
"stable-diffusion",
"stablediffusion",
"pixart",
"playground",
"firefly",
"imagen",
"flux",
"image",
"img",
"canvas",
"painting",
"draw",
];
if image_keywords
.iter()
.any(|keyword| normalized.contains(keyword))
{
return true;
}
let chat_keywords = [
"claude", "chatgpt", "gpt-4", "gpt-3", "gpt-5", "o1", "o3", "o4", "gemini", "deepseek",
"qwen", "llama", "mistral", "command", "kimi", "doubao", "glm", "hunyuan", "spark",
"baichuan", "ernie", "wenxin", "grok",
];
!chat_keywords
.iter()
.any(|keyword| normalized.contains(keyword))
}
async fn status_handler(
State(state): State<AppState>,
headers: HeaderMap,
@@ -353,6 +406,47 @@ async fn models_handler(
})))
}
async fn image_models_handler(
State(state): State<AppState>,
headers: HeaderMap,
) -> Result<Json<Value>, (axum::http::StatusCode, Json<Value>)> {
let _auth = resolve_agent_auth(&headers, &state).await?;
let rows: Vec<String> = sqlx::query_scalar(
"SELECT model_pattern
FROM token_routes
WHERE COALESCE(enabled, 1) = 1
ORDER BY model_pattern COLLATE NOCASE",
)
.fetch_all(&*state.db)
.await
.map_err(|err| {
(
axum::http::StatusCode::INTERNAL_SERVER_ERROR,
Json(json!({ "error": err.to_string() })),
)
})?;
let mut seen = std::collections::HashSet::new();
let models: Vec<String> = rows
.into_iter()
.map(|value| value.trim().to_string())
.filter(|value| is_image_model_pattern(value))
.filter(|value| seen.insert(value.clone()))
.collect();
Ok(Json(json!({
"success": true,
"models": models,
"defaultModel": models
.iter()
.find(|model| model.to_ascii_lowercase().contains("image"))
.or_else(|| models.first())
.cloned()
.unwrap_or_else(|| "gpt-image-2".to_string()),
})))
}
#[derive(Debug, Default)]
struct ParsedSections {
characters: Vec<CharacterSection>,
@@ -501,7 +595,12 @@ fn parse_sections(full_text: &str) -> ParsedSections {
sections
}
async fn call_image_gen(api_base: &str, api_key: &str, prompt: &str) -> Result<String, String> {
async fn call_image_gen(
api_base: &str,
api_key: &str,
image_model: &str,
prompt: &str,
) -> Result<String, String> {
let client = Client::new();
let url = format!("{}/images/generations", api_base.trim_end_matches('/'));
@@ -509,7 +608,7 @@ async fn call_image_gen(api_base: &str, api_key: &str, prompt: &str) -> Result<S
.post(&url)
.bearer_auth(api_key)
.json(&json!({
"model": "gpt-image-2",
"model": image_model,
"prompt": prompt,
"n": 1,
"size": "1024x1024"
@@ -587,6 +686,19 @@ fn attachment_summary(attachment: &AgentAttachment) -> String {
.unwrap_or("application/octet-stream");
let name = attachment.name.trim();
let name = if name.is_empty() { "untitled" } else { name };
let reference = attachment
.reference_label
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("@素材");
let media_kind = attachment
.media_kind
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("file");
let header = format!("\n\n[素材引用: {reference}, 名称: {name}, 类型: {mime}, 分类: {media_kind}]");
if let Some(text) = attachment
.text
@@ -594,9 +706,9 @@ fn attachment_summary(attachment: &AgentAttachment) -> String {
.map(str::trim)
.filter(|value| !value.is_empty())
{
format!("\n\n[附件: {name}, 类型: {mime}]\n{text}")
format!("{header}\n{text}")
} else {
format!("\n\n[附件: {name}, 类型: {mime}]")
header
}
}
@@ -684,6 +796,12 @@ async fn chat_handler(
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
let requested_image_model = body
.image_model
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
let messages: Vec<Value> = body
.messages
@@ -699,6 +817,7 @@ async fn chat_handler(
"comfyui-agent: resolved upstream",
);
let agent_model = requested_model.unwrap_or_else(|| state.config.comfyui.agent_model.clone());
let image_model = requested_image_model.unwrap_or_else(|| "gpt-image-2".to_string());
let (tx, rx) = unbounded_channel::<Result<Event, Infallible>>();
tokio::spawn(async move {
@@ -813,12 +932,35 @@ async fn chat_handler(
continue;
};
let content = parsed
let delta = parsed
.get("choices")
.and_then(Value::as_array)
.and_then(|choices| choices.first())
.and_then(|choice| choice.get("delta"))
.and_then(|delta| delta.get("content"))
.cloned()
.unwrap_or(Value::Null);
let reasoning = delta
.get("reasoning_content")
.or_else(|| delta.get("reasoning"))
.and_then(Value::as_str)
.unwrap_or_default()
.to_string();
if !reasoning.is_empty() {
if !send_sse(
&tx,
json!({
"type": "reasoning",
"content": reasoning,
}),
) {
return;
}
}
let content = delta
.get("content")
.and_then(Value::as_str)
.unwrap_or_default()
.to_string();
@@ -893,7 +1035,7 @@ async fn chat_handler(
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
}
match call_image_gen(&api_base, &api_key, prompt).await {
match call_image_gen(&api_base, &api_key, &image_model, prompt).await {
Ok(url) if !url.is_empty() => {
send_sse(
&tx,
@@ -927,7 +1069,7 @@ async fn chat_handler(
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
match call_image_gen(&api_base, &api_key, &scene.prompt).await {
match call_image_gen(&api_base, &api_key, &image_model, &scene.prompt).await {
Ok(url) if !url.is_empty() => {
send_sse(
&tx,
@@ -1191,8 +1333,10 @@ mod tests {
messages: vec![AgentMessage {
role: "user".to_string(),
content: "hello".to_string(),
attachments: Vec::new(),
}],
model: None,
image_model: None,
}),
)
.await;
src/routes/digital_human.rs
+160
View File
@@ -0,0 +1,160 @@
use crate::{
auth::jwt::validate_jwt,
db::models::downstream_api_key::DownstreamApiKey,
error::{AppError, AppResult},
AppState,
};
use axum::{
extract::State,
http::{header, HeaderMap},
routing::{get, post},
Json, Router,
};
use serde_json::{json, Value};
#[derive(Debug, Clone)]
struct DigitalHumanAuth {
outbound_key: String,
}
pub fn routes() -> Router<AppState> {
Router::new()
.route("/api/digital-human/status", get(status))
.route("/api/digital-human/generate", post(generate))
}
fn local_v1_base(state: &AppState) -> String {
let host = if state.config.server.host == "0.0.0.0" || state.config.server.host == "::" {
"127.0.0.1".to_string()
} else {
state.config.server.host.clone()
};
format!("http://{}:{}/v1", host, state.config.server.port)
}
fn extract_bearer_token(headers: &HeaderMap) -> Option<&str> {
headers
.get(header::AUTHORIZATION)
.and_then(|value| value.to_str().ok())
.and_then(|value| {
value
.strip_prefix("Bearer ")
.or_else(|| value.strip_prefix("bearer "))
})
.map(str::trim)
.filter(|token| !token.is_empty())
}
async fn user_downstream_key(state: &AppState, user_id: i64) -> Result<DownstreamApiKey, String> {
sqlx::query_as::<_, DownstreamApiKey>(
"SELECT *
FROM downstream_api_keys
WHERE user_id = ?
AND COALESCE(enabled, 1) = 1
AND (
expires_at IS NULL
OR expires_at = ''
OR datetime(expires_at) > datetime('now')
)
ORDER BY updated_at DESC, created_at DESC, id DESC
LIMIT 1",
)
.bind(user_id)
.fetch_optional(&*state.db)
.await
.map_err(|err| err.to_string())?
.ok_or_else(|| "当前用户没有可用的 API Key，请先在用户中心创建或启用 API Key".to_string())
}
async fn resolve_auth(headers: &HeaderMap, state: &AppState) -> Result<DigitalHumanAuth, AppError> {
let token = extract_bearer_token(headers)
.ok_or_else(|| AppError::Unauthorized("Missing Authorization header".into()))?;
if token == state.config.auth.auth_token || token == state.config.auth.proxy_token {
return Ok(DigitalHumanAuth {
outbound_key: state.config.auth.proxy_token.clone(),
});
}
let claims = validate_jwt(&state.config.auth.jwt_secret, token)
.map_err(|_| AppError::Unauthorized("Invalid or expired token".into()))?;
if claims.role != "admin" && claims.role != "user" {
return Err(AppError::Forbidden("Forbidden".into()));
}
let user_id = claims
.sub
.parse::<i64>()
.map_err(|_| AppError::Unauthorized("Invalid user ID in token".into()))?;
match user_downstream_key(state, user_id).await {
Ok(key) => Ok(DigitalHumanAuth {
outbound_key: key.key,
}),
Err(_) if claims.role == "admin" => Ok(DigitalHumanAuth {
outbound_key: state.config.auth.proxy_token.clone(),
}),
Err(message) => Err(AppError::Forbidden(message)),
}
}
async fn status(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
let _ = resolve_auth(&headers, &state).await?;
Ok(Json(json!({
"success": true,
"mode": "boosapi-video-route",
"capabilities": [
"text_to_avatar_video",
"image_to_avatar_video",
"audio_driven_avatar",
"multi_person_prompt",
"segmented_long_video"
]
})))
}
async fn generate(
State(state): State<AppState>,
headers: HeaderMap,
Json(mut body): Json<Value>,
) -> AppResult<Json<Value>> {
let auth = resolve_auth(&headers, &state).await?;
if body.get("model").and_then(Value::as_str).unwrap_or("").trim().is_empty() {
return Err(AppError::BadRequest("model is required".into()));
}
if body
.get("prompt")
.and_then(Value::as_str)
.unwrap_or("")
.trim()
.is_empty()
{
return Err(AppError::BadRequest("prompt is required".into()));
}
if let Some(object) = body.as_object_mut() {
object.insert("task_type".to_string(), json!("digital_human"));
}
let response = reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(300))
.build()
.unwrap_or_else(|_| reqwest::Client::new())
.post(format!("{}/videos", local_v1_base(&state)))
.bearer_auth(auth.outbound_key)
.json(&body)
.send()
.await
.map_err(|err| AppError::UpstreamError(err.to_string()))?;
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(json!({
"success": true,
"request": body,
"response": value
})))
}
src/routes/mod.rs
+138 -66
View File
@@ -1,22 +1,25 @@
use crate::AppState;
use axum::{
body::Body,
http::{Request, StatusCode},
response::Response,
routing::get,
Router,
http::{header, HeaderValue, Request, StatusCode},
response::{IntoResponse, Response},
routing::{any, get},
Json, Router,
};
use include_dir::{include_dir, Dir};
use serde_json::{json, Value};
use std::convert::Infallible;
use tower::service_fn;
use tower_http::services::ServeDir;
static EMBEDDED_WEB: Dir<'_> = include_dir!("$CARGO_MANIFEST_DIR/dist/web");
mod admin;
mod auth;
mod checkin;
#[cfg(feature = "multi_tenant_comfy")]
mod comfy_proxy;
mod comfyui;
mod comfyui_agent;
mod digital_human;
mod events;
mod health;
mod monitor;
@@ -31,19 +34,12 @@ mod tasks;
mod update_center;
mod user;
mod user_api_keys;
mod workspace;
// Two implementations of /api/video-agent/*: legacy singleton (default)
// vs per-user pool. Exactly one is mounted, picked at compile time by
// the `multi_tenant_comfy` feature flag.
#[cfg(feature = "multi_tenant_comfy")]
mod video_agent;
#[cfg(not(feature = "multi_tenant_comfy"))]
#[path = "video_agent_legacy.rs"]
mod video_agent;
mod workflow_ui;
mod workspace;
pub fn build_router(state: AppState) -> Router {
let web_dir = std::env::var("BOOSAPI_WEB_DIR").unwrap_or_else(|_| "dist/web".to_string());
// Build the API router
let api_routes = Router::new()
.route("/healthz", get(health::health_check))
@@ -51,12 +47,12 @@ pub fn build_router(state: AppState) -> Router {
.merge(proxy::routes())
.merge(picture_agent::routes())
.merge(comfyui_agent::routes())
.merge(comfyui::routes());
// /comfy/* reverse proxy is only mounted when the per-user pool is
// active; without it there's no per-user container to forward to.
#[cfg(feature = "multi_tenant_comfy")]
let api_routes = api_routes.merge(comfy_proxy::routes());
.merge(digital_human::routes())
.merge(comfyui::routes())
.merge(crate::workflow::routes::routes())
.route("/comfy", any(removed_comfy_workbench))
.route("/comfy/", any(removed_comfy_workbench))
.route("/comfy/{*path}", any(removed_comfy_workbench));
let api_routes = api_routes
.merge(auth::routes())
@@ -67,6 +63,7 @@ pub fn build_router(state: AppState) -> Router {
.merge(organizations::routes())
.merge(search::routes())
.merge(site_announcements::routes())
.merge(workflow_ui::routes())
.merge(stats::routes())
.merge(tasks::routes())
.merge(update_center::routes())
@@ -76,50 +73,125 @@ pub fn build_router(state: AppState) -> Router {
.merge(video_agent::routes())
.with_state(state);
// Static file serving with SPA fallback
let web_dir_owned = web_dir.clone();
let static_service = ServeDir::new(&web_dir)
.append_index_html_on_directories(true)
.fallback(service_fn(move |req: Request<Body>| {
let web_dir = web_dir_owned.clone();
async move {
let path = req.uri().path();
let router = Router::new().merge(api_routes);
match std::env::var("BOOSAPI_WEB_DIR")
.ok()
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
{
Some(web_dir) => {
let web_dir_owned = web_dir.clone();
let static_service = ServeDir::new(&web_dir)
.append_index_html_on_directories(true)
.fallback(service_fn(move |req: Request<Body>| {
let web_dir = web_dir_owned.clone();
async move {
let path = req.uri().path();
if is_api_path(path) {
return Ok::<_, Infallible>(api_not_found_response(path));
}
// API routes that don't match → return JSON error, not index.html
if path.starts_with("/api/") || path.starts_with("/v1/") {
return Ok::<_, Infallible>(
Response::builder()
.status(StatusCode::NOT_FOUND)
.header("Content-Type", "application/json; charset=utf-8")
.body(Body::from(
serde_json::json!({"error": "Not found", "path": path}).to_string(),
))
.unwrap(),
);
}
// SPA fallback: serve index.html for client-side routing
let index_path = format!("{}/index.html", web_dir);
match tokio::fs::read_to_string(&index_path).await {
Ok(html) => Ok::<_, Infallible>(
Response::builder()
.status(StatusCode::OK)
.header("Content-Type", "text/html; charset=utf-8")
.body(Body::from(html))
.unwrap(),
),
Err(_) => Ok::<_, Infallible>(
Response::builder()
.status(StatusCode::NOT_FOUND)
.header("Content-Type", "text/plain; charset=utf-8")
.body(Body::from(format!("404: {} not found", path)))
.unwrap(),
),
}
}
}));
Router::new()
.merge(api_routes)
.fallback_service(static_service)
let index_path = format!("{}/index.html", web_dir);
match tokio::fs::read_to_string(&index_path).await {
Ok(html) => Ok::<_, Infallible>(html_response(html)),
Err(_) => Ok::<_, Infallible>(
Response::builder()
.status(StatusCode::NOT_FOUND)
.header(header::CONTENT_TYPE, "text/plain; charset=utf-8")
.body(Body::from(format!("404: {} not found", path)))
.unwrap_or_else(|_| StatusCode::NOT_FOUND.into_response()),
),
}
}
}));
router.fallback_service(static_service)
}
None => router.fallback(embedded_static_handler),
}
}
fn api_not_found_response(path: &str) -> Response {
Response::builder()
.status(StatusCode::NOT_FOUND)
.header(header::CONTENT_TYPE, "application/json; charset=utf-8")
.body(Body::from(
serde_json::json!({"error": "Not found", "path": path}).to_string(),
))
.unwrap_or_else(|_| StatusCode::NOT_FOUND.into_response())
}
async fn embedded_static_handler(req: Request<Body>) -> Response {
let path = req.uri().path();
if is_api_path(path) {
return api_not_found_response(path);
}
let asset_path = embedded_asset_path(path);
if let Some(file) = EMBEDDED_WEB.get_file(&asset_path) {
return embedded_file_response(&asset_path, file.contents());
}
if let Some(file) = EMBEDDED_WEB.get_file("index.html") {
return embedded_file_response("index.html", file.contents());
}
Response::builder()
.status(StatusCode::NOT_FOUND)
.header(header::CONTENT_TYPE, "text/plain; charset=utf-8")
.body(Body::from("embedded dist/web/index.html not found"))
.unwrap_or_else(|_| StatusCode::NOT_FOUND.into_response())
}
fn is_api_path(path: &str) -> bool {
path.starts_with("/api/") || path.starts_with("/v1/")
}
fn embedded_asset_path(path: &str) -> String {
let trimmed = path.trim_start_matches('/');
if trimmed.is_empty() || trimmed.ends_with('/') {
format!("{trimmed}index.html")
} else {
trimmed.to_string()
}
}
fn html_response(html: String) -> Response {
Response::builder()
.status(StatusCode::OK)
.header(header::CONTENT_TYPE, "text/html; charset=utf-8")
.body(Body::from(html))
.unwrap_or_else(|_| StatusCode::INTERNAL_SERVER_ERROR.into_response())
}
fn embedded_file_response(path: &str, bytes: &[u8]) -> Response {
let content_type = mime_guess::from_path(path)
.first_or_octet_stream()
.essence_str()
.to_string();
let mut response = Response::builder()
.status(StatusCode::OK)
.header(header::CONTENT_TYPE, content_type)
.body(Body::from(bytes.to_vec()))
.unwrap_or_else(|_| StatusCode::INTERNAL_SERVER_ERROR.into_response());
if path.starts_with("assets/") || path.contains(".") && !path.ends_with("index.html") {
response.headers_mut().insert(
header::CACHE_CONTROL,
HeaderValue::from_static("public, max-age=31536000, immutable"),
);
}
response
}
async fn removed_comfy_workbench() -> (StatusCode, Json<Value>) {
(
StatusCode::GONE,
Json(json!({
"error": {
"message": "内置 Docker ComfyUI 工作台已移除。请配置外部 comfyui.base_url 后使用 /api/comfyui/proxy，或打开 /workflows 使用 ComfyUI 风格工作流画布。",
"type": "ComfyUIWorkbenchRemoved"
},
"redirect": "/workflows"
})),
)
}
src/routes/picture_agent.rs
+584 -29
View File
@@ -1,12 +1,13 @@
use crate::{
auth::jwt::validate_jwt,
db::models::background_task::BackgroundTask,
db::models::downstream_api_key::DownstreamApiKey,
error::{AppError, AppResult},
AppState,
};
use axum::{
body::Body,
extract::{Multipart, Query, State},
extract::{Multipart, Path, Query, State},
http::{header, HeaderMap, Response, StatusCode},
response::IntoResponse,
routing::{get, post},
@@ -19,6 +20,11 @@ use std::time::Duration;
const TEMPLATE_GALLERY_URL: &str =
"https://weiyyin.github.io/linux-do-gpt-image/linux_do_2036927_prompts_images.filtered.json";
const TASK_TYPE_PICTURE_GENERATION: &str = "picture-agent-generation";
const TEMPLATE_GALLERY_TIMEOUT_SECS: u64 = 30;
const PICTURE_GENERATION_TIMEOUT_SECS: u64 = 300;
const PICTURE_GENERATION_RETRY_SECS: u64 = 180;
const STALE_RUNNING_JOB_MINUTES: i64 = 12;
#[derive(Debug, Clone)]
struct PictureAgentAuth {
@@ -33,6 +39,13 @@ pub fn routes() -> Router<AppState> {
.route("/api/picture-agent/status", get(status_handler))
.route("/api/picture-agent/models", get(models_handler))
.route("/api/picture-agent/templates", get(templates_handler))
.route(
"/api/picture-agent/jobs",
get(list_jobs).post(create_generation_job),
)
.route("/api/picture-agent/jobs/clear", post(clear_jobs))
.route("/api/picture-agent/jobs/{id}", get(get_job))
.route("/api/picture-agent/jobs/{id}/cancel", post(cancel_job))
.route(
"/api/picture-agent/images/generations",
post(images_generations),
@@ -54,6 +67,11 @@ struct GalleryItem {
image_host_url: Option<String>,
}
#[derive(Debug, Deserialize)]
struct JobListQuery {
limit: Option<i64>,
}
fn local_v1_base(state: &AppState) -> String {
let host = if state.config.server.host == "0.0.0.0" || state.config.server.host == "::" {
"127.0.0.1".to_string()
@@ -149,6 +167,52 @@ fn is_exact_model_pattern(value: &str) -> bool {
&& !normalized.contains('?')
}
fn is_picture_model_pattern(value: &str) -> bool {
if !is_exact_model_pattern(value) {
return false;
}
let normalized = value.trim().to_ascii_lowercase();
let image_keywords = [
"gpt-image",
"dall-e",
"dalle",
"midjourney",
"recraft",
"kolors",
"cogview",
"seedream",
"sdxl",
"stable-diffusion",
"stablediffusion",
"pixart",
"playground",
"firefly",
"imagen",
"flux",
"image",
"img",
"canvas",
"painting",
"draw",
];
if image_keywords
.iter()
.any(|keyword| normalized.contains(keyword))
{
return true;
}
let chat_keywords = [
"claude", "chatgpt", "gpt-4", "gpt-3", "gpt-5", "o1", "o3", "o4", "gemini", "deepseek",
"qwen", "llama", "mistral", "command", "kimi", "doubao", "glm", "hunyuan", "spark",
"baichuan", "ernie", "wenxin", "grok",
];
!chat_keywords
.iter()
.any(|keyword| normalized.contains(keyword))
}
async fn status_handler(
State(state): State<AppState>,
headers: HeaderMap,
@@ -183,7 +247,7 @@ async fn models_handler(
let models: Vec<String> = rows
.into_iter()
.map(|value| value.trim().to_string())
.filter(|value| is_exact_model_pattern(value))
.filter(|value| is_picture_model_pattern(value))
.filter(|value| seen.insert(value.clone()))
.collect();
@@ -199,6 +263,30 @@ async fn models_handler(
})))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn picture_model_filter_keeps_image_models() {
assert!(is_picture_model_pattern("gpt-image-2"));
assert!(is_picture_model_pattern("dall-e-3"));
assert!(is_picture_model_pattern("flux-schnell"));
assert!(is_picture_model_pattern("replicate/recraft-v3"));
assert!(is_picture_model_pattern("google/imagen-3"));
assert!(is_picture_model_pattern("midjourney"));
}
#[test]
fn picture_model_filter_rejects_chat_models() {
assert!(!is_picture_model_pattern("gpt-4o"));
assert!(!is_picture_model_pattern("claude-3.7-sonnet"));
assert!(!is_picture_model_pattern("gemini-2.5-pro"));
assert!(!is_picture_model_pattern("deepseek-chat"));
assert!(!is_picture_model_pattern("qwen-plus"));
}
}
fn prompt_title(prompt: &str, post_number: Option<i64>) -> String {
let clean = prompt
.split_whitespace()
@@ -255,8 +343,8 @@ async fn templates_handler(
Query(query): Query<TemplateQuery>,
) -> AppResult<Json<Value>> {
let _auth = resolve_picture_agent_auth(&headers, &state).await?;
let limit = query.limit.unwrap_or(48).clamp(1, 96);
let items = proxy_client()
let limit = query.limit.unwrap_or(144).clamp(1, 240);
let items = template_client()
.get(TEMPLATE_GALLERY_URL)
.send()
.await
@@ -311,9 +399,20 @@ async fn templates_handler(
})))
}
fn proxy_client() -> Client {
fn template_client() -> Client {
Client::builder()
.timeout(Duration::from_secs(180))
.timeout(Duration::from_secs(TEMPLATE_GALLERY_TIMEOUT_SECS))
.build()
.unwrap_or_else(|_| Client::new())
}
fn picture_generation_client() -> Client {
picture_generation_client_with_timeout(PICTURE_GENERATION_TIMEOUT_SECS)
}
fn picture_generation_client_with_timeout(timeout_secs: u64) -> Client {
Client::builder()
.timeout(Duration::from_secs(timeout_secs))
.build()
.unwrap_or_else(|_| Client::new())
}
@@ -336,6 +435,473 @@ async fn archive_picture_agent_response(
.await;
}
async fn execute_generation_payload(
state: &AppState,
auth: &PictureAgentAuth,
body: Value,
) -> AppResult<(Value, Option<String>, Option<String>)> {
execute_generation_payload_with_timeout(state, auth, body, PICTURE_GENERATION_TIMEOUT_SECS)
.await
}
fn classify_upstream_failure(status: Option<StatusCode>, text: &str) -> String {
let status_text = status
.map(|s| format!("{s}"))
.unwrap_or_else(|| "unknown".to_string());
let lower = text.to_ascii_lowercase();
if lower.contains("524") || lower.contains("timeout") {
"图片上游超时，请稍后重试或缩短提示词".to_string()
} else if lower.contains("502") || lower.contains("bad gateway") {
"图片上游网关异常，请稍后重试".to_string()
} else if lower.contains("503") || lower.contains("service unavailable") {
"图片上游暂时不可用，请稍后重试".to_string()
} else if status.is_some_and(|s| s == StatusCode::GATEWAY_TIMEOUT) {
"图片上游请求超时，请稍后重试".to_string()
} else {
format!("Upstream returned {status_text}: {text}")
}
}
async fn execute_generation_payload_with_timeout(
state: &AppState,
auth: &PictureAgentAuth,
body: Value,
timeout_secs: u64,
) -> AppResult<(Value, Option<String>, Option<String>)> {
let prompt = body
.get("prompt")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_string);
let model = body
.get("model")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_string);
let requested_model = model
.clone()
.ok_or_else(|| AppError::BadRequest("model is required".into()))?;
let response = picture_generation_client_with_timeout(timeout_secs)
.post(format!("{}/images/generations", local_v1_base(state)))
.bearer_auth(auth.outbound_key.clone())
.json(&body)
.send()
.await
.map_err(|err| {
let hint = if err.is_timeout() {
format!(
"图片上游在 {} 分钟内没有返回，请稍后重试、缩短提示词或切换图片模型通道",
PICTURE_GENERATION_TIMEOUT_SECS / 60
)
} else if err.is_connect() {
"无法连接本地 BoosAPI 图片路由，请确认后端服务仍在运行".to_string()
} else {
err.to_string()
};
AppError::UpstreamError(hint)
})?;
let status = response.status();
let text = response
.text()
.await
.map_err(|err| AppError::UpstreamError(err.to_string()))?;
if !status.is_success() {
return Err(AppError::UpstreamError(classify_upstream_failure(
Some(status),
&text,
)));
}
let payload: Value = serde_json::from_str(&text).map_err(|err| {
AppError::UpstreamError(format!(
"Invalid JSON from upstream: {err}; body: {}",
text.chars().take(500).collect::<String>()
))
})?;
Ok((payload, prompt, Some(requested_model)))
}
async fn mark_stale_picture_jobs_failed(state: &AppState) -> AppResult<()> {
sqlx::query(
"UPDATE background_tasks
SET status = 'failed',
message = '生图任务已中断',
error = '后台任务长时间没有进展，可能是后端重启、上游超时或图片通道无响应。请重新提交。',
updated_at = datetime('now'),
finished_at = datetime('now')
WHERE task_type = ?
AND status IN ('pending', 'running')
AND datetime(COALESCE(updated_at, created_at)) < datetime('now', ?)",
)
.bind(TASK_TYPE_PICTURE_GENERATION)
.bind(format!("-{STALE_RUNNING_JOB_MINUTES} minutes"))
.execute(&*state.db)
.await?;
Ok(())
}
fn job_visible_to_auth(task: &BackgroundTask, auth: &PictureAgentAuth) -> bool {
if auth.role == "admin" && auth.user_id.is_none() {
return true;
}
let Some(raw) = task.result_json.as_deref() else {
return false;
};
let Ok(value) = serde_json::from_str::<Value>(raw) else {
return false;
};
value.get("userId").and_then(Value::as_i64) == auth.user_id
}
fn task_to_json(task: BackgroundTask) -> Value {
let result = task
.result_json
.as_deref()
.and_then(|raw| serde_json::from_str::<Value>(raw).ok())
.unwrap_or_else(|| json!({}));
json!({
"id": task.id,
"type": task.task_type,
"title": task.title,
"status": task.status,
"message": task.message,
"error": task.error,
"progress": task.progress.unwrap_or(0),
"progressTotal": task.progress_total.unwrap_or(4),
"result": result,
"createdAt": task.created_at,
"updatedAt": task.updated_at,
"startedAt": task.started_at,
"finishedAt": task.finished_at,
})
}
fn is_terminal_job_status(status: &str) -> bool {
matches!(status, "succeeded" | "failed" | "cancelled")
}
async fn load_picture_job(state: &AppState, id: &str, auth: &PictureAgentAuth) -> AppResult<Value> {
mark_stale_picture_jobs_failed(state).await?;
let task: BackgroundTask =
sqlx::query_as("SELECT * FROM background_tasks WHERE id = ? AND task_type = ? LIMIT 1")
.bind(id)
.bind(TASK_TYPE_PICTURE_GENERATION)
.fetch_optional(&*state.db)
.await?
.ok_or_else(|| AppError::NotFound("Picture job not found".into()))?;
if !job_visible_to_auth(&task, auth) {
return Err(AppError::NotFound("Picture job not found".into()));
}
Ok(task_to_json(task))
}
async fn update_job_progress(
state: &AppState,
id: &str,
progress: i64,
message: &str,
) -> AppResult<()> {
sqlx::query(
"UPDATE background_tasks
SET progress = ?, message = ?, updated_at = datetime('now')
WHERE id = ? AND status IN ('pending', 'running')",
)
.bind(progress)
.bind(message)
.bind(id)
.execute(&*state.db)
.await?;
Ok(())
}
async fn is_job_cancelled(state: &AppState, id: &str) -> AppResult<bool> {
let status: Option<String> =
sqlx::query_scalar("SELECT status FROM background_tasks WHERE id = ? LIMIT 1")
.bind(id)
.fetch_optional(&*state.db)
.await?;
Ok(status.as_deref() == Some("cancelled"))
}
async fn create_generation_job(
State(state): State<AppState>,
headers: HeaderMap,
Json(body): Json<Value>,
) -> AppResult<Json<Value>> {
let auth = resolve_picture_agent_auth(&headers, &state).await?;
let prompt = body
.get("prompt")
.and_then(Value::as_str)
.unwrap_or("")
.trim()
.to_string();
let model = body
.get("model")
.and_then(Value::as_str)
.unwrap_or("gpt-image-2")
.trim()
.to_string();
if prompt.is_empty() {
return Err(AppError::BadRequest("prompt is required".into()));
}
if model.is_empty() {
return Err(AppError::BadRequest("model is required".into()));
}
let task_id = format!("pic_{}", uuid::Uuid::new_v4().simple());
let title = prompt.chars().take(30).collect::<String>();
let initial_result = json!({
"userId": auth.user_id,
"role": auth.role,
"keyName": auth.key_name,
"prompt": prompt,
"model": model,
"mode": "generate",
"request": body,
});
sqlx::query(
"INSERT INTO background_tasks
(id, task_type, title, status, message, dedupe_key, result_json, progress, progress_total, started_at)
VALUES (?, ?, ?, 'running', '任务已提交，等待路由', ?, ?, 1, 4, datetime('now'))",
)
.bind(&task_id)
.bind(TASK_TYPE_PICTURE_GENERATION)
.bind(if title.trim().is_empty() { "生图任务" } else { &title })
.bind(format!("picture-agent:{}", auth.user_id.map(|id| id.to_string()).unwrap_or_else(|| "admin".into())))
.bind(initial_result.to_string())
.execute(&*state.db)
.await?;
let state_for_task = state.clone();
let task_id_for_spawn = task_id.clone();
let auth_for_task = auth.clone();
tokio::spawn(async move {
let result = async {
update_job_progress(
&state_for_task,
&task_id_for_spawn,
2,
"本地图片路由已确认，正在请求上游",
)
.await?;
let generation_result =
execute_generation_payload(&state_for_task, &auth_for_task, body.clone()).await;
let (payload, prompt, model) = match generation_result {
Ok(value) => value,
Err(AppError::UpstreamError(message))
if message.contains("超时") || message.contains("timeout") =>
{
update_job_progress(
&state_for_task,
&task_id_for_spawn,
2,
"上游超时，正在自动重试一次",
)
.await?;
execute_generation_payload_with_timeout(
&state_for_task,
&auth_for_task,
body,
PICTURE_GENERATION_RETRY_SECS,
)
.await?
}
Err(err) => return Err(err),
};
if is_job_cancelled(&state_for_task, &task_id_for_spawn).await? {
return Ok(());
}
update_job_progress(
&state_for_task,
&task_id_for_spawn,
3,
"图片已返回，正在归档到工作空间",
)
.await?;
archive_picture_agent_response(
&state_for_task,
&auth_for_task,
prompt.clone(),
model.clone(),
&payload,
)
.await;
let urls = payload
.get("data")
.and_then(Value::as_array)
.cloned()
.unwrap_or_default();
let final_result = json!({
"userId": auth_for_task.user_id,
"role": auth_for_task.role,
"keyName": auth_for_task.key_name,
"prompt": prompt,
"model": model,
"mode": "generate",
"response": payload,
"images": urls,
});
sqlx::query(
"UPDATE background_tasks
SET status = 'succeeded', message = '生图完成', result_json = ?, progress = 4,
updated_at = datetime('now'), finished_at = datetime('now')
WHERE id = ? AND status IN ('pending', 'running')",
)
.bind(final_result.to_string())
.bind(&task_id_for_spawn)
.execute(&*state_for_task.db)
.await?;
Ok::<(), AppError>(())
}
.await;
if let Err(err) = result {
let _ = sqlx::query(
"UPDATE background_tasks
SET status = 'failed', message = '生图失败', error = ?, updated_at = datetime('now'), finished_at = datetime('now')
WHERE id = ? AND status IN ('pending', 'running')",
)
.bind(err.to_string())
.bind(&task_id_for_spawn)
.execute(&*state_for_task.db)
.await;
}
});
let task = load_picture_job(&state, &task_id, &auth).await?;
Ok(Json(json!({ "success": true, "job": task })))
}
async fn cancel_job(
State(state): State<AppState>,
headers: HeaderMap,
Path(id): Path<String>,
) -> AppResult<Json<Value>> {
let auth = resolve_picture_agent_auth(&headers, &state).await?;
let task: BackgroundTask =
sqlx::query_as("SELECT * FROM background_tasks WHERE id = ? AND task_type = ? LIMIT 1")
.bind(&id)
.bind(TASK_TYPE_PICTURE_GENERATION)
.fetch_optional(&*state.db)
.await?
.ok_or_else(|| AppError::NotFound("Picture job not found".into()))?;
if !job_visible_to_auth(&task, &auth) {
return Err(AppError::NotFound("Picture job not found".into()));
}
if task.status != "pending" && task.status != "running" {
return Ok(Json(json!({
"success": true,
"job": task_to_json(task),
})));
}
sqlx::query(
"UPDATE background_tasks
SET status = 'cancelled',
message = '已取消',
error = NULL,
updated_at = datetime('now'),
finished_at = datetime('now')
WHERE id = ? AND task_type = ? AND status IN ('pending', 'running')",
)
.bind(&id)
.bind(TASK_TYPE_PICTURE_GENERATION)
.execute(&*state.db)
.await?;
let job = load_picture_job(&state, &id, &auth).await?;
Ok(Json(json!({ "success": true, "job": job })))
}
async fn list_jobs(
State(state): State<AppState>,
headers: HeaderMap,
Query(query): Query<JobListQuery>,
) -> AppResult<Json<Value>> {
let auth = resolve_picture_agent_auth(&headers, &state).await?;
mark_stale_picture_jobs_failed(&state).await?;
let limit = query.limit.unwrap_or(20).clamp(1, 100);
let tasks: Vec<BackgroundTask> = sqlx::query_as(
"SELECT * FROM background_tasks
WHERE task_type = ?
ORDER BY created_at DESC
LIMIT ?",
)
.bind(TASK_TYPE_PICTURE_GENERATION)
.bind(limit * 3)
.fetch_all(&*state.db)
.await?;
let jobs: Vec<Value> = tasks
.into_iter()
.filter(|task| job_visible_to_auth(task, &auth))
.take(limit as usize)
.map(task_to_json)
.collect();
Ok(Json(json!({ "success": true, "jobs": jobs })))
}
async fn clear_jobs(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
let auth = resolve_picture_agent_auth(&headers, &state).await?;
mark_stale_picture_jobs_failed(&state).await?;
let tasks: Vec<BackgroundTask> = sqlx::query_as(
"SELECT * FROM background_tasks
WHERE task_type = ?
ORDER BY created_at DESC",
)
.bind(TASK_TYPE_PICTURE_GENERATION)
.fetch_all(&*state.db)
.await?;
let mut cleared = 0_i64;
for task in tasks {
if !job_visible_to_auth(&task, &auth) || !is_terminal_job_status(&task.status) {
continue;
}
let affected = sqlx::query(
"DELETE FROM background_tasks
WHERE id = ? AND task_type = ?",
)
.bind(&task.id)
.bind(TASK_TYPE_PICTURE_GENERATION)
.execute(&*state.db)
.await?
.rows_affected();
cleared += i64::try_from(affected).unwrap_or(0);
}
Ok(Json(json!({
"success": true,
"cleared": cleared,
})))
}
async fn get_job(
State(state): State<AppState>,
headers: HeaderMap,
Path(id): Path<String>,
) -> AppResult<Json<Value>> {
let auth = resolve_picture_agent_auth(&headers, &state).await?;
let job = load_picture_job(&state, &id, &auth).await?;
Ok(Json(json!({ "success": true, "job": job })))
}
async fn json_response_and_archive(
state: &AppState,
auth: &PictureAgentAuth,
payload: Value,
prompt: Option<String>,
model: Option<String>,
) -> Response<Body> {
archive_picture_agent_response(state, auth, prompt, model, &payload).await;
Response::builder()
.status(StatusCode::OK)
.header(header::CONTENT_TYPE, "application/json; charset=utf-8")
.body(Body::from(payload.to_string()))
.unwrap_or_else(|_| StatusCode::INTERNAL_SERVER_ERROR.into_response())
}
async fn upstream_json_response_to_axum_and_archive(
state: &AppState,
auth: &PictureAgentAuth,
@@ -375,27 +941,8 @@ async fn images_generations(
Json(body): Json<Value>,
) -> AppResult<Response<Body>> {
let auth = resolve_picture_agent_auth(&headers, &state).await?;
let prompt = body
.get("prompt")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_string);
let model = body
.get("model")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_string);
let url = format!("{}/images/generations", local_v1_base(&state));
let response = proxy_client()
.post(url)
.bearer_auth(auth.outbound_key.clone())
.json(&body)
.send()
.await
.map_err(|err| AppError::UpstreamError(err.to_string()))?;
upstream_json_response_to_axum_and_archive(&state, &auth, response, prompt, model).await
let (payload, prompt, model) = execute_generation_payload(&state, &auth, body).await?;
Ok(json_response_and_archive(&state, &auth, payload, prompt, model).await)
}
async fn images_edits(
@@ -455,12 +1002,20 @@ async fn images_edits(
}
let url = format!("{}/images/edits", local_v1_base(&state));
let response = proxy_client()
let response = picture_generation_client()
.post(url)
.bearer_auth(auth.outbound_key.clone())
.multipart(form)
.send()
.await
.map_err(|err| AppError::UpstreamError(err.to_string()))?;
upstream_json_response_to_axum_and_archive(&state, &auth, response, prompt, model).await
let response = upstream_json_response_to_axum_and_archive(
&state,
&auth,
response,
prompt.clone(),
model.clone(),
)
.await?;
Ok(response)
}
src/routes/proxy/mod.rs
+162 -25
View File
@@ -103,10 +103,12 @@ async fn messages_handler(
let model = canonical.model.clone();
if canonical.stream {
let chunks = execute_messages_stream_with_conductor(state, key_id, model, canonical).await?;
let sse = Sse::new(stream::iter(chunks).map(|chunk| {
Ok::<_, Infallible>(axum::response::sse::Event::default().data(chunk))
}));
let chunks =
execute_messages_stream_with_conductor(state, key_id, model, canonical).await?;
let sse =
Sse::new(stream::iter(chunks).map(|chunk| {
Ok::<_, Infallible>(axum::response::sse::Event::default().data(chunk))
}));
Ok(sse.into_response())
} else {
let result = execute_messages_with_conductor(state, key_id, model, canonical).await?;
@@ -371,10 +373,53 @@ async fn images_generations(
.ok_or_else(|| AppError::BadRequest("model is required".into()))?
.to_string();
let route = resolve_route(&state, &model).await?;
let result = executors::execute_image_generation(state.clone(), route, body).await?;
record_billing(&state, key_id, &model, 0, 0, 0.0).await;
Ok(Json(result))
let started_at = Instant::now();
let selected = TokenRouter::new(state.db.clone())
.resolve_selected(&model)
.await?;
let body = body_with_selected_model(body, &selected);
let result = executors::execute_image_generation(
state.clone(),
selected_channel_to_route(selected.clone()),
body,
)
.await;
match result {
Ok(payload) => {
let cost = estimate_request_cost(&selected, 0, 0);
record_proxy_log_success(
&state,
key_id,
&selected,
&model,
Some(elapsed_ms_since(started_at)),
0,
0,
cost,
1,
false,
)
.await;
record_billing(&state, key_id, &model, 0, 0, cost).await;
Ok(Json(payload))
}
Err(err) => {
record_proxy_log_failure(
&state,
key_id,
Some(&selected),
&model,
upstream_status_from_error(&err),
Some(elapsed_ms_since(started_at)),
Some(&err.to_string()),
1,
false,
)
.await;
Err(err)
}
}
}
/// POST /v1/images/edits
@@ -452,21 +497,63 @@ async fn images_edits(
}
let model = model.ok_or_else(|| AppError::BadRequest("model is required".into()))?;
let route = resolve_route(&state, &model).await?;
let started_at = Instant::now();
let selected = TokenRouter::new(state.db.clone())
.resolve_selected(&model)
.await?;
// Override the multipart `model` field with the conductor-resolved
// upstream model name (route may map "dall-e-3" → upstream-specific id).
if route.model != model {
if selected.actual_model != model {
for field in fields.iter_mut() {
if field.name == "model" && field.bytes.is_none() {
field.text = Some(route.model.clone());
field.text = Some(selected.actual_model.clone());
}
}
}
let result = executors::execute_image_edit(state.clone(), route, fields).await?;
record_billing(&state, key_id, &model, 0, 0, 0.0).await;
Ok(Json(result))
let result = executors::execute_image_edit(
state.clone(),
selected_channel_to_route(selected.clone()),
fields,
)
.await;
match result {
Ok(payload) => {
let cost = estimate_request_cost(&selected, 0, 0);
record_proxy_log_success(
&state,
key_id,
&selected,
&model,
Some(elapsed_ms_since(started_at)),
0,
0,
cost,
1,
false,
)
.await;
record_billing(&state, key_id, &model, 0, 0, cost).await;
Ok(Json(payload))
}
Err(err) => {
record_proxy_log_failure(
&state,
key_id,
Some(&selected),
&model,
upstream_status_from_error(&err),
Some(elapsed_ms_since(started_at)),
Some(&err.to_string()),
1,
false,
)
.await;
Err(err)
}
}
}
/// GET /v1/models
@@ -564,6 +651,23 @@ async fn record_billing(
}
}
fn estimate_request_cost(
selected: &SelectedChannel,
prompt_tokens: i64,
completion_tokens: i64,
) -> f64 {
let unit_cost = selected.account_unit_cost.unwrap_or(0.0).max(0.0);
if unit_cost <= 0.0 {
return 0.0;
}
let total_tokens = (prompt_tokens.max(0) + completion_tokens.max(0)) as f64;
if total_tokens <= 0.0 {
unit_cost
} else {
total_tokens / 1_000_000.0 * unit_cost
}
}
async fn record_proxy_log_success(
state: &AppState,
key_id: Option<i64>,
@@ -699,6 +803,31 @@ fn body_with_selected_model(mut body: Value, selected: &SelectedChannel) -> Valu
body
}
fn ensure_stream_usage(mut body: Value) -> Value {
if let Some(object) = body.as_object_mut() {
match object.get_mut("stream_options") {
Some(options) => {
if let Some(map) = options.as_object_mut() {
map.entry("include_usage".to_string())
.or_insert_with(|| Value::Bool(true));
} else {
object.insert(
"stream_options".to_string(),
serde_json::json!({ "include_usage": true }),
);
}
}
None => {
object.insert(
"stream_options".to_string(),
serde_json::json!({ "include_usage": true }),
);
}
}
}
body
}
fn build_conductor(state: &AppState) -> ProxyConductor {
let db = state.db.clone();
let select_db = db.clone();
@@ -797,6 +926,7 @@ async fn execute_messages_with_conductor(
attempts,
} => {
let (prompt, completion) = extract_usage(&response);
let cost = estimate_request_cost(&selected, prompt, completion);
record_proxy_log_success(
&billing_state,
key_id,
@@ -805,7 +935,7 @@ async fn execute_messages_with_conductor(
Some(elapsed_ms_since(started_at)),
prompt,
completion,
0.0,
cost,
attempts,
false,
)
@@ -816,7 +946,7 @@ async fn execute_messages_with_conductor(
&billing_model,
prompt,
completion,
0.0,
cost,
)
.await;
Ok(response)
@@ -938,6 +1068,7 @@ async fn execute_messages_stream_with_conductor(
}
};
let cost = estimate_request_cost(&selected, 0, 0);
record_proxy_log_success(
&billing_state,
key_id,
@@ -946,12 +1077,12 @@ async fn execute_messages_stream_with_conductor(
Some(elapsed_ms_since(started_at)),
0,
0,
0.0,
cost,
attempts,
true,
)
.await;
record_billing(&billing_state, key_id, &billing_model, 0, 0, 0.0).await;
record_billing(&billing_state, key_id, &billing_model, 0, 0, cost).await;
Ok(chunks)
}
@@ -1005,6 +1136,7 @@ async fn execute_chat_completions_with_conductor(
attempts,
} => {
let (prompt, completion) = extract_usage(&response);
let cost = estimate_request_cost(&selected, prompt, completion);
record_proxy_log_success(
&billing_state,
key_id,
@@ -1013,7 +1145,7 @@ async fn execute_chat_completions_with_conductor(
Some(elapsed_ms_since(started_at)),
prompt,
completion,
0.0,
cost,
attempts,
false,
)
@@ -1024,7 +1156,7 @@ async fn execute_chat_completions_with_conductor(
&billing_model,
prompt,
completion,
0.0,
cost,
)
.await;
Ok(response)
@@ -1073,7 +1205,8 @@ async fn execute_chat_completions_stream_with_conductor(
downstream_policy: None,
attempt_fn: Box::new(move |ctx| {
let state = state.clone();
let body = body_with_selected_model(body.clone(), &ctx.selected);
let body =
ensure_stream_usage(body_with_selected_model(body.clone(), &ctx.selected));
async move {
match executors::execute_chat_completions_stream(
state,
@@ -1160,8 +1293,11 @@ async fn execute_chat_completions_stream_with_conductor(
.iter()
.rev()
.find_map(|chunk| {
let payload = chunk.strip_prefix("data: ")?;
if payload.trim() == "[DONE]" {
let payload = chunk
.strip_prefix("data: ")
.or_else(|| chunk.strip_prefix("data:"))?
.trim();
if payload == "[DONE]" {
return None;
}
let value: Value = serde_json::from_str(payload).ok()?;
@@ -1173,6 +1309,7 @@ async fn execute_chat_completions_stream_with_conductor(
}
})
.unwrap_or((0, 0));
let cost = estimate_request_cost(&selected, prompt, completion);
record_proxy_log_success(
&billing_state,
key_id,
@@ -1181,7 +1318,7 @@ async fn execute_chat_completions_stream_with_conductor(
Some(elapsed_ms_since(started_at)),
prompt,
completion,
0.0,
cost,
attempts,
true,
)
@@ -1192,7 +1329,7 @@ async fn execute_chat_completions_stream_with_conductor(
&billing_model,
prompt,
completion,
0.0,
cost,
)
.await;
Ok(chunks)
src/routes/site_announcements.rs
+15 -13
View File
@@ -49,13 +49,13 @@ pub fn routes() -> Router<AppState> {
put(update_announcement).delete(delete_announcement),
)
.route("/api/site-announcements/{id}/read", post(mark_read_admin))
.route("/api/site-announcements/read-all", post(mark_all_read_admin))
.route(
"/api/site-announcements/read-all",
post(mark_all_read_admin),
)
.route("/api/user/announcements", get(list_user_announcements))
.route("/api/user/announcements/{id}/read", post(mark_read_user))
.route(
"/api/user/announcements/read-all",
post(mark_all_read_user),
)
.route("/api/user/announcements/read-all", post(mark_all_read_user))
}
async fn list_admin_announcements(
@@ -95,11 +95,7 @@ async fn list_admin_announcements(
if let Some(status) = status {
query = query.bind(status);
}
let rows = query
.bind(limit)
.bind(offset)
.fetch_all(&*state.db)
.await?;
let rows = query.bind(limit).bind(offset).fetch_all(&*state.db).await?;
Ok(Json(Value::Array(
rows.into_iter().map(|row| row.into_admin_json()).collect(),
@@ -135,7 +131,9 @@ async fn list_user_announcements(
sql.push_str(" AND r.read_at IS NULL");
}
sql.push_str(" ORDER BY COALESCE(a.published_at, a.created_at) DESC, a.id DESC LIMIT ? OFFSET ?");
sql.push_str(
" ORDER BY COALESCE(a.published_at, a.created_at) DESC, a.id DESC LIMIT ? OFFSET ?",
);
let rows: Vec<UserAnnouncementRow> = sqlx::query_as(&sql)
.bind(user_id)
@@ -412,10 +410,14 @@ fn normalize_payload(payload: AnnouncementPayload) -> AppResult<NormalizedAnnoun
let title = payload.title.trim().to_string();
let content = payload.content.trim().to_string();
if title.is_empty() {
return Err(crate::error::AppError::BadRequest("title is required".into()));
return Err(crate::error::AppError::BadRequest(
"title is required".into(),
));
}
if content.is_empty() {
return Err(crate::error::AppError::BadRequest("content is required".into()));
return Err(crate::error::AppError::BadRequest(
"content is required".into(),
));
}
let level = payload
src/routes/user/mod.rs
+371 -93
View File
@@ -1,5 +1,6 @@
use crate::{
auth::jwt::validate_jwt,
db::models::background_task::BackgroundTask,
db::models::{downstream_api_key::DownstreamApiKey, user::User},
error::{AppError, AppResult},
AppState,
@@ -12,6 +13,7 @@ use axum::{
};
use serde::Deserialize;
use serde_json::{json, Value};
use sqlx::Row;
pub fn routes() -> Router<AppState> {
Router::new()
@@ -20,6 +22,305 @@ pub fn routes() -> Router<AppState> {
.route("/v1/user/api-keys/{id}", delete(delete_api_key))
.route("/api/user/logs", get(get_user_logs))
.route("/api/user/summary", get(get_user_summary))
.route(
"/api/user/models/marketplace",
get(get_user_models_marketplace),
)
}
async fn user_picture_log_items(state: &AppState, user_id: i64) -> AppResult<Vec<Value>> {
let tasks: Vec<BackgroundTask> = sqlx::query_as(
"SELECT *
FROM background_tasks
WHERE task_type = ?
AND json_extract(result_json, '$.userId') = ?
ORDER BY created_at DESC, id DESC",
)
.bind("picture-agent-generation")
.bind(user_id)
.fetch_all(&*state.db)
.await?;
Ok(tasks
.into_iter()
.map(|task| {
let result = task
.result_json
.as_deref()
.and_then(|raw| serde_json::from_str::<Value>(raw).ok())
.unwrap_or_else(|| json!({}));
let status = match task.status.as_str() {
"succeeded" => "success",
"failed" => "failed",
"cancelled" => "cancelled",
"pending" => "pending",
_ => "running",
};
json!({
"id": task.id,
"route_id": Value::Null,
"channel_id": Value::Null,
"account_id": Value::Null,
"downstream_api_key_id": Value::Null,
"site_id": Value::Null,
"username": Value::Null,
"site_name": Value::Null,
"site_url": Value::Null,
"downstream_key_name": result.get("keyName").cloned().unwrap_or(Value::Null),
"downstream_key_group_name": Value::Null,
"downstream_key_tags": Value::Null,
"model_requested": result.get("model").cloned().unwrap_or(Value::Null),
"model_actual": result.get("model").cloned().unwrap_or(Value::Null),
"status": status,
"http_status": if status == "success" { 200 } else if status == "running" || status == "pending" { 202 } else { 500 },
"latency_ms": Value::Null,
"first_byte_latency_ms": Value::Null,
"prompt_tokens": Value::Null,
"completion_tokens": Value::Null,
"total_tokens": Value::Null,
"estimated_cost": 0.0,
"error_message": task.error.clone(),
"retry_count": 0,
"is_stream": false,
"client_family": "picture-agent",
"client_app_id": "picture-agent",
"client_app_name": "Picture Agent",
"client_confidence": "high",
"created_at": task.created_at,
"source": "picture-agent",
"task_type": task.task_type,
"title": task.title,
"message": task.message,
"error": task.error,
"progress": task.progress,
"progress_total": task.progress_total,
"status_label": status,
"result": result,
})
})
.collect())
}
async fn user_usage_log_items(state: &AppState, user_id: i64) -> AppResult<Vec<Value>> {
let rows = sqlx::query(
r#"
SELECT
l.id,
l.model,
l.cost,
l.prompt_tokens,
l.completion_tokens,
l.total_tokens,
l.created_at,
k.name AS downstream_key_name
FROM usage_logs l
JOIN downstream_api_keys k ON k.id = l.downstream_key_id
WHERE k.user_id = ?
ORDER BY l.created_at DESC, l.id DESC
LIMIT 5000
"#,
)
.bind(user_id)
.fetch_all(&*state.db)
.await?;
Ok(rows
.into_iter()
.map(|row| {
let model = row
.try_get::<Option<String>, _>("model")
.ok()
.flatten()
.unwrap_or_default();
let model_trimmed = model.trim().to_string();
let source_model = model_trimmed
.split(':')
.next_back()
.unwrap_or(&model_trimmed)
.trim()
.to_string();
json!({
"id": format!("usage-{}", row.try_get::<i64, _>("id").unwrap_or_default()),
"route_id": Value::Null,
"channel_id": Value::Null,
"account_id": Value::Null,
"downstream_api_key_id": Value::Null,
"site_id": Value::Null,
"username": Value::Null,
"site_name": Value::Null,
"site_url": Value::Null,
"downstream_key_name": row.try_get::<Option<String>, _>("downstream_key_name").ok().flatten(),
"downstream_key_group_name": Value::Null,
"downstream_key_tags": Value::Null,
"model_requested": source_model,
"model_actual": source_model,
"status": "success",
"http_status": 200,
"latency_ms": Value::Null,
"first_byte_latency_ms": Value::Null,
"prompt_tokens": row.try_get::<Option<i64>, _>("prompt_tokens").ok().flatten(),
"completion_tokens": row.try_get::<Option<i64>, _>("completion_tokens").ok().flatten(),
"total_tokens": row.try_get::<Option<i64>, _>("total_tokens").ok().flatten(),
"estimated_cost": row.try_get::<Option<f64>, _>("cost").ok().flatten().unwrap_or(0.0),
"error_message": Value::Null,
"retry_count": 0,
"is_stream": false,
"client_family": "picture-agent",
"client_app_id": "picture-agent",
"client_app_name": "Picture Agent",
"client_confidence": "high",
"created_at": row.try_get::<Option<String>, _>("created_at").ok().flatten(),
"source": "usage-logs",
"task_type": "usage-log",
"title": "生图助手",
"message": model_trimmed,
"error": Value::Null,
"progress": Value::Null,
"progress_total": Value::Null,
})
})
.collect())
}
fn user_log_text(item: &Value) -> String {
[
item.get("model_requested").and_then(Value::as_str),
item.get("model_actual").and_then(Value::as_str),
item.get("downstream_key_name").and_then(Value::as_str),
item.get("site_name").and_then(Value::as_str),
item.get("site_url").and_then(Value::as_str),
item.get("error_message").and_then(Value::as_str),
item.get("message").and_then(Value::as_str),
item.get("title").and_then(Value::as_str),
item.get("source").and_then(Value::as_str),
]
.into_iter()
.flatten()
.collect::<Vec<_>>()
.join(" ")
.to_ascii_lowercase()
}
fn user_log_status(item: &Value) -> String {
item.get("status")
.and_then(Value::as_str)
.unwrap_or_default()
.to_ascii_lowercase()
}
fn user_log_created_at(item: &Value) -> String {
item.get("created_at")
.and_then(Value::as_str)
.unwrap_or_default()
.to_string()
}
fn user_log_cost(item: &Value) -> f64 {
item.get("estimated_cost")
.and_then(Value::as_f64)
.unwrap_or(0.0)
}
fn user_log_tokens(item: &Value) -> i64 {
item.get("total_tokens")
.and_then(Value::as_i64)
.unwrap_or(0)
}
fn matches_user_log_filters(
item: &Value,
q: Option<&str>,
status: Option<&str>,
date_from: Option<&str>,
date_to: Option<&str>,
) -> bool {
if let Some(q) = q {
let q = q.to_ascii_lowercase();
if !user_log_text(item).contains(&q) {
return false;
}
}
if let Some(status) = status {
let item_status = user_log_status(item);
let normalized = status.to_ascii_lowercase();
if normalized == "success" {
if item_status != "success" {
return false;
}
} else if normalized == "failed" {
if !matches!(item_status.as_str(), "failed" | "cancelled") {
return false;
}
} else if item_status != normalized {
return false;
}
}
let created_at = user_log_created_at(item);
if let Some(date_from) = date_from {
if !created_at.is_empty() && created_at.as_str() < date_from {
return false;
}
}
if let Some(date_to) = date_to {
if !created_at.is_empty() && created_at.as_str() > date_to {
return false;
}
}
true
}
fn sort_user_logs_desc(items: &mut [Value]) {
items.sort_by(|a, b| {
let created_cmp = user_log_created_at(b).cmp(&user_log_created_at(a));
if created_cmp != std::cmp::Ordering::Equal {
return created_cmp;
}
let id_a = a.get("id").map(Value::to_string).unwrap_or_default();
let id_b = b.get("id").map(Value::to_string).unwrap_or_default();
id_b.cmp(&id_a)
});
}
async fn load_user_log_items(state: &AppState, user_id: i64) -> AppResult<Vec<Value>> {
let mut items: Vec<Value> = sqlx::query_as::<_, crate::routes::stats::ProxyLogRow>(
r#"
SELECT
l.id, l.route_id, l.channel_id, l.account_id, l.downstream_api_key_id,
l.model_requested, l.model_actual, l.status, l.http_status,
l.latency_ms, l.first_byte_latency_ms,
l.prompt_tokens, l.completion_tokens, l.total_tokens, l.estimated_cost,
l.error_message, l.retry_count, l.is_stream,
l.client_family, l.client_app_id, l.client_app_name, l.client_confidence,
l.created_at,
a.site_id AS site_id,
COALESCE(NULLIF(a.username, ''), CASE WHEN a.id IS NULL THEN NULL ELSE '账号 #' || a.id END) AS username,
s.name AS site_name,
s.url AS site_url,
k.name AS downstream_key_name,
k.group_name AS downstream_key_group_name,
k.tags AS downstream_key_tags
FROM proxy_logs l
JOIN downstream_api_keys k ON l.downstream_api_key_id = k.id
LEFT JOIN accounts a ON a.id = l.account_id
LEFT JOIN sites s ON s.id = a.site_id
WHERE k.user_id = ?
ORDER BY l.created_at DESC
LIMIT 5000
"#,
)
.bind(user_id)
.fetch_all(&*state.db)
.await?
.into_iter()
.map(|row| serde_json::to_value(row).unwrap_or(json!({})))
.collect();
items.extend(user_usage_log_items(state, user_id).await?);
items.extend(user_picture_log_items(state, user_id).await?);
sort_user_logs_desc(&mut items);
Ok(items)
}
#[derive(Debug, Deserialize)]
@@ -190,98 +491,40 @@ async fn get_user_logs(
let offset = params.offset.unwrap_or((page - 1) * page_size).max(0);
let q = params.q.as_deref().map(str::trim).filter(|s| !s.is_empty());
let q_like = q.map(|s| format!("%{s}%"));
let status = params
.status
.as_deref()
.filter(|s| !s.is_empty() && *s != "all");
let date_from = params.date_from.as_deref().filter(|s| !s.is_empty());
let date_to = params.date_to.as_deref().filter(|s| !s.is_empty());
let items: Vec<Value> = sqlx::query_as::<_, crate::routes::stats::ProxyLogRow>(
r#"
SELECT
l.id, l.route_id, l.channel_id, l.account_id, l.downstream_api_key_id,
l.model_requested, l.model_actual, l.status, l.http_status,
l.latency_ms, l.first_byte_latency_ms,
l.prompt_tokens, l.completion_tokens, l.total_tokens, l.estimated_cost,
l.error_message, l.retry_count, l.is_stream,
l.client_family, l.client_app_id, l.client_app_name, l.client_confidence,
l.created_at,
a.site_id AS site_id,
COALESCE(NULLIF(a.username, ''), CASE WHEN a.id IS NULL THEN NULL ELSE '账号 #' || a.id END) AS username,
s.name AS site_name,
s.url AS site_url,
k.name AS downstream_key_name,
k.group_name AS downstream_key_group_name,
k.tags AS downstream_key_tags
FROM proxy_logs l
JOIN downstream_api_keys k ON l.downstream_api_key_id = k.id
LEFT JOIN accounts a ON a.id = l.account_id
LEFT JOIN sites s ON s.id = a.site_id
WHERE k.user_id = ?
AND (? IS NULL OR l.model_requested LIKE ?)
AND (? IS NULL OR l.status = ?)
AND (? IS NULL OR l.created_at >= ?)
AND (? IS NULL OR l.created_at <= ?)
ORDER BY l.created_at DESC
LIMIT ? OFFSET ?
"#,
)
.bind(user_id)
.bind(q_like.as_deref())
.bind(q_like.as_deref())
.bind(status)
.bind(status)
.bind(date_from)
.bind(date_from)
.bind(date_to)
.bind(date_to)
.bind(page_size)
.bind(offset)
.fetch_all(&*state.db)
.await?
.into_iter()
.map(|row| serde_json::to_value(row).unwrap_or(json!({})))
.collect();
let (total, success_count, total_cost, total_tokens): (i64, i64, f64, i64) = sqlx::query_as(
r#"
SELECT
COUNT(*),
SUM(CASE WHEN l.status = 'success' THEN 1 ELSE 0 END),
COALESCE(SUM(l.estimated_cost), 0.0),
COALESCE(SUM(l.total_tokens), 0)
FROM proxy_logs l
JOIN downstream_api_keys k ON l.downstream_api_key_id = k.id
WHERE k.user_id = ?
AND (? IS NULL OR l.model_requested LIKE ?)
AND (? IS NULL OR l.status = ?)
AND (? IS NULL OR l.created_at >= ?)
AND (? IS NULL OR l.created_at <= ?)
"#,
)
.bind(user_id)
.bind(q_like.as_deref())
.bind(q_like.as_deref())
.bind(status)
.bind(status)
.bind(date_from)
.bind(date_from)
.bind(date_to)
.bind(date_to)
.fetch_one(&*state.db)
.await?;
let items = load_user_log_items(&state, user_id).await?;
let filtered_items: Vec<Value> = items
.into_iter()
.filter(|item| matches_user_log_filters(item, q, status, date_from, date_to))
.collect();
let total = filtered_items.len() as i64;
let success_count = filtered_items
.iter()
.filter(|item| user_log_status(item).as_str() == "success")
.count() as i64;
let failed_count = total - success_count;
let total_cost = filtered_items.iter().map(user_log_cost).sum::<f64>();
let total_tokens = filtered_items.iter().map(user_log_tokens).sum::<i64>();
let paged_items = filtered_items
.into_iter()
.skip(offset as usize)
.take(page_size as usize)
.collect::<Vec<_>>();
Ok(Json(json!({
"items": items,
"items": paged_items,
"total": total,
"page": page,
"page_size": page_size,
"summary": {
"total": total,
"success": success_count,
"failed": total - success_count,
"failed": failed_count,
"total_cost": total_cost,
"total_tokens": total_tokens,
}
@@ -293,22 +536,57 @@ async fn get_user_summary(
headers: HeaderMap,
) -> AppResult<Json<Value>> {
let user_id = current_user_id(&headers, &state)?;
let stats: (i64, i64, f64) = sqlx::query_as(
"SELECT COUNT(*),
SUM(CASE WHEN l.status = 'success' THEN 1 ELSE 0 END),
COALESCE(SUM(l.estimated_cost), 0.0)
FROM proxy_logs l
JOIN downstream_api_keys k ON l.downstream_api_key_id = k.id
WHERE k.user_id = ?",
)
.bind(user_id)
.fetch_one(&*state.db)
.await?;
let items = load_user_log_items(&state, user_id).await?;
let total = items.len() as i64;
let success_count = items
.iter()
.filter(|item| user_log_status(item).as_str() == "success")
.count() as i64;
let total_cost = items.iter().map(user_log_cost).sum::<f64>();
Ok(Json(json!({
"totalRequests": stats.0,
"successRequests": stats.1,
"totalCost": stats.2,
"totalRequests": total,
"successRequests": success_count,
"totalCost": total_cost,
})))
}
async fn get_user_models_marketplace(
State(state): State<AppState>,
headers: HeaderMap,
Query(query): Query<crate::routes::admin::models_marketplace::MarketplaceQuery>,
) -> AppResult<Json<Value>> {
let _user_id = current_user_id(&headers, &state)?;
let value = crate::routes::admin::models_marketplace::marketplace_inner(&state.db).await?;
let models = value
.get("models")
.and_then(Value::as_array)
.map(|items| {
items
.iter()
.map(|model| {
let brand = model.get("name").and_then(Value::as_str).unwrap_or_default();
json!({
"name": brand,
"accountCount": 0,
"tokenCount": 0,
"avgLatency": model.get("avgLatency").cloned().unwrap_or(Value::Null),
"successRate": model.get("successRate").cloned().unwrap_or(Value::Null),
"description": model.get("description").cloned().unwrap_or(Value::Null),
"tags": model.get("tags").cloned().unwrap_or_else(|| json!([])),
"supportedEndpointTypes": model.get("supportedEndpointTypes").cloned().unwrap_or_else(|| json!([])),
"pricingSources": [],
"accounts": [],
})
})
.collect::<Vec<_>>()
})
.unwrap_or_default();
Ok(Json(json!({
"models": models,
"meta": value.get("meta").cloned().unwrap_or_else(|| json!({})),
"userView": true,
"includePricing": query.include_pricing.is_some(),
})))
}
src/routes/video_agent.rs
-785
View File
@@ -1,785 +0,0 @@
//! `/api/video-agent/*` — per-user multi-tenant refactor.
//!
//! Every handler resolves the caller from the JWT in the `Authorization`
//! header (no anonymous access) and, where it makes sense, lazily starts
//! that user's ComfyUI container via [`ComfyUIPool`]. The LLM/image/TTS
//! helpers continue to call BoosAPI's own `/v1/*` loopback — they are
//! independent of ComfyUI and stay shared.
//!
//! Status responses preserve every existing JSON key the frontend reads
//! and add `containerState` so the new full-screen loading overlay can
//! poll until the container is ready.
use axum::{
body::Body,
extract::State,
http::{header, HeaderMap, Response, StatusCode},
response::{
sse::{Event, Sse},
IntoResponse,
},
routing::{get, post},
Json, Router,
};
use futures::stream;
use serde_json::{json, Value};
use std::convert::Infallible;
use std::time::{Duration, Instant};
use tracing::info;
use crate::{
auth::jwt::validate_jwt,
error::{AppError, AppResult},
services::comfyui_pool::{ContainerState, PoolError},
AppState,
};
const COLD_START_HOLD: Duration = Duration::from_secs(30);
const COLD_START_POLL: Duration = Duration::from_millis(500);
const ADMIN_WORKSPACE_USER_ID: i64 = 0;
const SCRIPT_ANALYSIS_PROMPT: &str = r#"你是一个 AI 视频生成助手。分析用户提供的脚本，提取以下结构化信息，以 JSON 格式返回：
{
"title": "视频标题",
"characters": [
{
"name": "角色名",
"gender": "男/女",
"age": "年龄描述",
"appearance": "外貌详细描述（发型、脸型、服装、体型等）",
"personality": "性格特征"
}
],
"scenes": [
{
"id": 1,
"description": "场景描述",
"environment": "环境详细描述（时间、地点、天气、光线等）",
"characters": ["出现的角色名"],
"action": "场景中发生的动作",
"dialogue": [
{"character": "角色名", "line": "台词", "emotion": "情感语气"}
]
}
],
"style": "整体视觉风格（动画/写实/像素等）"
}
只返回 JSON，不要其他文字。"#;
const CHARACTER_PROMPT_GEN: &str = r#"根据角色描述生成 4 组稳定的图像生成提示词（正面、背面、左侧面、右侧面），用于 AI 图像生成模型。
每组提示词应包含：人物姿态、服装、光照、背景。
返回 JSON 格式：
{
"front": "正面提示词",
"back": "背面提示词",
"left": "左侧面提示词",
"right": "右侧面提示词"
}
只返回 JSON。"#;
const SCENE_PROMPT_GEN: &str = r#"根据场景描述生成图像生成提示词，用于 AI 图像生成模型创建场景背景。
包含：环境、光照、色调、氛围。
返回 JSON：
{
"prompt": "场景生成提示词",
"negative_prompt": "负面提示词"
}
只返回 JSON。"#;
pub fn routes() -> Router<AppState> {
Router::new()
.route("/api/video-agent/status", get(status))
.route("/api/video-agent/start", post(start_comfyui))
.route("/api/video-agent/stop", post(stop_comfyui))
.route("/api/video-agent/restart", post(restart_comfyui))
.route("/api/video-agent/chat", post(chat))
.route("/api/video-agent/analyze", post(analyze))
.route(
"/api/video-agent/character-prompts",
post(character_prompts),
)
.route(
"/api/video-agent/generate-character",
post(generate_character),
)
.route("/api/video-agent/scene-prompt", post(scene_prompt))
.route("/api/video-agent/generate-scene", post(generate_scene))
.route("/api/video-agent/generate-tts", post(generate_tts))
.route("/api/video-agent/full-workflow", post(full_workflow))
}
// ── Auth ───────────────────────────────────────────────────────────────
/// Resolve the caller workspace from `Authorization: Bearer <token>`.
/// User JWTs map to their real user id; admin/global proxy tokens map to
/// workspace 0 so the admin console can use the same Docker-backed ComfyUI.
fn resolve_user_id(headers: &HeaderMap, state: &AppState) -> AppResult<i64> {
let auth = headers
.get(header::AUTHORIZATION)
.and_then(|v| v.to_str().ok())
.and_then(|v| v.strip_prefix("Bearer "))
.ok_or_else(|| AppError::Unauthorized("Missing or invalid Authorization header".into()))?;
if auth == state.config.auth.auth_token || auth == state.config.auth.proxy_token {
return Ok(ADMIN_WORKSPACE_USER_ID);
}
let claims = validate_jwt(&state.config.auth.jwt_secret, auth)
.map_err(|_| AppError::Unauthorized("Invalid or expired token".into()))?;
claims
.sub
.parse::<i64>()
.map_err(|_| AppError::Unauthorized("Invalid user ID in token".into()))
}
// ── Pool helpers ───────────────────────────────────────────────────────
/// Map a [`PoolError`] to the appropriate [`AppError`].
fn pool_err_to_app(e: PoolError) -> AppError {
match e {
PoolError::AtCapacity(_) => AppError::ServiceUnavailable("comfyui pool at capacity".into()),
PoolError::StartTimeout(_) => {
AppError::GatewayTimeout("comfyui container start timeout".into())
}
other => AppError::Internal(anyhow::anyhow!(other)),
}
}
/// Lazily start the caller's container and wait for readiness up to 30 s.
/// Returns the host port for the loopback HTTP API.
async fn ensure_container_ready(state: &AppState, user_id: i64) -> AppResult<u16> {
state
.comfyui_pool
.get_or_start(user_id)
.await
.map_err(pool_err_to_app)?;
let deadline = Instant::now() + COLD_START_HOLD;
loop {
match state.comfyui_pool.state(user_id).await {
ContainerState::Ready => break,
ContainerState::Failed | ContainerState::Stopped => {
return Err(AppError::UpstreamError(
"comfyui container failed to start".into(),
));
}
ContainerState::Starting => {
if Instant::now() >= deadline {
return Err(AppError::GatewayTimeout(
"comfyui cold-start timeout".into(),
));
}
tokio::time::sleep(COLD_START_POLL).await;
}
}
}
state.comfyui_pool.touch(user_id).await;
state
.comfyui_pool
.entry(user_id)
.await
.map(|e| e.host_port)
.ok_or_else(|| AppError::UpstreamError("container vanished after ready".into()))
}
/// JSON status response. Preserves the keys the frontend already reads
/// (`status`, `pid`, `uptime_secs`, `mode`, `model_backend`, `base_url`,
/// `version`, `last_error`) and adds `containerState` (task 5.2).
async fn pool_status_json(state: &AppState, user_id: i64) -> Value {
let cs = state.comfyui_pool.state(user_id).await;
let entry = state.comfyui_pool.entry(user_id).await;
let port = entry.as_ref().map(|e| e.host_port);
let last_error = entry.as_ref().and_then(|e| e.last_error.clone());
let status = match cs {
ContainerState::Starting => "starting",
ContainerState::Ready => "running",
ContainerState::Stopped => "stopped",
ContainerState::Failed => "unhealthy",
};
let base_url = port
.map(|p| format!("http://127.0.0.1:{}", p))
.unwrap_or_default();
let model_backend = format!(
"{}/v1",
state
.comfyui_pool
.config()
.boosapi_base
.trim_end_matches("/v1")
.trim_end_matches('/')
);
json!({
"status": status,
"pid": Value::Null,
"uptime_secs": 0,
"mode": "container",
"model_backend": model_backend,
"base_url": base_url,
"version": Value::Null,
"last_error": last_error,
"containerState": cs.as_str(),
})
}
// ── Lifecycle handlers ─────────────────────────────────────────────────
async fn status(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
let user_id = resolve_user_id(&headers, &state)?;
Ok(Json(pool_status_json(&state, user_id).await))
}
async fn start_comfyui(
State(state): State<AppState>,
headers: HeaderMap,
) -> AppResult<Json<Value>> {
let user_id = resolve_user_id(&headers, &state)?;
// Idempotent: get_or_start returns the existing entry if already up.
state
.comfyui_pool
.get_or_start(user_id)
.await
.map_err(pool_err_to_app)?;
Ok(Json(json!({
"success": true,
"comfyui": pool_status_json(&state, user_id).await,
})))
}
async fn stop_comfyui(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
let user_id = resolve_user_id(&headers, &state)?;
state
.comfyui_pool
.stop(user_id)
.await
.map_err(pool_err_to_app)?;
Ok(Json(json!({
"success": true,
"comfyui": pool_status_json(&state, user_id).await,
})))
}
async fn restart_comfyui(
State(state): State<AppState>,
headers: HeaderMap,
) -> AppResult<Json<Value>> {
let user_id = resolve_user_id(&headers, &state)?;
// Stop is best-effort: if the container is already gone, that's fine.
let _ = state.comfyui_pool.stop(user_id).await;
state
.comfyui_pool
.get_or_start(user_id)
.await
.map_err(pool_err_to_app)?;
Ok(Json(json!({
"success": true,
"comfyui": pool_status_json(&state, user_id).await,
})))
}
// ── BoosAPI loopback helpers ───────────────────────────────────────────
//
// LLM, image, and TTS calls go to BoosAPI's own `/v1/*` endpoints — they
// are not ComfyUI traffic. We pull the base URL from the same config the
// pool uses so there's exactly one source of truth.
fn boosapi_v1_base(state: &AppState) -> String {
let raw = state
.comfyui_pool
.config()
.boosapi_base
.trim_end_matches('/');
if raw.ends_with("/v1") {
raw.to_string()
} else {
format!("{}/v1", raw)
}
}
async fn chat(
State(state): State<AppState>,
headers: HeaderMap,
Json(body): Json<Value>,
) -> AppResult<Sse<impl futures::Stream<Item = Result<Event, Infallible>>>> {
let user_id = resolve_user_id(&headers, &state)?;
// Lazily warm the container so a subsequent generate-* call doesn't
// pay the full cold start.
let _ = state.comfyui_pool.get_or_start(user_id).await;
state.comfyui_pool.touch(user_id).await;
let prompt = body
.get("message")
.and_then(Value::as_str)
.or_else(|| {
body.get("messages")
.and_then(Value::as_array)
.and_then(|messages| messages.last())
.and_then(|message| message.get("content"))
.and_then(Value::as_str)
})
.unwrap_or("生成一个短视频工作流");
let workflow = build_basic_workflow(prompt);
let events = vec![
json!({
"type": "text",
"content": "已收到需求，正在生成 ComfyUI workflow..."
}),
json!({
"type": "workflow_generated",
"workflow": workflow
}),
json!({
"type": "done"
}),
];
Ok(Sse::new(stream::iter(events.into_iter().map(|event| {
Ok(Event::default().data(event.to_string()))
}))))
}
fn build_basic_workflow(prompt: &str) -> Value {
json!({
"boosapi": {
"kind": "video-agent-workflow",
"version": 1,
"prompt": prompt,
},
"nodes": {},
})
}
fn required_str<'a>(body: &'a Value, key: &str) -> AppResult<&'a str> {
body.get(key)
.and_then(Value::as_str)
.filter(|s| !s.trim().is_empty())
.ok_or_else(|| AppError::BadRequest(format!("{key} is required")))
}
fn required_value<'a>(body: &'a Value, key: &str) -> AppResult<&'a Value> {
body.get(key)
.filter(|v| !v.is_null())
.ok_or_else(|| AppError::BadRequest(format!("{key} is required")))
}
fn strip_json_markers(s: &str) -> String {
let mut text = s.trim().to_string();
if text.starts_with("```json") {
text = text.trim_start_matches("```json").to_string();
} else if text.starts_with("```") {
text = text.trim_start_matches("```").to_string();
}
if text.ends_with("```") {
text = text.trim_end_matches("```").to_string();
}
text.trim().to_string()
}
fn parse_llm_json(raw: &str) -> AppResult<Value> {
let cleaned = strip_json_markers(raw);
serde_json::from_str::<Value>(&cleaned).map_err(|e| {
AppError::Internal(anyhow::anyhow!(
"failed to parse LLM JSON response: {}; response: {}",
e,
cleaned
))
})
}
async fn call_llm(state: &AppState, prompt: &str, system: &str) -> AppResult<String> {
let client = reqwest::Client::new();
let url = format!("{}/chat/completions", boosapi_v1_base(state));
info!("video-agent calling BoosAPI LLM endpoint: {}", url);
let resp = client
.post(url)
.bearer_auth(&state.config.auth.proxy_token)
.json(&json!({
"model": &state.config.comfyui.agent_model,
"messages": [
{ "role": "system", "content": system },
{ "role": "user", "content": prompt }
],
"temperature": 0.7
}))
.send()
.await
.map_err(|e| AppError::Internal(anyhow::anyhow!("LLM request failed: {}", e)))?;
let status = resp.status();
let data: Value = resp
.json()
.await
.map_err(|e| AppError::Internal(anyhow::anyhow!("LLM response parse failed: {}", e)))?;
if !status.is_success() {
return Err(AppError::Internal(anyhow::anyhow!(
"LLM upstream returned {}: {}",
status,
data
)));
}
Ok(data
.get("choices")
.and_then(Value::as_array)
.and_then(|choices| choices.first())
.and_then(|choice| choice.get("message"))
.and_then(|message| message.get("content"))
.and_then(Value::as_str)
.unwrap_or("")
.to_string())
}
async fn call_image_gen(state: &AppState, prompt: &str) -> AppResult<String> {
let client = reqwest::Client::new();
let url = format!("{}/images/generations", boosapi_v1_base(state));
info!("video-agent calling BoosAPI image endpoint: {}", url);
let resp = client
.post(url)
.bearer_auth(&state.config.auth.proxy_token)
.json(&json!({
"model": "gpt-image-2",
"prompt": prompt,
"n": 1,
"size": "1024x1024"
}))
.send()
.await
.map_err(|e| AppError::Internal(anyhow::anyhow!("image request failed: {}", e)))?;
let status = resp.status();
let data: Value = resp
.json()
.await
.map_err(|e| AppError::Internal(anyhow::anyhow!("image response parse failed: {}", e)))?;
if !status.is_success() {
return Err(AppError::Internal(anyhow::anyhow!(
"image upstream returned {}: {}",
status,
data
)));
}
Ok(data
.get("data")
.and_then(Value::as_array)
.and_then(|items| items.first())
.and_then(|item| item.get("url"))
.and_then(Value::as_str)
.unwrap_or("")
.to_string())
}
async fn call_tts(state: &AppState, text: &str, voice: &str) -> AppResult<Vec<u8>> {
let client = reqwest::Client::new();
let url = format!("{}/audio/speech", boosapi_v1_base(state));
info!("video-agent calling BoosAPI TTS endpoint: {}", url);
let resp = client
.post(url)
.bearer_auth(&state.config.auth.proxy_token)
.json(&json!({
"model": "tts-1",
"input": text,
"voice": voice
}))
.send()
.await
.map_err(|e| AppError::Internal(anyhow::anyhow!("TTS request failed: {}", e)))?;
let status = resp.status();
if !status.is_success() {
let body = resp.text().await.unwrap_or_default();
return Err(AppError::Internal(anyhow::anyhow!(
"TTS upstream returned {}: {}",
status,
body
)));
}
let bytes = resp
.bytes()
.await
.map_err(|e| AppError::Internal(anyhow::anyhow!("TTS response read failed: {}", e)))?;
Ok(bytes.to_vec())
}
// ── Per-request handlers (5.5) ─────────────────────────────────────────
async fn analyze(
State(state): State<AppState>,
headers: HeaderMap,
Json(body): Json<Value>,
) -> AppResult<Json<Value>> {
let user_id = resolve_user_id(&headers, &state)?;
state.comfyui_pool.touch(user_id).await;
let script = required_str(&body, "script")?;
let result = call_llm(&state, script, SCRIPT_ANALYSIS_PROMPT).await?;
let parsed = parse_llm_json(&result)?;
Ok(Json(json!({
"success": true,
"data": parsed
})))
}
async fn character_prompts(
State(state): State<AppState>,
headers: HeaderMap,
Json(body): Json<Value>,
) -> AppResult<Json<Value>> {
let user_id = resolve_user_id(&headers, &state)?;
state.comfyui_pool.touch(user_id).await;
let character = required_value(&body, "character")?;
let name = character
.get("name")
.and_then(Value::as_str)
.unwrap_or_default();
let appearance = character
.get("appearance")
.and_then(Value::as_str)
.unwrap_or_default();
let gender = character
.get("gender")
.and_then(Value::as_str)
.unwrap_or_default();
let age = character
.get("age")
.and_then(Value::as_str)
.unwrap_or_default();
let desc = format!("{name}: {appearance}, {gender}, {age}");
let result = call_llm(&state, &desc, CHARACTER_PROMPT_GEN).await?;
let prompts = parse_llm_json(&result)?;
Ok(Json(json!({
"success": true,
"data": prompts
})))
}
async fn generate_character(
State(state): State<AppState>,
headers: HeaderMap,
Json(body): Json<Value>,
) -> AppResult<Json<Value>> {
let user_id = resolve_user_id(&headers, &state)?;
// Image generation can take a while; ensure the container is warm so
// downstream ComfyUI workflow reuse doesn't pay another cold start.
let _ = ensure_container_ready(&state, user_id).await?;
let prompts = required_value(&body, "prompts")?;
let prompt_obj = prompts
.as_object()
.ok_or_else(|| AppError::BadRequest("prompts must be an object".to_string()))?;
let mut results = serde_json::Map::new();
for (angle, prompt_value) in prompt_obj {
let prompt = prompt_value.as_str().ok_or_else(|| {
AppError::BadRequest(format!("prompt for angle '{angle}' must be a string"))
})?;
let url = call_image_gen(&state, prompt).await?;
results.insert(angle.clone(), Value::String(url));
}
state.comfyui_pool.touch(user_id).await;
Ok(Json(json!({
"success": true,
"data": Value::Object(results)
})))
}
async fn scene_prompt(
State(state): State<AppState>,
headers: HeaderMap,
Json(body): Json<Value>,
) -> AppResult<Json<Value>> {
let user_id = resolve_user_id(&headers, &state)?;
state.comfyui_pool.touch(user_id).await;
let scene = required_value(&body, "scene")?;
let environment = scene
.get("environment")
.and_then(Value::as_str)
.unwrap_or_default();
let description = scene
.get("description")
.and_then(Value::as_str)
.unwrap_or_default();
let desc = format!("{environment}. {description}");
let result = call_llm(&state, &desc, SCENE_PROMPT_GEN).await?;
let prompt = parse_llm_json(&result)?;
Ok(Json(json!({
"success": true,
"data": prompt
})))
}
async fn generate_scene(
State(state): State<AppState>,
headers: HeaderMap,
Json(body): Json<Value>,
) -> AppResult<Json<Value>> {
let user_id = resolve_user_id(&headers, &state)?;
let _ = ensure_container_ready(&state, user_id).await?;
let prompt = required_str(&body, "prompt")?;
let url = call_image_gen(&state, prompt).await?;
state.comfyui_pool.touch(user_id).await;
Ok(Json(json!({
"success": true,
"data": {
"url": url
}
})))
}
async fn generate_tts(
State(state): State<AppState>,
headers: HeaderMap,
Json(body): Json<Value>,
) -> AppResult<Response<Body>> {
let user_id = resolve_user_id(&headers, &state)?;
state.comfyui_pool.touch(user_id).await;
let text = required_str(&body, "text")?;
let voice = body
.get("voice")
.and_then(Value::as_str)
.filter(|s| !s.trim().is_empty())
.unwrap_or("alloy");
let audio = call_tts(&state, text, voice).await?;
Ok((
StatusCode::OK,
[(header::CONTENT_TYPE, "audio/mpeg")],
audio,
)
.into_response())
}
async fn full_workflow(
State(state): State<AppState>,
headers: HeaderMap,
Json(body): Json<Value>,
) -> AppResult<Json<Value>> {
let user_id = resolve_user_id(&headers, &state)?;
let _ = ensure_container_ready(&state, user_id).await?;
let script = required_str(&body, "script")?;
let requested_style = body.get("style").and_then(Value::as_str);
let analysis_raw = call_llm(&state, script, SCRIPT_ANALYSIS_PROMPT).await?;
let analysis = parse_llm_json(&analysis_raw)?;
let mut characters = Vec::new();
if let Some(chars) = analysis.get("characters").and_then(Value::as_array) {
for character in chars {
let name = character
.get("name")
.and_then(Value::as_str)
.unwrap_or_default();
let appearance = character
.get("appearance")
.and_then(Value::as_str)
.unwrap_or_default();
let gender = character
.get("gender")
.and_then(Value::as_str)
.unwrap_or_default();
let age = character
.get("age")
.and_then(Value::as_str)
.unwrap_or_default();
let desc = format!("{name}: {appearance}, {gender}, {age}");
let prompts_raw = call_llm(&state, &desc, CHARACTER_PROMPT_GEN).await?;
let prompts = parse_llm_json(&prompts_raw)?;
let mut item = character.clone();
if let Some(obj) = item.as_object_mut() {
obj.insert("imagePrompts".to_string(), prompts);
}
characters.push(item);
}
}
let mut scenes = Vec::new();
if let Some(scene_items) = analysis.get("scenes").and_then(Value::as_array) {
for scene in scene_items {
let environment = scene
.get("environment")
.and_then(Value::as_str)
.unwrap_or_default();
let description = scene
.get("description")
.and_then(Value::as_str)
.unwrap_or_default();
let desc = format!("{environment}. {description}");
let scene_prompt_raw = call_llm(&state, &desc, SCENE_PROMPT_GEN).await?;
let image_prompt = parse_llm_json(&scene_prompt_raw)?;
let mut item = scene.clone();
if let Some(obj) = item.as_object_mut() {
obj.insert("imagePrompt".to_string(), image_prompt);
}
scenes.push(item);
}
}
let style = analysis
.get("style")
.cloned()
.or_else(|| requested_style.map(|s| Value::String(s.to_string())))
.unwrap_or(Value::Null);
state.comfyui_pool.touch(user_id).await;
Ok(Json(json!({
"success": true,
"data": {
"title": analysis.get("title").cloned().unwrap_or(Value::Null),
"style": style,
"characters": characters,
"scenes": scenes
}
})))
}
src/routes/video_agent_legacy.rs
+34 -22
View File
@@ -93,10 +93,7 @@ pub fn routes() -> Router<AppState> {
async fn status(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
check_auth(&headers, &state)?;
Ok(Json(
serde_json::to_value(state.comfyui.status_response().await)
.unwrap_or_else(|_| json!({ "status": "unknown" })),
))
Ok(Json(external_comfyui_status(&state).await))
}
async fn start_comfyui(
@@ -104,27 +101,19 @@ async fn start_comfyui(
headers: HeaderMap,
) -> AppResult<Json<Value>> {
check_auth(&headers, &state)?;
state
.comfyui
.start()
.await
.map_err(|err| AppError::Internal(anyhow::anyhow!(err)))?;
Ok(Json(json!({
"success": true,
"comfyui": state.comfyui.status_response().await
"message": "内置 Docker ComfyUI 工作台已移除；请配置外部 ComfyUI 或使用 /workflows 工作流画布。",
"comfyui": external_comfyui_status_value(&state).await
})))
}
async fn stop_comfyui(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
check_auth(&headers, &state)?;
state
.comfyui
.stop()
.await
.map_err(|err| AppError::Internal(anyhow::anyhow!(err)))?;
Ok(Json(json!({
"success": true,
"comfyui": state.comfyui.status_response().await
"message": "内置 Docker ComfyUI 工作台已移除；BoosAPI 不再管理 ComfyUI 进程生命周期。",
"comfyui": external_comfyui_status_value(&state).await
})))
}
@@ -133,14 +122,10 @@ async fn restart_comfyui(
headers: HeaderMap,
) -> AppResult<Json<Value>> {
check_auth(&headers, &state)?;
state
.comfyui
.restart()
.await
.map_err(|err| AppError::Internal(anyhow::anyhow!(err)))?;
Ok(Json(json!({
"success": true,
"comfyui": state.comfyui.status_response().await
"message": "内置 Docker ComfyUI 工作台已移除；请在外部 ComfyUI 服务侧重启。",
"comfyui": external_comfyui_status_value(&state).await
})))
}
@@ -187,6 +172,33 @@ fn local_api_base(state: &AppState) -> String {
state.comfyui.model_backend_url()
}
async fn external_comfyui_status(state: &AppState) -> Value {
json!({
"success": true,
"comfyui": external_comfyui_status_value(state).await,
})
}
async fn external_comfyui_status_value(state: &AppState) -> Value {
let status = state.comfyui.status_response().await;
let reachable = state.comfyui.health_check().await;
json!({
"status": if reachable { "running" } else { "unavailable" },
"pid": null,
"uptime_secs": 0,
"mode": "external",
"model_backend": status.model_backend,
"base_url": status.base_url,
"version": status.version,
"last_error": if reachable {
Value::Null
} else {
json!("外部 ComfyUI 不可达；内置 Docker ComfyUI 工作台已移除，请配置 comfyui.base_url 或使用 /workflows 工作流画布。")
},
"dockerManaged": false,
})
}
fn check_auth(headers: &HeaderMap, state: &AppState) -> AppResult<()> {
let expected = format!("Bearer {}", state.config.auth.auth_token);
src/routes/workflow_ui.rs
+2607
View File
File diff suppressed because it is too large Load Diff
src/routes/workspace.rs
+31 -6
View File
@@ -36,11 +36,23 @@ struct AssetQuery {
pub fn routes() -> Router<AppState> {
Router::new()
.route("/api/workspace/assets", get(list_admin_assets))
.route("/api/workspace/assets/{public_id}", get(get_admin_asset).delete(delete_admin_asset))
.route("/api/workspace/assets/{public_id}/file", get(get_admin_asset_file))
.route(
"/api/workspace/assets/{public_id}",
get(get_admin_asset).delete(delete_admin_asset),
)
.route(
"/api/workspace/assets/{public_id}/file",
get(get_admin_asset_file),
)
.route("/api/user/workspace/assets", get(list_user_assets))
.route("/api/user/workspace/assets/{public_id}", get(get_user_asset).delete(delete_user_asset))
.route("/api/user/workspace/assets/{public_id}/file", get(get_user_asset_file))
.route(
"/api/user/workspace/assets/{public_id}",
get(get_user_asset).delete(delete_user_asset),
)
.route(
"/api/user/workspace/assets/{public_id}/file",
get(get_user_asset_file),
)
}
fn bearer(headers: &HeaderMap) -> Option<&str> {
@@ -56,7 +68,11 @@ fn bearer(headers: &HeaderMap) -> Option<&str> {
.filter(|value| !value.is_empty())
}
fn resolve_caller(headers: &HeaderMap, state: &AppState, admin_required: bool) -> AppResult<WorkspaceCaller> {
fn resolve_caller(
headers: &HeaderMap,
state: &AppState,
admin_required: bool,
) -> AppResult<WorkspaceCaller> {
let token = bearer(headers)
.ok_or_else(|| AppError::Unauthorized("Missing Authorization header".into()))?;
@@ -79,7 +95,16 @@ fn resolve_caller(headers: &HeaderMap, state: &AppState, admin_required: bool) -
Err(AppError::Forbidden("Workspace access denied".into()))
}
fn normalize_query(query: &AssetQuery) -> (i64, i64, Option<String>, Option<String>, Option<String>, Option<i64>) {
fn normalize_query(
query: &AssetQuery,
) -> (
i64,
i64,
Option<String>,
Option<String>,
Option<String>,
Option<i64>,
) {
let page_size = query.page_size.unwrap_or(36).clamp(1, 120);
let page = query.page.unwrap_or(1).max(1);
let offset = (page - 1) * page_size;