From dac29696e36cf2ef2f3dddad4a79d4762ae7ede6 Mon Sep 17 00:00:00 2001 From: zyycn Date: Sun, 14 Jun 2026 04:07:55 +0800 Subject: [PATCH] fix: complete account lifecycle Co-authored-by: Codex --- src/admin/http/accounts/create.rs | 9 ++ src/admin/http/usage.rs | 4 + src/codex/accounts/jwt.rs | 71 ++++++++++++++++ src/codex/accounts/mod.rs | 1 + src/codex/accounts/model.rs | 2 + src/codex/accounts/pool.rs | 95 +++++++++++++++++++++- src/codex/accounts/repository.rs | 23 ++++-- src/codex/accounts/service/import.rs | 17 +++- src/codex/accounts/service/mod.rs | 1 + src/codex/accounts/service/runtime_pool.rs | 1 + src/codex/serving/dispatch/mod.rs | 6 +- src/codex/serving/dispatch/usage.rs | 22 +++++ src/codex/serving/responses.rs | 40 +++++++++ src/platform/storage/db.rs | 20 ++++- src/platform/storage/schema.sql | 1 + tests/admin/accounts/cookies_quota.rs | 17 ++-- tests/admin/usage_stats_route.rs | 8 +- tests/codex_accounts/pool_scheduling.rs | 56 +++++++++++++ tests/codex_accounts/repository.rs | 4 + tests/codex_serving/chat_completions.rs | 1 + tests/codex_serving/responses_http_sse.rs | 45 ++++++++++ tests/runtime/startup.rs | 1 + tests/support/admin_accounts.rs | 1 + tests/support/upstream.rs | 2 + 24 files changed, 429 insertions(+), 19 deletions(-) create mode 100644 src/codex/accounts/jwt.rs diff --git a/src/admin/http/accounts/create.rs b/src/admin/http/accounts/create.rs index f84b90f42..62479b0fb 100644 --- a/src/admin/http/accounts/create.rs +++ b/src/admin/http/accounts/create.rs @@ -417,6 +417,15 @@ fn account_entry_from_value(value: &Value, fallback_label: Option) -> Ac ), token, refresh_token, + access_token_expires_at: first_string( + value, + &[ + &["accessTokenExpiresAt"], + &["access_token_expires_at"], + &["credentials", "accessTokenExpiresAt"], + &["credentials", "access_token_expires_at"], + ], + ), status: first_string(value, &[&["status"]]), } } diff --git a/src/admin/http/usage.rs b/src/admin/http/usage.rs index f7a3ff289..46b85e827 100644 --- a/src/admin/http/usage.rs +++ b/src/admin/http/usage.rs @@ -33,6 +33,7 @@ pub struct AdminUsageStatsData { pub label: Option, pub plan_type: Option, pub request_count: i64, + pub empty_response_count: i64, pub input_tokens: i64, pub output_tokens: i64, pub cached_tokens: i64, @@ -47,6 +48,7 @@ impl From for AdminUsageStatsData { label: usage.label, plan_type: usage.plan_type, request_count: usage.request_count, + empty_response_count: usage.empty_response_count, input_tokens: usage.input_tokens, output_tokens: usage.output_tokens, cached_tokens: usage.cached_tokens, @@ -60,6 +62,7 @@ impl From for AdminUsageStatsData { pub struct AdminUsageStatsSummaryData { pub account_count: i64, pub request_count: i64, + pub empty_response_count: i64, pub input_tokens: i64, pub output_tokens: i64, pub cached_tokens: i64, @@ -70,6 +73,7 @@ impl From for AdminUsageStatsSummaryData { Self { account_count: summary.account_count, request_count: summary.request_count, + empty_response_count: summary.empty_response_count, input_tokens: summary.input_tokens, output_tokens: summary.output_tokens, cached_tokens: summary.cached_tokens, diff --git a/src/codex/accounts/jwt.rs b/src/codex/accounts/jwt.rs new file mode 100644 index 000000000..c90200e49 --- /dev/null +++ b/src/codex/accounts/jwt.rs @@ -0,0 +1,71 @@ +use base64::Engine; +use chrono::{DateTime, Utc}; +use serde_json::Value; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum JwtExpiry { + Expired, + Valid, + MissingOrInvalid, +} + +pub fn jwt_expiry(token: &str, now: DateTime) -> JwtExpiry { + let Some(exp) = jwt_exp(token) else { + return JwtExpiry::MissingOrInvalid; + }; + if now.timestamp() >= exp { + JwtExpiry::Expired + } else { + JwtExpiry::Valid + } +} + +fn jwt_exp(token: &str) -> Option { + let mut parts = token.split('.'); + let _header = parts.next()?; + let payload = parts.next()?; + let _signature = parts.next()?; + if parts.next().is_some() { + return None; + } + + let decoded = base64::engine::general_purpose::URL_SAFE_NO_PAD + .decode(payload) + .ok()?; + let value = serde_json::from_slice::(&decoded).ok()?; + value.get("exp")?.as_i64() +} + +#[cfg(test)] +mod tests { + use base64::Engine; + use chrono::{TimeZone, Utc}; + use serde_json::json; + + use super::{jwt_expiry, JwtExpiry}; + + #[test] + fn jwt_expiry_should_detect_expired_valid_and_invalid_tokens() { + let now = Utc.with_ymd_and_hms(2026, 6, 14, 8, 0, 0).unwrap(); + + assert_eq!( + jwt_expiry(&test_jwt(now.timestamp() - 1), now), + JwtExpiry::Expired + ); + assert_eq!( + jwt_expiry(&test_jwt(now.timestamp() + 1), now), + JwtExpiry::Valid + ); + assert_eq!(jwt_expiry("not-a-jwt", now), JwtExpiry::MissingOrInvalid); + } + + fn test_jwt(exp: i64) -> String { + let header = json!({"alg": "none", "typ": "JWT"}); + let payload = json!({ "exp": exp }); + format!("{}.{}.", jwt_part(&header), jwt_part(&payload)) + } + + fn jwt_part(value: &serde_json::Value) -> String { + base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(serde_json::to_vec(value).unwrap()) + } +} diff --git a/src/codex/accounts/mod.rs b/src/codex/accounts/mod.rs index c0e6adfbc..36f4f5ff0 100644 --- a/src/codex/accounts/mod.rs +++ b/src/codex/accounts/mod.rs @@ -1,5 +1,6 @@ pub mod cf_path_block; pub mod cookies; +pub mod jwt; pub mod lifecycle; pub mod model; pub mod models; diff --git a/src/codex/accounts/model.rs b/src/codex/accounts/model.rs index abb07da8e..32f67768a 100644 --- a/src/codex/accounts/model.rs +++ b/src/codex/accounts/model.rs @@ -27,6 +27,7 @@ pub struct Account { pub quota_cooldown_until: Option>, pub cloudflare_cooldown_until: Option>, pub request_count: u64, + pub empty_response_count: u64, pub window_request_count: u64, pub window_started_at: Option>, pub window_reset_at: Option>, @@ -52,6 +53,7 @@ impl Account { quota_cooldown_until: None, cloudflare_cooldown_until: None, request_count: 0, + empty_response_count: 0, window_request_count: 0, window_started_at: None, window_reset_at: None, diff --git a/src/codex/accounts/pool.rs b/src/codex/accounts/pool.rs index 3e426243e..1a833eb97 100644 --- a/src/codex/accounts/pool.rs +++ b/src/codex/accounts/pool.rs @@ -5,7 +5,10 @@ use std::{ use chrono::{DateTime, Duration, Utc}; -use crate::codex::accounts::model::{Account, AccountStatus}; +use crate::codex::accounts::{ + jwt::{jwt_expiry, JwtExpiry}, + model::{Account, AccountStatus}, +}; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum RotationStrategy { @@ -155,6 +158,7 @@ impl AccountPool { }; account.last_used_at = None; account.request_count = 0; + account.empty_response_count = 0; account.window_request_count = 0; account.window_started_at = None; account.window_reset_at = None; @@ -202,6 +206,7 @@ impl AccountPool { pub fn acquire_with(&mut self, request: AccountAcquireRequest) -> Option { self.cleanup_stale_slots(request.now); + self.refresh_account_statuses(request.now); let candidates = self.candidates(&request); let selected = if let Some(preferred_account_id) = &request.preferred_account_id { candidates @@ -240,6 +245,7 @@ impl AccountPool { pub fn capacity_summary(&mut self, now: DateTime) -> AccountCapacitySummary { self.cleanup_stale_slots(now); + self.refresh_account_statuses(now); let active_accounts = self .accounts .values() @@ -386,6 +392,29 @@ impl AccountPool { }); } + fn refresh_account_statuses(&mut self, now: DateTime) { + let account_ids = self.accounts.keys().cloned().collect::>(); + for account_id in account_ids { + self.refresh_account_status(&account_id, now); + } + } + + fn refresh_account_status(&mut self, account_id: &str, now: DateTime) { + let mut should_clear_slots = false; + if let Some(account) = self.accounts.get_mut(account_id) { + if account.status == AccountStatus::Active && access_token_expired(account, now) { + account.status = AccountStatus::Expired; + should_clear_slots = true; + } else { + refresh_quota_window(account, now); + refresh_cloudflare_cooldown(account, now); + } + } + if should_clear_slots { + self.slots.remove(account_id); + } + } + fn is_model_allowed(&self, account: &Account, model: &str) -> bool { let Some(allowed_plans) = self.options.model_plan_allowlist.get(model) else { return true; @@ -405,6 +434,70 @@ impl AccountPool { } } +fn access_token_expired(account: &Account, now: DateTime) -> bool { + if account + .access_token_expires_at + .is_some_and(|expires_at| now >= expires_at) + { + return true; + } + + match jwt_expiry(&account.access_token, now) { + JwtExpiry::Expired => true, + JwtExpiry::Valid => false, + JwtExpiry::MissingOrInvalid => account.access_token_expires_at.is_none(), + } +} + +fn refresh_quota_window(account: &mut Account, now: DateTime) { + let window_expired = account + .window_reset_at + .is_some_and(|reset_at| now >= reset_at); + let cooldown_expired = account + .quota_cooldown_until + .is_some_and(|cooldown_until| now >= cooldown_until); + + if window_expired { + account.window_request_count = 0; + account.window_started_at = Some(now); + account.window_reset_at = + next_window_reset_at(account.window_reset_at, account.limit_window_seconds, now); + } + + if account.quota_limit_reached && (cooldown_expired || window_expired) { + account.quota_limit_reached = false; + account.quota_cooldown_until = None; + } +} + +fn next_window_reset_at( + reset_at: Option>, + limit_window_seconds: Option, + now: DateTime, +) -> Option> { + let reset_at = reset_at?; + let window_seconds = limit_window_seconds?; + if window_seconds == 0 { + return None; + } + + let elapsed_seconds = now.signed_duration_since(reset_at).num_seconds().max(0) as u64; + let windows_to_advance = elapsed_seconds / window_seconds + 1; + let advance_seconds = window_seconds + .saturating_mul(windows_to_advance) + .min(i64::MAX as u64); + Some(reset_at + Duration::seconds(advance_seconds as i64)) +} + +fn refresh_cloudflare_cooldown(account: &mut Account, now: DateTime) { + if account + .cloudflare_cooldown_until + .is_some_and(|cooldown_until| now >= cooldown_until) + { + account.cloudflare_cooldown_until = None; + } +} + #[derive(Debug, Clone, PartialEq, Eq)] struct LeastUsedGroupKey { quota_limited: bool, diff --git a/src/codex/accounts/repository.rs b/src/codex/accounts/repository.rs index e98589b74..18f97e61e 100644 --- a/src/codex/accounts/repository.rs +++ b/src/codex/accounts/repository.rs @@ -93,12 +93,14 @@ pub struct UsageDelta { pub input_tokens: i64, pub output_tokens: i64, pub cached_tokens: i64, + pub empty_response_count: i64, } #[derive(Debug, Clone, PartialEq, Eq)] pub struct AccountUsageRecord { pub account_id: String, pub request_count: i64, + pub empty_response_count: i64, pub input_tokens: i64, pub output_tokens: i64, pub cached_tokens: i64, @@ -112,6 +114,7 @@ pub struct AccountUsageListRecord { pub label: Option, pub plan_type: Option, pub request_count: i64, + pub empty_response_count: i64, pub input_tokens: i64, pub output_tokens: i64, pub cached_tokens: i64, @@ -130,6 +133,7 @@ pub struct AccountQuotaSnapshot { pub struct AccountUsageSummary { pub account_count: i64, pub request_count: i64, + pub empty_response_count: i64, pub input_tokens: i64, pub output_tokens: i64, pub cached_tokens: i64, @@ -322,7 +326,7 @@ impl AccountRepository { pub async fn list_pool_accounts(&self) -> AccountRepositoryResult> { let rows = sqlx::query( - "select accounts.id, email, accounts.account_id, user_id, label, plan_type, access_token_cipher, refresh_token_cipher, access_token_expires_at, status, added_at, updated_at, quota_limit_reached, quota_cooldown_until, cloudflare_cooldown_until, coalesce(account_usage.request_count, 0) as usage_request_count, account_usage.last_used_at as usage_last_used_at from accounts left join account_usage on account_usage.account_id = accounts.id order by added_at desc, accounts.id desc", + "select accounts.id, email, accounts.account_id, user_id, label, plan_type, access_token_cipher, refresh_token_cipher, access_token_expires_at, status, added_at, updated_at, quota_limit_reached, quota_cooldown_until, cloudflare_cooldown_until, coalesce(account_usage.request_count, 0) as usage_request_count, coalesce(account_usage.empty_response_count, 0) as usage_empty_response_count, account_usage.last_used_at as usage_last_used_at from accounts left join account_usage on account_usage.account_id = accounts.id order by added_at desc, accounts.id desc", ) .fetch_all(&self.pool) .await?; @@ -353,6 +357,7 @@ impl AccountRepository { row.get::, _>("cloudflare_cooldown_until"), )?, request_count: row.get::("usage_request_count").max(0) as u64, + empty_response_count: row.get::("usage_empty_response_count").max(0) as u64, window_request_count: 0, window_started_at: None, window_reset_at: None, @@ -511,9 +516,10 @@ impl AccountRepository { ) -> AccountRepositoryResult<()> { let now = Utc::now().to_rfc3339(); sqlx::query( - "insert into account_usage (account_id, request_count, input_tokens, output_tokens, cached_tokens, last_used_at) values (?, 1, ?, ?, ?, ?) on conflict(account_id) do update set request_count = request_count + 1, input_tokens = input_tokens + excluded.input_tokens, output_tokens = output_tokens + excluded.output_tokens, cached_tokens = cached_tokens + excluded.cached_tokens, last_used_at = excluded.last_used_at", + "insert into account_usage (account_id, request_count, empty_response_count, input_tokens, output_tokens, cached_tokens, last_used_at) values (?, 1, ?, ?, ?, ?, ?) on conflict(account_id) do update set request_count = request_count + 1, empty_response_count = empty_response_count + excluded.empty_response_count, input_tokens = input_tokens + excluded.input_tokens, output_tokens = output_tokens + excluded.output_tokens, cached_tokens = cached_tokens + excluded.cached_tokens, last_used_at = excluded.last_used_at", ) .bind(account_id) + .bind(usage.empty_response_count) .bind(usage.input_tokens) .bind(usage.output_tokens) .bind(usage.cached_tokens) @@ -528,7 +534,7 @@ impl AccountRepository { account_id: &str, ) -> AccountRepositoryResult> { let row = sqlx::query( - "select account_id, request_count, input_tokens, output_tokens, cached_tokens, last_used_at from account_usage where account_id = ?", + "select account_id, request_count, empty_response_count, input_tokens, output_tokens, cached_tokens, last_used_at from account_usage where account_id = ?", ) .bind(account_id) .fetch_optional(&self.pool) @@ -684,6 +690,7 @@ fn usage_from_row(row: &sqlx::sqlite::SqliteRow) -> AccountRepositoryResult AccountRepositoryResult { let row = sqlx::query( - "select count(*) as account_count, coalesce(sum(request_count), 0) as request_count, coalesce(sum(input_tokens), 0) as input_tokens, coalesce(sum(output_tokens), 0) as output_tokens, coalesce(sum(cached_tokens), 0) as cached_tokens from account_usage", + "select count(*) as account_count, coalesce(sum(request_count), 0) as request_count, coalesce(sum(empty_response_count), 0) as empty_response_count, coalesce(sum(input_tokens), 0) as input_tokens, coalesce(sum(output_tokens), 0) as output_tokens, coalesce(sum(cached_tokens), 0) as cached_tokens from account_usage", ) .fetch_one(&self.pool) .await?; Ok(AccountUsageSummary { account_count: row.get("account_count"), request_count: row.get("request_count"), + empty_response_count: row.get("empty_response_count"), input_tokens: row.get("input_tokens"), output_tokens: row.get("output_tokens"), cached_tokens: row.get("cached_tokens"), @@ -778,7 +786,7 @@ impl AccountUsageRepository { pub async fn reset_account(&self, account_id: &str) -> AccountRepositoryResult<()> { sqlx::query( - "insert into account_usage (account_id, request_count, input_tokens, output_tokens, cached_tokens, last_used_at) values (?, 0, 0, 0, 0, null) on conflict(account_id) do update set request_count = 0, input_tokens = 0, output_tokens = 0, cached_tokens = 0, last_used_at = null", + "insert into account_usage (account_id, request_count, empty_response_count, input_tokens, output_tokens, cached_tokens, last_used_at) values (?, 0, 0, 0, 0, 0, null) on conflict(account_id) do update set request_count = 0, empty_response_count = 0, input_tokens = 0, output_tokens = 0, cached_tokens = 0, last_used_at = null", ) .bind(account_id) .execute(&self.pool) @@ -796,6 +804,7 @@ fn usage_list_from_row( label: row.get("label"), plan_type: row.get("plan_type"), request_count: row.get("request_count"), + empty_response_count: row.get("empty_response_count"), input_tokens: row.get("input_tokens"), output_tokens: row.get("output_tokens"), cached_tokens: row.get("cached_tokens"), diff --git a/src/codex/accounts/service/import.rs b/src/codex/accounts/service/import.rs index aab2af7a6..b0f843074 100644 --- a/src/codex/accounts/service/import.rs +++ b/src/codex/accounts/service/import.rs @@ -176,6 +176,12 @@ impl AccountService { let label = empty_to_none(entry.label); let plan_type = empty_to_none(entry.plan_type); let refresh_token = empty_to_none(entry.refresh_token); + let access_token_expires_at = entry + .access_token_expires_at + .as_deref() + .map(parse_import_datetime) + .transpose() + .map_err(StoreImportAccountError::Invalid)?; let access_token = access_token.to_string(); let now = Utc::now().to_rfc3339(); let pool_account = Account { @@ -187,12 +193,13 @@ impl AccountService { plan_type: plan_type.clone(), access_token: access_token.clone(), refresh_token: refresh_token.clone(), - access_token_expires_at: None, + access_token_expires_at, status, quota_limit_reached: false, quota_cooldown_until: None, cloudflare_cooldown_until: None, request_count: 0, + empty_response_count: 0, window_request_count: 0, window_started_at: None, window_reset_at: None, @@ -209,7 +216,7 @@ impl AccountService { plan_type, access_token: SecretString::new(access_token.into()), refresh_token: refresh_token.map(|token| SecretString::new(token.into())), - access_token_expires_at: None, + access_token_expires_at, status, }; repo.insert(account) @@ -283,6 +290,12 @@ pub(super) fn empty_to_none(value: Option) -> Option { .filter(|value| !value.is_empty()) } +fn parse_import_datetime(value: &str) -> Result, String> { + DateTime::parse_from_rfc3339(value) + .map(|value| value.with_timezone(&Utc)) + .map_err(|_| "Invalid accessTokenExpiresAt".to_string()) +} + pub(super) fn normalized_account_id(id: Option) -> String { id.and_then(|id| empty_to_none(Some(id))) .unwrap_or_else(|| format!("acct_{}", Uuid::new_v4().simple())) diff --git a/src/codex/accounts/service/mod.rs b/src/codex/accounts/service/mod.rs index b66d719bc..a7cab1424 100644 --- a/src/codex/accounts/service/mod.rs +++ b/src/codex/accounts/service/mod.rs @@ -93,6 +93,7 @@ pub struct AccountImportEntry { pub plan_type: Option, pub token: Option, pub refresh_token: Option, + pub access_token_expires_at: Option, pub status: Option, } diff --git a/src/codex/accounts/service/runtime_pool.rs b/src/codex/accounts/service/runtime_pool.rs index d7e631fd8..19f59796a 100644 --- a/src/codex/accounts/service/runtime_pool.rs +++ b/src/codex/accounts/service/runtime_pool.rs @@ -55,6 +55,7 @@ pub(super) fn pool_account_from_stored(account: StoredAccount) -> Account { quota_cooldown_until: None, cloudflare_cooldown_until: None, request_count: 0, + empty_response_count: 0, window_request_count: 0, window_started_at: None, window_reset_at: None, diff --git a/src/codex/serving/dispatch/mod.rs b/src/codex/serving/dispatch/mod.rs index 84677d347..bc1d5bf6a 100644 --- a/src/codex/serving/dispatch/mod.rs +++ b/src/codex/serving/dispatch/mod.rs @@ -65,7 +65,7 @@ use self::{ refresh::refresh_account_after_unauthorized, routing::request_domain, stream::{completed_response_metadata, ensure_stream_metadata, responses_sse_failure}, - usage::{record_request_attempt, record_usage_with_deps}, + usage::{record_empty_response_with_deps, record_request_attempt, record_usage_with_deps}, }; #[derive(Clone)] @@ -239,6 +239,10 @@ impl CodexUpstreamService { record_usage_with_deps(&self.deps, account_id, usage).await } + pub(crate) async fn record_empty_response(&self, account_id: &str) -> Result<(), ()> { + record_empty_response_with_deps(&self.deps, account_id).await + } + pub(crate) async fn record_response_affinity( &self, request: &CodexResponsesRequest, diff --git a/src/codex/serving/dispatch/usage.rs b/src/codex/serving/dispatch/usage.rs index 16e4ac216..da74c9b98 100644 --- a/src/codex/serving/dispatch/usage.rs +++ b/src/codex/serving/dispatch/usage.rs @@ -18,6 +18,7 @@ pub(super) async fn record_usage_with_deps( input_tokens: u64_to_i64_saturating(usage.input_tokens), output_tokens: u64_to_i64_saturating(usage.output_tokens), cached_tokens: u64_to_i64_saturating(usage.cached_tokens), + empty_response_count: 0, }, ) .await @@ -37,6 +38,27 @@ pub(super) async fn record_request_attempt( input_tokens: 0, output_tokens: 0, cached_tokens: 0, + empty_response_count: 0, + }, + ) + .await + .map_err(|_| ()) +} + +pub(super) async fn record_empty_response_with_deps( + deps: &CodexUpstreamDependencies, + account_id: &str, +) -> Result<(), ()> { + let Some(repo) = deps.account_repository.as_ref() else { + return Ok(()); + }; + repo.record_usage( + account_id, + UsageDelta { + input_tokens: 0, + output_tokens: 0, + cached_tokens: 0, + empty_response_count: 1, }, ) .await diff --git a/src/codex/serving/responses.rs b/src/codex/serving/responses.rs index 9692118fe..d59d3a752 100644 --- a/src/codex/serving/responses.rs +++ b/src/codex/serving/responses.rs @@ -168,6 +168,26 @@ impl ResponsesService { Ok(response) => { // 检测空响应 if let Ok(CollectedResponse::Empty) = completed_response_json(&response.body) { + if self + .upstream + .record_empty_response(&account.id) + .await + .is_err() + { + self.upstream + .log_response( + &log_context, + StatusCode::OK, + EventLevel::Warn, + "v1 responses 记录空响应计数失败", + json!({ + "stream": false, + "emptyResponse": true, + "emptyResponseStoreError": true, + }), + ) + .await; + } empty_response_retries += 1; if empty_response_retries <= MAX_EMPTY_RETRIES { self.upstream @@ -338,6 +358,26 @@ impl ResponsesService { (StatusCode::OK, Json(body)).into_response() } Ok(CollectedResponse::Empty) => { + if self + .upstream + .record_empty_response(&account.id) + .await + .is_err() + { + self.upstream + .log_response( + &log_context, + StatusCode::OK, + EventLevel::Warn, + "v1 responses 记录空响应计数失败", + json!({ + "stream": false, + "emptyResponse": true, + "emptyResponseStoreError": true, + }), + ) + .await; + } self.upstream .log_response( &log_context, diff --git a/src/platform/storage/db.rs b/src/platform/storage/db.rs index da24dd544..4cbc509af 100644 --- a/src/platform/storage/db.rs +++ b/src/platform/storage/db.rs @@ -1,6 +1,6 @@ use std::{fs, path::Path, str::FromStr, time::Duration}; -use sqlx::{sqlite::SqliteConnectOptions, SqlitePool}; +use sqlx::{sqlite::SqliteConnectOptions, Row, SqlitePool}; pub async fn connect_sqlite(database_url: &str) -> Result { ensure_sqlite_parent_dir(database_url)?; @@ -20,6 +20,24 @@ pub async fn connect_sqlite(database_url: &str) -> Result Result<(), sqlx::Error> { let schema = include_str!("schema.sql"); sqlx::raw_sql(schema).execute(pool).await?; + ensure_account_usage_schema(pool).await?; + Ok(()) +} + +async fn ensure_account_usage_schema(pool: &SqlitePool) -> Result<(), sqlx::Error> { + let rows = sqlx::query("pragma table_info(account_usage)") + .fetch_all(pool) + .await?; + let exists = rows + .iter() + .any(|row| row.get::("name") == "empty_response_count"); + if !exists { + sqlx::query( + "alter table account_usage add column empty_response_count integer not null default 0", + ) + .execute(pool) + .await?; + } Ok(()) } diff --git a/src/platform/storage/schema.sql b/src/platform/storage/schema.sql index 8a33fa1e2..b632b8bf5 100644 --- a/src/platform/storage/schema.sql +++ b/src/platform/storage/schema.sql @@ -82,6 +82,7 @@ create index if not exists idx_account_refresh_leases_expires on account_refresh create table if not exists account_usage ( account_id text primary key references accounts(id) on delete cascade, request_count integer not null default 0, + empty_response_count integer not null default 0, input_tokens integer not null default 0, output_tokens integer not null default 0, cached_tokens integer not null default 0, diff --git a/tests/admin/accounts/cookies_quota.rs b/tests/admin/accounts/cookies_quota.rs index 113c3f650..aff4fab6f 100644 --- a/tests/admin/accounts/cookies_quota.rs +++ b/tests/admin/accounts/cookies_quota.rs @@ -17,8 +17,8 @@ use codex_proxy_rs::codex::gateway::oauth::{RefreshFailure, TokenPair}; use crate::support::{ admin_accounts::{ admin_accounts_test_app, admin_accounts_test_app_with_base_url, - admin_accounts_test_app_with_refresher, import_test_account, FailingTokenRefresher, - StaticTokenRefresher, + admin_accounts_test_app_with_refresher, import_test_account, test_jwt, + FailingTokenRefresher, StaticTokenRefresher, }, response_json, }; @@ -150,12 +150,19 @@ async fn admin_account_cookies_should_require_existing_account_and_non_empty_coo #[tokio::test] async fn admin_account_refresh_should_update_tokens_and_runtime_pool_without_returning_secrets() { + let refreshed_access_token = test_jwt( + Some("chatgpt-account"), + Some("chatgpt-user"), + Some("refresh@example.com"), + Some("plus"), + 3600, + ); let (app, state, pool, _dir) = admin_accounts_test_app_with_refresher( "admin-account-refresh.sqlite", 28, StaticTokenRefresher { result: Ok(TokenPair { - access_token: "new-admin-refresh-access".to_string(), + access_token: refreshed_access_token.clone(), refresh_token: Some("new-admin-refresh-rt".to_string()), }), calls: Arc::new(Mutex::new(Vec::new())), @@ -183,7 +190,7 @@ async fn admin_account_refresh_should_update_tokens_and_runtime_pool_without_ret assert_eq!(body["data"]["result"], "alive"); assert_eq!(body["data"]["previousStatus"], "active"); let serialized = serde_json::to_string(&body).unwrap(); - assert!(!serialized.contains("new-admin-refresh-access")); + assert!(!serialized.contains(&refreshed_access_token)); assert!(!serialized.contains("new-admin-refresh-rt")); let stored: (String, String) = sqlx::query_as( @@ -203,7 +210,7 @@ async fn admin_account_refresh_should_update_tokens_and_runtime_pool_without_ret .acquire_runtime_account("gpt-5.5") .await .unwrap(); - assert_eq!(acquired.access_token, "new-admin-refresh-access"); + assert_eq!(acquired.access_token, refreshed_access_token); } #[tokio::test] diff --git a/tests/admin/usage_stats_route.rs b/tests/admin/usage_stats_route.rs index b6febbbdd..1579a9ee2 100644 --- a/tests/admin/usage_stats_route.rs +++ b/tests/admin/usage_stats_route.rs @@ -110,6 +110,7 @@ async fn admin_usage_stats_should_cursor_page_account_usage() { assert_eq!(first_body["data"].as_array().unwrap().len(), 1); assert_eq!(first_body["data"][0]["accountId"], "acct_b"); assert_eq!(first_body["data"][0]["requestCount"], 2); + assert_eq!(first_body["data"][0]["emptyResponseCount"], 1); assert_eq!(first_body["data"][0]["inputTokens"], 7); assert_eq!(first_body["page"]["limit"], 1); let cursor = first_body["page"]["nextCursor"].as_str().unwrap(); @@ -160,6 +161,7 @@ async fn admin_usage_stats_summary_should_return_usage_totals() { assert_eq!(body["code"], 200); assert_eq!(body["data"]["accountCount"], 2); assert_eq!(body["data"]["requestCount"], 5); + assert_eq!(body["data"]["emptyResponseCount"], 1); assert_eq!(body["data"]["inputTokens"], 19); assert_eq!(body["data"]["outputTokens"], 8); assert_eq!(body["data"]["cachedTokens"], 3); @@ -212,10 +214,11 @@ async fn seed_usage(pool: &sqlx::SqlitePool) { .unwrap(); } sqlx::query( - "insert into account_usage (account_id, request_count, input_tokens, output_tokens, cached_tokens, last_used_at) values (?, ?, ?, ?, ?, ?)", + "insert into account_usage (account_id, request_count, empty_response_count, input_tokens, output_tokens, cached_tokens, last_used_at) values (?, ?, ?, ?, ?, ?, ?)", ) .bind("acct_a") .bind(3_i64) + .bind(0_i64) .bind(12_i64) .bind(5_i64) .bind(1_i64) @@ -224,10 +227,11 @@ async fn seed_usage(pool: &sqlx::SqlitePool) { .await .unwrap(); sqlx::query( - "insert into account_usage (account_id, request_count, input_tokens, output_tokens, cached_tokens, last_used_at) values (?, ?, ?, ?, ?, ?)", + "insert into account_usage (account_id, request_count, empty_response_count, input_tokens, output_tokens, cached_tokens, last_used_at) values (?, ?, ?, ?, ?, ?, ?)", ) .bind("acct_b") .bind(2_i64) + .bind(1_i64) .bind(7_i64) .bind(3_i64) .bind(2_i64) diff --git a/tests/codex_accounts/pool_scheduling.rs b/tests/codex_accounts/pool_scheduling.rs index b0dfe35f0..bdda5b8d5 100644 --- a/tests/codex_accounts/pool_scheduling.rs +++ b/tests/codex_accounts/pool_scheduling.rs @@ -6,6 +6,8 @@ use codex_proxy_rs::codex::accounts::{ pool::{AccountAcquireRequest, AccountPool, AccountPoolOptions, RotationStrategy}, }; +use crate::support::admin_accounts::test_jwt; + #[test] fn account_pool_should_respect_max_concurrent_slots_per_account() { let mut pool = AccountPool::with_options(AccountPoolOptions { @@ -76,6 +78,60 @@ fn account_pool_should_reuse_quota_limited_accounts_after_cooldown() { ); } +#[test] +fn acquire_should_skip_accounts_with_expired_token_metadata() { + let now = fixed_time(); + let mut expired = Account::test("expired", AccountStatus::Active); + expired.access_token_expires_at = Some(now - Duration::seconds(1)); + let mut pool = AccountPool::default(); + pool.insert(expired); + + let acquired = pool.acquire_with(AccountAcquireRequest::new("gpt-5.5", now)); + + assert!(acquired.is_none()); +} + +#[test] +fn acquire_should_skip_accounts_with_expired_jwt_when_metadata_is_missing() { + let now = Utc::now(); + let mut expired = Account::test("expired", AccountStatus::Active); + expired.access_token = test_jwt(None, None, None, None, -60); + expired.access_token_expires_at = None; + let mut pool = AccountPool::default(); + pool.insert(expired); + + let acquired = pool.acquire_with(AccountAcquireRequest::new("gpt-5.5", now)); + + assert!(acquired.is_none()); +} + +#[test] +fn acquire_should_refresh_expired_cooldowns_before_selecting_account() { + let now = fixed_time(); + let mut account = Account::test("acct_a", AccountStatus::Active); + account.quota_limit_reached = true; + account.quota_cooldown_until = Some(now - Duration::seconds(1)); + account.cloudflare_cooldown_until = Some(now - Duration::seconds(1)); + account.window_request_count = 7; + account.window_reset_at = Some(now - Duration::seconds(1)); + account.limit_window_seconds = Some(60); + let mut pool = AccountPool::default(); + pool.insert(account); + + let acquired = pool + .acquire_with(AccountAcquireRequest::new("gpt-5.5", now)) + .unwrap(); + + assert!(!acquired.account.quota_limit_reached); + assert!(acquired.account.quota_cooldown_until.is_none()); + assert!(acquired.account.cloudflare_cooldown_until.is_none()); + assert_eq!(acquired.account.window_request_count, 1); + assert!(acquired + .account + .window_reset_at + .is_some_and(|reset| reset > now)); +} + #[test] fn account_pool_should_prefer_configured_tier_priority() { let mut pool = AccountPool::with_options(AccountPoolOptions { diff --git a/tests/codex_accounts/repository.rs b/tests/codex_accounts/repository.rs index d5daa0b8d..cf4789416 100644 --- a/tests/codex_accounts/repository.rs +++ b/tests/codex_accounts/repository.rs @@ -167,6 +167,7 @@ async fn account_repository_should_accumulate_usage_counters() { input_tokens: 10, output_tokens: 4, cached_tokens: 3, + empty_response_count: 0, }, ) .await @@ -177,6 +178,7 @@ async fn account_repository_should_accumulate_usage_counters() { input_tokens: 8, output_tokens: 2, cached_tokens: 1, + empty_response_count: 1, }, ) .await @@ -187,6 +189,7 @@ async fn account_repository_should_accumulate_usage_counters() { assert_eq!(usage.input_tokens, 18); assert_eq!(usage.output_tokens, 6); assert_eq!(usage.cached_tokens, 4); + assert_eq!(usage.empty_response_count, 1); assert!(usage.last_used_at.is_some()); } @@ -218,6 +221,7 @@ async fn account_repository_should_load_usage_count_into_runtime_pool_accounts() input_tokens: 1, output_tokens: 2, cached_tokens: 3, + empty_response_count: 0, }, ) .await diff --git a/tests/codex_serving/chat_completions.rs b/tests/codex_serving/chat_completions.rs index 05282e89a..5b6465847 100644 --- a/tests/codex_serving/chat_completions.rs +++ b/tests/codex_serving/chat_completions.rs @@ -148,6 +148,7 @@ async fn build_imported_app_with_accounts( "accountId": account.account_id, "token": account.token, "refreshToken": account.refresh_token, + "accessTokenExpiresAt": "2999-01-01T00:00:00Z", "status": "active" }) }) diff --git a/tests/codex_serving/responses_http_sse.rs b/tests/codex_serving/responses_http_sse.rs index 9972e194a..96456576d 100644 --- a/tests/codex_serving/responses_http_sse.rs +++ b/tests/codex_serving/responses_http_sse.rs @@ -31,6 +31,7 @@ const DONE_ITEM_COMPLETED_SSE: &str = include_str!("../fixtures/responses/http_sse/done_item_completed.sse"); const STREAM_USAGE_SSE: &str = include_str!("../fixtures/responses/http_sse/stream_usage.sse"); const DEFAULT_STREAM_SSE: &str = include_str!("../fixtures/responses/http_sse/default_stream.sse"); +const EMPTY_COMPLETED_SSE: &str = "event: response.completed\ndata: {\"response\":{\"id\":\"resp_empty\",\"object\":\"response\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":0}}}\n\n"; #[tokio::test] async fn v1_responses_should_honor_explicit_http_sse_transport() { @@ -143,6 +144,50 @@ async fn v1_responses_should_use_imported_account_and_record_usage() { assert_eq!(event.4["usage"]["inputTokens"], 7); } +#[tokio::test] +async fn v1_responses_should_record_empty_response_attempts() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(EMPTY_COMPLETED_SSE), + ) + .mount(&server) + .await; + let imported = build_imported_app(server.uri()).await; + + let response = imported + .app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header( + "authorization", + format!("Bearer {}", imported.client_api_key), + ) + .header("content-type", "application/json") + .body(Body::from( + r#"{"model":"gpt-5.5","input":[],"stream":false}"#, + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::BAD_GATEWAY); + let usage: (i64, i64, i64, i64) = sqlx::query_as( + "select request_count, empty_response_count, input_tokens, output_tokens from account_usage where account_id = ?", + ) + .bind("acct_imported") + .fetch_one(&imported.pool) + .await + .unwrap(); + assert_eq!(usage, (3, 3, 0, 0)); +} + #[tokio::test] async fn v1_responses_should_forward_parity_fields_and_context_headers() { let server = MockServer::start().await; diff --git a/tests/runtime/startup.rs b/tests/runtime/startup.rs index 6f34cf6e7..91a13a7dc 100644 --- a/tests/runtime/startup.rs +++ b/tests/runtime/startup.rs @@ -106,6 +106,7 @@ async fn app_state_should_restore_account_pool_from_sqlite_accounts() { input_tokens: 10, output_tokens: 3, cached_tokens: 2, + empty_response_count: 0, }, ) .await diff --git a/tests/support/admin_accounts.rs b/tests/support/admin_accounts.rs index 67062a1d4..a1a875fe5 100644 --- a/tests/support/admin_accounts.rs +++ b/tests/support/admin_accounts.rs @@ -273,6 +273,7 @@ pub async fn import_test_account(app: &Router, session_id: &str, account_id: &st "planType": "plus", "token": format!("access-{account_id}"), "refreshToken": format!("refresh-{account_id}"), + "accessTokenExpiresAt": "2999-01-01T00:00:00Z", "status": "active" }] }); diff --git a/tests/support/upstream.rs b/tests/support/upstream.rs index 0ab014864..b9c3ebca7 100644 --- a/tests/support/upstream.rs +++ b/tests/support/upstream.rs @@ -167,6 +167,7 @@ pub async fn build_imported_app_with_accounts( "accountId": account.account_id, "token": account.token, "refreshToken": account.refresh_token, + "accessTokenExpiresAt": "2999-01-01T00:00:00Z", "status": "active" }) }) @@ -267,6 +268,7 @@ where "accountId": account.account_id, "token": account.token, "refreshToken": account.refresh_token, + "accessTokenExpiresAt": "2999-01-01T00:00:00Z", "status": "active" }) })