starred/codex-proxy-rs
1
0
Fork 0
mirror of https://github.com/zyycn/codex-proxy-rs.git synced 2026-11-13 16:07:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: clean account routing and affinity fallback

Browse Source 
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Zyy
2026-07-07 14:01:35 +08:00
co-authored by Codex
parent a6c75098be
commit 64f2802434
16 changed files with 226 additions and 360 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/src/admin/settings/routes.rs
+1 -4
View File
@@ -16,9 +16,8 @@ use crate::{
upstream::models::config::ModelConfig,
};
const ALLOWED_SETTINGS_KEYS: [&str; 7] = [
const ALLOWED_SETTINGS_KEYS: [&str; 6] = [
"modelAliases",
"modelAccountRoutes",
"refreshMarginSeconds",
"refreshConcurrency",
"maxConcurrentPerAccount",
@@ -30,7 +29,6 @@ const ALLOWED_SETTINGS_KEYS: [&str; 7] = [
#[serde(rename_all = "camelCase")]
struct AdminSettingsData {
model_aliases: BTreeMap<String, String>,
model_account_routes: BTreeMap<String, Vec<String>>,
refresh_margin_seconds: u64,
refresh_concurrency: u32,
max_concurrent_per_account: usize,
@@ -42,7 +40,6 @@ impl AdminSettingsData {
fn from_config(config: &AppConfig) -> Self {
Self {
model_aliases: config.model_aliases.clone(),
model_account_routes: config.model_account_routes.clone(),
refresh_margin_seconds: config.auth.refresh_margin_seconds,
refresh_concurrency: config.auth.refresh_concurrency,
max_concurrent_per_account: config.auth.max_concurrent_per_account,
backend/src/config/schema.rs
-3
View File
@@ -24,9 +24,6 @@ pub struct AppConfig {
/// 运行时模型别名，由数据库设置加载，配置文件不承载该字段。
#[serde(skip)]
pub model_aliases: BTreeMap<String, String>,
/// 运行时模型账号路由，由 `model_account_routes` 表加载，配置文件不承载该字段。
#[serde(skip)]
pub model_account_routes: BTreeMap<String, Vec<String>>,
/// 认证与刷新配置，由固定默认值和数据库运行时设置承载，配置文件不承载该字段。
#[serde(skip)]
pub auth: AuthConfig,
backend/src/config/settings.rs
+4 -109
View File
@@ -26,8 +26,6 @@ const RUNTIME_SETTINGS_ID: i64 = 1;
pub struct AdminSettings {
/// 模型别名映射。
pub model_aliases: BTreeMap<String, String>,
/// 模型到账号 ID 的显式路由。
pub model_account_routes: BTreeMap<String, Vec<String>>,
/// 访问令牌过期前多少秒开始刷新。
pub refresh_margin_seconds: u64,
/// 访问令牌刷新并发数。
@@ -44,7 +42,6 @@ impl Default for AdminSettings {
fn default() -> Self {
Self {
model_aliases: BTreeMap::new(),
model_account_routes: BTreeMap::new(),
refresh_margin_seconds: 300,
refresh_concurrency: 2,
max_concurrent_per_account: 3,
@@ -60,8 +57,6 @@ impl Default for AdminSettings {
pub struct AdminSettingsPatch {
/// 模型别名映射。
pub model_aliases: Option<BTreeMap<String, String>>,
/// 模型到账号 ID 的显式路由。
pub model_account_routes: Option<BTreeMap<String, Vec<String>>>,
/// 访问令牌过期前多少秒开始刷新。
pub refresh_margin_seconds: Option<u64>,
/// 访问令牌刷新并发数。
@@ -87,9 +82,6 @@ impl SettingsService {
if let Some(model_aliases) = patch.model_aliases {
current.model_aliases = validate_model_aliases(model_aliases)?;
}
if let Some(model_account_routes) = patch.model_account_routes {
current.model_account_routes = validate_model_account_routes(model_account_routes)?;
}
if let Some(refresh_margin_seconds) = patch.refresh_margin_seconds {
current.refresh_margin_seconds =
positive_u64("refreshMarginSeconds", refresh_margin_seconds)?;
@@ -151,34 +143,6 @@ fn validate_model_aliases(
Ok(normalized)
}
fn validate_model_account_routes(
routes: BTreeMap<String, Vec<String>>,
) -> Result<BTreeMap<String, Vec<String>>, SettingsServiceError> {
let mut normalized = BTreeMap::new();
for (model, account_ids) in routes {
let model = non_empty_string("modelAccountRoutes", &model)?;
let mut normalized_account_ids = Vec::new();
for account_id in account_ids {
let account_id = non_empty_string("modelAccountRoutes", &account_id)?;
if normalized_account_ids.contains(&account_id) {
return Err(invalid_field(
"modelAccountRoutes",
format!("duplicate account id `{account_id}`"),
));
}
normalized_account_ids.push(account_id);
}
if normalized_account_ids.is_empty() {
return Err(invalid_field(
"modelAccountRoutes",
"account list must not be empty",
));
}
normalized.insert(model, normalized_account_ids);
}
Ok(normalized)
}
fn validate_rotation_strategy(strategy: &str) -> Result<String, SettingsServiceError> {
let strategy = non_empty_string("rotationStrategy", strategy)?;
if ROTATION_STRATEGIES.contains(&strategy.as_str()) {
@@ -394,7 +358,6 @@ pub enum RuntimeSettingsError {
fn admin_settings_from_config(config: &AppConfig) -> AdminSettings {
AdminSettings {
model_aliases: config.model_aliases.clone(),
model_account_routes: config.model_account_routes.clone(),
refresh_margin_seconds: config.auth.refresh_margin_seconds,
refresh_concurrency: config.auth.refresh_concurrency,
max_concurrent_per_account: config.auth.max_concurrent_per_account,
@@ -405,7 +368,6 @@ fn admin_settings_from_config(config: &AppConfig) -> AdminSettings {
fn apply_admin_settings_to_config(config: &mut AppConfig, settings: AdminSettings) {
config.model_aliases = settings.model_aliases;
config.model_account_routes = settings.model_account_routes;
config.auth.refresh_margin_seconds = settings.refresh_margin_seconds;
config.auth.refresh_concurrency = settings.refresh_concurrency;
config.auth.max_concurrent_per_account = settings.max_concurrent_per_account;
@@ -418,8 +380,7 @@ async fn insert_runtime_settings_if_missing(
settings: &AdminSettings,
) -> Result<(), RuntimeSettingsError> {
let model_aliases_json = serde_json::to_string(&settings.model_aliases)?;
let mut tx = pool.begin().await?;
let result = sqlx::query(
sqlx::query(
r"
insert or ignore into runtime_settings (
id,
@@ -449,12 +410,8 @@ insert or ignore into runtime_settings (
)
.bind(&settings.rotation_strategy)
.bind(Utc::now().to_rfc3339())
.execute(&mut *tx)
.execute(pool)
.await?;
if result.rows_affected() > 0 {
replace_model_account_routes(&mut tx, &settings.model_account_routes).await?;
}
tx.commit().await?;
Ok(())
}
@@ -463,7 +420,6 @@ async fn save_runtime_settings(
settings: &AdminSettings,
) -> Result<(), RuntimeSettingsError> {
let model_aliases_json = serde_json::to_string(&settings.model_aliases)?;
let mut tx = pool.begin().await?;
sqlx::query(
r"
insert into runtime_settings (
@@ -502,10 +458,8 @@ on conflict(id) do update set
)
.bind(&settings.rotation_strategy)
.bind(Utc::now().to_rfc3339())
.execute(&mut *tx)
.execute(pool)
.await?;
replace_model_account_routes(&mut tx, &settings.model_account_routes).await?;
tx.commit().await?;
Ok(())
}
@@ -525,9 +479,7 @@ where id = ?",
.bind(RUNTIME_SETTINGS_ID)
.fetch_one(pool)
.await?;
let mut settings = runtime_settings_from_row(&row)?;
settings.model_account_routes = load_model_account_routes(pool).await?;
Ok(settings)
runtime_settings_from_row(&row)
}
async fn admin_api_key_status(
@@ -566,7 +518,6 @@ fn runtime_settings_from_row(row: &SqliteRow) -> Result<AdminSettings, RuntimeSe
let model_aliases_json: String = row.get("model_aliases_json");
Ok(AdminSettings {
model_aliases: serde_json::from_str(&model_aliases_json)?,
model_account_routes: BTreeMap::new(),
refresh_margin_seconds: positive_i64_to_u64(
"refreshMarginSeconds",
row.get("refresh_margin_seconds"),
@@ -587,61 +538,6 @@ fn runtime_settings_from_row(row: &SqliteRow) -> Result<AdminSettings, RuntimeSe
})
}
async fn replace_model_account_routes(
tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
routes: &BTreeMap<String, Vec<String>>,
) -> Result<(), RuntimeSettingsError> {
sqlx::query("delete from model_account_routes")
.execute(&mut **tx)
.await?;
let now = Utc::now().to_rfc3339();
for (model, account_ids) in routes {
for (priority, account_id) in account_ids.iter().enumerate() {
sqlx::query(
r"
insert into model_account_routes (
model,
account_id,
priority,
enabled,
created_at,
updated_at
) values (?, ?, ?, 1, ?, ?)",
)
.bind(model)
.bind(account_id)
.bind(
i64::try_from(priority)
.map_err(|_| stored_field_error("modelAccountRoutes", "out of range"))?,
)
.bind(&now)
.bind(&now)
.execute(&mut **tx)
.await?;
}
}
Ok(())
}
async fn load_model_account_routes(
pool: &SqlitePool,
) -> Result<BTreeMap<String, Vec<String>>, RuntimeSettingsError> {
let rows: Vec<(String, String)> = sqlx::query_as(
r"
select model, account_id
from model_account_routes
where enabled = 1
order by model asc, priority asc, account_id asc",
)
.fetch_all(pool)
.await?;
let mut routes = BTreeMap::<String, Vec<String>>::new();
for (model, account_id) in rows {
routes.entry(model).or_default().push(account_id);
}
Ok(routes)
}
/// 从当前运行配置生成账号池调度参数。
pub fn account_pool_options_from_config(config: &AppConfig) -> AccountPoolOptions {
AccountPoolOptions {
@@ -649,7 +545,6 @@ pub fn account_pool_options_from_config(config: &AppConfig) -> AccountPoolOption
rotation_strategy: rotation_strategy_from_config(&config.auth.rotation_strategy),
skip_quota_limited: config.quota.skip_exhausted,
tier_priority: config.auth.tier_priority.clone(),
model_account_routes: config.model_account_routes.clone(),
..AccountPoolOptions::default()
}
}
backend/src/infra/migrations/0001_initial.sql
-18
View File
@@ -176,24 +176,6 @@ on usage_time_buckets(bucket_start);
create index idx_usage_time_buckets_model_bucket
on usage_time_buckets(model, bucket_start);
-- ============================================
-- Model Account Routes
-- ============================================
create table model_account_routes (
model text not null,
account_id text not null references accounts(id) on delete cascade,
priority integer not null default 0 check (priority >= 0),
enabled integer not null default 1 check (enabled in (0, 1)),
created_at text not null,
updated_at text not null,
primary key (model, account_id)
);
create index idx_model_account_routes_account
on model_account_routes(account_id);
create index idx_model_account_routes_enabled_model
on model_account_routes(enabled, model, priority asc, account_id);
-- ============================================
-- Account Cookies
-- ============================================
backend/src/proxy/dispatch/responses/service.rs
+28 -3
View File
@@ -388,6 +388,7 @@ impl ResponseDispatchService {
let mut excluded_account_ids = Vec::new();
let mut exhausted_accounts = AccountExhaustionTracker::default();
let mut history_recovery_used = false;
let mut next_required_account_id: Option<String> = None;
let mut empty_response_retries = 0u8;
let mut quota_verify_attempts = 0usize;
let mut trace = ResponseDispatchTrace::default();
@@ -401,6 +402,12 @@ impl ResponseDispatchService {
let mut attempt_acquire_request = acquire_request
.clone()
.with_exclude_account_ids(excluded_account_ids.iter().cloned());
if let Some(account_id) = next_required_account_id.take() {
attempt_acquire_request =
attempt_acquire_request.with_required_account_id(account_id);
}
let allow_quota_verify_account_retry =
attempt_acquire_request.required_account_id.is_none();
attempt_acquire_request.now = Utc::now();
let Some(acquired) = self
.account_pool
@@ -441,6 +448,7 @@ impl ResponseDispatchService {
request_id,
excluded_account_ids: &mut excluded_account_ids,
verify_attempts: &mut quota_verify_attempts,
allow_retry_with_another_account: allow_quota_verify_account_retry,
},
acquired,
)
@@ -452,7 +460,8 @@ impl ResponseDispatchService {
.record_rate_limited(None, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE);
continue;
}
QuotaVerificationDecision::MaxAttemptsReached => {
QuotaVerificationDecision::MaxAttemptsReached
| QuotaVerificationDecision::RequiredAccountUnavailable => {
exhausted_accounts.record_rate_limited(
Some(&acquired_account_id),
QUOTA_VERIFY_LIMIT_REACHED_MESSAGE,
@@ -568,6 +577,7 @@ impl ResponseDispatchService {
self.recover_request_history(&mut request, &mut implicit_resume)
.await;
history_recovery_used = true;
next_required_account_id = Some(release_account_id);
continue;
}
if is_model_unsupported_sse_failure(failure) {
@@ -665,6 +675,7 @@ impl ResponseDispatchService {
self.recover_request_history(&mut request, &mut implicit_resume)
.await;
history_recovery_used = true;
next_required_account_id = Some(release_account_id);
}
Err(error) if is_auth_upstream_error(&error) => {
let upstream_error = upstream_error_body(&error);
@@ -932,6 +943,7 @@ impl ResponseDispatchService {
let mut excluded_account_ids = Vec::new();
let mut exhausted_accounts = AccountExhaustionTracker::default();
let mut history_recovery_used = false;
let mut next_required_account_id: Option<String> = None;
let mut quota_verify_attempts = 0usize;
let mut trace = ResponseDispatchTrace::default();
macro_rules! return_stream_dispatch_error {
@@ -978,6 +990,12 @@ impl ResponseDispatchService {
let mut attempt_acquire_request = acquire_request
.clone()
.with_exclude_account_ids(excluded_account_ids.iter().cloned());
if let Some(account_id) = next_required_account_id.take() {
attempt_acquire_request =
attempt_acquire_request.with_required_account_id(account_id);
}
let allow_quota_verify_account_retry =
attempt_acquire_request.required_account_id.is_none();
attempt_acquire_request.now = Utc::now();
let Some(acquired) = self
.account_pool
@@ -1000,6 +1018,7 @@ impl ResponseDispatchService {
request_id,
excluded_account_ids: &mut excluded_account_ids,
verify_attempts: &mut quota_verify_attempts,
allow_retry_with_another_account: allow_quota_verify_account_retry,
},
acquired,
)
@@ -1011,7 +1030,8 @@ impl ResponseDispatchService {
.record_rate_limited(None, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE);
continue;
}
QuotaVerificationDecision::MaxAttemptsReached => {
QuotaVerificationDecision::MaxAttemptsReached
| QuotaVerificationDecision::RequiredAccountUnavailable => {
exhausted_accounts.record_rate_limited(
Some(&acquired_account_id),
QUOTA_VERIFY_LIMIT_REACHED_MESSAGE,
@@ -1091,6 +1111,7 @@ impl ResponseDispatchService {
self.recover_request_history(&mut request, &mut implicit_resume)
.await;
history_recovery_used = true;
next_required_account_id = Some(release_account_id);
continue;
}
Err(ResponseDispatchError::Upstream(error))
@@ -1261,6 +1282,7 @@ impl ResponseDispatchService {
self.recover_request_history(&mut request, &mut implicit_resume)
.await;
history_recovery_used = true;
next_required_account_id = Some(release_account_id);
continue;
}
if is_model_unsupported_sse_failure(&failure) {
@@ -1403,6 +1425,7 @@ impl ResponseDispatchService {
self.recover_request_history(&mut request, &mut implicit_resume)
.await;
history_recovery_used = true;
next_required_account_id = Some(release_account_id);
}
Err(error) if is_auth_upstream_error(&error) => {
self.account_pool.release(&release_account_id).await;
@@ -1548,6 +1571,7 @@ impl ResponseDispatchService {
request_id,
excluded_account_ids: &mut excluded_account_ids,
verify_attempts: &mut quota_verify_attempts,
allow_retry_with_another_account: true,
},
acquired,
)
@@ -1559,7 +1583,8 @@ impl ResponseDispatchService {
.record_rate_limited(None, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE);
continue;
}
QuotaVerificationDecision::MaxAttemptsReached => {
QuotaVerificationDecision::MaxAttemptsReached
| QuotaVerificationDecision::RequiredAccountUnavailable => {
exhausted_accounts.record_rate_limited(
Some(&acquired_account_id),
QUOTA_VERIFY_LIMIT_REACHED_MESSAGE,
backend/src/proxy/dispatch/upstream.rs
+6 -1
View File
@@ -34,6 +34,7 @@ pub(crate) const QUOTA_VERIFY_LIMIT_REACHED_MESSAGE: &str =
pub(crate) enum QuotaVerificationDecision {
Ready(Box<AcquiredAccount>),
RetryWithAnotherAccount,
RequiredAccountUnavailable,
MaxAttemptsReached,
}
@@ -45,6 +46,7 @@ pub(crate) struct QuotaVerificationContext<'a> {
pub request_id: &'a str,
pub excluded_account_ids: &'a mut Vec<String>,
pub verify_attempts: &'a mut usize,
pub allow_retry_with_another_account: bool,
}
pub(crate) async fn verify_acquired_quota_if_required(
@@ -111,9 +113,12 @@ pub(crate) async fn verify_acquired_quota_if_required(
quota_verify_required = true,
quota_verify_result = "limit_reached",
verify_attempts = *context.verify_attempts,
retry_with_another_account = !max_attempts_reached,
retry_with_another_account = context.allow_retry_with_another_account && !max_attempts_reached,
"quota verification reported exhausted account before upstream request"
);
if !context.allow_retry_with_another_account {
return QuotaVerificationDecision::RequiredAccountUnavailable;
}
if max_attempts_reached {
return QuotaVerificationDecision::MaxAttemptsReached;
}
backend/src/upstream/accounts/pool.rs
+22 -14
View File
@@ -62,8 +62,6 @@ pub struct AccountPoolOptions {
pub model_plan_allowlist: BTreeMap<String, Vec<String>>,
/// 已成功拉取过模型列表的订阅计划。
pub fetched_model_plan_types: BTreeSet<String>,
/// 模型到指定账号 ID 的显式路由。
pub model_account_routes: BTreeMap<String, Vec<String>>,
}
impl Default for AccountPoolOptions {
@@ -76,7 +74,6 @@ impl Default for AccountPoolOptions {
tier_priority: Vec::new(),
model_plan_allowlist: BTreeMap::new(),
fetched_model_plan_types: BTreeSet::new(),
model_account_routes: BTreeMap::new(),
}
}
}
@@ -88,6 +85,8 @@ pub struct AccountAcquireRequest {
pub model: String,
/// 本次请求需要排除的账号 ID。
pub exclude_account_ids: Vec<String>,
/// 本次请求必须使用的账号 ID，用于同请求内的同账号重试。
pub required_account_id: Option<String>,
/// 会话亲和性建议的优先账号 ID。
pub preferred_account_id: Option<String>,
/// 本次调度使用的当前时间。
@@ -100,6 +99,7 @@ impl AccountAcquireRequest {
Self {
model: model.into(),
exclude_account_ids: Vec::new(),
required_account_id: None,
preferred_account_id: None,
now,
}
@@ -119,6 +119,12 @@ impl AccountAcquireRequest {
self.preferred_account_id = Some(account_id.into());
self
}
/// 设置本次调度必须使用的账号 ID。
pub fn with_required_account_id(mut self, account_id: impl Into<String>) -> Self {
self.required_account_id = Some(account_id.into());
self
}
}
/// 成功获取到的账号及调度元数据。
@@ -480,7 +486,15 @@ impl AccountPool {
self.cleanup_stale_slots(request.now);
let refreshed_accounts = self.refresh_account_statuses(request.now);
let candidates = self.candidates(request);
let selected = if let Some(preferred_account_id) = &request.preferred_account_id {
let selected = if let Some(required_account_id) = &request.required_account_id {
candidates.first().cloned().map(|account| {
(
account,
self.previous_slot_at(required_account_id),
"required",
)
})
} else if let Some(preferred_account_id) = &request.preferred_account_id {
let preferred = candidates
.iter()
.find(|account| account.id == *preferred_account_id)
@@ -742,6 +756,10 @@ impl AccountPool {
.filter(|account| self.is_base_available(account, request))
.cloned()
.collect::<Vec<_>>();
if let Some(required_account_id) = &request.required_account_id {
candidates.retain(|account| account.id == *required_account_id);
return candidates;
}
if let Some(best_tier) = self.best_available_tier(&candidates) {
candidates.retain(|account| account.plan_type.as_deref() == Some(best_tier.as_str()));
}
@@ -756,7 +774,6 @@ impl AccountPool {
request.now,
self.options.skip_quota_limited,
)
&& self.is_model_account_allowed(account, &request.model)
&& self.is_model_allowed(account, &request.model)
&& !request
.exclude_account_ids
@@ -853,15 +870,6 @@ impl AccountPool {
|| !self.options.fetched_model_plan_types.contains(plan_type)
}
fn is_model_account_allowed(&self, account: &Account, model: &str) -> bool {
let Some(account_ids) = self.options.model_account_routes.get(model) else {
return true;
};
account_ids
.iter()
.any(|account_id| account_id == &account.id)
}
fn best_available_tier(&self, candidates: &[Account]) -> Option<String> {
self.options.tier_priority.iter().find_map(|tier| {
candidates
backend/tests/admin/accounts/lifecycle.rs
-17
View File
@@ -118,16 +118,6 @@ async fn seed_account_related_rows(pool: &SqlitePool, account_id: &str) {
.execute(pool)
.await
.unwrap();
sqlx::query(
"insert into model_account_routes (model, account_id, priority, created_at, updated_at) values (?, ?, 0, ?, ?)",
)
.bind("gpt-5.5")
.bind(account_id)
.bind("2026-06-18T00:00:00Z")
.bind("2026-06-18T00:00:00Z")
.execute(pool)
.await
.unwrap();
sqlx::query("insert into account_cookies (id, account_id, domain, name, value, updated_at) values (?, ?, ?, ?, ?, ?)")
.bind("cookie_lifecycle")
.bind(account_id)
@@ -180,12 +170,6 @@ async fn assert_account_related_rows_deleted(pool: &SqlitePool, account_id: &str
.fetch_one(pool)
.await
.unwrap();
let model_routes: i64 =
sqlx::query_scalar("select count(*) from model_account_routes where account_id = ?")
.bind(account_id)
.fetch_one(pool)
.await
.unwrap();
let leases: i64 =
sqlx::query_scalar("select count(*) from account_refresh_leases where account_id = ?")
.bind(account_id)
@@ -202,7 +186,6 @@ async fn assert_account_related_rows_deleted(pool: &SqlitePool, account_id: &str
assert_eq!(accounts, 0);
assert_eq!(usage, 0);
assert_eq!(model_usage, 0);
assert_eq!(model_routes, 0);
assert_eq!(cookies, 0);
assert_eq!(leases, 0);
assert_eq!(affinities, 0);
backend/tests/admin/settings/mod.rs
-32
View File
@@ -21,13 +21,6 @@ fn settings_service_should_apply_retained_settings_patch() {
]
.into(),
),
model_account_routes: Some(
[(
" gpt-5.5 ".to_string(),
vec![" acct_a ".to_string(), "acct_b".to_string()],
)]
.into(),
),
refresh_margin_seconds: Some(600),
refresh_concurrency: Some(3),
max_concurrent_per_account: Some(5),
@@ -39,10 +32,6 @@ fn settings_service_should_apply_retained_settings_patch() {
assert_eq!(current.model_aliases["claude-sonnet"], "gpt-5.5");
assert_eq!(current.model_aliases["openai-fast"], "openai:gpt-4o");
assert_eq!(
current.model_account_routes["gpt-5.5"],
vec!["acct_a".to_string(), "acct_b".to_string()]
);
assert_eq!(current.refresh_margin_seconds, 600);
assert_eq!(current.refresh_concurrency, 3);
assert_eq!(current.rotation_strategy, "round_robin");
@@ -50,27 +39,6 @@ fn settings_service_should_apply_retained_settings_patch() {
assert_eq!(current.request_interval_ms, 125);
}
#[test]
fn settings_service_should_reject_invalid_model_account_routes() {
let mut current = AdminSettings::default();
let error = SettingsService::apply_patch(
&mut current,
AdminSettingsPatch {
model_account_routes: Some(
[(
"gpt-5.5".to_string(),
vec!["acct_a".to_string(), "acct_a".to_string()],
)]
.into(),
),
..AdminSettingsPatch::default()
},
)
.expect_err("duplicate model route account should be rejected");
assert_eq!(error.field(), "modelAccountRoutes");
}
#[test]
fn settings_service_should_reject_invalid_patch_values() {
let mut current = AdminSettings::default();
backend/tests/admin/settings/routes.rs
+1 -63
View File
@@ -66,7 +66,6 @@ async fn admin_settings_should_return_runtime_fields() {
.unwrap();
let body = response_json(response).await;
assert_eq!(body["data"]["modelAliases"]["codex-fast"], "gpt-5.5");
assert_eq!(body["data"]["modelAccountRoutes"], json!({}));
assert_eq!(body["data"]["refreshMarginSeconds"], 240);
assert_eq!(body["data"]["refreshConcurrency"], 2);
assert_eq!(body["data"]["rotationStrategy"], "least_used");
@@ -267,8 +266,6 @@ async fn admin_settings_update_should_persist_runtime_settings_to_database() {
let config = test_config(url);
let pool = connect_sqlite(&config.database.url).await.unwrap();
seed_admin_session(&pool, "session_1").await;
seed_account(&pool, "acct_route_a").await;
seed_account(&pool, "acct_route_b").await;
let stores = BackgroundTaskStores {
accounts: SqliteAccountStore::new(pool.clone()),
admin_sessions: SqliteAdminSessionStore::new(pool.clone()),
@@ -303,9 +300,6 @@ async fn admin_settings_update_should_persist_runtime_settings_to_database() {
"gpt-5.2": "gpt-5.5",
"claude-sonnet": "gpt-5.5"
},
"modelAccountRoutes": {
"gpt-5.5": ["acct_route_a", "acct_route_b"]
},
"refreshMarginSeconds": 900,
"refreshConcurrency": 4,
"maxConcurrentPerAccount": 7,
@@ -320,10 +314,6 @@ async fn admin_settings_update_should_persist_runtime_settings_to_database() {
assert_eq!(response.status(), StatusCode::OK);
let body = response_json(response).await;
assert_eq!(body["data"]["modelAliases"]["gpt-5.2"], "gpt-5.5");
assert_eq!(
body["data"]["modelAccountRoutes"]["gpt-5.5"][0],
"acct_route_a"
);
assert_eq!(body["data"]["refreshMarginSeconds"], 900);
assert_eq!(body["data"]["refreshConcurrency"], 4);
@@ -357,29 +347,12 @@ async fn admin_settings_update_should_persist_runtime_settings_to_database() {
assert_eq!(row.3, 7);
assert_eq!(row.4, 80);
assert_eq!(row.5, "round_robin");
let route_rows: Vec<(String, String, i64)> = sqlx::query_as(
"select model, account_id, priority from model_account_routes order by model, priority",
)
.fetch_all(&pool)
.await
.unwrap();
assert_eq!(
route_rows,
vec![
("gpt-5.5".to_string(), "acct_route_a".to_string(), 0),
("gpt-5.5".to_string(), "acct_route_b".to_string(), 1),
]
);
assert!(!dir.path().join("config.yaml").exists());
let restarted_config = RuntimeSettingsService::load_or_initialize_config(config.clone(), &pool)
.await
.unwrap();
assert_eq!(restarted_config.model_aliases["gpt-5.2"], "gpt-5.5");
assert_eq!(
restarted_config.model_account_routes["gpt-5.5"],
vec!["acct_route_a".to_string(), "acct_route_b".to_string()]
);
assert_eq!(restarted_config.auth.refresh_margin_seconds, 900);
assert_eq!(restarted_config.auth.refresh_concurrency, 4);
assert_eq!(restarted_config.auth.rotation_strategy, "round_robin");
@@ -433,10 +406,6 @@ async fn admin_settings_update_should_apply_runtime_services() {
"modelAliases": {
"runtime-alias": "gpt-5.5"
},
"modelAccountRoutes": {
"gpt-5.5": ["acct_runtime_b"],
"gpt-5.6": ["acct_runtime_a", "acct_runtime_b"]
},
"refreshMarginSeconds": 120,
"refreshConcurrency": 3,
"maxConcurrentPerAccount": 2,
@@ -466,15 +435,6 @@ async fn admin_settings_update_should_apply_runtime_services() {
assert_eq!(capacity.total_slots, 4);
assert_eq!(capacity.max_concurrent_per_account, 2);
let route_request = AccountAcquireRequest::new("gpt-5.5", Utc::now());
let routed = services
.account_pool
.acquire_with(&route_request)
.await
.unwrap();
assert_eq!(routed.account.id, "acct_runtime_b");
services.account_pool.release(&routed.account.id).await;
let rotation_request = AccountAcquireRequest::new("gpt-5.6", Utc::now());
let first = services
.account_pool
@@ -486,15 +446,13 @@ async fn admin_settings_update_should_apply_runtime_services() {
.acquire_with(&rotation_request)
.await
.unwrap();
assert_eq!(first.account.id, "acct_runtime_b");
assert_eq!(second.account.id, "acct_runtime_a");
assert_ne!(first.account.id, second.account.id);
let repeated = services
.account_pool
.acquire_with(&AccountAcquireRequest::new("gpt-5.5", Utc::now()))
.await
.unwrap();
assert_eq!(repeated.account.id, "acct_runtime_b");
let started_at = Instant::now();
services
.account_pool
@@ -574,26 +532,6 @@ async fn admin_settings_test_app(db_name: &str) -> (axum::Router, tempfile::Temp
)
}
async fn seed_account(pool: &SqlitePool, account_id: &str) {
sqlx::query(
r"
insert into accounts (
id,
access_token,
status,
added_at,
updated_at
) values (?, ?, 'active', ?, ?)",
)
.bind(account_id)
.bind("access-token")
.bind("2026-06-18T00:00:00Z")
.bind("2026-06-18T00:00:00Z")
.execute(pool)
.await
.unwrap();
}
async fn seed_schedulable_account(pool: &SqlitePool, account_id: &str) {
sqlx::query(
r"
backend/tests/infra/storage_schema/mod.rs
+2 -67
View File
@@ -10,7 +10,7 @@ async fn sqlite_schema_should_create_current_tables() {
let pool = connect_sqlite(&url).await.unwrap();
let rows: Vec<(String,)> = sqlx::query_as(
"select name from sqlite_master where type = 'table' and name in ('schema_migrations', 'admin_users', 'admin_sessions', 'client_api_keys', 'runtime_settings', 'accounts', 'account_refresh_leases', 'account_usage', 'account_model_usage', 'usage_time_buckets', 'model_account_routes', 'account_cookies', 'fingerprints', 'fingerprint_update_history', 'usage_records', 'model_plan_snapshots', 'session_affinities') order by name",
"select name from sqlite_master where type = 'table' and name in ('schema_migrations', 'admin_users', 'admin_sessions', 'client_api_keys', 'runtime_settings', 'accounts', 'account_refresh_leases', 'account_usage', 'account_model_usage', 'usage_time_buckets', 'account_cookies', 'fingerprints', 'fingerprint_update_history', 'usage_records', 'model_plan_snapshots', 'session_affinities') order by name",
)
.fetch_all(&pool)
.await
@@ -28,7 +28,6 @@ async fn sqlite_schema_should_create_current_tables() {
"client_api_keys",
"fingerprint_update_history",
"fingerprints",
"model_account_routes",
"model_plan_snapshots",
"runtime_settings",
"schema_migrations",
@@ -156,7 +155,7 @@ async fn sqlite_schema_should_persist_diagnostic_filter_indexes() {
let pool = connect_sqlite(&url).await.unwrap();
let rows: Vec<(String,)> = sqlx::query_as(
"select name from sqlite_master where type = 'index' and name in ('idx_account_cookies_expires', 'idx_usage_records_level_created', 'idx_usage_records_model_created', 'idx_usage_records_route_created', 'idx_usage_records_status_created', 'idx_usage_records_upstream_status_created', 'idx_session_affinities_active_order', 'idx_model_account_routes_account', 'idx_model_account_routes_enabled_model', 'idx_usage_time_buckets_bucket', 'idx_usage_time_buckets_model_bucket') order by name",
"select name from sqlite_master where type = 'index' and name in ('idx_account_cookies_expires', 'idx_usage_records_level_created', 'idx_usage_records_model_created', 'idx_usage_records_route_created', 'idx_usage_records_status_created', 'idx_usage_records_upstream_status_created', 'idx_session_affinities_active_order', 'idx_usage_time_buckets_bucket', 'idx_usage_time_buckets_model_bucket') order by name",
)
.fetch_all(&pool)
.await
@@ -166,8 +165,6 @@ async fn sqlite_schema_should_persist_diagnostic_filter_indexes() {
rows.into_iter().map(|row| row.0).collect::<Vec<_>>(),
[
"idx_account_cookies_expires",
"idx_model_account_routes_account",
"idx_model_account_routes_enabled_model",
"idx_session_affinities_active_order",
"idx_usage_records_level_created",
"idx_usage_records_model_created",
@@ -246,33 +243,6 @@ async fn sqlite_schema_should_persist_usage_time_bucket_columns() {
);
}
#[tokio::test]
async fn sqlite_schema_should_persist_model_account_routes_columns() {
let dir = tempfile::tempdir().unwrap();
let db = dir.path().join("model-account-routes.sqlite");
let url = format!("sqlite://{}", db.display());
let pool = connect_sqlite(&url).await.unwrap();
let rows: Vec<(String,)> = sqlx::query_as(
"select name from pragma_table_info('model_account_routes') where name in ('model', 'account_id', 'priority', 'enabled', 'created_at', 'updated_at') order by name",
)
.fetch_all(&pool)
.await
.unwrap();
assert_eq!(
rows.into_iter().map(|row| row.0).collect::<Vec<_>>(),
[
"account_id",
"created_at",
"enabled",
"model",
"priority",
"updated_at",
]
);
}
#[tokio::test]
async fn sqlite_schema_should_persist_plain_account_cookie_value_column() {
let dir = tempfile::tempdir().unwrap();
@@ -553,22 +523,9 @@ async fn sqlite_schema_should_reject_non_boolean_flags() {
)
.execute(&pool)
.await;
sqlx::query(
"insert into accounts (id, access_token, status, added_at, updated_at) values ('acct_route', 'access-token', 'active', '2026-06-14T00:00:00Z', '2026-06-14T00:00:00Z')",
)
.execute(&pool)
.await
.unwrap();
let model_route_result = sqlx::query(
"insert into model_account_routes (model, account_id, enabled, created_at, updated_at) values ('gpt-5.5', 'acct_route', 2, '2026-06-14T00:00:00Z', '2026-06-14T00:00:00Z')",
)
.execute(&pool)
.await;
assert!(api_key_result.is_err());
assert!(account_result.is_err());
assert!(quota_verify_result.is_err());
assert!(model_route_result.is_err());
}
#[tokio::test]
@@ -599,28 +556,6 @@ async fn sqlite_schema_should_reject_invalid_runtime_settings() {
assert!(invalid_strategy_result.is_err());
}
#[tokio::test]
async fn sqlite_schema_should_reject_invalid_model_account_route_priority() {
let dir = tempfile::tempdir().unwrap();
let db = dir.path().join("model-route-check.sqlite");
let url = format!("sqlite://{}", db.display());
let pool = connect_sqlite(&url).await.unwrap();
sqlx::query(
"insert into accounts (id, access_token, status, added_at, updated_at) values ('acct_route', 'access-token', 'active', '2026-06-14T00:00:00Z', '2026-06-14T00:00:00Z')",
)
.execute(&pool)
.await
.unwrap();
let result = sqlx::query(
"insert into model_account_routes (model, account_id, priority, created_at, updated_at) values ('gpt-5.5', 'acct_route', -1, '2026-06-14T00:00:00Z', '2026-06-14T00:00:00Z')",
)
.execute(&pool)
.await;
assert!(result.is_err());
}
#[tokio::test]
async fn sqlite_schema_should_reject_negative_usage_time_bucket_counts() {
let dir = tempfile::tempdir().unwrap();
backend/tests/proxy/dispatch/chat_upstream/mod.rs
+2 -2
View File
@@ -27,8 +27,8 @@ use codex_proxy_rs::{
runtime::services::{BackgroundTaskStores, Services, UsageRecordOptions},
runtime::state::AppState,
upstream::accounts::{
cookies::SqliteCookieStore, model::AccountStatus, store::SqliteAccountStore,
token_refresh::RefreshLeaseStore,
cookies::SqliteCookieStore, model::AccountStatus, pool::AccountAcquireRequest,
store::SqliteAccountStore, token_refresh::RefreshLeaseStore,
},
upstream::fingerprint::FingerprintRepository,
};
backend/tests/proxy/dispatch/chat_upstream/responses_websocket.rs
+131
View File
@@ -1580,6 +1580,106 @@ async fn responses_websocket_non_stream_previous_response_not_found_should_strip
assert!(second_payload.get("previous_response_id").is_none());
}
#[tokio::test]
async fn responses_websocket_history_recovery_should_stay_on_fallback_account_when_preferred_is_full(
) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let base_url = format!("http://{}", listener.local_addr().unwrap());
let upstream = tokio::spawn(async move {
let first_payload = accept_websocket_response_with_authorization_and_message(
&listener,
"Bearer access-secondary",
WEBSOCKET_PREVIOUS_RESPONSE_NOT_FOUND.trim().to_string(),
)
.await;
let second_payload = accept_websocket_response_with_authorization_and_message(
&listener,
"Bearer access-secondary",
websocket_completed_response("resp_after_fallback_history_strip", 3, 1),
)
.await;
(first_payload, second_payload)
});
let (app, state, api_key, _pool, _dir) =
test_app_with_two_accounts_and_state_config(base_url, |config| {
config.auth.max_concurrent_per_account = 1;
})
.await;
seed_primary_affinity_and_hold_primary_slot(&state).await;
let response = app
.oneshot(responses_json_request(
&api_key,
&json!({
"model": "gpt-5.5",
"input": [],
"stream": false,
"previous_response_id": "resp_owned_by_primary"
}),
))
.await
.unwrap();
let body = response_json(response).await;
let (first_payload, second_payload) = upstream.await.unwrap();
assert_eq!(body["id"], "resp_after_fallback_history_strip");
assert_eq!(
first_payload["previous_response_id"],
"resp_owned_by_primary"
);
assert!(second_payload.get("previous_response_id").is_none());
}
#[tokio::test]
async fn responses_websocket_stream_history_recovery_should_stay_on_fallback_account_when_preferred_is_full(
) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let base_url = format!("http://{}", listener.local_addr().unwrap());
let upstream = tokio::spawn(async move {
let first_payload = accept_websocket_response_with_authorization_and_message(
&listener,
"Bearer access-secondary",
WEBSOCKET_PREVIOUS_RESPONSE_NOT_FOUND.trim().to_string(),
)
.await;
let second_payload = accept_websocket_response_with_authorization_and_message(
&listener,
"Bearer access-secondary",
websocket_completed_response("resp_stream_after_fallback_history_strip", 3, 1),
)
.await;
(first_payload, second_payload)
});
let (app, state, api_key, _pool, _dir) =
test_app_with_two_accounts_and_state_config(base_url, |config| {
config.auth.max_concurrent_per_account = 1;
})
.await;
seed_primary_affinity_and_hold_primary_slot(&state).await;
let response = app
.oneshot(responses_json_request(
&api_key,
&json!({
"model": "gpt-5.5",
"input": [],
"stream": true,
"previous_response_id": "resp_owned_by_primary"
}),
))
.await
.unwrap();
let body = response_text(response).await;
let (first_payload, second_payload) = upstream.await.unwrap();
assert!(body.contains("resp_stream_after_fallback_history_strip"));
assert_eq!(
first_payload["previous_response_id"],
"resp_owned_by_primary"
);
assert!(second_payload.get("previous_response_id").is_none());
}
#[tokio::test]
async fn responses_websocket_stream_previous_response_not_found_should_strip_history_and_retry_same_account(
) {
@@ -2264,3 +2364,34 @@ async fn responses_stream_with_previous_response_id_should_record_websocket_audi
Some(44.0)
);
}
async fn seed_primary_affinity_and_hold_primary_slot(state: &AppState) {
state
.services
.session_affinity
.record(
"resp_owned_by_primary".to_string(),
SessionAffinityEntry {
account_id: "acct_primary".to_string(),
conversation_id: "conv_primary".to_string(),
turn_state: Some("turn-primary".to_string()),
instructions_hash: None,
input_tokens: None,
function_call_ids: Vec::new(),
variant_hash: None,
created_at: Utc::now(),
},
)
.await
.unwrap();
let held_primary_slot = state
.services
.account_pool
.acquire_with(
&AccountAcquireRequest::new("gpt-5.5", Utc::now())
.with_required_account_id("acct_primary"),
)
.await
.unwrap();
assert_eq!(held_primary_slot.account.id, "acct_primary");
}
backend/tests/support/config.rs
+1 -4
View File
@@ -1,5 +1,3 @@
use std::collections::BTreeMap;
use codex_proxy_rs::config::schema::{
AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, FingerprintConfig,
LoggingConfig, QuotaConfig, ServerConfig, TlsConfig, WebSocketPoolConfig,
@@ -15,8 +13,7 @@ pub(crate) fn test_config(database_url: String) -> AppConfig {
api: ApiConfig {
base_url: "https://chatgpt.com/backend-api".to_string(),
},
model_aliases: BTreeMap::new(),
model_account_routes: BTreeMap::new(),
model_aliases: std::collections::BTreeMap::new(),
auth: AuthConfig {
refresh_margin_seconds: 300,
refresh_enabled: true,
backend/tests/upstream/accounts/account_pool/selection.rs
+27 -22
View File
@@ -445,28 +445,6 @@ fn account_pool_distinct_plan_accounts_should_filter_like_model_refresh() {
assert_eq!(pool.release("plus-ok").unwrap().model, None);
}
#[test]
fn account_pool_should_filter_by_model_account_routes() {
let mut model_routes = BTreeMap::new();
model_routes.insert("gpt-5.5".to_string(), vec!["acct_b".to_string()]);
let mut pool = AccountPool::with_options(AccountPoolOptions {
max_concurrent_per_account: 1,
model_account_routes: model_routes,
..AccountPoolOptions::default()
});
pool.insert(crate::support::accounts::test_account(
"acct_a",
AccountStatus::Active,
));
pool.insert(crate::support::accounts::test_account(
"acct_b",
AccountStatus::Active,
));
assert_eq!(acquire_account(&mut pool, "gpt-5.5").unwrap().id, "acct_b");
assert!(acquire_account(&mut pool, "gpt-5.5").is_none());
}
#[test]
fn account_pool_should_exclude_requested_account_ids() {
let mut pool = AccountPool::default();
@@ -511,6 +489,33 @@ fn account_pool_should_prefer_session_affinity_account_when_available() {
assert_eq!(acquired.account.id, "acct_b");
}
#[test]
fn account_pool_should_not_fallback_when_required_account_is_unavailable() {
let mut pool = AccountPool::with_options(AccountPoolOptions {
max_concurrent_per_account: 1,
..AccountPoolOptions::default()
});
pool.insert(crate::support::accounts::test_account(
"acct_a",
AccountStatus::Active,
));
pool.insert(crate::support::accounts::test_account(
"acct_b",
AccountStatus::Active,
));
assert!(pool
.acquire_with(
&AccountAcquireRequest::new("gpt-5.5", fixed_time()).with_required_account_id("acct_a"),
)
.is_some());
let acquired = pool.acquire_with(
&AccountAcquireRequest::new("gpt-5.5", fixed_time()).with_required_account_id("acct_a"),
);
assert!(acquired.is_none());
}
#[test]
fn account_pool_should_cleanup_stale_slots_before_acquire() {
let now = fixed_time();
frontend/src/views/api-keys/components/ApiKeyUseModal.vue
+1 -1
View File
@@ -49,7 +49,7 @@ windows_wsl_setup_acknowledged = true
[model_providers.OpenAI]
name = "OpenAI"
base_url = "${props.serviceRootUrl}"
base_url = "${props.serviceRootUrl}/v1"
wire_api = "responses"
requires_openai_auth = true