diff --git a/codex-rs/login/src/auth_accounts.rs b/codex-rs/login/src/auth_accounts.rs new file mode 100644 index 000000000000..712f96323bb6 --- /dev/null +++ b/codex-rs/login/src/auth_accounts.rs @@ -0,0 +1,541 @@ +use crate::token_data::TokenData; +use chrono::DateTime; +use chrono::Utc; +use codex_app_server_protocol::AuthMode; +use rand::RngCore; +use serde::Deserialize; +use serde::Serialize; +use std::fs; +use std::fs::File; +use std::fs::OpenOptions; +use std::io; +use std::io::Read; +use std::io::Write; +use std::path::Path; +use std::path::PathBuf; + +const ACCOUNTS_FILE_NAME: &str = "auth_accounts.json"; + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct StoredAccount { + pub id: String, + pub mode: AuthMode, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub label: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub openai_api_key: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub tokens: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub last_refresh: Option>, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub created_at: Option>, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub last_used_at: Option>, +} + +#[derive(Debug, Clone, Deserialize, PartialEq, Serialize)] +struct AccountsFile { + #[serde(default = "default_version")] + version: u32, + + #[serde(default, skip_serializing_if = "Option::is_none")] + active_account_id: Option, + + #[serde(default, skip_serializing_if = "Vec::is_empty")] + accounts: Vec, +} + +impl Default for AccountsFile { + fn default() -> Self { + Self { + version: default_version(), + active_account_id: None, + accounts: Vec::new(), + } + } +} + +fn default_version() -> u32 { + 1 +} + +fn accounts_file_path(codex_home: &Path) -> PathBuf { + codex_home.join(ACCOUNTS_FILE_NAME) +} + +fn read_accounts_file(path: &Path) -> io::Result { + match File::open(path) { + Ok(mut file) => { + let mut contents = String::new(); + file.read_to_string(&mut contents)?; + let (parsed, repaired) = parse_accounts_file(&contents)?; + if repaired { + write_accounts_file(path, &parsed)?; + } + Ok(parsed) + } + Err(err) if err.kind() == io::ErrorKind::NotFound => Ok(AccountsFile::default()), + Err(err) => Err(err), + } +} + +fn parse_accounts_file(contents: &str) -> io::Result<(AccountsFile, bool)> { + if contents.trim().is_empty() { + return Ok((AccountsFile::default(), false)); + } + + match serde_json::from_str(contents) { + Ok(parsed) => Ok((parsed, false)), + Err(original_err) => { + let mut latest: Option = None; + let mut recovered_count = 0; + let stream = serde_json::Deserializer::from_str(contents).into_iter::(); + for value in stream { + match value { + Ok(parsed) => { + recovered_count += 1; + latest = Some(parsed); + } + Err(_) => return Err(original_err.into()), + } + } + + match (recovered_count, latest) { + (count, Some(parsed)) if count > 1 => Ok((parsed, true)), + _ => Err(original_err.into()), + } + } + } +} + +fn write_accounts_file(path: &Path, data: &AccountsFile) -> io::Result<()> { + let parent = path.parent().ok_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidInput, + format!("accounts path has no parent: {}", path.display()), + ) + })?; + fs::create_dir_all(parent)?; + + let raw = serde_json::to_string_pretty(data).map_err(io::Error::other)?; + let (tmp_path, mut file) = create_accounts_tmp_file(path)?; + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + file.set_permissions(fs::Permissions::from_mode(0o600))?; + } + file.write_all(raw.as_bytes())?; + file.flush()?; + file.sync_all()?; + drop(file); + if let Err(err) = replace_accounts_file(&tmp_path, path) { + let _ = fs::remove_file(&tmp_path); + return Err(err); + } + Ok(()) +} + +fn create_accounts_tmp_file(path: &Path) -> io::Result<(PathBuf, fs::File)> { + let file_name = path + .file_name() + .and_then(|name| name.to_str()) + .unwrap_or(ACCOUNTS_FILE_NAME); + let parent = path.parent().unwrap_or_else(|| Path::new(".")); + for attempt in 0..100 { + let pid = std::process::id(); + let tmp_path = parent.join(format!(".{file_name}.{pid}.{attempt}.tmp")); + let mut options = OpenOptions::new(); + options.write(true).create_new(true); + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt; + options.mode(0o600); + } + match options.open(&tmp_path) { + Ok(file) => return Ok((tmp_path, file)), + Err(err) if err.kind() == io::ErrorKind::AlreadyExists => continue, + Err(err) => return Err(err), + } + } + Err(io::Error::new( + io::ErrorKind::AlreadyExists, + format!( + "failed to allocate temporary accounts path for {}", + path.display() + ), + )) +} + +fn replace_accounts_file(src: &Path, dst: &Path) -> io::Result<()> { + #[cfg(not(windows))] + { + fs::rename(src, dst) + } + #[cfg(windows)] + { + match replace_file_windows(src, dst) { + Ok(()) => Ok(()), + Err(err) if err.kind() == io::ErrorKind::NotFound => fs::rename(src, dst), + Err(err) => Err(err), + } + } +} + +#[cfg(windows)] +fn replace_file_windows(src: &Path, dst: &Path) -> io::Result<()> { + use std::ffi::c_void; + use std::os::windows::ffi::OsStrExt; + + const REPLACEFILE_IGNORE_MERGE_ERRORS: u32 = 0x0000_0002; + + unsafe extern "system" { + fn ReplaceFileW( + lp_replaced_file_name: *const u16, + lp_replacement_file_name: *const u16, + lp_backup_file_name: *const u16, + dw_replace_flags: u32, + lp_exclude: *mut c_void, + lp_reserved: *mut c_void, + ) -> i32; + } + + let dst_wide: Vec = dst.as_os_str().encode_wide().chain(Some(0)).collect(); + let src_wide: Vec = src.as_os_str().encode_wide().chain(Some(0)).collect(); + let replaced = unsafe { + ReplaceFileW( + dst_wide.as_ptr(), + src_wide.as_ptr(), + std::ptr::null(), + REPLACEFILE_IGNORE_MERGE_ERRORS, + std::ptr::null_mut(), + std::ptr::null_mut(), + ) + }; + if replaced == 0 { + return Err(io::Error::last_os_error()); + } + Ok(()) +} + +fn normalize_email(email: &str) -> String { + email.trim().to_ascii_lowercase() +} + +fn now() -> DateTime { + Utc::now() +} + +fn next_id() -> String { + let mut bytes = [0_u8; 16]; + rand::rng().fill_bytes(&mut bytes); + bytes.iter().map(|byte| format!("{byte:02x}")).collect() +} + +fn is_supported_chatgpt_mode(mode: AuthMode) -> bool { + matches!(mode, AuthMode::Chatgpt | AuthMode::ChatgptAuthTokens) +} + +fn chatgpt_account_id(tokens: &TokenData) -> Option<&str> { + tokens + .account_id + .as_deref() + .or(tokens.id_token.chatgpt_account_id.as_deref()) +} + +fn match_chatgpt_account(existing: &StoredAccount, tokens: &TokenData) -> bool { + if !is_supported_chatgpt_mode(existing.mode) { + return false; + } + + let existing_tokens = match &existing.tokens { + Some(tokens) => tokens, + None => return false, + }; + + match ( + chatgpt_account_id(existing_tokens), + chatgpt_account_id(tokens), + ) { + (Some(a), Some(b)) if a == b => { + match ( + existing_tokens.id_token.email.as_ref(), + tokens.id_token.email.as_ref(), + ) { + (Some(a), Some(b)) => normalize_email(a) == normalize_email(b), + _ => true, + } + } + (Some(_), Some(_)) => false, + _ => false, + } +} + +fn match_api_key_account(existing: &StoredAccount, api_key: &str) -> bool { + existing.mode == AuthMode::ApiKey + && existing + .openai_api_key + .as_ref() + .is_some_and(|stored| stored == api_key) +} + +fn touch_account(account: &mut StoredAccount, used: bool) { + if account.created_at.is_none() { + account.created_at = Some(now()); + } + if used { + account.last_used_at = Some(now()); + } +} + +fn upsert_account( + mut data: AccountsFile, + mut new_account: StoredAccount, +) -> (AccountsFile, StoredAccount) { + let existing_idx = match new_account.mode { + AuthMode::ApiKey => new_account.openai_api_key.as_ref().and_then(|api_key| { + data.accounts + .iter() + .position(|acc| match_api_key_account(acc, api_key)) + }), + AuthMode::Chatgpt | AuthMode::ChatgptAuthTokens => { + new_account.tokens.as_ref().and_then(|tokens| { + data.accounts + .iter() + .position(|acc| match_chatgpt_account(acc, tokens)) + }) + } + AuthMode::AgentIdentity | AuthMode::PersonalAccessToken => None, + }; + + if let Some(idx) = existing_idx { + let mut account = data.accounts[idx].clone(); + if new_account.label.is_some() { + account.label = new_account.label; + } + if new_account.last_refresh.is_some() { + account.last_refresh = new_account.last_refresh; + } + if let Some(tokens) = new_account.tokens { + account.tokens = Some(tokens); + } + if let Some(api_key) = new_account.openai_api_key { + account.openai_api_key = Some(api_key); + } + if let Some(last_used) = new_account.last_used_at { + account.last_used_at = Some(last_used); + } + data.accounts[idx] = account.clone(); + return (data, account); + } + + if new_account.created_at.is_none() { + new_account.created_at = Some(now()); + } + + data.accounts.push(new_account.clone()); + (data, new_account) +} + +fn select_fallback_active_account(data: &mut AccountsFile) { + if let Some(account) = data.accounts.first_mut() { + data.active_account_id = Some(account.id.clone()); + touch_account(account, true); + } else { + data.active_account_id = None; + } +} + +pub fn list_accounts(codex_home: &Path) -> io::Result> { + let path = accounts_file_path(codex_home); + let data = read_accounts_file(&path)?; + Ok(data.accounts) +} + +pub fn get_active_account_id(codex_home: &Path) -> io::Result> { + let path = accounts_file_path(codex_home); + let data = read_accounts_file(&path)?; + Ok(data.active_account_id) +} + +pub fn find_account(codex_home: &Path, account_id: &str) -> io::Result> { + let path = accounts_file_path(codex_home); + let data = read_accounts_file(&path)?; + Ok(data.accounts.into_iter().find(|acc| acc.id == account_id)) +} + +pub fn update_account_last_refresh( + codex_home: &Path, + account_id: &str, + last_refresh: DateTime, +) -> io::Result> { + let path = accounts_file_path(codex_home); + let mut data = read_accounts_file(&path)?; + + let Some(account) = data.accounts.iter_mut().find(|acc| acc.id == account_id) else { + return Ok(None); + }; + account.last_refresh = Some(last_refresh); + let updated = account.clone(); + write_accounts_file(&path, &data)?; + Ok(Some(updated)) +} + +pub fn set_active_account_id( + codex_home: &Path, + account_id: Option, +) -> io::Result> { + let path = accounts_file_path(codex_home); + let mut data = read_accounts_file(&path)?; + + let mut updated = None; + if let Some(id) = account_id + && let Some(account) = data.accounts.iter_mut().find(|acc| acc.id == id) + { + data.active_account_id = Some(id); + touch_account(account, true); + updated = Some(account.clone()); + } else { + data.active_account_id = None; + } + write_accounts_file(&path, &data)?; + Ok(updated) +} + +pub fn remove_account(codex_home: &Path, account_id: &str) -> io::Result> { + let path = accounts_file_path(codex_home); + let mut data = read_accounts_file(&path)?; + + let removed = if let Some(pos) = data.accounts.iter().position(|acc| acc.id == account_id) { + Some(data.accounts.remove(pos)) + } else { + None + }; + + if data + .active_account_id + .as_ref() + .is_some_and(|active| active == account_id) + { + select_fallback_active_account(&mut data); + } + + write_accounts_file(&path, &data)?; + Ok(removed) +} + +pub fn remove_account_matching_credentials( + codex_home: &Path, + mode: AuthMode, + openai_api_key: Option<&str>, + tokens: Option<&TokenData>, +) -> io::Result> { + let path = accounts_file_path(codex_home); + let mut data = read_accounts_file(&path)?; + + let removed = match mode { + AuthMode::ApiKey => openai_api_key.and_then(|api_key| { + data.accounts + .iter() + .position(|account| match_api_key_account(account, api_key)) + }), + AuthMode::Chatgpt | AuthMode::ChatgptAuthTokens => tokens.and_then(|tokens| { + data.accounts + .iter() + .position(|account| match_chatgpt_account(account, tokens)) + }), + AuthMode::AgentIdentity | AuthMode::PersonalAccessToken => None, + } + .map(|pos| data.accounts.remove(pos)); + + if let Some(removed) = &removed + && data + .active_account_id + .as_ref() + .is_some_and(|active| active == &removed.id) + { + select_fallback_active_account(&mut data); + } + + write_accounts_file(&path, &data)?; + Ok(removed) +} + +pub fn upsert_api_key_account( + codex_home: &Path, + api_key: String, + label: Option, + make_active: bool, +) -> io::Result { + let path = accounts_file_path(codex_home); + let data = read_accounts_file(&path)?; + + let new_account = StoredAccount { + id: next_id(), + mode: AuthMode::ApiKey, + label, + openai_api_key: Some(api_key), + tokens: None, + last_refresh: None, + created_at: None, + last_used_at: None, + }; + + let (mut data, mut stored) = upsert_account(data, new_account); + + if make_active { + data.active_account_id = Some(stored.id.clone()); + if let Some(account) = data.accounts.iter_mut().find(|acc| acc.id == stored.id) { + touch_account(account, true); + stored = account.clone(); + } + } + + write_accounts_file(&path, &data)?; + Ok(stored) +} + +pub fn upsert_chatgpt_account( + codex_home: &Path, + tokens: TokenData, + last_refresh: DateTime, + label: Option, + make_active: bool, +) -> io::Result { + let path = accounts_file_path(codex_home); + let data = read_accounts_file(&path)?; + + let new_account = StoredAccount { + id: next_id(), + mode: AuthMode::Chatgpt, + label, + openai_api_key: None, + tokens: Some(tokens), + last_refresh: Some(last_refresh), + created_at: None, + last_used_at: None, + }; + + let (mut data, mut stored) = upsert_account(data, new_account); + + if make_active { + data.active_account_id = Some(stored.id.clone()); + if let Some(account) = data.accounts.iter_mut().find(|acc| acc.id == stored.id) { + touch_account(account, true); + stored = account.clone(); + } + } + + write_accounts_file(&path, &data)?; + Ok(stored) +} + +#[cfg(test)] +#[path = "auth_accounts_tests.rs"] +mod tests; diff --git a/codex-rs/login/src/auth_accounts_tests.rs b/codex-rs/login/src/auth_accounts_tests.rs new file mode 100644 index 000000000000..3b03dba086a3 --- /dev/null +++ b/codex-rs/login/src/auth_accounts_tests.rs @@ -0,0 +1,552 @@ +use super::*; +use crate::auth_profiles::record_auth_profile_login; +use crate::token_data::IdTokenInfo; +use base64::Engine; +use pretty_assertions::assert_eq; +use serde::Serialize; +use tempfile::TempDir; + +fn make_chatgpt_tokens(account_id: Option<&str>, email: Option<&str>) -> TokenData { + fn fake_jwt(account_id: Option<&str>, email: Option<&str>) -> String { + #[derive(Serialize)] + struct Header { + alg: &'static str, + typ: &'static str, + } + + let header = Header { + alg: "none", + typ: "JWT", + }; + let payload = serde_json::json!({ + "email": email, + "https://api.openai.com/auth": { + "chatgpt_account_id": account_id.unwrap_or("acct"), + "chatgpt_user_id": "user-12345", + "user_id": "user-12345" + } + }); + let b64 = |value: &serde_json::Value| { + base64::engine::general_purpose::URL_SAFE_NO_PAD + .encode(serde_json::to_vec(value).expect("json to vec")) + }; + let header_b64 = b64(&serde_json::to_value(header).expect("header value")); + let payload_b64 = b64(&payload); + let signature_b64 = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(b"sig"); + format!("{header_b64}.{payload_b64}.{signature_b64}") + } + + TokenData { + id_token: IdTokenInfo { + email: email.map(str::to_string), + chatgpt_plan_type: None, + chatgpt_user_id: Some("user-12345".to_string()), + chatgpt_account_id: account_id.map(str::to_string), + chatgpt_account_is_fedramp: false, + raw_jwt: fake_jwt(account_id, email), + }, + access_token: "access".to_string(), + refresh_token: "refresh".to_string(), + account_id: account_id.map(str::to_string), + } +} + +fn make_chatgpt_tokens_with_claim_only_account_id( + account_id: Option<&str>, + email: Option<&str>, +) -> TokenData { + let mut tokens = make_chatgpt_tokens(account_id, email); + tokens.account_id = None; + tokens +} + +#[test] +fn missing_accounts_file_defaults_to_empty_state() { + let temp = TempDir::new().expect("tempdir"); + + assert_eq!( + Vec::::new(), + list_accounts(temp.path()).expect("list accounts") + ); + assert_eq!( + None, + get_active_account_id(temp.path()).expect("active account id") + ); +} + +#[test] +fn empty_accounts_file_defaults_to_empty_state() { + let temp = TempDir::new().expect("tempdir"); + fs::write(accounts_file_path(temp.path()), "\n \t").expect("write empty accounts file"); + + assert_eq!( + Vec::::new(), + list_accounts(temp.path()).expect("list accounts") + ); + assert_eq!( + None, + get_active_account_id(temp.path()).expect("active account id") + ); +} + +#[test] +fn upsert_api_key_creates_dedupes_and_sets_active() { + let temp = TempDir::new().expect("tempdir"); + let first = upsert_api_key_account( + temp.path(), + "sk-test".to_string(), + Some("Work".to_string()), + /*make_active*/ true, + ) + .expect("upsert api key"); + + let second = upsert_api_key_account( + temp.path(), + "sk-test".to_string(), + Some("Updated".to_string()), + /*make_active*/ false, + ) + .expect("upsert same key"); + + assert_eq!(first.id, second.id); + assert_eq!(Some("Updated"), second.label.as_deref()); + assert_eq!( + Some(first.id.as_str()), + get_active_account_id(temp.path()) + .expect("active id") + .as_deref() + ); + + let mut expected = second; + expected.last_used_at = first.last_used_at; + assert_eq!( + vec![expected], + list_accounts(temp.path()).expect("list accounts") + ); +} + +#[test] +fn upsert_chatgpt_dedupes_by_account_id_and_email() { + let temp = TempDir::new().expect("tempdir"); + let first = upsert_chatgpt_account( + temp.path(), + make_chatgpt_tokens(Some("acct-1"), Some("USER@example.com")), + Utc::now(), + /*label*/ None, + /*make_active*/ true, + ) + .expect("insert chatgpt"); + + let second = upsert_chatgpt_account( + temp.path(), + make_chatgpt_tokens(Some("acct-1"), Some("user@example.com")), + Utc::now(), + Some("Personal".to_string()), + /*make_active*/ false, + ) + .expect("update chatgpt"); + + assert_eq!(first.id, second.id); + assert_eq!(Some("Personal"), second.label.as_deref()); + assert_eq!(1, list_accounts(temp.path()).expect("list accounts").len()); +} + +#[test] +fn upsert_chatgpt_dedupes_by_id_token_account_id_without_email() { + let temp = TempDir::new().expect("tempdir"); + let first = upsert_chatgpt_account( + temp.path(), + make_chatgpt_tokens_with_claim_only_account_id(Some("acct-1"), None), + Utc::now(), + /*label*/ None, + /*make_active*/ true, + ) + .expect("insert chatgpt"); + + let second_tokens = make_chatgpt_tokens_with_claim_only_account_id(Some("acct-1"), None); + let second = upsert_chatgpt_account( + temp.path(), + second_tokens.clone(), + Utc::now(), + Some("Workspace".to_string()), + /*make_active*/ false, + ) + .expect("update chatgpt"); + + assert_eq!(first.id, second.id); + assert_eq!(1, list_accounts(temp.path()).expect("list accounts").len()); + assert_eq!( + Some(second), + remove_account_matching_credentials( + temp.path(), + AuthMode::Chatgpt, + /*openai_api_key*/ None, + Some(&second_tokens), + ) + .expect("remove by claim-only account id") + ); +} + +#[test] +fn chatgpt_accounts_with_same_email_but_different_ids_are_distinct() { + let temp = TempDir::new().expect("tempdir"); + let personal = upsert_chatgpt_account( + temp.path(), + make_chatgpt_tokens(Some("acct-personal"), Some("user@example.com")), + Utc::now(), + /*label*/ None, + /*make_active*/ true, + ) + .expect("insert personal account"); + let team = upsert_chatgpt_account( + temp.path(), + make_chatgpt_tokens(Some("acct-team"), Some("user@example.com")), + Utc::now(), + /*label*/ None, + /*make_active*/ false, + ) + .expect("insert team account"); + + assert_ne!(personal.id, team.id); + assert_eq!(2, list_accounts(temp.path()).expect("list accounts").len()); +} + +#[test] +fn set_active_account_id_records_and_touches_account() { + let temp = TempDir::new().expect("tempdir"); + let stored = upsert_api_key_account( + temp.path(), + "sk-test".to_string(), + /*label*/ None, + /*make_active*/ false, + ) + .expect("upsert api key"); + assert_eq!(None, stored.last_used_at); + + let activated = set_active_account_id(temp.path(), Some(stored.id.clone())) + .expect("set active") + .expect("activated account"); + + assert_eq!(stored.id, activated.id); + assert!(activated.last_used_at.is_some()); + assert_eq!( + Some(stored.id.as_str()), + get_active_account_id(temp.path()) + .expect("active id") + .as_deref() + ); +} + +#[test] +fn set_active_account_id_does_not_persist_missing_id() { + let temp = TempDir::new().expect("tempdir"); + let stored = upsert_api_key_account( + temp.path(), + "sk-test".to_string(), + /*label*/ None, + /*make_active*/ true, + ) + .expect("upsert api key"); + + assert_eq!( + None, + set_active_account_id(temp.path(), Some("missing".to_string())) + .expect("set missing active") + ); + + assert_eq!(None, get_active_account_id(temp.path()).expect("active id")); + assert_eq!( + vec![stored], + list_accounts(temp.path()).expect("list accounts") + ); +} + +#[test] +fn remove_account_clears_active() { + let temp = TempDir::new().expect("tempdir"); + let stored = upsert_chatgpt_account( + temp.path(), + make_chatgpt_tokens(Some("acct-remove"), Some("user@example.com")), + Utc::now(), + /*label*/ None, + /*make_active*/ true, + ) + .expect("insert chatgpt"); + + assert_eq!( + Some(stored.id.as_str()), + get_active_account_id(temp.path()) + .expect("active id") + .as_deref() + ); + + assert_eq!( + Some(stored.clone()), + remove_account(temp.path(), &stored.id).expect("remove") + ); + assert_eq!(None, get_active_account_id(temp.path()).expect("active id")); + assert_eq!( + Vec::::new(), + list_accounts(temp.path()).expect("list accounts") + ); +} + +#[test] +fn remove_account_promotes_remaining_account_when_active_is_removed() { + let temp = TempDir::new().expect("tempdir"); + let active = upsert_api_key_account( + temp.path(), + "sk-active".to_string(), + /*label*/ None, + /*make_active*/ true, + ) + .expect("insert active account"); + let fallback = upsert_api_key_account( + temp.path(), + "sk-fallback".to_string(), + /*label*/ None, + /*make_active*/ false, + ) + .expect("insert fallback account"); + + let active_id = active.id.clone(); + assert_eq!( + Some(active), + remove_account(temp.path(), &active_id).expect("remove active account") + ); + assert_eq!( + Some(fallback.id.as_str()), + get_active_account_id(temp.path()) + .expect("active id") + .as_deref() + ); + let promoted = find_account(temp.path(), &fallback.id) + .expect("find promoted account") + .expect("promoted account"); + assert!(promoted.last_used_at.is_some()); +} + +#[test] +fn remove_account_matching_credentials_removes_api_key_or_chatgpt_account() { + let temp = TempDir::new().expect("tempdir"); + let api = upsert_api_key_account( + temp.path(), + "sk-test".to_string(), + /*label*/ None, + /*make_active*/ true, + ) + .expect("insert api account"); + let tokens = make_chatgpt_tokens(Some("acct-chatgpt"), Some("user@example.com")); + let chatgpt = upsert_chatgpt_account( + temp.path(), + tokens.clone(), + Utc::now(), + /*label*/ None, + /*make_active*/ false, + ) + .expect("insert chatgpt account"); + + assert_eq!( + Some(api), + remove_account_matching_credentials( + temp.path(), + AuthMode::ApiKey, + Some("sk-test"), + /*tokens*/ None, + ) + .expect("remove api key") + ); + assert_eq!( + Some(chatgpt.id.as_str()), + get_active_account_id(temp.path()) + .expect("active id") + .as_deref() + ); + let removed_chatgpt = remove_account_matching_credentials( + temp.path(), + AuthMode::Chatgpt, + /*openai_api_key*/ None, + Some(&tokens), + ) + .expect("remove chatgpt") + .expect("removed chatgpt"); + assert_eq!(chatgpt.id, removed_chatgpt.id); + assert_eq!( + Vec::::new(), + list_accounts(temp.path()).expect("list accounts") + ); +} + +#[test] +fn remove_account_matching_credentials_promotes_remaining_account() { + let temp = TempDir::new().expect("tempdir"); + let active_tokens = make_chatgpt_tokens(Some("acct-active"), Some("active@example.com")); + let active = upsert_chatgpt_account( + temp.path(), + active_tokens.clone(), + Utc::now(), + /*label*/ None, + /*make_active*/ true, + ) + .expect("insert active chatgpt account"); + let fallback = upsert_api_key_account( + temp.path(), + "sk-fallback".to_string(), + /*label*/ None, + /*make_active*/ false, + ) + .expect("insert fallback account"); + + assert_eq!( + Some(active), + remove_account_matching_credentials( + temp.path(), + AuthMode::Chatgpt, + /*openai_api_key*/ None, + Some(&active_tokens), + ) + .expect("remove active chatgpt") + ); + assert_eq!( + Some(fallback.id.as_str()), + get_active_account_id(temp.path()) + .expect("active id") + .as_deref() + ); + let promoted = find_account(temp.path(), &fallback.id) + .expect("find promoted account") + .expect("promoted account"); + assert!(promoted.last_used_at.is_some()); +} + +#[test] +fn update_account_last_refresh_updates_only_target_account() { + let temp = TempDir::new().expect("tempdir"); + let first = upsert_chatgpt_account( + temp.path(), + make_chatgpt_tokens(Some("acct-first"), Some("first@example.com")), + Utc::now(), + /*label*/ None, + /*make_active*/ false, + ) + .expect("insert first"); + let second = upsert_chatgpt_account( + temp.path(), + make_chatgpt_tokens(Some("acct-second"), Some("second@example.com")), + Utc::now(), + /*label*/ None, + /*make_active*/ false, + ) + .expect("insert second"); + let refresh = Utc::now(); + + let updated = update_account_last_refresh(temp.path(), &second.id, refresh) + .expect("update refresh") + .expect("updated account"); + + assert_eq!(Some(refresh), updated.last_refresh); + assert_eq!( + None, + update_account_last_refresh(temp.path(), "missing", refresh).expect("missing update") + ); + let first_after = find_account(temp.path(), &first.id) + .expect("find first") + .expect("first account"); + assert_eq!(first.last_refresh, first_after.last_refresh); +} + +#[test] +fn account_store_ignores_existing_auth_profiles_until_import_exists() { + let temp = TempDir::new().expect("tempdir"); + record_auth_profile_login( + temp.path(), + "work", + Some("acct-profile".to_string()), + Some("profile@example.com".to_string()), + ) + .expect("record auth profile login"); + + assert_eq!( + Vec::::new(), + list_accounts(temp.path()).expect("list accounts") + ); + assert_eq!(None, get_active_account_id(temp.path()).expect("active id")); +} + +#[test] +fn recovers_from_trailing_json_documents_by_keeping_latest_accounts_file() { + let temp = TempDir::new().expect("tempdir"); + let path = accounts_file_path(temp.path()); + + let first = AccountsFile { + version: default_version(), + active_account_id: Some("first-active".to_string()), + accounts: vec![StoredAccount { + id: "first-active".to_string(), + mode: AuthMode::ApiKey, + label: Some("first".to_string()), + openai_api_key: Some("sk-first".to_string()), + tokens: None, + last_refresh: None, + created_at: None, + last_used_at: None, + }], + }; + let second = AccountsFile { + version: default_version(), + active_account_id: Some("second-active".to_string()), + accounts: vec![StoredAccount { + id: "second-active".to_string(), + mode: AuthMode::ApiKey, + label: Some("second".to_string()), + openai_api_key: Some("sk-second".to_string()), + tokens: None, + last_refresh: None, + created_at: None, + last_used_at: None, + }], + }; + + let first_json = serde_json::to_string_pretty(&first).expect("serialize first"); + let second_json = serde_json::to_string_pretty(&second).expect("serialize second"); + fs::write(&path, format!("{first_json}\n{second_json}\n")) + .expect("write corrupt accounts file"); + + assert_eq!( + second.accounts, + list_accounts(temp.path()).expect("recover accounts") + ); + assert_eq!( + Some("second-active"), + get_active_account_id(temp.path()) + .expect("active id") + .as_deref() + ); + assert_eq!( + second_json, + fs::read_to_string(&path).expect("read repaired accounts file") + ); +} + +#[cfg(unix)] +#[test] +fn saved_accounts_file_is_private() { + use std::os::unix::fs::PermissionsExt; + + let temp = TempDir::new().expect("tempdir"); + + upsert_api_key_account( + temp.path(), + "sk-test".to_string(), + /*label*/ None, + /*make_active*/ true, + ) + .expect("upsert api account"); + + let mode = fs::metadata(accounts_file_path(temp.path())) + .expect("metadata") + .permissions() + .mode() + & 0o777; + assert_eq!(0o600, mode); +} diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index 5825c910c8f7..bad5599f18a5 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -1,4 +1,5 @@ pub mod auth; +pub mod auth_accounts; pub mod auth_env_telemetry; pub mod auth_profiles; pub mod token_data; @@ -46,6 +47,16 @@ pub use auth::logout_with_revoke; pub use auth::read_codex_access_token_from_env; pub use auth::read_openai_api_key_from_env; pub use auth::save_auth; +pub use auth_accounts::StoredAccount; +pub use auth_accounts::find_account; +pub use auth_accounts::get_active_account_id; +pub use auth_accounts::list_accounts; +pub use auth_accounts::remove_account; +pub use auth_accounts::remove_account_matching_credentials; +pub use auth_accounts::set_active_account_id; +pub use auth_accounts::update_account_last_refresh; +pub use auth_accounts::upsert_api_key_account; +pub use auth_accounts::upsert_chatgpt_account; pub use auth_env_telemetry::AuthEnvTelemetry; pub use auth_env_telemetry::collect_auth_env_telemetry; pub use auth_profiles::AuthProfileEntry;