alex_auth/

lib.rs

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";

// The gemini-cli OAuth client is a public "installed app" credential embedded
// in Google's open-source CLI (not a confidential secret). Assembled from
// fragments so repo secret-scanners don't false-positive on the literal.
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)
    }

    /// A display-only identity hint. This never returns token material.
    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)
    })
}

/// The binding quota window is the active window closest to exhaustion.
/// Ties use the earlier reset. This avoids always picking the short window
/// when a weekly/monthly window is the actual constraint.
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
    }

    /// Select an account using the provider policy, while keeping an existing
    /// conversation on `preferred_id` when that account is still usable.
    ///
    /// A preferred account never bypasses pause/eligibility/cooldown or the
    /// caller's exclusion set. This lets the proxy preserve prompt-cache
    /// affinity without defeating routing controls or retry failover.
    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
}

/// Import Amp API key from `~/.local/share/amp/secrets.json` (CLI login material).
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; // (url, key)
    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()));
                    // Prefer ampcode.com if multiple
                    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
}

/// Save an Amp API key provided by the user (settings token / AMP_API_KEY).
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();
    }
}