feat: Rust refactor — full backend rewrite (Phases 1-6)

Phase 1-3: Scaffold + types + auth (JWT, Argon2, API keys)
Phase 4a: Admin API (sites/accounts/tokens/downstream-keys/settings)
Phase 5: Proxy core (/v1/chat, /v1/images, /v1/models)
Phase 6: ComfyUI proxy routes migrated

- 80+ source files, all cargo check + 18 tests green
- Drop-in compatible with Node.js API surface
This commit is contained in:
boos
2026-05-17 13:20:52 +08:00
parent ad0bb86e31
commit 7987833ba9
61 changed files with 3661 additions and 1 deletions
+195
View File
@@ -0,0 +1,195 @@
use crate::{
db::models::account::Account,
error::{AppError, AppResult},
AppState,
};
use axum::{
extract::{Path, Query, State},
http::StatusCode,
Json,
};
use serde::Deserialize;
use serde_json::{json, Value};
use tracing::info;
#[derive(Debug, Deserialize)]
pub struct ListAccountsQuery {
pub site_id: Option<i64>,
}
#[derive(Debug, Deserialize)]
pub struct CreateAccountRequest {
pub site_id: i64,
pub username: Option<String>,
pub access_token: String,
pub api_token: Option<String>,
pub status: Option<String>,
pub checkin_enabled: Option<bool>,
pub extra_config: Option<String>,
pub oauth_provider: Option<String>,
pub oauth_account_key: Option<String>,
pub oauth_project_id: Option<String>,
}
#[derive(Debug, Deserialize)]
pub struct UpdateAccountRequest {
pub username: Option<String>,
pub access_token: Option<String>,
pub api_token: Option<String>,
pub status: Option<String>,
pub checkin_enabled: Option<bool>,
pub extra_config: Option<String>,
pub oauth_provider: Option<String>,
pub oauth_account_key: Option<String>,
pub oauth_project_id: Option<String>,
}
/// GET /api/accounts?site_id=
pub async fn list_accounts(
State(state): State<AppState>,
Query(query): Query<ListAccountsQuery>,
) -> AppResult<Json<Value>> {
info!("Listing accounts, site_id filter: {:?}", query.site_id);
let accounts: Vec<Account> = if let Some(site_id) = query.site_id {
sqlx::query_as::<_, Account>(
"SELECT * FROM accounts WHERE site_id = ? ORDER BY sort_order, id",
)
.bind(site_id)
.fetch_all(&*state.db)
.await?
} else {
sqlx::query_as::<_, Account>("SELECT * FROM accounts ORDER BY site_id, sort_order, id")
.fetch_all(&*state.db)
.await?
};
Ok(Json(json!({ "data": accounts })))
}
/// POST /api/accounts
pub async fn create_account(
State(state): State<AppState>,
Json(payload): Json<CreateAccountRequest>,
) -> AppResult<(StatusCode, Json<Value>)> {
info!("Creating account for site_id: {}", payload.site_id);
if payload.access_token.trim().is_empty() {
return Err(AppError::BadRequest("access_token is required".into()));
}
let site_exists: Option<(i64,)> = sqlx::query_as("SELECT id FROM sites WHERE id = ?")
.bind(payload.site_id)
.fetch_optional(&*state.db)
.await?;
if site_exists.is_none() {
return Err(AppError::BadRequest(format!(
"Site with id {} not found",
payload.site_id
)));
}
let status = payload.status.unwrap_or_else(|| "active".to_string());
let result = sqlx::query(
"INSERT INTO accounts (site_id, username, access_token, api_token, status, checkin_enabled, extra_config, oauth_provider, oauth_account_key, oauth_project_id) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"
)
.bind(payload.site_id)
.bind(&payload.username)
.bind(&payload.access_token)
.bind(&payload.api_token)
.bind(&status)
.bind(payload.checkin_enabled)
.bind(&payload.extra_config)
.bind(&payload.oauth_provider)
.bind(&payload.oauth_account_key)
.bind(&payload.oauth_project_id)
.execute(&*state.db)
.await?;
let account_id = result.last_insert_rowid();
let account: Account = sqlx::query_as::<_, Account>("SELECT * FROM accounts WHERE id = ?")
.bind(account_id)
.fetch_one(&*state.db)
.await?;
Ok((StatusCode::CREATED, Json(json!({ "data": account }))))
}
/// GET /api/accounts/{id}
pub async fn get_account(
State(state): State<AppState>,
Path(id): Path<i64>,
) -> AppResult<Json<Value>> {
info!("Getting account: {}", id);
let account: Account = sqlx::query_as::<_, Account>("SELECT * FROM accounts WHERE id = ?")
.bind(id)
.fetch_optional(&*state.db)
.await?
.ok_or_else(|| AppError::NotFound(format!("Account with id {} not found", id)))?;
Ok(Json(json!({ "data": account })))
}
/// PUT /api/accounts/{id}
pub async fn update_account(
State(state): State<AppState>,
Path(id): Path<i64>,
Json(payload): Json<UpdateAccountRequest>,
) -> AppResult<Json<Value>> {
info!("Updating account: {}", id);
let existing: Account = sqlx::query_as::<_, Account>("SELECT * FROM accounts WHERE id = ?")
.bind(id)
.fetch_optional(&*state.db)
.await?
.ok_or_else(|| AppError::NotFound(format!("Account with id {} not found", id)))?;
sqlx::query(
"UPDATE accounts SET username = ?, access_token = ?, api_token = ?, status = ?, checkin_enabled = ?, extra_config = ?, oauth_provider = ?, oauth_account_key = ?, oauth_project_id = ?, updated_at = datetime('now') WHERE id = ?"
)
.bind(payload.username.or(existing.username))
.bind(payload.access_token.as_ref().unwrap_or(&existing.access_token))
.bind(payload.api_token.or(existing.api_token))
.bind(payload.status.unwrap_or(existing.status))
.bind(payload.checkin_enabled.or(existing.checkin_enabled))
.bind(payload.extra_config.or(existing.extra_config))
.bind(payload.oauth_provider.or(existing.oauth_provider))
.bind(payload.oauth_account_key.or(existing.oauth_account_key))
.bind(payload.oauth_project_id.or(existing.oauth_project_id))
.bind(id)
.execute(&*state.db)
.await?;
let account: Account = sqlx::query_as::<_, Account>("SELECT * FROM accounts WHERE id = ?")
.bind(id)
.fetch_one(&*state.db)
.await?;
Ok(Json(json!({ "data": account })))
}
/// DELETE /api/accounts/{id}
pub async fn delete_account(
State(state): State<AppState>,
Path(id): Path<i64>,
) -> AppResult<Json<Value>> {
info!("Deleting account: {}", id);
let affected = sqlx::query("DELETE FROM accounts WHERE id = ?")
.bind(id)
.execute(&*state.db)
.await?
.rows_affected();
if affected == 0 {
return Err(AppError::NotFound(format!(
"Account with id {} not found",
id
)));
}
Ok(Json(json!({ "data": { "deleted": true, "id": id } })))
}
@@ -0,0 +1,196 @@
use crate::{
db::models::downstream_api_key::DownstreamApiKey,
error::{AppError, AppResult},
AppState,
};
use axum::{
extract::{Path, State},
http::StatusCode,
Json,
};
use serde::Deserialize;
use serde_json::{json, Value};
use tracing::info;
use uuid::Uuid;
#[derive(Debug, Deserialize)]
pub struct CreateDownstreamKeyRequest {
pub name: String,
pub description: Option<String>,
pub group_name: Option<String>,
pub tags: Option<String>,
pub enabled: Option<bool>,
pub expires_at: Option<String>,
pub max_cost: Option<f64>,
pub max_requests: Option<i64>,
pub supported_models: Option<Vec<String>>,
pub allowed_route_ids: Option<String>,
}
#[derive(Debug, Deserialize)]
pub struct UpdateDownstreamKeyRequest {
pub name: Option<String>,
pub description: Option<String>,
pub group_name: Option<String>,
pub tags: Option<String>,
pub enabled: Option<bool>,
pub expires_at: Option<String>,
pub max_cost: Option<f64>,
pub max_requests: Option<i64>,
pub supported_models: Option<Vec<String>>,
pub allowed_route_ids: Option<String>,
}
/// Helper to validate models from a JSON string against a Vec<String>
fn validate_supported_models(models: &[String]) -> AppResult<()> {
for model in models {
if model.trim().is_empty() {
return Err(AppError::BadRequest(
"supported_models contains empty model name".into(),
));
}
}
Ok(())
}
/// GET /api/downstream-api-keys
pub async fn list_downstream_keys(State(state): State<AppState>) -> AppResult<Json<Value>> {
info!("Listing downstream API keys");
let keys: Vec<DownstreamApiKey> =
sqlx::query_as::<_, DownstreamApiKey>("SELECT * FROM downstream_api_keys ORDER BY name")
.fetch_all(&*state.db)
.await?;
Ok(Json(json!({ "data": keys })))
}
/// POST /api/downstream-api-keys
pub async fn create_downstream_key(
State(state): State<AppState>,
Json(payload): Json<CreateDownstreamKeyRequest>,
) -> AppResult<(StatusCode, Json<Value>)> {
info!("Creating downstream API key: {}", payload.name);
if payload.name.trim().is_empty() {
return Err(AppError::BadRequest("name is required".into()));
}
// Auto-generate key
let key = format!("sk-{}", Uuid::new_v4().to_string().replace('-', ""));
// Serialize supported_models to JSON string if provided
let supported_models: Option<String> = if let Some(ref models) = payload.supported_models {
validate_supported_models(models)?;
Some(
serde_json::to_string(models)
.map_err(|e| AppError::BadRequest(format!("Invalid supported_models: {}", e)))?,
)
} else {
None
};
let enabled = payload.enabled.unwrap_or(true);
let result = sqlx::query(
"INSERT INTO downstream_api_keys (name, key, description, group_name, tags, enabled, expires_at, max_cost, max_requests, supported_models, allowed_route_ids) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"
)
.bind(&payload.name)
.bind(&key)
.bind(&payload.description)
.bind(&payload.group_name)
.bind(&payload.tags)
.bind(enabled)
.bind(&payload.expires_at)
.bind(payload.max_cost)
.bind(payload.max_requests)
.bind(&supported_models)
.bind(&payload.allowed_route_ids)
.execute(&*state.db)
.await?;
let key_id = result.last_insert_rowid();
let api_key: DownstreamApiKey =
sqlx::query_as::<_, DownstreamApiKey>("SELECT * FROM downstream_api_keys WHERE id = ?")
.bind(key_id)
.fetch_one(&*state.db)
.await?;
Ok((StatusCode::CREATED, Json(json!({ "data": api_key }))))
}
/// PUT /api/downstream-api-keys/{id}
pub async fn update_downstream_key(
State(state): State<AppState>,
Path(id): Path<i64>,
Json(payload): Json<UpdateDownstreamKeyRequest>,
) -> AppResult<Json<Value>> {
info!("Updating downstream API key: {}", id);
let existing: DownstreamApiKey =
sqlx::query_as::<_, DownstreamApiKey>("SELECT * FROM downstream_api_keys WHERE id = ?")
.bind(id)
.fetch_optional(&*state.db)
.await?
.ok_or_else(|| {
AppError::NotFound(format!("Downstream key with id {} not found", id))
})?;
let supported_models: Option<String> = if let Some(ref models) = payload.supported_models {
validate_supported_models(models)?;
Some(
serde_json::to_string(models)
.map_err(|e| AppError::BadRequest(format!("Invalid supported_models: {}", e)))?,
)
} else {
existing.supported_models.clone()
};
sqlx::query(
"UPDATE downstream_api_keys SET name = ?, description = ?, group_name = ?, tags = ?, enabled = ?, expires_at = ?, max_cost = ?, max_requests = ?, supported_models = ?, allowed_route_ids = ?, updated_at = datetime('now') WHERE id = ?"
)
.bind(payload.name.unwrap_or(existing.name))
.bind(payload.description.or(existing.description))
.bind(payload.group_name.or(existing.group_name))
.bind(payload.tags.or(existing.tags))
.bind(payload.enabled.unwrap_or(existing.enabled.unwrap_or(true)))
.bind(payload.expires_at.or(existing.expires_at))
.bind(payload.max_cost.or(existing.max_cost))
.bind(payload.max_requests.or(existing.max_requests))
.bind(&supported_models)
.bind(payload.allowed_route_ids.or(existing.allowed_route_ids))
.bind(id)
.execute(&*state.db)
.await?;
let api_key: DownstreamApiKey =
sqlx::query_as::<_, DownstreamApiKey>("SELECT * FROM downstream_api_keys WHERE id = ?")
.bind(id)
.fetch_one(&*state.db)
.await?;
Ok(Json(json!({ "data": api_key })))
}
/// DELETE /api/downstream-api-keys/{id}
pub async fn delete_downstream_key(
State(state): State<AppState>,
Path(id): Path<i64>,
) -> AppResult<Json<Value>> {
info!("Deleting downstream API key: {}", id);
let affected = sqlx::query("DELETE FROM downstream_api_keys WHERE id = ?")
.bind(id)
.execute(&*state.db)
.await?
.rows_affected();
if affected == 0 {
return Err(AppError::NotFound(format!(
"Downstream key with id {} not found",
id
)));
}
Ok(Json(json!({ "data": { "deleted": true, "id": id } })))
}
+158
View File
@@ -0,0 +1,158 @@
use crate::{
db::models::setting::Setting,
error::{AppError, AppResult},
AppState,
};
use axum::{
extract::State,
http::HeaderMap,
middleware,
response::Response,
routing::{get, put},
Json, Router,
};
use serde_json::{json, Value};
use tracing::info;
mod accounts;
mod downstream_keys;
mod sites;
mod tokens;
/// Middleware to require admin authentication via Bearer token.
async fn require_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)
}
/// Build all admin routes. Returns Router<AppState> (without calling .with_state).
pub fn routes(state: AppState) -> Router<AppState> {
let sites_routes = Router::new()
.route(
"/api/sites",
get(sites::list_sites).post(sites::create_site),
)
.route(
"/api/sites/{id}",
get(sites::get_site)
.put(sites::update_site)
.delete(sites::delete_site),
);
let accounts_routes = Router::new()
.route(
"/api/accounts",
get(accounts::list_accounts).post(accounts::create_account),
)
.route(
"/api/accounts/{id}",
get(accounts::get_account)
.put(accounts::update_account)
.delete(accounts::delete_account),
);
let tokens_routes = Router::new()
.route(
"/api/tokens",
get(tokens::list_tokens).post(tokens::create_token),
)
.route(
"/api/tokens/refresh-route",
get(tokens::refresh_route).post(tokens::refresh_route),
)
.route("/api/tokens/route-decision", get(tokens::route_decision))
.route(
"/api/tokens/{id}",
get(tokens::get_token)
.put(tokens::update_token)
.delete(tokens::delete_token),
);
let downstream_routes = Router::new()
.route(
"/api/downstream-api-keys",
get(downstream_keys::list_downstream_keys).post(downstream_keys::create_downstream_key),
)
.route(
"/api/downstream-api-keys/{id}",
put(downstream_keys::update_downstream_key)
.delete(downstream_keys::delete_downstream_key),
);
let settings_routes =
Router::new().route("/api/settings", get(get_settings).put(update_setting));
Router::new()
.merge(sites_routes)
.merge(accounts_routes)
.merge(tokens_routes)
.merge(downstream_routes)
.merge(settings_routes)
.route_layer(middleware::from_fn_with_state(state, require_admin_auth))
}
// ── Settings handlers ────────────────────────────────────────────────
/// GET /api/settings — return all settings as a key-value map.
async fn get_settings(State(state): State<AppState>) -> AppResult<Json<Value>> {
info!("Getting all settings");
let rows: Vec<Setting> =
sqlx::query_as::<_, Setting>("SELECT key, value FROM settings ORDER BY key")
.fetch_all(&*state.db)
.await?;
let map: std::collections::BTreeMap<String, Option<String>> =
rows.into_iter().map(|s| (s.key, s.value)).collect();
Ok(Json(json!({ "data": map })))
}
/// PUT /api/settings — upsert a setting by key.
/// Expects JSON body: { "key": "...", "value": "..." }
async fn update_setting(
State(state): State<AppState>,
Json(payload): Json<serde_json::Value>,
) -> AppResult<Json<Value>> {
let key = payload
.get("key")
.and_then(|k| k.as_str())
.ok_or_else(|| AppError::BadRequest("key is required".into()))?;
let value = payload.get("value").and_then(|v| {
if v.is_null() {
None
} else {
v.as_str().map(|s| s.to_string())
}
});
info!("Upserting setting: {}", key);
sqlx::query("INSERT OR REPLACE INTO settings (key, value) VALUES (?, ?)")
.bind(key)
.bind(&value)
.execute(&*state.db)
.await?;
Ok(Json(json!({
"data": {
"key": key,
"value": value
}
})))
}
+168
View File
@@ -0,0 +1,168 @@
use crate::{
db::models::site::Site,
error::{AppError, AppResult},
AppState,
};
use axum::{
extract::{Path, State},
http::StatusCode,
Json,
};
use serde::Deserialize;
use serde_json::{json, Value};
use tracing::info;
#[derive(Debug, Deserialize)]
pub struct CreateSiteRequest {
pub name: String,
pub url: String,
pub platform: String,
pub status: Option<String>,
pub api_key: Option<String>,
pub proxy_url: Option<String>,
pub use_system_proxy: Option<bool>,
pub custom_headers: Option<String>,
pub global_weight: Option<f64>,
}
#[derive(Debug, Deserialize)]
pub struct UpdateSiteRequest {
pub name: Option<String>,
pub url: Option<String>,
pub platform: Option<String>,
pub status: Option<String>,
pub api_key: Option<String>,
pub proxy_url: Option<String>,
pub use_system_proxy: Option<bool>,
pub custom_headers: Option<String>,
pub global_weight: Option<f64>,
}
/// GET /api/sites
pub async fn list_sites(State(state): State<AppState>) -> AppResult<Json<Value>> {
info!("Listing all sites");
let sites: Vec<Site> =
sqlx::query_as::<_, Site>("SELECT * FROM sites ORDER BY sort_order, name")
.fetch_all(&*state.db)
.await?;
Ok(Json(json!({ "data": sites })))
}
/// POST /api/sites
pub async fn create_site(
State(state): State<AppState>,
Json(payload): Json<CreateSiteRequest>,
) -> AppResult<(StatusCode, Json<Value>)> {
info!("Creating site: {}", payload.name);
if payload.name.trim().is_empty() {
return Err(AppError::BadRequest("name is required".into()));
}
if payload.url.trim().is_empty() {
return Err(AppError::BadRequest("url is required".into()));
}
if payload.platform.trim().is_empty() {
return Err(AppError::BadRequest("platform is required".into()));
}
let status = payload.status.unwrap_or_else(|| "active".to_string());
let result = sqlx::query(
"INSERT INTO sites (name, url, platform, status, api_key, proxy_url, use_system_proxy, custom_headers, global_weight) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)"
)
.bind(&payload.name)
.bind(&payload.url)
.bind(&payload.platform)
.bind(&status)
.bind(&payload.api_key)
.bind(&payload.proxy_url)
.bind(payload.use_system_proxy)
.bind(&payload.custom_headers)
.bind(payload.global_weight)
.execute(&*state.db)
.await?;
let site_id = result.last_insert_rowid();
let site: Site = sqlx::query_as::<_, Site>("SELECT * FROM sites WHERE id = ?")
.bind(site_id)
.fetch_one(&*state.db)
.await?;
Ok((StatusCode::CREATED, Json(json!({ "data": site }))))
}
/// GET /api/sites/{id}
pub async fn get_site(
State(state): State<AppState>,
Path(id): Path<i64>,
) -> AppResult<Json<Value>> {
info!("Getting site: {}", id);
let site: Site = sqlx::query_as::<_, Site>("SELECT * FROM sites WHERE id = ?")
.bind(id)
.fetch_optional(&*state.db)
.await?
.ok_or_else(|| AppError::NotFound(format!("Site with id {} not found", id)))?;
Ok(Json(json!({ "data": site })))
}
/// PUT /api/sites/{id}
pub async fn update_site(
State(state): State<AppState>,
Path(id): Path<i64>,
Json(payload): Json<UpdateSiteRequest>,
) -> AppResult<Json<Value>> {
info!("Updating site: {}", id);
let existing: Site = sqlx::query_as::<_, Site>("SELECT * FROM sites WHERE id = ?")
.bind(id)
.fetch_optional(&*state.db)
.await?
.ok_or_else(|| AppError::NotFound(format!("Site with id {} not found", id)))?;
sqlx::query(
"UPDATE sites SET name = ?, url = ?, platform = ?, status = ?, api_key = ?, proxy_url = ?, use_system_proxy = ?, custom_headers = ?, global_weight = ?, updated_at = datetime('now') WHERE id = ?"
)
.bind(payload.name.unwrap_or(existing.name))
.bind(payload.url.unwrap_or(existing.url))
.bind(payload.platform.unwrap_or(existing.platform))
.bind(payload.status.unwrap_or(existing.status))
.bind(payload.api_key.or(existing.api_key))
.bind(payload.proxy_url.or(existing.proxy_url))
.bind(payload.use_system_proxy.or(existing.use_system_proxy))
.bind(payload.custom_headers.or(existing.custom_headers))
.bind(payload.global_weight.or(existing.global_weight))
.bind(id)
.execute(&*state.db)
.await?;
let site: Site = sqlx::query_as::<_, Site>("SELECT * FROM sites WHERE id = ?")
.bind(id)
.fetch_one(&*state.db)
.await?;
Ok(Json(json!({ "data": site })))
}
/// DELETE /api/sites/{id}
pub async fn delete_site(
State(state): State<AppState>,
Path(id): Path<i64>,
) -> AppResult<Json<Value>> {
info!("Deleting site: {}", id);
let affected = sqlx::query("DELETE FROM sites WHERE id = ?")
.bind(id)
.execute(&*state.db)
.await?
.rows_affected();
if affected == 0 {
return Err(AppError::NotFound(format!("Site with id {} not found", id)));
}
Ok(Json(json!({ "data": { "deleted": true, "id": id } })))
}
+286
View File
@@ -0,0 +1,286 @@
use crate::{
db::models::account::AccountToken,
error::{AppError, AppResult},
AppState,
};
use axum::{
extract::{Path, Query, State},
http::StatusCode,
Json,
};
use serde::Deserialize;
use serde_json::{json, Value};
use tracing::info;
#[derive(Debug, Deserialize)]
pub struct ListTokensQuery {
pub account_id: Option<i64>,
}
#[derive(Debug, Deserialize)]
pub struct CreateTokenRequest {
pub account_id: i64,
pub name: String,
pub token: String,
pub source: Option<String>,
pub token_group: Option<String>,
pub enabled: Option<bool>,
pub is_default: Option<bool>,
}
#[derive(Debug, Deserialize)]
pub struct UpdateTokenRequest {
pub name: Option<String>,
pub token: Option<String>,
pub source: Option<String>,
pub token_group: Option<String>,
pub enabled: Option<bool>,
pub is_default: Option<bool>,
}
#[derive(Debug, Deserialize)]
pub struct RefreshRouteRequest {
pub model_pattern: Option<String>,
pub route_id: Option<i64>,
}
/// GET /api/tokens?account_id=
pub async fn list_tokens(
State(state): State<AppState>,
Query(query): Query<ListTokensQuery>,
) -> AppResult<Json<Value>> {
info!("Listing tokens, account_id filter: {:?}", query.account_id);
let tokens: Vec<AccountToken> = if let Some(account_id) = query.account_id {
sqlx::query_as::<_, AccountToken>(
"SELECT * FROM account_tokens WHERE account_id = ? ORDER BY id",
)
.bind(account_id)
.fetch_all(&*state.db)
.await?
} else {
sqlx::query_as::<_, AccountToken>("SELECT * FROM account_tokens ORDER BY account_id, id")
.fetch_all(&*state.db)
.await?
};
Ok(Json(json!({ "data": tokens })))
}
/// POST /api/tokens
pub async fn create_token(
State(state): State<AppState>,
Json(payload): Json<CreateTokenRequest>,
) -> AppResult<(StatusCode, Json<Value>)> {
info!("Creating token for account_id: {}", payload.account_id);
// Validate
if payload.name.trim().is_empty() {
return Err(AppError::BadRequest("name is required".into()));
}
if payload.token.trim().is_empty() {
return Err(AppError::BadRequest("token is required".into()));
}
// Verify account exists
let account_exists: Option<(i64,)> = sqlx::query_as("SELECT id FROM accounts WHERE id = ?")
.bind(payload.account_id)
.fetch_optional(&*state.db)
.await?;
if account_exists.is_none() {
return Err(AppError::BadRequest(format!(
"Account with id {} not found",
payload.account_id
)));
}
let source = payload.source.unwrap_or_else(|| "manual".to_string());
let enabled = payload.enabled.unwrap_or(true);
let is_default = payload.is_default.unwrap_or(false);
let result = sqlx::query(
"INSERT INTO account_tokens (account_id, name, token, source, token_group, enabled, is_default) VALUES (?, ?, ?, ?, ?, ?, ?)"
)
.bind(payload.account_id)
.bind(&payload.name)
.bind(&payload.token)
.bind(&source)
.bind(&payload.token_group)
.bind(enabled)
.bind(is_default)
.execute(&*state.db)
.await?;
let token_id = result.last_insert_rowid();
let token: AccountToken =
sqlx::query_as::<_, AccountToken>("SELECT * FROM account_tokens WHERE id = ?")
.bind(token_id)
.fetch_one(&*state.db)
.await?;
Ok((StatusCode::CREATED, Json(json!({ "data": token }))))
}
/// GET /api/tokens/{id}
pub async fn get_token(
State(state): State<AppState>,
Path(id): Path<i64>,
) -> AppResult<Json<Value>> {
info!("Getting token: {}", id);
let token: AccountToken =
sqlx::query_as::<_, AccountToken>("SELECT * FROM account_tokens WHERE id = ?")
.bind(id)
.fetch_optional(&*state.db)
.await?
.ok_or_else(|| AppError::NotFound(format!("Token with id {} not found", id)))?;
Ok(Json(json!({ "data": token })))
}
/// PUT /api/tokens/{id}
pub async fn update_token(
State(state): State<AppState>,
Path(id): Path<i64>,
Json(payload): Json<UpdateTokenRequest>,
) -> AppResult<Json<Value>> {
info!("Updating token: {}", id);
let existing: AccountToken =
sqlx::query_as::<_, AccountToken>("SELECT * FROM account_tokens WHERE id = ?")
.bind(id)
.fetch_optional(&*state.db)
.await?
.ok_or_else(|| AppError::NotFound(format!("Token with id {} not found", id)))?;
sqlx::query(
"UPDATE account_tokens SET name = ?, token = ?, source = ?, token_group = ?, enabled = ?, is_default = ?, updated_at = datetime('now') WHERE id = ?"
)
.bind(payload.name.unwrap_or(existing.name))
.bind(payload.token.as_ref().unwrap_or(&existing.token))
.bind(payload.source.unwrap_or(existing.source.unwrap_or_else(|| "manual".to_string())))
.bind(payload.token_group.or(existing.token_group))
.bind(payload.enabled.unwrap_or(existing.enabled.unwrap_or(true)))
.bind(payload.is_default.unwrap_or(existing.is_default.unwrap_or(false)))
.bind(id)
.execute(&*state.db)
.await?;
let token: AccountToken =
sqlx::query_as::<_, AccountToken>("SELECT * FROM account_tokens WHERE id = ?")
.bind(id)
.fetch_one(&*state.db)
.await?;
Ok(Json(json!({ "data": token })))
}
/// DELETE /api/tokens/{id}
pub async fn delete_token(
State(state): State<AppState>,
Path(id): Path<i64>,
) -> AppResult<Json<Value>> {
info!("Deleting token: {}", id);
let affected = sqlx::query("DELETE FROM account_tokens WHERE id = ?")
.bind(id)
.execute(&*state.db)
.await?
.rows_affected();
if affected == 0 {
return Err(AppError::NotFound(format!(
"Token with id {} not found",
id
)));
}
Ok(Json(json!({ "data": { "deleted": true, "id": id } })))
}
/// POST /api/tokens/refresh-route
pub async fn refresh_route(
State(state): State<AppState>,
Json(payload): Json<RefreshRouteRequest>,
) -> AppResult<Json<Value>> {
info!("Refreshing route decision");
if let Some(route_id) = payload.route_id {
sqlx::query(
"UPDATE token_routes SET decision_refreshed_at = datetime('now'), updated_at = datetime('now') WHERE id = ?"
)
.bind(route_id)
.execute(&*state.db)
.await?;
let route: Option<crate::db::models::token_route::TokenRoute> =
sqlx::query_as("SELECT * FROM token_routes WHERE id = ?")
.bind(route_id)
.fetch_optional(&*state.db)
.await?;
return Ok(Json(
json!({ "data": { "refreshed": true, "route": route } }),
));
}
if let Some(ref model_pattern) = payload.model_pattern {
let result = sqlx::query(
"UPDATE token_routes SET decision_refreshed_at = datetime('now'), updated_at = datetime('now') WHERE model_pattern LIKE ?"
)
.bind(format!("%{}%", model_pattern))
.execute(&*state.db)
.await?;
return Ok(Json(json!({
"data": {
"refreshed": true,
"affected_routes": result.rows_affected()
}
})));
}
// Refresh all routes
let result = sqlx::query(
"UPDATE token_routes SET decision_refreshed_at = datetime('now'), updated_at = datetime('now')"
)
.execute(&*state.db)
.await?;
Ok(Json(json!({
"data": {
"refreshed": true,
"affected_routes": result.rows_affected()
}
})))
}
/// GET /api/tokens/route-decision
pub async fn route_decision(
State(state): State<AppState>,
Query(query): Query<RefreshRouteRequest>,
) -> AppResult<Json<Value>> {
info!("Getting route decisions");
let routes: Vec<crate::db::models::token_route::TokenRoute> =
if let Some(route_id) = query.route_id {
sqlx::query_as("SELECT * FROM token_routes WHERE id = ? AND enabled = 1")
.bind(route_id)
.fetch_optional(&*state.db)
.await?
.map(|r| vec![r])
.unwrap_or_default()
} else if let Some(ref model_pattern) = query.model_pattern {
sqlx::query_as("SELECT * FROM token_routes WHERE model_pattern LIKE ? AND enabled = 1")
.bind(format!("%{}%", model_pattern))
.fetch_all(&*state.db)
.await?
} else {
sqlx::query_as("SELECT * FROM token_routes WHERE enabled = 1 ORDER BY model_pattern")
.fetch_all(&*state.db)
.await?
};
Ok(Json(json!({ "data": routes })))
}
+315
View File
@@ -0,0 +1,315 @@
// ComfyUI proxy routes — forward to local ComfyUI instance
//
// These routes proxy requests to a local ComfyUI server for AI video/image workflow editing.
// All business routes require downstream API key or admin token authentication.
// The /api/comfyui/proxy/* route serves static assets (public).
use crate::{
auth::api_key::{ authorize_downstream_token, DownstreamAuthResult },
error::{ AppError, AppResult },
AppState,
};
use axum::{
extract::{ Path, Query, State },
http::{ HeaderMap, StatusCode },
response::{ Html, IntoResponse, Response },
routing::{ get, post },
Json, Router,
};
use reqwest::Client;
use serde_json::Value;
use std::collections::HashMap;
use tracing::info;
/// Default ComfyUI backend address.
const DEFAULT_COMFYUI_BASE: &str = "http://127.0.0.1:8188";
/// Shared HTTP client
fn http_client() -> Client {
Client::builder()
.timeout(std::time::Duration::from_secs(120))
.build()
.expect("Failed to build HTTP client")
}
// ── Auth helper (same pattern as proxy module) ────────────────────────
async fn check_auth(state: &AppState, headers: &HeaderMap) -> Result<(), AppError> {
let auth = headers
.get("Authorization")
.and_then(|v| v.to_str().ok())
.and_then(|v| v.strip_prefix("Bearer "))
.and_then(|v| v.trim().split_whitespace().next())
.ok_or_else(|| AppError::Unauthorized("Missing Authorization header".into()))?;
// Admin token passes
if auth == state.config.auth.auth_token {
return Ok(());
}
// Downstream API key
if auth.starts_with("sk-") {
let result = authorize_downstream_token(auth, &state.db).await
.map_err(|e| AppError::Unauthorized(e))?;
match result {
DownstreamAuthResult::Managed { key, .. } => {
if !key.enabled.unwrap_or(true) {
return Err(AppError::Forbidden("API key is disabled".into()));
}
Ok(())
}
DownstreamAuthResult::Global { .. } => Ok(()),
}
} else {
Err(AppError::Unauthorized("Invalid token".into()))
}
}
/// Get ComfyUI base URL from env or default
fn comfyui_base() -> String {
std::env::var("COMFYUI_BASE").unwrap_or_else(|_| DEFAULT_COMFYUI_BASE.to_string())
}
// ── Route builder ────────────────────────────────────────────────────
pub fn routes() -> Router<AppState> {
Router::new()
// Business routes (require auth)
.route("/api/comfyui/prompt", post(proxy_post_prompt))
.route("/api/comfyui/history", get(proxy_get_history))
.route("/api/comfyui/history/{id}", get(proxy_get_history_by_id))
.route("/api/comfyui/queue", get(proxy_get_queue))
.route("/api/comfyui/view", get(proxy_get_view))
// Static proxy (public — assets/JS/CSS)
.route("/api/comfyui/proxy/*path", get(proxy_static))
// HTML page (public)
.route("/comfyui", get(comfyui_page))
}
// ── Business route handlers (require auth) ───────────────────────────
/// POST /api/comfyui/prompt
async fn proxy_post_prompt(
State(state): State<AppState>,
headers: HeaderMap,
Json(body): Json<Value>,
) -> AppResult<Response> {
info!("POST /api/comfyui/prompt");
check_auth(&state, &headers).await?;
proxy_to_comfyui("POST", "/prompt", None, Some(body)).await
}
/// GET /api/comfyui/history
async fn proxy_get_history(
State(state): State<AppState>,
headers: HeaderMap,
) -> AppResult<Response> {
info!("GET /api/comfyui/history");
check_auth(&state, &headers).await?;
proxy_to_comfyui("GET", "/history", None, None).await
}
/// GET /api/comfyui/history/{id}
async fn proxy_get_history_by_id(
State(state): State<AppState>,
headers: HeaderMap,
Path(id): Path<String>,
) -> AppResult<Response> {
info!("GET /api/comfyui/history/{}", id);
check_auth(&state, &headers).await?;
let path = format!("/history/{}", urlencoding(&id));
proxy_to_comfyui("GET", &path, None, None).await
}
/// GET /api/comfyui/queue
async fn proxy_get_queue(
State(state): State<AppState>,
headers: HeaderMap,
) -> AppResult<Response> {
info!("GET /api/comfyui/queue");
check_auth(&state, &headers).await?;
proxy_to_comfyui("GET", "/queue", None, None).await
}
/// GET /api/comfyui/view?filename=...&type=...&subfolder=...
async fn proxy_get_view(
State(state): State<AppState>,
headers: HeaderMap,
Query(params): Query<HashMap<String, String>>,
) -> AppResult<Response> {
info!("GET /api/comfyui/view");
check_auth(&state, &headers).await?;
let qs = build_query_string(&params);
let path = format!("/view{}", qs);
proxy_to_comfyui_binary("GET", &path).await
}
// ── Static proxy (public — assets/JS/CSS for iframe) ─────────────────
/// GET /api/comfyui/proxy/{path} — proxy everything (HTML, assets)
async fn proxy_static(
Path(path): Path<String>,
) -> AppResult<Response> {
info!("GET /api/comfyui/proxy/{}", path);
let target_path = if path.is_empty() || path == "/" {
"/".to_string()
} else {
format!("/{}", path)
};
let client = http_client();
let url = format!("{}{}", comfyui_base(), target_path);
let resp = client.get(&url).send().await
.map_err(|e| AppError::UpstreamError(format!("ComfyUI unreachable: {e}")))?;
let status = resp.status();
let content_type = resp.headers()
.get("content-type")
.and_then(|v| v.to_str().ok())
.unwrap_or("application/octet-stream")
.to_string();
if content_type.contains("html") {
let text = resp.text().await
.map_err(|e| AppError::UpstreamError(e.to_string()))?;
// Rebrand: ComfyUI → BoosAPI
let rebranded = text
.replace("ComfyUI", "BoosAPI")
.replace("comfyui", "BoosAPI");
// Rewrite URLs
let rewritten = rewrite_proxy_urls(&rebranded);
Ok((StatusCode::from_u16(status.as_u16()).unwrap(), [(axum::http::header::CONTENT_TYPE, "text/html; charset=utf-8")], rewritten).into_response())
} else if content_type.contains("text") || content_type.contains("json") {
let text = resp.text().await
.map_err(|e| AppError::UpstreamError(e.to_string()))?;
Ok((StatusCode::from_u16(status.as_u16()).unwrap(), [(axum::http::header::CONTENT_TYPE, content_type)], text).into_response())
} else {
let bytes = resp.bytes().await
.map_err(|e| AppError::UpstreamError(e.to_string()))?;
Ok((StatusCode::from_u16(status.as_u16()).unwrap(), [(axum::http::header::CONTENT_TYPE, content_type)], bytes).into_response())
}
}
/// GET /comfyui — HTML page with iframe
async fn comfyui_page() -> Html<&'static str> {
Html(r##"<!DOCTYPE html>
<html lang="zh-CN">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>BoosAPI - ComfyUI</title>
<style>
* { margin: 0; padding: 0; box-sizing: border-box; }
html, body { width: 100%; height: 100%; overflow: hidden; background: #1a1a2e; }
iframe { width: 100%; height: 100%; border: none; }
</style>
</head>
<body>
<iframe src="/api/comfyui/proxy/" title="BoosAPI ComfyUI"></iframe>
</body>
</html>"##)
}
// ── Helpers ──────────────────────────────────────────────────────────
/// Rewrite URLs in HTML so asset paths go through the proxy
fn rewrite_proxy_urls(html: &str) -> String {
// Replace src="/...", href="/...", action="/..." with /api/comfyui/proxy/...
let mut result = String::with_capacity(html.len() * 12 / 10);
let mut chars = html.char_indices().peekable();
while let Some((_, ch)) = chars.next() {
if ch == 's' || ch == 'h' || ch == 'a' {
// Look ahead for rc="/", ref="/", ction="/"
let snippet: String = chars.clone().take(6).map(|(_, c)| c).collect();
let full = format!("{}{}", ch, snippet);
if full.starts_with("src=\"/") || full.starts_with("href=\"/") || full.starts_with("action=\"/") {
let attr_end = if full.starts_with("src=\"") { 4 } else if full.starts_with("href=\"") { 5 } else { 7 };
result.push_str(&full[..attr_end]);
result.push_str("/api/comfyui/proxy/");
// Skip the rest of the attribute
for _ in 0..snippet.len() { chars.next(); }
continue;
}
}
result.push(ch);
}
result
}
/// Build query string from params
fn build_query_string(params: &HashMap<String, String>) -> String {
if params.is_empty() {
return String::new();
}
let mut parts: Vec<String> = params
.iter()
.map(|(k, v)| format!("{}={}", urlencoding(k), urlencoding(v)))
.collect();
parts.sort(); // deterministic ordering
format!("?{}", parts.join("&"))
}
/// Simple URL encoding (handles the chars we need)
fn urlencoding(s: &str) -> String {
s.replace('%', "%25")
.replace(' ', "%20")
.replace('?', "%3F")
.replace('&', "%26")
.replace('#', "%23")
.replace('/', "%2F")
}
/// Proxy to ComfyUI — returns JSON/text response
async fn proxy_to_comfyui(
method: &str,
path: &str,
_query: Option<&HashMap<String, String>>,
body: Option<Value>,
) -> AppResult<Response> {
let client = http_client();
let url = format!("{}{}", comfyui_base(), path);
let req = match method {
"POST" => {
let body_str = body.map(|b| b.to_string()).unwrap_or_else(|| "{}".to_string());
client.post(&url)
.header("Content-Type", "application/json")
.body(body_str)
}
_ => client.get(&url),
};
let resp = req.send().await
.map_err(|e| AppError::UpstreamError(format!("ComfyUI unreachable: {e}")))?;
let status = resp.status();
let text = resp.text().await
.map_err(|e| AppError::UpstreamError(e.to_string()))?;
Ok((StatusCode::from_u16(status.as_u16()).unwrap(), [(axum::http::header::CONTENT_TYPE, "application/json")], text).into_response())
}
/// Proxy to ComfyUI — returns binary response (for images)
async fn proxy_to_comfyui_binary(_method: &str, path: &str) -> AppResult<Response> {
let client = http_client();
let url = format!("{}{}", comfyui_base(), path);
let req = client.get(&url);
let resp = req.send().await
.map_err(|e| AppError::UpstreamError(format!("ComfyUI unreachable: {e}")))?;
let status = resp.status();
let content_type = resp.headers()
.get("content-type")
.and_then(|v| v.to_str().ok())
.unwrap_or("application/octet-stream")
.to_string();
let bytes = resp.bytes().await
.map_err(|e| AppError::UpstreamError(e.to_string()))?;
Ok((StatusCode::from_u16(status.as_u16()).unwrap(), [(axum::http::header::CONTENT_TYPE, content_type)], bytes).into_response())
}
+14
View File
@@ -0,0 +1,14 @@
use crate::{error::AppResult, AppState};
use axum::{extract::State, Json};
use serde_json::{json, Value};
pub async fn health_check(State(state): State<AppState>) -> AppResult<Json<Value>> {
// Simple DB ping
sqlx::query("SELECT 1").execute(&*state.db).await?;
Ok(Json(json!({
"status": "ok",
"database": "ok",
"version": env!("CARGO_PKG_VERSION"),
})))
}
+16
View File
@@ -0,0 +1,16 @@
use crate::AppState;
use axum::{routing::get, Router};
mod admin;
mod comfyui;
mod health;
mod proxy;
pub fn build_router(state: AppState) -> Router {
Router::new()
.route("/healthz", get(health::health_check))
.merge(admin::routes(state.clone()))
.merge(proxy::routes())
.merge(comfyui::routes())
.with_state(state)
}
+190
View File
@@ -0,0 +1,190 @@
// Proxy routes — OpenAI-compatible /v1/ endpoints
//
// These routes handle:
// POST /v1/chat/completions — streaming and non-streaming
// POST /v1/images/generations — image generation
// GET /v1/models — list available models
use crate::{
auth::api_key::{ authorize_downstream_token, DownstreamAuthResult },
error::{ AppError, AppResult },
proxy::executors::{ self, UpstreamRoute },
AppState,
};
use axum::{
extract::State,
http::HeaderMap,
response::{ IntoResponse, Response, Sse },
routing::{ get, post },
Json, Router,
};
use futures::stream;
use futures::StreamExt;
use serde_json::{ json, Value };
use std::convert::Infallible;
use tracing::info;
// ── Auth types ────────────────────────────────────────────────────────
#[derive(Clone)]
pub struct ProxyAuth {
pub token: String,
pub key_id: Option<i64>,
pub source: DownstreamSource,
}
#[derive(Clone, PartialEq)]
pub enum DownstreamSource {
Managed { key_id: i64 },
Global,
}
// ── Route definitions ─────────────────────────────────────────────────
/// Build the /v1/ proxy router
pub fn routes() -> Router<AppState> {
Router::new()
.route("/v1/chat/completions", post(chat_completions))
.route("/v1/images/generations", post(images_generations))
.route("/v1/models", get(list_models))
}
// ── Handlers ───────────────────────────────────────────────────────────
/// POST /v1/chat/completions
async fn chat_completions(
State(state): State<AppState>,
headers: HeaderMap,
Json(body): Json<Value>,
) -> Result<Response, AppError> {
info!("POST /v1/chat/completions");
check_auth(&state, &headers).await?;
let model = body
.get("model")
.and_then(|m| m.as_str())
.ok_or_else(|| AppError::BadRequest("model is required".into()))?
.to_string();
let route = resolve_route(&state, &model).await?;
let stream = body.get("stream").and_then(|s| s.as_bool()).unwrap_or(false);
if stream {
let chunks = executors::execute_chat_completions_stream(state, route, body).await?;
let sse = Sse::new(stream::iter(chunks).map(|chunk| {
Ok::<_, Infallible>(axum::response::sse::Event::default().data(chunk))
}));
Ok(sse.into_response())
} else {
let result = executors::execute_chat_completions(state, route, body).await?;
Ok(Json(result).into_response())
}
}
/// POST /v1/images/generations
async fn images_generations(
State(state): State<AppState>,
headers: HeaderMap,
Json(body): Json<Value>,
) -> AppResult<Json<Value>> {
info!("POST /v1/images/generations");
check_auth(&state, &headers).await?;
let model = body
.get("model")
.and_then(|m| m.as_str())
.ok_or_else(|| AppError::BadRequest("model is required".into()))?
.to_string();
let route = resolve_route(&state, &model).await?;
let result = executors::execute_image_generation(state, route, body).await?;
Ok(Json(result))
}
/// GET /v1/models
async fn list_models(
State(state): State<AppState>,
headers: HeaderMap,
) -> AppResult<Json<Value>> {
info!("GET /v1/models");
check_auth(&state, &headers).await?;
let rows: Vec<(String, String)> = sqlx::query_as(
"SELECT tr.name, tr.model FROM token_routes tr WHERE tr.enabled = 1 ORDER BY tr.name"
)
.fetch_all(&*state.db)
.await?;
let models: Vec<Value> = rows
.into_iter()
.map(|(name, model)| {
json!({
"id": model,
"object": "model",
"created": 0,
"owned_by": name,
})
})
.collect();
Ok(Json(json!({
"object": "list",
"data": models,
})))
}
// ── Auth helper ───────────────────────────────────────────────────────
async fn check_auth(state: &AppState, headers: &HeaderMap) -> Result<(), AppError> {
let auth = headers
.get("Authorization")
.and_then(|v| v.to_str().ok())
.and_then(|v| v.strip_prefix("Bearer "))
.and_then(|v| v.trim().split_whitespace().next())
.ok_or_else(|| AppError::Unauthorized("Missing Authorization header".into()))?;
let result = authorize_downstream_token(auth, &state.db).await
.map_err(|e| AppError::Unauthorized(e))?;
match result {
DownstreamAuthResult::Managed { key, .. } => {
if !key.enabled.unwrap_or(true) {
return Err(AppError::Forbidden("API key is disabled".into()));
}
Ok(())
}
DownstreamAuthResult::Global { .. } => Ok(()),
}
}
// ── Route Resolution ──────────────────────────────────────────────────
/// Resolve upstream endpoint + API key for a given model.
async fn resolve_route(state: &AppState, model: &str) -> AppResult<UpstreamRoute> {
let row: Option<(String, String)> = sqlx::query_as(
"SELECT s.base_url, COALESCE(a.api_key, s.default_key, '') as api_key
FROM token_routes tr
JOIN sites s ON s.id = tr.site_id
LEFT JOIN accounts a ON a.id = tr.account_id AND a.site_id = s.id
WHERE tr.enabled = 1 AND tr.model = ?
LIMIT 1"
)
.bind(model)
.fetch_optional(&*state.db)
.await?;
match row {
Some((base_url, api_key)) => Ok(UpstreamRoute {
base_url,
api_key,
model: model.to_string(),
}),
None => Err(AppError::BadRequest(format!(
"No route configured for model '{}'",
model
))),
}
}
+1
View File
@@ -0,0 +1 @@
// Routes - user