boos/BoosAPI
2
0
Fork 2
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: refine BoosAPI console and user workspace

Browse Source
This commit is contained in:
boos
2026-06-09 16:48:18 +08:00
parent aaba6f152e
commit 59106e6a96
284 changed files with 59886 additions and 5594 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/services/billing/mod.rs
+358
View File
@@ -4,6 +4,7 @@ use crate::{
};
use chrono::Timelike;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use sqlx::{FromRow, Row, SqlitePool};
use std::sync::Arc;
@@ -22,6 +23,23 @@ pub struct UsageSummary {
pub total_cost: f64,
}
#[derive(Debug, Clone)]
pub struct ComfyTaskBillingEvent {
pub key_id: Option<i64>,
pub model: String,
pub task_type: &'static str,
pub client_family: &'static str,
pub client_app_id: &'static str,
pub client_app_name: &'static str,
pub prompt_tokens: i64,
pub completion_tokens: i64,
pub cost: f64,
pub content_summary: Option<String>,
pub request_content: Option<String>,
pub response_content: Option<String>,
pub error_message: Option<String>,
}
impl UsageTracker {
pub fn new(db: Arc<SqlitePool>) -> Self {
Self { db }
@@ -119,6 +137,112 @@ impl UsageTracker {
Ok(())
}
pub async fn record_comfy_task_success(
&self,
event: ComfyTaskBillingEvent,
) -> AppResult<Value> {
Self::ensure_schema(&self.db).await?;
Self::ensure_proxy_log_task_columns(&self.db).await?;
let prompt_tokens = event.prompt_tokens.max(0);
let completion_tokens = event.completion_tokens.max(0);
let total_tokens = prompt_tokens + completion_tokens;
let cost = event.cost.max(0.0);
let billing_details = comfy_task_billing_details(&event, "success", cost);
let mut tx = self.db.begin().await?;
if let Some(key_id) = event.key_id {
sqlx::query(
"INSERT INTO usage_logs
(downstream_key_id, model, prompt_tokens, completion_tokens, total_tokens, cost)
VALUES (?, ?, ?, ?, ?, ?)",
)
.bind(key_id)
.bind(&event.model)
.bind(prompt_tokens)
.bind(completion_tokens)
.bind(total_tokens)
.bind(cost)
.execute(&mut *tx)
.await?;
sqlx::query(
"UPDATE downstream_api_keys
SET used_requests = COALESCE(used_requests, 0) + 1,
used_cost = COALESCE(used_cost, 0) + ?,
last_used_at = datetime('now'),
updated_at = datetime('now')
WHERE id = ?",
)
.bind(cost)
.bind(key_id)
.execute(&mut *tx)
.await?;
}
sqlx::query(
"INSERT INTO proxy_logs (
downstream_api_key_id, model_requested, model_actual, status, http_status,
prompt_tokens, completion_tokens, total_tokens, estimated_cost,
retry_count, is_stream, client_family, client_app_id, client_app_name,
client_confidence, billing_details, content_summary, request_content,
response_content
) VALUES (?, ?, ?, 'success', 200, ?, ?, ?, ?, 0, 0, ?, ?, ?, 'high', ?, ?, ?, ?)",
)
.bind(event.key_id)
.bind(&event.model)
.bind(&event.model)
.bind(prompt_tokens)
.bind(completion_tokens)
.bind(total_tokens)
.bind(cost)
.bind(event.client_family)
.bind(event.client_app_id)
.bind(event.client_app_name)
.bind(billing_details.to_string())
.bind(event.content_summary.as_deref())
.bind(event.request_content.as_deref())
.bind(event.response_content.as_deref())
.execute(&mut *tx)
.await?;
tx.commit().await?;
Ok(billing_details)
}
pub async fn record_comfy_task_failure(
&self,
event: ComfyTaskBillingEvent,
) -> AppResult<Value> {
Self::ensure_schema(&self.db).await?;
Self::ensure_proxy_log_task_columns(&self.db).await?;
let billing_details = comfy_task_billing_details(&event, "failed", 0.0);
sqlx::query(
"INSERT INTO proxy_logs (
downstream_api_key_id, model_requested, model_actual, status, http_status,
prompt_tokens, completion_tokens, total_tokens, estimated_cost, error_message,
retry_count, is_stream, client_family, client_app_id, client_app_name,
client_confidence, billing_details, content_summary, request_content,
response_content
) VALUES (?, ?, ?, 'failed', 502, 0, 0, 0, 0.0, ?, 0, 0, ?, ?, ?, 'high', ?, ?, ?, ?)",
)
.bind(event.key_id)
.bind(&event.model)
.bind(&event.model)
.bind(event.error_message.as_deref().unwrap_or("Comfy task failed"))
.bind(event.client_family)
.bind(event.client_app_id)
.bind(event.client_app_name)
.bind(billing_details.to_string())
.bind(event.content_summary.as_deref())
.bind(event.request_content.as_deref())
.bind(event.response_content.as_deref())
.execute(&*self.db)
.await?;
Ok(billing_details)
}
pub async fn create_daily_usage_snapshot(&self) -> AppResult<()> {
Self::create_daily_usage_snapshot_for_db(&self.db).await
}
@@ -485,6 +609,112 @@ impl UsageTracker {
.await?;
Ok(())
}
async fn ensure_proxy_log_task_columns(db: &SqlitePool) -> AppResult<()> {
for (name, definition) in [
("client_family", "TEXT"),
("client_app_id", "TEXT"),
("client_app_name", "TEXT"),
("client_confidence", "TEXT"),
("is_stream", "INTEGER"),
("first_byte_latency_ms", "INTEGER"),
("billing_details", "TEXT"),
("content_summary", "TEXT"),
("request_content", "TEXT"),
("response_content", "TEXT"),
] {
ensure_table_column(db, "proxy_logs", name, definition).await?;
}
Ok(())
}
}
async fn ensure_table_column(
db: &SqlitePool,
table: &str,
name: &str,
definition: &str,
) -> AppResult<()> {
let info_sql = format!("PRAGMA table_info({table})");
let rows = sqlx::query(&info_sql).fetch_all(db).await?;
if rows.is_empty() {
return Ok(());
}
let exists = rows
.iter()
.filter_map(|row| row.try_get::<String, _>("name").ok())
.any(|column| column == name);
if exists {
return Ok(());
}
let alter_sql = format!("ALTER TABLE {table} ADD COLUMN {name} {definition}");
sqlx::query(&alter_sql).execute(db).await?;
Ok(())
}
fn comfy_task_billing_details(
event: &ComfyTaskBillingEvent,
status: &str,
total_cost: f64,
) -> Value {
let prompt_tokens = event.prompt_tokens.max(0);
let completion_tokens = event.completion_tokens.max(0);
let total_tokens = prompt_tokens + completion_tokens;
json!({
"source": "comfy-native-task",
"taskType": event.task_type,
"status": status,
"currency": "USD",
"charged": status == "success" && event.key_id.is_some(),
"totalCost": total_cost,
"usage": {
"promptTokens": prompt_tokens,
"completionTokens": completion_tokens,
"totalTokens": total_tokens,
"cacheReadTokens": 0,
"cacheCreationTokens": 0,
"billablePromptTokens": prompt_tokens,
"promptTokensIncludeCache": false,
},
"pricing": {
"modelRatio": 1,
"completionRatio": 1,
"cacheRatio": 0,
"cacheCreationRatio": 0,
"groupRatio": 1,
},
"breakdown": {
"inputPerMillion": 0,
"outputPerMillion": 0,
"cacheReadPerMillion": 0,
"cacheCreationPerMillion": 0,
"inputCost": 0,
"outputCost": total_cost,
"cacheReadCost": 0,
"cacheCreationCost": 0,
"totalCost": total_cost,
},
"steps": [
{
"name": "submit_workflow",
"status": status,
"model": event.model,
"cost": 0.0,
},
{
"name": "generate_outputs",
"status": status,
"model": event.model,
"cost": total_cost,
},
{
"name": "archive_outputs",
"status": status,
"model": event.model,
"cost": 0.0,
}
]
})
}
async fn touch_projection_checkpoint(
@@ -614,6 +844,56 @@ mod tests {
pool
}
async fn create_proxy_logs(pool: &SqlitePool) {
sqlx::query(
"CREATE TABLE proxy_logs (
id INTEGER PRIMARY KEY,
downstream_api_key_id INTEGER,
model_requested TEXT,
model_actual TEXT,
status TEXT,
http_status INTEGER,
prompt_tokens INTEGER,
completion_tokens INTEGER,
total_tokens INTEGER,
estimated_cost REAL,
error_message TEXT,
retry_count INTEGER DEFAULT 0,
is_stream INTEGER,
client_family TEXT,
client_app_id TEXT,
client_app_name TEXT,
client_confidence TEXT,
billing_details TEXT,
content_summary TEXT,
request_content TEXT,
response_content TEXT,
created_at TEXT DEFAULT (datetime('now'))
)",
)
.execute(pool)
.await
.unwrap();
}
fn test_comfy_event(cost: f64) -> ComfyTaskBillingEvent {
ComfyTaskBillingEvent {
key_id: Some(1),
model: "seedance-comfy".to_string(),
task_type: "money-printer-video",
client_family: "money-printer",
client_app_id: "money-printer",
client_app_name: "Money Printer",
prompt_tokens: 0,
completion_tokens: 0,
cost,
content_summary: Some("topic".to_string()),
request_content: Some(json!({ "topic": "topic" }).to_string()),
response_content: Some(json!({ "ok": true }).to_string()),
error_message: None,
}
}
#[tokio::test]
async fn records_and_aggregates_usage() {
let tracker = UsageTracker::new(Arc::new(pool().await));
@@ -662,6 +942,84 @@ mod tests {
assert_eq!(count, 1);
}
#[tokio::test]
async fn records_comfy_task_success_as_usage_and_proxy_log() {
let pool = pool().await;
create_proxy_logs(&pool).await;
let tracker = UsageTracker::new(Arc::new(pool));
let billing = tracker
.record_comfy_task_success(test_comfy_event(0.35))
.await
.unwrap();
assert_eq!(billing["source"], "comfy-native-task");
assert_eq!(billing["charged"], true);
assert_eq!(billing["totalCost"], 0.35);
let usage: (i64, f64) = sqlx::query_as(
"SELECT COUNT(*), COALESCE(SUM(cost), 0.0) FROM usage_logs WHERE downstream_key_id = 1",
)
.fetch_one(&*tracker.db)
.await
.unwrap();
assert_eq!(usage, (1, 0.35));
let key_usage: (i64, f64) =
sqlx::query_as("SELECT used_requests, used_cost FROM downstream_api_keys WHERE id = 1")
.fetch_one(&*tracker.db)
.await
.unwrap();
assert_eq!(key_usage, (1, 0.35));
let proxy_log: (String, f64, String, String) = sqlx::query_as(
"SELECT status, estimated_cost, client_app_id, billing_details FROM proxy_logs WHERE downstream_api_key_id = 1",
)
.fetch_one(&*tracker.db)
.await
.unwrap();
assert_eq!(proxy_log.0, "success");
assert_eq!(proxy_log.1, 0.35);
assert_eq!(proxy_log.2, "money-printer");
assert!(proxy_log.3.contains("\"steps\""));
}
#[tokio::test]
async fn records_comfy_task_failure_without_debiting_usage() {
let pool = pool().await;
create_proxy_logs(&pool).await;
let tracker = UsageTracker::new(Arc::new(pool));
let mut event = test_comfy_event(0.35);
event.error_message = Some("Not found".to_string());
let billing = tracker.record_comfy_task_failure(event).await.unwrap();
assert_eq!(billing["charged"], false);
assert_eq!(billing["totalCost"], 0.0);
let usage_count: i64 =
sqlx::query_scalar("SELECT COUNT(*) FROM usage_logs WHERE downstream_key_id = 1")
.fetch_one(&*tracker.db)
.await
.unwrap();
assert_eq!(usage_count, 0);
let key_usage: (i64, f64) =
sqlx::query_as("SELECT used_requests, used_cost FROM downstream_api_keys WHERE id = 1")
.fetch_one(&*tracker.db)
.await
.unwrap();
assert_eq!(key_usage, (0, 0.0));
let proxy_log: (String, f64, String) = sqlx::query_as(
"SELECT status, estimated_cost, error_message FROM proxy_logs WHERE downstream_api_key_id = 1",
)
.fetch_one(&*tracker.db)
.await
.unwrap();
assert_eq!(proxy_log, ("failed".to_string(), 0.0, "Not found".to_string()));
}
#[tokio::test]
async fn projects_proxy_logs_to_site_and_model_usage() {
let pool = pool().await;
src/services/oauth/mod.rs
+523 -50
View File
@@ -1,13 +1,16 @@
use crate::error::{AppError, AppResult};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use serde_json::{json, Map, Value};
use sqlx::{FromRow, SqlitePool};
use std::sync::Arc;
pub const SYSTEM_PROXY_URL_SETTING_KEY: &str = "runtime.system_proxy_url";
#[derive(Clone)]
pub struct OAuthService {
db: Arc<SqlitePool>,
client: reqwest::Client,
#[cfg(test)]
client_override: Option<reqwest::Client>,
}
#[derive(Debug, Clone, Serialize, Deserialize, FromRow)]
@@ -26,6 +29,7 @@ pub struct OAuthRouteUnit {
pub struct RefreshedToken {
pub account_id: i64,
pub access_token: String,
pub refresh_token: Option<String>,
pub expires_at: Option<String>,
}
@@ -33,7 +37,16 @@ impl OAuthService {
pub fn new(db: Arc<SqlitePool>) -> Self {
Self {
db,
client: reqwest::Client::new(),
#[cfg(test)]
client_override: None,
}
}
#[cfg(test)]
fn with_client(db: Arc<SqlitePool>, client: reqwest::Client) -> Self {
Self {
db,
client_override: Some(client),
}
}
@@ -55,6 +68,7 @@ impl OAuthService {
return Ok(RefreshedToken {
account_id,
access_token: current_access_token,
refresh_token,
expires_at: token_expiry(extra_config.as_deref()),
});
}
@@ -63,21 +77,24 @@ impl OAuthService {
.filter(|token| !token.trim().is_empty())
.ok_or_else(|| AppError::BadRequest("OAuth account has no refresh token".into()))?;
let mut refreshed = self.perform_refresh(site_id, &refresh_token).await?;
let (proxy_url, use_system_proxy) = oauth_proxy_settings(extra_config.as_deref());
let client = self
.configured_http_client(proxy_url.as_deref(), use_system_proxy)
.await?;
let mut refreshed = self
.perform_refresh(&client, site_id, &refresh_token)
.await?;
refreshed.account_id = account_id;
let next_extra_config =
merge_refreshed_extra_config(extra_config.as_deref(), &refreshed, &refresh_token)?;
sqlx::query(
"UPDATE accounts
SET access_token = ?, extra_config = ?, updated_at = datetime('now')
SET access_token = ?, api_token = ?, extra_config = ?, updated_at = datetime('now')
WHERE id = ? AND site_id = ?",
)
.bind(&refreshed.access_token)
.bind(
serde_json::to_string(&serde_json::json!({
"expires_at": refreshed.expires_at,
"refreshed_at": chrono::Utc::now().to_rfc3339(),
}))
.map_err(|err| AppError::Internal(err.into()))?,
)
.bind(refreshed.refresh_token.as_deref().unwrap_or(&refresh_token))
.bind(next_extra_config)
.bind(account_id)
.bind(site_id)
.execute(&*self.db)
@@ -90,6 +107,19 @@ impl OAuthService {
token_expired(extra_config)
}
async fn configured_http_client(
&self,
proxy_url: Option<&str>,
use_system_proxy: bool,
) -> AppResult<reqwest::Client> {
#[cfg(test)]
if let Some(client) = &self.client_override {
return Ok(client.clone());
}
oauth_http_client(&*self.db, proxy_url, use_system_proxy).await
}
pub async fn list_route_units(&self, site_id: Option<i64>) -> AppResult<Vec<OAuthRouteUnit>> {
if let Some(site_id) = site_id {
sqlx::query_as("SELECT * FROM oauth_route_units WHERE site_id = ? ORDER BY id")
@@ -132,29 +162,37 @@ impl OAuthService {
async fn perform_refresh(
&self,
client: &reqwest::Client,
site_id: i64,
refresh_token: &str,
) -> AppResult<RefreshedToken> {
let token_endpoint: Option<String> = sqlx::query_scalar(
"SELECT json_extract(custom_headers, '$.oauth_token_url') FROM sites WHERE id = ?",
)
.bind(site_id)
.fetch_optional(&*self.db)
.await?;
let raw_custom_headers: Option<String> =
sqlx::query_scalar("SELECT custom_headers FROM sites WHERE id = ?")
.bind(site_id)
.fetch_optional(&*self.db)
.await?;
let Some(token_endpoint) = token_endpoint.filter(|url| !url.trim().is_empty()) else {
let token_config = oauth_token_config(raw_custom_headers.as_deref())?;
let Some(token_endpoint) = token_config.token_url.filter(|url| !url.trim().is_empty())
else {
return Err(AppError::BadRequest(
"OAuth token endpoint is not configured for this site".into(),
));
};
let mut form = vec![
("grant_type".to_string(), "refresh_token".to_string()),
("refresh_token".to_string(), refresh_token.to_string()),
];
if let Some(client_id) = token_config.client_id {
form.push(("client_id".to_string(), client_id));
}
if let Some(client_secret) = token_config.client_secret {
form.push(("client_secret".to_string(), client_secret));
}
let response: Value = self
.client
let response: Value = client
.post(token_endpoint)
.form(&[
("grant_type", "refresh_token"),
("refresh_token", refresh_token),
])
.form(&form)
.send()
.await
.map_err(|err| AppError::UpstreamError(format!("OAuth refresh failed: {err}")))?
@@ -162,46 +200,481 @@ impl OAuthService {
.await
.map_err(|err| AppError::UpstreamError(format!("Invalid OAuth response: {err}")))?;
let access_token = response
.get("access_token")
.and_then(Value::as_str)
let access_token = string_field(&response, &["access_token", "accessToken"])
.ok_or_else(|| AppError::UpstreamError("OAuth response missing access_token".into()))?
.to_string();
let expires_at = response
.get("expires_at")
.and_then(Value::as_str)
.map(str::to_string);
let refresh_token =
string_field(&response, &["refresh_token", "refreshToken"]).map(str::to_string);
let expires_at = token_expiry_from_response(&response);
Ok(RefreshedToken {
account_id: 0,
access_token,
refresh_token,
expires_at,
})
}
}
pub async fn system_proxy_url(db: &SqlitePool) -> AppResult<Option<String>> {
let raw: Option<String> = sqlx::query_scalar("SELECT value FROM settings WHERE key = ?")
.bind(SYSTEM_PROXY_URL_SETTING_KEY)
.fetch_optional(db)
.await?;
Ok(raw
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty()))
}
pub async fn oauth_http_client(
db: &SqlitePool,
proxy_url: Option<&str>,
use_system_proxy: bool,
) -> AppResult<reqwest::Client> {
let resolved_proxy = proxy_url
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
let resolved_proxy = match resolved_proxy {
Some(proxy_url) => Some(proxy_url),
None if use_system_proxy => system_proxy_url(db).await?,
None => None,
};
let mut builder = reqwest::Client::builder();
if let Some(proxy_url) = resolved_proxy {
let proxy = reqwest::Proxy::all(&proxy_url)
.map_err(|err| AppError::BadRequest(format!("Invalid OAuth proxy URL: {err}")))?;
builder = builder.proxy(proxy);
}
builder
.build()
.map_err(|err| AppError::Internal(err.into()))
}
#[derive(Debug, Clone, Default)]
struct OAuthTokenConfig {
token_url: Option<String>,
client_id: Option<String>,
client_secret: Option<String>,
}
fn string_field<'a>(value: &'a Value, keys: &[&str]) -> Option<&'a str> {
keys.iter()
.find_map(|key| value.get(*key))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
}
fn oauth_token_config(raw_custom_headers: Option<&str>) -> AppResult<OAuthTokenConfig> {
let Some(raw) = raw_custom_headers
.map(str::trim)
.filter(|value| !value.is_empty())
else {
return Ok(OAuthTokenConfig::default());
};
let value = serde_json::from_str::<Value>(raw)
.map_err(|err| AppError::BadRequest(format!("Invalid site custom_headers JSON: {err}")))?;
let object = value
.as_object()
.ok_or_else(|| AppError::BadRequest("site custom_headers must be a JSON object".into()))?;
Ok(OAuthTokenConfig {
token_url: [
"oauth_token_url",
"oauth_token_endpoint",
"token_url",
"tokenUrl",
"tokenEndpoint",
]
.iter()
.find_map(|key| object.get(*key))
.and_then(|value| value.as_str())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
client_id: ["oauth_client_id", "client_id", "clientId", "oauthClientId"]
.iter()
.find_map(|key| object.get(*key))
.and_then(|value| value.as_str())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
client_secret: [
"oauth_client_secret",
"client_secret",
"clientSecret",
"oauthClientSecret",
]
.iter()
.find_map(|key| object.get(*key))
.and_then(|value| value.as_str())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
})
}
fn token_expiry_from_response(response: &Value) -> Option<String> {
for key in [
"expires_at",
"expiresAt",
"tokenExpiresAt",
"token_expires_at",
] {
if let Some(expires_at) = response.get(key).and_then(expiry_datetime_from_value) {
return Some(expires_at.to_rfc3339());
}
}
let expires_in = response
.get("expires_in")
.or_else(|| response.get("expiresIn"))
.and_then(|value| {
value
.as_i64()
.or_else(|| value.as_str()?.trim().parse::<i64>().ok())
})
.filter(|value| *value > 0)?;
Some((chrono::Utc::now() + chrono::Duration::seconds(expires_in)).to_rfc3339())
}
fn expiry_datetime_from_value(value: &Value) -> Option<chrono::DateTime<chrono::Utc>> {
if let Some(number) = value.as_i64() {
return expiry_datetime_from_unixish(number);
}
let text = value.as_str()?.trim();
if text.is_empty() {
return None;
}
if let Ok(parsed) = chrono::DateTime::parse_from_rfc3339(text) {
return Some(parsed.with_timezone(&chrono::Utc));
}
text.parse::<i64>()
.ok()
.and_then(expiry_datetime_from_unixish)
}
fn expiry_datetime_from_unixish(value: i64) -> Option<chrono::DateTime<chrono::Utc>> {
if value <= 0 {
return None;
}
let millis = if value < 10_000_000_000 {
value.checked_mul(1000)?
} else {
value
};
chrono::DateTime::<chrono::Utc>::from_timestamp_millis(millis)
}
fn parse_extra_config_object(raw_extra_config: Option<&str>) -> AppResult<Map<String, Value>> {
let Some(raw) = raw_extra_config
.map(str::trim)
.filter(|value| !value.is_empty())
else {
return Ok(Map::new());
};
let value = serde_json::from_str::<Value>(raw)
.map_err(|err| AppError::BadRequest(format!("Invalid extra_config JSON: {err}")))?;
value
.as_object()
.cloned()
.ok_or_else(|| AppError::BadRequest("extra_config must be a JSON object".into()))
}
fn oauth_proxy_settings(raw_extra_config: Option<&str>) -> (Option<String>, bool) {
let Ok(config) = parse_extra_config_object(raw_extra_config) else {
return (None, false);
};
let proxy_url = config
.get("proxyUrl")
.or_else(|| config.get("proxy_url"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
let use_system_proxy = config
.get("useSystemProxy")
.and_then(Value::as_bool)
.or_else(|| config.get("use_system_proxy").and_then(Value::as_bool))
.unwrap_or(false);
(proxy_url, use_system_proxy)
}
fn merge_refreshed_extra_config(
raw_extra_config: Option<&str>,
refreshed: &RefreshedToken,
previous_refresh_token: &str,
) -> AppResult<Option<String>> {
let mut config = parse_extra_config_object(raw_extra_config)?;
config.insert(
"refreshed_at".to_string(),
json!(chrono::Utc::now().to_rfc3339()),
);
if let Some(expires_at) = refreshed.expires_at.as_deref() {
config.insert("expires_at".to_string(), json!(expires_at));
if let Ok(parsed) = chrono::DateTime::parse_from_rfc3339(expires_at) {
config.insert(
"tokenExpiresAt".to_string(),
json!(parsed.timestamp_millis()),
);
}
}
config.insert(
"refreshToken".to_string(),
json!(refreshed
.refresh_token
.as_deref()
.unwrap_or(previous_refresh_token)),
);
if config.is_empty() {
Ok(None)
} else {
serde_json::to_string(&Value::Object(config))
.map(Some)
.map_err(|err| AppError::Internal(err.into()))
}
}
fn token_expired(extra_config: Option<&str>) -> bool {
let Some(extra_config) = extra_config else {
return true;
};
let Ok(value) = serde_json::from_str::<Value>(extra_config) else {
return true;
};
let Some(expires_at) = value.get("expires_at").and_then(Value::as_str) else {
return true;
};
chrono::DateTime::parse_from_rfc3339(expires_at)
token_expiry_datetime(extra_config)
.map(|expires_at| expires_at <= chrono::Utc::now() + chrono::Duration::minutes(5))
.unwrap_or(true)
}
fn token_expiry(extra_config: Option<&str>) -> Option<String> {
extra_config
.and_then(|config| serde_json::from_str::<Value>(config).ok())
.and_then(|value| {
value
.get("expires_at")
.and_then(Value::as_str)
.map(str::to_string)
})
token_expiry_datetime(extra_config).map(|expires_at| expires_at.to_rfc3339())
}
fn token_expiry_datetime(extra_config: Option<&str>) -> Option<chrono::DateTime<chrono::Utc>> {
let value = extra_config.and_then(|config| serde_json::from_str::<Value>(config).ok())?;
if let Some(expires_at) = value
.get("expires_at")
.or_else(|| value.get("expiresAt"))
.and_then(Value::as_str)
{
return chrono::DateTime::parse_from_rfc3339(expires_at)
.map(|value| value.with_timezone(&chrono::Utc))
.ok();
}
let millis = value
.get("tokenExpiresAt")
.or_else(|| value.get("token_expires_at"))
.and_then(expiry_datetime_from_value)?;
Some(millis)
}
#[cfg(test)]
mod tests {
use super::*;
use sqlx::sqlite::SqlitePoolOptions;
use wiremock::matchers::{body_string_contains, method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
async fn test_pool() -> SqlitePool {
let pool = SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.expect("sqlite pool");
sqlx::query(
"CREATE TABLE settings (
key TEXT PRIMARY KEY NOT NULL,
value TEXT
)",
)
.execute(&pool)
.await
.expect("create settings");
sqlx::query(
"CREATE TABLE sites (
id INTEGER PRIMARY KEY,
custom_headers TEXT
)",
)
.execute(&pool)
.await
.expect("create sites");
sqlx::query(
"CREATE TABLE accounts (
id INTEGER PRIMARY KEY,
site_id INTEGER NOT NULL,
access_token TEXT NOT NULL,
api_token TEXT,
extra_config TEXT,
oauth_provider TEXT,
updated_at TEXT
)",
)
.execute(&pool)
.await
.expect("create accounts");
pool
}
#[tokio::test]
async fn refresh_token_preserves_extra_config_and_updates_refresh_token() {
let pool = test_pool().await;
let mock = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/oauth/token"))
.and(body_string_contains("grant_type=refresh_token"))
.and(body_string_contains("refresh_token=refresh-old"))
.and(body_string_contains("client_id=client-a"))
.and(body_string_contains("client_secret=secret-a"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"access_token": "access-new",
"refresh_token": "refresh-new",
"expires_at": 1893456000_i64
})))
.mount(&mock)
.await;
sqlx::query("INSERT INTO sites (id, custom_headers) VALUES (1, ?)")
.bind(
json!({
"oauth_token_url": format!("{}/oauth/token", mock.uri()),
"oauth_client_id": "client-a",
"oauth_client_secret": "secret-a"
})
.to_string(),
)
.execute(&pool)
.await
.expect("insert site");
sqlx::query(
"INSERT INTO accounts (id, site_id, access_token, api_token, extra_config, oauth_provider)
VALUES (7, 1, 'access-old', 'refresh-old', ?, 'codex')",
)
.bind(
json!({
"credentialMode": "oauth",
"proxyUrl": "socks5://127.0.0.1:1080",
"expires_at": "2000-01-01T00:00:00Z"
})
.to_string(),
)
.execute(&pool)
.await
.expect("insert account");
let service = OAuthService::with_client(Arc::new(pool.clone()), reqwest::Client::new());
let refreshed = service.refresh_token(7, 1).await.expect("refresh token");
assert_eq!(refreshed.access_token, "access-new");
assert_eq!(refreshed.refresh_token.as_deref(), Some("refresh-new"));
assert!(refreshed.expires_at.is_some());
let row: (String, Option<String>, Option<String>) = sqlx::query_as(
"SELECT access_token, api_token, extra_config FROM accounts WHERE id = 7",
)
.fetch_one(&pool)
.await
.expect("account row");
assert_eq!(row.0, "access-new");
assert_eq!(row.1.as_deref(), Some("refresh-new"));
let extra: Value =
serde_json::from_str(row.2.as_deref().expect("extra config")).expect("extra json");
assert_eq!(
extra.get("credentialMode").and_then(Value::as_str),
Some("oauth")
);
assert_eq!(
extra.get("proxyUrl").and_then(Value::as_str),
Some("socks5://127.0.0.1:1080")
);
assert_eq!(
extra.get("refreshToken").and_then(Value::as_str),
Some("refresh-new")
);
assert_eq!(
extra.get("expires_at").and_then(Value::as_str),
Some("2030-01-01T00:00:00+00:00")
);
assert_eq!(
extra.get("tokenExpiresAt").and_then(Value::as_i64),
Some(1893456000000_i64)
);
}
#[test]
fn token_expiry_from_response_accepts_string_token_expires_at() {
assert_eq!(
token_expiry_from_response(&json!({ "tokenExpiresAt": "1893456000" })).as_deref(),
Some("2030-01-01T00:00:00+00:00")
);
assert_eq!(
token_expiry_from_response(&json!({ "token_expires_at": 1893456000000_i64 }))
.as_deref(),
Some("2030-01-01T00:00:00+00:00")
);
}
#[tokio::test]
async fn refresh_token_keeps_current_token_when_not_expired() {
let pool = test_pool().await;
sqlx::query("INSERT INTO sites (id, custom_headers) VALUES (1, '{}')")
.execute(&pool)
.await
.expect("insert site");
let future = (chrono::Utc::now() + chrono::Duration::hours(1)).to_rfc3339();
sqlx::query(
"INSERT INTO accounts (id, site_id, access_token, api_token, extra_config, oauth_provider)
VALUES (7, 1, 'access-current', 'refresh-current', ?, 'codex')",
)
.bind(json!({ "expires_at": future }).to_string())
.execute(&pool)
.await
.expect("insert account");
let service = OAuthService::with_client(Arc::new(pool), reqwest::Client::new());
let refreshed = service.refresh_token(7, 1).await.expect("current token");
assert_eq!(refreshed.access_token, "access-current");
assert_eq!(refreshed.refresh_token.as_deref(), Some("refresh-current"));
}
#[tokio::test]
async fn refresh_token_keeps_current_token_when_token_expires_at_is_future() {
let pool = test_pool().await;
sqlx::query("INSERT INTO sites (id, custom_headers) VALUES (1, '{}')")
.execute(&pool)
.await
.expect("insert site");
let future_millis = (chrono::Utc::now() + chrono::Duration::hours(1)).timestamp_millis();
sqlx::query(
"INSERT INTO accounts (id, site_id, access_token, api_token, extra_config, oauth_provider)
VALUES (7, 1, 'access-current', 'refresh-current', ?, 'codex')",
)
.bind(json!({ "tokenExpiresAt": future_millis }).to_string())
.execute(&pool)
.await
.expect("insert account");
let service = OAuthService::with_client(Arc::new(pool), reqwest::Client::new());
let refreshed = service.refresh_token(7, 1).await.expect("current token");
assert_eq!(refreshed.access_token, "access-current");
assert_eq!(refreshed.refresh_token.as_deref(), Some("refresh-current"));
assert!(refreshed
.expires_at
.as_deref()
.is_some_and(|value| value.contains('T')));
}
#[tokio::test]
async fn system_proxy_url_reads_trimmed_runtime_setting() {
let pool = test_pool().await;
assert_eq!(system_proxy_url(&pool).await.unwrap(), None);
sqlx::query("INSERT INTO settings (key, value) VALUES (?, ?)")
.bind(SYSTEM_PROXY_URL_SETTING_KEY)
.bind(" http://127.0.0.1:7890 ")
.execute(&pool)
.await
.expect("insert system proxy");
assert_eq!(
system_proxy_url(&pool).await.unwrap().as_deref(),
Some("http://127.0.0.1:7890")
);
}
}
src/services/route_rebuild.rs
+111 -5
View File
@@ -51,6 +51,7 @@ pub struct RebuildSummary {
struct AvailRow {
account_id: i64,
model_name: String,
oauth_route_unit_id: Option<i64>,
}
/// Run the rebuild against the live database. The whole thing runs inside a
@@ -75,10 +76,11 @@ pub(crate) async fn rebuild_routes_inner(
// lets us drop inactive accounts and apply site_disabled_models in one
// query — keeping the in-memory set small for installs with hundreds of
// accounts.
let rows: Vec<(i64, String)> = sqlx::query_as(
"SELECT ma.account_id, ma.model_name
let rows: Vec<(i64, String, Option<i64>)> = sqlx::query_as(
"SELECT ma.account_id, ma.model_name, rum.unit_id
FROM model_availability ma
JOIN accounts a ON a.id = ma.account_id
LEFT JOIN oauth_route_unit_members rum ON rum.account_id = a.id
WHERE COALESCE(ma.available, 0) = 1
AND a.status = 'active'
AND NOT EXISTS (
@@ -93,9 +95,10 @@ pub(crate) async fn rebuild_routes_inner(
let avail: Vec<AvailRow> = rows
.into_iter()
.map(|(account_id, model_name)| AvailRow {
.map(|(account_id, model_name, oauth_route_unit_id)| AvailRow {
account_id,
model_name,
oauth_route_unit_id,
})
.collect();
@@ -112,11 +115,16 @@ pub(crate) async fn rebuild_routes_inner(
// 2. Group by model name. Each group becomes one token_routes row.
let mut by_model: HashMap<String, Vec<i64>> = HashMap::new();
let mut route_unit_by_key: HashMap<(i64, String), Option<i64>> = HashMap::new();
for row in &avail {
by_model
.entry(row.model_name.clone())
.or_default()
.push(row.account_id);
route_unit_by_key.insert(
(row.account_id, row.model_name.clone()),
row.oauth_route_unit_id,
);
}
// 3. Pre-load existing routes by model_pattern so we can reuse rows. The
@@ -211,21 +219,41 @@ pub(crate) async fn rebuild_routes_inner(
.fetch_optional(&mut *tx)
.await?;
let oauth_route_unit_id = route_unit_by_key.get(key).copied().flatten();
sqlx::query(
"INSERT INTO route_channels
(route_id, account_id, token_id, source_model,
priority, weight, enabled, manual_override)
VALUES (?, ?, ?, ?, 0, 10, 1, 0)",
oauth_route_unit_id, priority, weight, enabled, manual_override)
VALUES (?, ?, ?, ?, ?, 0, 10, 1, 0)",
)
.bind(route_id)
.bind(key.0)
.bind(token_id)
.bind(&key.1)
.bind(oauth_route_unit_id)
.execute(&mut *tx)
.await?;
summary.created_channels += 1;
}
// Keep OAuth route-unit annotations in sync for existing auto-managed channels.
for (key, ch_id) in &auto_by_key {
if !desired.contains(key) {
continue;
}
let oauth_route_unit_id = route_unit_by_key.get(key).copied().flatten();
sqlx::query(
"UPDATE route_channels
SET oauth_route_unit_id = ?
WHERE id = ?
AND COALESCE(manual_override, 0) = 0",
)
.bind(oauth_route_unit_id)
.bind(ch_id)
.execute(&mut *tx)
.await?;
}
// Delete auto channels that are no longer desired (account dropped the
// model, or the model is now in site_disabled_models).
for (key, ch_id) in &auto_by_key {
@@ -379,6 +407,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,
@@ -391,6 +420,20 @@ mod tests {
enabled INTEGER DEFAULT 1,
is_default INTEGER DEFAULT 0
)",
"CREATE TABLE oauth_route_units (
id INTEGER PRIMARY KEY AUTOINCREMENT,
site_id INTEGER NOT NULL,
provider TEXT NOT NULL,
name TEXT NOT NULL,
strategy TEXT NOT NULL DEFAULT 'round_robin',
enabled INTEGER DEFAULT 1
)",
"CREATE TABLE oauth_route_unit_members (
id INTEGER PRIMARY KEY AUTOINCREMENT,
unit_id INTEGER NOT NULL,
account_id INTEGER NOT NULL,
sort_order INTEGER DEFAULT 0
)",
] {
sqlx::query(ddl)
.execute(&pool)
@@ -524,6 +567,69 @@ mod tests {
assert_eq!(route_count, 1);
}
#[tokio::test]
async fn rebuild_marks_auto_channels_with_oauth_route_unit_id() {
let pool = fixture_pool().await;
sqlx::query(
"INSERT INTO accounts (id, site_id, status) VALUES
(1, 1, 'active'),
(2, 1, 'active')",
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO model_availability (account_id, model_name, available) VALUES
(1, 'gpt-oauth', 1),
(2, 'gpt-oauth', 1)",
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO oauth_route_units (id, site_id, provider, name, strategy)
VALUES (9, 1, 'codex', 'Codex Pool', 'round_robin')",
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO oauth_route_unit_members (unit_id, account_id, sort_order) VALUES
(9, 1, 0),
(9, 2, 1)",
)
.execute(&pool)
.await
.unwrap();
let summary = rebuild_routes_inner(&pool).await.expect("rebuild");
assert_eq!(summary.created_channels, 2);
let unit_ids: Vec<Option<i64>> = sqlx::query_scalar(
"SELECT oauth_route_unit_id FROM route_channels ORDER BY account_id",
)
.fetch_all(&pool)
.await
.unwrap();
assert_eq!(unit_ids, vec![Some(9), Some(9)]);
sqlx::query("DELETE FROM oauth_route_unit_members WHERE unit_id = 9")
.execute(&pool)
.await
.unwrap();
let summary = rebuild_routes_inner(&pool).await.expect("second rebuild");
assert_eq!(summary.created_channels, 0);
let unit_ids: Vec<Option<i64>> = sqlx::query_scalar(
"SELECT oauth_route_unit_id FROM route_channels ORDER BY account_id",
)
.fetch_all(&pool)
.await
.unwrap();
assert_eq!(unit_ids, vec![None, None]);
}
/// User-reported bug: empty DB still shows hundreds of routes on the
/// `/routes` page, and clicking "rebuild" does nothing. Reproduces by
/// pre-seeding `token_routes` rows that no account / availability row