forked from boos/BoosAPI
feat: refine BoosAPI console and user workspace
This commit is contained in:
+255
-11
@@ -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
File diff suppressed because it is too large
Load Diff
+63
-15
@@ -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
@@ -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
@@ -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
@@ -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");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
File diff suppressed because it is too large
Load Diff
@@ -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
@@ -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
@@ -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(¶ms);
|
||||
|
||||
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(¶ms);
|
||||
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(¶ms);
|
||||
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(¶ms);
|
||||
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(¶ms);
|
||||
|
||||
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
File diff suppressed because it is too large
Load Diff
@@ -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
@@ -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"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user