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

Integrate ComfyUI agents and video workflows

Browse Source
This commit is contained in:
boos
2026-05-22 04:46:15 +08:00
parent 3c12738138
commit c1074ddf71
35 changed files with 2527 additions and 213 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/routes/admin/account_tokens.rs
+2 -1
View File
@@ -457,7 +457,8 @@ pub async fn get_account_token_value(
if is_masked_token(&token.token) || token.value_status.eq_ignore_ascii_case("masked") {
return Err(AppError::BadRequest(
"当前仅保存了脱敏令牌，无法展开/复制。请在站点重新生成并同步，或手动更新为完整令牌。".into(),
"当前仅保存了脱敏令牌，无法展开/复制。请在站点重新生成并同步，或手动更新为完整令牌。"
.into(),
));
}
src/routes/admin/accounts.rs
+40 -18
View File
@@ -170,15 +170,23 @@ pub async fn list_accounts(
.unwrap_or_default();
let health_by_account: std::collections::HashMap<
i64,
(String, Option<i64>, Option<String>, Option<i64>, Option<String>),
(
String,
Option<i64>,
Option<String>,
Option<i64>,
Option<String>,
),
> = latest_health_rows
.into_iter()
.map(|(account_id, status, http_status, error_message, latency_ms, created_at)| {
(
account_id,
(status, http_status, error_message, latency_ms, created_at),
)
})
.map(
|(account_id, status, http_status, error_message, latency_ms, created_at)| {
(
account_id,
(status, http_status, error_message, latency_ms, created_at),
)
},
)
.collect();
let enriched: Vec<Value> = accounts
@@ -200,7 +208,13 @@ pub async fn list_accounts(
fn account_response_value(
account: &Account,
site: Option<&Site>,
runtime_health: Option<&(String, Option<i64>, Option<String>, Option<i64>, Option<String>)>,
runtime_health: Option<&(
String,
Option<i64>,
Option<String>,
Option<i64>,
Option<String>,
)>,
) -> Value {
let credential_mode = resolve_account_credential_mode(account);
let mut row = serde_json::to_value(account).unwrap_or(json!({}));
@@ -220,7 +234,9 @@ fn account_response_value(
serde_json::to_value(site).unwrap_or(Value::Null),
);
}
let runtime_health_json = if account.status == "disabled" || site.is_some_and(|s| s.status == "disabled") {
let runtime_health_json = if account.status == "disabled"
|| site.is_some_and(|s| s.status == "disabled")
{
json!({
"state": "disabled",
"reason": "账号或站点已禁用",
@@ -230,7 +246,9 @@ fn account_response_value(
"state": "unhealthy",
"reason": "连接凭证已过期，请更新凭证",
})
} else if let Some((status, http_status, error_message, latency_ms, created_at)) = runtime_health {
} else if let Some((status, http_status, error_message, latency_ms, created_at)) =
runtime_health
{
let ok = status == "success";
json!({
"state": if ok { "healthy" } else { "unhealthy" },
@@ -597,7 +615,10 @@ pub async fn update_account(
// the previous value and decide whether routing needs to be rebuilt.
// `rebuild_routes` filters on `accounts.status = 'active'`, so any flip
// into or out of `active` changes which channels should exist.
let new_status = payload.status.clone().unwrap_or_else(|| existing.status.clone());
let new_status = payload
.status
.clone()
.unwrap_or_else(|| existing.status.clone());
let status_changed = new_status != existing.status;
sqlx::query(
@@ -747,12 +768,10 @@ pub async fn batch_accounts(
Err(e) => Err(format!("DB error: {}", e)),
}
}
"refreshBalance" => {
match mark_account_balance_refreshed(&state, *id).await {
Ok(_) => Ok(()),
Err(e) => Err(e),
}
}
"refreshBalance" => match mark_account_balance_refreshed(&state, *id).await {
Ok(_) => Ok(()),
Err(e) => Err(e),
},
other => Err(format!("不支持的批量操作: {}", other)),
};
@@ -1088,7 +1107,10 @@ pub async fn refresh_account_models(
.fetch_optional(&*state.db)
.await?
.ok_or_else(|| {
AppError::NotFound(format!("Site {} not found for account {}", account.site_id, id))
AppError::NotFound(format!(
"Site {} not found for account {}",
account.site_id, id
))
})?;
let base_url = resolve_site_api_base_url(&state, &site).await?;
src/routes/admin/downstream_keys.rs
+2 -6
View File
@@ -208,9 +208,7 @@ pub async fn update_downstream_key(
// The agent's outbound calls in Local mode bear this key — disabling it
// would silently break image / TTS / chat for every user.
if is_system_key(&existing.name, existing.user_id) && payload.enabled == Some(false) {
return Err(AppError::Forbidden(
"system key cannot be disabled".into(),
));
return Err(AppError::Forbidden("system key cannot be disabled".into()));
}
let supported_models: Option<String> = if let Some(ref models) = payload.supported_models {
@@ -270,9 +268,7 @@ pub async fn delete_downstream_key(
})?;
if is_system_key(&existing.name, existing.user_id) {
return Err(AppError::Forbidden(
"system key cannot be deleted".into(),
));
return Err(AppError::Forbidden("system key cannot be deleted".into()));
}
let affected = sqlx::query("DELETE FROM downstream_api_keys WHERE id = ?")
src/routes/admin/mod.rs
+5 -8
View File
@@ -71,8 +71,7 @@ pub fn routes(state: AppState) -> Router<AppState> {
)
.route(
"/api/sites/{id}/disabled-models",
get(sites::list_site_disabled_models)
.put(sites::update_site_disabled_models),
get(sites::list_site_disabled_models).put(sites::update_site_disabled_models),
);
let accounts_routes = Router::new()
@@ -299,12 +298,10 @@ mod tests {
.await
.unwrap();
sqlx::query(
"CREATE TABLE settings (key TEXT PRIMARY KEY, value TEXT)",
)
.execute(&pool)
.await
.unwrap();
sqlx::query("CREATE TABLE settings (key TEXT PRIMARY KEY, value TEXT)")
.execute(&pool)
.await
.unwrap();
sqlx::query("INSERT INTO settings (key, value) VALUES ('site.title', 'BoosAPI')")
.execute(&pool)
.await
src/routes/admin/model_candidates.rs
+9 -6
View File
@@ -147,12 +147,15 @@ pub async fn token_candidates(State(state): State<AppState>) -> AppResult<Json<V
if !seen.insert(row.account_id) {
continue;
}
models_without_token.entry(model_name).or_default().push(json!({
"accountId": row.account_id,
"username": row.username,
"siteId": row.site_id,
"siteName": row.site_name,
}));
models_without_token
.entry(model_name)
.or_default()
.push(json!({
"accountId": row.account_id,
"username": row.username,
"siteId": row.site_id,
"siteName": row.site_name,
}));
}
// Convert HashMap -> serde_json::Map so the response keys serialize
src/routes/admin/models_marketplace.rs
+64 -23
View File
@@ -64,9 +64,7 @@ pub async fn get_marketplace(
/// `AppState` and axum extractors so unit tests can drive it with an
/// in-memory SQLite pool (mirrors the `rebuild_routes_inner` pattern in
/// `services::route_rebuild`).
pub(crate) async fn marketplace_inner(
pool: &sqlx::Pool<sqlx::Sqlite>,
) -> AppResult<Value> {
pub(crate) async fn marketplace_inner(pool: &sqlx::Pool<sqlx::Sqlite>) -> AppResult<Value> {
// 1. Active (account, model) pairs. Filtering on accounts.status =
// 'active' here keeps the rest of the aggregation small.
let avail_rows: Vec<AvailRow> = sqlx::query_as(
@@ -137,10 +135,8 @@ pub(crate) async fn marketplace_inner(
.get(&acc.site_id)
.cloned()
.unwrap_or_default();
let acc_tokens: &[TokenRow] = tokens_by_account
.get(aid)
.map(Vec::as_slice)
.unwrap_or(&[]);
let acc_tokens: &[TokenRow] =
tokens_by_account.get(aid).map(Vec::as_slice).unwrap_or(&[]);
token_count += acc_tokens.len();
let token_json: Vec<Value> = acc_tokens
.iter()
@@ -253,26 +249,41 @@ mod tests {
let pool = fixture_pool().await;
sqlx::query("INSERT INTO sites (id, name) VALUES (1, 'site-a')")
.execute(&pool).await.unwrap();
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO accounts (id, site_id, username, balance, status) VALUES
(1, 1, 'alice', 10.0, 'active'),
(2, 1, 'bob', 5.0, 'disabled'),
(3, 1, 'carol', 7.0, 'active')",
).execute(&pool).await.unwrap();
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO model_availability (account_id, model_name, available) VALUES
(1, 'gpt-x', 1),
(2, 'gpt-x', 1),
(3, 'gpt-x', 0)",
).execute(&pool).await.unwrap();
)
.execute(&pool)
.await
.unwrap();
let value = marketplace_inner(&pool).await.expect("inner ok");
let models = value.get("models").and_then(Value::as_array).unwrap();
assert_eq!(models.len(), 1, "exactly one model should survive filtering");
assert_eq!(
models.len(),
1,
"exactly one model should survive filtering"
);
let m = &models[0];
assert_eq!(m["name"], "gpt-x");
assert_eq!(m["accountCount"], 1, "only account 1 (active + available) counts");
assert_eq!(
m["accountCount"], 1,
"only account 1 (active + available) counts"
);
let accounts = m["accounts"].as_array().unwrap();
assert_eq!(accounts.len(), 1);
assert_eq!(accounts[0]["username"], "alice");
@@ -288,18 +299,28 @@ mod tests {
async fn aggregates_accounts_and_tokens_per_model() {
let pool = fixture_pool().await;
sqlx::query("INSERT INTO sites (id, name) VALUES (1, 'openai-mirror'), (2, 'anthropic-mirror')")
.execute(&pool).await.unwrap();
sqlx::query(
"INSERT INTO sites (id, name) VALUES (1, 'openai-mirror'), (2, 'anthropic-mirror')",
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO accounts (id, site_id, username, balance, status) VALUES
(1, 1, 'alice', 100.5, 'active'),
(2, 2, 'bob', 50.0, 'active')",
).execute(&pool).await.unwrap();
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO model_availability (account_id, model_name, available) VALUES
(1, 'gpt-x', 1),
(2, 'gpt-x', 1)",
).execute(&pool).await.unwrap();
)
.execute(&pool)
.await
.unwrap();
// Account 1: two enabled tokens (one default), one disabled token (must
// be excluded). Account 2: one enabled non-default token.
sqlx::query(
@@ -308,7 +329,10 @@ mod tests {
(11, 1, 'backup', 1, 0),
(12, 1, 'retired', 0, 0),
(20, 2, 'only-one', 1, 0)",
).execute(&pool).await.unwrap();
)
.execute(&pool)
.await
.unwrap();
let value = marketplace_inner(&pool).await.expect("inner ok");
let models = value["models"].as_array().unwrap();
@@ -316,7 +340,10 @@ mod tests {
let m = &models[0];
assert_eq!(m["name"], "gpt-x");
assert_eq!(m["accountCount"], 2);
assert_eq!(m["tokenCount"], 3, "2 from account 1 + 1 from account 2; disabled token excluded");
assert_eq!(
m["tokenCount"], 3,
"2 from account 1 + 1 from account 2; disabled token excluded"
);
// Stage 1 placeholders.
assert!(m["avgLatency"].is_null());
assert!(m["successRate"].is_null());
@@ -330,12 +357,19 @@ mod tests {
// Find each account by id rather than relying on insertion order in
// the HashMap-driven account_ids vec.
let alice = accounts.iter().find(|a| a["id"] == 1).expect("alice present");
let alice = accounts
.iter()
.find(|a| a["id"] == 1)
.expect("alice present");
assert_eq!(alice["site"], "openai-mirror");
assert_eq!(alice["username"], "alice");
assert_eq!(alice["balance"], 100.5);
let alice_tokens = alice["tokens"].as_array().unwrap();
assert_eq!(alice_tokens.len(), 2, "disabled 'retired' token must not be exposed");
assert_eq!(
alice_tokens.len(),
2,
"disabled 'retired' token must not be exposed"
);
let primary = alice_tokens.iter().find(|t| t["id"] == 10).unwrap();
assert_eq!(primary["name"], "primary");
assert_eq!(primary["isDefault"], true);
@@ -356,16 +390,23 @@ mod tests {
let pool = fixture_pool().await;
sqlx::query("INSERT INTO sites (id, name) VALUES (1, 's')")
.execute(&pool).await.unwrap();
.execute(&pool)
.await
.unwrap();
sqlx::query("INSERT INTO accounts (id, site_id, status) VALUES (1, 1, 'active')")
.execute(&pool).await.unwrap();
.execute(&pool)
.await
.unwrap();
// Insert in scrambled order; expect output sorted lexicographically.
sqlx::query(
"INSERT INTO model_availability (account_id, model_name, available) VALUES
(1, 'gpt-z', 1),
(1, 'gpt-a', 1),
(1, 'gpt-m', 1)",
).execute(&pool).await.unwrap();
)
.execute(&pool)
.await
.unwrap();
let value = marketplace_inner(&pool).await.expect("inner ok");
let names: Vec<&str> = value["models"]
src/routes/admin/sites.rs
+19 -16
View File
@@ -369,20 +369,20 @@ pub async fn list_site_disabled_models(
) -> AppResult<Json<Value>> {
// Validate the site exists so we return a clean 404 instead of an empty list
// for unknown ids.
let exists: Option<i64> =
sqlx::query_scalar("SELECT id FROM sites WHERE id = ?")
.bind(id)
.fetch_optional(&*state.db)
.await?;
let exists: Option<i64> = sqlx::query_scalar("SELECT id FROM sites WHERE id = ?")
.bind(id)
.fetch_optional(&*state.db)
.await?;
if exists.is_none() {
return Err(AppError::NotFound("site not found".into()));
}
let models: Vec<String> =
sqlx::query_scalar("SELECT model_name FROM site_disabled_models WHERE site_id = ? ORDER BY model_name")
.bind(id)
.fetch_all(&*state.db)
.await?;
let models: Vec<String> = sqlx::query_scalar(
"SELECT model_name FROM site_disabled_models WHERE site_id = ? ORDER BY model_name",
)
.bind(id)
.fetch_all(&*state.db)
.await?;
Ok(Json(json!({ "models": models })))
}
@@ -393,11 +393,10 @@ pub async fn update_site_disabled_models(
Path(id): Path<i64>,
Json(payload): Json<UpdateDisabledModelsRequest>,
) -> AppResult<Json<Value>> {
let exists: Option<i64> =
sqlx::query_scalar("SELECT id FROM sites WHERE id = ?")
.bind(id)
.fetch_optional(&*state.db)
.await?;
let exists: Option<i64> = sqlx::query_scalar("SELECT id FROM sites WHERE id = ?")
.bind(id)
.fetch_optional(&*state.db)
.await?;
if exists.is_none() {
return Err(AppError::NotFound("site not found".into()));
}
@@ -428,7 +427,11 @@ pub async fn update_site_disabled_models(
}
tx.commit().await?;
info!(site_id = id, count = normalised.len(), "updated site disabled models");
info!(
site_id = id,
count = normalised.len(),
"updated site disabled models"
);
Ok(Json(json!({ "models": normalised })))
}
src/routes/comfy_proxy.rs
+6 -5
View File
@@ -127,10 +127,7 @@ fn comfy_session_cookie(token: &str) -> String {
// ── HTTP forwarder ─────────────────────────────────────────────────────
async fn http_handler(
State(state): State<AppState>,
req: Request<Body>,
) -> Response<Body> {
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()))
@@ -159,7 +156,11 @@ async fn http_handler(
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 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 {
src/routes/comfyui_agent.rs
+4 -1
View File
@@ -618,7 +618,10 @@ fn agent_message_to_upstream(message: AgentMessage) -> Value {
.map(str::trim)
.filter(|value| value.starts_with("data:image/"))
{
if !content_parts.iter().any(|part: &Value| part["type"] == "text") {
if !content_parts
.iter()
.any(|part: &Value| part["type"] == "text")
{
content_parts.push(json!({
"type": "text",
"text": text,
src/routes/mod.rs
+2
View File
@@ -22,6 +22,7 @@ mod health;
mod monitor;
mod oauth;
mod organizations;
mod picture_agent;
mod proxy;
mod search;
mod site_announcements;
@@ -47,6 +48,7 @@ pub fn build_router(state: AppState) -> Router {
.route("/healthz", get(health::health_check))
.merge(admin::routes(state.clone()))
.merge(proxy::routes())
.merge(picture_agent::routes())
.merge(comfyui_agent::routes())
.merge(comfyui::routes());
src/routes/picture_agent.rs
+417
View File
@@ -0,0 +1,417 @@
use crate::{
auth::jwt::validate_jwt,
db::models::downstream_api_key::DownstreamApiKey,
error::{AppError, AppResult},
AppState,
};
use axum::{
body::Body,
extract::{Multipart, Query, State},
http::{header, HeaderMap, Response, StatusCode},
response::IntoResponse,
routing::{get, post},
Json, Router,
};
use reqwest::Client;
use serde::Deserialize;
use serde_json::{json, Value};
use std::time::Duration;
const TEMPLATE_GALLERY_URL: &str =
"https://weiyyin.github.io/linux-do-gpt-image/linux_do_2036927_prompts_images.filtered.json";
#[derive(Debug, Clone)]
struct PictureAgentAuth {
user_id: Option<i64>,
role: String,
outbound_key: String,
key_name: String,
}
pub fn routes() -> Router<AppState> {
Router::new()
.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/images/generations",
post(images_generations),
)
.route("/api/picture-agent/images/edits", post(images_edits))
}
#[derive(Debug, Deserialize)]
struct TemplateQuery {
limit: Option<usize>,
}
#[derive(Debug, Deserialize)]
struct GalleryItem {
post_number: Option<i64>,
username: Option<String>,
prompt: Option<String>,
images: Option<Vec<String>>,
image_host_url: Option<String>,
}
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> {
let auth_header = headers
.get(header::AUTHORIZATION)
.and_then(|value| value.to_str().ok())
.unwrap_or_default();
auth_header
.strip_prefix("Bearer ")
.or_else(|| auth_header.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_picture_agent_auth(
headers: &HeaderMap,
state: &AppState,
) -> Result<PictureAgentAuth, 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(PictureAgentAuth {
user_id: None,
role: "admin".to_string(),
outbound_key: state.config.auth.proxy_token.clone(),
key_name: "global-proxy".to_string(),
});
}
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(PictureAgentAuth {
user_id: Some(user_id),
role: claims.role,
outbound_key: key.key,
key_name: key.name,
}),
Err(_) if claims.role == "admin" => Ok(PictureAgentAuth {
user_id: Some(user_id),
role: claims.role,
outbound_key: state.config.auth.proxy_token.clone(),
key_name: "global-proxy".to_string(),
}),
Err(message) => Err(AppError::Forbidden(message)),
}
}
fn is_exact_model_pattern(value: &str) -> bool {
let normalized = value.trim();
!normalized.is_empty()
&& !normalized.to_ascii_lowercase().starts_with("re:")
&& !normalized.contains('*')
&& !normalized.contains('?')
}
async fn status_handler(
State(state): State<AppState>,
headers: HeaderMap,
) -> AppResult<Json<Value>> {
let auth = resolve_picture_agent_auth(&headers, &state).await?;
Ok(Json(json!({
"success": true,
"mode": "local",
"role": auth.role,
"userId": auth.user_id,
"keyName": auth.key_name,
"upstreamLabel": "BoosAPI 本地图片路由",
})))
}
async fn models_handler(
State(state): State<AppState>,
headers: HeaderMap,
) -> AppResult<Json<Value>> {
let _auth = resolve_picture_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?;
let mut seen = std::collections::HashSet::new();
let models: Vec<String> = rows
.into_iter()
.map(|value| value.trim().to_string())
.filter(|value| is_exact_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()),
})))
}
fn prompt_title(prompt: &str, post_number: Option<i64>) -> String {
let clean = prompt
.split_whitespace()
.collect::<Vec<_>>()
.join(" ")
.trim()
.to_string();
let mut title = clean.chars().take(18).collect::<String>();
if title.trim().is_empty() {
title = post_number
.map(|value| format!("#{value}"))
.unwrap_or_else(|| "参考模板".to_string());
}
if clean.chars().count() > 18 {
title.push('…');
}
title
}
fn infer_template_tag(prompt: &str) -> &'static str {
let lower = prompt.to_ascii_lowercase();
if prompt.contains("海报") || lower.contains("poster") {
"海报"
} else if prompt.contains("漫画") || lower.contains("comic") {
"漫画"
} else if prompt.contains("图鉴") || prompt.contains("信息图") || lower.contains("infographic")
{
"图鉴"
} else if prompt.contains("贴纸") || lower.contains("sticker") || lower.contains("mascot") {
"贴纸"
} else if prompt.contains("分镜") || lower.contains("storyboard") {
"分镜"
} else {
"参考"
}
}
fn infer_template_size(prompt: &str) -> &'static str {
let lower = prompt.to_ascii_lowercase();
if lower.contains("16:9") || prompt.contains("横版") {
"1536x1024"
} else if lower.contains("9:16") || lower.contains("--ar 9:16") || prompt.contains("竖版") {
"1024x1536"
} else if lower.contains("1:1") || prompt.contains("方图") {
"1024x1024"
} else {
"auto"
}
}
async fn templates_handler(
State(state): State<AppState>,
headers: HeaderMap,
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()
.get(TEMPLATE_GALLERY_URL)
.send()
.await
.map_err(|err| AppError::UpstreamError(format!("failed to fetch templates: {err}")))?
.error_for_status()
.map_err(|err| AppError::UpstreamError(format!("template gallery returned {err}")))?
.json::<Vec<GalleryItem>>()
.await
.map_err(|err| AppError::UpstreamError(format!("failed to parse templates: {err}")))?;
let templates: Vec<Value> = items
.into_iter()
.filter_map(|item| {
let prompt = item.prompt?.trim().to_string();
if prompt.chars().count() < 20 {
return None;
}
let image = item
.images
.unwrap_or_default()
.into_iter()
.find(|value| value.starts_with("http"))
.or_else(|| {
item.image_host_url
.filter(|value| value.starts_with("http"))
})?;
let post_number = item.post_number;
Some(json!({
"id": post_number
.map(|value| format!("linux-do-{value}"))
.unwrap_or_else(|| format!("linux-do-{}", prompt_title(&prompt, None))),
"title": prompt_title(&prompt, post_number),
"tag": infer_template_tag(&prompt),
"size": infer_template_size(&prompt),
"image": image,
"prompt": prompt,
"source": {
"name": "Linux Do GPT Image Gallery",
"url": "https://weiyyin.github.io/linux-do-gpt-image/prompt_gallery.html",
"postNumber": post_number,
"username": item.username,
},
}))
})
.take(limit)
.collect();
Ok(Json(json!({
"success": true,
"sourceUrl": "https://weiyyin.github.io/linux-do-gpt-image/prompt_gallery.html",
"templates": templates,
})))
}
fn proxy_client() -> Client {
Client::builder()
.timeout(Duration::from_secs(180))
.build()
.unwrap_or_else(|_| Client::new())
}
async fn upstream_response_to_axum(response: reqwest::Response) -> AppResult<Response<Body>> {
let status = StatusCode::from_u16(response.status().as_u16())
.unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
let content_type = response
.headers()
.get(header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.map(str::to_string)
.unwrap_or_else(|| "application/json; charset=utf-8".to_string());
let bytes = response
.bytes()
.await
.map_err(|err| AppError::UpstreamError(err.to_string()))?;
Ok(Response::builder()
.status(status)
.header(header::CONTENT_TYPE, content_type)
.body(Body::from(bytes))
.unwrap_or_else(|_| StatusCode::INTERNAL_SERVER_ERROR.into_response()))
}
async fn images_generations(
State(state): State<AppState>,
headers: HeaderMap,
Json(body): Json<Value>,
) -> AppResult<Response<Body>> {
let auth = resolve_picture_agent_auth(&headers, &state).await?;
let url = format!("{}/images/generations", local_v1_base(&state));
let response = proxy_client()
.post(url)
.bearer_auth(auth.outbound_key)
.json(&body)
.send()
.await
.map_err(|err| AppError::UpstreamError(err.to_string()))?;
upstream_response_to_axum(response).await
}
async fn images_edits(
State(state): State<AppState>,
headers: HeaderMap,
mut multipart: Multipart,
) -> AppResult<Response<Body>> {
let auth = resolve_picture_agent_auth(&headers, &state).await?;
let mut form = reqwest::multipart::Form::new();
while let Some(field) = multipart
.next_field()
.await
.map_err(|err| AppError::BadRequest(format!("invalid multipart body: {err}")))?
{
let name = field
.name()
.ok_or_else(|| AppError::BadRequest("multipart field missing name".into()))?
.to_string();
let filename = field.file_name().map(str::to_string);
let content_type = field.content_type().map(str::to_string);
let is_file = filename.is_some()
|| content_type
.as_deref()
.map(|ct| !ct.starts_with("text/") && ct != "application/x-www-form-urlencoded")
.unwrap_or(false);
if is_file {
let bytes = field
.bytes()
.await
.map_err(|err| AppError::BadRequest(format!("failed to read upload: {err}")))?;
let mut part = reqwest::multipart::Part::bytes(bytes.to_vec());
if let Some(filename) = filename {
part = part.file_name(filename);
}
if let Some(content_type) = content_type {
part = part
.mime_str(&content_type)
.map_err(|err| AppError::BadRequest(format!("invalid content type: {err}")))?;
}
form = form.part(name, part);
} else {
let text = field
.text()
.await
.map_err(|err| AppError::BadRequest(format!("failed to read field: {err}")))?;
form = form.text(name, text);
}
}
let url = format!("{}/images/edits", local_v1_base(&state));
let response = proxy_client()
.post(url)
.bearer_auth(auth.outbound_key)
.multipart(form)
.send()
.await
.map_err(|err| AppError::UpstreamError(err.to_string()))?;
upstream_response_to_axum(response).await
}
src/routes/proxy/mod.rs
+112 -23
View File
@@ -17,7 +17,7 @@ use crate::{
AppState,
};
use axum::{
extract::{Multipart, State},
extract::{Multipart, Query, State},
http::HeaderMap,
response::{IntoResponse, Response, Sse},
routing::{get, post},
@@ -60,6 +60,7 @@ pub fn routes() -> Router<AppState> {
.route("/v1/images/generations", post(images_generations))
.route("/v1/images/edits", post(images_edits))
.route("/v1/videos", post(videos_handler))
.route("/v1/videos/{id}", get(video_task_handler))
.route("/v1/models", get(list_models))
// ── ComfyUI shim aliases ──────────────────────────────────────
// ComfyUI's OpenAI nodes hit /proxy/openai/... paths (see
@@ -115,8 +116,7 @@ async fn responses_handler(
.unwrap_or(false);
if stream {
let chunks =
execute_responses_stream_with_conductor(state, key_id, model, body).await?;
let chunks = execute_responses_stream_with_conductor(state, key_id, model, body).await?;
let sse =
Sse::new(stream::iter(chunks).map(|chunk| {
Ok::<_, Infallible>(axum::response::sse::Event::default().data(chunk))
@@ -276,8 +276,7 @@ async fn audio_translations(
let route = resolve_route(&state, &model).await?;
let result =
executors::execute_proxy_request(state.clone(), route, "/audio/translations", body)
.await?;
executors::execute_proxy_request(state.clone(), route, "/audio/translations", body).await?;
record_billing(&state, key_id, &model, 0, 0, 0.0).await;
Ok(Json(result))
}
@@ -299,7 +298,33 @@ async fn videos_handler(
.to_string();
let route = resolve_route(&state, &model).await?;
let result = executors::execute_proxy_request(state.clone(), route, "/videos", body).await?;
let result = executors::execute_video_generation(state.clone(), route, body).await?;
record_billing(&state, key_id, &model, 0, 0, 0.0).await;
Ok(Json(result))
}
/// GET /v1/videos/{id} — inspect video generation task.
async fn video_task_handler(
State(state): State<AppState>,
headers: HeaderMap,
axum::extract::Path(id): axum::extract::Path<String>,
Query(query): Query<std::collections::HashMap<String, String>>,
) -> AppResult<Json<Value>> {
info!("GET /v1/videos/{id}");
let key_id = check_auth(&state, &headers).await?;
let model = query
.get("model")
.map(String::as_str)
.unwrap_or("default")
.to_string();
let route = resolve_route(&state, &model).await?;
let suffix = if route.base_url.contains("ark.cn-beijing.volces.com") {
format!("/contents/generations/tasks/{id}")
} else {
format!("/videos/{id}")
};
let result = executors::execute_proxy_get(route, &suffix).await?;
record_billing(&state, key_id, &model, 0, 0, 0.0).await;
Ok(Json(result))
}
@@ -426,7 +451,22 @@ async fn list_models(State(state): State<AppState>, headers: HeaderMap) -> AppRe
let _ = check_auth(&state, &headers).await?;
let rows: Vec<(String, String)> = sqlx::query_as(
"SELECT tr.model_pattern, COALESCE(tr.model_mapping, tr.model_pattern) FROM token_routes tr WHERE tr.enabled = 1 ORDER BY tr.id",
"SELECT DISTINCT
tr.model_pattern,
COALESCE(NULLIF(TRIM(tr.model_mapping), ''), tr.model_pattern)
FROM token_routes tr
JOIN route_channels rc ON rc.route_id = tr.id
JOIN accounts a ON a.id = rc.account_id
JOIN sites s ON s.id = a.site_id
LEFT JOIN account_tokens at ON at.id = rc.token_id AND at.enabled = 1
WHERE tr.enabled = 1
AND rc.enabled = 1
AND a.status = 'active'
AND s.status = 'active'
AND TRIM(COALESCE(s.url, '')) <> ''
AND (rc.cooldown_until IS NULL OR rc.cooldown_until <= datetime('now'))
AND (rc.token_id IS NULL OR at.id IS NOT NULL)
ORDER BY tr.id",
)
.fetch_all(&*state.db)
.await?;
@@ -671,6 +711,17 @@ fn build_conductor(state: &AppState) -> ProxyConductor {
}
}
fn upstream_status_from_error(err: &AppError) -> Option<u16> {
let message = err.to_string();
let marker = "Upstream returned ";
let start = message.find(marker)? + marker.len();
message
.get(start..)?
.split_whitespace()
.next()
.and_then(|value| value.parse::<u16>().ok())
}
async fn execute_chat_completions_with_conductor(
state: AppState,
key_id: Option<i64>,
@@ -703,7 +754,7 @@ async fn execute_chat_completions_with_conductor(
},
Err(err) => AttemptResult::Failure {
action: crate::proxy::conductor::AttemptFailureAction::Failover,
status: None,
status: upstream_status_from_error(&err),
raw_error_text: Some(err.to_string()),
},
}
@@ -734,7 +785,15 @@ async fn execute_chat_completions_with_conductor(
false,
)
.await;
record_billing(&billing_state, key_id, &billing_model, prompt, completion, 0.0).await;
record_billing(
&billing_state,
key_id,
&billing_model,
prompt,
completion,
0.0,
)
.await;
Ok(response)
}
ExecuteResult::Failure {
@@ -797,7 +856,7 @@ async fn execute_chat_completions_stream_with_conductor(
},
Err(err) => AttemptResult::Failure {
action: crate::proxy::conductor::AttemptFailureAction::Failover,
status: None,
status: upstream_status_from_error(&err),
raw_error_text: Some(err.to_string()),
},
}
@@ -814,14 +873,20 @@ async fn execute_chat_completions_stream_with_conductor(
response,
attempts,
} => match response {
Value::Array(chunks) => (selected, attempts, chunks
.into_iter()
.map(|chunk| {
chunk.as_str().map(str::to_string).ok_or_else(|| {
AppError::UpstreamError("Invalid streaming response from conductor".into())
})
})
.collect::<AppResult<Vec<String>>>()?),
Value::Array(chunks) => (
selected,
attempts,
chunks
.into_iter()
.map(|chunk| {
chunk.as_str().map(str::to_string).ok_or_else(|| {
AppError::UpstreamError(
"Invalid streaming response from conductor".into(),
)
})
})
.collect::<AppResult<Vec<String>>>()?,
),
_ => {
return Err(AppError::UpstreamError(
"Invalid streaming response from conductor".into(),
@@ -888,7 +953,15 @@ async fn execute_chat_completions_stream_with_conductor(
true,
)
.await;
record_billing(&billing_state, key_id, &billing_model, prompt, completion, 0.0).await;
record_billing(
&billing_state,
key_id,
&billing_model,
prompt,
completion,
0.0,
)
.await;
Ok(chunks)
}
@@ -941,7 +1014,7 @@ async fn execute_responses_with_conductor(
},
Err(err) => AttemptResult::Failure {
action: crate::proxy::conductor::AttemptFailureAction::Failover,
status: None,
status: upstream_status_from_error(&err),
raw_error_text: Some(err.to_string()),
},
}
@@ -954,7 +1027,15 @@ async fn execute_responses_with_conductor(
let response = conductor_result_to_app_result(result)?;
let (prompt, completion) = extract_usage(&response);
record_billing(&billing_state, key_id, &billing_model, prompt, completion, 0.0).await;
record_billing(
&billing_state,
key_id,
&billing_model,
prompt,
completion,
0.0,
)
.await;
Ok(response)
}
@@ -988,7 +1069,7 @@ async fn execute_responses_stream_with_conductor(
},
Err(err) => AttemptResult::Failure {
action: crate::proxy::conductor::AttemptFailureAction::Failover,
status: None,
status: upstream_status_from_error(&err),
raw_error_text: Some(err.to_string()),
},
}
@@ -1037,7 +1118,15 @@ async fn execute_responses_stream_with_conductor(
}
})
.unwrap_or((0, 0));
record_billing(&billing_state, key_id, &billing_model, prompt, completion, 0.0).await;
record_billing(
&billing_state,
key_id,
&billing_model,
prompt,
completion,
0.0,
)
.await;
Ok(chunks)
}
src/routes/stats.rs
+8 -8
View File
@@ -67,10 +67,9 @@ async fn dashboard(
)
.fetch_one(&*state.db)
.await?;
let total_balance: f64 =
sqlx::query_scalar("SELECT COALESCE(SUM(balance), 0.0) FROM accounts")
.fetch_one(&*state.db)
.await?;
let total_balance: f64 = sqlx::query_scalar("SELECT COALESCE(SUM(balance), 0.0) FROM accounts")
.fetch_one(&*state.db)
.await?;
let total_used: f64 =
sqlx::query_scalar("SELECT COALESCE(SUM(balance_used), 0.0) FROM accounts")
.fetch_one(&*state.db)
@@ -93,10 +92,11 @@ async fn dashboard(
sqlx::query_scalar("SELECT COUNT(*) FROM accounts WHERE status = 'active'")
.fetch_one(&*state.db)
.await?;
let today_checkin_total: i64 =
sqlx::query_scalar("SELECT COUNT(*) FROM checkin_logs WHERE date(created_at) = date('now')")
.fetch_one(&*state.db)
.await?;
let today_checkin_total: i64 = sqlx::query_scalar(
"SELECT COUNT(*) FROM checkin_logs WHERE date(created_at) = date('now')",
)
.fetch_one(&*state.db)
.await?;
let today_checkin_success: i64 = sqlx::query_scalar(
"SELECT COUNT(*) FROM checkin_logs
WHERE status = 'success' AND date(created_at) = date('now')",
src/routes/video_agent.rs
+16 -8
View File
@@ -120,9 +120,7 @@ fn resolve_user_id(headers: &HeaderMap, state: &AppState) -> AppResult<i64> {
.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())
})?;
.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);
@@ -142,9 +140,7 @@ fn resolve_user_id(headers: &HeaderMap, state: &AppState) -> AppResult<i64> {
/// 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::AtCapacity(_) => AppError::ServiceUnavailable("comfyui pool at capacity".into()),
PoolError::StartTimeout(_) => {
AppError::GatewayTimeout("comfyui container start timeout".into())
}
@@ -211,7 +207,15 @@ async fn pool_status_json(state: &AppState, user_id: i64) -> Value {
.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('/'));
let model_backend = format!(
"{}/v1",
state
.comfyui_pool
.config()
.boosapi_base
.trim_end_matches("/v1")
.trim_end_matches('/')
);
json!({
"status": status,
@@ -288,7 +292,11 @@ async fn restart_comfyui(
// 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('/');
let raw = state
.comfyui_pool
.config()
.boosapi_base
.trim_end_matches('/');
if raw.ends_with("/v1") {
raw.to_string()
} else {