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/proxy/token_router/router.rs
+734 -33
View File
@@ -2,21 +2,26 @@ use super::types::{RouteDecision, RouteMatch};
use crate::error::{AppError, AppResult};
use crate::proxy::conductor::{FailureInfo, SelectedChannel, SuccessMetrics};
use crate::proxy::executors::UpstreamRoute;
use crate::services::oauth::OAuthService;
use once_cell::sync::Lazy;
use parking_lot::RwLock;
use regex::Regex;
use sqlx::{FromRow, SqlitePool};
use std::{
collections::HashMap,
collections::{HashMap, HashSet},
sync::Arc,
time::{Duration, Instant},
};
const ROUTE_CACHE_TTL: Duration = Duration::from_secs(2);
const DEFAULT_FAILURE_COOLDOWN_MAX_SECONDS: i64 = 30 * 24 * 60 * 60;
const FAILURE_COOLDOWN_MAX_SETTING_KEY: &str = "runtime.token_router_failure_cooldown_max_sec";
static ROUTE_ROWS_CACHE: Lazy<RwLock<Option<TimedCache<Vec<RouteRow>>>>> =
Lazy::new(|| RwLock::new(None));
static CHANNEL_ROWS_CACHE: Lazy<RwLock<HashMap<i64, TimedCache<Vec<ChannelRow>>>>> =
type DbCacheKey = usize;
static ROUTE_ROWS_CACHE: Lazy<RwLock<HashMap<DbCacheKey, TimedCache<Vec<RouteRow>>>>> =
Lazy::new(|| RwLock::new(HashMap::new()));
static CHANNEL_ROWS_CACHE: Lazy<RwLock<HashMap<(DbCacheKey, i64), TimedCache<Vec<ChannelRow>>>>> =
Lazy::new(|| RwLock::new(HashMap::new()));
#[derive(Debug, Clone)]
@@ -32,16 +37,10 @@ impl<T> TimedCache<T> {
}
pub fn invalidate_route_cache() {
*ROUTE_ROWS_CACHE.write() = None;
ROUTE_ROWS_CACHE.write().clear();
CHANNEL_ROWS_CACHE.write().clear();
}
fn invalidate_channel_cache_for_channel(channel_id: i64) {
CHANNEL_ROWS_CACHE
.write()
.retain(|_, cache| !cache.value.iter().any(|row| row.channel_id == channel_id));
}
#[derive(Clone)]
pub struct TokenRouter {
db: Arc<SqlitePool>,
@@ -60,15 +59,43 @@ struct ChannelRow {
channel_id: i64,
route_id: i64,
account_id: i64,
oauth_route_unit_id: Option<i64>,
oauth_route_unit_strategy: Option<String>,
site_id: i64,
base_url: String,
site_api_key: Option<String>,
account_access_token: String,
account_api_token: Option<String>,
oauth_provider: Option<String>,
account_unit_cost: Option<f64>,
token_name: Option<String>,
token_value: Option<String>,
source_model: Option<String>,
priority: i64,
weight: i64,
consecutive_fail_count: i64,
last_selected_at: Option<String>,
}
#[derive(Debug, Clone)]
struct RouteUnitCandidate {
strategy: RouteUnitStrategy,
rows: Vec<ChannelRow>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum RouteUnitStrategy {
RoundRobin,
StickUntilUnavailable,
}
impl RouteUnitStrategy {
fn from_db(value: Option<&str>) -> Self {
match value.unwrap_or("round_robin") {
"stick_until_unavailable" => Self::StickUntilUnavailable,
_ => Self::RoundRobin,
}
}
}
impl TokenRouter {
@@ -76,6 +103,20 @@ impl TokenRouter {
Self { db }
}
fn db_cache_key(&self) -> DbCacheKey {
Arc::as_ptr(&self.db) as usize
}
fn invalidate_channel_cache_for_channel(&self, channel_id: i64) {
let db_cache_key = self.db_cache_key();
CHANNEL_ROWS_CACHE
.write()
.retain(|(cache_db_key, _), cache| {
*cache_db_key != db_cache_key
|| !cache.value.iter().any(|row| row.channel_id == channel_id)
});
}
pub async fn resolve(&self, model: &str) -> AppResult<UpstreamRoute> {
let decision = self.resolve_decision(model, &[]).await?;
Ok(UpstreamRoute {
@@ -161,12 +202,27 @@ impl TokenRouter {
exclude_channel_ids: &[i64],
) -> AppResult<Option<SelectedChannel>> {
let rows = self.channel_rows(route.id).await?;
let row = select_row_with_route_units(rows, exclude_channel_ids);
let row = rows
.into_iter()
.find(|row| !exclude_channel_ids.contains(&row.channel_id));
let Some(row) = row else {
return Ok(None);
};
let oauth_access_token = if row
.oauth_provider
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.is_some()
{
let service = OAuthService::new(self.db.clone());
let refreshed = service.refresh_token(row.account_id, row.site_id).await?;
self.invalidate_channel_cache_for_channel(row.channel_id);
Some(refreshed.access_token)
} else {
None
};
Ok(row.map(|row| {
Ok(Some({
let actual_model = route
.model_mapping
.clone()
@@ -177,11 +233,13 @@ impl TokenRouter {
.filter(|value| !value.trim().is_empty())
})
.unwrap_or_else(|| requested_model.to_string());
let api_key = row
.token_value
.clone()
.or(row.account_api_token.clone())
.unwrap_or(row.account_access_token.clone());
let api_key = row.token_value.clone().unwrap_or_else(|| {
oauth_access_token.unwrap_or_else(|| {
row.account_api_token
.clone()
.unwrap_or(row.account_access_token.clone())
})
});
SelectedChannel {
channel_id: row.channel_id,
@@ -203,7 +261,8 @@ impl TokenRouter {
}
async fn route_rows(&self) -> AppResult<Vec<RouteRow>> {
if let Some(cache) = ROUTE_ROWS_CACHE.read().as_ref() {
let db_cache_key = self.db_cache_key();
if let Some(cache) = ROUTE_ROWS_CACHE.read().get(&db_cache_key) {
if cache.fresh() {
return Ok(cache.value.clone());
}
@@ -218,15 +277,19 @@ impl TokenRouter {
.fetch_all(&*self.db)
.await?;
*ROUTE_ROWS_CACHE.write() = Some(TimedCache {
loaded_at: Instant::now(),
value: rows.clone(),
});
ROUTE_ROWS_CACHE.write().insert(
db_cache_key,
TimedCache {
loaded_at: Instant::now(),
value: rows.clone(),
},
);
Ok(rows)
}
async fn channel_rows(&self, route_id: i64) -> AppResult<Vec<ChannelRow>> {
if let Some(cache) = CHANNEL_ROWS_CACHE.read().get(&route_id) {
let db_cache_key = self.db_cache_key();
if let Some(cache) = CHANNEL_ROWS_CACHE.read().get(&(db_cache_key, route_id)) {
if cache.fresh() {
return Ok(cache.value.clone());
}
@@ -237,6 +300,8 @@ impl TokenRouter {
rc.id AS channel_id,
rc.route_id AS route_id,
rc.account_id AS account_id,
rc.oauth_route_unit_id AS oauth_route_unit_id,
oru.strategy AS oauth_route_unit_strategy,
a.site_id AS site_id,
COALESCE((
SELECT sae.url
@@ -250,18 +315,25 @@ impl TokenRouter {
s.api_key AS site_api_key,
a.access_token AS account_access_token,
a.api_token AS account_api_token,
a.oauth_provider AS oauth_provider,
a.unit_cost AS account_unit_cost,
at.name AS token_name,
at.token AS token_value,
rc.source_model AS source_model
rc.source_model AS source_model,
COALESCE(rc.priority, 0) AS priority,
COALESCE(rc.weight, 10) AS weight,
rc.consecutive_fail_count AS consecutive_fail_count,
rc.last_selected_at AS last_selected_at
FROM route_channels rc
JOIN accounts a ON a.id = rc.account_id
JOIN sites s ON s.id = a.site_id
LEFT JOIN oauth_route_units oru ON oru.id = rc.oauth_route_unit_id AND COALESCE(oru.enabled, 1) = 1
LEFT JOIN account_tokens at ON at.id = rc.token_id AND at.enabled = 1
WHERE rc.route_id = ?
AND rc.enabled = 1
AND a.status = 'active'
AND s.status = 'active'
AND (rc.oauth_route_unit_id IS NULL OR oru.id IS NOT NULL)
AND (rc.cooldown_until IS NULL OR rc.cooldown_until <= datetime('now'))
AND (rc.token_id IS NULL OR at.id IS NOT NULL)
ORDER BY
@@ -276,7 +348,7 @@ impl TokenRouter {
.await?;
CHANNEL_ROWS_CACHE.write().insert(
route_id,
(db_cache_key, route_id),
TimedCache {
loaded_at: Instant::now(),
value: rows.clone(),
@@ -303,12 +375,14 @@ impl TokenRouter {
.bind(channel_id)
.execute(&*self.db)
.await?;
invalidate_channel_cache_for_channel(channel_id);
self.invalidate_channel_cache_for_channel(channel_id);
self.record_route_unit_member_success(channel_id, &metrics)
.await?;
Ok(())
}
pub async fn record_failure(&self, channel_id: i64, info: FailureInfo) -> AppResult<()> {
let cooldown_seconds = cooldown_seconds_for_failure(&info);
let cooldown_seconds = self.cooldown_seconds_for_failure(&info).await?;
sqlx::query(
"UPDATE route_channels
SET fail_count = COALESCE(fail_count, 0) + 1,
@@ -326,9 +400,156 @@ impl TokenRouter {
.bind(channel_id)
.execute(&*self.db)
.await?;
invalidate_channel_cache_for_channel(channel_id);
self.invalidate_channel_cache_for_channel(channel_id);
self.record_route_unit_member_failure(channel_id, info, cooldown_seconds)
.await?;
Ok(())
}
async fn cooldown_seconds_for_failure(&self, info: &FailureInfo) -> AppResult<i64> {
let raw_max: Option<String> =
sqlx::query_scalar("SELECT value FROM settings WHERE key = ?")
.bind(FAILURE_COOLDOWN_MAX_SETTING_KEY)
.fetch_optional(&*self.db)
.await?;
let max_seconds = raw_max
.and_then(|value| value.trim().parse::<i64>().ok())
.unwrap_or(DEFAULT_FAILURE_COOLDOWN_MAX_SECONDS)
.max(0);
Ok(cooldown_seconds_for_failure(info).min(max_seconds))
}
async fn record_route_unit_member_success(
&self,
channel_id: i64,
metrics: &SuccessMetrics,
) -> AppResult<()> {
sqlx::query(
"UPDATE oauth_route_unit_members
SET success_count = COALESCE(success_count, 0) + 1,
total_latency_ms = COALESCE(total_latency_ms, 0) + ?,
total_cost = COALESCE(total_cost, 0) + ?,
consecutive_fail_count = 0,
cooldown_level = 0,
cooldown_until = NULL,
last_used_at = datetime('now'),
last_selected_at = datetime('now'),
updated_at = datetime('now')
WHERE account_id = (
SELECT account_id FROM route_channels WHERE id = ?
)
AND unit_id = (
SELECT oauth_route_unit_id FROM route_channels WHERE id = ?
)",
)
.bind(metrics.latency_ms.unwrap_or(0) as i64)
.bind(metrics.cost.unwrap_or(0.0))
.bind(channel_id)
.bind(channel_id)
.execute(&*self.db)
.await?;
Ok(())
}
async fn record_route_unit_member_failure(
&self,
channel_id: i64,
_info: FailureInfo,
cooldown_seconds: i64,
) -> AppResult<()> {
sqlx::query(
"UPDATE oauth_route_unit_members
SET fail_count = COALESCE(fail_count, 0) + 1,
consecutive_fail_count = consecutive_fail_count + 1,
cooldown_level = cooldown_level + 1,
cooldown_until = CASE
WHEN ? > 0 THEN datetime('now', '+' || ? || ' seconds')
ELSE cooldown_until
END,
last_fail_at = datetime('now'),
updated_at = datetime('now')
WHERE account_id = (
SELECT account_id FROM route_channels WHERE id = ?
)
AND unit_id = (
SELECT oauth_route_unit_id FROM route_channels WHERE id = ?
)",
)
.bind(cooldown_seconds)
.bind(cooldown_seconds)
.bind(channel_id)
.bind(channel_id)
.execute(&*self.db)
.await?;
Ok(())
}
}
fn select_row_with_route_units(
rows: Vec<ChannelRow>,
exclude_channel_ids: &[i64],
) -> Option<ChannelRow> {
let excluded: HashSet<i64> = exclude_channel_ids.iter().copied().collect();
let mut candidates: Vec<Result<ChannelRow, RouteUnitCandidate>> = Vec::new();
let mut unit_candidate_indexes: HashMap<i64, usize> = HashMap::new();
for row in rows {
if excluded.contains(&row.channel_id) {
continue;
}
let Some(unit_id) = row.oauth_route_unit_id else {
candidates.push(Ok(row));
continue;
};
if let Some(index) = unit_candidate_indexes.get(&unit_id).copied() {
if let Some(Err(candidate)) = candidates.get_mut(index) {
candidate.rows.push(row);
}
continue;
}
let strategy = RouteUnitStrategy::from_db(row.oauth_route_unit_strategy.as_deref());
let index = candidates.len();
candidates.push(Err(RouteUnitCandidate {
strategy,
rows: vec![row],
}));
unit_candidate_indexes.insert(unit_id, index);
}
candidates
.into_iter()
.find_map(|candidate| match candidate {
Ok(row) => Some(row),
Err(candidate) => select_route_unit_member(candidate),
})
}
fn select_route_unit_member(candidate: RouteUnitCandidate) -> Option<ChannelRow> {
let mut rows = candidate.rows;
match candidate.strategy {
RouteUnitStrategy::RoundRobin => {
rows.sort_by(|a, b| {
a.last_selected_at
.cmp(&b.last_selected_at)
.then_with(|| a.consecutive_fail_count.cmp(&b.consecutive_fail_count))
.then_with(|| b.weight.cmp(&a.weight))
.then_with(|| a.channel_id.cmp(&b.channel_id))
});
}
RouteUnitStrategy::StickUntilUnavailable => {
rows.sort_by(|a, b| {
b.priority
.cmp(&a.priority)
.then_with(|| b.weight.cmp(&a.weight))
.then_with(|| a.consecutive_fail_count.cmp(&b.consecutive_fail_count))
.then_with(|| a.channel_id.cmp(&b.channel_id))
});
}
}
rows.into_iter().next()
}
fn cooldown_seconds_for_failure(info: &FailureInfo) -> i64 {
@@ -393,14 +614,27 @@ fn wildcard_matches(model: &str, pattern: &str) -> bool {
mod tests {
use super::*;
use sqlx::sqlite::SqlitePoolOptions;
use tokio::sync::Mutex;
static TEST_CACHE_LOCK: Lazy<Mutex<()>> = Lazy::new(|| Mutex::new(()));
async fn test_pool() -> SqlitePool {
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 sites (
id INTEGER PRIMARY KEY,
@@ -408,7 +642,8 @@ mod tests {
url TEXT NOT NULL,
platform TEXT NOT NULL DEFAULT 'openai',
status TEXT NOT NULL DEFAULT 'active',
api_key TEXT
api_key TEXT,
custom_headers TEXT
)",
)
.execute(&pool)
@@ -432,8 +667,11 @@ 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'
status TEXT NOT NULL DEFAULT 'active',
updated_at TEXT
)",
)
.execute(&pool)
@@ -457,6 +695,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,
@@ -476,6 +715,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();
// select_channel_for_route's COALESCE subquery joins through
// site_api_endpoints. The table must exist or the SELECT errors
// out — we leave it empty and the COALESCE falls back to sites.url.
@@ -501,16 +775,25 @@ mod tests {
.execute(&pool)
.await
.unwrap();
sqlx::query("INSERT INTO accounts (id, site_id, access_token, api_token) VALUES (2, 1, 'access-key-2', NULL)")
.execute(&pool)
.await
.unwrap();
sqlx::query("INSERT INTO account_tokens (id, account_id, name, token) VALUES (1, 1, 'primary', 'token-key')")
.execute(&pool)
.await
.unwrap();
sqlx::query("INSERT INTO account_tokens (id, account_id, name, token) VALUES (2, 2, 'secondary', 'token-key-2')")
.execute(&pool)
.await
.unwrap();
pool
}
#[tokio::test]
async fn resolves_exact_model_route() {
let _guard = TEST_CACHE_LOCK.lock().await;
let pool = test_pool().await;
sqlx::query("INSERT INTO token_routes (id, model_pattern, model_mapping, route_mode) VALUES (1, 'gpt-5.5', 'upstream-gpt-5.5', 'exact')")
.execute(&pool)
@@ -530,8 +813,57 @@ mod tests {
assert_eq!(decision.selected.api_key, "token-key");
}
#[tokio::test]
async fn route_cache_is_scoped_to_database_instance() {
let _guard = TEST_CACHE_LOCK.lock().await;
let first_pool = test_pool().await;
let second_pool = test_pool().await;
sqlx::query("INSERT INTO token_routes (id, model_pattern, route_mode) VALUES (1, 'gpt-alpha', 'exact')")
.execute(&first_pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO route_channels (id, route_id, account_id, token_id) VALUES (1, 1, 1, 1)",
)
.execute(&first_pool)
.await
.unwrap();
sqlx::query("INSERT INTO token_routes (id, model_pattern, route_mode) VALUES (1, 'gpt-beta', 'exact')")
.execute(&second_pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO route_channels (id, route_id, account_id, token_id) VALUES (1, 1, 1, 1)",
)
.execute(&second_pool)
.await
.unwrap();
let first_router = TokenRouter::new(Arc::new(first_pool));
let second_router = TokenRouter::new(Arc::new(second_pool));
let first = first_router
.resolve_decision("gpt-alpha", &[])
.await
.unwrap();
assert_eq!(first.match_kind, RouteMatch::Exact);
let second = second_router
.resolve_decision("gpt-beta", &[])
.await
.unwrap();
assert_eq!(second.match_kind, RouteMatch::Exact);
assert!(second_router
.resolve_decision("gpt-alpha", &[])
.await
.is_err());
}
#[tokio::test]
async fn resolves_wildcard_route_and_skips_excluded_channel() {
let _guard = TEST_CACHE_LOCK.lock().await;
let pool = test_pool().await;
sqlx::query("INSERT INTO token_routes (id, model_pattern, route_mode) VALUES (1, 'gpt-*', 'pattern')")
.execute(&pool)
@@ -551,4 +883,373 @@ mod tests {
assert_eq!(decision.match_kind, RouteMatch::Wildcard);
assert_eq!(decision.selected.channel_id, 2);
}
#[tokio::test]
async fn oauth_channel_uses_access_token_instead_of_refresh_token() {
let _guard = TEST_CACHE_LOCK.lock().await;
let pool = test_pool().await;
sqlx::query(
"INSERT INTO token_routes (id, model_pattern, route_mode) VALUES (1, 'gpt-*', 'pattern')",
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"UPDATE accounts
SET access_token = 'oauth-access',
api_token = 'oauth-refresh',
oauth_provider = 'codex',
extra_config = ?
WHERE id = 1",
)
.bind(
serde_json::json!({
"credentialMode": "oauth",
"expires_at": (chrono::Utc::now() + chrono::Duration::hours(1)).to_rfc3339()
})
.to_string(),
)
.execute(&pool)
.await
.unwrap();
sqlx::query("INSERT INTO route_channels (id, route_id, account_id, token_id) VALUES (1, 1, 1, NULL)")
.execute(&pool)
.await
.unwrap();
let router = TokenRouter::new(Arc::new(pool));
let decision = router.resolve_decision("gpt-4o", &[]).await.unwrap();
assert_eq!(decision.selected.api_key, "oauth-access");
}
#[tokio::test]
async fn oauth_channel_refreshes_expired_access_token_before_selection() {
use wiremock::matchers::{body_string_contains, method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
let _guard = TEST_CACHE_LOCK.lock().await;
let pool = test_pool().await;
let mock = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/oauth/token"))
.and(body_string_contains("refresh_token=oauth-refresh-old"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"access_token": "oauth-access-new",
"refresh_token": "oauth-refresh-new",
"expires_in": 3600
})))
.mount(&mock)
.await;
sqlx::query(
"UPDATE sites
SET api_key = NULL,
platform = 'codex',
custom_headers = ?
WHERE id = 1",
)
.bind(
serde_json::json!({
"oauth_token_url": format!("{}/oauth/token", mock.uri())
})
.to_string(),
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"UPDATE accounts
SET access_token = 'oauth-access-old',
api_token = 'oauth-refresh-old',
oauth_provider = 'codex',
extra_config = ?
WHERE id = 1",
)
.bind(
serde_json::json!({
"credentialMode": "oauth",
"expires_at": "2000-01-01T00:00:00Z"
})
.to_string(),
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO token_routes (id, model_pattern, route_mode) VALUES (1, 'gpt-*', 'pattern')",
)
.execute(&pool)
.await
.unwrap();
sqlx::query("INSERT INTO route_channels (id, route_id, account_id, token_id) VALUES (1, 1, 1, NULL)")
.execute(&pool)
.await
.unwrap();
let router = TokenRouter::new(Arc::new(pool.clone()));
let decision = router.resolve_decision("gpt-4o", &[]).await.unwrap();
assert_eq!(decision.selected.api_key, "oauth-access-new");
let row: (String, Option<String>, Option<String>) = sqlx::query_as(
"SELECT access_token, api_token, extra_config FROM accounts WHERE id = 1",
)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(row.0, "oauth-access-new");
assert_eq!(row.1.as_deref(), Some("oauth-refresh-new"));
let extra: serde_json::Value = serde_json::from_str(row.2.as_deref().unwrap()).unwrap();
assert_eq!(
extra
.get("refreshToken")
.and_then(serde_json::Value::as_str),
Some("oauth-refresh-new")
);
}
#[tokio::test]
async fn route_unit_selects_one_member_and_failover_stays_inside_pool() {
let _guard = TEST_CACHE_LOCK.lock().await;
let pool = test_pool().await;
sqlx::query("INSERT INTO token_routes (id, model_pattern, route_mode) VALUES (1, 'gpt-*', 'pattern')")
.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 (id, unit_id, account_id, sort_order) VALUES
(91, 9, 1, 0),
(92, 9, 2, 1)",
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO route_channels (id, route_id, account_id, token_id, oauth_route_unit_id, priority) VALUES
(1, 1, 1, 1, 9, 10),
(2, 1, 2, 2, 9, 10),
(3, 1, 1, 1, NULL, 1)",
)
.execute(&pool)
.await
.unwrap();
let router = TokenRouter::new(Arc::new(pool.clone()));
let first = router.resolve_decision("gpt-4o", &[]).await.unwrap();
assert_eq!(first.selected.channel_id, 1);
let failover = router
.resolve_next_selected("gpt-4o", &[first.selected.channel_id])
.await
.unwrap();
assert_eq!(failover.channel_id, 2);
assert_eq!(failover.api_key, "token-key-2");
router
.record_success(
failover.channel_id,
SuccessMetrics {
latency_ms: Some(123),
cost: Some(0.25),
metadata: HashMap::new(),
},
)
.await
.unwrap();
let member_success: (i64, i64, f64) = sqlx::query_as(
"SELECT success_count, total_latency_ms, total_cost
FROM oauth_route_unit_members
WHERE unit_id = 9 AND account_id = 2",
)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(member_success, (1, 123, 0.25));
router
.record_failure(
first.selected.channel_id,
FailureInfo {
action: "failover".to_string(),
status: Some(429),
raw_error_text: Some("rate limited".to_string()),
},
)
.await
.unwrap();
let member_failure: (i64, i64) = sqlx::query_as(
"SELECT fail_count, consecutive_fail_count
FROM oauth_route_unit_members
WHERE unit_id = 9 AND account_id = 1",
)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(member_failure, (1, 1));
}
#[tokio::test]
async fn route_unit_round_robin_prefers_least_recent_member() {
let _guard = TEST_CACHE_LOCK.lock().await;
let pool = test_pool().await;
sqlx::query("INSERT INTO token_routes (id, model_pattern, route_mode) VALUES (1, 'gpt-*', 'pattern')")
.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 (id, unit_id, account_id, sort_order) VALUES
(91, 9, 1, 0),
(92, 9, 2, 1)",
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO route_channels (
id, route_id, account_id, token_id, oauth_route_unit_id, priority, last_selected_at
) VALUES
(1, 1, 1, 1, 9, 10, '2026-01-02 00:00:00'),
(2, 1, 2, 2, 9, 10, '2026-01-01 00:00:00')",
)
.execute(&pool)
.await
.unwrap();
let router = TokenRouter::new(Arc::new(pool.clone()));
let first = router.resolve_decision("gpt-4o", &[]).await.unwrap();
assert_eq!(first.selected.channel_id, 2);
router
.record_success(
first.selected.channel_id,
SuccessMetrics {
latency_ms: Some(88),
cost: Some(0.12),
metadata: HashMap::new(),
},
)
.await
.unwrap();
let next = router.resolve_decision("gpt-4o", &[]).await.unwrap();
assert_eq!(next.selected.channel_id, 1);
}
#[tokio::test]
async fn route_unit_stick_strategy_keeps_primary_until_it_cools_down() {
let _guard = TEST_CACHE_LOCK.lock().await;
let pool = test_pool().await;
sqlx::query("INSERT INTO token_routes (id, model_pattern, route_mode) VALUES (1, 'gpt-*', 'pattern')")
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO oauth_route_units (id, site_id, provider, name, strategy)
VALUES (9, 1, 'codex', 'Codex Pool', 'stick_until_unavailable')",
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO oauth_route_unit_members (id, unit_id, account_id, sort_order) VALUES
(91, 9, 1, 0),
(92, 9, 2, 1)",
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO route_channels (
id, route_id, account_id, token_id, oauth_route_unit_id, priority, weight, last_selected_at
) VALUES
(1, 1, 1, 1, 9, 20, 10, '2026-01-02 00:00:00'),
(2, 1, 2, 2, 9, 10, 100, NULL)",
)
.execute(&pool)
.await
.unwrap();
let router = TokenRouter::new(Arc::new(pool.clone()));
let first = router.resolve_decision("gpt-4o", &[]).await.unwrap();
assert_eq!(first.selected.channel_id, 1);
let again = router.resolve_decision("gpt-4o", &[]).await.unwrap();
assert_eq!(again.selected.channel_id, 1);
router
.record_failure(
first.selected.channel_id,
FailureInfo {
action: "failover".to_string(),
status: Some(429),
raw_error_text: Some("rate limited".to_string()),
},
)
.await
.unwrap();
let failover = router.resolve_decision("gpt-4o", &[]).await.unwrap();
assert_eq!(failover.selected.channel_id, 2);
}
#[tokio::test]
async fn record_failure_clamps_cooldown_to_runtime_setting() {
let _guard = TEST_CACHE_LOCK.lock().await;
let pool = test_pool().await;
sqlx::query("INSERT INTO settings (key, value) VALUES (?, ?)")
.bind(FAILURE_COOLDOWN_MAX_SETTING_KEY)
.bind("30")
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO token_routes (id, model_pattern, route_mode) VALUES (1, 'gpt-4o', 'exact')",
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO route_channels (id, route_id, account_id, token_id) VALUES (1, 1, 1, 1)",
)
.execute(&pool)
.await
.unwrap();
let router = TokenRouter::new(Arc::new(pool.clone()));
router
.record_failure(
1,
FailureInfo {
action: "failover".to_string(),
status: Some(429),
raw_error_text: Some("rate limited".to_string()),
},
)
.await
.unwrap();
let seconds_until: i64 = sqlx::query_scalar(
"SELECT CAST(strftime('%s', cooldown_until) AS INTEGER) - CAST(strftime('%s', 'now') AS INTEGER)
FROM route_channels
WHERE id = 1",
)
.fetch_one(&pool)
.await
.unwrap();
assert!(
(0..=35).contains(&seconds_until),
"cooldown should be clamped near 30 seconds, got {seconds_until}"
);
}
}