1use std::collections::HashMap;
2use std::path::{Path, PathBuf};
3use std::sync::atomic::{AtomicUsize, Ordering};
4use std::sync::Arc;
5use std::time::{SystemTime, UNIX_EPOCH};
6
7use alex_core::Provider;
8use anyhow::{anyhow, bail, Context, Result};
9use base64::Engine;
10use serde::{Deserialize, Serialize};
11use serde_json::{json, Value};
12use tokio::sync::{Mutex, RwLock};
13
14pub mod login;
15pub mod sessions;
16
17pub const ANTHROPIC_CLIENT_ID: &str = "9d1c250a-e61b-44d9-88ed-5944d1962f5e";
18pub const ANTHROPIC_TOKEN_URL: &str = "https://console.anthropic.com/v1/oauth/token";
19pub const OPENAI_CLIENT_ID: &str = "app_EMoamEEZ73f0CkXaXp7hrann";
20pub const OPENAI_TOKEN_URL: &str = "https://auth.openai.com/oauth/token";
21pub const XAI_CLIENT_ID: &str = "b1a00492-073a-47ea-816f-4c329264a828";
22pub const XAI_TOKEN_URL: &str = "https://auth.x.ai/oauth2/token";
23pub const GEMINI_CLIENT_ID: &str =
24 "681255809395-oo8ft2oprdrnp9e3aqf6av3hmdib135j.apps.googleusercontent.com";
25pub const GOOGLE_TOKEN_URL: &str = "https://oauth2.googleapis.com/token";
26
27pub fn gemini_client_secret() -> String {
31 ["GOCSPX", "4uHgMPm", "1o7Sk", "geV6Cu5clXFsxl"].join("-")
32}
33
34const REFRESH_MARGIN_MS: i64 = 120_000;
35
36pub fn now_ms() -> i64 {
37 SystemTime::now()
38 .duration_since(UNIX_EPOCH)
39 .unwrap()
40 .as_millis() as i64
41}
42
43#[derive(Debug, Clone, Serialize, Deserialize)]
44pub struct Account {
45 pub id: String,
46 pub provider: Provider,
47 pub kind: String,
48 #[serde(default)]
49 pub label: Option<String>,
50 #[serde(default)]
51 pub access_token: Option<String>,
52 #[serde(default)]
53 pub refresh_token: Option<String>,
54 #[serde(default)]
55 pub id_token: Option<String>,
56 #[serde(default)]
57 pub api_key: Option<String>,
58 #[serde(default)]
59 pub expires_at_ms: Option<i64>,
60 #[serde(default)]
61 pub last_refresh_ms: Option<i64>,
62 #[serde(default)]
63 pub account_meta: Value,
64 #[serde(default)]
65 pub cooldown_until_ms: Option<i64>,
66 #[serde(default = "default_status")]
67 pub status: String,
68}
69
70fn default_status() -> String {
71 "active".to_string()
72}
73
74impl Account {
75 pub fn chatgpt_account_id(&self) -> Option<String> {
76 self.account_meta
77 .get("account_id")
78 .and_then(|v| v.as_str())
79 .map(String::from)
80 }
81
82 fn needs_refresh(&self) -> bool {
83 self.kind == "oauth"
84 && match self.expires_at_ms {
85 Some(exp) => exp < now_ms() + REFRESH_MARGIN_MS,
86 None => true,
87 }
88 }
89}
90
91fn oauth_rank(account: &Account, prefer_oauth: bool) -> u8 {
92 if (account.kind == "oauth") == prefer_oauth {
93 0
94 } else {
95 1
96 }
97}
98
99pub fn rfc3339_to_ms(s: &str) -> Option<i64> {
100 chrono::DateTime::parse_from_rfc3339(s)
101 .ok()
102 .map(|dt| dt.timestamp_millis())
103}
104
105pub fn jwt_exp_ms(token: &str) -> Option<i64> {
106 let payload = token.split('.').nth(1)?;
107 let bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD
108 .decode(payload)
109 .ok()?;
110 let v: Value = serde_json::from_slice(&bytes).ok()?;
111 v.get("exp")?.as_i64().map(|s| s * 1000)
112}
113
114pub struct Vault {
115 dir: PathBuf,
116 accounts: RwLock<HashMap<String, Account>>,
117 locks: Mutex<HashMap<String, Arc<Mutex<()>>>>,
118 rr_counter: AtomicUsize,
119 http: reqwest::Client,
120}
121
122impl Vault {
123 pub fn open(dir: PathBuf) -> Result<Self> {
124 std::fs::create_dir_all(&dir)?;
125 let mut accounts = HashMap::new();
126 for entry in std::fs::read_dir(&dir)? {
127 let path = entry?.path();
128 if path.extension().map(|e| e == "json").unwrap_or(false) {
129 match std::fs::read_to_string(&path)
130 .map_err(anyhow::Error::from)
131 .and_then(|s| serde_json::from_str::<Account>(&s).map_err(Into::into))
132 {
133 Ok(acct) => {
134 accounts.insert(acct.id.clone(), acct);
135 }
136 Err(e) => tracing::warn!("skipping unreadable account file {path:?}: {e}"),
137 }
138 }
139 }
140 Ok(Self {
141 dir,
142 accounts: RwLock::new(accounts),
143 locks: Mutex::new(HashMap::new()),
144 rr_counter: AtomicUsize::new(0),
145 http: reqwest::Client::new(),
146 })
147 }
148
149 pub async fn list(&self) -> Vec<Account> {
150 let mut v: Vec<Account> = self.accounts.read().await.values().cloned().collect();
151 v.sort_by(|a, b| a.id.cmp(&b.id));
152 v
153 }
154
155 pub async fn upsert(&self, account: Account) -> Result<()> {
156 write_account_file(&self.dir, &account)?;
157 self.accounts
158 .write()
159 .await
160 .insert(account.id.clone(), account);
161 Ok(())
162 }
163
164 pub async fn remove(&self, id: &str) -> Result<bool> {
165 let existed = self.accounts.write().await.remove(id).is_some();
166 let path = self.dir.join(format!("{id}.json"));
167 if path.exists() {
168 std::fs::remove_file(&path)?;
169 }
170 Ok(existed)
171 }
172
173 pub async fn account_for(&self, provider: Provider, prefer_oauth: bool) -> Result<Account> {
174 let now = now_ms();
175 let candidates: Vec<Account> = {
176 let map = self.accounts.read().await;
177 let mut v: Vec<Account> = map
178 .values()
179 .filter(|a| a.provider == provider && a.status == "active")
180 .cloned()
181 .collect();
182 v.sort_by_key(|a| (oauth_rank(a, prefer_oauth), a.id.clone()));
183 v
184 };
185 if candidates.is_empty() {
186 bail!(
187 "no active {} account; run `alexandria auth import`",
188 provider.as_str()
189 );
190 }
191 let ready: Vec<&Account> = candidates
192 .iter()
193 .filter(|a| a.cooldown_until_ms.map(|c| c <= now).unwrap_or(true))
194 .collect();
195 let account = if ready.is_empty() {
196 let account = candidates
197 .iter()
198 .min_by_key(|a| a.cooldown_until_ms.unwrap_or(i64::MAX))
199 .unwrap()
200 .clone();
201 tracing::warn!(
202 "all {} accounts cooling down; using {} (soonest expiry) in degraded mode",
203 provider.as_str(),
204 account.id
205 );
206 account
207 } else {
208 let top_rank = oauth_rank(ready[0], prefer_oauth);
209 let group: Vec<&&Account> = ready
210 .iter()
211 .filter(|a| oauth_rank(a, prefer_oauth) == top_rank)
212 .collect();
213 let idx = self.rr_counter.fetch_add(1, Ordering::Relaxed) % group.len();
214 (*group[idx]).clone()
215 };
216 if account.needs_refresh() {
217 return self.refresh(&account.id, false).await;
218 }
219 Ok(account)
220 }
221
222 pub async fn mark_cooldown(&self, id: &str, until_ms: i64) -> Result<()> {
223 let mut account = self
224 .accounts
225 .read()
226 .await
227 .get(id)
228 .cloned()
229 .ok_or_else(|| anyhow!("unknown account {id}"))?;
230 account.cooldown_until_ms = Some(until_ms);
231 self.upsert(account).await
232 }
233
234 pub async fn refresh(&self, id: &str, force: bool) -> Result<Account> {
235 let lock = {
236 let mut locks = self.locks.lock().await;
237 locks
238 .entry(id.to_string())
239 .or_insert_with(|| Arc::new(Mutex::new(())))
240 .clone()
241 };
242 let _guard = lock.lock().await;
243
244 let account = self
245 .accounts
246 .read()
247 .await
248 .get(id)
249 .cloned()
250 .ok_or_else(|| anyhow!("unknown account {id}"))?;
251 if !force && !account.needs_refresh() {
252 return Ok(account);
253 }
254 tracing::info!("refreshing oauth token for {id}");
255 let result = match (account.provider, account.refresh_token.clone()) {
256 (Provider::Anthropic, Some(rt)) => self.refresh_anthropic(&rt).await,
257 (Provider::Openai, Some(rt)) => self.refresh_openai(&rt).await,
258 (Provider::Xai, Some(rt)) => {
259 let client_id = account
260 .account_meta
261 .get("oidc_client_id")
262 .and_then(|v| v.as_str())
263 .unwrap_or(XAI_CLIENT_ID)
264 .to_string();
265 self.refresh_xai(&rt, &client_id).await
266 }
267 (Provider::Gemini, Some(rt)) => self.refresh_gemini(&rt).await,
268 (_, None) => Err(anyhow!("account {id} has no refresh token")),
269 };
270 let refreshed = match result {
271 Ok(r) => r,
272 Err(e) => {
273 tracing::warn!("refresh failed for {id}: {e}; re-importing from native creds");
274 if let Some(fresh) = self.reimport_native(&account).await {
275 return Ok(fresh);
276 }
277 return Err(e);
278 }
279 };
280
281 let mut updated = account.clone();
282 if let Some(t) = refreshed.access_token {
283 updated.access_token = Some(t);
284 }
285 if let Some(t) = refreshed.refresh_token {
286 updated.refresh_token = Some(t);
287 }
288 if let Some(t) = refreshed.id_token {
289 updated.id_token = Some(t);
290 }
291 updated.expires_at_ms = refreshed
292 .expires_in
293 .map(|s| now_ms() + s * 1000)
294 .or_else(|| updated.access_token.as_deref().and_then(jwt_exp_ms));
295 updated.last_refresh_ms = Some(now_ms());
296 self.upsert(updated.clone()).await?;
297 Ok(updated)
298 }
299
300 async fn reimport_native(&self, stale: &Account) -> Option<Account> {
301 match stale.provider {
302 Provider::Anthropic => import_claude(self).await,
303 Provider::Openai => import_codex(self).await,
304 Provider::Gemini => import_gemini(self).await,
305 Provider::Xai => import_grok(self).await,
306 };
307 let fresh = self.accounts.read().await.get(&stale.id).cloned()?;
308 let changed = fresh.access_token != stale.access_token;
309 let valid = match fresh.expires_at_ms {
310 Some(exp) => exp > now_ms() + REFRESH_MARGIN_MS,
311 None => false,
312 };
313 if changed && valid {
314 tracing::info!("recovered {} from native credential source", stale.id);
315 Some(fresh)
316 } else {
317 None
318 }
319 }
320
321 async fn refresh_anthropic(&self, refresh_token: &str) -> Result<RefreshedTokens> {
322 let resp = self
323 .http
324 .post(ANTHROPIC_TOKEN_URL)
325 .json(&json!({
326 "grant_type": "refresh_token",
327 "refresh_token": refresh_token,
328 "client_id": ANTHROPIC_CLIENT_ID,
329 }))
330 .send()
331 .await?;
332 parse_token_response(resp).await
333 }
334
335 async fn refresh_openai(&self, refresh_token: &str) -> Result<RefreshedTokens> {
336 let resp = self
337 .http
338 .post(OPENAI_TOKEN_URL)
339 .json(&json!({
340 "client_id": OPENAI_CLIENT_ID,
341 "grant_type": "refresh_token",
342 "refresh_token": refresh_token,
343 "scope": "openid profile email",
344 }))
345 .send()
346 .await?;
347 parse_token_response(resp).await
348 }
349
350 async fn refresh_xai(&self, refresh_token: &str, client_id: &str) -> Result<RefreshedTokens> {
351 let resp = self
352 .http
353 .post(XAI_TOKEN_URL)
354 .form(&[
355 ("grant_type", "refresh_token"),
356 ("refresh_token", refresh_token),
357 ("client_id", client_id),
358 ])
359 .send()
360 .await?;
361 parse_token_response(resp).await
362 }
363
364 async fn refresh_gemini(&self, refresh_token: &str) -> Result<RefreshedTokens> {
365 let resp = self
366 .http
367 .post(GOOGLE_TOKEN_URL)
368 .form(&[
369 ("grant_type", "refresh_token"),
370 ("refresh_token", refresh_token),
371 ("client_id", GEMINI_CLIENT_ID),
372 ("client_secret", &gemini_client_secret()),
373 ])
374 .send()
375 .await?;
376 parse_token_response(resp).await
377 }
378
379 pub async fn set_account_meta(&self, id: &str, key: &str, value: Value) -> Result<()> {
380 let mut account = self
381 .accounts
382 .read()
383 .await
384 .get(id)
385 .cloned()
386 .ok_or_else(|| anyhow!("unknown account {id}"))?;
387 if !account.account_meta.is_object() {
388 account.account_meta = json!({});
389 }
390 account.account_meta[key] = value;
391 self.upsert(account).await
392 }
393}
394
395#[derive(Debug, Deserialize)]
396struct RefreshedTokens {
397 access_token: Option<String>,
398 refresh_token: Option<String>,
399 id_token: Option<String>,
400 expires_in: Option<i64>,
401}
402
403async fn parse_token_response(resp: reqwest::Response) -> Result<RefreshedTokens> {
404 let status = resp.status();
405 let text = resp.text().await?;
406 if !status.is_success() {
407 bail!("token refresh failed ({status}): {text}");
408 }
409 serde_json::from_str(&text).context("bad token response")
410}
411
412fn write_account_file(dir: &Path, account: &Account) -> Result<()> {
413 let path = dir.join(format!("{}.json", account.id));
414 let data = serde_json::to_string_pretty(account)?;
415 std::fs::write(&path, data)?;
416 #[cfg(unix)]
417 {
418 use std::os::unix::fs::PermissionsExt;
419 std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o600))?;
420 }
421 Ok(())
422}
423
424pub struct ImportOutcome {
425 pub source: String,
426 pub imported: Vec<String>,
427 pub note: Option<String>,
428}
429
430pub async fn import_all(vault: &Vault, source: &str) -> Result<Vec<ImportOutcome>> {
431 let mut outcomes = Vec::new();
432 if source == "all" || source == "claude" {
433 outcomes.push(import_claude(vault).await);
434 }
435 if source == "all" || source == "codex" {
436 outcomes.push(import_codex(vault).await);
437 }
438 if source == "all" || source == "gemini" {
439 outcomes.push(import_gemini(vault).await);
440 }
441 if source == "all" || source == "grok" || source == "xai" {
442 outcomes.push(import_grok(vault).await);
443 }
444 if outcomes.is_empty() {
445 bail!("unknown source '{source}' (expected claude|codex|gemini|grok|xai|all)");
446 }
447 Ok(outcomes)
448}
449
450fn home() -> PathBuf {
451 dirs::home_dir().expect("no home dir")
452}
453
454async fn import_claude(vault: &Vault) -> ImportOutcome {
455 let mut outcome = ImportOutcome {
456 source: "claude".into(),
457 imported: vec![],
458 note: None,
459 };
460 let path = home().join(".claude/.credentials.json");
461 let raw = if path.exists() {
462 std::fs::read_to_string(&path).ok()
463 } else {
464 claude_keychain()
465 };
466 let Some(raw) = raw else {
467 outcome.note = Some("no ~/.claude/.credentials.json and no Keychain entry".into());
468 return outcome;
469 };
470 let Ok(v) = serde_json::from_str::<Value>(&raw) else {
471 outcome.note = Some("could not parse claude credentials".into());
472 return outcome;
473 };
474 let oauth = &v["claudeAiOauth"];
475 let Some(access) = oauth["accessToken"].as_str() else {
476 outcome.note = Some("no claudeAiOauth.accessToken found".into());
477 return outcome;
478 };
479 let account = Account {
480 id: "anthropic-oauth".into(),
481 provider: Provider::Anthropic,
482 kind: "oauth".into(),
483 label: oauth["subscriptionType"]
484 .as_str()
485 .map(|s| format!("claude-code ({s})")),
486 access_token: Some(access.to_string()),
487 refresh_token: oauth["refreshToken"].as_str().map(String::from),
488 id_token: None,
489 api_key: None,
490 expires_at_ms: oauth["expiresAt"].as_i64(),
491 last_refresh_ms: Some(now_ms()),
492 account_meta: json!({"scopes": oauth["scopes"].clone()}),
493 cooldown_until_ms: None,
494 status: "active".into(),
495 };
496 match vault.upsert(account).await {
497 Ok(()) => outcome.imported.push("anthropic-oauth".into()),
498 Err(e) => outcome.note = Some(format!("failed to save: {e}")),
499 }
500 outcome
501}
502
503fn claude_keychain() -> Option<String> {
504 let out = std::process::Command::new("security")
505 .args([
506 "find-generic-password",
507 "-s",
508 "Claude Code-credentials",
509 "-w",
510 ])
511 .output()
512 .ok()?;
513 if !out.status.success() {
514 return None;
515 }
516 Some(String::from_utf8_lossy(&out.stdout).trim().to_string())
517}
518
519async fn import_codex(vault: &Vault) -> ImportOutcome {
520 let mut outcome = ImportOutcome {
521 source: "codex".into(),
522 imported: vec![],
523 note: None,
524 };
525 let path = home().join(".codex/auth.json");
526 let Ok(raw) = std::fs::read_to_string(&path) else {
527 outcome.note = Some("no ~/.codex/auth.json".into());
528 return outcome;
529 };
530 let Ok(v) = serde_json::from_str::<Value>(&raw) else {
531 outcome.note = Some("could not parse codex auth.json".into());
532 return outcome;
533 };
534 if let Some(access) = v["tokens"]["access_token"].as_str() {
535 let account = Account {
536 id: "openai-oauth".into(),
537 provider: Provider::Openai,
538 kind: "oauth".into(),
539 label: Some("codex (chatgpt)".into()),
540 access_token: Some(access.to_string()),
541 refresh_token: v["tokens"]["refresh_token"].as_str().map(String::from),
542 id_token: v["tokens"]["id_token"].as_str().map(String::from),
543 api_key: None,
544 expires_at_ms: jwt_exp_ms(access),
545 last_refresh_ms: Some(now_ms()),
546 account_meta: json!({"account_id": v["tokens"]["account_id"].clone()}),
547 cooldown_until_ms: None,
548 status: "active".into(),
549 };
550 match vault.upsert(account).await {
551 Ok(()) => outcome.imported.push("openai-oauth".into()),
552 Err(e) => outcome.note = Some(format!("failed to save: {e}")),
553 }
554 }
555 if let Some(key) = v["OPENAI_API_KEY"].as_str() {
556 let account = Account {
557 id: "openai-api-key".into(),
558 provider: Provider::Openai,
559 kind: "api_key".into(),
560 label: Some("codex (api key)".into()),
561 access_token: None,
562 refresh_token: None,
563 id_token: None,
564 api_key: Some(key.to_string()),
565 expires_at_ms: None,
566 last_refresh_ms: None,
567 account_meta: Value::Null,
568 cooldown_until_ms: None,
569 status: "active".into(),
570 };
571 match vault.upsert(account).await {
572 Ok(()) => outcome.imported.push("openai-api-key".into()),
573 Err(e) => outcome.note = Some(format!("failed to save: {e}")),
574 }
575 }
576 if outcome.imported.is_empty() && outcome.note.is_none() {
577 outcome.note = Some("auth.json had neither tokens nor OPENAI_API_KEY".into());
578 }
579 outcome
580}
581
582async fn import_gemini(vault: &Vault) -> ImportOutcome {
583 let mut outcome = ImportOutcome {
584 source: "gemini".into(),
585 imported: vec![],
586 note: None,
587 };
588 let path = home().join(".gemini/oauth_creds.json");
589 let Ok(raw) = std::fs::read_to_string(&path) else {
590 outcome.note = Some("no ~/.gemini/oauth_creds.json".into());
591 return outcome;
592 };
593 let Ok(v) = serde_json::from_str::<Value>(&raw) else {
594 outcome.note = Some("could not parse gemini oauth_creds.json".into());
595 return outcome;
596 };
597 let Some(access) = v["access_token"].as_str() else {
598 outcome.note = Some("no access_token in gemini creds".into());
599 return outcome;
600 };
601 let account = Account {
602 id: "gemini-oauth".into(),
603 provider: Provider::Gemini,
604 kind: "oauth".into(),
605 label: Some("gemini-cli".into()),
606 access_token: Some(access.to_string()),
607 refresh_token: v["refresh_token"].as_str().map(String::from),
608 id_token: v["id_token"].as_str().map(String::from),
609 api_key: None,
610 expires_at_ms: v["expiry_date"].as_i64(),
611 last_refresh_ms: Some(now_ms()),
612 account_meta: Value::Null,
613 cooldown_until_ms: None,
614 status: "active".into(),
615 };
616 match vault.upsert(account).await {
617 Ok(()) => outcome.imported.push("gemini-oauth".into()),
618 Err(e) => outcome.note = Some(format!("failed to save: {e}")),
619 }
620 outcome
621}
622
623async fn import_grok(vault: &Vault) -> ImportOutcome {
624 let mut outcome = ImportOutcome {
625 source: "grok".into(),
626 imported: vec![],
627 note: None,
628 };
629 let path = home().join(".grok/auth.json");
630 let Ok(raw) = std::fs::read_to_string(&path) else {
631 outcome.note = Some("no ~/.grok/auth.json".into());
632 return outcome;
633 };
634 let accounts = grok_accounts_from_json(&raw);
635 if accounts.is_empty() {
636 outcome.note = Some("no usable entries in grok auth.json".into());
637 return outcome;
638 }
639 for account in accounts {
640 let id = account.id.clone();
641 match vault.upsert(account).await {
642 Ok(()) => outcome.imported.push(id),
643 Err(e) => outcome.note = Some(format!("failed to save: {e}")),
644 }
645 }
646 outcome
647}
648
649pub fn grok_accounts_from_json(raw: &str) -> Vec<Account> {
650 let Ok(v) = serde_json::from_str::<Value>(raw) else {
651 return vec![];
652 };
653 let Some(map) = v.as_object() else {
654 return vec![];
655 };
656 let mut accounts = Vec::new();
657 for entry in map.values() {
658 let Some(key) = entry["key"].as_str() else {
659 continue;
660 };
661 let idx = accounts.len();
662 let id = if idx == 0 {
663 "xai-oauth".to_string()
664 } else {
665 format!("xai-oauth-{}", idx + 1)
666 };
667 let email = entry["email"].as_str().unwrap_or("unknown");
668 accounts.push(Account {
669 id,
670 provider: Provider::Xai,
671 kind: "oauth".into(),
672 label: Some(format!("grok ({email})")),
673 access_token: Some(key.to_string()),
674 refresh_token: entry["refresh_token"].as_str().map(String::from),
675 id_token: None,
676 api_key: None,
677 expires_at_ms: entry["expires_at"].as_str().and_then(rfc3339_to_ms),
678 last_refresh_ms: Some(now_ms()),
679 account_meta: json!({
680 "oidc_issuer": entry["oidc_issuer"].clone(),
681 "oidc_client_id": entry["oidc_client_id"].clone(),
682 "user_id": entry["user_id"].clone(),
683 }),
684 cooldown_until_ms: None,
685 status: "active".into(),
686 });
687 }
688 accounts
689}
690
691#[cfg(test)]
692mod tests {
693 use super::*;
694
695 fn temp_dir(name: &str) -> PathBuf {
696 let nanos = SystemTime::now()
697 .duration_since(UNIX_EPOCH)
698 .unwrap()
699 .as_nanos();
700 std::env::temp_dir().join(format!("alexandria-auth-{name}-{nanos}-{}", std::process::id()))
701 }
702
703 fn api_key_account(id: &str, provider: Provider) -> Account {
704 Account {
705 id: id.into(),
706 provider,
707 kind: "api_key".into(),
708 label: None,
709 access_token: None,
710 refresh_token: None,
711 id_token: None,
712 api_key: Some(format!("sk-{id}")),
713 expires_at_ms: None,
714 last_refresh_ms: None,
715 account_meta: Value::Null,
716 cooldown_until_ms: None,
717 status: "active".into(),
718 }
719 }
720
721 #[test]
722 fn rfc3339_parse() {
723 assert_eq!(rfc3339_to_ms("1970-01-01T00:00:01Z"), Some(1000));
724 assert_eq!(
725 rfc3339_to_ms("2001-09-09T01:46:40Z"),
726 Some(1_000_000_000_000)
727 );
728 assert_eq!(
729 rfc3339_to_ms("2001-09-09T03:46:40+02:00"),
730 Some(1_000_000_000_000)
731 );
732 assert_eq!(rfc3339_to_ms("not a timestamp"), None);
733 }
734
735 #[test]
736 fn grok_json_parse() {
737 let dir = temp_dir("grok");
738 std::fs::create_dir_all(&dir).unwrap();
739 let path = dir.join("auth.json");
740 let raw = json!({
741 "https://auth.x.ai::b1a00492-0000-0000-0000-000000000000": {
742 "key": "bearer-token-abc",
743 "auth_mode": "oauth",
744 "create_time": "2026-01-01T00:00:00Z",
745 "user_id": "user-1",
746 "email": "user@x.com",
747 "refresh_token": "refresh-abc",
748 "expires_at": "2026-07-07T00:00:00Z",
749 "oidc_issuer": "https://auth.x.ai",
750 "oidc_client_id": "client-1"
751 },
752 "https://auth.x.ai::c2b11503-0000-0000-0000-000000000000": {
753 "key": "bearer-token-def",
754 "email": "second@x.com",
755 "refresh_token": "refresh-def",
756 "expires_at": "2026-08-01T00:00:00Z"
757 }
758 })
759 .to_string();
760 std::fs::write(&path, &raw).unwrap();
761 let read_back = std::fs::read_to_string(&path).unwrap();
762 let accounts = grok_accounts_from_json(&read_back);
763 assert_eq!(accounts.len(), 2);
764 assert_eq!(accounts[0].id, "xai-oauth");
765 assert_eq!(accounts[1].id, "xai-oauth-2");
766 assert_eq!(accounts[0].provider, Provider::Xai);
767 assert_eq!(accounts[0].kind, "oauth");
768 assert_eq!(accounts[0].label.as_deref(), Some("grok (user@x.com)"));
769 assert_eq!(accounts[0].access_token.as_deref(), Some("bearer-token-abc"));
770 assert_eq!(accounts[0].refresh_token.as_deref(), Some("refresh-abc"));
771 assert_eq!(
772 accounts[0].expires_at_ms,
773 rfc3339_to_ms("2026-07-07T00:00:00Z")
774 );
775 assert!(accounts[0].expires_at_ms.is_some());
776 assert_eq!(
777 accounts[0].account_meta["oidc_issuer"].as_str(),
778 Some("https://auth.x.ai")
779 );
780 assert_eq!(
781 accounts[0].account_meta["oidc_client_id"].as_str(),
782 Some("client-1")
783 );
784 assert_eq!(accounts[0].account_meta["user_id"].as_str(), Some("user-1"));
785 assert!(grok_accounts_from_json("not json").is_empty());
786 assert!(grok_accounts_from_json("[1,2]").is_empty());
787 std::fs::remove_dir_all(&dir).ok();
788 }
789
790 #[tokio::test]
791 async fn round_robin_and_cooldown() {
792 let dir = temp_dir("pool");
793 let vault = Vault::open(dir.clone()).unwrap();
794 vault
795 .upsert(api_key_account("openai-key-a", Provider::Openai))
796 .await
797 .unwrap();
798 vault
799 .upsert(api_key_account("openai-key-b", Provider::Openai))
800 .await
801 .unwrap();
802
803 let mut picks = Vec::new();
804 for _ in 0..4 {
805 picks.push(vault.account_for(Provider::Openai, false).await.unwrap().id);
806 }
807 assert!(picks.contains(&"openai-key-a".to_string()));
808 assert!(picks.contains(&"openai-key-b".to_string()));
809 for pair in picks.windows(2) {
810 assert_ne!(pair[0], pair[1]);
811 }
812
813 vault
814 .mark_cooldown("openai-key-a", now_ms() + 60_000)
815 .await
816 .unwrap();
817 for _ in 0..4 {
818 let picked = vault.account_for(Provider::Openai, false).await.unwrap();
819 assert_eq!(picked.id, "openai-key-b");
820 }
821
822 vault
823 .mark_cooldown("openai-key-b", now_ms() + 120_000)
824 .await
825 .unwrap();
826 let degraded = vault.account_for(Provider::Openai, false).await.unwrap();
827 assert_eq!(degraded.id, "openai-key-a");
828
829 std::fs::remove_dir_all(&dir).ok();
830 }
831}