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_exchange, generate_pkce,
10 openai_authorize_url, wait_for_openai_callback, xai_device_poll_once, xai_device_start,
11 xai_upsert_from_tokens, XaiDevicePoll, OPENAI_CALLBACK_ADDR, OPENAI_REDIRECT_URI,
12};
13use crate::{now_ms, Vault};
14
15const SESSION_TTL_MS: i64 = 30 * 60 * 1000;
16
17#[derive(Debug, Clone, PartialEq)]
18pub enum LoginPhase {
19 Pending,
20 Done { account_id: String },
21 Failed { error: String },
22}
23
24pub struct LoginSession {
25 pub id: String,
26 pub provider: String,
27 pub mode: &'static str,
28 pub authorize_url: Option<String>,
29 pub user_code: Option<String>,
30 pub verification_uri: Option<String>,
31 pub verification_uri_complete: Option<String>,
32 pub created_ms: i64,
33 pub expires_at_ms: i64,
34 verifier: Option<String>,
35 pub phase: LoginPhase,
36}
37
38impl LoginSession {
39 pub fn snapshot(&self) -> Value {
40 let (state, account_id, error) = match &self.phase {
41 LoginPhase::Pending => ("pending", None, None),
42 LoginPhase::Done { account_id } => ("done", Some(account_id.clone()), None),
43 LoginPhase::Failed { error } => ("failed", None, Some(error.clone())),
44 };
45 json!({
46 "login_id": self.id,
47 "provider": self.provider,
48 "mode": self.mode,
49 "state": state,
50 "account_id": account_id,
51 "error": error,
52 "authorize_url": self.authorize_url,
53 "user_code": self.user_code,
54 "verification_uri": self.verification_uri,
55 "verification_uri_complete": self.verification_uri_complete,
56 "expires_at_ms": self.expires_at_ms,
57 })
58 }
59}
60
61type SharedSession = Arc<Mutex<LoginSession>>;
62
63#[derive(Default)]
64pub struct LoginManager {
65 sessions: Mutex<HashMap<String, SharedSession>>,
66}
67
68fn random_id() -> String {
69 use rand::RngCore;
70 let mut bytes = [0u8; 12];
71 rand::thread_rng().fill_bytes(&mut bytes);
72 bytes.iter().map(|b| format!("{b:02x}")).collect()
73}
74
75impl LoginManager {
76 pub async fn start(&self, vault: Arc<Vault>, provider: &str) -> Result<Value> {
77 self.prune().await;
78 let id = random_id();
79 let shared = match provider {
80 "claude" | "anthropic" => Arc::new(Mutex::new(self.start_claude(&id))),
81 "codex" | "openai" | "chatgpt" => self.start_codex(&id, vault.clone()).await?,
82 "grok" | "xai" => self.start_grok(&id, vault.clone()).await?,
83 "gemini" | "google" => self.start_gemini(&id, vault.clone()).await?,
84 other => bail!("unknown provider '{other}' (expected claude|codex|grok|gemini)"),
85 };
86 let snapshot = shared.lock().await.snapshot();
87 self.sessions.lock().await.insert(id, shared);
88 Ok(snapshot)
89 }
90
91 pub async fn status(&self, login_id: &str) -> Option<Value> {
92 let session = self.sessions.lock().await.get(login_id).cloned()?;
93 let snapshot = session.lock().await.snapshot();
94 Some(snapshot)
95 }
96
97 pub async fn complete(&self, vault: Arc<Vault>, login_id: &str, input: &str) -> Result<Value> {
98 let session = self
99 .sessions
100 .lock()
101 .await
102 .get(login_id)
103 .cloned()
104 .context("unknown or expired login session")?;
105 let mut session = session.lock().await;
106 if session.mode != "paste" {
107 bail!(
108 "login session '{}' does not take a pasted code (mode: {})",
109 login_id,
110 session.mode
111 );
112 }
113 if session.phase != LoginPhase::Pending {
114 return Ok(session.snapshot());
115 }
116 let verifier = session.verifier.clone().context("session has no verifier")?;
117 match claude_exchange(&vault, &verifier, input).await {
118 Ok(account_id) => session.phase = LoginPhase::Done { account_id },
119 Err(e) => session.phase = LoginPhase::Failed { error: e.to_string() },
120 }
121 Ok(session.snapshot())
122 }
123
124 fn start_claude(&self, id: &str) -> LoginSession {
125 let pkce = generate_pkce();
126 let url = anthropic_authorize_url(&pkce.challenge, &pkce.verifier);
127 LoginSession {
128 id: id.to_string(),
129 provider: "claude".into(),
130 mode: "paste",
131 authorize_url: Some(url),
132 user_code: None,
133 verification_uri: None,
134 verification_uri_complete: None,
135 created_ms: now_ms(),
136 expires_at_ms: now_ms() + SESSION_TTL_MS,
137 verifier: Some(pkce.verifier),
138 phase: LoginPhase::Pending,
139 }
140 }
141
142 async fn start_codex(&self, id: &str, vault: Arc<Vault>) -> Result<SharedSession> {
143 let listener = tokio::net::TcpListener::bind(OPENAI_CALLBACK_ADDR)
144 .await
145 .with_context(|| {
146 format!("binding {OPENAI_CALLBACK_ADDR} for the oauth callback (is another login in progress?)")
147 })?;
148 let pkce = generate_pkce();
149 let state = random_id();
150 let url = openai_authorize_url(&pkce.challenge, &state);
151 let session = LoginSession {
152 id: id.to_string(),
153 provider: "codex".into(),
154 mode: "loopback",
155 authorize_url: Some(url),
156 user_code: None,
157 verification_uri: Some(OPENAI_REDIRECT_URI.into()),
158 verification_uri_complete: None,
159 created_ms: now_ms(),
160 expires_at_ms: now_ms() + SESSION_TTL_MS,
161 verifier: None,
162 phase: LoginPhase::Pending,
163 };
164 let shared = Arc::new(Mutex::new(session));
165 let worker = shared.clone();
166 let verifier = pkce.verifier;
167 tokio::spawn(async move {
168 let deadline = std::time::Duration::from_millis(SESSION_TTL_MS as u64);
169 let phase = match tokio::time::timeout(
170 deadline,
171 wait_for_openai_callback(&listener, &state),
172 )
173 .await
174 {
175 Ok(Ok(code)) => match codex_exchange(&vault, &verifier, &code).await {
176 Ok(account_id) => LoginPhase::Done { account_id },
177 Err(e) => LoginPhase::Failed { error: e.to_string() },
178 },
179 Ok(Err(e)) => LoginPhase::Failed { error: e.to_string() },
180 Err(_) => LoginPhase::Failed {
181 error: "timed out waiting for the browser callback".into(),
182 },
183 };
184 worker.lock().await.phase = phase;
185 });
186 Ok(shared)
187 }
188
189 async fn start_gemini(&self, id: &str, vault: Arc<Vault>) -> Result<SharedSession> {
190 let (listener, port) = crate::login::bind_loopback().await?;
191 let redirect_uri = format!("http://localhost:{port}{}", crate::login::GEMINI_CALLBACK_PATH);
192 let pkce = generate_pkce();
193 let state = random_id();
194 let url = crate::login::gemini_authorize_url(&pkce.challenge, &state, &redirect_uri);
195 let session = LoginSession {
196 id: id.to_string(),
197 provider: "gemini".into(),
198 mode: "loopback",
199 authorize_url: Some(url),
200 user_code: None,
201 verification_uri: Some(redirect_uri.clone()),
202 verification_uri_complete: None,
203 created_ms: now_ms(),
204 expires_at_ms: now_ms() + SESSION_TTL_MS,
205 verifier: None,
206 phase: LoginPhase::Pending,
207 };
208 let shared = Arc::new(Mutex::new(session));
209 let worker = shared.clone();
210 let verifier = pkce.verifier;
211 tokio::spawn(async move {
212 let deadline = std::time::Duration::from_millis(SESSION_TTL_MS as u64);
213 let phase = match tokio::time::timeout(
214 deadline,
215 crate::login::wait_for_loopback_callback(
216 &listener,
217 &state,
218 crate::login::GEMINI_CALLBACK_PATH,
219 ),
220 )
221 .await
222 {
223 Ok(Ok(code)) => {
224 match crate::login::gemini_exchange(&vault, &verifier, &redirect_uri, &code)
225 .await
226 {
227 Ok(account_id) => LoginPhase::Done { account_id },
228 Err(e) => LoginPhase::Failed { error: e.to_string() },
229 }
230 }
231 Ok(Err(e)) => LoginPhase::Failed { error: e.to_string() },
232 Err(_) => LoginPhase::Failed {
233 error: "timed out waiting for the browser callback".into(),
234 },
235 };
236 worker.lock().await.phase = phase;
237 });
238 Ok(shared)
239 }
240
241 async fn start_grok(&self, id: &str, vault: Arc<Vault>) -> Result<SharedSession> {
242 let http = reqwest::Client::new();
243 let start = xai_device_start(&http).await?;
244 let session = LoginSession {
245 id: id.to_string(),
246 provider: "grok".into(),
247 mode: "device",
248 authorize_url: start
249 .verification_uri_complete
250 .clone()
251 .or_else(|| Some(start.verification_uri.clone())),
252 user_code: Some(start.user_code.clone()),
253 verification_uri: Some(start.verification_uri.clone()),
254 verification_uri_complete: start.verification_uri_complete.clone(),
255 created_ms: now_ms(),
256 expires_at_ms: now_ms() + start.expires_in * 1000,
257 verifier: None,
258 phase: LoginPhase::Pending,
259 };
260 let shared = Arc::new(Mutex::new(session));
261 let worker = shared.clone();
262 tokio::spawn(async move {
263 let deadline = now_ms() + start.expires_in * 1000;
264 let mut interval = start.interval.max(1) as u64;
265 let phase = loop {
266 if now_ms() > deadline {
267 break LoginPhase::Failed {
268 error: "device code expired before authorization completed".into(),
269 };
270 }
271 tokio::time::sleep(std::time::Duration::from_secs(interval)).await;
272 match xai_device_poll_once(&http, &start.device_code).await {
273 XaiDevicePoll::Pending => continue,
274 XaiDevicePoll::SlowDown => {
275 interval += 5;
276 continue;
277 }
278 XaiDevicePoll::Done(tokens) => {
279 break match xai_upsert_from_tokens(&vault, &tokens).await {
280 Ok(account_id) => LoginPhase::Done { account_id },
281 Err(e) => LoginPhase::Failed { error: e.to_string() },
282 };
283 }
284 XaiDevicePoll::Failed(e) => break LoginPhase::Failed { error: e },
285 }
286 };
287 worker.lock().await.phase = phase;
288 });
289 Ok(shared)
290 }
291
292 async fn prune(&self) {
293 let now = now_ms();
294 self.sessions
295 .lock()
296 .await
297 .retain(|_, s| match s.try_lock() {
298 Ok(s) => s.expires_at_ms > now - SESSION_TTL_MS,
299 Err(_) => true,
300 });
301 }
302}
303
304#[cfg(test)]
305mod tests {
306 use super::*;
307 use std::path::PathBuf;
308 use std::time::{SystemTime, UNIX_EPOCH};
309
310 fn temp_vault(name: &str) -> (PathBuf, Arc<Vault>) {
311 let nanos = SystemTime::now()
312 .duration_since(UNIX_EPOCH)
313 .unwrap()
314 .as_nanos();
315 let dir = std::env::temp_dir().join(format!(
316 "alexandria-sessions-{name}-{nanos}-{}",
317 std::process::id()
318 ));
319 let vault = Arc::new(Vault::open(dir.clone()).unwrap());
320 (dir, vault)
321 }
322
323 #[tokio::test]
324 async fn claude_session_lifecycle() {
325 let (dir, vault) = temp_vault("claude");
326 let mgr = LoginManager::default();
327 let snap = mgr.start(vault.clone(), "claude").await.unwrap();
328 assert_eq!(snap["mode"], "paste");
329 assert_eq!(snap["state"], "pending");
330 let url = snap["authorize_url"].as_str().unwrap();
331 assert!(url.starts_with("https://claude.ai/oauth/authorize"));
332 let id = snap["login_id"].as_str().unwrap();
333 let status = mgr.status(id).await.unwrap();
334 assert_eq!(status["state"], "pending");
335 let bad = mgr.complete(vault.clone(), id, "code#wrong-state").await.unwrap();
336 assert_eq!(bad["state"], "failed");
337 assert!(bad["error"].as_str().unwrap().contains("state mismatch"));
338 assert!(mgr.status("nope").await.is_none());
339 std::fs::remove_dir_all(&dir).ok();
340 }
341
342 #[tokio::test]
343 async fn unknown_provider_rejected() {
344 let (dir, vault) = temp_vault("unknown");
345 let mgr = LoginManager::default();
346 assert!(mgr.start(vault, "hal9000").await.is_err());
347 std::fs::remove_dir_all(&dir).ok();
348 }
349
350 #[tokio::test]
351 async fn gemini_starts_loopback_oauth() {
352 let (dir, vault) = temp_vault("gemini");
353 let mgr = LoginManager::default();
354 let snap = mgr.start(vault, "gemini").await.unwrap();
355 assert_eq!(snap["mode"], "loopback");
356 assert_eq!(snap["state"], "pending");
357 let url = snap["authorize_url"].as_str().unwrap();
358 assert!(url.starts_with("https://accounts.google.com/o/oauth2/v2/auth"));
359 assert!(url.contains("code_challenge"));
360 assert!(url.contains("access_type=offline"));
361 std::fs::remove_dir_all(&dir).ok();
362 }
363}