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
+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
))),
}
}