feat: refine BoosAPI console and user workspace

This commit is contained in:
boos
2026-06-09 16:48:18 +08:00
parent aaba6f152e
commit 59106e6a96
284 changed files with 59886 additions and 5594 deletions
+255 -11
View File
@@ -13,6 +13,8 @@ use serde_json::{json, Value};
use std::time::Duration;
use tracing::info;
const RUNTIME_HEALTH_DEGRADED_LATENCY_MS: i64 = 5_000;
#[derive(Debug, Deserialize)]
pub struct ListAccountsQuery {
pub site_id: Option<i64>,
@@ -151,10 +153,23 @@ pub async fn list_accounts(
let site_by_id: std::collections::HashMap<i64, &Site> =
sites.iter().map(|s| (s.id, s)).collect();
let latest_health_rows: Vec<(i64, String, Option<i64>, Option<String>, Option<i64>, Option<String>)> =
let latest_health_rows: Vec<(
i64,
String,
Option<i64>,
Option<String>,
Option<i64>,
Option<String>,
Option<String>,
Option<String>,
Option<i64>,
Option<i64>,
i64,
)> =
sqlx::query_as(
r#"
SELECT l.account_id, l.status, l.http_status, l.error_message, l.latency_ms, l.created_at
SELECT l.account_id, l.status, l.http_status, l.error_message, l.latency_ms, l.created_at,
l.model_requested, l.model_actual, l.route_id, l.channel_id, l.id
FROM proxy_logs l
INNER JOIN (
SELECT account_id, MAX(id) AS id
@@ -176,14 +191,42 @@ pub async fn list_accounts(
Option<String>,
Option<i64>,
Option<String>,
Option<String>,
Option<String>,
Option<i64>,
Option<i64>,
i64,
),
> = 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,
model_requested,
model_actual,
route_id,
channel_id,
log_id,
)| {
(
account_id,
(status, http_status, error_message, latency_ms, created_at),
(
status,
http_status,
error_message,
latency_ms,
created_at,
model_requested,
model_actual,
route_id,
channel_id,
log_id,
),
)
},
)
@@ -214,6 +257,11 @@ fn account_response_value(
Option<String>,
Option<i64>,
Option<String>,
Option<String>,
Option<String>,
Option<i64>,
Option<i64>,
i64,
)>,
) -> Value {
let credential_mode = resolve_account_credential_mode(account);
@@ -246,16 +294,46 @@ fn account_response_value(
"state": "unhealthy",
"reason": "连接凭证已过期,请更新凭证",
})
} else if let Some((status, http_status, error_message, latency_ms, created_at)) =
} else if let Some((
status,
http_status,
error_message,
latency_ms,
created_at,
model_requested,
model_actual,
route_id,
channel_id,
log_id,
)) =
runtime_health
{
let ok = status == "success";
let is_degraded_latency =
ok && latency_ms.is_some_and(|ms| ms >= RUNTIME_HEALTH_DEGRADED_LATENCY_MS);
let model = model_actual
.as_ref()
.filter(|value| !value.trim().is_empty())
.or_else(|| model_requested.as_ref().filter(|value| !value.trim().is_empty()));
json!({
"state": if ok { "healthy" } else { "unhealthy" },
"state": if !ok {
"unhealthy"
} else if is_degraded_latency {
"degraded"
} else {
"healthy"
},
"reason": if ok {
match latency_ms {
Some(ms) => format!("最近代理请求成功,用时 {}ms", ms),
None => "最近代理请求成功".to_string(),
if is_degraded_latency {
match latency_ms {
Some(ms) => format!("最近代理请求成功但延迟偏高,用时 {}ms", ms),
None => "最近代理请求成功但延迟偏高".to_string(),
}
} else {
match latency_ms {
Some(ms) => format!("最近代理请求成功,用时 {}ms", ms),
None => "最近代理请求成功".to_string(),
}
}
} else {
error_message
@@ -265,11 +343,17 @@ fn account_response_value(
"latencyMs": latency_ms,
"httpStatus": http_status,
"checkedAt": created_at,
"modelRequested": model_requested,
"modelActual": model_actual,
"recentModel": model,
"routeId": route_id,
"channelId": channel_id,
"logId": log_id,
})
} else if credential_mode == "apikey" || account.status == "active" {
json!({
"state": "degraded",
"reason": "尚无代理请求记录,等待首次调用后更新",
"state": "unknown",
"reason": "尚无代理请求记录,首次调用后自动更新",
})
} else {
json!({
@@ -1266,3 +1350,163 @@ pub async fn refresh_account_models(
}
})))
}
#[cfg(test)]
mod tests {
use super::*;
fn account_fixture(status: &str, credential_mode: &str) -> Account {
Account {
id: 7,
site_id: 3,
username: Some("tester".to_string()),
access_token: "sk-test".to_string(),
api_token: None,
balance: Some(0.0),
balance_used: Some(0.0),
quota: None,
unit_cost: None,
value_score: None,
status: status.to_string(),
checkin_enabled: Some(false),
last_checkin_at: None,
last_balance_refresh: None,
extra_config: Some(json!({ "credentialMode": credential_mode }).to_string()),
is_pinned: Some(false),
sort_order: Some(0),
oauth_provider: None,
oauth_account_key: None,
oauth_project_id: None,
created_at: None,
updated_at: None,
}
}
fn site_fixture(status: &str) -> Site {
Site {
id: 3,
name: "Demo".to_string(),
platform: "openai".to_string(),
url: "https://example.com".to_string(),
status: status.to_string(),
api_key: None,
proxy_url: None,
use_system_proxy: None,
custom_headers: None,
external_checkin_url: None,
global_weight: None,
post_refresh_probe_enabled: None,
post_refresh_probe_model: None,
post_refresh_probe_scope: None,
post_refresh_probe_latency_threshold_ms: None,
is_pinned: Some(false),
sort_order: Some(0),
created_at: None,
updated_at: None,
}
}
#[test]
fn active_apikey_without_proxy_logs_is_pending_detection_not_degraded() {
let account = account_fixture("active", "apikey");
let site = site_fixture("active");
let value = account_response_value(&account, Some(&site), None);
let runtime_health = value
.get("runtimeHealth")
.expect("runtimeHealth should be present");
assert_eq!(
runtime_health.get("state").and_then(Value::as_str),
Some("unknown")
);
assert_eq!(
runtime_health.get("reason").and_then(Value::as_str),
Some("尚无代理请求记录,首次调用后自动更新")
);
}
#[test]
fn runtime_health_includes_recent_proxy_context_for_diagnostics() {
let account = account_fixture("active", "apikey");
let site = site_fixture("active");
let latest = (
"failed".to_string(),
Some(404),
Some("Not found".to_string()),
Some(1820),
Some("2026-06-08 10:00:00".to_string()),
Some("gpt-4o-mini".to_string()),
Some("gpt-4o-mini-2026".to_string()),
Some(77),
Some(88),
101,
);
let value = account_response_value(&account, Some(&site), Some(&latest));
let runtime_health = value
.get("runtimeHealth")
.expect("runtimeHealth should be present");
assert_eq!(
runtime_health.get("state").and_then(Value::as_str),
Some("unhealthy")
);
assert_eq!(
runtime_health.get("recentModel").and_then(Value::as_str),
Some("gpt-4o-mini-2026")
);
assert_eq!(
runtime_health.get("modelRequested").and_then(Value::as_str),
Some("gpt-4o-mini")
);
assert_eq!(
runtime_health.get("routeId").and_then(Value::as_i64),
Some(77)
);
assert_eq!(
runtime_health.get("channelId").and_then(Value::as_i64),
Some(88)
);
assert_eq!(
runtime_health.get("logId").and_then(Value::as_i64),
Some(101)
);
}
#[test]
fn successful_but_slow_proxy_log_marks_runtime_health_degraded() {
let account = account_fixture("active", "apikey");
let site = site_fixture("active");
let latest = (
"success".to_string(),
Some(200),
None,
Some(7928),
Some("2026-06-08 10:00:00".to_string()),
Some("gpt-4o-mini".to_string()),
Some("gpt-4o-mini".to_string()),
Some(77),
Some(88),
101,
);
let value = account_response_value(&account, Some(&site), Some(&latest));
let runtime_health = value
.get("runtimeHealth")
.expect("runtimeHealth should be present");
assert_eq!(
runtime_health.get("state").and_then(Value::as_str),
Some("degraded")
);
assert_eq!(
runtime_health.get("reason").and_then(Value::as_str),
Some("最近代理请求成功但延迟偏高,用时 7928ms")
);
assert_eq!(
runtime_health.get("latencyMs").and_then(Value::as_i64),
Some(7928)
);
}
}
+1669 -54
View File
File diff suppressed because it is too large Load Diff
+63 -15
View File
@@ -24,7 +24,7 @@ mod sites;
mod tokens;
mod users;
/// Middleware to require admin authentication via config token or admin JWT.
/// Middleware to require admin authentication via the configured admin API token.
async fn require_admin_auth(
State(state): State<AppState>,
headers: HeaderMap,
@@ -37,20 +37,8 @@ async fn require_admin_auth(
.and_then(|v| v.strip_prefix("Bearer "))
.ok_or_else(|| AppError::Unauthorized("Missing or invalid Authorization header".into()))?;
if auth_header == state.config.auth.auth_token {
return Ok(next.run(request).await);
}
match crate::auth::jwt::validate_jwt(&state.config.auth.jwt_secret, auth_header) {
Ok(claims) if claims.role == "admin" => {}
Ok(_) => {
return Err(AppError::Forbidden(
"Admin account is required for this route".into(),
));
}
Err(_) => {
return Err(AppError::Unauthorized("Invalid admin token".into()));
}
if auth_header != state.config.auth.auth_token {
return Err(AppError::Unauthorized("Invalid admin token".into()));
}
Ok(next.run(request).await)
@@ -73,6 +61,15 @@ 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),
)
.route(
"/api/sites/{id}/available-models",
get(sites::list_site_available_models),
)
.route("/api/sites/{id}/probe-now", post(sites::probe_site_now))
.route(
"/api/sites/{id}/probe-stream",
get(sites::probe_site_stream),
);
let accounts_routes = Router::new()
@@ -289,9 +286,15 @@ async fn update_setting(
#[cfg(test)]
mod tests {
use super::*;
use crate::auth::jwt::generate_jwt;
use crate::config::{AppConfig, AuthConfig, DbConfig, ProxyConfig, ServerConfig};
use crate::services::builtin_keys::BuiltinKeys;
use axum::{
body::Body,
http::{Request, StatusCode},
};
use sqlx::sqlite::SqlitePoolOptions;
use tower::ServiceExt as _;
async fn settings_state() -> AppState {
let pool = SqlitePoolOptions::new()
@@ -414,4 +417,49 @@ mod tests {
.expect("update ok");
assert_eq!(resp.0["data"]["value"], "Renamed");
}
#[tokio::test]
async fn admin_routes_accept_config_token_and_reject_admin_jwt() {
let state = settings_state().await;
let app = routes(state.clone()).with_state(state.clone());
let token_response = app
.clone()
.oneshot(
Request::builder()
.uri("/api/settings/auth/info")
.header("Authorization", "Bearer auth")
.body(Body::empty())
.unwrap(),
)
.await
.expect("token request");
assert_eq!(token_response.status(), StatusCode::OK);
let admin_jwt = generate_jwt(&state.config.auth.jwt_secret, "1", "admin", 1)
.expect("admin jwt");
let jwt_response = app
.clone()
.oneshot(
Request::builder()
.uri("/api/settings/auth/info")
.header("Authorization", format!("Bearer {admin_jwt}"))
.body(Body::empty())
.unwrap(),
)
.await
.expect("jwt request");
assert_eq!(jwt_response.status(), StatusCode::UNAUTHORIZED);
let missing_response = app
.oneshot(
Request::builder()
.uri("/api/settings/auth/info")
.body(Body::empty())
.unwrap(),
)
.await
.expect("missing auth request");
assert_eq!(missing_response.status(), StatusCode::UNAUTHORIZED);
}
}
+894 -2
View File
@@ -1,16 +1,27 @@
use crate::{
db::models::site::Site,
error::{AppError, AppResult},
proxy::{conductor::SuccessMetrics, executors, token_router::TokenRouter},
services::oauth::oauth_http_client,
AppState,
};
use axum::{
extract::{Path, State},
extract::{Path, Query, State},
http::StatusCode,
response::sse::{Event, Sse},
Json,
};
use futures::Stream;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use tracing::info;
use std::{
collections::HashSet,
convert::Infallible,
time::{Duration, Instant},
};
use tokio::sync::mpsc::{unbounded_channel, UnboundedSender};
use tokio_stream::wrappers::UnboundedReceiverStream;
use tracing::{info, warn};
#[derive(Debug, Deserialize)]
pub struct DetectSiteRequest {
@@ -398,12 +409,89 @@ async fn site_api_endpoints(state: &AppState, site_id: i64) -> AppResult<Vec<Val
.collect())
}
async fn resolve_site_api_base_url(state: &AppState, site: &Site) -> AppResult<String> {
let endpoint: Option<String> = sqlx::query_scalar(
"SELECT url
FROM site_api_endpoints
WHERE site_id = ?
AND COALESCE(enabled, 1) = 1
AND (cooldown_until IS NULL OR cooldown_until <= datetime('now'))
ORDER BY COALESCE(sort_order, 0), id
LIMIT 1",
)
.bind(site.id)
.fetch_optional(&*state.db)
.await?;
let raw_url = endpoint.as_deref().unwrap_or(&site.url).trim();
if raw_url.is_empty() {
return Err(AppError::BadRequest("site url is required".into()));
}
let normalized = raw_url.trim_end_matches('/').to_string();
reqwest::Url::parse(&normalized)
.map_err(|err| AppError::BadRequest(format!("Invalid site API URL: {}", err)))?;
Ok(normalized)
}
#[derive(Debug, Deserialize)]
pub struct UpdateDisabledModelsRequest {
#[serde(default)]
pub models: Vec<String>,
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ProbeSiteOptions {
pub scope: Option<String>,
pub model_name: Option<String>,
pub latency_threshold_ms: Option<i64>,
}
#[derive(Debug, Clone)]
struct SiteProbeModel {
name: String,
source_account_id: i64,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct SiteProbeModelResult {
model_name: String,
status: String,
latency_ms: Option<i64>,
latency_exceeded: bool,
reason: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
http_status: Option<u16>,
#[serde(skip_serializing_if = "Option::is_none")]
stage: Option<&'static str>,
#[serde(skip_serializing_if = "Option::is_none")]
source: Option<&'static str>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct SiteProbeSummary {
site_id: i64,
scope: String,
models_count: usize,
probed: usize,
supported: usize,
unsupported: usize,
disabled_models: Vec<String>,
results: Vec<SiteProbeModelResult>,
}
#[derive(Debug, Clone)]
struct SiteProbePlan {
site: Site,
base_url: String,
scope: String,
latency_threshold_ms: i64,
models: Vec<SiteProbeModel>,
}
/// GET /api/sites/{id}/disabled-models — list disabled model names for a site.
pub async fn list_site_disabled_models(
State(state): State<AppState>,
@@ -477,6 +565,727 @@ pub async fn update_site_disabled_models(
Ok(Json(json!({ "models": normalised })))
}
/// GET /api/sites/{id}/available-models — list currently known available models for a site.
pub async fn list_site_available_models(
State(state): State<AppState>,
Path(id): Path<i64>,
) -> AppResult<Json<Value>> {
ensure_site_exists(&state, id).await?;
let models = site_available_model_names(&state, id).await?;
Ok(Json(json!({ "models": models })))
}
/// POST /api/sites/{id}/probe-now — actively probe one or all known models for a site.
pub async fn probe_site_now(
State(state): State<AppState>,
Path(id): Path<i64>,
Json(options): Json<ProbeSiteOptions>,
) -> AppResult<Json<Value>> {
let plan = build_probe_plan(&state, id, options).await?;
let summary = run_site_probe_plan(state, plan, None).await;
let summary = summary?;
Ok(Json(json!({ "success": true, "probe": summary })))
}
/// GET /api/sites/{id}/probe-stream — stream active site probe progress as SSE.
pub async fn probe_site_stream(
State(state): State<AppState>,
Path(id): Path<i64>,
Query(options): Query<ProbeSiteOptions>,
) -> AppResult<Sse<impl Stream<Item = Result<Event, Infallible>>>> {
let plan = build_probe_plan(&state, id, options).await?;
let (tx, rx) = unbounded_channel::<Result<Event, Infallible>>();
tokio::spawn(async move {
if let Err(err) = run_site_probe_plan(state, plan, Some(tx.clone())).await {
send_probe_sse(&tx, "error", probe_error_payload(&err, None, None));
}
});
Ok(Sse::new(UnboundedReceiverStream::new(rx)))
}
async fn ensure_site_exists(state: &AppState, id: i64) -> AppResult<()> {
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()));
}
Ok(())
}
async fn site_available_model_names(state: &AppState, site_id: i64) -> AppResult<Vec<String>> {
let rows: Vec<String> = sqlx::query_scalar(
"SELECT DISTINCT ma.model_name
FROM model_availability ma
JOIN accounts a ON a.id = ma.account_id
WHERE a.site_id = ?
AND a.status = 'active'
AND COALESCE(ma.available, 0) = 1
AND TRIM(COALESCE(ma.model_name, '')) <> ''
ORDER BY ma.model_name ASC",
)
.bind(site_id)
.fetch_all(&*state.db)
.await?;
Ok(rows)
}
async fn build_probe_plan(
state: &AppState,
site_id: i64,
options: ProbeSiteOptions,
) -> AppResult<SiteProbePlan> {
let site: Site = sqlx::query_as::<_, Site>("SELECT * FROM sites WHERE id = ?")
.bind(site_id)
.fetch_optional(&*state.db)
.await?
.ok_or_else(|| AppError::NotFound("site not found".into()))?;
if site.status != "active" {
return Err(AppError::BadRequest("站点未启用,无法执行探测".into()));
}
let scope = match options
.scope
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("single")
{
"all" => "all",
"single" => "single",
_ => return Err(AppError::BadRequest("scope must be single or all".into())),
}
.to_string();
let latency_threshold_ms = options.latency_threshold_ms.unwrap_or(0).max(0);
let base_url = resolve_site_api_base_url(state, &site).await?;
let requested_model = options
.model_name
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty());
let models = if scope == "all" {
site_probe_models(state, site_id, None).await?
} else if let Some(model_name) = requested_model {
site_probe_models(state, site_id, Some(model_name))
.await?
.into_iter()
.next()
.map(|model| vec![model])
.unwrap_or_else(|| {
vec![SiteProbeModel {
name: model_name.to_string(),
source_account_id: 0,
}]
})
} else {
site_probe_models(state, site_id, None)
.await?
.into_iter()
.take(1)
.collect()
};
if models.is_empty() {
return Err(AppError::BadRequest(
"当前站点还没有可探测模型,请先刷新账号模型列表".into(),
));
}
Ok(SiteProbePlan {
site,
base_url,
scope,
latency_threshold_ms,
models,
})
}
async fn site_probe_models(
state: &AppState,
site_id: i64,
model_name: Option<&str>,
) -> AppResult<Vec<SiteProbeModel>> {
let rows: Vec<(String, i64)> = if let Some(model_name) = model_name {
sqlx::query_as(
"SELECT ma.model_name, MIN(a.id) AS source_account_id
FROM model_availability ma
JOIN accounts a ON a.id = ma.account_id
WHERE a.site_id = ?
AND a.status = 'active'
AND COALESCE(ma.available, 0) = 1
AND ma.model_name = ?
GROUP BY ma.model_name
ORDER BY ma.model_name ASC",
)
.bind(site_id)
.bind(model_name)
.fetch_all(&*state.db)
.await?
} else {
sqlx::query_as(
"SELECT ma.model_name, MIN(a.id) AS source_account_id
FROM model_availability ma
JOIN accounts a ON a.id = ma.account_id
WHERE a.site_id = ?
AND a.status = 'active'
AND COALESCE(ma.available, 0) = 1
AND TRIM(COALESCE(ma.model_name, '')) <> ''
GROUP BY ma.model_name
ORDER BY ma.model_name ASC",
)
.bind(site_id)
.fetch_all(&*state.db)
.await?
};
Ok(rows
.into_iter()
.map(|(name, source_account_id)| SiteProbeModel {
name,
source_account_id,
})
.collect())
}
async fn run_site_probe_plan(
state: AppState,
plan: SiteProbePlan,
tx: Option<UnboundedSender<Result<Event, Infallible>>>,
) -> AppResult<SiteProbeSummary> {
if let Some(tx) = &tx {
send_probe_sse(
tx,
"start",
json!({
"siteId": plan.site.id,
"scope": plan.scope,
"modelsCount": plan.models.len(),
"latencyThresholdMs": plan.latency_threshold_ms,
}),
);
}
let mut results = Vec::with_capacity(plan.models.len());
let mut disabled_models = Vec::new();
let mut disabled_seen = HashSet::new();
for model in &plan.models {
let result = probe_site_model(&state, &plan, model).await;
if let Some(tx) = &tx {
send_probe_sse(tx, "model", json!(&result));
}
if result.status == "unsupported" {
if disabled_seen.insert(result.model_name.clone()) {
disabled_models.push(result.model_name.clone());
}
if let Err(err) = disable_site_model(&state, plan.site.id, &result.model_name).await {
warn!(
site_id = plan.site.id,
model = %result.model_name,
error = %err,
"failed to add probed model to site disabled list"
);
} else if let Some(tx) = &tx {
send_probe_sse(
tx,
"action",
json!({ "action": "disabled", "modelName": result.model_name }),
);
}
}
results.push(result);
}
if !disabled_models.is_empty() {
spawn_route_rebuild(state.clone());
}
let supported = results
.iter()
.filter(|result| result.status == "supported")
.count();
let unsupported = results
.iter()
.filter(|result| result.status == "unsupported")
.count();
let summary = SiteProbeSummary {
site_id: plan.site.id,
scope: plan.scope,
models_count: results.len(),
probed: results.len(),
supported,
unsupported,
disabled_models,
results,
};
if let Some(tx) = &tx {
send_probe_sse(
tx,
"complete",
json!({
"siteId": summary.site_id,
"scope": summary.scope,
"modelsCount": summary.models_count,
"probed": summary.probed,
"supported": summary.supported,
"unsupported": summary.unsupported,
"disabledModels": summary.disabled_models,
}),
);
}
Ok(summary)
}
async fn probe_site_model(
state: &AppState,
plan: &SiteProbePlan,
model: &SiteProbeModel,
) -> SiteProbeModelResult {
let started_at = Instant::now();
let mut selected = match resolve_probe_channel(state, plan, model).await {
Ok(selected) => selected,
Err(err) => {
let diagnostics = probe_model_error_diagnostics(&err, "model");
return SiteProbeModelResult {
model_name: model.name.clone(),
status: "unsupported".to_string(),
latency_ms: None,
latency_exceeded: false,
reason: Some(probe_error_message(&err)),
http_status: diagnostics.http_status,
stage: Some(diagnostics.stage),
source: Some(diagnostics.source),
};
}
};
selected.base_url = plan.base_url.clone();
let request_body = json!({
"model": selected.actual_model,
"messages": [
{ "role": "user", "content": "ping" }
],
"max_tokens": 1,
"stream": false,
});
let mut probe_state = state.clone();
match site_probe_http_client(state, &plan.site).await {
Ok(client) => probe_state.http = client,
Err(err) => {
let diagnostics = probe_model_error_diagnostics(&err, "connect");
return SiteProbeModelResult {
model_name: model.name.clone(),
status: "unsupported".to_string(),
latency_ms: None,
latency_exceeded: false,
reason: Some(probe_error_message(&err)),
http_status: diagnostics.http_status,
stage: Some(diagnostics.stage),
source: Some(diagnostics.source),
};
}
}
let timeout_secs = state.config.proxy.upstream_timeout_secs.max(10);
let response = match tokio::time::timeout(
Duration::from_secs(timeout_secs),
executors::execute_chat_completions(
probe_state,
selected_channel_to_upstream_route(&selected),
request_body,
),
)
.await
{
Ok(result) => result,
Err(_) => Err(AppError::GatewayTimeout(format!(
"探测请求超时({}s)",
timeout_secs
))),
};
let latency_ms = started_at.elapsed().as_millis().min(i64::MAX as u128) as i64;
let latency_exceeded = plan.latency_threshold_ms > 0 && latency_ms > plan.latency_threshold_ms;
match response {
Ok(_) if latency_exceeded => {
if let Err(err) =
update_model_probe_latency(state, selected.account_id, &model.name, latency_ms, true)
.await
{
warn!(error = %err, model = %model.name, "failed to save probe latency");
}
SiteProbeModelResult {
model_name: model.name.clone(),
status: "unsupported".to_string(),
latency_ms: Some(latency_ms),
latency_exceeded: true,
reason: Some(format!("响应延迟 {}ms 超过阈值 {}ms", latency_ms, plan.latency_threshold_ms)),
http_status: None,
stage: Some("model"),
source: Some("probe"),
}
}
Ok(_) => {
if let Some(channel_id) = selected_channel_id(&selected) {
let router = TokenRouter::new(state.db.clone());
if let Err(err) = router
.record_success(
channel_id,
SuccessMetrics {
latency_ms: Some(latency_ms as u64),
cost: None,
metadata: Default::default(),
},
)
.await
{
warn!(channel_id, error = %err, "failed to record probe success");
}
}
if let Err(err) =
update_model_probe_latency(state, selected.account_id, &model.name, latency_ms, true)
.await
{
warn!(error = %err, model = %model.name, "failed to save probe latency");
}
SiteProbeModelResult {
model_name: model.name.clone(),
status: "supported".to_string(),
latency_ms: Some(latency_ms),
latency_exceeded: false,
reason: None,
http_status: None,
stage: Some("model"),
source: Some("upstream"),
}
}
Err(err) => {
let diagnostics = probe_model_error_diagnostics(&err, "model");
if let Some(channel_id) = selected_channel_id(&selected) {
let router = TokenRouter::new(state.db.clone());
if let Err(record_err) = router
.record_failure(
channel_id,
crate::proxy::conductor::FailureInfo {
action: "terminal".to_string(),
status: diagnostics.http_status,
raw_error_text: Some(probe_error_message(&err)),
},
)
.await
{
warn!(channel_id, error = %record_err, "failed to record probe failure");
}
}
if let Err(update_err) =
update_model_probe_latency(state, selected.account_id, &model.name, latency_ms, false)
.await
{
warn!(error = %update_err, model = %model.name, "failed to save failed probe state");
}
SiteProbeModelResult {
model_name: model.name.clone(),
status: "unsupported".to_string(),
latency_ms: Some(latency_ms),
latency_exceeded: false,
reason: Some(probe_error_message(&err)),
http_status: diagnostics.http_status,
stage: Some(diagnostics.stage),
source: Some(diagnostics.source),
}
}
}
}
async fn resolve_probe_channel(
state: &AppState,
plan: &SiteProbePlan,
model: &SiteProbeModel,
) -> AppResult<crate::proxy::conductor::SelectedChannel> {
let router = TokenRouter::new(state.db.clone());
if let Ok(selected) = router.resolve_selected(&model.name).await {
if selected.site_id == plan.site.id {
return Ok(selected);
}
}
let row: Option<(i64, String, Option<String>, Option<String>)> = sqlx::query_as(
"SELECT id, access_token, api_token, oauth_provider
FROM accounts
WHERE id = ?
AND site_id = ?
AND status = 'active'
LIMIT 1",
)
.bind(model.source_account_id)
.bind(plan.site.id)
.fetch_optional(&*state.db)
.await?;
let (account_id, access_token, api_token, oauth_provider) = match row {
Some(row) => row,
None => {
let fallback: Option<(i64, String, Option<String>, Option<String>)> = sqlx::query_as(
"SELECT id, access_token, api_token, oauth_provider
FROM accounts
WHERE site_id = ?
AND status = 'active'
ORDER BY sort_order, id
LIMIT 1",
)
.bind(plan.site.id)
.fetch_optional(&*state.db)
.await?;
fallback.ok_or_else(|| AppError::BadRequest("当前站点没有启用中的账号,无法探测".into()))?
}
};
let api_key = if oauth_provider
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.is_some()
{
crate::services::oauth::OAuthService::new(state.db.clone())
.refresh_token(account_id, plan.site.id)
.await?
.access_token
} else {
api_token.unwrap_or(access_token)
}
.trim()
.to_string();
let api_key = if api_key.is_empty() {
plan.site.api_key.clone().unwrap_or_default()
} else {
api_key
};
if api_key.trim().is_empty() {
return Err(AppError::BadRequest("当前站点账号缺少可用 Token,无法探测".into()));
}
Ok(crate::proxy::conductor::SelectedChannel {
channel_id: 0,
route_id: None,
site_id: plan.site.id,
account_id: Some(account_id),
account_unit_cost: None,
base_url: plan.base_url.clone(),
api_key,
model: model.name.clone(),
actual_model: model.name.clone(),
token_name: None,
token_value: None,
})
}
async fn site_probe_http_client(state: &AppState, site: &Site) -> AppResult<reqwest::Client> {
let has_custom_proxy = site
.proxy_url
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.is_some();
if has_custom_proxy || site.use_system_proxy.unwrap_or(false) {
return oauth_http_client(
&state.db,
site.proxy_url.as_deref(),
site.use_system_proxy.unwrap_or(false),
)
.await;
}
Ok(state.http.clone())
}
fn selected_channel_to_upstream_route(
selected: &crate::proxy::conductor::SelectedChannel,
) -> crate::proxy::executors::UpstreamRoute {
crate::proxy::executors::UpstreamRoute {
base_url: selected.base_url.clone(),
api_key: selected.api_key.clone(),
model: selected.actual_model.clone(),
}
}
fn selected_channel_id(selected: &crate::proxy::conductor::SelectedChannel) -> Option<i64> {
(selected.channel_id > 0).then_some(selected.channel_id)
}
async fn update_model_probe_latency(
state: &AppState,
account_id: Option<i64>,
model_name: &str,
latency_ms: i64,
available: bool,
) -> AppResult<()> {
let Some(account_id) = account_id else {
return Ok(());
};
sqlx::query(
"UPDATE model_availability
SET latency_ms = ?,
available = ?,
checked_at = datetime('now')
WHERE account_id = ?
AND model_name = ?",
)
.bind(latency_ms)
.bind(available)
.bind(account_id)
.bind(model_name)
.execute(&*state.db)
.await?;
Ok(())
}
async fn disable_site_model(state: &AppState, site_id: i64, model_name: &str) -> AppResult<()> {
sqlx::query(
"INSERT INTO site_disabled_models (site_id, model_name)
VALUES (?, ?)
ON CONFLICT(site_id, model_name) DO NOTHING",
)
.bind(site_id)
.bind(model_name)
.execute(&*state.db)
.await?;
Ok(())
}
fn send_probe_sse(
tx: &UnboundedSender<Result<Event, Infallible>>,
event_type: &str,
payload: Value,
) -> bool {
tx.send(Ok(Event::default()
.event(event_type)
.json_data(payload)
.unwrap_or_else(|_| Event::default().event(event_type).data("{}"))))
.is_ok()
}
fn probe_error_message(err: &AppError) -> String {
match err {
AppError::NotFound(message)
| AppError::BadRequest(message)
| AppError::Forbidden(message)
| AppError::Unauthorized(message)
| AppError::ServiceUnavailable(message)
| AppError::GatewayTimeout(message) => message.clone(),
AppError::UpstreamError(message) => simplify_upstream_error_message(message),
AppError::RateLimited => "触发频率限制".to_string(),
AppError::RateLimitedWithRetry { message, .. } => message.clone(),
AppError::Internal(err) => err.to_string(),
AppError::Database(err) => err.to_string(),
}
}
fn probe_error_payload(err: &AppError, site_id: Option<i64>, model_name: Option<&str>) -> Value {
let message = probe_error_message(err);
let diagnostics = probe_model_error_diagnostics(err, "stream");
json!({
"message": message,
"httpStatus": diagnostics.http_status,
"siteId": site_id,
"modelName": model_name,
"stage": diagnostics.stage,
"source": diagnostics.source,
})
}
fn simplify_upstream_error_message(message: &str) -> String {
let marker = "Upstream returned ";
if let Some(start) = message.find(marker).map(|index| index + marker.len()) {
let tail = &message[start..];
if let Some((status, body)) = tail.split_once(':') {
let body = body.trim();
return if body.is_empty() {
format!("上游返回 HTTP {}", status.trim())
} else {
format!(
"上游返回 HTTP {}:{}",
status.trim(),
body.chars().take(180).collect::<String>()
)
};
}
}
message.chars().take(220).collect()
}
fn upstream_status_from_error(err: &AppError) -> Option<u16> {
match err {
AppError::UpstreamError(message) => extract_http_status_from_text(message),
_ => None,
}
}
fn extract_http_status_from_text(message: &str) -> Option<u16> {
message
.split(|ch: char| !ch.is_ascii_digit())
.find_map(|token| {
if token.len() != 3 {
return None;
}
let status = token.parse::<u16>().ok()?;
if (100..=599).contains(&status) {
Some(status)
} else {
None
}
})
}
#[derive(Debug, Clone, Copy)]
struct ProbeErrorDiagnostics {
http_status: Option<u16>,
stage: &'static str,
source: &'static str,
}
fn probe_model_error_diagnostics(
err: &AppError,
fallback_stage: &'static str,
) -> ProbeErrorDiagnostics {
if matches!(err, AppError::UpstreamError(_)) {
return ProbeErrorDiagnostics {
http_status: upstream_status_from_error(err),
stage: "model",
source: "upstream",
};
}
ProbeErrorDiagnostics {
http_status: None,
stage: fallback_stage,
source: "probe",
}
}
fn spawn_route_rebuild(state: AppState) {
tokio::spawn(async move {
match crate::services::route_rebuild::rebuild_routes(&state).await {
Ok(summary) => {
info!(
created_routes = summary.created_routes,
reused_routes = summary.reused_routes,
created_channels = summary.created_channels,
deleted_channels = summary.deleted_channels,
"route rebuild completed after site probe"
);
}
Err(err) => {
warn!(error = %err, "route rebuild after site probe failed");
}
}
});
}
fn normalize_site_url(raw: &str) -> AppResult<String> {
let trimmed = raw.trim().trim_end_matches('/').to_string();
if trimmed.is_empty() {
@@ -645,3 +1454,86 @@ async fn replace_site_api_endpoints(
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn probe_options_accept_camel_case_payload() {
let options: ProbeSiteOptions = serde_json::from_value(json!({
"scope": "single",
"modelName": "gpt-4o",
"latencyThresholdMs": 1800
}))
.expect("probe options should deserialize from frontend payload");
assert_eq!(options.scope.as_deref(), Some("single"));
assert_eq!(options.model_name.as_deref(), Some("gpt-4o"));
assert_eq!(options.latency_threshold_ms, Some(1800));
}
#[test]
fn simplify_upstream_error_message_names_http_status() {
let message = simplify_upstream_error_message(
"Upstream returned 404 Not Found: {\"error\":\"model missing\"}",
);
assert_eq!(
message,
"上游返回 HTTP 404 Not Found:{\"error\":\"model missing\"}"
);
}
#[test]
fn upstream_status_from_error_extracts_probe_failure_status() {
let err = AppError::UpstreamError("Upstream returned 503 Service Unavailable: busy".into());
assert_eq!(upstream_status_from_error(&err), Some(503));
}
#[test]
fn upstream_status_from_error_extracts_localized_probe_failure_status() {
let err = AppError::UpstreamError("上游返回 HTTP 404 Not Found:model missing".into());
assert_eq!(upstream_status_from_error(&err), Some(404));
}
#[test]
fn site_probe_model_result_serializes_structured_diagnostics() {
let result = SiteProbeModelResult {
model_name: "gpt-4o".to_string(),
status: "unsupported".to_string(),
latency_ms: Some(88),
latency_exceeded: false,
reason: Some("Not found".to_string()),
http_status: Some(404),
stage: Some("model"),
source: Some("upstream"),
};
let payload = serde_json::to_value(&result).expect("model result should serialize");
assert_eq!(payload.get("httpStatus").and_then(Value::as_u64), Some(404));
assert_eq!(payload.get("stage").and_then(Value::as_str), Some("model"));
assert_eq!(payload.get("source").and_then(Value::as_str), Some("upstream"));
assert_eq!(payload.get("modelName").and_then(Value::as_str), Some("gpt-4o"));
}
#[test]
fn probe_error_payload_preserves_upstream_status_for_sse_diagnostics() {
let err = AppError::UpstreamError("Upstream returned 404 Not Found: model missing".into());
let payload = probe_error_payload(&err, Some(7), Some("gpt-4o"));
assert_eq!(payload.get("httpStatus").and_then(Value::as_u64), Some(404));
assert_eq!(payload.get("stage").and_then(Value::as_str), Some("model"));
assert_eq!(payload.get("source").and_then(Value::as_str), Some("upstream"));
assert_eq!(payload.get("siteId").and_then(Value::as_i64), Some(7));
assert_eq!(payload.get("modelName").and_then(Value::as_str), Some("gpt-4o"));
assert_eq!(
payload.get("message").and_then(Value::as_str),
Some("上游返回 HTTP 404 Not Found:model missing")
);
}
}
+150 -1
View File
@@ -324,7 +324,10 @@ pub async fn get_user_usage(
// Recent logs
let recent_logs: Vec<Value> = sqlx::query(
r#"
SELECT l.*
SELECT
l.*,
k.id AS downstream_api_key_id,
k.name AS downstream_key_name
FROM usage_logs l
JOIN downstream_api_keys k ON l.downstream_key_id = k.id
WHERE k.user_id = ?
@@ -341,6 +344,11 @@ pub async fn get_user_usage(
json!({
"id": row.get::<i64, _>("id"),
"model": row.get::<Option<String>, _>("model"),
"downstreamApiKeyId": row.get::<Option<i64>, _>("downstream_api_key_id"),
"downstreamKeyName": row.get::<Option<String>, _>("downstream_key_name"),
"status": "success",
"httpStatus": 200,
"source": "usage-logs",
"promptTokens": row.get::<Option<i64>, _>("prompt_tokens").unwrap_or(0),
"completionTokens": row.get::<Option<i64>, _>("completion_tokens").unwrap_or(0),
"totalTokens": row.get::<Option<i64>, _>("total_tokens").unwrap_or(0),
@@ -699,3 +707,144 @@ pub async fn update_user_key(
Ok(Json(json!({ "data": updated })))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{
config::{AppConfig, AuthConfig, DbConfig, ProxyConfig, ServerConfig},
services::{builtin_keys::BuiltinKeys, runtime_auth::RuntimeAuthSettings},
};
use axum::extract::State;
use sqlx::sqlite::SqlitePoolOptions;
async fn test_state() -> AppState {
let pool = SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.expect("in-memory database");
sqlx::query(
"CREATE TABLE users (
id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
username TEXT NOT NULL,
email TEXT NOT NULL,
password_hash TEXT NOT NULL,
role TEXT NOT NULL DEFAULT 'user',
status TEXT NOT NULL DEFAULT 'active',
created_at TEXT DEFAULT (datetime('now')),
updated_at TEXT DEFAULT (datetime('now'))
)",
)
.execute(&pool)
.await
.expect("users table");
sqlx::query(
"CREATE TABLE downstream_api_keys (
id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
name TEXT NOT NULL,
key TEXT NOT NULL,
description TEXT,
group_name TEXT,
tags TEXT,
enabled INTEGER DEFAULT 1,
expires_at TEXT,
max_cost REAL,
used_cost REAL DEFAULT 0,
max_requests INTEGER,
used_requests INTEGER DEFAULT 0,
supported_models TEXT,
allowed_route_ids TEXT,
site_weight_multipliers TEXT,
excluded_site_ids TEXT,
excluded_credential_refs TEXT,
last_used_at TEXT,
user_id INTEGER,
created_at TEXT DEFAULT (datetime('now')),
updated_at TEXT DEFAULT (datetime('now'))
)",
)
.execute(&pool)
.await
.expect("downstream keys table");
AppState::new(
AppConfig {
server: ServerConfig {
host: "127.0.0.1".into(),
port: 0,
log_level: "info".into(),
cors_origins: vec![],
},
db: DbConfig {
url: "sqlite::memory:".into(),
max_connections: 1,
ssl: false,
runtime_db_url: None,
runtime_db_dialect: None,
runtime_db_ssl: None,
},
auth: AuthConfig {
jwt_secret: "test-secret".into(),
jwt_expiry_hours: 1,
auth_token: "auth".into(),
proxy_token: "proxy".into(),
},
proxy: ProxyConfig {
upstream_timeout_secs: 30,
first_byte_timeout_secs: 30,
max_retries: 3,
proxy_log_retention_days: 7,
template_proxy_url: None,
},
comfyui: Default::default(),
},
pool,
BuiltinKeys {
comfyui_agent: (1, "sk-comfyui-test".into()),
},
RuntimeAuthSettings::from_token_for_test("proxy"),
)
}
#[tokio::test]
async fn update_user_persists_admin_role_for_existing_user() {
let state = test_state().await;
let user_id = sqlx::query(
"INSERT INTO users (username, email, password_hash, role, status)
VALUES (?, ?, ?, 'user', 'active')",
)
.bind("Regular")
.bind("regular@example.test")
.bind("hash")
.execute(&*state.db)
.await
.expect("insert user")
.last_insert_rowid();
let Json(value) = update_user(
State(state.clone()),
Path(user_id),
Json(UpdateUserRequest {
username: None,
email: None,
password: None,
role: Some("admin".into()),
status: None,
}),
)
.await
.expect("update user");
assert_eq!(value["data"]["role"], "admin");
let role: String = sqlx::query_scalar("SELECT role FROM users WHERE id = ?")
.bind(user_id)
.fetch_one(&*state.db)
.await
.expect("read role");
assert_eq!(role, "admin");
}
}
+813 -13
View File
@@ -10,14 +10,14 @@ use crate::{
};
use axum::{
extract::{ConnectInfo, Path, State},
http::HeaderMap,
http::{header, HeaderMap},
routing::{get, post},
Json, Router,
};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use std::net::SocketAddr;
use tracing::info;
use std::net::{IpAddr, SocketAddr};
use tracing::{info, warn};
// ── Request / Response types ─────────────────────────────────────────
@@ -70,7 +70,6 @@ pub fn routes() -> Router<AppState> {
.route("/api/users/register", post(register))
.route("/api/auth/login", post(login))
.route("/api/users/login", post(login))
.route("/api/auth/login-api-key", post(login_api_key))
.route("/api/users/login-api-key", post(login_api_key))
.route("/api/auth/me", get(me))
.route("/api/users/me", get(me).patch(update_me))
@@ -229,24 +228,82 @@ async fn register(
/// POST /api/auth/login
async fn login(
State(state): State<AppState>,
ConnectInfo(addr): ConnectInfo<SocketAddr>,
headers: HeaderMap,
Json(payload): Json<LoginRequest>,
) -> AppResult<Json<Value>> {
info!("Login: {}", payload.email);
let client_ip = client_ip_from_headers(&headers, addr);
let user_agent = user_agent_from_headers(&headers);
let email = payload.email.trim().to_string();
let user: User = sqlx::query_as::<_, User>("SELECT * FROM users WHERE email = ?")
.bind(&payload.email)
let user: User = if let Some(user) = sqlx::query_as::<_, User>("SELECT * FROM users WHERE email = ?")
.bind(&email)
.fetch_optional(&*state.db)
.await?
.ok_or_else(|| AppError::Unauthorized("Invalid email or password".into()))?;
{
user
} else {
audit_user_login(&state, LoginAuditEvent {
method: LoginMethod::Password,
success: false,
user: None,
identifier: Some(&email),
api_key: None,
client_ip: &client_ip,
user_agent: user_agent.as_deref(),
reason: Some("Invalid email or password"),
})
.await;
return Err(AppError::Unauthorized("Invalid email or password".into()));
};
if user.status != "active" {
audit_user_login(&state, LoginAuditEvent {
method: LoginMethod::Password,
success: false,
user: Some(&user),
identifier: Some(&email),
api_key: None,
client_ip: &client_ip,
user_agent: user_agent.as_deref(),
reason: Some("Account is disabled"),
})
.await;
return Err(AppError::Forbidden("Account is disabled".into()));
}
if user.role == "admin" {
audit_user_login(&state, LoginAuditEvent {
method: LoginMethod::Password,
success: false,
user: Some(&user),
identifier: Some(&email),
api_key: None,
client_ip: &client_ip,
user_agent: user_agent.as_deref(),
reason: Some("Admin users must sign in with the admin API token"),
})
.await;
return Err(AppError::Forbidden(
"Admin users must sign in with the admin API token".into(),
));
}
let valid = verify_password(&payload.password, &user.password_hash)
.map_err(|e| AppError::Internal(anyhow::anyhow!("Password verification failed: {e}")))?;
if !valid {
audit_user_login(&state, LoginAuditEvent {
method: LoginMethod::Password,
success: false,
user: Some(&user),
identifier: Some(&email),
api_key: None,
client_ip: &client_ip,
user_agent: user_agent.as_deref(),
reason: Some("Invalid email or password"),
})
.await;
return Err(AppError::Unauthorized("Invalid email or password".into()));
}
@@ -258,6 +315,18 @@ async fn login(
)
.map_err(|e| AppError::Internal(anyhow::anyhow!("Failed to generate JWT: {e}")))?;
audit_user_login(&state, LoginAuditEvent {
method: LoginMethod::Password,
success: true,
user: Some(&user),
identifier: Some(&email),
api_key: None,
client_ip: &client_ip,
user_agent: user_agent.as_deref(),
reason: None,
})
.await;
Ok(Json(json!({
"token": token,
"user": {
@@ -288,10 +357,25 @@ async fn login(
/// be enabled, and the owning user must be active.
async fn login_api_key(
State(state): State<AppState>,
ConnectInfo(addr): ConnectInfo<SocketAddr>,
headers: HeaderMap,
Json(payload): Json<LoginApiKeyRequest>,
) -> AppResult<Json<Value>> {
let client_ip = client_ip_from_headers(&headers, addr);
let user_agent = user_agent_from_headers(&headers);
let key = payload.api_key.trim();
if key.is_empty() {
audit_user_login(&state, LoginAuditEvent {
method: LoginMethod::ApiKey,
success: false,
user: None,
identifier: None,
api_key: None,
client_ip: &client_ip,
user_agent: user_agent.as_deref(),
reason: Some("api_key is required"),
})
.await;
return Err(AppError::BadRequest("api_key is required".into()));
}
@@ -302,20 +386,91 @@ async fn login_api_key(
.fetch_optional(&*state.db)
.await?;
let api_key = row.ok_or_else(|| AppError::Unauthorized("Invalid API key".into()))?;
let user_id = api_key
.user_id
.ok_or_else(|| AppError::Unauthorized("API key is not bound to a user".into()))?;
let api_key = if let Some(api_key) = row {
api_key
} else {
let masked_key = mask_api_key(key);
audit_user_login(&state, LoginAuditEvent {
method: LoginMethod::ApiKey,
success: false,
user: None,
identifier: Some(masked_key.as_str()),
api_key: None,
client_ip: &client_ip,
user_agent: user_agent.as_deref(),
reason: Some("Invalid API key"),
})
.await;
return Err(AppError::Unauthorized("Invalid API key".into()));
};
let user_id = if let Some(user_id) = api_key.user_id {
user_id
} else {
audit_user_login(&state, LoginAuditEvent {
method: LoginMethod::ApiKey,
success: false,
user: None,
identifier: Some(&api_key.name),
api_key: Some(&api_key),
client_ip: &client_ip,
user_agent: user_agent.as_deref(),
reason: Some("API key is not bound to a user"),
})
.await;
return Err(AppError::Unauthorized("API key is not bound to a user".into()));
};
let user: User = sqlx::query_as::<_, User>("SELECT * FROM users WHERE id = ?")
let user: User = if let Some(user) = sqlx::query_as::<_, User>("SELECT * FROM users WHERE id = ?")
.bind(user_id)
.fetch_optional(&*state.db)
.await?
.ok_or_else(|| AppError::Unauthorized("User not found".into()))?;
{
user
} else {
audit_user_login(&state, LoginAuditEvent {
method: LoginMethod::ApiKey,
success: false,
user: None,
identifier: Some(&api_key.name),
api_key: Some(&api_key),
client_ip: &client_ip,
user_agent: user_agent.as_deref(),
reason: Some("User not found"),
})
.await;
return Err(AppError::Unauthorized("User not found".into()));
};
if user.status != "active" {
audit_user_login(&state, LoginAuditEvent {
method: LoginMethod::ApiKey,
success: false,
user: Some(&user),
identifier: Some(&user.email),
api_key: Some(&api_key),
client_ip: &client_ip,
user_agent: user_agent.as_deref(),
reason: Some("Account is disabled"),
})
.await;
return Err(AppError::Forbidden("Account is disabled".into()));
}
if user.role == "admin" {
audit_user_login(&state, LoginAuditEvent {
method: LoginMethod::ApiKey,
success: false,
user: Some(&user),
identifier: Some(&user.email),
api_key: Some(&api_key),
client_ip: &client_ip,
user_agent: user_agent.as_deref(),
reason: Some("Admin users must sign in with the admin API token"),
})
.await;
return Err(AppError::Forbidden(
"Admin users must sign in with the admin API token".into(),
));
}
info!("Login via API key: user_id={}", user.id);
@@ -327,6 +482,18 @@ async fn login_api_key(
)
.map_err(|e| AppError::Internal(anyhow::anyhow!("Failed to generate JWT: {e}")))?;
audit_user_login(&state, LoginAuditEvent {
method: LoginMethod::ApiKey,
success: true,
user: Some(&user),
identifier: Some(&user.email),
api_key: Some(&api_key),
client_ip: &client_ip,
user_agent: user_agent.as_deref(),
reason: None,
})
.await;
Ok(Json(json!({
"token": token,
"user": {
@@ -503,6 +670,207 @@ async fn get_user(
})))
}
// ── User login audit helpers ──────────────────────────────────────────
#[derive(Debug, Clone, Copy)]
enum LoginMethod {
Password,
ApiKey,
}
impl LoginMethod {
fn label(self) -> &'static str {
match self {
LoginMethod::Password => "password",
LoginMethod::ApiKey => "api_key",
}
}
fn title_label(self) -> &'static str {
match self {
LoginMethod::Password => "密码登录",
LoginMethod::ApiKey => "API Key 登录",
}
}
}
struct LoginAuditEvent<'a> {
method: LoginMethod,
success: bool,
user: Option<&'a User>,
identifier: Option<&'a str>,
api_key: Option<&'a DownstreamApiKey>,
client_ip: &'a str,
user_agent: Option<&'a str>,
reason: Option<&'a str>,
}
async fn audit_user_login(state: &AppState, audit: LoginAuditEvent<'_>) {
let event_type = "user_login";
let level = if audit.success { "info" } else { "warning" };
let title = if audit.success {
format!("用户登录成功:{}", display_user_for_audit(audit.user, audit.identifier))
} else {
format!("用户登录失败:{}", display_user_for_audit(audit.user, audit.identifier))
};
let mut message_parts = vec![
format!("方式={}", audit.method.label()),
format!("IP={}", audit.client_ip),
];
if let Some(user) = audit.user {
message_parts.push(format!("用户ID={}", user.id));
message_parts.push(format!("邮箱={}", user.email));
} else if let Some(identifier) = audit.identifier {
message_parts.push(format!("标识={}", identifier));
}
if let Some(api_key) = audit.api_key {
message_parts.push(format!("密钥={}", api_key.name));
message_parts.push(format!("密钥ID={}", api_key.id));
}
if let Some(reason) = audit.reason.filter(|value| !value.trim().is_empty()) {
message_parts.push(format!("原因={reason}"));
}
if let Some(user_agent) = audit.user_agent.filter(|value| !value.trim().is_empty()) {
message_parts.push(format!("UA={}", truncate_for_audit(user_agent, 180)));
}
if let Err(err) = sqlx::query(
"INSERT INTO events (type, title, message, level, read, related_id, related_type)
VALUES (?, ?, ?, ?, 0, ?, ?)",
)
.bind(event_type)
.bind(title)
.bind(message_parts.join(" · "))
.bind(level)
.bind(audit.user.map(|user| user.id))
.bind(if audit.user.is_some() { Some("user") } else { None::<&str> })
.execute(&*state.db)
.await
{
warn!(
%err,
method = audit.method.label(),
success = audit.success,
client_ip = audit.client_ip,
"failed to write user login audit event",
);
} else {
info!(
method = audit.method.label(),
method_title = audit.method.title_label(),
success = audit.success,
client_ip = audit.client_ip,
"user login audit recorded",
);
}
}
fn display_user_for_audit(user: Option<&User>, identifier: Option<&str>) -> String {
if let Some(user) = user {
if !user.email.trim().is_empty() {
return user.email.clone();
}
if !user.username.trim().is_empty() {
return user.username.clone();
}
return format!("用户 #{}", user.id);
}
identifier
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("未知用户")
.to_string()
}
fn user_agent_from_headers(headers: &HeaderMap) -> Option<String> {
headers
.get(header::USER_AGENT)
.and_then(|value| value.to_str().ok())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| truncate_for_audit(value, 240))
}
fn client_ip_from_headers(headers: &HeaderMap, fallback: SocketAddr) -> String {
for name in [
"cf-connecting-ip",
"x-real-ip",
"x-forwarded-for",
"forwarded",
] {
if let Some(value) = header_ip_value(headers, name) {
return value;
}
}
fallback.ip().to_string()
}
fn header_ip_value(headers: &HeaderMap, name: &str) -> Option<String> {
let raw = headers.get(name)?.to_str().ok()?.trim();
if raw.is_empty() {
return None;
}
if name.eq_ignore_ascii_case("forwarded") {
return forwarded_header_ip(raw);
}
raw.split(',')
.filter_map(normalize_ip_token)
.next()
}
fn forwarded_header_ip(raw: &str) -> Option<String> {
raw.split(',')
.flat_map(|entry| entry.split(';'))
.filter_map(|part| {
let (key, value) = part.split_once('=')?;
if !key.trim().eq_ignore_ascii_case("for") {
return None;
}
normalize_ip_token(value)
})
.next()
}
fn normalize_ip_token(raw: &str) -> Option<String> {
let mut token = raw.trim().trim_matches('"').trim_matches('\'').trim();
if token.is_empty() || token.eq_ignore_ascii_case("unknown") {
return None;
}
if token.starts_with('[') {
if let Some(end) = token.find(']') {
token = &token[1..end];
}
} else if let Some((host, port)) = token.rsplit_once(':') {
if host.matches(':').count() == 0 && port.parse::<u16>().is_ok() {
token = host;
}
}
token
.parse::<IpAddr>()
.map(|ip| ip.to_string())
.ok()
}
fn truncate_for_audit(value: &str, max_chars: usize) -> String {
if value.chars().count() <= max_chars {
return value.to_string();
}
let mut truncated = value.chars().take(max_chars.saturating_sub(1)).collect::<String>();
truncated.push('…');
truncated
}
fn mask_api_key(key: &str) -> String {
let trimmed = key.trim();
let tail = trimmed.chars().rev().take(4).collect::<Vec<_>>();
if tail.is_empty() {
return "sk-****".to_string();
}
let suffix = tail.into_iter().rev().collect::<String>();
format!("••••{suffix}")
}
// ── Inline JWT extract helper for this module ──────────────────────────
fn extract_jwt(state: &AppState, headers: &HeaderMap) -> AppResult<(i64, String)> {
@@ -522,3 +890,435 @@ fn extract_jwt(state: &AppState, headers: &HeaderMap) -> AppResult<(i64, String)
Ok((uid, claims.role))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{
auth::{jwt::validate_jwt, password::hash_password},
config::{AppConfig, AuthConfig, DbConfig, ProxyConfig, ServerConfig},
services::{builtin_keys::BuiltinKeys, runtime_auth::RuntimeAuthSettings},
};
use sqlx::sqlite::SqlitePoolOptions;
async fn test_state() -> AppState {
let pool = SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.expect("in-memory database");
sqlx::query(
"CREATE TABLE users (
id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
username TEXT NOT NULL,
email TEXT NOT NULL,
password_hash TEXT NOT NULL,
role TEXT NOT NULL DEFAULT 'user',
status TEXT NOT NULL DEFAULT 'active',
created_at TEXT DEFAULT (datetime('now')),
updated_at TEXT DEFAULT (datetime('now'))
)",
)
.execute(&pool)
.await
.expect("users table");
sqlx::query(
"CREATE TABLE downstream_api_keys (
id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
name TEXT NOT NULL,
key TEXT NOT NULL,
description TEXT,
group_name TEXT,
tags TEXT,
enabled INTEGER DEFAULT 1,
expires_at TEXT,
max_cost REAL,
used_cost REAL DEFAULT 0,
max_requests INTEGER,
used_requests INTEGER DEFAULT 0,
supported_models TEXT,
allowed_route_ids TEXT,
site_weight_multipliers TEXT,
excluded_site_ids TEXT,
excluded_credential_refs TEXT,
last_used_at TEXT,
user_id INTEGER,
created_at TEXT DEFAULT (datetime('now')),
updated_at TEXT DEFAULT (datetime('now'))
)",
)
.execute(&pool)
.await
.expect("downstream keys table");
sqlx::query(
"CREATE TABLE events (
id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
type TEXT NOT NULL,
title TEXT NOT NULL,
message TEXT,
level TEXT NOT NULL DEFAULT 'info',
read INTEGER DEFAULT 0,
related_id INTEGER,
related_type TEXT,
created_at TEXT DEFAULT (datetime('now'))
)",
)
.execute(&pool)
.await
.expect("events table");
AppState::new(
AppConfig {
server: ServerConfig {
host: "127.0.0.1".into(),
port: 0,
log_level: "info".into(),
cors_origins: vec![],
},
db: DbConfig {
url: "sqlite::memory:".into(),
max_connections: 1,
ssl: false,
runtime_db_url: None,
runtime_db_dialect: None,
runtime_db_ssl: None,
},
auth: AuthConfig {
jwt_secret: "test-secret".into(),
jwt_expiry_hours: 1,
auth_token: "auth".into(),
proxy_token: "proxy".into(),
},
proxy: ProxyConfig {
upstream_timeout_secs: 30,
first_byte_timeout_secs: 30,
max_retries: 3,
proxy_log_retention_days: 7,
template_proxy_url: None,
},
comfyui: Default::default(),
},
pool,
BuiltinKeys {
comfyui_agent: (1, "sk-comfyui-test".into()),
},
RuntimeAuthSettings::from_token_for_test("proxy"),
)
}
fn test_addr() -> ConnectInfo<SocketAddr> {
ConnectInfo(SocketAddr::from(([127, 0, 0, 1], 12345)))
}
fn empty_headers() -> HeaderMap {
HeaderMap::new()
}
fn headers_with(name: &'static str, value: &'static str) -> HeaderMap {
let mut headers = HeaderMap::new();
headers.insert(name, value.parse().expect("valid test header"));
headers
}
async fn latest_audit_event(state: &AppState) -> (String, String, String, i64, String) {
sqlx::query_as::<_, (String, String, String, i64, String)>(
"SELECT type, title, COALESCE(message, ''), COALESCE(related_id, 0), COALESCE(related_type, '')
FROM events
ORDER BY id DESC
LIMIT 1",
)
.fetch_one(&*state.db)
.await
.expect("latest audit event")
}
#[tokio::test]
async fn password_login_rejects_admin_users() {
let state = test_state().await;
let password_hash = hash_password("admin-password").expect("password hash");
sqlx::query(
"INSERT INTO users (username, email, password_hash, role, status)
VALUES (?, ?, ?, 'admin', 'active')",
)
.bind("Admin")
.bind("admin@example.test")
.bind(password_hash)
.execute(&*state.db)
.await
.expect("insert admin user");
let err = login(
State(state),
test_addr(),
empty_headers(),
Json(LoginRequest {
email: "admin@example.test".into(),
password: "admin-password".into(),
}),
)
.await
.expect_err("admin password login must be rejected");
assert!(matches!(err, AppError::Forbidden(message) if message.contains("admin API token")));
}
#[tokio::test]
async fn password_login_allows_regular_users() {
let state = test_state().await;
let password_hash = hash_password("user-password").expect("password hash");
let user_id = sqlx::query(
"INSERT INTO users (username, email, password_hash, role, status)
VALUES (?, ?, ?, 'user', 'active')",
)
.bind("Regular User")
.bind("user@example.test")
.bind(password_hash)
.execute(&*state.db)
.await
.expect("insert user")
.last_insert_rowid();
let Json(value) = login(
State(state.clone()),
test_addr(),
empty_headers(),
Json(LoginRequest {
email: "user@example.test".into(),
password: "user-password".into(),
}),
)
.await
.expect("regular user password login succeeds");
assert_eq!(value["user"]["role"], "user");
let token = value["token"].as_str().expect("token");
let claims = validate_jwt(&state.config.auth.jwt_secret, token).expect("valid jwt");
assert_eq!(claims.sub, user_id.to_string());
assert_eq!(claims.role, "user");
}
#[tokio::test]
async fn password_login_audits_forwarded_client_ip() {
let state = test_state().await;
let password_hash = hash_password("user-password").expect("password hash");
let user_id = sqlx::query(
"INSERT INTO users (username, email, password_hash, role, status)
VALUES (?, ?, ?, 'user', 'active')",
)
.bind("Regular User")
.bind("audit@example.test")
.bind(password_hash)
.execute(&*state.db)
.await
.expect("insert user")
.last_insert_rowid();
let Json(_) = login(
State(state.clone()),
test_addr(),
headers_with("x-forwarded-for", "203.0.113.8, 10.0.0.1"),
Json(LoginRequest {
email: "audit@example.test".into(),
password: "user-password".into(),
}),
)
.await
.expect("regular user password login succeeds");
let (event_type, title, message, related_id, related_type) = latest_audit_event(&state).await;
assert_eq!(event_type, "user_login");
assert!(title.contains("用户登录成功"));
assert!(title.contains("audit@example.test"));
assert!(message.contains("方式=password"));
assert!(message.contains("IP=203.0.113.8"));
assert_eq!(related_id, user_id);
assert_eq!(related_type, "user");
}
#[tokio::test]
async fn api_key_login_rejects_admin_role() {
let state = test_state().await;
let password_hash = hash_password("unused-password").expect("password hash");
let user_id = sqlx::query(
"INSERT INTO users (username, email, password_hash, role, status)
VALUES (?, ?, ?, 'admin', 'active')",
)
.bind("Admin")
.bind("admin@example.test")
.bind(password_hash)
.execute(&*state.db)
.await
.expect("insert admin user")
.last_insert_rowid();
sqlx::query(
"INSERT INTO downstream_api_keys (name, key, user_id, enabled)
VALUES (?, ?, ?, 1)",
)
.bind("Admin default")
.bind("sk-admin-bound")
.bind(user_id)
.execute(&*state.db)
.await
.expect("insert bound api key");
let err = login_api_key(
State(state),
test_addr(),
empty_headers(),
Json(LoginApiKeyRequest {
api_key: "sk-admin-bound".into(),
}),
)
.await
.expect_err("admin api key login must be rejected");
assert!(matches!(err, AppError::Forbidden(message) if message.contains("admin API token")));
assert!(user_id > 0);
}
#[tokio::test]
async fn api_key_login_rejects_latest_admin_role_after_promotion() {
let state = test_state().await;
let password_hash = hash_password("unused-password").expect("password hash");
let user_id = sqlx::query(
"INSERT INTO users (username, email, password_hash, role, status)
VALUES (?, ?, ?, 'user', 'active')",
)
.bind("Promoted User")
.bind("promoted@example.test")
.bind(password_hash)
.execute(&*state.db)
.await
.expect("insert regular user")
.last_insert_rowid();
sqlx::query(
"INSERT INTO downstream_api_keys (name, key, user_id, enabled)
VALUES (?, ?, ?, 1)",
)
.bind("Existing user key")
.bind("sk-promoted-bound")
.bind(user_id)
.execute(&*state.db)
.await
.expect("insert bound api key");
sqlx::query("UPDATE users SET role = 'admin' WHERE id = ?")
.bind(user_id)
.execute(&*state.db)
.await
.expect("promote user to admin");
let err = login_api_key(
State(state),
test_addr(),
empty_headers(),
Json(LoginApiKeyRequest {
api_key: "sk-promoted-bound".into(),
}),
)
.await
.expect_err("api key login must use latest role and reject admin");
assert!(matches!(err, AppError::Forbidden(message) if message.contains("admin API token")));
assert!(user_id > 0);
}
#[tokio::test]
async fn api_key_login_allows_regular_user_role() {
let state = test_state().await;
let password_hash = hash_password("unused-password").expect("password hash");
let user_id = sqlx::query(
"INSERT INTO users (username, email, password_hash, role, status)
VALUES (?, ?, ?, 'user', 'active')",
)
.bind("Regular User")
.bind("regular@example.test")
.bind(password_hash)
.execute(&*state.db)
.await
.expect("insert regular user")
.last_insert_rowid();
sqlx::query(
"INSERT INTO downstream_api_keys (name, key, user_id, enabled)
VALUES (?, ?, ?, 1)",
)
.bind("Regular default")
.bind("sk-user-bound")
.bind(user_id)
.execute(&*state.db)
.await
.expect("insert bound api key");
let Json(value) = login_api_key(
State(state.clone()),
test_addr(),
empty_headers(),
Json(LoginApiKeyRequest {
api_key: "sk-user-bound".into(),
}),
)
.await
.expect("regular api key login succeeds");
assert_eq!(value["user"]["role"], "user");
assert_eq!(value["data"]["user"]["role"], "user");
let token = value["token"].as_str().expect("token");
let claims = validate_jwt(&state.config.auth.jwt_secret, token).expect("valid jwt");
assert_eq!(claims.sub, user_id.to_string());
assert_eq!(claims.role, "user");
}
#[tokio::test]
async fn api_key_login_audits_forwarded_header_ip_and_key_name() {
let state = test_state().await;
let password_hash = hash_password("unused-password").expect("password hash");
let user_id = sqlx::query(
"INSERT INTO users (username, email, password_hash, role, status)
VALUES (?, ?, ?, 'user', 'active')",
)
.bind("Regular User")
.bind("key-audit@example.test")
.bind(password_hash)
.execute(&*state.db)
.await
.expect("insert regular user")
.last_insert_rowid();
sqlx::query(
"INSERT INTO downstream_api_keys (name, key, user_id, enabled)
VALUES (?, ?, ?, 1)",
)
.bind("Mobile Login Key")
.bind("sk-user-audit-bound")
.bind(user_id)
.execute(&*state.db)
.await
.expect("insert bound api key");
let Json(_) = login_api_key(
State(state.clone()),
test_addr(),
headers_with("forwarded", "for=\"198.51.100.24:443\";proto=https"),
Json(LoginApiKeyRequest {
api_key: "sk-user-audit-bound".into(),
}),
)
.await
.expect("regular api key login succeeds");
let (event_type, title, message, related_id, related_type) = latest_audit_event(&state).await;
assert_eq!(event_type, "user_login");
assert!(title.contains("用户登录成功"));
assert!(message.contains("方式=api_key"));
assert!(message.contains("IP=198.51.100.24"));
assert!(message.contains("密钥=Mobile Login Key"));
assert_eq!(related_id, user_id);
assert_eq!(related_type, "user");
}
}
+109
View File
@@ -6,6 +6,7 @@
use crate::{
auth::api_key::{authorize_downstream_token, DownstreamAuthResult},
auth::jwt::validate_jwt,
error::{AppError, AppResult},
AppState,
};
@@ -47,6 +48,15 @@ async fn check_auth(state: &AppState, headers: &HeaderMap) -> Result<(), AppErro
return Ok(());
}
if let Ok(claims) = validate_jwt(&state.config.auth.jwt_secret, auth) {
if claims.role == "admin" {
return Ok(());
}
return Err(AppError::Forbidden(
"Admin account is required for this route".into(),
));
}
// Downstream API key
if auth.starts_with("sk-") {
let proxy_token = state.runtime_proxy_token();
@@ -69,6 +79,105 @@ async fn check_auth(state: &AppState, headers: &HeaderMap) -> Result<(), AppErro
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{
auth::jwt::generate_jwt,
config::{AppConfig, AuthConfig, ComfyUIConfig, DbConfig, ProxyConfig, ServerConfig},
services::{builtin_keys::BuiltinKeys, runtime_auth::RuntimeAuthSettings},
};
use axum::http::HeaderValue;
use sqlx::sqlite::SqlitePoolOptions;
async fn test_state() -> AppState {
let pool = SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.expect("sqlite pool");
AppState::new(
AppConfig {
server: ServerConfig {
host: "127.0.0.1".to_string(),
port: 4000,
log_level: "info".to_string(),
cors_origins: Vec::new(),
},
db: DbConfig {
url: "sqlite::memory:".to_string(),
max_connections: 1,
ssl: false,
runtime_db_url: None,
runtime_db_dialect: None,
runtime_db_ssl: None,
},
auth: AuthConfig {
jwt_secret: "test-jwt-secret".to_string(),
jwt_expiry_hours: 24,
auth_token: "test-admin-token".to_string(),
proxy_token: "sk-test-proxy".to_string(),
},
proxy: ProxyConfig {
upstream_timeout_secs: 30,
first_byte_timeout_secs: 30,
max_retries: 1,
proxy_log_retention_days: 7,
template_proxy_url: None,
},
comfyui: ComfyUIConfig::default(),
},
pool,
BuiltinKeys {
comfyui_agent: (0, "sk-comfyui-test".to_string()),
},
RuntimeAuthSettings::from_token_for_test("sk-test-proxy"),
)
}
fn bearer_headers(token: &str) -> HeaderMap {
let mut headers = HeaderMap::new();
headers.insert(
"Authorization",
HeaderValue::from_str(&format!("Bearer {token}")).expect("valid bearer header"),
);
headers
}
#[tokio::test]
async fn comfyui_auth_accepts_admin_jwt() {
let state = test_state().await;
let token = generate_jwt(
&state.config.auth.jwt_secret,
"42",
"admin",
state.config.auth.jwt_expiry_hours,
)
.expect("admin jwt");
check_auth(&state, &bearer_headers(&token))
.await
.expect("admin jwt can access comfyui proxy routes");
}
#[tokio::test]
async fn comfyui_auth_rejects_user_jwt() {
let state = test_state().await;
let token = generate_jwt(
&state.config.auth.jwt_secret,
"43",
"user",
state.config.auth.jwt_expiry_hours,
)
.expect("user jwt");
let result = check_auth(&state, &bearer_headers(&token)).await;
assert!(matches!(result, Err(AppError::Forbidden(_))));
}
}
/// Get ComfyUI base URL from embedded/external runtime configuration.
fn comfyui_base(state: &AppState) -> String {
state.comfyui.internal_base_url()
+585 -10
View File
@@ -2,6 +2,7 @@ use crate::{
auth::jwt::validate_jwt,
db::models::{background_task::BackgroundTask, downstream_api_key::DownstreamApiKey},
error::{AppError, AppResult},
services::billing::{ComfyTaskBillingEvent, UsageTracker},
services::workspace::{self, WorkspaceAssetInput},
workflow::storage,
AppState,
@@ -9,7 +10,7 @@ use crate::{
use axum::{
body::Bytes,
extract::{Multipart, Path, Query, State},
http::{header, HeaderMap, Method},
http::{header, HeaderMap, Method, StatusCode},
routing::{get, post},
Json, Router,
};
@@ -22,11 +23,61 @@ const DIGITAL_HUMAN_TIMEOUT_SECS: u64 = 900;
const DIGITAL_HUMAN_POLL_INTERVAL_SECS: u64 = 5;
const DIGITAL_HUMAN_MAX_POLLS: usize = 180;
fn normalize_video_model(value: &str) -> String {
match value.trim() {
"seedance-1-0-lite-t2v" => "seedance-1-0-lite-t2v-250428".to_string(),
"seedance-1-0-lite-i2v" => "seedance-1-0-pro-i2v-250528".to_string(),
other => other.to_string(),
}
}
fn configured_task_cost(value: f64) -> f64 {
if value.is_finite() {
value.max(0.0)
} else {
0.0
}
}
fn truncate_chars(value: &str, limit: usize) -> String {
value.chars().take(limit).collect()
}
fn comfy_task_event(
auth: &DigitalHumanAuth,
model: String,
prompt_tokens: i64,
completion_tokens: i64,
cost: f64,
script: &str,
request: &Value,
response: Option<&Value>,
error_message: Option<String>,
) -> ComfyTaskBillingEvent {
ComfyTaskBillingEvent {
key_id: auth.key_id,
model,
task_type: TASK_TYPE_DIGITAL_HUMAN,
client_family: "digital-human",
client_app_id: "digital-human",
client_app_name: "Digital Human",
prompt_tokens,
completion_tokens,
cost,
content_summary: Some(truncate_chars(script, 240)),
request_content: Some(request.to_string()),
response_content: response.map(Value::to_string),
error_message,
}
}
#[derive(Debug, Clone)]
struct DigitalHumanAuth {
user_id: Option<i64>,
role: String,
outbound_key: String,
key_id: Option<i64>,
key_name: String,
}
#[derive(Debug, Deserialize)]
@@ -110,6 +161,8 @@ async fn resolve_auth(headers: &HeaderMap, state: &AppState) -> Result<DigitalHu
user_id: None,
role: "admin".to_string(),
outbound_key: proxy_token,
key_id: None,
key_name: "global-proxy".to_string(),
});
}
@@ -128,21 +181,26 @@ async fn resolve_auth(headers: &HeaderMap, state: &AppState) -> Result<DigitalHu
user_id: Some(user_id),
role: claims.role,
outbound_key: key.key,
key_id: Some(key.id),
key_name: key.name,
}),
Err(_) if claims.role == "admin" => Ok(DigitalHumanAuth {
user_id: Some(user_id),
role: claims.role,
outbound_key: state.runtime_proxy_token(),
key_id: None,
key_name: "global-proxy".to_string(),
}),
Err(message) => Err(AppError::Forbidden(message)),
}
}
async fn status(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
let _ = resolve_auth(&headers, &state).await?;
let auth = resolve_auth(&headers, &state).await?;
Ok(Json(json!({
"success": true,
"mode": if state.config.comfyui.uses_native_backend() { "comfy-native" } else { "boosapi-video-route" },
"keyName": auth.key_name,
"runninghub": {
"configured": state.config.comfyui.api_key.as_deref().map(str::trim).is_some_and(|value| !value.is_empty()),
"endpoint": "https://www.runninghub.cn/task/openapi/ai-app/run"
@@ -486,6 +544,7 @@ fn default_mapping() -> Value {
json!({
"prompt": { "nodeId": "", "field": "prompt" },
"script": { "nodeId": "", "field": "text" },
"model": { "nodeId": "", "field": "model" },
"avatar": { "nodeId": "", "field": "image" },
"referenceAudio": { "nodeId": "", "field": "audio" },
"voiceReference": { "nodeId": "", "field": "audio" },
@@ -520,11 +579,106 @@ fn apply_node_input(prompt: &mut Value, node_id: &str, field: &str, value: Value
Ok(())
}
fn graph_node_type(node: &Value) -> String {
node.get("type")
.or_else(|| node.get("class_type"))
.and_then(Value::as_str)
.unwrap_or_default()
.to_string()
}
fn graph_node_params(node: &Value) -> serde_json::Map<String, Value> {
node.get("params")
.and_then(Value::as_object)
.cloned()
.or_else(|| node.get("inputs").and_then(Value::as_object).cloned())
.unwrap_or_default()
}
fn normalize_comfy_value(class_type: &str, field: &str, value: Value) -> Value {
match (class_type, field) {
("KSampler", "cfg") => value
.as_f64()
.filter(|value| (0.0..=50.0).contains(value))
.map_or(json!(8.0), |value| json!(value)),
("KSampler", "denoise") => value
.as_f64()
.map(|value| if value > 1.0 { value / 100.0 } else { value })
.filter(|value| (0.0..=1.0).contains(value))
.map_or(json!(1.0), |value| json!(value)),
("KSampler", "sampler_name") => {
if value.is_string() {
value
} else {
json!("euler")
}
}
("KSampler", "scheduler") => {
if value.is_string() {
value
} else {
json!("normal")
}
}
("KSampler", "seed") => value.as_i64().map_or(json!(0), |value| json!(value.max(0))),
("EmptyLatentImage", "width") | ("EmptyLatentImage", "height") => value
.as_i64()
.filter(|value| (64..=4096).contains(value))
.map_or(json!(1024), |value| json!(value)),
("EmptyLatentImage", "batch_size") => value
.as_i64()
.filter(|value| (1..=16).contains(value))
.map_or(json!(1), |value| json!(value)),
_ => value,
}
}
fn split_endpoint(value: &str) -> Option<(&str, &str)> {
let (node_id, port) = value.split_once('.')?;
Some((node_id.trim(), port.trim()))
}
fn graph_to_comfy_prompt(graph: &Value) -> Option<Value> {
let nodes = graph.get("nodes")?.as_array()?;
let mut prompt = serde_json::Map::new();
for node in nodes {
let id = scalar_string(node.get("id").or_else(|| node.get("nodeId")))?;
let class_type = graph_node_type(node);
if class_type.trim().is_empty() {
continue;
}
let mut inputs = serde_json::Map::new();
for (field, value) in graph_node_params(node) {
inputs.insert(field.clone(), normalize_comfy_value(&class_type, &field, value));
}
prompt.insert(id, json!({
"class_type": class_type,
"inputs": inputs,
}));
}
for edge in graph.get("edges").and_then(Value::as_array).into_iter().flatten() {
let Some(from) = edge.get("from").and_then(Value::as_str) else { continue; };
let Some(to) = edge.get("to").and_then(Value::as_str) else { continue; };
let Some((from_id, _from_port)) = split_endpoint(from) else { continue; };
let Some((to_id, to_field)) = split_endpoint(to) else { continue; };
if let Some(inputs) = prompt
.get_mut(to_id)
.and_then(Value::as_object_mut)
.and_then(|node| node.get_mut("inputs"))
.and_then(Value::as_object_mut)
{
inputs.insert(to_field.to_string(), json!([from_id, 0]));
}
}
Some(Value::Object(prompt))
}
fn workflow_prompt_from_graph(graph: &Value) -> Value {
graph.get("prompt")
.cloned()
.or_else(|| graph.get("workflow").cloned())
.or_else(|| graph.get("comfyPrompt").cloned())
.or_else(|| graph_to_comfy_prompt(graph))
.unwrap_or_else(|| graph.clone())
}
@@ -583,6 +737,16 @@ fn inject_digital_human_inputs(
&mapping_field(mapping, &["script", "scriptNode", "script_node"], "text"),
json!(product_script),
)?;
if let Some(model) = optional_string(request, &["comfyApiModel", "comfy_api_model", "apiModel", "api_model"]) {
if let Some(model_node) = mapping_node_id(mapping, &["model", "modelNode", "apiModel"]) {
apply_node_input(
prompt,
&model_node,
&mapping_field(mapping, &["model", "modelNode", "apiModel"], "model"),
json!(normalize_video_model(&model)),
)?;
}
}
apply_node_input(
prompt,
&avatar_node,
@@ -677,6 +841,13 @@ fn native_request(state: &AppState, method: Method, path: &str) -> reqwest::Requ
request
}
fn is_pending_history_status(status: StatusCode) -> bool {
matches!(
status,
StatusCode::NOT_FOUND | StatusCode::ACCEPTED | StatusCode::NO_CONTENT
)
}
async fn native_json(
state: &AppState,
method: Method,
@@ -703,6 +874,35 @@ async fn native_json(
Ok(value)
}
async fn native_json_optional(
state: &AppState,
method: Method,
path: &str,
body: Option<Value>,
) -> AppResult<Option<Value>> {
let mut request = native_request(state, method, path);
if let Some(body) = body {
request = request.json(&body);
}
let response = request
.send()
.await
.map_err(|err| AppError::UpstreamError(format!("Comfy backend unreachable: {err}")))?;
let status = response.status();
if is_pending_history_status(status) {
return Ok(None);
}
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(Some(value))
}
async fn submit_prompt(state: &AppState, prompt: Value) -> AppResult<(String, Value)> {
let body = json!({
"client_id": format!("boosapi-digital-human-{}", uuid::Uuid::new_v4().simple()),
@@ -720,13 +920,13 @@ async fn submit_prompt(state: &AppState, prompt: Value) -> AppResult<(String, Va
Ok((prompt_id, value))
}
async fn prompt_history(state: &AppState, prompt_id: &str) -> AppResult<Value> {
async fn prompt_history_optional(state: &AppState, prompt_id: &str) -> AppResult<Option<Value>> {
let path = if state.config.comfyui.uses_cloud_backend() {
format!("/jobs/{prompt_id}")
} else {
format!("/history/{prompt_id}")
};
native_json(state, Method::GET, &path, None).await
native_json_optional(state, Method::GET, &path, None).await
}
fn prompt_failure_message(value: &Value) -> Option<String> {
@@ -746,12 +946,24 @@ fn prompt_failure_message(value: &Value) -> Option<String> {
fn output_objects(value: &Value) -> Vec<Value> {
let mut out = Vec::new();
fn push_unique(out: &mut Vec<Value>, item: Value) {
let identity = output_identity(&item);
if out.iter().any(|existing| output_identity(existing) == identity) {
return;
}
out.push(item);
}
fn walk(value: &Value, out: &mut Vec<Value>) {
match value {
Value::Object(object) => {
if output_filename(value).is_some() || output_remote_url(value).is_some() {
push_unique(out, value.clone());
}
for key in ["videos", "gifs", "images", "audio"] {
if let Some(items) = object.get(key).and_then(Value::as_array) {
out.extend(items.iter().filter(|item| item.is_object()).cloned());
for item in items {
walk(item, out);
}
}
}
for value in object.values() {
@@ -763,6 +975,11 @@ fn output_objects(value: &Value) -> Vec<Value> {
walk(item, out);
}
}
Value::String(value) => {
if let Some(item) = output_item_from_reference(value) {
push_unique(out, item);
}
}
_ => {}
}
}
@@ -770,10 +987,112 @@ fn output_objects(value: &Value) -> Vec<Value> {
out
}
fn output_identity(item: &Value) -> String {
output_remote_url(item)
.or_else(|| output_filename(item))
.unwrap_or_else(|| item.to_string())
}
fn looks_like_media_reference(value: &str) -> bool {
let trimmed = value.trim();
if !(trimmed.starts_with("http://") || trimmed.starts_with("https://") || trimmed.starts_with('/')) {
return false;
}
let lower = trimmed
.split(['?', '#'])
.next()
.unwrap_or(trimmed)
.to_ascii_lowercase();
lower.contains("/view")
|| lower.ends_with(".mp4")
|| lower.ends_with(".webm")
|| lower.ends_with(".mov")
|| lower.ends_with(".m4v")
|| lower.ends_with(".gif")
|| lower.ends_with(".png")
|| lower.ends_with(".jpg")
|| lower.ends_with(".jpeg")
|| lower.ends_with(".webp")
|| lower.ends_with(".mp3")
|| lower.ends_with(".wav")
|| lower.ends_with(".m4a")
}
fn output_item_from_reference(value: &str) -> Option<Value> {
let trimmed = value.trim();
if !looks_like_media_reference(trimmed) {
return None;
}
let filename = trimmed
.split(['?', '#'])
.next()
.unwrap_or(trimmed)
.rsplit('/')
.next()
.unwrap_or("digital-human-output")
.to_string();
Some(json!({
"url": trimmed,
"filename": filename,
"type": "output"
}))
}
fn output_filename(item: &Value) -> Option<String> {
["filename", "fileName", "name", "file"]
.iter()
.find_map(|key| scalar_string(item.get(*key)))
}
fn output_remote_url(item: &Value) -> Option<String> {
[
"url",
"remoteUrl",
"remote_url",
"videoUrl",
"video_url",
"audioUrl",
"audio_url",
"imageUrl",
"image_url",
"fileUrl",
"file_url",
"downloadUrl",
"download_url",
"video",
"audio",
"image",
]
.iter()
.find_map(|key| scalar_string(item.get(*key)))
.filter(|value| looks_like_media_reference(value))
}
fn output_kind(item: &Value) -> &'static str {
let filename = scalar_string(item.get("filename")).unwrap_or_default();
let lower = filename.to_ascii_lowercase();
if lower.ends_with(".mp4") || lower.ends_with(".webm") || lower.ends_with(".mov") {
let explicit = scalar_string(item.get("kind"))
.or_else(|| scalar_string(item.get("mediaType")))
.or_else(|| scalar_string(item.get("mime")))
.unwrap_or_default()
.to_ascii_lowercase();
if explicit.contains("video") {
return "video";
}
if explicit.contains("audio") {
return "audio";
}
if explicit.contains("image") {
return "image";
}
let lower = output_filename(item)
.or_else(|| output_remote_url(item))
.unwrap_or_default()
.to_ascii_lowercase();
if lower.ends_with(".mp4")
|| lower.ends_with(".webm")
|| lower.ends_with(".mov")
|| lower.ends_with(".m4v")
|| lower.ends_with(".gif")
{
"video"
} else if lower.ends_with(".mp3") || lower.ends_with(".wav") || lower.ends_with(".m4a") {
"audio"
@@ -784,7 +1103,7 @@ fn output_kind(item: &Value) -> &'static str {
fn output_query(item: &Value) -> HashMap<String, String> {
let mut params = HashMap::new();
if let Some(filename) = scalar_string(item.get("filename")) {
if let Some(filename) = output_filename(item) {
params.insert("filename".to_string(), filename);
}
if let Some(file_type) = scalar_string(item.get("type")) {
@@ -829,6 +1148,56 @@ async fn archive_output(
prompt_id: &str,
script: &str,
) -> Option<Value> {
if let Some(remote_url) = output_remote_url(item) {
let filename = output_filename(item).unwrap_or_else(|| {
remote_url
.split(['?', '#'])
.next()
.unwrap_or(&remote_url)
.rsplit('/')
.next()
.unwrap_or("digital-human-output")
.to_string()
});
let kind = output_kind(item).to_string();
let input = WorkspaceAssetInput {
user_id: auth.user_id,
source: "digital-human".to_string(),
kind: kind.clone(),
title: Some(filename.chars().take(80).collect()),
prompt: Some(script.chars().take(1000).collect()),
model: Some("comfy-native".to_string()),
remote_url: Some(remote_url.clone()),
metadata: json!({
"workspaceAsset": { "role": "output" },
"digitalHuman": {
"promptId": prompt_id,
"filename": filename,
"raw": item,
}
}),
};
return match workspace::record_remote_asset(state, input).await {
Ok(asset) => Some(json!({
"publicId": asset.public_id,
"kind": kind,
"filename": filename,
"status": asset.status,
"remoteUrl": asset.remote_url,
"fileName": asset.file_name,
})),
Err(err) => {
tracing::warn!(%err, "failed to archive remote digital human output");
Some(json!({
"kind": kind,
"filename": filename,
"status": "failed",
"remoteUrl": remote_url,
"error": err.to_string(),
}))
}
};
}
let params = output_query(item);
let filename = params.get("filename")?.clone();
let kind = output_kind(item).to_string();
@@ -963,6 +1332,9 @@ async fn create_job(
Json(body): Json<Value>,
) -> AppResult<Json<Value>> {
let auth = resolve_auth(&headers, &state).await?;
if let Some(key_id) = auth.key_id {
UsageTracker::new(state.db.clone()).check_rate_limit(key_id).await?;
}
if !state.config.comfyui.uses_native_backend() || state.config.comfyui.uses_runninghub_backend()
{
return Err(AppError::BadRequest(
@@ -985,8 +1357,11 @@ async fn create_job(
let initial_result = json!({
"userId": auth.user_id,
"role": auth.role,
"keyId": auth.key_id,
"keyName": auth.key_name,
"templateId": template_id,
"templateName": workflow.name,
"script": script,
"audioRole": optional_string(&body, &["audioRole", "audio_role"]).unwrap_or_else(|| "both".to_string()),
"request": body,
"mapping": mapping,
@@ -1008,6 +1383,14 @@ async fn create_job(
let task_id_for_spawn = task_id.clone();
let auth_for_task = auth.clone();
tokio::spawn(async move {
let failure_body = body.clone();
let failure_script = script.clone();
let billing_model = optional_string(
&failure_body,
&["comfyApiModel", "comfy_api_model", "apiModel", "api_model", "model"],
)
.map(|value| normalize_video_model(&value))
.unwrap_or_else(|| workflow.name.clone());
let result = async {
update_progress(&state_for_task, &task_id_for_spawn, 2, "正在提交 Comfy 工作流")
.await?;
@@ -1015,8 +1398,11 @@ async fn create_job(
let mid_result = json!({
"userId": auth_for_task.user_id,
"role": auth_for_task.role,
"keyId": auth_for_task.key_id,
"keyName": auth_for_task.key_name,
"templateId": template_id,
"templateName": workflow.name,
"script": script,
"promptId": prompt_id,
"submitResponse": submit_response,
"request": body,
@@ -1036,7 +1422,12 @@ async fn create_job(
}
update_progress(&state_for_task, &task_id_for_spawn, 3, "Comfy 正在生成视频")
.await?;
let history = prompt_history(&state_for_task, &prompt_id).await?;
let Some(history) = prompt_history_optional(&state_for_task, &prompt_id).await?
else {
tokio::time::sleep(Duration::from_secs(DIGITAL_HUMAN_POLL_INTERVAL_SECS))
.await;
continue;
};
if let Some(message) = prompt_failure_message(&history) {
return Err(AppError::UpstreamError(message));
}
@@ -1063,17 +1454,62 @@ async fn create_job(
assets.push(asset);
}
}
let tracker = UsageTracker::new(state_for_task.db.clone());
let billing_response = json!({
"promptId": prompt_id,
"history": history,
"outputs": outputs,
"assets": assets,
});
let billing = match tracker
.record_comfy_task_success(comfy_task_event(
&auth_for_task,
billing_model.clone(),
0,
0,
configured_task_cost(state_for_task.config.comfyui.digital_human_video_cost),
&script,
&body,
Some(&billing_response),
None,
))
.await
{
Ok(details) => details,
Err(err) => {
tracing::warn!(
%err,
task_id = %task_id_for_spawn,
"failed to record digital human comfy billing",
);
json!({
"source": "comfy-native-task",
"taskType": TASK_TYPE_DIGITAL_HUMAN,
"status": "success",
"charged": false,
"failed": true,
"error": err.to_string(),
"totalCost": 0.0,
"steps": []
})
}
};
let final_result = json!({
"userId": auth_for_task.user_id,
"role": auth_for_task.role,
"keyId": auth_for_task.key_id,
"keyName": auth_for_task.key_name,
"templateId": template_id,
"templateName": workflow.name,
"model": billing_model.clone(),
"script": script,
"promptId": prompt_id,
"request": body,
"mapping": mapping,
"history": history,
"outputs": outputs,
"assets": assets,
"billing": billing,
});
sqlx::query(
"UPDATE background_tasks
@@ -1093,6 +1529,27 @@ async fn create_job(
}
.await;
if let Err(err) = result {
let tracker = UsageTracker::new(state_for_task.db.clone());
if let Err(log_err) = tracker
.record_comfy_task_failure(comfy_task_event(
&auth_for_task,
billing_model,
0,
0,
0.0,
&failure_script,
&failure_body,
None,
Some(err.to_string()),
))
.await
{
tracing::warn!(
%log_err,
task_id = %task_id_for_spawn,
"failed to record digital human comfy failure log",
);
}
let _ = sqlx::query(
"UPDATE background_tasks
SET status = 'failed', message = '数字人任务失败', error = ?, updated_at = datetime('now'), finished_at = datetime('now')
@@ -1278,3 +1735,121 @@ async fn upload_asset(
}
})))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn workflow_prompt_from_graph_converts_node_edge_workflow_to_comfy_prompt() {
let graph = json!({
"nodes": [
{ "id": "1", "type": "LoadImage", "params": { "image": "avatar.png" } },
{ "id": "2", "type": "VHS_VideoCombine", "params": { "filename_prefix": "digital-human" } }
],
"edges": [
{ "from": "1.IMAGE", "to": "2.images" }
]
});
let prompt = workflow_prompt_from_graph(&graph);
assert_eq!(prompt.pointer("/1/class_type").and_then(Value::as_str), Some("LoadImage"));
assert_eq!(prompt.pointer("/2/class_type").and_then(Value::as_str), Some("VHS_VideoCombine"));
assert_eq!(prompt.pointer("/2/inputs/images"), Some(&json!(["1", 0])));
}
#[test]
fn output_objects_collects_cloud_and_native_media_shapes() {
let history = json!({
"data": {
"outputs": [
"https://cdn.example.com/avatar-final.mp4",
{ "audioUrl": "https://cdn.example.com/voice.wav" }
]
},
"outputs": {
"9": {
"images": [
{ "filename": "cover.webp", "subfolder": "previews", "type": "output" }
]
}
}
});
let outputs = output_objects(&history);
assert_eq!(outputs.len(), 3);
assert!(outputs.iter().any(|item| output_remote_url(item).as_deref() == Some("https://cdn.example.com/avatar-final.mp4")));
assert!(outputs.iter().any(|item| output_kind(item) == "audio"));
assert!(outputs.iter().any(|item| output_filename(item).as_deref() == Some("cover.webp")));
}
#[test]
fn output_objects_deduplicates_repeated_remote_urls() {
let history = json!({
"outputs": {
"1": {
"videoUrl": "https://cdn.example.com/final.mp4",
"videos": [
{ "url": "https://cdn.example.com/final.mp4", "filename": "final.mp4" }
]
}
}
});
let outputs = output_objects(&history);
assert_eq!(outputs.len(), 1);
assert_eq!(output_remote_url(&outputs[0]).as_deref(), Some("https://cdn.example.com/final.mp4"));
}
#[test]
fn history_polling_treats_not_ready_statuses_as_pending() {
assert!(is_pending_history_status(StatusCode::NOT_FOUND));
assert!(is_pending_history_status(StatusCode::ACCEPTED));
assert!(is_pending_history_status(StatusCode::NO_CONTENT));
assert!(!is_pending_history_status(StatusCode::OK));
assert!(!is_pending_history_status(StatusCode::BAD_GATEWAY));
}
#[test]
fn inject_digital_human_inputs_writes_optional_comfy_api_model() {
let mut prompt = json!({
"1": { "class_type": "PromptTextNode", "inputs": { "prompt": "" } },
"2": { "class_type": "TTSVoiceNode", "inputs": { "text": "" } },
"3": { "class_type": "LoadImage", "inputs": { "image": "" } },
"4": { "class_type": "LoadAudio", "inputs": { "audio": "" } },
"5": { "class_type": "ByteDanceDigitalHumanNode", "inputs": { "model": "seedance-1-0-lite-t2v-250428" } }
});
let mapping = json!({
"prompt": { "nodeId": "1", "field": "prompt" },
"script": { "nodeId": "2", "field": "text" },
"avatar": { "nodeId": "3", "field": "image" },
"referenceAudio": { "nodeId": "4", "field": "audio" },
"model": { "nodeId": "5", "field": "model" }
});
let request = json!({
"avatarUrl": "https://cdn.example.com/avatar.png",
"referenceAudioUrl": "https://cdn.example.com/voice.wav",
"productScript": "介绍新款咖啡机",
"prompt": "真人主播讲解商品",
"audioRole": "both",
"comfyApiModel": "seedance-1-0-lite-t2v"
});
let script = inject_digital_human_inputs(&mut prompt, &mapping, &request)
.expect("digital human inputs should be injected");
assert_eq!(script, "介绍新款咖啡机");
assert_eq!(
prompt.pointer("/5/inputs/model").and_then(Value::as_str),
Some("seedance-1-0-lite-t2v-250428")
);
assert_eq!(
prompt.pointer("/1/inputs/prompt").and_then(Value::as_str),
Some("真人主播讲解商品")
);
assert_eq!(
prompt.pointer("/3/inputs/image").and_then(Value::as_str),
Some("https://cdn.example.com/avatar.png")
);
}
}
+595 -10
View File
@@ -1,8 +1,16 @@
// Events routes — system event notifications
use crate::{db::models::event::Event, error::AppResult, services::runtime_logs, AppState};
use crate::{
db::models::event::Event,
error::{AppError, AppResult},
services::runtime_logs,
AppState,
};
use axum::{
extract::{Path, Query, State},
http::HeaderMap,
middleware,
response::Response,
routing::{get, post},
Json, Router,
};
@@ -20,12 +28,35 @@ pub struct ListEventsQuery {
read: Option<String>, // "true" | "false"
}
pub fn routes() -> Router<AppState> {
pub fn routes(state: AppState) -> Router<AppState> {
Router::new()
.route("/api/events", get(list_events).delete(clear_events))
.route("/api/events/count", get(unread_count))
.route("/api/events/{id}/read", post(mark_read))
.route("/api/events/read-all", post(mark_all_read))
.route_layer(middleware::from_fn_with_state(
state,
require_events_admin_auth,
))
}
async fn require_events_admin_auth(
State(state): State<AppState>,
headers: HeaderMap,
request: axum::http::Request<axum::body::Body>,
next: axum::middleware::Next,
) -> Result<Response, AppError> {
let auth_header = headers
.get("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_header != state.config.auth.auth_token {
return Err(AppError::Unauthorized("Invalid admin token".into()));
}
Ok(next.run(request).await)
}
/// GET /api/events
@@ -36,15 +67,26 @@ async fn list_events(
debug!("Listing events");
let limit = q.limit.unwrap_or(30).clamp(1, 500);
let offset = q.offset.unwrap_or(0).max(0);
let filter_type = q
.filter_type
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_string);
let read_filter = q.read.as_deref();
let source_limit = offset
.saturating_add(limit)
.saturating_add(200)
.clamp(50, 1000);
let mut sql = "SELECT id, type, title, message, level, read, related_id, related_type, created_at FROM events WHERE 1=1".to_string();
let mut params: Vec<String> = vec![];
if let Some(ref t) = q.filter_type {
if let Some(ref t) = filter_type {
sql.push_str(" AND type = ?");
params.push(t.clone());
}
if let Some(ref r) = q.read {
if let Some(r) = read_filter {
if r == "true" {
sql.push_str(" AND read = 1");
} else if r == "false" {
@@ -61,18 +103,34 @@ async fn list_events(
let events: Vec<Event> = query.fetch_all(&*state.db).await?;
let mut rows: Vec<EventListItem> = events.into_iter().map(EventListItem::from).collect();
if q.filter_type
.as_deref()
.is_none_or(|filter_type| filter_type == "runtime")
{
if should_include_source(filter_type.as_deref(), "runtime") {
rows.extend(
runtime_logs::list()
.into_iter()
.filter(|entry| q.read.as_deref() != Some("false") || !entry.is_read)
.filter(|entry| read_filter != Some("false") || !entry.is_read)
.map(EventListItem::from),
);
}
if should_include_read_only_source(filter_type.as_deref(), read_filter, "checkin") {
rows.extend(load_checkin_log_events(&state, source_limit).await?);
}
if should_include_read_only_source(filter_type.as_deref(), read_filter, "proxy") {
rows.extend(load_proxy_log_events(&state, source_limit).await?);
}
if should_include_read_only_source(filter_type.as_deref(), read_filter, "balance") {
rows.extend(load_balance_log_events(&state, source_limit).await?);
}
if should_include_read_only_source(filter_type.as_deref(), read_filter, "token") {
rows.extend(load_token_log_events(&state, source_limit).await?);
}
if should_include_read_only_source(filter_type.as_deref(), read_filter, "site_notice") {
rows.extend(load_site_notice_events(&state, source_limit).await?);
}
if should_include_read_only_source(filter_type.as_deref(), read_filter, "status") {
rows.extend(load_status_log_events(&state, source_limit).await?);
}
rows.sort_by(|a, b| {
b.created_at
.cmp(&a.created_at)
@@ -89,6 +147,253 @@ async fn list_events(
Ok(Json(json!({ "data": page })))
}
fn should_include_source(filter_type: Option<&str>, source: &str) -> bool {
filter_type.is_none_or(|value| value == source)
}
fn should_include_read_only_source(
filter_type: Option<&str>,
read_filter: Option<&str>,
source: &str,
) -> bool {
should_include_source(filter_type, source) && read_filter != Some("false")
}
async fn load_checkin_log_events(state: &AppState, limit: i64) -> AppResult<Vec<EventListItem>> {
let rows = sqlx::query_as::<_, EventListItem>(
r#"
SELECT
(-10000000000 - c.id) AS id,
'checkin' AS event_type,
CASE
WHEN lower(COALESCE(c.status, '')) IN ('success', 'ok', 'completed') THEN '签到成功'
WHEN lower(COALESCE(c.status, '')) IN ('skipped', 'skip') THEN '签到跳过'
ELSE '签到失败'
END || ':' || COALESCE(NULLIF(s.name, ''), '账号 #' || c.account_id) AS title,
COALESCE(NULLIF(a.username, ''), '账号 #' || c.account_id)
|| ' · 状态=' || COALESCE(NULLIF(c.status, ''), '-')
|| CASE WHEN c.reward IS NOT NULL AND trim(c.reward) != '' THEN ' · 奖励=' || c.reward ELSE '' END
|| CASE WHEN c.message IS NOT NULL AND trim(c.message) != '' THEN ' · ' || c.message ELSE '' END AS message,
CASE
WHEN lower(COALESCE(c.status, '')) IN ('success', 'ok', 'completed') THEN 'info'
WHEN lower(COALESCE(c.status, '')) IN ('skipped', 'skip') THEN 'warning'
ELSE 'error'
END AS level,
1 AS is_read,
c.id AS related_id,
'checkin_log' AS related_type,
c.created_at AS created_at
FROM checkin_logs c
LEFT JOIN accounts a ON a.id = c.account_id
LEFT JOIN sites s ON s.id = a.site_id
ORDER BY c.created_at DESC
LIMIT ?
"#,
)
.bind(limit)
.fetch_all(&*state.db)
.await?;
Ok(rows)
}
async fn load_proxy_log_events(state: &AppState, limit: i64) -> AppResult<Vec<EventListItem>> {
let rows = sqlx::query_as::<_, EventListItem>(
r#"
SELECT
(-20000000000 - l.id) AS id,
'proxy' AS event_type,
CASE
WHEN lower(COALESCE(l.status, '')) = 'success' THEN '代理请求成功'
ELSE '代理请求失败'
END || ':' || COALESCE(NULLIF(l.model_requested, ''), NULLIF(l.model_actual, ''), '未知模型') AS title,
'站点=' || COALESCE(NULLIF(s.name, ''), '-')
|| ' · 账号=' || COALESCE(NULLIF(a.username, ''), CASE WHEN l.account_id IS NOT NULL THEN '账号 #' || l.account_id ELSE '-' END)
|| ' · 状态=' || COALESCE(NULLIF(l.status, ''), '-')
|| CASE WHEN l.http_status IS NOT NULL THEN ' · HTTP ' || l.http_status ELSE '' END
|| CASE WHEN l.latency_ms IS NOT NULL THEN ' · ' || l.latency_ms || 'ms' ELSE '' END
|| CASE WHEN l.total_tokens IS NOT NULL THEN ' · Tokens=' || l.total_tokens ELSE '' END
|| CASE WHEN l.estimated_cost IS NOT NULL THEN ' · 成本=' || printf('$%.6f', l.estimated_cost) ELSE '' END
|| CASE WHEN l.error_message IS NOT NULL AND trim(l.error_message) != '' THEN ' · ' || l.error_message ELSE '' END AS message,
CASE WHEN lower(COALESCE(l.status, '')) = 'success' THEN 'info' ELSE 'error' END AS level,
1 AS is_read,
l.id AS related_id,
'proxy_log' AS related_type,
l.created_at AS created_at
FROM proxy_logs l
LEFT JOIN accounts a ON a.id = l.account_id
LEFT JOIN sites s ON s.id = a.site_id
ORDER BY l.created_at DESC
LIMIT ?
"#,
)
.bind(limit)
.fetch_all(&*state.db)
.await?;
Ok(rows)
}
async fn load_balance_log_events(state: &AppState, limit: i64) -> AppResult<Vec<EventListItem>> {
let rows = sqlx::query_as::<_, EventListItem>(
r#"
SELECT
(-30000000000 - a.id) AS id,
'balance' AS event_type,
'余额已刷新:' || COALESCE(NULLIF(s.name, ''), '账号 #' || a.id) AS title,
COALESCE(NULLIF(a.username, ''), '账号 #' || a.id)
|| ' · 当前余额=' || printf('%.6f', COALESCE(a.balance, 0))
|| ' · 已用=' || printf('%.6f', COALESCE(a.balance_used, 0))
|| ' · 额度=' || printf('%.6f', COALESCE(a.quota, 0)) AS message,
'info' AS level,
1 AS is_read,
a.id AS related_id,
'account' AS related_type,
a.last_balance_refresh AS created_at
FROM accounts a
LEFT JOIN sites s ON s.id = a.site_id
WHERE a.last_balance_refresh IS NOT NULL
AND trim(a.last_balance_refresh) != ''
ORDER BY a.last_balance_refresh DESC
LIMIT ?
"#,
)
.bind(limit)
.fetch_all(&*state.db)
.await?;
Ok(rows)
}
async fn load_token_log_events(state: &AppState, limit: i64) -> AppResult<Vec<EventListItem>> {
let rows = sqlx::query_as::<_, EventListItem>(
r#"
SELECT
(-40000000000 - t.id) AS id,
'token' AS event_type,
CASE
WHEN COALESCE(t.updated_at, t.created_at, '') != COALESCE(t.created_at, '') THEN '令牌已更新'
ELSE '令牌已创建'
END || ':' || COALESCE(NULLIF(t.name, ''), 'Token #' || t.id) AS title,
COALESCE(NULLIF(s.name, ''), '-')
|| ' · ' || COALESCE(NULLIF(a.username, ''), '账号 #' || t.account_id)
|| ' · 状态=' || COALESCE(NULLIF(t.value_status, ''), '-')
|| CASE WHEN COALESCE(t.enabled, 1) = 1 THEN ' · 已启用' ELSE ' · 已禁用' END
|| CASE WHEN t.token_group IS NOT NULL AND trim(t.token_group) != '' THEN ' · 分组=' || t.token_group ELSE '' END
|| CASE WHEN t.source IS NOT NULL AND trim(t.source) != '' THEN ' · 来源=' || t.source ELSE '' END AS message,
CASE
WHEN COALESCE(t.enabled, 1) = 1
AND lower(COALESCE(t.value_status, 'ready')) IN ('ready', 'valid') THEN 'info'
ELSE 'warning'
END AS level,
1 AS is_read,
t.id AS related_id,
'account_token' AS related_type,
COALESCE(t.updated_at, t.created_at) AS created_at
FROM account_tokens t
LEFT JOIN accounts a ON a.id = t.account_id
LEFT JOIN sites s ON s.id = a.site_id
WHERE COALESCE(t.updated_at, t.created_at) IS NOT NULL
AND trim(COALESCE(t.updated_at, t.created_at)) != ''
ORDER BY COALESCE(t.updated_at, t.created_at) DESC
LIMIT ?
"#,
)
.bind(limit)
.fetch_all(&*state.db)
.await?;
Ok(rows)
}
async fn load_site_notice_events(state: &AppState, limit: i64) -> AppResult<Vec<EventListItem>> {
let rows = sqlx::query_as::<_, EventListItem>(
r#"
SELECT
(-50000000000 - n.id) AS id,
'site_notice' AS event_type,
'站点公告:' || COALESCE(NULLIF(n.title, ''), '未命名公告') AS title,
COALESCE(NULLIF(s.name, ''), '站点 #' || n.site_id)
|| CASE WHEN n.content IS NOT NULL AND trim(n.content) != '' THEN ' · ' || substr(n.content, 1, 420) ELSE '' END AS message,
COALESCE(NULLIF(n.level, ''), 'info') AS level,
1 AS is_read,
n.id AS related_id,
'site_announcement' AS related_type,
COALESCE(n.upstream_created_at, n.first_seen_at) AS created_at
FROM site_announcements n
LEFT JOIN sites s ON s.id = n.site_id
ORDER BY COALESCE(n.upstream_created_at, n.first_seen_at) DESC
LIMIT ?
"#,
)
.bind(limit)
.fetch_all(&*state.db)
.await?;
Ok(rows)
}
async fn load_status_log_events(state: &AppState, limit: i64) -> AppResult<Vec<EventListItem>> {
let per_source_limit = limit.clamp(50, 1000);
let mut rows = sqlx::query_as::<_, EventListItem>(
r#"
SELECT
(-60000000000 - s.id) AS id,
'status' AS event_type,
'站点状态:' || COALESCE(NULLIF(s.name, ''), '站点 #' || s.id) AS title,
'平台=' || COALESCE(NULLIF(s.platform, ''), '-')
|| ' · 状态=' || COALESCE(NULLIF(s.status, ''), '-')
|| CASE WHEN s.url IS NOT NULL AND trim(s.url) != '' THEN ' · ' || s.url ELSE '' END AS message,
CASE
WHEN lower(COALESCE(s.status, 'active')) IN ('active', 'enabled', 'healthy') THEN 'info'
WHEN lower(COALESCE(s.status, '')) IN ('disabled', 'inactive') THEN 'warning'
ELSE 'error'
END AS level,
1 AS is_read,
s.id AS related_id,
'site' AS related_type,
COALESCE(s.updated_at, s.created_at) AS created_at
FROM sites s
WHERE COALESCE(s.updated_at, s.created_at) IS NOT NULL
AND trim(COALESCE(s.updated_at, s.created_at)) != ''
ORDER BY COALESCE(s.updated_at, s.created_at) DESC
LIMIT ?
"#,
)
.bind(per_source_limit)
.fetch_all(&*state.db)
.await?;
rows.extend(
sqlx::query_as::<_, EventListItem>(
r#"
SELECT
(-70000000000 - a.id) AS id,
'status' AS event_type,
'账号状态:' || COALESCE(NULLIF(s.name, ''), '站点 #' || a.site_id) AS title,
COALESCE(NULLIF(a.username, ''), '账号 #' || a.id)
|| ' · 状态=' || COALESCE(NULLIF(a.status, ''), '-')
|| CASE WHEN a.oauth_provider IS NOT NULL AND trim(a.oauth_provider) != '' THEN ' · OAuth=' || a.oauth_provider ELSE '' END AS message,
CASE
WHEN lower(COALESCE(a.status, 'active')) IN ('active', 'enabled', 'healthy') THEN 'info'
WHEN lower(COALESCE(a.status, '')) IN ('disabled', 'inactive') THEN 'warning'
ELSE 'error'
END AS level,
1 AS is_read,
a.id AS related_id,
'account' AS related_type,
COALESCE(a.updated_at, a.created_at) AS created_at
FROM accounts a
LEFT JOIN sites s ON s.id = a.site_id
WHERE COALESCE(a.updated_at, a.created_at) IS NOT NULL
AND trim(COALESCE(a.updated_at, a.created_at)) != ''
ORDER BY COALESCE(a.updated_at, a.created_at) DESC
LIMIT ?
"#,
)
.bind(per_source_limit)
.fetch_all(&*state.db)
.await?,
);
Ok(rows)
}
/// GET /api/events/count — unread count
async fn unread_count(State(state): State<AppState>) -> AppResult<Json<Value>> {
let count: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM events WHERE read = 0")
@@ -126,7 +431,7 @@ async fn clear_events(State(state): State<AppState>) -> AppResult<Json<Value>> {
Ok(Json(json!({ "success": true })))
}
#[derive(Debug, Clone, Serialize)]
#[derive(Debug, Clone, Serialize, sqlx::FromRow)]
#[serde(rename_all = "camelCase")]
struct EventListItem {
id: i64,
@@ -173,3 +478,283 @@ impl From<runtime_logs::RuntimeLogEntry> for EventListItem {
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{
auth::jwt::generate_jwt,
config::{AppConfig, AuthConfig, DbConfig, ProxyConfig, ServerConfig},
services::{builtin_keys::BuiltinKeys, runtime_auth::RuntimeAuthSettings},
};
use axum::{
body::Body,
http::{Request, StatusCode},
};
use axum::extract::{Query, State};
use sqlx::sqlite::SqlitePoolOptions;
use tower::ServiceExt as _;
async fn test_state() -> AppState {
let pool = SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.expect("memory sqlite");
sqlx::migrate!("./migrations")
.run(&pool)
.await
.expect("run migrations");
AppState::new(
AppConfig {
server: ServerConfig {
host: "127.0.0.1".into(),
port: 0,
log_level: "info".into(),
cors_origins: vec![],
},
db: DbConfig {
url: "sqlite::memory:".into(),
max_connections: 1,
ssl: false,
runtime_db_url: None,
runtime_db_dialect: None,
runtime_db_ssl: None,
},
auth: AuthConfig {
jwt_secret: "secret".into(),
jwt_expiry_hours: 1,
auth_token: "auth".into(),
proxy_token: "proxy".into(),
},
proxy: ProxyConfig {
upstream_timeout_secs: 30,
first_byte_timeout_secs: 30,
max_retries: 3,
proxy_log_retention_days: 7,
template_proxy_url: None,
},
comfyui: Default::default(),
},
pool,
BuiltinKeys {
comfyui_agent: (1, "sk-comfyui-test".into()),
},
RuntimeAuthSettings::from_token_for_test("proxy"),
)
}
async fn seed_event_sources(state: &AppState) {
sqlx::query(
"INSERT INTO sites (id, name, url, platform, status)
VALUES (1, '主站', 'https://example.test', 'new-api', 'active')",
)
.execute(&*state.db)
.await
.unwrap();
sqlx::query(
"INSERT INTO accounts
(id, site_id, username, access_token, balance, balance_used, quota, last_balance_refresh, status)
VALUES
(10, 1, 'alice', 'access-token', 12.5, 1.25, 20.0, '2026-06-08T09:30:00Z', 'active')",
)
.execute(&*state.db)
.await
.unwrap();
sqlx::query(
"INSERT INTO account_tokens
(id, account_id, name, token, source, token_group, value_status, enabled, created_at, updated_at)
VALUES
(20, 10, '默认 Key', 'sk-account-token', 'manual', 'prod', 'ready', 1, '2026-06-08T09:00:00Z', '2026-06-08T09:40:00Z')",
)
.execute(&*state.db)
.await
.unwrap();
sqlx::query(
"INSERT INTO checkin_logs (id, account_id, status, message, reward, created_at)
VALUES (30, 10, 'success', 'Daily checkin completed', '+$0.0100', '2026-06-08T09:50:00Z')",
)
.execute(&*state.db)
.await
.unwrap();
sqlx::query(
"INSERT INTO proxy_logs
(id, account_id, model_requested, model_actual, status, http_status, latency_ms, total_tokens, estimated_cost, error_message, created_at)
VALUES
(40, 10, 'gpt-test', 'gpt-test', 'failed', 404, 812, 33, 0.0012, 'model not found', '2026-06-08T10:00:00Z')",
)
.execute(&*state.db)
.await
.unwrap();
sqlx::query(
"INSERT INTO site_announcements
(id, site_id, platform, source_key, title, content, level, first_seen_at, read_at)
VALUES
(50, 1, 'new-api', 'notice-1', '维护通知', '上游将维护 30 分钟', 'warning', '2026-06-08T10:10:00Z', NULL)",
)
.execute(&*state.db)
.await
.unwrap();
}
async fn list_for_type(state: AppState, filter_type: Option<&str>) -> Value {
let Json(body) = list_events(
State(state),
Query(ListEventsQuery {
limit: Some(50),
offset: Some(0),
filter_type: filter_type.map(str::to_string),
read: None,
}),
)
.await
.unwrap();
body
}
#[tokio::test]
async fn events_router_requires_admin_api_token() {
let state = test_state().await;
let app = routes(state.clone()).with_state(state.clone());
let ok_response = app
.clone()
.oneshot(
Request::builder()
.uri("/api/events")
.header("Authorization", "Bearer auth")
.body(Body::empty())
.unwrap(),
)
.await
.expect("admin token request");
assert_eq!(ok_response.status(), StatusCode::OK);
let user_jwt = generate_jwt(&state.config.auth.jwt_secret, "1", "user", 1)
.expect("user jwt");
let user_response = app
.clone()
.oneshot(
Request::builder()
.uri("/api/events")
.header("Authorization", format!("Bearer {user_jwt}"))
.body(Body::empty())
.unwrap(),
)
.await
.expect("user jwt request");
assert_eq!(user_response.status(), StatusCode::UNAUTHORIZED);
let missing_response = app
.oneshot(
Request::builder()
.uri("/api/events")
.body(Body::empty())
.unwrap(),
)
.await
.expect("missing auth request");
assert_eq!(missing_response.status(), StatusCode::UNAUTHORIZED);
}
fn data_types(body: &Value) -> Vec<String> {
body["data"]
.as_array()
.unwrap()
.iter()
.filter_map(|row| row.get("type").and_then(Value::as_str).map(str::to_string))
.collect()
}
#[tokio::test]
async fn list_events_aggregates_business_log_sources() {
let state = test_state().await;
seed_event_sources(&state).await;
let all = list_for_type(state.clone(), None).await;
let all_types = data_types(&all);
assert!(all_types.contains(&"checkin".to_string()));
assert!(all_types.contains(&"proxy".to_string()));
assert!(all_types.contains(&"balance".to_string()));
assert!(all_types.contains(&"token".to_string()));
assert!(all_types.contains(&"site_notice".to_string()));
assert!(all_types.contains(&"status".to_string()));
for (filter_type, expected_title) in [
("checkin", "签到成功"),
("proxy", "代理请求失败"),
("balance", "余额已刷新"),
("token", "令牌已更新"),
("site_notice", "站点公告"),
("status", "站点状态"),
] {
let body = list_for_type(state.clone(), Some(filter_type)).await;
let rows = body["data"].as_array().unwrap();
assert!(!rows.is_empty(), "{filter_type} should return rows");
assert!(rows
.iter()
.all(|row| row.get("type").and_then(Value::as_str) == Some(filter_type)));
assert!(rows.iter().any(|row| row
.get("title")
.and_then(Value::as_str)
.unwrap_or_default()
.contains(expected_title)));
}
}
#[tokio::test]
async fn unread_filter_excludes_read_only_virtual_logs() {
let state = test_state().await;
seed_event_sources(&state).await;
let Json(checkin_body) = list_events(
State(state.clone()),
Query(ListEventsQuery {
limit: Some(50),
offset: Some(0),
filter_type: Some("checkin".into()),
read: Some("false".into()),
}),
)
.await
.unwrap();
assert!(checkin_body["data"].as_array().unwrap().is_empty());
let Json(notice_body) = list_events(
State(state),
Query(ListEventsQuery {
limit: Some(50),
offset: Some(0),
filter_type: Some("site_notice".into()),
read: Some("false".into()),
}),
)
.await
.unwrap();
assert!(notice_body["data"].as_array().unwrap().is_empty());
}
#[tokio::test]
async fn status_snapshots_are_available_in_all_logs_and_status_filter() {
let state = test_state().await;
seed_event_sources(&state).await;
let all = list_for_type(state.clone(), None).await;
assert!(data_types(&all).contains(&"status".to_string()));
let status = list_for_type(state, Some("status")).await;
let rows = status["data"].as_array().unwrap();
assert!(rows.iter().any(|row| row["title"]
.as_str()
.unwrap_or_default()
.contains("站点状态")));
assert!(rows.iter().any(|row| row["title"]
.as_str()
.unwrap_or_default()
.contains("账号状态")));
assert!(rows
.iter()
.all(|row| row.get("type").and_then(Value::as_str) == Some("status")));
}
}
+30 -3
View File
@@ -27,7 +27,6 @@ mod comfyui_agent;
mod digital_human;
mod events;
mod health;
mod monitor;
mod money_printer;
mod oauth;
mod organizations;
@@ -62,9 +61,8 @@ pub fn build_router(state: AppState) -> Router {
let api_routes = api_routes
.merge(auth::routes())
.merge(events::routes())
.merge(events::routes(state.clone()))
.merge(checkin::routes())
.merge(monitor::routes())
.merge(oauth::routes())
.merge(organizations::routes())
.merge(search::routes())
@@ -221,6 +219,8 @@ fn is_api_path(path: &str) -> bool {
fn should_serve_website(path: &str) -> bool {
path == "/"
|| path == "/new"
|| path.starts_with("/new/")
|| path.starts_with("/_next/")
|| path == "/favicon.ico"
|| path == "/favicon.png"
@@ -230,6 +230,9 @@ fn should_serve_website(path: &str) -> bool {
fn embedded_asset_path(path: &str) -> String {
let trimmed = path.trim_start_matches('/');
if trimmed == "new" || trimmed == "new/" || trimmed == "new/index.html" {
return "new.html".to_string();
}
if trimmed.is_empty() || trimmed.ends_with('/') {
format!("{trimmed}index.html")
} else {
@@ -269,3 +272,27 @@ async fn removed_comfy_workbench() -> (StatusCode, Json<Value>) {
})),
)
}
#[cfg(test)]
mod tests {
use super::{embedded_asset_path, should_serve_website};
#[test]
fn experimental_website_route_is_served_from_website_export() {
assert!(should_serve_website("/new"));
assert!(should_serve_website("/new/"));
assert!(should_serve_website("/new/index.html"));
assert_eq!(embedded_asset_path("/new"), "new.html");
assert_eq!(embedded_asset_path("/new/"), "new.html");
assert_eq!(embedded_asset_path("/new/index.html"), "new.html");
}
#[test]
fn application_routes_still_use_web_spa() {
assert!(!should_serve_website("/boos"));
assert!(!should_serve_website("/dashboard"));
assert!(!should_serve_website("/user/recharge"));
assert!(!should_serve_website("/logs"));
assert!(!should_serve_website("/events"));
}
}
+557 -11
View File
@@ -2,13 +2,14 @@ use crate::{
auth::jwt::validate_jwt,
db::models::{background_task::BackgroundTask, downstream_api_key::DownstreamApiKey},
error::{AppError, AppResult},
services::billing::{ComfyTaskBillingEvent, UsageTracker},
services::workspace::{self, WorkspaceAssetInput},
workflow::storage,
AppState,
};
use axum::{
extract::{Path, Query, State},
http::{header, HeaderMap, Method},
http::{header, HeaderMap, Method, StatusCode},
routing::{get, post},
Json, Router,
};
@@ -23,11 +24,60 @@ const VIDEO_MAX_POLLS: usize = 120;
const COMFY_POLL_INTERVAL_SECS: u64 = 5;
const COMFY_MAX_POLLS: usize = 180;
fn normalize_video_model(value: &str) -> String {
match value.trim() {
"seedance-1-0-lite-t2v" => "seedance-1-0-lite-t2v-250428".to_string(),
"seedance-1-0-lite-i2v" => "seedance-1-0-pro-i2v-250528".to_string(),
other => other.to_string(),
}
}
fn configured_task_cost(value: f64) -> f64 {
if value.is_finite() {
value.max(0.0)
} else {
0.0
}
}
fn truncate_chars(value: &str, limit: usize) -> String {
value.chars().take(limit).collect()
}
fn comfy_task_event(
auth: &MoneyPrinterAuth,
model: String,
prompt_tokens: i64,
completion_tokens: i64,
cost: f64,
topic: &str,
request: &Value,
response: Option<&Value>,
error_message: Option<String>,
) -> ComfyTaskBillingEvent {
ComfyTaskBillingEvent {
key_id: auth.key_id,
model,
task_type: TASK_TYPE_MONEY_PRINTER,
client_family: "money-printer",
client_app_id: "money-printer",
client_app_name: "Money Printer",
prompt_tokens,
completion_tokens,
cost,
content_summary: Some(truncate_chars(topic, 240)),
request_content: Some(request.to_string()),
response_content: response.map(Value::to_string),
error_message,
}
}
#[derive(Debug, Clone)]
struct MoneyPrinterAuth {
user_id: Option<i64>,
role: String,
outbound_key: String,
key_id: Option<i64>,
key_name: String,
}
@@ -94,6 +144,7 @@ async fn resolve_auth(headers: &HeaderMap, state: &AppState) -> Result<MoneyPrin
user_id: None,
role: "admin".to_string(),
outbound_key: proxy_token,
key_id: None,
key_name: "global-proxy".to_string(),
});
}
@@ -113,12 +164,14 @@ async fn resolve_auth(headers: &HeaderMap, state: &AppState) -> Result<MoneyPrin
user_id: Some(user_id),
role: claims.role,
outbound_key: key.key,
key_id: Some(key.id),
key_name: key.name,
}),
Err(_) if claims.role == "admin" => Ok(MoneyPrinterAuth {
user_id: Some(user_id),
role: claims.role,
outbound_key: state.runtime_proxy_token(),
key_id: None,
key_name: "global-proxy".to_string(),
}),
Err(message) => Err(AppError::Forbidden(message)),
@@ -150,6 +203,11 @@ async fn status(State(state): State<AppState>, headers: HeaderMap) -> AppResult<
"label": "内置 Cloud Seedance 文生视频",
"nodeType": "ByteDanceTextToVideoNode",
"model": "seedance-1-0-lite-t2v-250428",
"models": [
"seedance-1-0-lite-t2v-250428",
"seedance-1-0-pro-t2v-250528",
"seedance-1-0-pro-i2v-250528"
],
"duration": { "min": 3, "max": 12, "default": 5 },
"resolutions": ["480p", "720p", "1080p"],
"ratios": ["9:16", "16:9", "1:1", "4:3", "3:4", "21:9"]
@@ -420,6 +478,13 @@ fn native_request(state: &AppState, method: Method, path: &str) -> reqwest::Requ
request
}
fn is_pending_history_status(status: StatusCode) -> bool {
matches!(
status,
StatusCode::NOT_FOUND | StatusCode::ACCEPTED | StatusCode::NO_CONTENT
)
}
async fn native_json(
state: &AppState,
method: Method,
@@ -446,6 +511,35 @@ async fn native_json(
Ok(value)
}
async fn native_json_optional(
state: &AppState,
method: Method,
path: &str,
body: Option<Value>,
) -> AppResult<Option<Value>> {
let mut request = native_request(state, method, path);
if let Some(body) = body {
request = request.json(&body);
}
let response = request
.send()
.await
.map_err(|err| AppError::UpstreamError(format!("Comfy backend unreachable: {err}")))?;
let status = response.status();
if is_pending_history_status(status) {
return Ok(None);
}
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(Some(value))
}
fn encode_query(value: &str) -> String {
let mut out = String::new();
for byte in value.bytes() {
@@ -488,6 +582,7 @@ fn default_comfy_mapping() -> Value {
json!({
"prompt": { "nodeId": "", "field": "prompt" },
"negativePrompt": { "nodeId": "", "field": "text" },
"model": { "nodeId": "", "field": "model" },
"ratio": { "nodeId": "", "field": "ratio" },
"resolution": { "nodeId": "", "field": "resolution" },
"duration": { "nodeId": "", "field": "duration" },
@@ -649,6 +744,7 @@ fn infer_comfy_mapping(prompt: &Value) -> Value {
let ratio_field = first_existing_input_field(object, &video_node, &["ratio", "model.ratio", "aspect_ratio"], "ratio");
let resolution_field = first_existing_input_field(object, &video_node, &["model.resolution", "resolution", "size"], "resolution");
let duration_field = first_existing_input_field(object, &video_node, &["model.duration", "duration", "length"], "duration");
let model_field = first_existing_input_field(object, &video_node, &["model", "model_name", "ckpt_name"], "model");
json!({
"prompt": {
"nodeId": prompt_node,
@@ -658,6 +754,10 @@ fn infer_comfy_mapping(prompt: &Value) -> Value {
"nodeId": negative_node,
"field": "text"
},
"model": {
"nodeId": video_node,
"field": model_field
},
"ratio": {
"nodeId": video_node,
"field": ratio_field
@@ -869,6 +969,15 @@ fn inject_comfy_video_inputs(prompt: &mut Value, mapping: &Value, body: &Value,
json!(negative),
)?;
}
if let Some(model) = string_from(body, &["comfyApiModel", "comfy_api_model", "apiModel", "api_model"]) {
apply_optional_node_input(
prompt,
mapping,
&["model", "modelNode", "apiModel"],
"model",
json!(normalize_video_model(&model)),
);
}
if let Some(value) = ["ratio", "aspectRatio", "aspect_ratio"].iter().find_map(|key| body.get(*key).cloned()) {
apply_optional_node_input(prompt, mapping, &["ratio", "ratioNode"], "ratio", value);
}
@@ -1042,13 +1151,13 @@ async fn submit_comfy_prompt(state: &AppState, prompt: Value) -> AppResult<(Stri
Ok((prompt_id, value))
}
async fn comfy_history(state: &AppState, prompt_id: &str) -> AppResult<Value> {
async fn comfy_history_optional(state: &AppState, prompt_id: &str) -> AppResult<Option<Value>> {
let path = if state.config.comfyui.uses_cloud_backend() {
format!("/jobs/{prompt_id}")
} else {
format!("/history/{prompt_id}")
};
native_json(state, Method::GET, &path, None).await
native_json_optional(state, Method::GET, &path, None).await
}
fn comfy_failure_message(value: &Value) -> Option<String> {
@@ -1071,12 +1180,24 @@ fn comfy_failure_message(value: &Value) -> Option<String> {
fn comfy_output_objects(value: &Value) -> Vec<Value> {
let mut out = Vec::new();
fn push_unique(out: &mut Vec<Value>, item: Value) {
let identity = output_identity(&item);
if out.iter().any(|existing| output_identity(existing) == identity) {
return;
}
out.push(item);
}
fn walk(value: &Value, out: &mut Vec<Value>) {
match value {
Value::Object(object) => {
if output_filename(value).is_some() || output_remote_url(value).is_some() {
push_unique(out, value.clone());
}
for key in ["videos", "gifs", "images", "audio"] {
if let Some(items) = object.get(key).and_then(Value::as_array) {
out.extend(items.iter().filter(|item| item.is_object()).cloned());
for item in items {
walk(item, out);
}
}
}
for value in object.values() {
@@ -1088,6 +1209,11 @@ fn comfy_output_objects(value: &Value) -> Vec<Value> {
walk(item, out);
}
}
Value::String(value) => {
if let Some(item) = output_item_from_reference(value) {
push_unique(out, item);
}
}
_ => {}
}
}
@@ -1095,9 +1221,112 @@ fn comfy_output_objects(value: &Value) -> Vec<Value> {
out
}
fn output_identity(item: &Value) -> String {
output_remote_url(item)
.or_else(|| output_filename(item))
.unwrap_or_else(|| item.to_string())
}
fn looks_like_media_reference(value: &str) -> bool {
let trimmed = value.trim();
if !(trimmed.starts_with("http://") || trimmed.starts_with("https://") || trimmed.starts_with('/')) {
return false;
}
let lower = trimmed
.split(['?', '#'])
.next()
.unwrap_or(trimmed)
.to_ascii_lowercase();
lower.contains("/view")
|| lower.ends_with(".mp4")
|| lower.ends_with(".webm")
|| lower.ends_with(".mov")
|| lower.ends_with(".m4v")
|| lower.ends_with(".gif")
|| lower.ends_with(".png")
|| lower.ends_with(".jpg")
|| lower.ends_with(".jpeg")
|| lower.ends_with(".webp")
|| lower.ends_with(".mp3")
|| lower.ends_with(".wav")
|| lower.ends_with(".m4a")
}
fn output_item_from_reference(value: &str) -> Option<Value> {
let trimmed = value.trim();
if !looks_like_media_reference(trimmed) {
return None;
}
let filename = trimmed
.split(['?', '#'])
.next()
.unwrap_or(trimmed)
.rsplit('/')
.next()
.unwrap_or("comfy-output")
.to_string();
Some(json!({
"url": trimmed,
"filename": filename,
"type": "output"
}))
}
fn output_filename(item: &Value) -> Option<String> {
["filename", "fileName", "name", "file"]
.iter()
.find_map(|key| scalar_string(item.get(*key)))
}
fn output_remote_url(item: &Value) -> Option<String> {
[
"url",
"remoteUrl",
"remote_url",
"videoUrl",
"video_url",
"audioUrl",
"audio_url",
"imageUrl",
"image_url",
"fileUrl",
"file_url",
"downloadUrl",
"download_url",
"video",
"audio",
"image",
]
.iter()
.find_map(|key| scalar_string(item.get(*key)))
.filter(|value| looks_like_media_reference(value))
}
fn output_kind(item: &Value) -> &'static str {
let filename = scalar_string(item.get("filename")).unwrap_or_default().to_ascii_lowercase();
if filename.ends_with(".mp4") || filename.ends_with(".webm") || filename.ends_with(".mov") {
let explicit = scalar_string(item.get("kind"))
.or_else(|| scalar_string(item.get("mediaType")))
.or_else(|| scalar_string(item.get("mime")))
.unwrap_or_default()
.to_ascii_lowercase();
if explicit.contains("video") {
return "video";
}
if explicit.contains("audio") {
return "audio";
}
if explicit.contains("image") {
return "image";
}
let filename = output_filename(item)
.or_else(|| output_remote_url(item))
.unwrap_or_default()
.to_ascii_lowercase();
if filename.ends_with(".mp4")
|| filename.ends_with(".webm")
|| filename.ends_with(".mov")
|| filename.ends_with(".m4v")
|| filename.ends_with(".gif")
{
"video"
} else if filename.ends_with(".mp3") || filename.ends_with(".wav") || filename.ends_with(".m4a") {
"audio"
@@ -1107,7 +1336,7 @@ fn output_kind(item: &Value) -> &'static str {
}
fn output_query(item: &Value) -> Option<String> {
let filename = scalar_string(item.get("filename"))?;
let filename = output_filename(item)?;
let file_type = scalar_string(item.get("type")).unwrap_or_else(|| "output".to_string());
let subfolder = scalar_string(item.get("subfolder")).unwrap_or_default();
let mut parts = vec![
@@ -1128,6 +1357,37 @@ async fn archive_comfy_output(
prompt: &str,
title: &str,
) -> Option<Value> {
if let Some(remote_url) = output_remote_url(item) {
let kind = output_kind(item).to_string();
let input = WorkspaceAssetInput {
user_id: auth.user_id,
source: "money-printer".to_string(),
kind: kind.clone(),
title: Some(title.chars().take(80).collect()),
prompt: Some(prompt.chars().take(1000).collect()),
model: Some("comfy-api".to_string()),
remote_url: Some(remote_url.clone()),
metadata: json!({
"workspaceAsset": { "role": "output" },
"moneyPrinter": { "backend": "comfy", "promptId": prompt_id, "raw": item },
}),
};
return match workspace::record_remote_asset(state, input).await {
Ok(asset) => Some(json!({
"publicId": asset.public_id,
"kind": kind,
"status": asset.status,
"remoteUrl": asset.remote_url,
"fileName": asset.file_name,
})),
Err(err) => Some(json!({
"kind": kind,
"status": "failed",
"remoteUrl": remote_url,
"error": err.to_string(),
})),
};
}
let query = output_query(item)?;
let kind = output_kind(item).to_string();
let path = if kind == "video" { "/viewvideo" } else { "/view" };
@@ -1195,9 +1455,13 @@ async fn submit_comfy_video(
}
let video_prompt = inject_comfy_video_inputs(&mut prompt, &mapping, body, plan)?;
let (prompt_id, response) = submit_comfy_prompt(state, prompt).await?;
let model = string_from(body, &["comfyApiModel", "comfy_api_model", "apiModel", "api_model"])
.map(|value| normalize_video_model(&value))
.unwrap_or_else(|| "comfy-api".to_string());
Ok(json!({
"backend": "comfy",
"promptId": prompt_id,
"model": model,
"prompt": video_prompt,
"mapping": mapping,
"response": response,
@@ -1211,7 +1475,8 @@ async fn submit_video(
plan: &Value,
) -> AppResult<Value> {
let model = string_from(body, &["videoModel", "video_model", "model"])
.unwrap_or_else(|| "seedance-1-0-lite-t2v".to_string());
.map(|value| normalize_video_model(&value))
.unwrap_or_else(|| "seedance-1-0-lite-t2v-250428".to_string());
let prompt = string_from(body, &["videoPrompt", "video_prompt"]).unwrap_or_else(|| {
plan.get("video_prompt")
.and_then(Value::as_str)
@@ -1352,6 +1617,9 @@ async fn create_job(
Json(body): Json<Value>,
) -> AppResult<Json<Value>> {
let auth = resolve_auth(&headers, &state).await?;
if let Some(key_id) = auth.key_id {
UsageTracker::new(state.db.clone()).check_rate_limit(key_id).await?;
}
let topic = string_from(&body, &["topic", "prompt"])
.ok_or_else(|| AppError::BadRequest("topic is required".into()))?;
let task_id = format!("mp_{}", uuid::Uuid::new_v4().simple());
@@ -1359,6 +1627,7 @@ async fn create_job(
let initial_result = json!({
"userId": auth.user_id,
"role": auth.role,
"keyId": auth.key_id,
"keyName": auth.key_name,
"topic": topic,
"request": body,
@@ -1380,6 +1649,23 @@ async fn create_job(
let task_id_for_spawn = task_id.clone();
let auth_for_task = auth.clone();
tokio::spawn(async move {
let failure_body = body.clone();
let failure_topic = topic.clone();
let failure_backend = string_from(&failure_body, &["videoBackend", "video_backend"])
.unwrap_or_else(|| "boosapi".to_string())
.to_ascii_lowercase();
let failure_model = string_from(
&failure_body,
&["comfyApiModel", "comfy_api_model", "apiModel", "api_model", "videoModel", "video_model", "model"],
)
.map(|value| normalize_video_model(&value))
.unwrap_or_else(|| {
if failure_backend == "comfy" {
"comfy-api".to_string()
} else {
"video".to_string()
}
});
let result = async {
update_progress(&state_for_task, &task_id_for_spawn, 2, "正在生成脚本和分镜").await?;
let script = generate_script(&state_for_task, &auth_for_task, &body).await?;
@@ -1390,6 +1676,7 @@ async fn create_job(
let mid_result = json!({
"userId": auth_for_task.user_id,
"role": auth_for_task.role,
"keyId": auth_for_task.key_id,
"keyName": auth_for_task.key_name,
"topic": topic,
"request": body,
@@ -1426,15 +1713,21 @@ async fn create_job(
.and_then(Value::as_str)
.unwrap_or_default()
.to_string();
let model = submitted
.get("model")
.and_then(Value::as_str)
.unwrap_or("comfy-api")
.to_string();
let submitted_result = json!({
"userId": auth_for_task.user_id,
"role": auth_for_task.role,
"keyId": auth_for_task.key_id,
"keyName": auth_for_task.key_name,
"topic": topic,
"request": body,
"script": script,
"videoBackend": video_backend,
"model": "comfy-api",
"model": model,
"prompt": prompt,
"video": submitted,
});
@@ -1449,7 +1742,11 @@ async fn create_job(
return Ok(());
}
update_progress(&state_for_task, &task_id_for_spawn, 4, "Comfy 正在生成视频").await?;
let history = comfy_history(&state_for_task, &prompt_id).await?;
let Some(history) = comfy_history_optional(&state_for_task, &prompt_id).await?
else {
tokio::time::sleep(Duration::from_secs(COMFY_POLL_INTERVAL_SECS)).await;
continue;
};
if let Some(message) = comfy_failure_message(&history) {
return Err(AppError::UpstreamError(message));
}
@@ -1478,7 +1775,7 @@ async fn create_job(
if final_history.is_null() {
return Err(AppError::GatewayTimeout("Comfy 视频任务等待超时".into()));
}
(submitted, final_history, None, prompt, "comfy-api".to_string())
(submitted, final_history, None, prompt, model)
} else {
update_progress(&state_for_task, &task_id_for_spawn, 3, "正在提交 BoosAPI 视频模型").await?;
let submitted = submit_video(&state_for_task, &auth_for_task, &body, &plan).await?;
@@ -1524,9 +1821,46 @@ async fn create_job(
}
(submitted, final_video_response, video_url, prompt, model)
};
let billing = if video_backend == "comfy" {
let tracker = UsageTracker::new(state_for_task.db.clone());
let event = comfy_task_event(
&auth_for_task,
model.clone(),
0,
0,
configured_task_cost(state_for_task.config.comfyui.money_printer_video_cost),
&topic,
&body,
Some(&final_video_response),
None,
);
match tracker.record_comfy_task_success(event).await {
Ok(details) => details,
Err(err) => {
tracing::warn!(
%err,
task_id = %task_id_for_spawn,
"failed to record money printer comfy billing",
);
json!({
"source": "comfy-native-task",
"taskType": TASK_TYPE_MONEY_PRINTER,
"status": "success",
"charged": false,
"failed": true,
"error": err.to_string(),
"totalCost": 0.0,
"steps": []
})
}
}
} else {
Value::Null
};
let final_result = json!({
"userId": auth_for_task.user_id,
"role": auth_for_task.role,
"keyId": auth_for_task.key_id,
"keyName": auth_for_task.key_name,
"topic": topic,
"request": body,
@@ -1538,6 +1872,7 @@ async fn create_job(
"finalVideoResponse": final_video_response,
"videoUrl": video_url,
"assets": assets,
"billing": billing,
});
sqlx::query(
"UPDATE background_tasks
@@ -1553,6 +1888,27 @@ async fn create_job(
}
.await;
if let Err(err) = result {
if failure_backend == "comfy" {
let tracker = UsageTracker::new(state_for_task.db.clone());
let event = comfy_task_event(
&auth_for_task,
failure_model,
0,
0,
0.0,
&failure_topic,
&failure_body,
None,
Some(err.to_string()),
);
if let Err(log_err) = tracker.record_comfy_task_failure(event).await {
tracing::warn!(
%log_err,
task_id = %task_id_for_spawn,
"failed to record money printer comfy failure log",
);
}
}
let _ = sqlx::query(
"UPDATE background_tasks
SET status = 'failed', message = '短视频生成失败', error = ?, updated_at = datetime('now'), finished_at = datetime('now')
@@ -1660,3 +2016,193 @@ async fn clear_jobs(State(state): State<AppState>, headers: HeaderMap) -> AppRes
}
Ok(Json(json!({ "success": true, "cleared": cleared })))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{
config::{AppConfig, AuthConfig, ComfyUIBackend, ComfyUIConfig, DbConfig, ProxyConfig, ServerConfig},
services::{builtin_keys::BuiltinKeys, runtime_auth::RuntimeAuthSettings},
};
use axum::http::HeaderValue;
use sqlx::sqlite::SqlitePoolOptions;
async fn test_state_with_comfy_backend(backend: ComfyUIBackend) -> AppState {
let pool = SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.expect("in-memory database");
let comfyui = ComfyUIConfig {
enabled: true,
backend,
base_url: Some("https://cloud.comfy.org".to_string()),
api_key: Some("sk-comfy-cloud-test".to_string()),
..Default::default()
};
AppState::new(
AppConfig {
server: ServerConfig {
host: "127.0.0.1".into(),
port: 0,
log_level: "info".into(),
cors_origins: vec![],
},
db: DbConfig {
url: "sqlite::memory:".into(),
max_connections: 1,
ssl: false,
runtime_db_url: None,
runtime_db_dialect: None,
runtime_db_ssl: None,
},
auth: AuthConfig {
jwt_secret: "test-secret".into(),
jwt_expiry_hours: 1,
auth_token: "admin-token".into(),
proxy_token: "proxy".into(),
},
proxy: ProxyConfig {
upstream_timeout_secs: 30,
first_byte_timeout_secs: 30,
max_retries: 3,
proxy_log_retention_days: 7,
template_proxy_url: None,
},
comfyui,
},
pool,
BuiltinKeys {
comfyui_agent: (1, "sk-comfyui-test".into()),
},
RuntimeAuthSettings::from_token_for_test("proxy"),
)
}
#[tokio::test]
async fn status_exposes_comfy_cloud_api_model_candidates() {
let state = test_state_with_comfy_backend(ComfyUIBackend::Cloud).await;
let mut headers = HeaderMap::new();
headers.insert(
header::AUTHORIZATION,
HeaderValue::from_static("Bearer proxy"),
);
let Json(value) = status(State(state), headers)
.await
.expect("status response");
assert_eq!(value["builtinCloudTemplate"]["enabled"], true);
assert_eq!(
value["builtinCloudTemplate"]["model"],
"seedance-1-0-lite-t2v-250428"
);
let models = value["builtinCloudTemplate"]["models"]
.as_array()
.expect("model candidates");
assert!(models
.iter()
.any(|item| item.as_str() == Some("seedance-1-0-pro-t2v-250528")));
assert!(models
.iter()
.any(|item| item.as_str() == Some("seedance-1-0-pro-i2v-250528")));
}
#[test]
fn comfy_output_objects_collects_nested_filename_and_remote_url_outputs() {
let history = json!({
"status": "success",
"outputs": {
"12": {
"videos": [
{ "filename": "clip.mp4", "subfolder": "final", "type": "output" }
],
"images": [
"https://cdn.example.com/poster.png"
]
},
"13": {
"result": {
"videoUrl": "https://cdn.example.com/final.webm?token=1"
}
}
}
});
let outputs = comfy_output_objects(&history);
assert_eq!(outputs.len(), 3);
assert!(outputs.iter().any(|item| output_filename(item).as_deref() == Some("clip.mp4")));
assert!(outputs.iter().any(|item| output_remote_url(item).as_deref() == Some("https://cdn.example.com/poster.png")));
assert!(outputs.iter().any(|item| output_kind(item) == "video"));
}
#[test]
fn comfy_output_objects_deduplicates_repeated_media_references() {
let history = json!({
"outputs": {
"1": {
"videoUrl": "https://cdn.example.com/final.mp4",
"videos": [
{ "url": "https://cdn.example.com/final.mp4", "filename": "final.mp4" }
]
}
}
});
let outputs = comfy_output_objects(&history);
assert_eq!(outputs.len(), 1);
assert_eq!(output_remote_url(&outputs[0]).as_deref(), Some("https://cdn.example.com/final.mp4"));
}
#[test]
fn comfy_history_polling_treats_not_ready_statuses_as_pending() {
assert!(is_pending_history_status(StatusCode::NOT_FOUND));
assert!(is_pending_history_status(StatusCode::ACCEPTED));
assert!(is_pending_history_status(StatusCode::NO_CONTENT));
assert!(!is_pending_history_status(StatusCode::OK));
assert!(!is_pending_history_status(StatusCode::INTERNAL_SERVER_ERROR));
}
#[test]
fn comfy_api_model_is_injected_into_mapped_video_node() {
let mut prompt = json!({
"1": {
"class_type": "ByteDanceTextToVideoNode",
"inputs": {
"model": "seedance-1-0-lite-t2v-250428",
"prompt": "placeholder",
"resolution": "480p",
"aspect_ratio": "9:16",
"duration": 5
}
}
});
let mapping = json!({
"prompt": { "nodeId": "1", "field": "prompt" },
"model": { "nodeId": "1", "field": "model" },
"ratio": { "nodeId": "1", "field": "aspect_ratio" },
"resolution": { "nodeId": "1", "field": "resolution" },
"duration": { "nodeId": "1", "field": "duration" }
});
let body = json!({
"comfyApiModel": "seedance-1-0-lite-t2v",
"ratio": "16:9",
"resolution": "720p",
"duration": 8
});
let plan = json!({
"video_prompt": "cinematic product demo"
});
let prompt_text = inject_comfy_video_inputs(&mut prompt, &mapping, &body, &plan)
.expect("injecting Comfy video inputs should succeed");
assert_eq!(prompt_text, "cinematic product demo");
assert_eq!(prompt.pointer("/1/inputs/model").and_then(Value::as_str), Some("seedance-1-0-lite-t2v-250428"));
assert_eq!(prompt.pointer("/1/inputs/prompt").and_then(Value::as_str), Some("cinematic product demo"));
assert_eq!(prompt.pointer("/1/inputs/aspect_ratio").and_then(Value::as_str), Some("16:9"));
assert_eq!(prompt.pointer("/1/inputs/resolution").and_then(Value::as_str), Some("720p"));
assert_eq!(prompt.pointer("/1/inputs/duration").and_then(Value::as_i64), Some(8));
}
}
-297
View File
@@ -1,297 +0,0 @@
// Monitor routes — LDOH incident dashboard proxy with cookie auth
use crate::{
error::{AppError, AppResult},
AppState,
};
use axum::{
extract::{Path, State},
http::{header, HeaderMap},
response::Response,
routing::{get, post},
Json, Router,
};
use serde_json::{json, Value};
const MONITOR_AUTH_COOKIE: &str = "meta_monitor_auth";
const LDOH_BASE_URL: &str = "https://ldoh.105117.xyz";
const LDOH_COOKIE_SETTING_KEY: &str = "monitor_ldoh_cookie";
pub fn routes() -> Router<AppState> {
Router::new()
.route(
"/api/monitor/config",
get(get_monitor_config).put(update_monitor_config),
)
.route("/api/monitor/session", post(create_session))
.route("/monitor-proxy/ldoh", get(ldoh_proxy_handler))
.route("/monitor-proxy/ldoh/", get(ldoh_proxy_handler))
.route("/monitor-proxy/ldoh/{*path}", get(ldoh_proxy_handler))
}
/// GET /api/monitor/config — returns whether LDOH cookie is configured
async fn get_monitor_config(State(state): State<AppState>) -> AppResult<Json<Value>> {
let cookie: Option<(String,)> = sqlx::query_as("SELECT value FROM settings WHERE key = ?")
.bind(LDOH_COOKIE_SETTING_KEY)
.fetch_optional(&*state.db)
.await?;
let raw = cookie
.and_then(|(v,)| {
serde_json::from_str::<serde_json::Value>(&v)
.ok()
.and_then(|parsed| parsed.as_str().map(String::from))
})
.unwrap_or_default();
let configured = !raw.is_empty();
let masked = if configured {
mask_cookie_value(&raw)
} else {
String::new()
};
Ok(Json(json!({
"ldohCookieConfigured": configured,
"ldohCookieMasked": masked,
})))
}
/// PUT /api/monitor/config — save LDOH cookie
async fn update_monitor_config(
State(state): State<AppState>,
Json(body): Json<serde_json::Value>,
) -> AppResult<Json<Value>> {
let cookie = body
.get("ldohCookie")
.and_then(|v| v.as_str())
.unwrap_or("");
if cookie.is_empty() {
sqlx::query("INSERT OR REPLACE INTO settings (key, value) VALUES (?, '')")
.bind(LDOH_COOKIE_SETTING_KEY)
.execute(&*state.db)
.await?;
return Ok(Json(json!({
"success": true,
"message": "LDOH Cookie 已清空",
"ldohCookieConfigured": false,
})));
}
let normalized = normalize_ldoh_cookie(cookie);
if !normalized.starts_with("ld_auth_session=") || normalized.len() < 24 {
return Err(AppError::BadRequest("Cookie 格式无效".into()));
}
let escaped = serde_json::to_string(&normalized).unwrap_or_default();
sqlx::query("INSERT OR REPLACE INTO settings (key, value) VALUES (?, ?)")
.bind(LDOH_COOKIE_SETTING_KEY)
.bind(&escaped)
.execute(&*state.db)
.await?;
Ok(Json(json!({
"success": true,
"message": "LDOH Cookie 已保存",
"ldohCookieConfigured": true,
"ldohCookieMasked": mask_cookie_value(&normalized),
})))
}
/// POST /api/monitor/session — returns cookie value for client to set
async fn create_session(State(state): State<AppState>) -> Json<Value> {
let token = &state.config.auth.auth_token;
Json(json!({
"success": true,
"_cookie": format!("{}={}; Path=/; HttpOnly; SameSite=Lax; Max-Age=7200", MONITOR_AUTH_COOKIE, token),
}))
}
/// GET /monitor-proxy/ldoh/{*path} — reverse proxy to ldoh.105117.xyz
async fn ldoh_proxy_handler(
State(state): State<AppState>,
headers: HeaderMap,
Path(path): Path<String>,
) -> Result<Response, AppError> {
// Auth check via cookie
let cookies = parse_cookies(
headers
.get(header::COOKIE)
.and_then(|v| v.to_str().ok())
.unwrap_or(""),
);
if cookies.get(MONITOR_AUTH_COOKIE).map(|s| s.as_str()) != Some(&state.config.auth.auth_token) {
return Err(AppError::Unauthorized(
"Missing or invalid monitor session".into(),
));
}
// Get stored LDOH cookie
let stored: Option<(String,)> = sqlx::query_as("SELECT value FROM settings WHERE key = ?")
.bind(LDOH_COOKIE_SETTING_KEY)
.fetch_optional(&*state.db)
.await?;
let ldoh_cookie = stored
.and_then(|(v,)| {
serde_json::from_str::<serde_json::Value>(&v)
.ok()
.and_then(|parsed| parsed.as_str().map(String::from))
})
.unwrap_or_default();
if ldoh_cookie.is_empty() {
return Err(AppError::BadRequest("LDOH cookie not configured".into()));
}
// Build upstream URL
let clean_path = if path.is_empty() || path == "/" {
String::new()
} else {
format!("/{}", path)
};
let target_url = format!("{}{}", LDOH_BASE_URL, clean_path);
let client = reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(30))
.build()
.map_err(|e| AppError::Internal(anyhow::anyhow!("Failed to build client: {e}")))?;
let mut req = client
.get(&target_url)
.header("cookie", &ldoh_cookie)
.header(
"accept",
headers
.get("accept")
.and_then(|v| v.to_str().ok())
.unwrap_or("*/*"),
)
.header(
"accept-language",
headers
.get("accept-language")
.and_then(|v| v.to_str().ok())
.unwrap_or("zh-CN,zh;q=0.9,en;q=0.8"),
)
.header(
"user-agent",
headers
.get("user-agent")
.and_then(|v| v.to_str().ok())
.unwrap_or("boosapiMonitorProxy/1.0"),
);
if let Some(referer) = headers.get("referer").and_then(|v| v.to_str().ok()) {
req = req.header("referer", referer.replace("/monitor-proxy/ldoh", ""));
}
let resp = req
.send()
.await
.map_err(|e| AppError::UpstreamError(format!("LDOH unreachable: {e}")))?;
let status = resp.status();
let content_type = resp
.headers()
.get("content-type")
.and_then(|v| v.to_str().ok())
.unwrap_or("application/octet-stream")
.to_string();
let mut builder = axum::response::Response::builder().status(status);
if let Some(ctl) = resp
.headers()
.get("cache-control")
.and_then(|v| v.to_str().ok())
{
builder = builder.header("cache-control", ctl);
}
// Rewrite text content
if content_type.contains("text/html")
|| content_type.contains("javascript")
|| content_type.contains("text/css")
|| content_type.contains("application/json")
{
let text = resp
.text()
.await
.map_err(|e| AppError::UpstreamError(e.to_string()))?;
return builder
.header(header::CONTENT_TYPE, &content_type)
.body(axum::body::Body::from(rewrite_proxy_text(&text)))
.map_err(|e| AppError::Internal(anyhow::anyhow!("{e}")));
}
let bytes = resp
.bytes()
.await
.map_err(|e| AppError::UpstreamError(e.to_string()))?;
builder
.header(header::CONTENT_TYPE, &content_type)
.body(axum::body::Body::from(bytes))
.map_err(|e| AppError::Internal(anyhow::anyhow!("{e}")))
}
// ── Helpers ────────────────────────────────────────────────────────────
fn parse_cookies(raw: &str) -> std::collections::HashMap<String, String> {
let mut map = std::collections::HashMap::new();
for part in raw.split(';') {
let part = part.trim();
if let Some(idx) = part.find('=') {
let key = part[..idx].trim().to_string();
let val = part[idx + 1..].trim().to_string();
if !key.is_empty() {
map.insert(key, val);
}
}
}
map
}
fn mask_cookie_value(cookie_text: &str) -> String {
let raw = cookie_text.trim();
if raw.is_empty() {
return String::new();
}
let value = raw.split('=').nth(1).unwrap_or(raw);
if value.len() <= 10 {
format!("{}****", &value[..2.min(value.len())])
} else {
format!("{}****{}", &value[..6], &value[value.len() - 4..])
}
}
fn normalize_ldoh_cookie(raw: &str) -> String {
let trimmed = raw.trim().to_string();
if trimmed.contains("ld_auth_session=") {
if let Some(first) = trimmed.split(';').next() {
if first.trim().starts_with("ld_auth_session=") {
return first.trim().to_string();
}
}
}
format!("ld_auth_session={}", trimmed)
}
fn rewrite_proxy_text(text: &str) -> String {
text.replace("https://ldoh.105117.xyz/", "/monitor-proxy/ldoh/")
.replace(
"https:\\/\\/ldoh.105117.xyz\\/",
"\\/monitor-proxy\\/ldoh\\/",
)
.replace("src=\"/", "src=\"/monitor-proxy/ldoh/")
.replace("src='/", "src='/monitor-proxy/ldoh/")
.replace("href=\"/", "href=\"/monitor-proxy/ldoh/")
.replace("href='/", "href='/monitor-proxy/ldoh/")
.replace("action=\"/", "action=\"/monitor-proxy/ldoh/")
.replace("action='/", "action='/monitor-proxy/ldoh/")
.replace("\"/_next/", "\"/monitor-proxy/ldoh/_next/")
.replace("'/_next/", "'/monitor-proxy/ldoh/_next/")
.replace("\"\\\\/api/", "\"\\\\/monitor-proxy\\\\/ldoh\\\\/api/")
.replace("'/api/", "'/monitor-proxy/ldoh/api/")
.replace("\"/api/", "\"/monitor-proxy/ldoh/api/")
}
+4206 -117
View File
File diff suppressed because it is too large Load Diff
+6
View File
@@ -32,6 +32,7 @@ struct PictureAgentAuth {
user_id: Option<i64>,
role: String,
outbound_key: String,
key_id: Option<i64>,
key_name: String,
}
@@ -129,6 +130,7 @@ async fn resolve_picture_agent_auth(
user_id: None,
role: "admin".to_string(),
outbound_key: proxy_token,
key_id: None,
key_name: "global-proxy".to_string(),
});
}
@@ -149,12 +151,14 @@ async fn resolve_picture_agent_auth(
user_id: Some(user_id),
role: claims.role,
outbound_key: key.key,
key_id: Some(key.id),
key_name: key.name,
}),
Err(_) if claims.role == "admin" => Ok(PictureAgentAuth {
user_id: Some(user_id),
role: claims.role,
outbound_key: state.runtime_proxy_token(),
key_id: None,
key_name: "global-proxy".to_string(),
}),
Err(message) => Err(AppError::Forbidden(message)),
@@ -656,6 +660,7 @@ async fn create_generation_job(
let initial_result = json!({
"userId": auth.user_id,
"role": auth.role,
"keyId": auth.key_id,
"keyName": auth.key_name,
"prompt": prompt,
"model": model,
@@ -737,6 +742,7 @@ async fn create_generation_job(
let final_result = json!({
"userId": auth_for_task.user_id,
"role": auth_for_task.role,
"keyId": auth_for_task.key_id,
"keyName": auth_for_task.key_name,
"prompt": prompt,
"model": model,
+57 -2
View File
@@ -1087,7 +1087,11 @@ where
}
let raw_error_text = Some(err.to_string());
if matches!(action, AttemptFailureAction::Terminal | AttemptFailureAction::Stop) || attempts >= 5 {
if matches!(
action,
AttemptFailureAction::Terminal | AttemptFailureAction::Stop
) || attempts >= 5
{
return Err(StreamOpenFailure {
selected: Some(selected),
attempts,
@@ -1661,7 +1665,10 @@ async fn execute_chat_completions_stream_with_conductor(
let request_log_body = body.clone();
let body_for_attempt = body.clone();
let opened = open_stream_with_failover(&state, &model, move |state, selected| {
let body = ensure_stream_usage(body_with_selected_model(body_for_attempt.clone(), &selected));
let body = ensure_stream_usage(body_with_selected_model(
body_for_attempt.clone(),
&selected,
));
async move {
executors::execute_chat_completions_stream_response(
state,
@@ -1883,12 +1890,22 @@ mod tests {
use sqlx::sqlite::SqlitePoolOptions;
async fn test_state(upstream_base_url: String) -> AppState {
crate::proxy::token_router::router::invalidate_route_cache();
let pool = SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.unwrap();
sqlx::query(
"CREATE TABLE settings (
key TEXT PRIMARY KEY NOT NULL,
value TEXT
)",
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"CREATE TABLE downstream_api_keys (
id INTEGER PRIMARY KEY,
@@ -1948,6 +1965,8 @@ mod tests {
site_id INTEGER NOT NULL,
access_token TEXT NOT NULL,
api_token TEXT,
oauth_provider TEXT,
extra_config TEXT,
unit_cost REAL,
status TEXT NOT NULL DEFAULT 'active'
)",
@@ -1973,6 +1992,7 @@ mod tests {
route_id INTEGER NOT NULL,
account_id INTEGER NOT NULL,
token_id INTEGER,
oauth_route_unit_id INTEGER,
source_model TEXT,
priority INTEGER DEFAULT 0,
weight INTEGER DEFAULT 10,
@@ -1992,6 +2012,41 @@ mod tests {
.execute(&pool)
.await
.unwrap();
sqlx::query(
"CREATE TABLE oauth_route_units (
id INTEGER PRIMARY KEY,
site_id INTEGER NOT NULL,
provider TEXT NOT NULL,
name TEXT NOT NULL,
strategy TEXT NOT NULL DEFAULT 'round_robin',
enabled INTEGER DEFAULT 1
)",
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"CREATE TABLE oauth_route_unit_members (
id INTEGER PRIMARY KEY,
unit_id INTEGER NOT NULL,
account_id INTEGER NOT NULL,
sort_order INTEGER DEFAULT 0,
success_count INTEGER DEFAULT 0,
fail_count INTEGER DEFAULT 0,
total_latency_ms INTEGER DEFAULT 0,
total_cost REAL DEFAULT 0,
last_used_at TEXT,
last_selected_at TEXT,
last_fail_at TEXT,
consecutive_fail_count INTEGER NOT NULL DEFAULT 0,
cooldown_level INTEGER NOT NULL DEFAULT 0,
cooldown_until TEXT,
updated_at TEXT
)",
)
.execute(&pool)
.await
.unwrap();
// The channel-selection query joins through site_api_endpoints to
// pick a per-site override URL. We don't need any rows — the
// COALESCE falls back to sites.url when this table is empty —
+462 -44
View File
@@ -8,8 +8,8 @@ use axum::{
};
use serde::Deserialize;
use serde_json::{json, Value};
use sqlx::Row;
use std::collections::{BTreeMap, HashMap};
use sqlx::{QueryBuilder, Row, Sqlite};
use std::collections::{BTreeMap, HashMap, HashSet};
use tracing::info;
#[derive(Debug, Deserialize)]
@@ -34,6 +34,324 @@ pub struct ProxyDebugTraceQuery {
status: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
enum ProxyLogClientFilter {
App(String),
Family(String),
Raw(String),
}
#[derive(Debug, Clone, Default)]
struct ProxyLogFilters {
status: Option<String>,
search_like: Option<String>,
client: Option<ProxyLogClientFilter>,
site_id: Option<i64>,
downstream_key_id: Option<i64>,
from: Option<String>,
to: Option<String>,
}
fn normalize_proxy_log_status(value: Option<&String>) -> Option<String> {
match value
.map(|item| item.trim().to_ascii_lowercase())
.as_deref()
{
Some("success") => Some("success".to_string()),
Some("failed") => Some("failed".to_string()),
_ => None,
}
}
fn normalize_positive_i64(value: Option<&String>) -> Option<i64> {
value
.and_then(|item| item.trim().parse::<i64>().ok())
.filter(|item| *item > 0)
}
fn normalize_non_empty(value: Option<&String>) -> Option<String> {
value
.map(|item| item.trim())
.filter(|item| !item.is_empty())
.map(ToOwned::to_owned)
}
fn parse_proxy_log_client_filter(value: Option<&String>) -> Option<ProxyLogClientFilter> {
let text = normalize_non_empty(value)?;
if let Some((scope, token)) = text.split_once(':') {
let token = token.trim();
if !token.is_empty() {
return match scope.trim().to_ascii_lowercase().as_str() {
"app" => Some(ProxyLogClientFilter::App(token.to_string())),
"family" => Some(ProxyLogClientFilter::Family(token.to_ascii_lowercase())),
_ => Some(ProxyLogClientFilter::Raw(text)),
};
}
}
Some(ProxyLogClientFilter::Raw(text))
}
fn parse_proxy_log_filters(q: &std::collections::HashMap<String, String>) -> ProxyLogFilters {
let search = normalize_non_empty(q.get("search").or_else(|| q.get("q")));
ProxyLogFilters {
status: normalize_proxy_log_status(q.get("status")),
search_like: search.map(|value| format!("%{value}%")),
client: parse_proxy_log_client_filter(q.get("client")),
site_id: normalize_positive_i64(q.get("siteId").or_else(|| q.get("site_id"))),
downstream_key_id: normalize_positive_i64(
q.get("downstreamKeyId")
.or_else(|| q.get("downstream_key_id"))
.or_else(|| q.get("keyId"))
.or_else(|| q.get("key_id")),
),
from: normalize_non_empty(q.get("from")),
to: normalize_non_empty(q.get("to")),
}
}
fn push_proxy_log_condition_prefix(
builder: &mut QueryBuilder<'_, Sqlite>,
condition_count: &mut usize,
) {
builder.push(if *condition_count == 0 {
" WHERE "
} else {
" AND "
});
*condition_count += 1;
}
fn push_proxy_log_raw_condition(
builder: &mut QueryBuilder<'_, Sqlite>,
condition_count: &mut usize,
condition: &str,
) {
push_proxy_log_condition_prefix(builder, condition_count);
builder.push(condition);
}
fn push_proxy_log_where(
builder: &mut QueryBuilder<'_, Sqlite>,
filters: &ProxyLogFilters,
) -> usize {
let mut condition_count = 0;
if let Some(status) = &filters.status {
push_proxy_log_condition_prefix(builder, &mut condition_count);
builder.push("l.status = ");
builder.push_bind(status.clone());
}
if let Some(search_like) = &filters.search_like {
push_proxy_log_condition_prefix(builder, &mut condition_count);
builder.push("(l.model_requested LIKE ");
builder.push_bind(search_like.clone());
builder.push(" OR l.model_actual LIKE ");
builder.push_bind(search_like.clone());
builder.push(" OR l.error_message LIKE ");
builder.push_bind(search_like.clone());
builder.push(" OR l.content_summary LIKE ");
builder.push_bind(search_like.clone());
builder.push(" OR dk.name LIKE ");
builder.push_bind(search_like.clone());
builder.push(" OR dk.group_name LIKE ");
builder.push_bind(search_like.clone());
builder.push(" OR dk.tags LIKE ");
builder.push_bind(search_like.clone());
builder.push(" OR a.username LIKE ");
builder.push_bind(search_like.clone());
builder.push(" OR s.name LIKE ");
builder.push_bind(search_like.clone());
builder.push(")");
}
if let Some(client) = &filters.client {
push_proxy_log_condition_prefix(builder, &mut condition_count);
match client {
ProxyLogClientFilter::App(value) => {
builder.push("(l.client_app_id = ");
builder.push_bind(value.clone());
builder.push(" OR l.client_app_name = ");
builder.push_bind(value.clone());
builder.push(")");
}
ProxyLogClientFilter::Family(value) => {
builder.push("LOWER(l.client_family) = ");
builder.push_bind(value.clone());
}
ProxyLogClientFilter::Raw(value) => {
builder.push("(l.client_app_id = ");
builder.push_bind(value.clone());
builder.push(" OR l.client_family = ");
builder.push_bind(value.clone());
builder.push(" OR l.client_app_name = ");
builder.push_bind(value.clone());
builder.push(")");
}
}
}
if let Some(site_id) = filters.site_id {
push_proxy_log_condition_prefix(builder, &mut condition_count);
builder.push("a.site_id = ");
builder.push_bind(site_id);
}
if let Some(downstream_key_id) = filters.downstream_key_id {
push_proxy_log_condition_prefix(builder, &mut condition_count);
builder.push("l.downstream_api_key_id = ");
builder.push_bind(downstream_key_id);
}
if let Some(from) = &filters.from {
push_proxy_log_condition_prefix(builder, &mut condition_count);
builder.push("datetime(l.created_at) >= datetime(");
builder.push_bind(from.clone());
builder.push(")");
}
if let Some(to) = &filters.to {
push_proxy_log_condition_prefix(builder, &mut condition_count);
builder.push("datetime(l.created_at) <= datetime(");
builder.push_bind(to.clone());
builder.push(")");
}
condition_count
}
fn push_proxy_log_joined_from(builder: &mut QueryBuilder<'_, Sqlite>) {
builder.push(
r#"
FROM proxy_logs l
LEFT JOIN accounts a ON a.id = l.account_id
LEFT JOIN sites s ON s.id = a.site_id
LEFT JOIN downstream_api_keys dk ON dk.id = l.downstream_api_key_id
"#,
);
}
fn normalize_owned_non_empty(value: Option<String>) -> Option<String> {
value
.map(|item| item.trim().to_string())
.filter(|item| !item.is_empty())
}
fn proxy_log_client_family_label(value: &str) -> String {
match value.trim().to_ascii_lowercase().as_str() {
"codex" => "Codex".to_string(),
"claude_code" => "Claude Code".to_string(),
"gemini_cli" => "Gemini CLI".to_string(),
"generic" => "通用".to_string(),
_ => value.trim().to_string(),
}
}
fn push_proxy_log_client_option(
options: &mut Vec<Value>,
seen: &mut HashSet<String>,
value: String,
label: String,
) {
if seen.insert(value.clone()) {
options.push(json!({ "value": value, "label": label }));
}
}
async fn build_proxy_log_client_options(
state: &AppState,
filters: &ProxyLogFilters,
) -> AppResult<Vec<Value>> {
let mut scoped = filters.clone();
scoped.client = None;
let mut options = Vec::new();
let mut seen = HashSet::new();
let mut app_query =
QueryBuilder::<Sqlite>::new("SELECT DISTINCT l.client_app_id, l.client_app_name ");
push_proxy_log_joined_from(&mut app_query);
let mut condition_count = push_proxy_log_where(&mut app_query, &scoped);
push_proxy_log_raw_condition(
&mut app_query,
&mut condition_count,
"((l.client_app_id IS NOT NULL AND TRIM(l.client_app_id) <> '') OR (l.client_app_name IS NOT NULL AND TRIM(l.client_app_name) <> ''))",
);
app_query.push(
" ORDER BY LOWER(COALESCE(NULLIF(l.client_app_name, ''), l.client_app_id)) ASC LIMIT 200",
);
let app_rows: Vec<(Option<String>, Option<String>)> =
app_query.build_query_as().fetch_all(&*state.db).await?;
for (client_app_id, client_app_name) in app_rows {
let app_id = normalize_owned_non_empty(client_app_id);
let app_name = normalize_owned_non_empty(client_app_name);
let Some(token) = app_id.clone().or_else(|| app_name.clone()) else {
continue;
};
let label_name = app_name.or(app_id).unwrap_or_else(|| token.clone());
push_proxy_log_client_option(
&mut options,
&mut seen,
format!("app:{token}"),
format!("应用 · {label_name}"),
);
}
let mut family_query = QueryBuilder::<Sqlite>::new("SELECT DISTINCT l.client_family ");
push_proxy_log_joined_from(&mut family_query);
let mut condition_count = push_proxy_log_where(&mut family_query, &scoped);
push_proxy_log_raw_condition(
&mut family_query,
&mut condition_count,
"l.client_family IS NOT NULL AND TRIM(l.client_family) <> ''",
);
family_query.push(" ORDER BY LOWER(l.client_family) ASC LIMIT 200");
let family_rows: Vec<(Option<String>,)> =
family_query.build_query_as().fetch_all(&*state.db).await?;
for (client_family,) in family_rows {
let Some(family) = normalize_owned_non_empty(client_family) else {
continue;
};
let token = family.to_ascii_lowercase();
push_proxy_log_client_option(
&mut options,
&mut seen,
format!("family:{token}"),
format!("协议 · {}", proxy_log_client_family_label(&family)),
);
}
Ok(options)
}
async fn build_proxy_log_site_options(
state: &AppState,
filters: &ProxyLogFilters,
) -> AppResult<Vec<Value>> {
let mut scoped = filters.clone();
scoped.site_id = None;
let mut site_query = QueryBuilder::<Sqlite>::new(
"SELECT DISTINCT s.id, COALESCE(NULLIF(s.name, ''), '站点 #' || s.id) AS name, s.status ",
);
push_proxy_log_joined_from(&mut site_query);
let mut condition_count = push_proxy_log_where(&mut site_query, &scoped);
push_proxy_log_raw_condition(&mut site_query, &mut condition_count, "s.id IS NOT NULL");
site_query
.push(" ORDER BY LOWER(COALESCE(NULLIF(s.name, ''), '站点 #' || s.id)) ASC LIMIT 200");
let rows: Vec<(i64, String, Option<String>)> =
site_query.build_query_as().fetch_all(&*state.db).await?;
Ok(rows
.into_iter()
.map(|(id, name, status)| json!({ "id": id, "name": name, "status": status }))
.collect())
}
pub fn routes() -> Router<AppState> {
Router::new()
.route("/api/stats/dashboard", get(dashboard))
@@ -417,48 +735,77 @@ async fn proxy_logs(
.and_then(|v| v.parse::<i64>().ok())
.unwrap_or(0)
.max(0);
let view = q.get("view").map(String::as_str);
let meta_only = matches!(view, Some("meta"));
let include_meta = !matches!(view, Some("query"));
let filters = parse_proxy_log_filters(&q);
let items: Vec<Value> = sqlx::query_as::<_, ProxyLogRow>(
let items: Vec<Value> = if meta_only {
Vec::new()
} else {
let mut list_query = QueryBuilder::<Sqlite>::new(
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.billing_details, l.content_summary, l.request_content, l.response_content,
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,
dk.name AS downstream_key_name,
dk.group_name AS downstream_key_group_name,
dk.tags AS downstream_key_tags
"#,
);
push_proxy_log_joined_from(&mut list_query);
push_proxy_log_where(&mut list_query, &filters);
list_query.push(
r#"
ORDER BY l.created_at DESC
"#,
);
list_query.push(" LIMIT ");
list_query.push_bind(limit);
list_query.push(" OFFSET ");
list_query.push_bind(offset);
list_query
.build_query_as::<ProxyLogRow>()
.fetch_all(&*state.db)
.await?
.into_iter()
.map(|row| serde_json::to_value(row).unwrap_or_else(|_| json!({})))
.collect()
};
let mut summary_query = QueryBuilder::<Sqlite>::new(
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.content_summary, l.request_content, l.response_content,
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,
dk.name AS downstream_key_name,
dk.group_name AS downstream_key_group_name,
dk.tags AS downstream_key_tags
FROM proxy_logs l
LEFT JOIN accounts a ON a.id = l.account_id
LEFT JOIN sites s ON s.id = a.site_id
LEFT JOIN downstream_api_keys dk ON dk.id = l.downstream_api_key_id
ORDER BY l.created_at DESC
LIMIT ? OFFSET ?
COUNT(l.id) AS total,
COALESCE(SUM(CASE WHEN l.status = 'success' THEN 1 ELSE 0 END), 0) AS success,
COALESCE(SUM(CASE WHEN l.status != 'success' OR l.status IS NULL THEN 1 ELSE 0 END), 0) AS failed,
COALESCE(SUM(l.estimated_cost), 0.0) AS cost,
COALESCE(SUM(l.total_tokens), 0) AS total_tokens
"#,
)
.bind(limit)
.bind(offset)
.fetch_all(&*state.db)
.await?
.into_iter()
.map(|row| serde_json::to_value(row).unwrap_or_else(|_| json!({})))
.collect();
let total: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM proxy_logs")
.fetch_one(&*state.db)
.await?;
let success: i64 =
sqlx::query_scalar("SELECT COUNT(*) FROM proxy_logs WHERE status = 'success'")
.fetch_one(&*state.db)
.await?;
);
push_proxy_log_joined_from(&mut summary_query);
push_proxy_log_where(&mut summary_query, &filters);
let (total, success, failed, cost, total_tokens): (i64, i64, i64, f64, i64) =
summary_query.build_query_as().fetch_one(&*state.db).await?;
let (client_options, sites) = if include_meta {
(
build_proxy_log_client_options(&state, &filters).await?,
build_proxy_log_site_options(&state, &filters).await?,
)
} else {
(Vec::new(), Vec::new())
};
Ok(Json(json!({
"items": items,
@@ -469,10 +816,12 @@ async fn proxy_logs(
"summary": {
"total": total,
"success": success,
"failed": total - success,
"failed": failed,
"cost": cost,
"totalTokens": total_tokens,
},
"clientOptions": [],
"sites": [],
"clientOptions": client_options,
"sites": sites,
})))
}
@@ -507,6 +856,7 @@ pub struct ProxyLogRow {
client_app_id: Option<String>,
client_app_name: Option<String>,
client_confidence: Option<String>,
billing_details: Option<String>,
content_summary: Option<String>,
request_content: Option<String>,
response_content: Option<String>,
@@ -526,7 +876,7 @@ async fn proxy_log_detail(
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.content_summary, l.request_content, l.response_content,
l.billing_details, l.content_summary, l.request_content, l.response_content,
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,
@@ -989,3 +1339,71 @@ async fn site_trend(
Ok(Json(json!({ "trend": trend })))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn proxy_log_filters_accept_downstream_key_aliases() {
let mut params = std::collections::HashMap::new();
params.insert("status".to_string(), "FAILED".to_string());
params.insert("q".to_string(), " gpt-4o ".to_string());
params.insert("site_id".to_string(), "9".to_string());
params.insert("key_id".to_string(), "17".to_string());
params.insert("client".to_string(), "family:codex".to_string());
let filters = parse_proxy_log_filters(&params);
assert_eq!(filters.status.as_deref(), Some("failed"));
assert_eq!(filters.search_like.as_deref(), Some("%gpt-4o%"));
assert_eq!(filters.site_id, Some(9));
assert_eq!(filters.downstream_key_id, Some(17));
assert_eq!(
filters.client,
Some(ProxyLogClientFilter::Family("codex".to_string())),
);
}
#[test]
fn proxy_log_filters_parse_client_scopes() {
let mut params = std::collections::HashMap::new();
params.insert("client".to_string(), " app:cherry_studio ".to_string());
let filters = parse_proxy_log_filters(&params);
assert_eq!(
filters.client,
Some(ProxyLogClientFilter::App("cherry_studio".to_string())),
);
params.insert("client".to_string(), "FAMILY:Codex".to_string());
let filters = parse_proxy_log_filters(&params);
assert_eq!(
filters.client,
Some(ProxyLogClientFilter::Family("codex".to_string())),
);
params.insert("client".to_string(), "legacy-client".to_string());
let filters = parse_proxy_log_filters(&params);
assert_eq!(
filters.client,
Some(ProxyLogClientFilter::Raw("legacy-client".to_string())),
);
}
#[test]
fn proxy_log_filters_drop_invalid_values() {
let mut params = std::collections::HashMap::new();
params.insert("status".to_string(), "warning".to_string());
params.insert("downstreamKeyId".to_string(), "-3".to_string());
params.insert("siteId".to_string(), "0".to_string());
params.insert("search".to_string(), " ".to_string());
let filters = parse_proxy_log_filters(&params);
assert!(filters.status.is_none());
assert!(filters.downstream_key_id.is_none());
assert!(filters.site_id.is_none());
assert!(filters.search_like.is_none());
}
}
+1414 -109
View File
File diff suppressed because it is too large Load Diff
+118 -6
View File
@@ -15,6 +15,7 @@ use std::convert::Infallible;
use tracing::info;
use crate::{
auth::jwt::validate_jwt,
error::{AppError, AppResult},
AppState,
};
@@ -200,15 +201,28 @@ async fn external_comfyui_status_value(state: &AppState) -> Value {
}
fn check_auth(headers: &HeaderMap, state: &AppState) -> AppResult<()> {
let expected = format!("Bearer {}", state.config.auth.auth_token);
let actual = headers
let token = headers
.get(header::AUTHORIZATION)
.and_then(|v| v.to_str().ok())
.unwrap_or("");
.and_then(|value| {
value
.strip_prefix("Bearer ")
.or_else(|| value.strip_prefix("bearer "))
})
.map(str::trim)
.filter(|value| !value.is_empty())
.ok_or_else(|| AppError::Unauthorized("Missing Authorization header".to_string()))?;
if actual != expected {
return Err(AppError::BadRequest("unauthorized".to_string()));
if token == state.config.auth.auth_token {
return Ok(());
}
let claims = validate_jwt(&state.config.auth.jwt_secret, token)
.map_err(|_| AppError::Unauthorized("Invalid admin token".to_string()))?;
if claims.role != "admin" {
return Err(AppError::Forbidden(
"Admin account is required for this route".to_string(),
));
}
Ok(())
@@ -633,3 +647,101 @@ async fn full_workflow(
}
})))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{
auth::jwt::generate_jwt,
config::{AppConfig, AuthConfig, ComfyUIConfig, DbConfig, ProxyConfig, ServerConfig},
services::{builtin_keys::BuiltinKeys, runtime_auth::RuntimeAuthSettings},
};
use axum::http::HeaderValue;
use sqlx::sqlite::SqlitePoolOptions;
async fn test_state() -> AppState {
let pool = SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.expect("sqlite pool");
AppState::new(
AppConfig {
server: ServerConfig {
host: "127.0.0.1".to_string(),
port: 4000,
log_level: "info".to_string(),
cors_origins: Vec::new(),
},
db: DbConfig {
url: "sqlite::memory:".to_string(),
max_connections: 1,
ssl: false,
runtime_db_url: None,
runtime_db_dialect: None,
runtime_db_ssl: None,
},
auth: AuthConfig {
jwt_secret: "test-jwt-secret".to_string(),
jwt_expiry_hours: 24,
auth_token: "test-admin-token".to_string(),
proxy_token: "sk-test-proxy".to_string(),
},
proxy: ProxyConfig {
upstream_timeout_secs: 30,
first_byte_timeout_secs: 30,
max_retries: 1,
proxy_log_retention_days: 7,
template_proxy_url: None,
},
comfyui: ComfyUIConfig::default(),
},
pool,
BuiltinKeys {
comfyui_agent: (0, "sk-comfyui-test".to_string()),
},
RuntimeAuthSettings::from_token_for_test("sk-test-proxy"),
)
}
fn bearer_headers(token: &str) -> HeaderMap {
let mut headers = HeaderMap::new();
headers.insert(
header::AUTHORIZATION,
HeaderValue::from_str(&format!("Bearer {token}")).expect("valid bearer header"),
);
headers
}
#[tokio::test]
async fn legacy_video_agent_auth_accepts_admin_jwt() {
let state = test_state().await;
let token = generate_jwt(
&state.config.auth.jwt_secret,
"42",
"admin",
state.config.auth.jwt_expiry_hours,
)
.expect("admin jwt");
check_auth(&bearer_headers(&token), &state)
.expect("admin jwt can access legacy video agent routes");
}
#[tokio::test]
async fn legacy_video_agent_auth_rejects_user_jwt() {
let state = test_state().await;
let token = generate_jwt(
&state.config.auth.jwt_secret,
"43",
"user",
state.config.auth.jwt_expiry_hours,
)
.expect("user jwt");
let result = check_auth(&bearer_headers(&token), &state);
assert!(matches!(result, Err(AppError::Forbidden(_))));
}
}
+158 -27
View File
@@ -30,7 +30,7 @@ use serde_json::{json, Map, Value};
use sqlx::Row;
use std::{
collections::{HashMap, HashSet},
path::PathBuf,
path::{Path as StdPath, PathBuf},
sync::{mpsc, Arc},
time::{Duration, Instant, SystemTime, UNIX_EPOCH},
};
@@ -46,6 +46,7 @@ const NATIVE_WORKFLOW_CACHE_MAX_AGE: Duration = Duration::from_secs(21600);
const WORKFLOW_UPLOAD_BODY_LIMIT_BYTES: usize = 256 * 1024 * 1024;
const OFFICIAL_WORKFLOW_TEMPLATES_BASE: &str =
"https://raw.githubusercontent.com/Comfy-Org/workflow_templates/main/templates";
const DEFAULT_TEXT_TO_VIDEO_MODEL: &str = "seedance-1-0-lite-t2v-250428";
static EMBEDDED_WORKFLOW_UI: Dir<'_> = include_dir!("$CARGO_MANIFEST_DIR/dist/workflow-ui");
static EMBEDDED_WORKFLOW_UI_EXTERNAL: Dir<'_> =
include_dir!("$CARGO_MANIFEST_DIR/dist/workflow-ui-external");
@@ -210,9 +211,9 @@ pub fn routes() -> Router<AppState> {
fn workflow_ui_root_for_config(config: &crate::config::ComfyUIConfig) -> PathBuf {
if let Ok(path) = std::env::var("BOOSAPI_WORKFLOW_UI_DIR") {
return PathBuf::from(path);
return resolve_workflow_ui_dir(PathBuf::from(path));
}
let preferred = if config.uses_cloud_backend() {
let preferred = resolve_workflow_ui_dir(if config.uses_cloud_backend() {
std::env::var("BOOSAPI_WORKFLOW_UI_CLOUD_DIR")
.map(PathBuf::from)
.unwrap_or_else(|_| PathBuf::from("dist/workflow-ui-cloud"))
@@ -222,11 +223,23 @@ fn workflow_ui_root_for_config(config: &crate::config::ComfyUIConfig) -> PathBuf
.unwrap_or_else(|_| PathBuf::from("dist/workflow-ui-external"))
} else {
PathBuf::from("dist/workflow-ui")
};
});
if workflow_ui_dir_is_usable(&preferred) {
preferred
} else {
PathBuf::from("dist/workflow-ui")
resolve_workflow_ui_dir(PathBuf::from("dist/workflow-ui"))
}
}
fn resolve_workflow_ui_dir(path: PathBuf) -> PathBuf {
if path.is_absolute() || workflow_ui_dir_is_usable(&path) {
return path;
}
let manifest_path = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join(&path);
if workflow_ui_dir_is_usable(&manifest_path) {
manifest_path
} else {
path
}
}
@@ -272,14 +285,24 @@ async fn workflow_ui_runtime(State(state): State<AppState>) -> Json<Value> {
} else {
"boosapi"
};
let frontend_dir = workflow_ui_root_for_config(&state.config.comfyui);
Json(json!({
"mode": mode,
"baseUrl": state.config.comfyui.target_base_url(),
"frontendDir": workflow_ui_root_for_config(&state.config.comfyui),
"frontendDir": frontend_dir,
"frontendBuild": workflow_ui_build_metadata(&frontend_dir),
"cloudApiKeyConfigured": workflow_cloud_api_key(&state.config.comfyui).is_some(),
}))
}
fn workflow_ui_build_metadata(dir: &StdPath) -> Value {
let path = dir.join("boosapi-workflow-ui.json");
std::fs::read_to_string(path)
.ok()
.and_then(|raw| serde_json::from_str::<Value>(&raw).ok())
.unwrap_or(Value::Null)
}
fn workflow_cloud_api_key(config: &crate::config::ComfyUIConfig) -> Option<&str> {
config
.api_key
@@ -675,7 +698,7 @@ fn text_to_video_workflow_graph() -> Value {
{ "name": "raw", "type": "JSON", "links": null, "slot_index": 2 }
],
"properties": { "Node name for S&R": "boosapi.text_to_video" },
"widgets_values": ["seedance-1-0-lite-t2v", "16:9", 5, "720p", 24, true, 5, 180]
"widgets_values": [DEFAULT_TEXT_TO_VIDEO_MODEL, "16:9", 5, "720p", 24, true, 5, 180]
},
{
"id": 3,
@@ -691,7 +714,7 @@ fn text_to_video_workflow_graph() -> Value {
{ "name": "raw", "type": "JSON", "links": null, "slot_index": 1 }
],
"properties": { "Node name for S&R": "boosapi.video_status" },
"widgets_values": ["seedance-1-0-lite-t2v"]
"widgets_values": [DEFAULT_TEXT_TO_VIDEO_MODEL]
},
{
"id": 4,
@@ -1073,11 +1096,10 @@ fn workflow_folder_models(folder: &str, models: &[String]) -> Vec<String> {
.cloned()
.collect(),
"clip" => models.to_vec(),
"video_models" => models
.iter()
.filter(|model| is_video_workflow_model(model))
.cloned()
.collect(),
// Comfy's video_models folder is a native Comfy API model folder. Keep it
// separate from BoosAPI route models so video apps submit real Comfy API
// model names instead of OpenAI-compatible route aliases.
"video_models" => Vec::new(),
"audio_tts_models" => models
.iter()
.filter(|model| is_audio_workflow_model(model))
@@ -1102,6 +1124,32 @@ fn workflow_folder_models(folder: &str, models: &[String]) -> Vec<String> {
}
}
fn workflow_folder_is_api_only(folder: &str) -> bool {
matches!(
folder.trim().to_ascii_lowercase().as_str(),
"boosapi" | "boosapi_models" | "api_models"
)
}
fn workflow_folder_has_route_models(folder: &str) -> bool {
matches!(
folder.trim().to_ascii_lowercase().as_str(),
"boosapi"
| "boosapi_models"
| "api_models"
| "image_models"
| "chat_models"
| "audio_tts_models"
| "embeddings"
| "3d_models"
| "vision_models"
)
}
fn workflow_folder_route_model_items(folder: &str, route_models: &[String]) -> Vec<Value> {
model_payload_items_from_names(workflow_folder_models(folder, route_models), 0)
}
fn combo_choices_from_node_input(
object_info: &Value,
node_type: &str,
@@ -1231,7 +1279,7 @@ fn inject_model_choices(defs: &mut Map<String, Value>, models: &[String]) {
{
optional.insert(
"model".to_string(),
comfy_combo_input(&video_models, "seedance-1-0-lite-t2v"),
comfy_combo_input(&video_models, DEFAULT_TEXT_TO_VIDEO_MODEL),
);
}
}
@@ -1788,7 +1836,7 @@ async fn serve_file(path: PathBuf) -> AppResult<Response> {
}
}
fn content_type_for_path(path: &PathBuf) -> &'static str {
fn content_type_for_path(path: &StdPath) -> &'static str {
match path
.extension()
.and_then(|ext| ext.to_str())
@@ -1809,15 +1857,15 @@ fn content_type_for_path(path: &PathBuf) -> &'static str {
}
}
fn cache_control_for_path(path: &PathBuf) -> &'static str {
fn cache_control_for_path(path: &StdPath) -> &'static str {
let ext = path
.extension()
.and_then(|ext| ext.to_str())
.unwrap_or("")
.to_ascii_lowercase();
match ext.as_str() {
"html" => "no-cache",
"js" | "mjs" | "css" | "woff2" | "png" | "jpg" | "jpeg" | "webp" | "svg" => {
"html" | "js" | "mjs" | "css" => "no-cache",
"woff2" | "png" | "jpg" | "jpeg" | "webp" | "svg" => {
"public, max-age=604800"
}
_ => "public, max-age=3600",
@@ -4792,15 +4840,28 @@ async fn experiment_models_in_folder(
) -> AppResult<Json<Value>> {
if state.config.comfyui.uses_cloud_backend() {
let _ = auth_caller_optional(&state, &headers);
let normalized_folder = folder.trim().to_ascii_lowercase();
if matches!(
normalized_folder.as_str(),
"boosapi" | "boosapi_models" | "api_models"
) {
if workflow_folder_has_route_models(&folder) {
let route_models = workflow_route_models(&state).await.unwrap_or_default();
let api_model_items =
model_payload_items_from_names(workflow_folder_models(&folder, &route_models), 0);
return Ok(Json(Value::Array(api_model_items)));
let api_model_items = workflow_folder_route_model_items(&folder, &route_models);
if workflow_folder_is_api_only(&folder) {
return Ok(Json(Value::Array(api_model_items)));
}
let mut native_items = Vec::<Value>::new();
let path = format!("/experiment/models/{}", simple_urlencode(&folder));
if let Ok(value) = native_workflow_json_cached(
&state,
Method::GET,
&path,
None,
native_workflow_metadata_ttl(&state),
)
.await
{
if let Value::Array(items) = value.0 {
native_items.extend(items);
}
}
return Ok(Json(merge_model_payload_items(api_model_items, native_items)));
}
let path = format!("/experiment/models/{}", simple_urlencode(&folder));
return native_workflow_json_cached(
@@ -4813,7 +4874,7 @@ async fn experiment_models_in_folder(
.await;
}
let route_models = workflow_route_models(&state).await.unwrap_or_default();
let api_model_items = model_payload_items_from_names(workflow_folder_models(&folder, &route_models), 0);
let api_model_items = workflow_folder_route_model_items(&folder, &route_models);
if state.config.comfyui.uses_native_backend() {
let _ = auth_caller_optional(&state, &headers);
@@ -5701,3 +5762,73 @@ fn parse_workflow_cookie(raw: &str) -> Option<&str> {
.map(|value| value.trim())
.filter(|value| !value.is_empty())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn video_model_folder_does_not_use_route_video_models() {
let route_models = vec![
"gpt-5.4".to_string(),
"gpt-image-2".to_string(),
"seedance-1-0-pro-t2v-250528".to_string(),
"wan2.1-i2v".to_string(),
"text-embedding-3-large".to_string(),
];
let names = workflow_folder_models("video_models", &route_models);
assert!(names.is_empty());
}
#[test]
fn api_backed_workflow_folders_are_not_native_only() {
assert!(!workflow_folder_has_route_models("video_models"));
assert!(workflow_folder_has_route_models("image_models"));
assert!(workflow_folder_has_route_models("api_models"));
assert!(workflow_folder_is_api_only("api_models"));
assert!(!workflow_folder_is_api_only("video_models"));
assert!(!workflow_folder_has_route_models("checkpoints"));
}
#[test]
fn video_model_folder_prefers_native_comfy_items_without_route_aliases() {
let route_models = vec![
"seedance-1-0-pro-t2v-250528".to_string(),
"wan2.1-t2v".to_string(),
];
let api_items = workflow_folder_route_model_items("video_models", &route_models);
let native_items = vec![
json!({ "name": "wan2.1-t2v", "pathIndex": 2 }),
json!({ "name": "local-video-model.safetensors", "pathIndex": 1 }),
];
let merged = merge_model_payload_items(api_items, native_items);
let names = merged
.as_array()
.expect("merged model payload should be an array")
.iter()
.filter_map(|item| item.get("name").and_then(Value::as_str))
.collect::<Vec<_>>();
assert_eq!(
names,
vec![
"wan2.1-t2v",
"local-video-model.safetensors",
]
);
}
#[test]
fn workflow_ui_runtime_assets_revalidate_after_install() {
assert_eq!(cache_control_for_path(StdPath::new("index.html")), "no-cache");
assert_eq!(cache_control_for_path(StdPath::new("assets/index.js")), "no-cache");
assert_eq!(cache_control_for_path(StdPath::new("assets/index.css")), "no-cache");
assert_eq!(
cache_control_for_path(StdPath::new("assets/logo.svg")),
"public, max-age=604800"
);
}
}