feat: enhance account listing with quota refresh and force option
This commit is contained in:
+25
-2
@@ -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(());
|
||||
|
||||
@@ -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)]
|
||||
|
||||
@@ -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
@@ -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
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user