diff --git a/src/account.rs b/src/account.rs index 10a04d9..517f2ab 100644 --- a/src/account.rs +++ b/src/account.rs @@ -53,9 +53,32 @@ pub fn add_api_key(key: String, base_url: Option, switch: bool) -> Resul Ok(()) } -pub fn list_accounts(json: bool) -> Result<()> { +pub async fn list_accounts(json: bool, force: bool) -> Result<()> { let home = paths::codex_home(None)?; - let store = Store::load(&home)?; + let mut store = Store::load(&home)?; + let quota_concurrency = store.settings.quota_concurrency; + let quota_max_age_seconds = if force { 0 } else { 60 }; + let ids = store + .accounts + .iter() + .map(|account| account.id.clone()) + .collect::>(); + let quota_report = crate::quota::refresh_stale_quotas( + &mut store, + &ids, + quota_max_age_seconds, + quota_concurrency, + ) + .await; + if !quota_report.errors.is_empty() { + eprintln!( + "提示: {} 个账号配额刷新失败,已使用本地缓存。", + quota_report.errors.len() + ); + } + if quota_report.changed { + store.save(&home)?; + } if json { println!("{}", serde_json::to_string_pretty(&store.accounts)?); return Ok(()); diff --git a/src/cli.rs b/src/cli.rs index 26f8056..589202f 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -19,6 +19,8 @@ pub enum Commands { List { #[arg(long)] json: bool, + #[arg(short, long)] + force: bool, }, /// Switch Codex auth.json to a saved account. Switch { @@ -124,6 +126,8 @@ pub enum AccountCommands { List { #[arg(long)] json: bool, + #[arg(short, long)] + force: bool, }, Current { #[arg(long)] diff --git a/src/config_store.rs b/src/config_store.rs index 705f301..d9882eb 100644 --- a/src/config_store.rs +++ b/src/config_store.rs @@ -24,6 +24,8 @@ pub struct Store { #[serde(default)] pub sync: SyncConfig, #[serde(default)] + pub settings: Settings, + #[serde(default)] pub accounts: Vec, } @@ -102,6 +104,20 @@ pub struct Quota { pub updated_at: i64, } +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Settings { + #[serde(default = "default_quota_concurrency")] + pub quota_concurrency: usize, +} + +impl Default for Settings { + fn default() -> Self { + Self { + quota_concurrency: default_quota_concurrency(), + } + } +} + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Home { pub name: String, @@ -223,6 +239,7 @@ impl Store { homes: Vec::new(), server: ServerConfig::default(), sync: SyncConfig::default(), + settings: Settings::default(), accounts: Vec::new(), }; store.ensure_default_home(codex_home); @@ -248,3 +265,7 @@ pub fn path_to_string(path: &Path) -> String { fn default_version() -> u32 { 1 } + +fn default_quota_concurrency() -> usize { + 4 +} diff --git a/src/main.rs b/src/main.rs index 4adbd3c..b0dd49f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -29,7 +29,10 @@ use crate::cli::{ #[tokio::main] async fn main() -> Result<()> { let cli = Cli::parse(); - match cli.command.unwrap_or(Commands::List { json: false }) { + match cli.command.unwrap_or(Commands::List { + json: false, + force: false, + }) { Commands::Login(args) => match args.command { Some(LoginCommands::Oauth { manual, @@ -46,7 +49,7 @@ async fn main() -> Result<()> { }) => account::import_auth(file, codex_home, switch), None => account::import_auth(args.file, args.codex_home, args.switch), }, - Commands::List { json } => account::list_accounts(json), + Commands::List { json, force } => account::list_accounts(json, force).await, Commands::Switch { account, codex_home, @@ -64,7 +67,7 @@ async fn main() -> Result<()> { Commands::Quota { account, all, json } => quota::quota_command(account, all, json).await, Commands::RefreshToken { account } => token::refresh_token_command(&account).await, Commands::Account { command } => match command { - AccountCommands::List { json } => account::list_accounts(json), + AccountCommands::List { json, force } => account::list_accounts(json, force).await, AccountCommands::Current { json } => account::current_account(json), AccountCommands::Show { account, json } => account::show_account(&account, json), AccountCommands::Remove { account } => account::remove_account(&account), diff --git a/src/quota.rs b/src/quota.rs index 968d6fb..b1ae150 100644 --- a/src/quota.rs +++ b/src/quota.rs @@ -42,9 +42,22 @@ struct QuotaDisplay<'a> { error: Option, } +pub struct QuotaRefreshReport { + pub errors: Vec<(String, String)>, + pub changed: bool, +} + +#[derive(Clone)] +struct QuotaFetchInput { + id: String, + access_token: String, + account_id: Option, +} + pub async fn quota_command(account: Option, all: bool, json: bool) -> Result<()> { let home = crate::paths::codex_home(None)?; let mut store = Store::load(&home)?; + let concurrency = store.settings.quota_concurrency; // Selection order matches CLI expectations: explicit --all, explicit // account, current account, then the first saved account. let ids = if all { @@ -74,14 +87,11 @@ pub async fn quota_command(account: Option, all: bool, json: bool) -> Re return Err(anyhow!("没有可查询的账号")); } - let mut errors = Vec::<(String, String)>::new(); - for id in &ids { - match refresh_one_quota(&mut store, id).await { - Ok(()) => {} - Err(error) => errors.push((id.clone(), error.to_string())), - } + let report = refresh_quotas(&mut store, &ids, None, concurrency).await; + let errors = report.errors; + if report.changed { + store.save(&home)?; } - store.save(&home)?; if json { let rows = ids @@ -138,9 +148,139 @@ pub async fn quota_command(account: Option, all: bool, json: bool) -> Re Ok(()) } -async fn refresh_one_quota(store: &mut Store, account_id: &str) -> Result<()> { +pub async fn refresh_stale_quotas( + store: &mut Store, + ids: &[String], + max_age_seconds: i64, + concurrency: usize, +) -> QuotaRefreshReport { + refresh_quotas(store, ids, Some(max_age_seconds), concurrency).await +} + +async fn refresh_quotas( + store: &mut Store, + ids: &[String], + max_age_seconds: Option, + concurrency: usize, +) -> QuotaRefreshReport { + let mut inputs = Vec::new(); + let mut errors = Vec::<(String, String)>::new(); + let mut changed = false; + + for id in ids { + match prepare_quota_fetch(store, id, max_age_seconds).await { + Ok((Some(input), token_changed)) => { + changed |= token_changed; + inputs.push(input); + } + Ok((None, token_changed)) => changed |= token_changed, + Err(error) => { + changed = true; + errors.push((id.clone(), error.to_string())); + } + } + } + + let mut retry_ids = Vec::new(); + for (id, result) in fetch_quotas_concurrently(inputs, concurrency).await { + match result { + Ok(quota) => { + apply_quota_result(store, &id, quota); + changed = true; + } + Err(error) if should_retry_with_refresh(&error.to_string()) => retry_ids.push(id), + Err(error) => errors.push((id, error.to_string())), + } + } + + for id in retry_ids { + match refresh_one_quota(store, &id).await { + Ok(()) => changed = true, + Err(error) => errors.push((id, error.to_string())), + } + } + + QuotaRefreshReport { errors, changed } +} + +async fn prepare_quota_fetch( + store: &mut Store, + account_id: &str, + max_age_seconds: Option, +) -> Result<(Option, bool)> { // Quota requests need a fresh access token and, when available, the // ChatGPT-Account-Id header for multi-account organizations. + let account = store + .find_account(account_id) + .ok_or_else(|| anyhow!("账号不存在: {account_id}"))?; + if account.auth_mode != AuthMode::Oauth { + return if max_age_seconds.is_some() { + Ok((None, false)) + } else { + Err(anyhow!("API Key 账号不支持 Codex OAuth 配额查询")) + }; + } + if let Some(max_age_seconds) = max_age_seconds { + let now = Utc::now().timestamp(); + let fresh = account + .quota + .as_ref() + .map(|quota| now - quota.updated_at < max_age_seconds) + .unwrap_or(false); + if fresh { + return Ok((None, false)); + } + } + + let token_changed = crate::token::refresh_account_if_needed(store, account_id).await?; + let account = store + .find_account(account_id) + .ok_or_else(|| anyhow!("账号不存在: {account_id}"))? + .clone(); + let tokens = account + .tokens + .as_ref() + .ok_or_else(|| anyhow!("OAuth 账号缺少 tokens: {account_id}"))?; + Ok(( + Some(QuotaFetchInput { + id: account_id.to_string(), + access_token: tokens.access_token.clone(), + account_id: account.account_id.clone(), + }), + token_changed, + )) +} + +async fn fetch_quotas_concurrently( + inputs: Vec, + concurrency: usize, +) -> Vec<(String, Result)> { + let concurrency = concurrency.max(1); + let mut results = Vec::new(); + for chunk in inputs.chunks(concurrency) { + let mut handles = Vec::new(); + for input in chunk.iter().cloned() { + handles.push(tokio::spawn(async move { + let id = input.id; + let result = + fetch_quota(input.access_token.as_str(), input.account_id.as_deref()).await; + (id, result) + })); + } + for handle in handles { + match handle.await { + Ok(result) => results.push(result), + Err(error) => results.push(( + "".to_string(), + Err(anyhow!("quota task join failed: {error}")), + )), + } + } + } + results +} + +async fn refresh_one_quota(store: &mut Store, account_id: &str) -> Result<()> { crate::token::refresh_account_if_needed(store, account_id).await?; let account = store .find_account(account_id) @@ -174,15 +314,19 @@ async fn refresh_one_quota(store: &mut Store, account_id: &str) -> Result<()> { Err(error) => return Err(error), }; + apply_quota_result(store, account_id, quota); + Ok(()) +} + +fn apply_quota_result(store: &mut Store, account_id: &str, quota: FetchQuotaResult) { let account = store .find_account_mut(account_id) - .ok_or_else(|| anyhow!("账号不存在: {account_id}"))?; + .expect("quota result references an existing account"); if let Some(plan) = quota.plan_type { account.plan_type = Some(plan); } account.quota = Some(quota.quota); account.updated_at = Utc::now().timestamp(); - Ok(()) } struct FetchQuotaResult {