1use std::collections::HashMap;
2use std::sync::Arc;
3
4use anyhow::{bail, Context, Result};
5use serde_json::{json, Value};
6use tokio::sync::Mutex;
7
8use crate::login::{
9 anthropic_authorize_url, claude_exchange, codex_device_exchange_auto,
10 codex_device_poll_once, codex_device_start, codex_exchange_named, generate_pkce,
11 openai_authorize_url, wait_for_openai_callback, xai_device_poll_once, xai_device_start,
12 xai_upsert_from_tokens, CodexDevicePoll, XaiDevicePoll, OPENAI_CALLBACK_ADDR,
13 OPENAI_DEVICE_VERIFICATION_URL, OPENAI_REDIRECT_URI,
14};
15use crate::{named_account_id, now_ms, Vault};
16
17const SESSION_TTL_MS: i64 = 30 * 60 * 1000;
18
19#[derive(Debug, Clone, PartialEq)]
20pub enum LoginPhase {
21 Pending,
22 Done { account_id: String },
23 Failed { error: String },
24}
25
26pub struct LoginSession {
27 pub id: String,
28 pub provider: String,
29 pub mode: &'static str,
30 pub authorize_url: Option<String>,
31 pub user_code: Option<String>,
32 pub verification_uri: Option<String>,
33 pub verification_uri_complete: Option<String>,
34 pub created_ms: i64,
35 pub expires_at_ms: i64,
36 account_name: String,
37 verifier: Option<String>,
38 pub phase: LoginPhase,
39}
40
41impl LoginSession {
42 pub fn snapshot(&self) -> Value {
43 let (state, account_id, error) = match &self.phase {
44 LoginPhase::Pending => ("pending", None, None),
45 LoginPhase::Done { account_id } => ("done", Some(account_id.clone()), None),
46 LoginPhase::Failed { error } => ("failed", None, Some(error.clone())),
47 };
48 json!({
49 "login_id": self.id,
50 "provider": self.provider,
51 "mode": self.mode,
52 "state": state,
53 "account_id": account_id,
54 "error": error,
55 "authorize_url": self.authorize_url,
56 "user_code": self.user_code,
57 "verification_uri": self.verification_uri,
58 "verification_uri_complete": self.verification_uri_complete,
59 "expires_at_ms": self.expires_at_ms,
60 })
61 }
62}
63
64type SharedSession = Arc<Mutex<LoginSession>>;
65
66#[derive(Default)]
67pub struct LoginManager {
68 sessions: Mutex<HashMap<String, SharedSession>>,
69}
70
71fn random_id() -> String {
72 use rand::RngCore;
73 let mut bytes = [0u8; 12];
74 rand::thread_rng().fill_bytes(&mut bytes);
75 bytes.iter().map(|b| format!("{b:02x}")).collect()
76}
77
78impl LoginManager {
79 pub async fn start(&self, vault: Arc<Vault>, provider: &str, account_name: &str) -> Result<Value> {
80 self.prune().await;
81 validate_account_name(account_name)?;
82 let id = random_id();
83 let shared = match provider {
84 "claude" | "anthropic" => Arc::new(Mutex::new(self.start_claude(&id, account_name))),
85 "codex" | "openai" | "chatgpt" => self.start_codex(&id, vault.clone(), account_name).await?,
86 "grok" | "xai" => self.start_grok(&id, vault.clone(), account_name).await?,
87 "gemini" | "google" => self.start_gemini(&id, vault.clone(), account_name).await?,
88 "amp" | "ampcode" => {
89 let imported = crate::import_amp(&vault).await;
91 let mut session = LoginSession {
92 id: id.clone(),
93 provider: "amp".into(),
94 mode: "import",
95 authorize_url: Some("https://ampcode.com/settings".into()),
96 user_code: None,
97 verification_uri: None,
98 verification_uri_complete: None,
99 created_ms: now_ms(),
100 expires_at_ms: now_ms() + SESSION_TTL_MS,
101 account_name: account_name.to_string(),
102 verifier: None,
103 phase: if let Some(aid) = imported.imported.first() {
104 LoginPhase::Done {
105 account_id: aid.clone(),
106 }
107 } else {
108 LoginPhase::Failed {
109 error: imported.note.unwrap_or_else(|| {
110 "run `amp login` then retry, or `alex auth amp-key <KEY>`".into()
111 }),
112 }
113 },
114 };
115 let _ = &mut session;
116 Arc::new(Mutex::new(session))
117 }
118 other => bail!("unknown provider '{other}' (expected claude|codex|grok|gemini|amp)"),
119 };
120 let snapshot = shared.lock().await.snapshot();
121 self.sessions.lock().await.insert(id, shared);
122 Ok(snapshot)
123 }
124
125 pub async fn start_auto(&self, vault: Arc<Vault>, provider: &str) -> Result<Value> {
126 self.prune().await;
127 if !matches!(provider, "codex" | "openai" | "chatgpt") {
128 bail!("automatic identity login is currently supported only for Codex");
129 }
130 let id = random_id();
131 let shared = self.start_codex_device(&id, vault).await?;
132 let snapshot = shared.lock().await.snapshot();
133 self.sessions.lock().await.insert(id, shared);
134 Ok(snapshot)
135 }
136
137 pub async fn status(&self, login_id: &str) -> Option<Value> {
138 let session = self.sessions.lock().await.get(login_id).cloned()?;
139 let snapshot = session.lock().await.snapshot();
140 Some(snapshot)
141 }
142
143 pub async fn complete(&self, vault: Arc<Vault>, login_id: &str, input: &str) -> Result<Value> {
144 let session = self
145 .sessions
146 .lock()
147 .await
148 .get(login_id)
149 .cloned()
150 .context("unknown or expired login session")?;
151 let mut session = session.lock().await;
152 if session.mode != "paste" {
153 bail!(
154 "login session '{}' does not take a pasted code (mode: {})",
155 login_id,
156 session.mode
157 );
158 }
159 if session.phase != LoginPhase::Pending {
160 return Ok(session.snapshot());
161 }
162 let verifier = session
163 .verifier
164 .clone()
165 .context("session has no verifier")?;
166 match claude_exchange(&vault, &verifier, input).await {
167 Ok(account_id) => match rename_login_account(&vault, &account_id, &session.account_name).await {
168 Ok(account_id) => session.phase = LoginPhase::Done { account_id },
169 Err(e) => session.phase = LoginPhase::Failed { error: e.to_string() },
170 },
171 Err(e) => session.phase = LoginPhase::Failed { error: e.to_string() },
172 }
173 Ok(session.snapshot())
174 }
175
176 fn start_claude(&self, id: &str, account_name: &str) -> LoginSession {
177 let pkce = generate_pkce();
178 let url = anthropic_authorize_url(&pkce.challenge, &pkce.verifier);
179 LoginSession {
180 id: id.to_string(),
181 provider: "claude".into(),
182 mode: "paste",
183 authorize_url: Some(url),
184 user_code: None,
185 verification_uri: None,
186 verification_uri_complete: None,
187 created_ms: now_ms(),
188 expires_at_ms: now_ms() + SESSION_TTL_MS,
189 account_name: account_name.to_string(),
190 verifier: Some(pkce.verifier),
191 phase: LoginPhase::Pending,
192 }
193 }
194
195 async fn start_codex(&self, id: &str, vault: Arc<Vault>, account_name: &str) -> Result<SharedSession> {
196 let listener = tokio::net::TcpListener::bind(OPENAI_CALLBACK_ADDR)
197 .await
198 .with_context(|| {
199 format!("binding {OPENAI_CALLBACK_ADDR} for the oauth callback (is another login in progress?)")
200 })?;
201 let pkce = generate_pkce();
202 let state = random_id();
203 let url = openai_authorize_url(&pkce.challenge, &state);
204 let session = LoginSession {
205 id: id.to_string(),
206 provider: "codex".into(),
207 mode: "loopback",
208 authorize_url: Some(url),
209 user_code: None,
210 verification_uri: Some(OPENAI_REDIRECT_URI.into()),
211 verification_uri_complete: None,
212 created_ms: now_ms(),
213 expires_at_ms: now_ms() + SESSION_TTL_MS,
214 account_name: account_name.to_string(),
215 verifier: None,
216 phase: LoginPhase::Pending,
217 };
218 let shared = Arc::new(Mutex::new(session));
219 let worker = shared.clone();
220 let verifier = pkce.verifier;
221 let account_name = account_name.to_string();
222 tokio::spawn(async move {
223 let deadline = std::time::Duration::from_millis(SESSION_TTL_MS as u64);
224 let phase = match tokio::time::timeout(
225 deadline,
226 wait_for_openai_callback(&listener, &state),
227 )
228 .await
229 {
230 Ok(Ok(code)) => match codex_exchange_named(&vault, &verifier, &code, &account_name).await {
231 Ok(account_id) => LoginPhase::Done { account_id },
232 Err(e) => LoginPhase::Failed { error: e.to_string() },
233 },
234 Ok(Err(e)) => LoginPhase::Failed { error: e.to_string() },
235 Err(_) => LoginPhase::Failed {
236 error: "timed out waiting for the browser callback".into(),
237 },
238 };
239 worker.lock().await.phase = phase;
240 });
241 Ok(shared)
242 }
243
244 async fn start_codex_device(&self, id: &str, vault: Arc<Vault>) -> Result<SharedSession> {
245 const DEVICE_TTL_MS: i64 = 15 * 60 * 1000;
246 let http = reqwest::Client::builder()
247 .timeout(std::time::Duration::from_secs(30))
248 .build()?;
249 let start = codex_device_start(&http).await?;
250 let session = LoginSession {
251 id: id.to_string(),
252 provider: "codex".into(),
253 mode: "device",
254 authorize_url: Some(OPENAI_DEVICE_VERIFICATION_URL.into()),
255 user_code: Some(start.user_code.clone()),
256 verification_uri: Some(OPENAI_DEVICE_VERIFICATION_URL.into()),
257 verification_uri_complete: None,
258 created_ms: now_ms(),
259 expires_at_ms: now_ms() + DEVICE_TTL_MS,
260 account_name: String::new(),
261 verifier: None,
262 phase: LoginPhase::Pending,
263 };
264 let shared = Arc::new(Mutex::new(session));
265 let worker = shared.clone();
266 tokio::spawn(async move {
267 let deadline = now_ms() + DEVICE_TTL_MS;
268 let phase = loop {
269 if now_ms() > deadline {
270 break LoginPhase::Failed {
271 error: "Codex device code expired before authorization completed".into(),
272 };
273 }
274 tokio::time::sleep(std::time::Duration::from_secs(start.interval_s)).await;
275 match codex_device_poll_once(&http, &start).await {
276 CodexDevicePoll::Pending => continue,
277 CodexDevicePoll::Done {
278 authorization_code,
279 code_verifier,
280 } => {
281 break match codex_device_exchange_auto(
282 &vault,
283 &authorization_code,
284 &code_verifier,
285 )
286 .await
287 {
288 Ok(account_id) => LoginPhase::Done { account_id },
289 Err(error) => LoginPhase::Failed {
290 error: error.to_string(),
291 },
292 }
293 }
294 CodexDevicePoll::Failed(error) => {
295 break LoginPhase::Failed { error }
296 }
297 }
298 };
299 worker.lock().await.phase = phase;
300 });
301 Ok(shared)
302 }
303
304 async fn start_gemini(&self, id: &str, vault: Arc<Vault>, account_name: &str) -> Result<SharedSession> {
305 let (listener, port) = crate::login::bind_loopback().await?;
306 let redirect_uri = format!(
307 "http://localhost:{port}{}",
308 crate::login::GEMINI_CALLBACK_PATH
309 );
310 let pkce = generate_pkce();
311 let state = random_id();
312 let url = crate::login::gemini_authorize_url(&pkce.challenge, &state, &redirect_uri);
313 let session = LoginSession {
314 id: id.to_string(),
315 provider: "gemini".into(),
316 mode: "loopback",
317 authorize_url: Some(url),
318 user_code: None,
319 verification_uri: Some(redirect_uri.clone()),
320 verification_uri_complete: None,
321 created_ms: now_ms(),
322 expires_at_ms: now_ms() + SESSION_TTL_MS,
323 account_name: account_name.to_string(),
324 verifier: None,
325 phase: LoginPhase::Pending,
326 };
327 let shared = Arc::new(Mutex::new(session));
328 let worker = shared.clone();
329 let verifier = pkce.verifier;
330 let account_name = account_name.to_string();
331 tokio::spawn(async move {
332 let deadline = std::time::Duration::from_millis(SESSION_TTL_MS as u64);
333 let phase = match tokio::time::timeout(
334 deadline,
335 crate::login::wait_for_loopback_callback(
336 &listener,
337 &state,
338 crate::login::GEMINI_CALLBACK_PATH,
339 ),
340 )
341 .await
342 {
343 Ok(Ok(code)) => {
344 match crate::login::gemini_exchange(&vault, &verifier, &redirect_uri, &code)
345 .await
346 {
347 Ok(account_id) => match rename_login_account(&vault, &account_id, &account_name).await {
348 Ok(account_id) => LoginPhase::Done { account_id },
349 Err(e) => LoginPhase::Failed { error: e.to_string() },
350 },
351 Err(e) => LoginPhase::Failed { error: e.to_string() },
352 }
353 }
354 Ok(Err(e)) => LoginPhase::Failed {
355 error: e.to_string(),
356 },
357 Err(_) => LoginPhase::Failed {
358 error: "timed out waiting for the browser callback".into(),
359 },
360 };
361 worker.lock().await.phase = phase;
362 });
363 Ok(shared)
364 }
365
366 async fn start_grok(&self, id: &str, vault: Arc<Vault>, account_name: &str) -> Result<SharedSession> {
367 let http = reqwest::Client::new();
368 let start = xai_device_start(&http).await?;
369 let session = LoginSession {
370 id: id.to_string(),
371 provider: "grok".into(),
372 mode: "device",
373 authorize_url: start
374 .verification_uri_complete
375 .clone()
376 .or_else(|| Some(start.verification_uri.clone())),
377 user_code: Some(start.user_code.clone()),
378 verification_uri: Some(start.verification_uri.clone()),
379 verification_uri_complete: start.verification_uri_complete.clone(),
380 created_ms: now_ms(),
381 expires_at_ms: now_ms() + start.expires_in * 1000,
382 account_name: account_name.to_string(),
383 verifier: None,
384 phase: LoginPhase::Pending,
385 };
386 let shared = Arc::new(Mutex::new(session));
387 let worker = shared.clone();
388 let account_name = account_name.to_string();
389 tokio::spawn(async move {
390 let deadline = now_ms() + start.expires_in * 1000;
391 let mut interval = start.interval.max(1) as u64;
392 let phase = loop {
393 if now_ms() > deadline {
394 break LoginPhase::Failed {
395 error: "device code expired before authorization completed".into(),
396 };
397 }
398 tokio::time::sleep(std::time::Duration::from_secs(interval)).await;
399 match xai_device_poll_once(&http, &start.device_code).await {
400 XaiDevicePoll::Pending => continue,
401 XaiDevicePoll::SlowDown => {
402 interval += 5;
403 continue;
404 }
405 XaiDevicePoll::Done(tokens) => {
406 break match xai_upsert_from_tokens(&vault, &tokens).await {
407 Ok(account_id) => match rename_login_account(&vault, &account_id, &account_name).await {
408 Ok(account_id) => LoginPhase::Done { account_id },
409 Err(e) => LoginPhase::Failed { error: e.to_string() },
410 },
411 Err(e) => LoginPhase::Failed { error: e.to_string() },
412 };
413 }
414 XaiDevicePoll::Failed(e) => break LoginPhase::Failed { error: e },
415 }
416 };
417 worker.lock().await.phase = phase;
418 });
419 Ok(shared)
420 }
421
422 async fn prune(&self) {
423 let now = now_ms();
424 self.sessions
425 .lock()
426 .await
427 .retain(|_, s| match s.try_lock() {
428 Ok(s) => s.expires_at_ms > now - SESSION_TTL_MS,
429 Err(_) => true,
430 });
431 }
432}
433
434fn validate_account_name(name: &str) -> Result<()> {
435 if name.is_empty() || name.len() > 32 || !name.chars().all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '_' || c == '-') {
436 bail!("account name must match [a-z0-9_-]{{1,32}}");
437 }
438 Ok(())
439}
440
441async fn rename_login_account(vault: &Vault, account_id: &str, account_name: &str) -> Result<String> {
442 if account_name == "default" {
443 return Ok(account_id.to_string());
444 }
445 let mut account = vault
446 .list()
447 .await
448 .into_iter()
449 .find(|account| account.id == account_id)
450 .context("login completed but the saved account could not be found")?;
451 if vault.has_account_name(account.provider, account_name).await {
452 bail!("{} account '{account_name}' already exists", account.provider.as_str());
453 }
454 let default_id = named_account_id(account.provider, &account.kind, "default");
455 let previous_default = vault
456 .list()
457 .await
458 .into_iter()
459 .find(|candidate| candidate.id == default_id && candidate.id != account.id);
460 account.name = account_name.to_string();
461 account.id = named_account_id(account.provider, &account.kind, account_name);
462 account.path = None;
463 let renamed_id = account.id.clone();
464 vault.upsert(account).await?;
465 if let Some(previous_default) = previous_default {
466 vault.upsert(previous_default).await?;
467 } else {
468 let _ = vault.remove(&default_id).await?;
469 }
470 Ok(renamed_id)
471}
472
473#[cfg(test)]
474mod tests {
475 use super::*;
476 use std::path::PathBuf;
477 use std::time::{SystemTime, UNIX_EPOCH};
478
479 fn temp_vault(name: &str) -> (PathBuf, Arc<Vault>) {
480 let nanos = SystemTime::now()
481 .duration_since(UNIX_EPOCH)
482 .unwrap()
483 .as_nanos();
484 let dir = std::env::temp_dir().join(format!(
485 "alexandria-sessions-{name}-{nanos}-{}",
486 std::process::id()
487 ));
488 let vault = Arc::new(Vault::open(dir.clone()).unwrap());
489 (dir, vault)
490 }
491
492 #[tokio::test]
493 async fn claude_session_lifecycle() {
494 let (dir, vault) = temp_vault("claude");
495 let mgr = LoginManager::default();
496 let snap = mgr.start(vault.clone(), "claude", "default").await.unwrap();
497 assert_eq!(snap["mode"], "paste");
498 assert_eq!(snap["state"], "pending");
499 let url = snap["authorize_url"].as_str().unwrap();
500 assert!(url.starts_with("https://claude.ai/oauth/authorize"));
501 let id = snap["login_id"].as_str().unwrap();
502 let status = mgr.status(id).await.unwrap();
503 assert_eq!(status["state"], "pending");
504 let bad = mgr
505 .complete(vault.clone(), id, "code#wrong-state")
506 .await
507 .unwrap();
508 assert_eq!(bad["state"], "failed");
509 assert!(bad["error"].as_str().unwrap().contains("state mismatch"));
510 assert!(mgr.status("nope").await.is_none());
511 std::fs::remove_dir_all(&dir).ok();
512 }
513
514 #[tokio::test]
515 async fn unknown_provider_rejected() {
516 let (dir, vault) = temp_vault("unknown");
517 let mgr = LoginManager::default();
518 assert!(mgr.start(vault, "hal9000", "default").await.is_err());
519 std::fs::remove_dir_all(&dir).ok();
520 }
521
522 #[tokio::test]
523 async fn gemini_starts_loopback_oauth() {
524 let (dir, vault) = temp_vault("gemini");
525 let mgr = LoginManager::default();
526 let snap = mgr.start(vault, "gemini", "default").await.unwrap();
527 assert_eq!(snap["mode"], "loopback");
528 assert_eq!(snap["state"], "pending");
529 let url = snap["authorize_url"].as_str().unwrap();
530 assert!(url.starts_with("https://accounts.google.com/o/oauth2/v2/auth"));
531 assert!(url.contains("code_challenge"));
532 assert!(url.contains("access_type=offline"));
533 std::fs::remove_dir_all(&dir).ok();
534 }
535}