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