use std::collections::{HashMap, HashSet};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, RwLock as StdRwLock};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use alex_core::Provider;
use anyhow::{anyhow, bail, Context, Result};
use base64::Engine;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use tokio::sync::{Mutex, RwLock};
pub mod login;
pub mod sessions;
pub const ANTHROPIC_CLIENT_ID: &str = "9d1c250a-e61b-44d9-88ed-5944d1962f5e";
pub const ANTHROPIC_TOKEN_URL: &str = "https://console.anthropic.com/v1/oauth/token";
pub const OPENAI_CLIENT_ID: &str = "app_EMoamEEZ73f0CkXaXp7hrann";
pub const OPENAI_TOKEN_URL: &str = "https://auth.openai.com/oauth/token";
pub const XAI_CLIENT_ID: &str = "b1a00492-073a-47ea-816f-4c329264a828";
pub const XAI_TOKEN_URL: &str = "https://auth.x.ai/oauth2/token";
pub const GEMINI_CLIENT_ID: &str =
"681255809395-oo8ft2oprdrnp9e3aqf6av3hmdib135j.apps.googleusercontent.com";
pub const GOOGLE_TOKEN_URL: &str = "https://oauth2.googleapis.com/token";
const ANTHROPIC_PROFILE_URL: &str = "https://api.anthropic.com/api/oauth/profile";
const XAI_USERINFO_URL: &str = "https://auth.x.ai/oauth2/userinfo";
pub fn gemini_client_secret() -> String {
["GOCSPX", "4uHgMPm", "1o7Sk", "geV6Cu5clXFsxl"].join("-")
}
const REFRESH_MARGIN_MS: i64 = 120_000;
pub fn now_ms() -> i64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_millis() as i64
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Account {
pub id: String,
pub provider: Provider,
pub kind: String,
#[serde(default = "default_account_name")]
pub name: String,
#[serde(default)]
pub description: Option<String>,
#[serde(default)]
pub paused: bool,
#[serde(default)]
pub label: Option<String>,
#[serde(default)]
pub access_token: Option<String>,
#[serde(default)]
pub refresh_token: Option<String>,
#[serde(default)]
pub id_token: Option<String>,
#[serde(default)]
pub api_key: Option<String>,
#[serde(default)]
pub expires_at_ms: Option<i64>,
#[serde(default)]
pub last_refresh_ms: Option<i64>,
#[serde(default)]
pub account_meta: Value,
#[serde(default)]
pub cooldown_until_ms: Option<i64>,
#[serde(default = "default_status")]
pub status: String,
#[serde(skip)]
pub path: Option<PathBuf>,
}
fn default_account_name() -> String {
"default".to_string()
}
fn default_status() -> String {
"active".to_string()
}
impl Account {
pub fn chatgpt_account_id(&self) -> Option<String> {
self.account_meta
.get("account_id")
.and_then(|v| v.as_str())
.map(String::from)
}
pub fn email(&self) -> Option<String> {
fn payload_email(token: &str) -> Option<String> {
let payload = crate::login::jwt_payload(token)?;
payload
.get("email")
.and_then(Value::as_str)
.and_then(normalize_email)
.or_else(|| {
payload
.get("https://api.openai.com/profile")
.and_then(|profile| profile.get("email"))
.and_then(Value::as_str)
.and_then(normalize_email)
})
}
self.account_meta
.get("email")
.and_then(Value::as_str)
.and_then(normalize_email)
.or_else(|| self.description.as_deref().and_then(normalize_email))
.or_else(|| {
(self.provider != Provider::Xai)
.then(|| self.id_token.as_deref().and_then(payload_email))
.flatten()
})
.or_else(|| {
(self.provider != Provider::Xai)
.then(|| self.access_token.as_deref().and_then(payload_email))
.flatten()
})
.or_else(|| {
self.label.as_deref().and_then(|label| {
label
.split(['(', ')', ' ', '·'])
.find_map(normalize_email)
})
})
}
fn needs_refresh(&self) -> bool {
self.kind == "oauth"
&& match self.expires_at_ms {
Some(exp) => exp < now_ms() + REFRESH_MARGIN_MS,
None => true,
}
}
}
pub(crate) fn normalize_email(value: &str) -> Option<String> {
let value = value.trim();
(value.contains('@') && !value.chars().any(char::is_whitespace))
.then(|| value.to_ascii_lowercase())
}
pub(crate) fn persist_account_email(account: &mut Account, email: &str) {
let Some(email) = normalize_email(email) else { return };
if !account.account_meta.is_object() {
account.account_meta = json!({});
}
account.account_meta["email"] = json!(email);
if account.description.is_none()
|| account
.description
.as_deref()
.and_then(normalize_email)
.is_some()
{
account.description = Some(email);
}
}
pub(crate) async fn fetch_provider_email(
http: &reqwest::Client,
provider: Provider,
access_token: &str,
) -> Option<String> {
let mut request = match provider {
Provider::Anthropic => http.get(ANTHROPIC_PROFILE_URL),
Provider::Xai => http.get(XAI_USERINFO_URL),
_ => return None,
}
.bearer_auth(access_token)
.timeout(Duration::from_secs(5));
if provider == Provider::Anthropic {
request = request
.header("content-type", "application/json")
.header("cache-control", "no-cache");
}
let response = request.send().await.ok()?;
if !response.status().is_success() {
return None;
}
let profile: Value = response.json().await.ok()?;
let email = match provider {
Provider::Anthropic => profile["account"]["email"].as_str(),
Provider::Xai if profile["email_verified"].as_bool() == Some(false) => None,
Provider::Xai => profile["email"].as_str(),
_ => None,
}?;
normalize_email(email)
}
fn oauth_rank(account: &Account, prefer_oauth: bool) -> u8 {
if (account.kind == "oauth") == prefer_oauth {
0
} else {
1
}
}
fn policy_rank(policy: &AccountPolicy, name: &str) -> usize {
policy.order.iter().position(|n| n == name).unwrap_or(usize::MAX / 2)
}
fn utilization_pct(account: &Account) -> Option<u8> {
account.account_meta
.get("rate_limit_pct")
.or_else(|| account.account_meta.get("utilization_pct"))
.and_then(|v| v.as_u64())
.map(|v| v.min(100) as u8)
}
fn codex_limit_snapshot(account: &Account) -> Option<&Value> {
account.account_meta.get("codex_limits")
}
fn codex_window_values(account: &Account) -> impl Iterator<Item = &Value> {
codex_limit_snapshot(account)
.and_then(|snapshot| snapshot.get("windows"))
.and_then(Value::as_array)
.into_iter()
.flatten()
}
pub fn codex_reserve_pct(account: &Account, policy: &AccountPolicy) -> u8 {
policy
.account_reserve_pct
.get(&account.name)
.or_else(|| policy.account_reserve_pct.get(&account.id))
.copied()
.or(policy.reserve_pct)
.unwrap_or(10)
.min(100)
}
pub fn codex_reserve_blocked(account: &Account, reserve_pct: u8, now_s: i64) -> bool {
if reserve_pct == 0 {
return false;
}
codex_window_values(account).any(|window| {
let future_window = window
.get("resets_at_s")
.and_then(Value::as_i64)
.map(|reset| reset > now_s)
.unwrap_or(true);
let used_pct = window.get("used_pct").and_then(Value::as_f64);
future_window && used_pct.map(|used| used >= 100.0 - reserve_pct as f64).unwrap_or(false)
})
}
pub fn codex_reset_selection(account: &Account, now_s: i64) -> Option<Value> {
let mut selected: Option<(f64, i64, Option<String>)> = None;
for window in codex_window_values(account) {
let Some(used_pct) = window.get("used_pct").and_then(Value::as_f64) else {
continue;
};
let Some(resets_at_s) = window.get("resets_at_s").and_then(Value::as_i64) else {
continue;
};
if resets_at_s <= now_s {
continue;
}
let label = window.get("window").and_then(Value::as_str).map(String::from);
let replace = selected
.as_ref()
.map(|(current_used, current_reset, _)| {
used_pct > *current_used
|| (used_pct == *current_used && resets_at_s < *current_reset)
})
.unwrap_or(true);
if replace {
selected = Some((used_pct, resets_at_s, label));
}
}
selected.map(|(used_pct, resets_at_s, window)| {
json!({"window": window, "used_pct": used_pct, "resets_at_s": resets_at_s})
})
}
fn codex_binding_reset(account: &Account, now_s: i64) -> Option<i64> {
codex_reset_selection(account, now_s)?["resets_at_s"].as_i64()
}
fn account_proxy_eligible(account: &Account, policy: &AccountPolicy) -> bool {
!policy.disabled.iter().any(|name| name == &account.name || name == &account.id)
}
fn sort_by_policy(accounts: &mut Vec<Account>, policy: &AccountPolicy, prefer_oauth: bool, rr: usize) {
let now_s = now_ms() / 1000;
match policy.mode {
AccountPolicyMode::RoundRobin => {
let _ = rr;
accounts.sort_by_key(|a| (
oauth_rank(a, prefer_oauth),
codex_reserve_blocked(a, codex_reserve_pct(a, policy), now_s),
policy_rank(policy, &a.name),
a.name.clone(),
a.id.clone(),
));
}
AccountPolicyMode::Threshold => {
let threshold = policy.threshold_pct.unwrap_or(80);
accounts.sort_by_key(|a| {
let over = utilization_pct(a).map(|p| p >= threshold).unwrap_or(false);
(oauth_rank(a, prefer_oauth), over, policy_rank(policy, &a.name), a.name.clone(), a.id.clone())
});
}
AccountPolicyMode::Priority => {
accounts.sort_by_key(|a| (
oauth_rank(a, prefer_oauth),
codex_reserve_blocked(a, codex_reserve_pct(a, policy), now_s),
policy_rank(policy, &a.name),
a.name.clone(),
a.id.clone(),
));
}
AccountPolicyMode::ResetFirst => {
accounts.sort_by_key(|a| (
oauth_rank(a, prefer_oauth),
codex_reserve_blocked(a, codex_reserve_pct(a, policy), now_s),
codex_binding_reset(a, now_s).unwrap_or(i64::MAX),
policy_rank(policy, &a.name),
a.name.clone(),
a.id.clone(),
));
}
}
}
pub fn rfc3339_to_ms(s: &str) -> Option<i64> {
chrono::DateTime::parse_from_rfc3339(s)
.ok()
.map(|dt| dt.timestamp_millis())
}
pub fn jwt_exp_ms(token: &str) -> Option<i64> {
let payload = token.split('.').nth(1)?;
let bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD
.decode(payload)
.ok()?;
let v: Value = serde_json::from_slice(&bytes).ok()?;
v.get("exp")?.as_i64().map(|s| s * 1000)
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum AccountPolicyMode {
Priority,
RoundRobin,
Threshold,
ResetFirst,
}
impl Default for AccountPolicyMode {
fn default() -> Self { Self::Priority }
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AccountPolicy {
#[serde(default)]
pub order: Vec<String>,
#[serde(default)]
pub mode: AccountPolicyMode,
#[serde(default)]
pub threshold_pct: Option<u8>,
#[serde(default)]
pub reserve_pct: Option<u8>,
#[serde(default)]
pub account_reserve_pct: HashMap<String, u8>,
#[serde(default = "default_allow_mid_thread_failover")]
pub allow_mid_thread_failover: bool,
#[serde(default)]
pub disabled: Vec<String>,
}
fn default_allow_mid_thread_failover() -> bool { true }
impl Default for AccountPolicy {
fn default() -> Self {
Self {
order: Vec::new(),
mode: AccountPolicyMode::default(),
threshold_pct: None,
reserve_pct: None,
account_reserve_pct: HashMap::new(),
allow_mid_thread_failover: true,
disabled: Vec::new(),
}
}
}
fn default_policy_for(provider: Provider) -> AccountPolicy {
if provider == Provider::Openai {
AccountPolicy {
mode: AccountPolicyMode::ResetFirst,
reserve_pct: Some(10),
..AccountPolicy::default()
}
} else {
AccountPolicy::default()
}
}
const ROUTING_POLICIES_FILE: &str = ".routing-policies";
pub struct Vault {
dir: PathBuf,
accounts: RwLock<HashMap<String, Account>>,
locks: Mutex<HashMap<String, Arc<Mutex<()>>>>,
rr_counter: AtomicUsize,
policies: StdRwLock<Vec<(Provider, AccountPolicy)>>,
http: reqwest::Client,
}
impl Vault {
pub fn open(dir: PathBuf) -> Result<Self> {
std::fs::create_dir_all(&dir)?;
let mut accounts = HashMap::new();
for entry in std::fs::read_dir(&dir)? {
let path = entry?.path();
if path.extension().map(|e| e == "json").unwrap_or(false) {
match std::fs::read_to_string(&path)
.map_err(anyhow::Error::from)
.and_then(|s| serde_json::from_str::<Account>(&s).map_err(Into::into))
{
Ok(mut acct) => {
if acct.name.is_empty() {
acct.name = "default".into();
}
acct.path = Some(path.clone());
accounts.insert(acct.id.clone(), acct);
}
Err(e) => tracing::warn!("skipping unreadable account file {path:?}: {e}"),
}
}
}
let policies = read_routing_policies(&dir).unwrap_or_else(|error| {
tracing::warn!(%error, "could not load persisted routing policies");
Vec::new()
});
Ok(Self {
dir,
accounts: RwLock::new(accounts),
locks: Mutex::new(HashMap::new()),
rr_counter: AtomicUsize::new(0),
policies: StdRwLock::new(policies),
http: reqwest::Client::new(),
})
}
pub async fn list(&self) -> Vec<Account> {
let mut v: Vec<Account> = self.accounts.read().await.values().cloned().collect();
v.sort_by(|a, b| (a.provider.as_str(), a.name.as_str(), a.kind.as_str()).cmp(&(b.provider.as_str(), b.name.as_str(), b.kind.as_str())));
v
}
pub fn set_policies_blocking(&self, policies: Vec<(Provider, AccountPolicy)>) {
*self.policies.write().expect("policy lock poisoned") = policies;
}
pub async fn set_policies(&self, policies: Vec<(Provider, AccountPolicy)>) {
self.set_policies_blocking(policies);
}
pub fn policy(&self, provider: Provider) -> AccountPolicy {
self.policies
.read()
.expect("policy lock poisoned")
.iter()
.find(|(candidate, _)| *candidate == provider)
.map(|(_, policy)| policy.clone())
.unwrap_or_else(|| default_policy_for(provider))
}
pub async fn set_policy_persisted(
&self,
provider: Provider,
policy: AccountPolicy,
) -> Result<()> {
let policies = {
let mut guard = self.policies.write().expect("policy lock poisoned");
if let Some((_, current)) = guard
.iter_mut()
.find(|(candidate, _)| *candidate == provider)
{
*current = policy;
} else {
guard.push((provider, policy));
}
guard.clone()
};
write_routing_policies(&self.dir, &policies)
}
pub async fn record_codex_limits(&self, id: &str, mut snapshot: Value) -> Result<()> {
let now = now_ms();
let mut account = self
.accounts
.read()
.await
.get(id)
.cloned()
.ok_or_else(|| anyhow!("unknown account {id}"))?;
if account.provider != Provider::Openai {
bail!("account {id} is not an OpenAI account");
}
if account
.account_meta
.get("codex_limits")
.and_then(|value| value.get("observed_at_ms"))
.and_then(Value::as_i64)
.map(|observed| now - observed < 30_000)
.unwrap_or(false)
{
return Ok(());
}
if let Some(object) = snapshot.as_object_mut() {
object.insert("observed_at_ms".into(), json!(now));
}
if !account.account_meta.is_object() {
account.account_meta = json!({});
}
account
.account_meta
.as_object_mut()
.expect("account_meta initialized as object")
.insert("codex_limits".into(), snapshot);
self.upsert(account).await
}
pub async fn pause(&self, provider: Provider, name: &str, paused: bool) -> Result<()> {
let account = self.accounts.read().await.values().find(|a| a.provider == provider && a.name == name).cloned().ok_or_else(|| anyhow!("unknown {} account '{name}'", provider.as_str()))?;
let mut account = account;
account.paused = paused;
self.upsert(account).await
}
pub async fn set_paused(&self, id: &str, paused: bool) -> Result<()> {
let mut account = self
.accounts
.read()
.await
.get(id)
.cloned()
.ok_or_else(|| anyhow!("unknown account {id}"))?;
account.paused = paused;
self.upsert(account).await
}
pub async fn has_account_name(&self, provider: Provider, name: &str) -> bool {
self.accounts.read().await.values().any(|a| a.provider == provider && a.name == name)
}
pub async fn upsert(&self, account: Account) -> Result<()> {
write_account_file(&self.dir, &account)?;
self.accounts
.write()
.await
.insert(account.id.clone(), account);
Ok(())
}
pub async fn remove(&self, id: &str) -> Result<bool> {
let existed = self.accounts.write().await.remove(id).is_some();
let path = self.dir.join(format!("{id}.json"));
if path.exists() {
std::fs::remove_file(&path)?;
}
Ok(existed)
}
pub async fn account_for(&self, provider: Provider, prefer_oauth: bool) -> Result<Account> {
let excluded = HashSet::new();
self.account_for_excluding(provider, prefer_oauth, &excluded)
.await
}
pub async fn account_for_excluding(
&self,
provider: Provider,
prefer_oauth: bool,
excluded: &HashSet<String>,
) -> Result<Account> {
self.account_for_excluding_preferred(provider, prefer_oauth, excluded, None)
.await
}
pub async fn account_for_excluding_preferred(
&self,
provider: Provider,
prefer_oauth: bool,
excluded: &HashSet<String>,
preferred_id: Option<&str>,
) -> Result<Account> {
self.account_for_excluding_preferred_mode(
provider,
prefer_oauth,
excluded,
preferred_id,
false,
)
.await
}
pub async fn account_for_excluding_preferred_mode(
&self,
provider: Provider,
prefer_oauth: bool,
excluded: &HashSet<String>,
preferred_id: Option<&str>,
preferred_ignores_cooldown: bool,
) -> Result<Account> {
let now = now_ms();
let candidates: Vec<Account> = {
let map = self.accounts.read().await;
let mut v: Vec<Account> = map
.values()
.filter(|a| {
a.provider == provider
&& a.status == "active"
&& !a.paused
&& !excluded.contains(&a.id)
})
.cloned()
.collect();
let policy = self.policy(provider);
v.retain(|account| account_proxy_eligible(account, &policy));
sort_by_policy(&mut v, &policy, prefer_oauth, self.rr_counter.load(Ordering::Relaxed));
v
};
if candidates.is_empty() {
bail!(
"no active {} account; run `alexandria auth import`",
provider.as_str()
);
}
let ready: Vec<&Account> = candidates
.iter()
.filter(|a| a.cooldown_until_ms.map(|c| c <= now).unwrap_or(true))
.collect();
let preferred = preferred_id.and_then(|preferred_id| {
candidates
.iter()
.find(|account| account.id == preferred_id)
.filter(|account| {
preferred_ignores_cooldown
|| account.cooldown_until_ms.map(|c| c <= now).unwrap_or(true)
})
.cloned()
});
let account = if let Some(preferred) = preferred {
preferred
} else if ready.is_empty() {
let account = candidates
.iter()
.min_by_key(|a| a.cooldown_until_ms.unwrap_or(i64::MAX))
.unwrap()
.clone();
tracing::warn!(
"all {} accounts cooling down; using {} (soonest expiry) in degraded mode",
provider.as_str(),
account.id
);
account
} else {
let policy = self.policy(provider);
let policy_mode = policy.mode.clone();
if policy_mode == AccountPolicyMode::RoundRobin {
let unblocked: Vec<&Account> = ready
.iter()
.copied()
.filter(|account| {
!codex_reserve_blocked(
account,
codex_reserve_pct(account, &policy),
now / 1000,
)
})
.collect();
let pool: Vec<&Account> = if unblocked.is_empty() {
ready.clone()
} else {
unblocked
};
let top_rank = oauth_rank(pool[0], prefer_oauth);
let group: Vec<&Account> = pool
.iter()
.copied()
.filter(|account| oauth_rank(account, prefer_oauth) == top_rank)
.collect();
let idx = self.rr_counter.fetch_add(1, Ordering::Relaxed) % group.len();
group[idx].clone()
} else {
ready[0].clone()
}
};
if account.needs_refresh() {
return self.refresh(&account.id, false).await;
}
Ok(account)
}
pub async fn mark_cooldown(&self, id: &str, until_ms: i64) -> Result<()> {
let mut account = self
.accounts
.read()
.await
.get(id)
.cloned()
.ok_or_else(|| anyhow!("unknown account {id}"))?;
account.cooldown_until_ms = Some(until_ms);
self.upsert(account).await
}
pub async fn refresh(&self, id: &str, force: bool) -> Result<Account> {
let lock = {
let mut locks = self.locks.lock().await;
locks
.entry(id.to_string())
.or_insert_with(|| Arc::new(Mutex::new(())))
.clone()
};
let _guard = lock.lock().await;
let account = self
.accounts
.read()
.await
.get(id)
.cloned()
.ok_or_else(|| anyhow!("unknown account {id}"))?;
if !force && !account.needs_refresh() {
return Ok(account);
}
tracing::info!("refreshing oauth token for {id}");
let result = match (account.provider, account.refresh_token.clone()) {
(Provider::Anthropic, Some(rt)) => self.refresh_anthropic(&rt).await,
(Provider::Openai, Some(rt)) => self.refresh_openai(&rt).await,
(Provider::Xai, Some(rt)) => {
let client_id = account
.account_meta
.get("oidc_client_id")
.and_then(|v| v.as_str())
.unwrap_or(XAI_CLIENT_ID)
.to_string();
self.refresh_xai(&rt, &client_id).await
}
(Provider::Gemini, Some(rt)) => self.refresh_gemini(&rt).await,
(Provider::Amp, _) => Err(anyhow!(
"amp accounts use a long-lived API key; re-run `alex auth import amp` or `alex auth amp-key`"
)),
(_, None) => Err(anyhow!("account {id} has no refresh token")),
};
let refreshed = match result {
Ok(r) => r,
Err(e) => {
tracing::warn!("refresh failed for {id}: {e}; re-importing from native creds");
if let Some(fresh) = self.reimport_native(&account).await {
return Ok(fresh);
}
return Err(e);
}
};
let mut updated = account.clone();
if let Some(t) = refreshed.access_token {
updated.access_token = Some(t);
}
if let Some(t) = refreshed.refresh_token {
updated.refresh_token = Some(t);
}
if let Some(t) = refreshed.id_token {
updated.id_token = Some(t);
}
updated.expires_at_ms = refreshed
.expires_in
.map(|s| now_ms() + s * 1000)
.or_else(|| updated.access_token.as_deref().and_then(jwt_exp_ms));
updated.last_refresh_ms = Some(now_ms());
if let Some(email) = updated.email() {
persist_account_email(&mut updated, &email);
} else {
if let Some(access_token) = updated.access_token.clone() {
if let Some(email) = fetch_provider_email(&self.http, updated.provider, &access_token).await {
persist_account_email(&mut updated, &email);
}
}
}
self.upsert(updated.clone()).await?;
Ok(updated)
}
async fn reimport_native(&self, stale: &Account) -> Option<Account> {
match stale.provider {
Provider::Anthropic => {
let _ = import_claude(self).await;
}
Provider::Openai => {
let _ = import_codex(self).await;
}
Provider::Gemini => {
let _ = import_gemini(self).await;
}
Provider::Xai => {
let _ = import_grok(self).await;
}
Provider::Amp => {
let _ = import_amp(self).await;
}
};
let fresh = self.accounts.read().await.get(&stale.id).cloned()?;
let changed = fresh.access_token != stale.access_token;
let valid = match fresh.expires_at_ms {
Some(exp) => exp > now_ms() + REFRESH_MARGIN_MS,
None => false,
};
if changed && valid {
tracing::info!("recovered {} from native credential source", stale.id);
Some(fresh)
} else {
None
}
}
async fn refresh_anthropic(&self, refresh_token: &str) -> Result<RefreshedTokens> {
let resp = self
.http
.post(ANTHROPIC_TOKEN_URL)
.json(&json!({
"grant_type": "refresh_token",
"refresh_token": refresh_token,
"client_id": ANTHROPIC_CLIENT_ID,
}))
.send()
.await?;
parse_token_response(resp).await
}
async fn refresh_openai(&self, refresh_token: &str) -> Result<RefreshedTokens> {
let resp = self
.http
.post(OPENAI_TOKEN_URL)
.json(&json!({
"client_id": OPENAI_CLIENT_ID,
"grant_type": "refresh_token",
"refresh_token": refresh_token,
"scope": "openid profile email",
}))
.send()
.await?;
parse_token_response(resp).await
}
async fn refresh_xai(&self, refresh_token: &str, client_id: &str) -> Result<RefreshedTokens> {
let resp = self
.http
.post(XAI_TOKEN_URL)
.form(&[
("grant_type", "refresh_token"),
("refresh_token", refresh_token),
("client_id", client_id),
])
.send()
.await?;
parse_token_response(resp).await
}
async fn refresh_gemini(&self, refresh_token: &str) -> Result<RefreshedTokens> {
let resp = self
.http
.post(GOOGLE_TOKEN_URL)
.form(&[
("grant_type", "refresh_token"),
("refresh_token", refresh_token),
("client_id", GEMINI_CLIENT_ID),
("client_secret", &gemini_client_secret()),
])
.send()
.await?;
parse_token_response(resp).await
}
pub async fn set_account_meta(&self, id: &str, key: &str, value: Value) -> Result<()> {
let mut account = self
.accounts
.read()
.await
.get(id)
.cloned()
.ok_or_else(|| anyhow!("unknown account {id}"))?;
if !account.account_meta.is_object() {
account.account_meta = json!({});
}
account.account_meta[key] = value;
self.upsert(account).await
}
}
#[derive(Debug, Deserialize)]
struct RefreshedTokens {
access_token: Option<String>,
refresh_token: Option<String>,
id_token: Option<String>,
expires_in: Option<i64>,
}
async fn parse_token_response(resp: reqwest::Response) -> Result<RefreshedTokens> {
let status = resp.status();
let text = resp.text().await?;
if !status.is_success() {
bail!("token refresh failed ({status}): {text}");
}
serde_json::from_str(&text).context("bad token response")
}
pub fn named_account_id(provider: Provider, kind: &str, name: &str) -> String {
if name == "default" {
format!("{}-{kind}", provider.as_str())
} else {
format!("{}-{kind}-{name}", provider.as_str())
}
}
fn account_path(dir: &Path, account: &Account) -> PathBuf {
account.path.clone().unwrap_or_else(|| dir.join(format!("{}.json", account.id)))
}
fn write_account_file(dir: &Path, account: &Account) -> Result<()> {
let path = account_path(dir, account);
let data = serde_json::to_string_pretty(account)?;
std::fs::write(&path, data)?;
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o600))?;
}
Ok(())
}
fn read_routing_policies(dir: &Path) -> Result<Vec<(Provider, AccountPolicy)>> {
let path = dir.join(ROUTING_POLICIES_FILE);
if !path.exists() {
return Ok(Vec::new());
}
let data = std::fs::read_to_string(&path)
.with_context(|| format!("reading {}", path.display()))?;
serde_json::from_str(&data).with_context(|| format!("parsing {}", path.display()))
}
fn write_routing_policies(
dir: &Path,
policies: &[(Provider, AccountPolicy)],
) -> Result<()> {
let path = dir.join(ROUTING_POLICIES_FILE);
let temp = dir.join(format!("{ROUTING_POLICIES_FILE}.tmp"));
let data = serde_json::to_vec_pretty(policies)?;
std::fs::write(&temp, data)?;
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
std::fs::set_permissions(&temp, std::fs::Permissions::from_mode(0o600))?;
}
std::fs::rename(&temp, &path)?;
Ok(())
}
pub struct ImportOutcome {
pub source: String,
pub imported: Vec<String>,
pub note: Option<String>,
}
pub async fn import_all(vault: &Vault, source: &str) -> Result<Vec<ImportOutcome>> {
let mut outcomes = Vec::new();
if source == "all" || source == "claude" {
outcomes.push(import_claude(vault).await);
}
if source == "all" || source == "codex" {
outcomes.push(import_codex(vault).await);
}
if source == "all" || source == "gemini" {
outcomes.push(import_gemini(vault).await);
}
if source == "all" || source == "grok" || source == "xai" {
outcomes.push(import_grok(vault).await);
}
if source == "all" || source == "amp" || source == "ampcode" {
outcomes.push(import_amp(vault).await);
}
if outcomes.is_empty() {
bail!("unknown source '{source}' (expected claude|codex|gemini|grok|xai|amp|all)");
}
Ok(outcomes)
}
fn home() -> PathBuf {
dirs::home_dir().expect("no home dir")
}
async fn import_claude(vault: &Vault) -> ImportOutcome {
let mut outcome = ImportOutcome {
source: "claude".into(),
imported: vec![],
note: None,
};
let path = home().join(".claude/.credentials.json");
let raw = if path.exists() {
std::fs::read_to_string(&path).ok()
} else {
claude_keychain()
};
let Some(raw) = raw else {
outcome.note = Some("no ~/.claude/.credentials.json and no Keychain entry".into());
return outcome;
};
let Ok(v) = serde_json::from_str::<Value>(&raw) else {
outcome.note = Some("could not parse claude credentials".into());
return outcome;
};
let oauth = &v["claudeAiOauth"];
let Some(access) = oauth["accessToken"].as_str() else {
outcome.note = Some("no claudeAiOauth.accessToken found".into());
return outcome;
};
let fallback_email = oauth["email"]
.as_str()
.or_else(|| v["email"].as_str())
.and_then(normalize_email)
.or_else(|| {
crate::login::jwt_payload(access)
.and_then(|payload| payload["email"].as_str().and_then(normalize_email))
});
let email = fetch_provider_email(&reqwest::Client::new(), Provider::Anthropic, access)
.await
.or(fallback_email);
let mut account = Account {
id: named_account_id(Provider::Anthropic, "oauth", "default"),
provider: Provider::Anthropic,
kind: "oauth".into(),
name: "default".into(),
description: email.clone(),
paused: false,
label: oauth["subscriptionType"]
.as_str()
.map(|s| format!("claude-code ({s})")),
access_token: Some(access.to_string()),
refresh_token: oauth["refreshToken"].as_str().map(String::from),
id_token: None,
api_key: None,
expires_at_ms: oauth["expiresAt"].as_i64(),
last_refresh_ms: Some(now_ms()),
account_meta: json!({"scopes": oauth["scopes"].clone()}),
cooldown_until_ms: None,
status: "active".into(),
path: None,
};
if let Some(email) = email {
persist_account_email(&mut account, &email);
}
match vault.upsert(account).await {
Ok(()) => outcome.imported.push("anthropic-oauth".into()),
Err(e) => outcome.note = Some(format!("failed to save: {e}")),
}
outcome
}
fn claude_keychain() -> Option<String> {
let out = std::process::Command::new("security")
.args([
"find-generic-password",
"-s",
"Claude Code-credentials",
"-w",
])
.output()
.ok()?;
if !out.status.success() {
return None;
}
Some(String::from_utf8_lossy(&out.stdout).trim().to_string())
}
async fn import_codex(vault: &Vault) -> ImportOutcome {
let mut outcome = ImportOutcome {
source: "codex".into(),
imported: vec![],
note: None,
};
let path = home().join(".codex/auth.json");
let Ok(raw) = std::fs::read_to_string(&path) else {
outcome.note = Some("no ~/.codex/auth.json".into());
return outcome;
};
let Ok(v) = serde_json::from_str::<Value>(&raw) else {
outcome.note = Some("could not parse codex auth.json".into());
return outcome;
};
if let Some(access) = v["tokens"]["access_token"].as_str() {
let account = Account {
id: named_account_id(Provider::Openai, "oauth", "default"),
provider: Provider::Openai,
kind: "oauth".into(),
name: "default".into(),
description: None,
paused: false,
label: Some("codex (chatgpt)".into()),
access_token: Some(access.to_string()),
refresh_token: v["tokens"]["refresh_token"].as_str().map(String::from),
id_token: v["tokens"]["id_token"].as_str().map(String::from),
api_key: None,
expires_at_ms: jwt_exp_ms(access),
last_refresh_ms: Some(now_ms()),
account_meta: json!({"account_id": v["tokens"]["account_id"].clone()}),
cooldown_until_ms: None,
status: "active".into(),
path: None,
};
match vault.upsert(account).await {
Ok(()) => outcome.imported.push("openai-oauth".into()),
Err(e) => outcome.note = Some(format!("failed to save: {e}")),
}
}
if let Some(key) = v["OPENAI_API_KEY"].as_str() {
let account = Account {
id: named_account_id(Provider::Openai, "api_key", "default"),
provider: Provider::Openai,
kind: "api_key".into(),
name: "default".into(),
description: None,
paused: false,
label: Some("codex (api key)".into()),
access_token: None,
refresh_token: None,
id_token: None,
api_key: Some(key.to_string()),
expires_at_ms: None,
last_refresh_ms: None,
account_meta: Value::Null,
cooldown_until_ms: None,
status: "active".into(),
path: None,
};
match vault.upsert(account).await {
Ok(()) => outcome.imported.push("openai-api-key".into()),
Err(e) => outcome.note = Some(format!("failed to save: {e}")),
}
}
if outcome.imported.is_empty() && outcome.note.is_none() {
outcome.note = Some("auth.json had neither tokens nor OPENAI_API_KEY".into());
}
outcome
}
async fn import_gemini(vault: &Vault) -> ImportOutcome {
let mut outcome = ImportOutcome {
source: "gemini".into(),
imported: vec![],
note: None,
};
let path = home().join(".gemini/oauth_creds.json");
let Ok(raw) = std::fs::read_to_string(&path) else {
outcome.note = Some("no ~/.gemini/oauth_creds.json".into());
return outcome;
};
let Ok(v) = serde_json::from_str::<Value>(&raw) else {
outcome.note = Some("could not parse gemini oauth_creds.json".into());
return outcome;
};
let Some(access) = v["access_token"].as_str() else {
outcome.note = Some("no access_token in gemini creds".into());
return outcome;
};
let account = Account {
id: named_account_id(Provider::Gemini, "oauth", "default"),
provider: Provider::Gemini,
kind: "oauth".into(),
name: "default".into(),
description: None,
paused: false,
label: Some("gemini-cli".into()),
access_token: Some(access.to_string()),
refresh_token: v["refresh_token"].as_str().map(String::from),
id_token: v["id_token"].as_str().map(String::from),
api_key: None,
expires_at_ms: v["expiry_date"].as_i64(),
last_refresh_ms: Some(now_ms()),
account_meta: Value::Null,
cooldown_until_ms: None,
status: "active".into(),
path: None,
};
match vault.upsert(account).await {
Ok(()) => outcome.imported.push("gemini-oauth".into()),
Err(e) => outcome.note = Some(format!("failed to save: {e}")),
}
outcome
}
pub async fn import_amp(vault: &Vault) -> ImportOutcome {
let mut outcome = ImportOutcome {
source: "amp".into(),
imported: vec![],
note: None,
};
let path = home().join(".local/share/amp/secrets.json");
if !path.exists() {
outcome.note = Some(
"no ~/.local/share/amp/secrets.json — run `amp login` or `alex auth amp-key <KEY>`"
.into(),
);
return outcome;
}
let Ok(raw) = std::fs::read_to_string(&path) else {
outcome.note = Some("could not read amp secrets.json".into());
return outcome;
};
let Ok(v) = serde_json::from_str::<Value>(&raw) else {
outcome.note = Some("could not parse amp secrets.json".into());
return outcome;
};
let Some(obj) = v.as_object() else {
outcome.note = Some("amp secrets.json is not an object".into());
return outcome;
};
let mut key: Option<(String, String)> = None; for (k, val) in obj {
if let Some(url) = k.strip_prefix("apiKey@") {
if let Some(s) = val.as_str() {
if !s.is_empty() {
key = Some((url.to_string(), s.to_string()));
if url.contains("ampcode.com") {
break;
}
}
}
}
}
let Some((amp_url, api_key)) = key else {
outcome.note = Some("no apiKey@… entry in amp secrets.json".into());
return outcome;
};
let account = Account {
id: "amp-api-key".into(),
provider: Provider::Amp,
kind: "api_key".into(),
name: default_account_name(),
description: None,
paused: false,
label: Some(format!("amp ({amp_url})")),
access_token: None,
refresh_token: None,
id_token: None,
api_key: Some(api_key),
expires_at_ms: None,
last_refresh_ms: Some(now_ms()),
account_meta: json!({ "amp_url": amp_url }),
cooldown_until_ms: None,
status: "active".into(),
path: None,
};
match vault.upsert(account).await {
Ok(()) => outcome.imported.push("amp-api-key".into()),
Err(e) => outcome.note = Some(format!("failed to save: {e}")),
}
outcome
}
pub async fn save_amp_api_key(vault: &Vault, api_key: &str) -> Result<String> {
let key = api_key.trim();
if key.is_empty() {
bail!("empty amp api key");
}
let account = Account {
id: "amp-api-key".into(),
provider: Provider::Amp,
kind: "api_key".into(),
name: default_account_name(),
description: None,
paused: false,
label: Some("amp (api key)".into()),
access_token: None,
refresh_token: None,
id_token: None,
api_key: Some(key.to_string()),
expires_at_ms: None,
last_refresh_ms: Some(now_ms()),
account_meta: json!({ "amp_url": "https://ampcode.com/" }),
cooldown_until_ms: None,
status: "active".into(),
path: None,
};
vault.upsert(account).await?;
Ok("amp-api-key".into())
}
async fn import_grok(vault: &Vault) -> ImportOutcome {
let mut outcome = ImportOutcome {
source: "grok".into(),
imported: vec![],
note: None,
};
let path = home().join(".grok/auth.json");
let Ok(raw) = std::fs::read_to_string(&path) else {
outcome.note = Some("no ~/.grok/auth.json".into());
return outcome;
};
let accounts = grok_accounts_from_json(&raw);
if accounts.is_empty() {
outcome.note = Some("no usable entries in grok auth.json".into());
return outcome;
}
for account in accounts {
let id = account.id.clone();
match vault.upsert(account).await {
Ok(()) => outcome.imported.push(id),
Err(e) => outcome.note = Some(format!("failed to save: {e}")),
}
}
outcome
}
pub fn grok_accounts_from_json(raw: &str) -> Vec<Account> {
let Ok(v) = serde_json::from_str::<Value>(raw) else {
return vec![];
};
let Some(map) = v.as_object() else {
return vec![];
};
let mut accounts = Vec::new();
for entry in map.values() {
let Some(key) = entry["key"].as_str() else {
continue;
};
let idx = accounts.len();
let id = if idx == 0 {
"xai-oauth".to_string()
} else {
format!("xai-oauth-{}", idx + 1)
};
let email = entry["email"].as_str().and_then(normalize_email);
let mut account = Account {
id,
provider: Provider::Xai,
kind: "oauth".into(),
name: if idx == 0 { "default".into() } else { format!("{}", idx + 1) },
description: email.clone(),
paused: false,
label: Some(email.as_ref().map(|email| format!("grok ({email})")).unwrap_or_else(|| "grok (oauth)".into())),
access_token: Some(key.to_string()),
refresh_token: entry["refresh_token"].as_str().map(String::from),
id_token: None,
api_key: None,
expires_at_ms: entry["expires_at"].as_str().and_then(rfc3339_to_ms),
last_refresh_ms: Some(now_ms()),
account_meta: json!({
"oidc_issuer": entry["oidc_issuer"].clone(),
"oidc_client_id": entry["oidc_client_id"].clone(),
"user_id": entry["user_id"].clone(),
}),
cooldown_until_ms: None,
status: "active".into(),
path: None,
};
if let Some(email) = email {
persist_account_email(&mut account, &email);
}
accounts.push(account);
}
accounts
}
#[cfg(test)]
mod tests {
use super::*;
fn temp_dir(name: &str) -> PathBuf {
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_nanos();
std::env::temp_dir().join(format!(
"alexandria-auth-{name}-{nanos}-{}",
std::process::id()
))
}
fn api_key_account(id: &str, provider: Provider) -> Account {
Account {
id: id.into(),
provider,
kind: "api_key".into(),
name: id.rsplit('-').next().unwrap_or("default").into(),
description: None,
paused: false,
label: None,
access_token: None,
refresh_token: None,
id_token: None,
api_key: Some(format!("sk-{id}")),
expires_at_ms: None,
last_refresh_ms: None,
account_meta: Value::Null,
cooldown_until_ms: None,
status: "active".into(),
path: None,
}
}
#[test]
fn rfc3339_parse() {
assert_eq!(rfc3339_to_ms("1970-01-01T00:00:01Z"), Some(1000));
assert_eq!(
rfc3339_to_ms("2001-09-09T01:46:40Z"),
Some(1_000_000_000_000)
);
assert_eq!(
rfc3339_to_ms("2001-09-09T03:46:40+02:00"),
Some(1_000_000_000_000)
);
assert_eq!(rfc3339_to_ms("not a timestamp"), None);
}
#[test]
fn account_email_prefers_metadata_and_falls_back_to_jwt_profile() {
let mut account = api_key_account("openai-oauth-default", Provider::Openai);
account.kind = "oauth".into();
account.api_key = None;
account.account_meta = json!({"email": "Primary@Example.com"});
assert_eq!(account.email().as_deref(), Some("primary@example.com"));
account.account_meta = json!({});
let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(
json!({"https://api.openai.com/profile": {"email": "Workspace@Example.com"}})
.to_string(),
);
account.id_token = Some(format!("header.{payload}.signature"));
assert_eq!(account.email().as_deref(), Some("workspace@example.com"));
}
#[test]
fn grok_json_parse() {
let dir = temp_dir("grok");
std::fs::create_dir_all(&dir).unwrap();
let path = dir.join("auth.json");
let raw = json!({
"https://auth.x.ai::b1a00492-0000-0000-0000-000000000000": {
"key": "bearer-token-abc",
"auth_mode": "oauth",
"create_time": "2026-01-01T00:00:00Z",
"user_id": "user-1",
"email": "user@x.com",
"refresh_token": "refresh-abc",
"expires_at": "2026-07-07T00:00:00Z",
"oidc_issuer": "https://auth.x.ai",
"oidc_client_id": "client-1"
},
"https://auth.x.ai::c2b11503-0000-0000-0000-000000000000": {
"key": "bearer-token-def",
"email": "second@x.com",
"refresh_token": "refresh-def",
"expires_at": "2026-08-01T00:00:00Z"
}
})
.to_string();
std::fs::write(&path, &raw).unwrap();
let read_back = std::fs::read_to_string(&path).unwrap();
let accounts = grok_accounts_from_json(&read_back);
assert_eq!(accounts.len(), 2);
assert_eq!(accounts[0].id, "xai-oauth");
assert_eq!(accounts[1].id, "xai-oauth-2");
assert_eq!(accounts[0].provider, Provider::Xai);
assert_eq!(accounts[0].kind, "oauth");
assert_eq!(accounts[0].label.as_deref(), Some("grok (user@x.com)"));
assert_eq!(accounts[0].description.as_deref(), Some("user@x.com"));
assert_eq!(accounts[0].account_meta["email"], "user@x.com");
assert_eq!(accounts[0].email().as_deref(), Some("user@x.com"));
assert_eq!(accounts[0].access_token.as_deref(), Some("bearer-token-abc"));
assert_eq!(accounts[0].refresh_token.as_deref(), Some("refresh-abc"));
assert_eq!(
accounts[0].expires_at_ms,
rfc3339_to_ms("2026-07-07T00:00:00Z")
);
assert!(accounts[0].expires_at_ms.is_some());
assert_eq!(
accounts[0].account_meta["oidc_issuer"].as_str(),
Some("https://auth.x.ai")
);
assert_eq!(
accounts[0].account_meta["oidc_client_id"].as_str(),
Some("client-1")
);
assert_eq!(accounts[0].account_meta["user_id"].as_str(), Some("user-1"));
assert!(grok_accounts_from_json("not json").is_empty());
assert!(grok_accounts_from_json("[1,2]").is_empty());
std::fs::remove_dir_all(&dir).ok();
}
#[tokio::test]
async fn round_robin_and_cooldown() {
let dir = temp_dir("pool");
let vault = Vault::open(dir.clone()).unwrap();
vault
.upsert(api_key_account("openai-key-a", Provider::Openai))
.await
.unwrap();
vault
.upsert(api_key_account("openai-key-b", Provider::Openai))
.await
.unwrap();
vault.set_policies(vec![(Provider::Openai, AccountPolicy { mode: AccountPolicyMode::RoundRobin, ..AccountPolicy::default() })]).await;
let mut picks = Vec::new();
for _ in 0..4 {
picks.push(vault.account_for(Provider::Openai, false).await.unwrap().id);
}
assert!(picks.contains(&"openai-key-a".to_string()));
assert!(picks.contains(&"openai-key-b".to_string()));
for pair in picks.windows(2) {
assert_ne!(pair[0], pair[1]);
}
vault
.mark_cooldown("openai-key-a", now_ms() + 60_000)
.await
.unwrap();
for _ in 0..4 {
let picked = vault.account_for(Provider::Openai, false).await.unwrap();
assert_eq!(picked.id, "openai-key-b");
}
vault
.mark_cooldown("openai-key-b", now_ms() + 120_000)
.await
.unwrap();
let degraded = vault.account_for(Provider::Openai, false).await.unwrap();
assert_eq!(degraded.id, "openai-key-a");
std::fs::remove_dir_all(&dir).ok();
}
#[tokio::test]
async fn exclusions_visit_each_active_account_once_even_in_degraded_mode() {
let dir = temp_dir("excluded-pool");
let vault = Vault::open(dir.clone()).unwrap();
for id in ["openai-key-a", "openai-key-b", "openai-key-c"] {
vault
.upsert(api_key_account(id, Provider::Openai))
.await
.unwrap();
}
let mut paused = api_key_account("openai-key-paused", Provider::Openai);
paused.paused = true;
vault.upsert(paused).await.unwrap();
let mut inactive = api_key_account("openai-key-inactive", Provider::Openai);
inactive.status = "disabled".into();
vault.upsert(inactive).await.unwrap();
let now = now_ms();
vault
.mark_cooldown("openai-key-a", now + 30_000)
.await
.unwrap();
vault
.mark_cooldown("openai-key-b", now + 20_000)
.await
.unwrap();
vault
.mark_cooldown("openai-key-c", now + 10_000)
.await
.unwrap();
let mut excluded = HashSet::new();
let mut selected = Vec::new();
for _ in 0..3 {
let account = vault
.account_for_excluding(Provider::Openai, false, &excluded)
.await
.unwrap();
assert!(excluded.insert(account.id.clone()));
selected.push(account.id);
}
assert_eq!(
selected,
vec!["openai-key-c", "openai-key-b", "openai-key-a"]
);
assert!(vault
.account_for_excluding(Provider::Openai, false, &excluded)
.await
.is_err());
assert!(!excluded.contains("openai-key-paused"));
assert!(!excluded.contains("openai-key-inactive"));
std::fs::remove_dir_all(&dir).ok();
}
#[tokio::test]
async fn preferred_account_never_bypasses_pause_policy_cooldown_or_exclusion() {
let dir = temp_dir("preferred-constraints");
let vault = Vault::open(dir.clone()).unwrap();
let mut first = api_key_account("openai-key-a", Provider::Openai);
first.name = "a".into();
let mut second = api_key_account("openai-key-b", Provider::Openai);
second.name = "b".into();
vault.upsert(first).await.unwrap();
vault.upsert(second.clone()).await.unwrap();
vault
.set_policies(vec![(
Provider::Openai,
AccountPolicy {
order: vec!["a".into(), "b".into()],
mode: AccountPolicyMode::Priority,
..AccountPolicy::default()
},
)])
.await;
let none = HashSet::new();
assert_eq!(
vault
.account_for_excluding_preferred(
Provider::Openai,
false,
&none,
Some("openai-key-b"),
)
.await
.unwrap()
.id,
"openai-key-b"
);
second.paused = true;
vault.upsert(second.clone()).await.unwrap();
assert_eq!(
vault
.account_for_excluding_preferred(
Provider::Openai,
false,
&none,
Some("openai-key-b"),
)
.await
.unwrap()
.id,
"openai-key-a"
);
second.paused = false;
second.cooldown_until_ms = Some(now_ms() + 60_000);
vault.upsert(second.clone()).await.unwrap();
assert_eq!(
vault
.account_for_excluding_preferred(
Provider::Openai,
false,
&none,
Some("openai-key-b"),
)
.await
.unwrap()
.id,
"openai-key-a"
);
assert_eq!(
vault
.account_for_excluding_preferred_mode(
Provider::Openai,
false,
&none,
Some("openai-key-b"),
true,
)
.await
.unwrap()
.id,
"openai-key-b"
);
second.cooldown_until_ms = None;
vault.upsert(second).await.unwrap();
vault
.set_policies(vec![(
Provider::Openai,
AccountPolicy {
disabled: vec!["b".into()],
..AccountPolicy::default()
},
)])
.await;
assert_eq!(
vault
.account_for_excluding_preferred(
Provider::Openai,
false,
&none,
Some("openai-key-b"),
)
.await
.unwrap()
.id,
"openai-key-a"
);
let excluded = HashSet::from(["openai-key-b".to_string()]);
assert_eq!(
vault
.account_for_excluding_preferred(
Provider::Openai,
false,
&excluded,
Some("openai-key-b"),
)
.await
.unwrap()
.id,
"openai-key-a"
);
std::fs::remove_dir_all(dir).ok();
}
#[tokio::test]
async fn vault_loads_multi_and_legacy_names() {
let dir = temp_dir("multi");
std::fs::create_dir_all(&dir).unwrap();
let mut legacy = api_key_account("openai-oauth", Provider::Openai);
legacy.kind = "oauth".into();
legacy.name = "default".into();
std::fs::write(dir.join("openai-oauth.json"), serde_json::to_string(&legacy).unwrap()).unwrap();
let mut work = api_key_account("openai-oauth-work", Provider::Openai);
work.kind = "oauth".into();
work.name = "work".into();
std::fs::write(dir.join("openai-oauth-work.json"), serde_json::to_string(&work).unwrap()).unwrap();
let vault = Vault::open(dir.clone()).unwrap();
let accounts = vault.list().await;
assert_eq!(accounts.len(), 2);
assert!(accounts.iter().any(|a| a.name == "default" && a.path.as_ref().unwrap().ends_with("openai-oauth.json")));
assert!(accounts.iter().any(|a| a.name == "work" && a.path.as_ref().unwrap().ends_with("openai-oauth-work.json")));
std::fs::remove_dir_all(&dir).ok();
}
#[tokio::test]
async fn policy_selection_priority_round_robin_threshold() {
let dir = temp_dir("policy");
let vault = Vault::open(dir.clone()).unwrap();
let mut work = api_key_account("openai-api_key-work", Provider::Openai);
work.name = "work".into();
let mut personal = api_key_account("openai-api_key-personal", Provider::Openai);
personal.name = "personal".into();
vault.upsert(work).await.unwrap();
vault.upsert(personal).await.unwrap();
vault.set_policies(vec![(Provider::Openai, AccountPolicy { order: vec!["work".into(), "personal".into()], mode: AccountPolicyMode::Priority, ..AccountPolicy::default() })]).await;
assert_eq!(vault.account_for(Provider::Openai, false).await.unwrap().name, "work");
vault.pause(Provider::Openai, "work", true).await.unwrap();
assert_eq!(vault.account_for(Provider::Openai, false).await.unwrap().name, "personal");
vault.pause(Provider::Openai, "work", false).await.unwrap();
vault.mark_cooldown("openai-api_key-work", now_ms() + 60_000).await.unwrap();
assert_eq!(vault.account_for(Provider::Openai, false).await.unwrap().name, "personal");
vault.mark_cooldown("openai-api_key-work", now_ms() - 1).await.unwrap();
vault.set_policies(vec![(Provider::Openai, AccountPolicy { order: vec!["work".into(), "personal".into()], mode: AccountPolicyMode::RoundRobin, ..AccountPolicy::default() })]).await;
let a = vault.account_for(Provider::Openai, false).await.unwrap().name;
let b = vault.account_for(Provider::Openai, false).await.unwrap().name;
assert_ne!(a, b);
let mut over = vault.list().await.into_iter().find(|a| a.name == "work").unwrap();
over.account_meta = json!({"rate_limit_pct": 90});
vault.upsert(over).await.unwrap();
vault.set_policies(vec![(Provider::Openai, AccountPolicy { order: vec!["work".into(), "personal".into()], mode: AccountPolicyMode::Threshold, threshold_pct: Some(80), ..AccountPolicy::default() })]).await;
assert_eq!(vault.account_for(Provider::Openai, false).await.unwrap().name, "personal");
std::fs::remove_dir_all(&dir).ok();
}
fn codex_limits(used_pct: f64, resets_at_s: i64) -> Value {
json!({
"codex_limits": {
"observed_at_ms": now_ms(),
"windows": [{
"window": "5h",
"used_pct": used_pct,
"resets_at_s": resets_at_s,
}],
}
})
}
#[test]
fn binding_reset_uses_most_consumed_active_window() {
let now_s = now_ms() / 1000;
let mut account = api_key_account("openai-api_key-binding", Provider::Openai);
account.account_meta = json!({
"codex_limits": {"windows": [
{"window": "5h", "used_pct": 25.0, "resets_at_s": now_s + 300},
{"window": "7d", "used_pct": 70.0, "resets_at_s": now_s + 500_000}
]}
});
let selected = codex_reset_selection(&account, now_s).unwrap();
assert_eq!(selected["window"], "7d");
assert_eq!(selected["used_pct"], 70.0);
account.account_meta["codex_limits"]["windows"][0]["used_pct"] = json!(70.0);
let tie = codex_reset_selection(&account, now_s).unwrap();
assert_eq!(tie["window"], "5h");
}
#[test]
fn legacy_policy_defaults_new_fields_without_losing_global_reserve() {
let policy: AccountPolicy = serde_json::from_value(json!({
"mode": "reset_first",
"reserve_pct": 17,
"order": ["personal"]
}))
.unwrap();
assert_eq!(policy.reserve_pct, Some(17));
assert!(policy.account_reserve_pct.is_empty());
assert!(policy.allow_mid_thread_failover);
}
#[tokio::test]
async fn reset_first_reserve_and_eligibility() {
let dir = temp_dir("reset-first");
let vault = Vault::open(dir.clone()).unwrap();
let now_s = now_ms() / 1000;
let mut soon = api_key_account("openai-api_key-soon", Provider::Openai);
soon.name = "soon".into();
soon.account_meta = codex_limits(20.0, now_s + 600);
let mut later = api_key_account("openai-api_key-later", Provider::Openai);
later.name = "later".into();
later.account_meta = codex_limits(20.0, now_s + 3600);
vault.upsert(soon.clone()).await.unwrap();
vault.upsert(later.clone()).await.unwrap();
assert_eq!(
vault.account_for(Provider::Openai, false).await.unwrap().name,
"soon"
);
soon.account_meta = codex_limits(95.0, now_s + 600);
vault.upsert(soon.clone()).await.unwrap();
assert_eq!(
vault.account_for(Provider::Openai, false).await.unwrap().name,
"later"
);
vault
.set_policies(vec![(
Provider::Openai,
AccountPolicy {
mode: AccountPolicyMode::ResetFirst,
reserve_pct: Some(10),
account_reserve_pct: HashMap::from([
("soon".into(), 0),
("later".into(), 10),
]),
..AccountPolicy::default()
},
)])
.await;
assert_eq!(
vault.account_for(Provider::Openai, false).await.unwrap().name,
"soon"
);
later.account_meta = codex_limits(95.0, now_s + 3600);
vault.upsert(later).await.unwrap();
assert_eq!(
vault.account_for(Provider::Openai, false).await.unwrap().name,
"soon"
);
vault
.set_policy_persisted(
Provider::Openai,
AccountPolicy {
mode: AccountPolicyMode::ResetFirst,
reserve_pct: Some(10),
disabled: vec!["soon".into()],
..AccountPolicy::default()
},
)
.await
.unwrap();
assert_eq!(
vault.account_for(Provider::Openai, false).await.unwrap().name,
"later"
);
std::fs::remove_dir_all(&dir).ok();
}
#[tokio::test]
async fn routing_policy_and_codex_limits_survive_reopen() {
let dir = temp_dir("routing-persist");
{
let vault = Vault::open(dir.clone()).unwrap();
let mut account = api_key_account("openai-api_key-work", Provider::Openai);
account.name = "work".into();
vault.upsert(account).await.unwrap();
vault
.set_policy_persisted(
Provider::Openai,
AccountPolicy {
order: vec!["work".into()],
mode: AccountPolicyMode::Priority,
reserve_pct: Some(15),
account_reserve_pct: HashMap::from([("work".into(), 23)]),
allow_mid_thread_failover: false,
..AccountPolicy::default()
},
)
.await
.unwrap();
vault
.record_codex_limits(
"openai-api_key-work",
json!({
"plan": "plus",
"windows": [{"window": "5h", "used_pct": 42.0, "resets_at_s": now_ms() / 1000 + 600}],
}),
)
.await
.unwrap();
}
let reopened = Vault::open(dir.clone()).unwrap();
let policy = reopened.policy(Provider::Openai);
assert_eq!(policy.mode, AccountPolicyMode::Priority);
assert_eq!(policy.reserve_pct, Some(15));
assert_eq!(policy.account_reserve_pct["work"], 23);
assert!(!policy.allow_mid_thread_failover);
assert_eq!(policy.order, vec!["work"]);
let account = reopened
.list()
.await
.into_iter()
.find(|account| account.name == "work")
.unwrap();
assert_eq!(
account.account_meta["codex_limits"]["windows"][0]["used_pct"],
42.0
);
assert!(account.account_meta["codex_limits"]["observed_at_ms"].is_i64());
std::fs::remove_dir_all(&dir).ok();
}
}