feat: enhance account listing with quota refresh and force option

This commit is contained in:
2026-04-30 10:07:09 +08:00
parent 7120a8a904
commit 01d3a09d1e
5 changed files with 210 additions and 15 deletions
+25 -2
View File
@@ -53,9 +53,32 @@ pub fn add_api_key(key: String, base_url: Option<String>, 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::<Vec<_>>();
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(());
+4
View File
@@ -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)]
+21
View File
@@ -24,6 +24,8 @@ pub struct Store {
#[serde(default)]
pub sync: SyncConfig,
#[serde(default)]
pub settings: Settings,
#[serde(default)]
pub accounts: Vec<Account>,
}
@@ -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
}
+6 -3
View File
@@ -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),
+154 -10
View File
@@ -42,9 +42,22 @@ struct QuotaDisplay<'a> {
error: Option<String>,
}
pub struct QuotaRefreshReport {
pub errors: Vec<(String, String)>,
pub changed: bool,
}
#[derive(Clone)]
struct QuotaFetchInput {
id: String,
access_token: String,
account_id: Option<String>,
}
pub async fn quota_command(account: Option<String>, 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<String>, 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<String>, 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<i64>,
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<i64>,
) -> Result<(Option<QuotaFetchInput>, 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<QuotaFetchInput>,
concurrency: usize,
) -> Vec<(String, Result<FetchQuotaResult>)> {
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((
"<unknown>".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 {