From 8076907c7bb5429243b1c5b1fc0ae295dd8a4708 Mon Sep 17 00:00:00 2001 From: blazz Date: Wed, 5 Aug 2026 08:28:50 -0400 Subject: [PATCH 1/3] feat(codex): port multi-account Codex domain core Port the Windows multi-account Codex core from ademisler/codexcontrol (MIT) into rust/src/codex_accounts: account manager, quota API, Codex Desktop restart, versioned stores, and models. MIT attribution added to NOTICE. --- NOTICE | 32 + rust/src/codex_accounts/account_manager.rs | 1034 ++++++++++++++++++++ rust/src/codex_accounts/api.rs | 853 ++++++++++++++++ rust/src/codex_accounts/codex_desktop.rs | 407 ++++++++ rust/src/codex_accounts/file_locations.rs | 178 ++++ rust/src/codex_accounts/mod.rs | 34 + rust/src/codex_accounts/models.rs | 669 +++++++++++++ rust/src/codex_accounts/stores.rs | 266 +++++ rust/src/lib.rs | 1 + 9 files changed, 3474 insertions(+) create mode 100644 NOTICE create mode 100644 rust/src/codex_accounts/account_manager.rs create mode 100644 rust/src/codex_accounts/api.rs create mode 100644 rust/src/codex_accounts/codex_desktop.rs create mode 100644 rust/src/codex_accounts/file_locations.rs create mode 100644 rust/src/codex_accounts/mod.rs create mode 100644 rust/src/codex_accounts/models.rs create mode 100644 rust/src/codex_accounts/stores.rs diff --git a/NOTICE b/NOTICE new file mode 100644 index 0000000000..64a81fc221 --- /dev/null +++ b/NOTICE @@ -0,0 +1,32 @@ +NOTICE + +This software includes portions derived from the +["codexcontrol"](https://github.com/ademisler/codexcontrol) project (MIT), +specifically the Windows multi-account Codex core ported to Rust under +`rust/src/codex_accounts/`. + +The derived code retains the upstream copyright and license notice: + + MIT License + + Copyright (c) 2026 Adem Isler + + Permission is hereby granted, free of charge, to any person obtaining a copy + of this software and associated documentation files (the "Software"), to deal + in the Software without restriction, including without limitation the rights + to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + copies of the Software, and to permit persons to whom the Software is + furnished to do so, subject to the following conditions: + + The above copyright notice and this permission notice shall be included in all + copies or substantial portions of the Software. + + THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + SOFTWARE. + +The original project URL: https://github.com/ademisler/codexcontrol \ No newline at end of file diff --git a/rust/src/codex_accounts/account_manager.rs b/rust/src/codex_accounts/account_manager.rs new file mode 100644 index 0000000000..52ad75e190 --- /dev/null +++ b/rust/src/codex_accounts/account_manager.rs @@ -0,0 +1,1034 @@ +//! Codex account management: discovery, authentication, switching, removal. +//! +//! Port of `windows/.../account_manager.py` (MIT). Manages isolated managed +//! homes under `managed-homes/`, discovers the ambient `~/.codex` identity, and +//! switches the active identity by swapping `auth.json` into the ambient home, +//! rewriting the Codex Desktop `creator_id` global state and backing up/restoring +//! the desktop session. + +use std::fs; +use std::path::{Path, PathBuf}; +use std::process::{Child, Command, Stdio}; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +use chrono::{DateTime, Utc}; +use thiserror::Error; +use uuid::Uuid; + +use super::api::{AuthBackedIdentity, CodexApiError, load_identity}; +use super::file_locations::{ + ambient_codex_home, auth_backups_directory, codex_desktop_session_root, + desktop_session_snapshot_path, ensure_directories, managed_homes_directory, +}; +use super::models::{CodexAccount, CodexAccountSource, utc_now}; + +/// Friendly account manager error. +#[derive(Debug, Error)] +pub enum CodexAccountManagerError { + #[error("{0}")] + Message(String), + #[error(transparent)] + Io(#[from] std::io::Error), +} + +impl From for CodexAccountManagerError { + fn from(value: CodexApiError) -> Self { + CodexAccountManagerError::Message(value.to_string()) + } +} + +/// Outcome of a `codex login` subprocess run. +#[derive(Debug, Clone)] +pub enum CodexLoginOutcome { + MissingBinary, + LaunchFailed(String), + TimedOut(String), + Cancelled, + Failed(String), + Success(String), +} + +impl CodexLoginOutcome { + pub fn as_str(&self) -> &'static str { + match self { + CodexLoginOutcome::MissingBinary => "missing_binary", + CodexLoginOutcome::LaunchFailed(_) => "launch_failed", + CodexLoginOutcome::TimedOut(_) => "timed_out", + CodexLoginOutcome::Cancelled => "cancelled", + CodexLoginOutcome::Failed(_) => "failed", + CodexLoginOutcome::Success(_) => "success", + } + } + + pub fn output(&self) -> &str { + match self { + CodexLoginOutcome::MissingBinary => "", + CodexLoginOutcome::LaunchFailed(output) + | CodexLoginOutcome::TimedOut(output) + | CodexLoginOutcome::Failed(output) + | CodexLoginOutcome::Success(output) => output, + CodexLoginOutcome::Cancelled => "", + } + } +} + +/// Result of a `codex login` subprocess run. +#[derive(Debug, Clone)] +pub struct CodexLoginResult { + pub outcome: CodexLoginOutcome, +} + +/// Handle around an in-flight `codex login` process, for cancellation. +#[derive(Debug, Default, Clone)] +pub struct ManagedLoginProcess { + inner: Arc>>, + cancelled: Arc, +} + +impl ManagedLoginProcess { + fn bind(&self, process: Child) { + *self.inner.lock().expect("login process lock") = Some(process); + self.cancelled.store(false, Ordering::SeqCst); + } + + fn clear(&self) { + *self.inner.lock().expect("login process lock") = None; + } + + pub fn is_cancelled(&self) -> bool { + self.cancelled.load(Ordering::SeqCst) + } + + pub fn cancel(&self) { + self.cancelled.store(true, Ordering::SeqCst); + let mut guard = self.inner.lock().expect("login process lock"); + if let Some(child) = guard.as_mut() { + let _ = child.kill(); + } + } +} + +/// Runs `codex login` inside an isolated `CODEX_HOME`. +pub struct CodexLoginRunner; + +impl CodexLoginRunner { + /// Resolve the `codex` executable, falling back to known install paths. + pub fn locate_codex_binary() -> Option { + if let Ok(found) = which::which("codex") { + return Some(found); + } + path_candidates() + .into_iter() + .find(|candidate| candidate.is_file()) + } + + pub fn run( + home_path: &Path, + timeout: Duration, + handle: Option<&ManagedLoginProcess>, + ) -> CodexLoginResult { + let active_handle = handle.cloned().unwrap_or_default(); + let Some(binary) = Self::locate_codex_binary() else { + return CodexLoginResult { + outcome: CodexLoginOutcome::MissingBinary, + }; + }; + + let mut command = Command::new(binary); + command + .arg("login") + .env("CODEX_HOME", home_path) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()); + let child = match command.spawn() { + Ok(child) => child, + Err(error) => { + return CodexLoginResult { + outcome: CodexLoginOutcome::LaunchFailed(error.to_string()), + }; + } + }; + active_handle.bind(child); + + let output = match wait_for_child(&active_handle, timeout) { + Some(output) => output, + None => { + let output = kill_and_drain(&active_handle); + active_handle.clear(); + return CodexLoginResult { + outcome: CodexLoginOutcome::TimedOut(combine_output(&output)), + }; + } + }; + + active_handle.clear(); + let combined = combine_output(&output); + if active_handle.is_cancelled() { + return CodexLoginResult { + outcome: CodexLoginOutcome::Cancelled, + }; + } + if output.status.success() { + return CodexLoginResult { + outcome: CodexLoginOutcome::Success(combined), + }; + } + CodexLoginResult { + outcome: CodexLoginOutcome::Failed(combined), + } + } +} + +fn path_candidates() -> Vec { + let local_app_data = std::env::var("LOCALAPPDATA") + .map(PathBuf::from) + .unwrap_or_else(|_| { + dirs::home_dir() + .unwrap_or_else(|| PathBuf::from(".")) + .join("AppData") + .join("Local") + }); + let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from(".")); + vec![ + local_app_data + .join("OpenAI") + .join("Codex") + .join("bin") + .join("codex.exe"), + home.join(".bun").join("bin").join("codex.exe"), + local_app_data + .join("Microsoft") + .join("WindowsApps") + .join("codex.exe"), + ] +} + +fn wait_for_child(handle: &ManagedLoginProcess, timeout: Duration) -> Option { + let deadline = Instant::now() + timeout; + loop { + if handle.is_cancelled() { + let output = take_child(handle)?.wait_with_output().ok(); + return output; + } + let polled = { + let mut guard = handle.inner.lock().expect("login process lock"); + match guard.as_mut().map(|child| child.try_wait()) { + Some(Ok(Some(_status))) => take_child(handle)?.wait_with_output().ok(), + Some(Err(_)) => take_child(handle)?.wait_with_output().ok(), + _ => None, + } + }; + if polled.is_some() { + return polled; + } + if Instant::now() >= deadline { + return None; + } + std::thread::sleep(Duration::from_millis(100)); + } +} + +fn take_child(handle: &ManagedLoginProcess) -> Option { + handle.inner.lock().expect("login process lock").take() +} + +fn kill_and_drain(handle: &ManagedLoginProcess) -> std::process::Output { + let mut child = take_child(handle).expect("login process present"); + let _ = child.kill(); + child + .wait_with_output() + .unwrap_or_else(|_| std::process::Output { + status: std::process::ExitStatus::default(), + stdout: Vec::new(), + stderr: Vec::new(), + }) +} + +fn combine_output(output: &std::process::Output) -> String { + let mut parts: Vec = Vec::new(); + for bytes in [&output.stdout, &output.stderr] { + let text = String::from_utf8_lossy(bytes); + let trimmed = text.trim(); + if !trimmed.is_empty() { + parts.push(trimmed.to_string()); + } + } + let merged = parts.join("\n"); + let merged = merged.trim(); + if merged.is_empty() { + "No output captured.".to_string() + } else { + merged.chars().take(4000).collect() + } +} + +/// Result of switching the active account. +#[derive(Debug, Clone, serde::Serialize)] +#[serde(rename_all = "camelCase")] +pub struct CodexSwitchResult { + pub materialized_account: Option, + pub backup_path: Option, + pub ambient_account: Option, + pub desktop_session_backup_path: Option, + pub desktop_session_restore_path: Option, + pub desktop_session_restore_exists: bool, +} + +/// Discovers, authenticates and switches Codex accounts. +#[derive(Debug, Default)] +pub struct CodexAccountManager; + +impl CodexAccountManager { + pub fn new() -> Self { + Self + } + + /// Start a `codex login` into a fresh managed home. + pub fn add_managed_account( + &self, + handle: Option<&ManagedLoginProcess>, + ) -> Result { + ensure_directories()?; + let home_path = managed_homes_directory().join(Uuid::new_v4().to_string()); + fs::create_dir_all(&home_path)?; + + match self.authenticate_account(&home_path, CodexAccountSource::ManagedByApp, None, handle) + { + Ok(account) => Ok(account), + Err(error) => { + let _ = fs::remove_dir_all(&home_path); + Err(error) + } + } + } + + /// Re-run `codex login` for an existing account. + pub fn reauthenticate( + &self, + account: &CodexAccount, + handle: Option<&ManagedLoginProcess>, + ) -> Result { + self.authenticate_account( + &account.codex_home_path, + account.source, + Some(account), + handle, + ) + } + + /// Remove app-owned managed homes matching this account. + pub fn remove_managed_files_if_owned( + &self, + account: &CodexAccount, + ) -> Result<(), CodexAccountManagerError> { + if !account.source.owns_files() { + return Ok(()); + } + + let root = fs::canonicalize(managed_homes_directory()) + .unwrap_or_else(|_| managed_homes_directory()); + let targets = self.managed_home_paths_matching(account)?; + + for target in targets { + let resolved = fs::canonicalize(&target).unwrap_or_else(|_| target.clone()); + let relative = resolved.strip_prefix(&root).map_err(|_| { + CodexAccountManagerError::Message( + "This path is not an app-managed home directory.".to_string(), + ) + })?; + if relative.as_os_str().is_empty() { + return Err(CodexAccountManagerError::Message( + "Refusing to remove the managed-homes root.".to_string(), + )); + } + if target.exists() { + fs::remove_dir_all(&target)?; + } + } + Ok(()) + } + + /// Discover managed homes and merge them against the stored accounts. + pub fn discover_managed_accounts( + &self, + existing: &[CodexAccount], + ) -> Result, CodexAccountManagerError> { + ensure_directories()?; + let mut discovered = Vec::new(); + let mut entries: Vec = fs::read_dir(managed_homes_directory())? + .filter_map(|entry| entry.ok().map(|entry| entry.path())) + .filter(|path| path.is_dir()) + .collect(); + entries.sort_by_key(|path| { + path.file_name() + .map(|name| name.to_string_lossy().to_lowercase()) + .unwrap_or_default() + }); + for home_path in entries { + if let Some(account) = self.discovered_managed_account(&home_path, existing) { + discovered.push(account); + } + } + Ok(discovered) + } + + /// Discover the ambient `~/.codex` account. + pub fn discover_ambient_account(&self, existing: &[CodexAccount]) -> Option { + let home_path = ambient_codex_home(); + let auth_path = home_path.join("auth.json"); + if !home_path.is_dir() || !auth_path.exists() { + return None; + } + let identity = load_identity(&home_path).ok()?; + if identity.email.is_none() && identity.provider_account_id.is_none() { + return None; + } + + let candidate = + candidate_account(identity.clone(), &home_path, CodexAccountSource::Ambient); + let matched = existing.iter().find(|account| candidate.matches(account)); + let discovered_at = directory_timestamp(&home_path); + Some(build_discovered_account( + matched, + identity, + home_path, + CodexAccountSource::Ambient, + discovered_at, + )) + } + + /// Identity of the currently active (ambient) account, if any. + pub fn load_active_identity(&self) -> Option { + let auth_path = ambient_codex_home().join("auth.json"); + if !auth_path.exists() { + return None; + } + load_identity(&ambient_codex_home()).ok() + } + + /// Switch the ambient identity to `target`, materializing the previous + /// ambient account as managed and preserving the desktop session. + pub fn switch_active_account( + &self, + target: &CodexAccount, + existing: &[CodexAccount], + ) -> Result { + ensure_directories()?; + + let target_auth_path = target.codex_home_path.join("auth.json"); + if !target_auth_path.exists() { + return Err(CodexAccountManagerError::Message( + "The selected account does not contain `auth.json`.".to_string(), + )); + } + + let ambient_account = self.discover_ambient_account(existing); + let session_root = codex_desktop_session_root(); + let mut materialized_account: Option = None; + if let Some(ambient) = &ambient_account { + let is_ambient = ambient.source == CodexAccountSource::Ambient; + if is_ambient && !ambient.matches(target) { + materialized_account = Some(self.materialize_as_managed(ambient)?); + } + } + + let mut desktop_session_backup_path: Option = None; + let mut desktop_session_restore_path: Option = None; + let mut desktop_session_restore_exists = false; + if session_root.is_some() { + if let Some(materialized) = &materialized_account { + desktop_session_backup_path = + Some(desktop_session_snapshot_path(&materialized.codex_home_path)); + } + let snapshot_path = desktop_session_snapshot_path(&target.codex_home_path); + desktop_session_restore_path = Some(snapshot_path.clone()); + desktop_session_restore_exists = path_has_children(&snapshot_path); + } + + fs::create_dir_all(ambient_codex_home())?; + let backup_path = self.backup_ambient_auth()?; + fs::copy(&target_auth_path, ambient_codex_home().join("auth.json"))?; + self.sync_ambient_global_state( + ambient_account + .as_ref() + .and_then(|account| account.provider_account_id.clone()), + self.target_account_id(target)?, + ); + + Ok(CodexSwitchResult { + materialized_account, + backup_path, + ambient_account: self.discover_ambient_account(existing), + desktop_session_backup_path, + desktop_session_restore_path, + desktop_session_restore_exists, + }) + } + + /// Copy the ambient account into an app-managed home. + pub fn materialize_as_managed( + &self, + account: &CodexAccount, + ) -> Result { + ensure_directories()?; + + let source_auth_path = account.codex_home_path.join("auth.json"); + if !source_auth_path.exists() { + return Err(CodexAccountManagerError::Message( + "The current active account does not contain `auth.json`.".to_string(), + )); + } + + let destination_home = managed_homes_directory().join(Uuid::new_v4().to_string()); + fs::create_dir_all(&destination_home)?; + fs::copy(&source_auth_path, destination_home.join("auth.json"))?; + + let now = utc_now(); + Ok(CodexAccount::new( + account.id, + account.nickname.clone(), + account.email_hint.clone(), + account.auth_subject.clone(), + account.provider_account_id.clone(), + destination_home, + CodexAccountSource::ManagedByApp, + account.created_at, + now, + Some(account.last_authenticated_at.unwrap_or(now)), + )) + } + + fn backup_ambient_auth(&self) -> Result, CodexAccountManagerError> { + ensure_directories()?; + let auth_path = ambient_codex_home().join("auth.json"); + if !auth_path.exists() { + return Ok(None); + } + let backup_path = + auth_backups_directory().join(format!("ambient-auth-{}.json", timestamp_slug())); + fs::copy(&auth_path, &backup_path)?; + Ok(Some(backup_path)) + } + + fn target_account_id(&self, target: &CodexAccount) -> Result, CodexApiError> { + if let Some(account_id) = &target.provider_account_id { + return Ok(Some(account_id.clone())); + } + let identity = load_identity(&target.codex_home_path)?; + Ok(identity.provider_account_id) + } + + fn sync_ambient_global_state( + &self, + previous_account_id: Option, + target_account_id: Option, + ) { + let Some(target_account_id) = target_account_id else { + return; + }; + for file_name in [".codex-global-state.json", ".codex-global-state.json.bak"] { + self.rewrite_creator_id( + &ambient_codex_home().join(file_name), + previous_account_id.as_deref(), + &target_account_id, + ); + } + } + + fn rewrite_creator_id( + &self, + path: &Path, + previous_account_id: Option<&str>, + target_account_id: &str, + ) { + if !path.exists() { + return; + } + let Ok(content) = fs::read_to_string(path) else { + return; + }; + let Ok(mut payload) = serde_json::from_str::(&content) else { + return; + }; + if !payload.is_object() { + return; + } + let Some(atom_state) = payload + .get_mut("electron-persisted-atom-state") + .and_then(|value| value.as_object_mut()) + else { + return; + }; + let Some(environment) = atom_state + .get_mut("environment") + .and_then(|value| value.as_object_mut()) + else { + return; + }; + let Some(creator_id) = environment.get("creator_id") else { + return; + }; + let Some(updated) = + updated_creator_id(creator_id.as_str(), previous_account_id, target_account_id) + else { + return; + }; + if updated == creator_id.as_str().unwrap_or_default() { + return; + } + environment.insert("creator_id".to_string(), serde_json::Value::String(updated)); + let Ok(encoded) = serde_json::to_string_pretty(&payload) else { + return; + }; + let _ = fs::write(path, format!("{encoded}\n")); + } + + fn managed_home_paths_matching( + &self, + account: &CodexAccount, + ) -> Result, CodexAccountManagerError> { + ensure_directories()?; + let mut targets: Vec = vec![ + std::path::absolute(&account.codex_home_path) + .unwrap_or_else(|_| account.codex_home_path.clone()), + ]; + let mut seen_keys: std::collections::HashSet = + [managed_home_key(targets[0].as_path())] + .into_iter() + .collect(); + + for entry in fs::read_dir(managed_homes_directory())? { + let Ok(entry) = entry else { + continue; + }; + let home_path = entry.path(); + if !home_path.is_dir() { + continue; + } + let Some(candidate) = + self.discovered_managed_account(&home_path, std::slice::from_ref(account)) + else { + continue; + }; + if !candidate.matches(account) { + continue; + } + let resolved = std::path::absolute(&home_path).unwrap_or_else(|_| home_path.clone()); + let key = managed_home_key(resolved.as_path()); + if seen_keys.contains(&key) { + continue; + } + targets.push(resolved); + seen_keys.insert(key); + } + Ok(targets) + } + + fn authenticate_account( + &self, + home_path: &Path, + source: CodexAccountSource, + existing: Option<&CodexAccount>, + handle: Option<&ManagedLoginProcess>, + ) -> Result { + let result = CodexLoginRunner::run(home_path, Duration::from_secs(180), handle); + + match &result.outcome { + CodexLoginOutcome::Cancelled => { + return Err(CodexAccountManagerError::Message( + "Account setup cancelled.".to_string(), + )); + } + CodexLoginOutcome::MissingBinary => { + return Err(CodexAccountManagerError::Message( + "The `codex` command could not be found.".to_string(), + )); + } + CodexLoginOutcome::TimedOut(_) => { + return Err(CodexAccountManagerError::Message( + "The Codex sign-in flow timed out.".to_string(), + )); + } + CodexLoginOutcome::LaunchFailed(output) => { + return Err(CodexAccountManagerError::Message(format!( + "Failed to start the Codex sign-in flow: {output}" + ))); + } + CodexLoginOutcome::Failed(output) => { + return Err(CodexAccountManagerError::Message(format!( + "The Codex sign-in flow did not complete.\n{output}" + ))); + } + CodexLoginOutcome::Success(_) => {} + } + + let identity = load_identity(home_path)?; + if identity.email.is_none() && identity.provider_account_id.is_none() { + return Err(CodexAccountManagerError::Message( + "Sign-in completed, but the account identity could not be read.".to_string(), + )); + } + + let now = utc_now(); + Ok(CodexAccount::new( + existing + .map(|account| account.id) + .unwrap_or_else(Uuid::new_v4), + existing.and_then(|account| account.nickname.clone()), + identity + .email + .or_else(|| existing.and_then(|account| account.email_hint.clone())), + identity + .auth_subject + .or_else(|| existing.and_then(|account| account.auth_subject.clone())), + identity + .provider_account_id + .or_else(|| existing.and_then(|account| account.provider_account_id.clone())), + home_path.to_path_buf(), + source, + existing.map(|account| account.created_at).unwrap_or(now), + now, + Some(now), + )) + } + + fn discovered_managed_account( + &self, + home_path: &Path, + existing: &[CodexAccount], + ) -> Option { + if !home_path.is_dir() { + return None; + } + let auth_path = home_path.join("auth.json"); + if !auth_path.exists() { + return None; + } + let identity = load_identity(home_path).ok()?; + if identity.email.is_none() && identity.provider_account_id.is_none() { + return None; + } + + let discovered_at = directory_timestamp(home_path); + let candidate = candidate_account( + identity.clone(), + home_path, + CodexAccountSource::ManagedByApp, + ); + let matched = existing.iter().find(|account| candidate.matches(account)); + Some(build_discovered_account( + matched, + identity, + home_path.to_path_buf(), + CodexAccountSource::ManagedByApp, + discovered_at, + )) + } +} + +fn candidate_account( + identity: AuthBackedIdentity, + home_path: &Path, + source: CodexAccountSource, +) -> CodexAccount { + CodexAccount::new( + Uuid::new_v4(), + None, + identity.email.clone(), + identity.auth_subject.clone(), + identity.provider_account_id.clone(), + home_path.to_path_buf(), + source, + utc_now(), + utc_now(), + None, + ) +} + +fn build_discovered_account( + matched: Option<&CodexAccount>, + identity: AuthBackedIdentity, + home_path: PathBuf, + source: CodexAccountSource, + discovered_at: DateTime, +) -> CodexAccount { + CodexAccount::new( + matched + .map(|account| account.id) + .unwrap_or_else(Uuid::new_v4), + matched.and_then(|account| account.nickname.clone()), + identity + .email + .or_else(|| matched.and_then(|account| account.email_hint.clone())), + identity + .auth_subject + .or_else(|| matched.and_then(|account| account.auth_subject.clone())), + identity + .provider_account_id + .or_else(|| matched.and_then(|account| account.provider_account_id.clone())), + home_path, + source, + matched + .map(|account| account.created_at) + .unwrap_or(discovered_at), + matched + .map(|account| account.updated_at.max(discovered_at)) + .unwrap_or(discovered_at), + matched + .and_then(|account| account.last_authenticated_at) + .or(Some(discovered_at)), + ) +} + +fn directory_timestamp(path: &Path) -> DateTime { + let auth_path = path.join("auth.json"); + if auth_path.exists() + && let Ok(metadata) = fs::metadata(&auth_path) + && let Ok(modified) = metadata.modified() + { + return modified.into(); + } + fs::metadata(path) + .and_then(|metadata| metadata.modified()) + .map(Into::into) + .unwrap_or_else(|_| utc_now()) +} + +fn managed_home_key(path: &Path) -> String { + std::path::absolute(path) + .unwrap_or_else(|_| path.to_path_buf()) + .to_string_lossy() + .to_lowercase() +} + +fn path_has_children(path: &Path) -> bool { + if !path.exists() || !path.is_dir() { + return false; + } + fs::read_dir(path) + .map(|mut entries| entries.next().is_some()) + .unwrap_or(false) +} + +fn timestamp_slug() -> String { + utc_now().format("%Y%m%d-%H%M%S").to_string() +} + +/// Compute the replacement `creator_id` for the target account. +fn updated_creator_id( + creator_id: Option<&str>, + previous_account_id: Option<&str>, + target_account_id: &str, +) -> Option { + let creator_id = creator_id?.trim(); + if creator_id.is_empty() { + return None; + } + if creator_id == target_account_id || creator_id.ends_with(&format!("__{target_account_id}")) { + return Some(creator_id.to_string()); + } + if let Some(previous) = previous_account_id + && creator_id.contains(previous) + { + return Some(creator_id.replace(previous, target_account_id)); + } + if looks_like_uuid(creator_id) { + return Some(target_account_id.to_string()); + } + if let Some((prefix, suffix)) = creator_id.rsplit_once("__") + && looks_like_uuid(suffix) + { + return Some(format!("{prefix}__{target_account_id}")); + } + None +} + +fn looks_like_uuid(value: &str) -> bool { + Uuid::parse_str(value.trim()).is_ok() +} + +#[cfg(test)] +mod tests { + use super::*; + use base64::Engine; + + /// Write an auth.json carrying a JWT identity for the given account id. + fn write_auth(home_path: &Path, email: &str, account_id: &str) { + let payload = serde_json::json!({ + "email": email, + "sub": format!("auth0|{account_id}"), + "https://api.openai.com/auth": { + "chatgpt_plan_type": "team", + "chatgpt_account_id": account_id, + }, + }); + let encoded = base64::engine::general_purpose::URL_SAFE_NO_PAD + .encode(serde_json::to_vec(&payload).unwrap()); + let auth_payload = serde_json::json!({ + "tokens": { + "access_token": format!("access-{account_id}"), + "refresh_token": format!("refresh-{account_id}"), + "id_token": format!("header.{encoded}.signature"), + "account_id": account_id, + }, + "last_refresh": "2026-04-23T00:00:00Z", + }); + std::fs::write( + home_path.join("auth.json"), + serde_json::to_vec_pretty(&auth_payload).unwrap(), + ) + .unwrap(); + } + + fn make_account(home_path: PathBuf, email: &str, account_id: &str) -> CodexAccount { + CodexAccount::new( + Uuid::new_v4(), + None, + Some(email.to_string()), + Some(format!("auth0|{account_id}")), + Some(account_id.to_string()), + home_path, + CodexAccountSource::ManagedByApp, + utc_now(), + utc_now(), + Some(utc_now()), + ) + } + + #[test] + fn remove_managed_account_removes_duplicate_homes_for_same_provider() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path(); + super::super::file_locations::with_app_support_directory(root.to_path_buf()); + + let account_id = "83c5ae92-f5ee-41f8-9528-199110d1d0f9"; + let first_home = root.join("managed-homes").join("first"); + let duplicate_home = root.join("managed-homes").join("duplicate"); + let other_home = root.join("managed-homes").join("other"); + for home in [&first_home, &duplicate_home, &other_home] { + std::fs::create_dir_all(home).unwrap(); + } + write_auth(&first_home, "user@example.com", account_id); + write_auth(&duplicate_home, "user@example.com", account_id); + write_auth(&other_home, "user@example.com", "different-provider"); + + let account = make_account(first_home.clone(), "user@example.com", account_id); + let manager = CodexAccountManager::new(); + manager.remove_managed_files_if_owned(&account).unwrap(); + + assert!(!first_home.exists()); + assert!(!duplicate_home.exists()); + assert!(other_home.exists()); + + super::super::file_locations::clear_app_support_directory_override(); + } + + #[test] + fn switch_active_account_updates_global_state_creator_id() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path(); + super::super::file_locations::with_app_support_directory(root.to_path_buf()); + + let old_account_id = "1ea93d04-5c50-42e3-857b-3db850785967"; + let new_account_id = "83c5ae92-f5ee-41f8-9528-199110d1d0f9"; + + let ambient_home = root.join(".codex"); + let target_home = root.join("managed-homes").join("target"); + let desktop_session_root = root.join("package-session"); + + std::fs::create_dir_all(&ambient_home).unwrap(); + std::fs::create_dir_all(&target_home).unwrap(); + std::fs::create_dir_all(&desktop_session_root).unwrap(); + + write_auth(&ambient_home, "old@example.com", old_account_id); + write_auth(&target_home, "new@example.com", new_account_id); + let target_session_dir = target_home.join("desktop-session").join("Network"); + std::fs::create_dir_all(&target_session_dir).unwrap(); + std::fs::write(target_session_dir.join("Cookies"), "cookie-data").unwrap(); + + let global_state = serde_json::json!({ + "electron-persisted-atom-state": { + "environment": { + "creator_id": format!("user-e9H3MsspGTF7UZJ8uaXuML55__{old_account_id}"), + } + } + }); + for file_name in [".codex-global-state.json", ".codex-global-state.json.bak"] { + std::fs::write( + ambient_home.join(file_name), + serde_json::to_vec_pretty(&global_state).unwrap(), + ) + .unwrap(); + } + + super::super::file_locations::with_ambient_codex_home(ambient_home.clone()); + super::super::file_locations::with_codex_desktop_session_root(desktop_session_root.clone()); + + let manager = CodexAccountManager::new(); + let target_account = make_account(target_home.clone(), "new@example.com", new_account_id); + let result = manager + .switch_active_account(&target_account, std::slice::from_ref(&target_account)) + .unwrap(); + + let ambient_auth: serde_json::Value = + serde_json::from_slice(&std::fs::read(ambient_home.join("auth.json")).unwrap()) + .unwrap(); + assert_eq!(ambient_auth["tokens"]["account_id"], new_account_id); + assert_eq!( + result + .ambient_account + .unwrap() + .provider_account_id + .as_deref(), + Some(new_account_id) + ); + let materialized = result.materialized_account.unwrap(); + assert_eq!( + materialized.provider_account_id.as_deref(), + Some(old_account_id) + ); + assert_eq!( + result.desktop_session_backup_path.unwrap(), + materialized.codex_home_path.join("desktop-session") + ); + assert_eq!( + result.desktop_session_restore_path.unwrap(), + target_home.join("desktop-session") + ); + assert!(result.desktop_session_restore_exists); + + let backup_files: Vec = std::fs::read_dir(root.join("auth-backups")) + .unwrap() + .map(|entry| entry.unwrap().path()) + .filter(|path| { + path.file_name() + .unwrap() + .to_string_lossy() + .starts_with("ambient-auth-") + }) + .collect(); + assert_eq!(backup_files.len(), 1); + let backup: serde_json::Value = + serde_json::from_slice(&std::fs::read(&backup_files[0]).unwrap()).unwrap(); + assert_eq!(backup["tokens"]["account_id"], old_account_id); + + for file_name in [".codex-global-state.json", ".codex-global-state.json.bak"] { + let payload: serde_json::Value = + serde_json::from_slice(&std::fs::read(ambient_home.join(file_name)).unwrap()) + .unwrap(); + let creator_id = payload["electron-persisted-atom-state"]["environment"]["creator_id"] + .as_str() + .unwrap(); + assert_eq!( + creator_id, + format!("user-e9H3MsspGTF7UZJ8uaXuML55__{new_account_id}") + ); + } + + super::super::file_locations::clear_app_support_directory_override(); + super::super::file_locations::clear_ambient_codex_home_override(); + super::super::file_locations::clear_codex_desktop_session_root_override(); + } +} diff --git a/rust/src/codex_accounts/api.rs b/rust/src/codex_accounts/api.rs new file mode 100644 index 0000000000..e2188a32ab --- /dev/null +++ b/rust/src/codex_accounts/api.rs @@ -0,0 +1,853 @@ +//! Codex API client: identity, OAuth refresh, quota fetch and recovery. +//! +//! Port of `windows/.../codex_api.py` (MIT). Reads a Codex home's `auth.json`, +//! refreshes tokens via the OpenAI OAuth endpoint, fetches `wham/usage` (or a +//! configured custom base URL) and normalizes the quota windows. + +use std::path::Path; + +use base64::Engine; +use chrono::{DateTime, Utc}; +use thiserror::Error; + +use super::models::{ + AccountUsageSnapshot, CreditsBalanceSnapshot, UsageWindowSnapshot, WindowRole, +}; +use crate::core::credentialed_http_client_builder; + +pub const REFRESH_ENDPOINT: &str = "https://auth.openai.com/oauth/token"; +pub const USAGE_DEFAULT_BASE: &str = "https://chatgpt.com/backend-api"; +pub const REFRESH_CLIENT_ID: &str = "app_EMoamEEZ73f0CkXaXp7hrann"; +const REQUEST_TIMEOUT_SECONDS: u64 = 30; +const UNAUTHORIZED_MESSAGE: &str = "The Codex usage API request returned unauthorized."; + +/// Friendly error surfaced to callers. +#[derive(Debug, Error)] +pub enum CodexApiError { + #[error("{0}")] + Message(String), + #[error("network error: {0}")] + Network(String), + #[error("failed to parse Codex payload: {0}")] + Parse(String), +} + +/// Identity derived from a Codex account's credentials. +#[derive(Debug, Clone)] +pub struct AuthBackedIdentity { + pub email: Option, + pub auth_subject: Option, + pub plan: Option, + pub provider_account_id: Option, +} + +/// Raw auth.json credentials. +#[derive(Debug, Clone)] +pub struct AuthCredentials { + pub access_token: String, + pub refresh_token: String, + pub id_token: Option, + pub account_id: Option, + pub last_refresh: Option>, +} + +impl AuthCredentials { + pub fn needs_refresh(&self) -> bool { + self.last_refresh + .is_none_or(|last| Utc::now() - last > chrono::TimeDelta::days(8)) + } +} + +/// Load the account identity from a Codex home's `auth.json`. +pub fn load_identity(codex_home_path: &Path) -> Result { + Ok(identity_from_credentials(&load_credentials( + codex_home_path, + )?)) +} + +/// Read and parse `auth.json`. +pub fn load_credentials(codex_home_path: &Path) -> Result { + let auth_path = codex_home_path.join("auth.json"); + let content = std::fs::read_to_string(&auth_path).map_err(|e| { + if e.kind() == std::io::ErrorKind::NotFound { + CodexApiError::Message("No `auth.json` was found for this account.".to_string()) + } else { + CodexApiError::Parse(format!("Failed to read the auth file: {e}")) + } + })?; + parse_credentials_json(&content) +} + +/// Parse `auth.json` contents, accepting `OPENAI_API_KEY` or a `tokens` object. +pub fn parse_credentials_json(content: &str) -> Result { + let json: serde_json::Value = serde_json::from_str(content) + .map_err(|e| CodexApiError::Parse(format!("Failed to parse the auth file: {e}")))?; + + if let Some(api_key) = json + .get("OPENAI_API_KEY") + .and_then(|v| v.as_str()) + .map(str::trim) + .filter(|s| !s.is_empty()) + { + return Ok(AuthCredentials { + access_token: api_key.to_string(), + refresh_token: String::new(), + id_token: None, + account_id: None, + last_refresh: None, + }); + } + + let tokens = json + .get("tokens") + .and_then(|v| v.as_object()) + .ok_or_else(|| { + CodexApiError::Message( + "The required token fields are missing from `auth.json`.".to_string(), + ) + })?; + + let access_token = tokens + .get("access_token") + .and_then(|v| v.as_str()) + .filter(|s| !s.is_empty()) + .ok_or_else(|| { + CodexApiError::Message( + "The required token fields are missing from `auth.json`.".to_string(), + ) + })? + .to_string(); + + let id_token = tokens + .get("id_token") + .and_then(|v| v.as_str()) + .map(str::to_string); + let account_id = tokens + .get("account_id") + .and_then(|v| v.as_str()) + .filter(|s| !s.is_empty()) + .map(str::to_string) + .or_else(|| account_id_from_id_token(id_token.as_deref())); + + Ok(AuthCredentials { + access_token, + refresh_token: tokens + .get("refresh_token") + .and_then(|v| v.as_str()) + .unwrap_or_default() + .to_string(), + id_token, + account_id, + last_refresh: json + .get("last_refresh") + .and_then(|v| v.as_str()) + .and_then(super::models::parse_datetime), + }) +} + +/// Save (possibly refreshed) credentials back to `auth.json`. +pub fn save_credentials( + codex_home_path: &Path, + credentials: &AuthCredentials, +) -> std::io::Result<()> { + let auth_path = codex_home_path.join("auth.json"); + let mut payload: serde_json::Value = std::fs::read_to_string(&auth_path) + .ok() + .and_then(|raw| serde_json::from_str(&raw).ok()) + .unwrap_or_else(|| serde_json::json!({})); + + let mut tokens = serde_json::Map::new(); + tokens.insert( + "access_token".to_string(), + serde_json::json!(credentials.access_token), + ); + tokens.insert( + "refresh_token".to_string(), + serde_json::json!(credentials.refresh_token), + ); + if let Some(id_token) = &credentials.id_token { + tokens.insert("id_token".to_string(), serde_json::json!(id_token)); + } + if let Some(account_id) = &credentials.account_id { + tokens.insert("account_id".to_string(), serde_json::json!(account_id)); + } + if let Some(obj) = payload.as_object_mut() { + obj.insert("tokens".to_string(), serde_json::Value::Object(tokens)); + obj.insert( + "last_refresh".to_string(), + serde_json::json!(Utc::now().to_rfc3339_opts(chrono::SecondsFormat::AutoSi, true)), + ); + } + std::fs::write(&auth_path, serde_json::to_vec_pretty(&payload)?) +} + +fn identity_from_credentials(credentials: &AuthCredentials) -> AuthBackedIdentity { + let payload = credentials + .id_token + .as_deref() + .and_then(jwt_payload) + .unwrap_or_default(); + let auth = payload + .get("https://api.openai.com/auth") + .and_then(|v| v.as_object()) + .cloned() + .unwrap_or_default(); + let profile = payload + .get("https://api.openai.com/profile") + .and_then(|v| v.as_object()) + .cloned() + .unwrap_or_default(); + + let email = normalize_string(payload.get("email").and_then(|v| v.as_str())) + .or_else(|| normalize_string(profile.get("email").and_then(|v| v.as_str()))); + let auth_subject = normalize_string(payload.get("sub").and_then(|v| v.as_str())); + let plan = normalize_string(auth.get("chatgpt_plan_type").and_then(|v| v.as_str())) + .or_else(|| normalize_string(payload.get("chatgpt_plan_type").and_then(|v| v.as_str()))); + let provider_account_id = normalize_string(credentials.account_id.as_deref()) + .or_else(|| normalize_string(auth.get("chatgpt_account_id").and_then(|v| v.as_str()))) + .or_else(|| normalize_string(payload.get("chatgpt_account_id").and_then(|v| v.as_str()))); + + AuthBackedIdentity { + email, + auth_subject, + plan, + provider_account_id, + } +} + +/// Minimal JWT payload extraction (base64url payload, no signature verification). +pub fn jwt_payload(token: &str) -> Option> { + let mut parts = token.split('.'); + let _header = parts.next()?; + let payload = parts.next()?; + let mut padded = payload.to_string(); + while padded.len() % 4 != 0 { + padded.push('='); + } + let decoded = base64::engine::general_purpose::URL_SAFE + .decode(padded.as_bytes()) + .ok()?; + serde_json::from_slice::(&decoded) + .ok()? + .as_object() + .cloned() +} + +fn account_id_from_id_token(id_token: Option<&str>) -> Option { + let payload = id_token.and_then(jwt_payload)?; + let auth = payload + .get("https://api.openai.com/auth") + .and_then(|v| v.as_object())?; + normalize_string(auth.get("chatgpt_account_id").and_then(|v| v.as_str())) + .or_else(|| normalize_string(payload.get("chatgpt_account_id").and_then(|v| v.as_str()))) +} + +fn normalize_string(value: Option<&str>) -> Option { + value + .map(str::trim) + .filter(|s| !s.is_empty()) + .map(str::to_string) +} + +fn string_value(value: &serde_json::Value, key: &str) -> Option { + value + .get(key) + .and_then(|v| v.as_str()) + .filter(|s| !s.is_empty()) + .map(str::to_string) +} + +// ── Quota fetching ────────────────────────────────────────────────────────── + +/// Client for live quota reads. Stateless per call; refresh decisions happen in +/// `CodexAccountApi::fetch_snapshot`. +pub struct CodexAccountApi { + client: reqwest::Client, +} + +impl CodexAccountApi { + pub fn new() -> Self { + let client = credentialed_http_client_builder() + .timeout(std::time::Duration::from_secs(REQUEST_TIMEOUT_SECONDS)) + .build() + .unwrap_or_else(|_| reqwest::Client::new()); + Self { client } + } + + /// Fetch a verified (or single) quota snapshot for the account at + /// `codex_home_path`, refreshing credentials when needed. + pub async fn fetch_snapshot( + &self, + codex_home_path: &Path, + email_hint: Option<&str>, + verify_live_data: bool, + ) -> Result { + let mut credentials = load_credentials(codex_home_path)?; + + if credentials.needs_refresh() + && !credentials.refresh_token.is_empty() + && let Ok(refreshed) = self.refresh(&credentials).await + { + let _ = save_credentials(codex_home_path, &refreshed); + credentials = refreshed; + } + + let result = self + .fetch_once(codex_home_path, &credentials, email_hint, verify_live_data) + .await; + if !matches!(&result, Err(CodexApiError::Message(msg)) if msg == UNAUTHORIZED_MESSAGE) + || credentials.refresh_token.is_empty() + { + return result; + } + + if let Ok(refreshed) = self.refresh(&credentials).await { + let _ = save_credentials(codex_home_path, &refreshed); + return self + .fetch_once(codex_home_path, &refreshed, email_hint, verify_live_data) + .await; + } + result + } + + async fn fetch_once( + &self, + codex_home_path: &Path, + credentials: &AuthCredentials, + email_hint: Option<&str>, + verify_live_data: bool, + ) -> Result { + if verify_live_data { + self.fetch_verified(codex_home_path, credentials, email_hint) + .await + } else { + self.fetch_single(codex_home_path, credentials, email_hint) + .await + } + } + + /// Fetch three reads and require equivalence (CodexControl accuracy model). + async fn fetch_verified( + &self, + codex_home_path: &Path, + credentials: &AuthCredentials, + email_hint: Option<&str>, + ) -> Result { + let first = self + .fetch_single(codex_home_path, credentials, email_hint) + .await?; + let second = self + .fetch_single(codex_home_path, credentials, email_hint) + .await?; + if is_equivalent(&first, &second) { + return Ok(second); + } + let third = self + .fetch_single(codex_home_path, credentials, email_hint) + .await?; + if is_equivalent(&first, &third) || is_equivalent(&second, &third) { + return Ok(third); + } + Err(CodexApiError::Message( + "Live API responses were inconsistent. The data could not be verified.".to_string(), + )) + } + + async fn fetch_single( + &self, + codex_home_path: &Path, + credentials: &AuthCredentials, + fallback_email: Option<&str>, + ) -> Result { + let identity = identity_from_credentials(credentials); + let response = self + .fetch_usage( + codex_home_path, + &credentials.access_token, + credentials.account_id.as_deref(), + ) + .await?; + let rate_limit = response.get("rate_limit").and_then(|v| v.as_object()); + let (primary_window, secondary_window) = make_normalized_windows(rate_limit); + let credits = response + .get("credits") + .and_then(|v| v.as_object()) + .map(make_credits); + + Ok(AccountUsageSnapshot { + email: identity.email.or_else(|| normalize_string(fallback_email)), + provider_account_id: identity + .provider_account_id + .or_else(|| credentials.account_id.clone()), + plan: normalize_string(response.get("plan_type").and_then(|v| v.as_str())) + .or(identity.plan), + allowed: rate_limit + .and_then(|r| r.get("allowed")) + .and_then(|v| v.as_bool()), + limit_reached: rate_limit + .and_then(|r| r.get("limit_reached")) + .and_then(|v| v.as_bool()), + primary_window, + secondary_window, + credits, + updated_at: Utc::now(), + }) + } + + async fn fetch_usage( + &self, + codex_home_path: &Path, + access_token: &str, + account_id: Option<&str>, + ) -> Result { + let url = resolve_usage_url(codex_home_path); + let mut request = self + .client + .get(&url) + .header("Authorization", format!("Bearer {access_token}")) + .header("User-Agent", "codex-cli") + .header("Accept", "application/json") + .header("Cache-Control", "no-cache, no-store, max-age=0") + .header("Pragma", "no-cache"); + if let Some(account_id) = account_id { + request = request.header("ChatGPT-Account-Id", account_id); + } + + let response = request + .send() + .await + .map_err(|e| CodexApiError::Network(e.to_string()))?; + if !response.status().is_success() { + let status = response.status(); + if status == reqwest::StatusCode::UNAUTHORIZED + || status == reqwest::StatusCode::FORBIDDEN + { + return Err(CodexApiError::Message(UNAUTHORIZED_MESSAGE.to_string())); + } + let body = response.text().await.unwrap_or_default().trim().to_string(); + let msg = if body.is_empty() { + format!("Codex API error {status}.") + } else { + format!("Codex API error {status}: {body}") + }; + return Err(CodexApiError::Message(msg)); + } + + let json: serde_json::Value = response + .json() + .await + .map_err(|e| CodexApiError::Parse(e.to_string()))?; + if !json.is_object() { + return Err(CodexApiError::Parse( + "The Codex API response was not in the expected format.".to_string(), + )); + } + Ok(json) + } + + /// Refresh an expired access token via the OpenAI OAuth endpoint. + pub async fn refresh( + &self, + credentials: &AuthCredentials, + ) -> Result { + if credentials.refresh_token.is_empty() { + return Err(CodexApiError::Message( + "No refresh token available for this account.".to_string(), + )); + } + let body = serde_json::json!({ + "client_id": REFRESH_CLIENT_ID, + "grant_type": "refresh_token", + "refresh_token": credentials.refresh_token, + "scope": "openid profile email", + }); + let response = self + .client + .post(REFRESH_ENDPOINT) + .json(&body) + .header("Content-Type", "application/json") + .header("Cache-Control", "no-cache, no-store, max-age=0") + .header("Pragma", "no-cache") + .send() + .await + .map_err(|e| CodexApiError::Network(e.to_string()))?; + + if response.status() == reqwest::StatusCode::UNAUTHORIZED { + let text = response.text().await.unwrap_or_default(); + let code = extract_error_code(&text).to_lowercase(); + let message = if code == "refresh_token_reused" { + "The refresh token can no longer be reused. Sign in again for this account." + } else if code == "refresh_token_invalidated" { + "The refresh token was revoked. Sign in again for this account." + } else { + "The refresh token has expired. Sign in again for this account." + }; + return Err(CodexApiError::Message(message.to_string())); + } + if !response.status().is_success() { + return Err(CodexApiError::Message( + "The Codex API response was not in the expected format.".to_string(), + )); + } + let payload: serde_json::Value = response + .json() + .await + .map_err(|e| CodexApiError::Parse(e.to_string()))?; + if !payload.is_object() { + return Err(CodexApiError::Message( + "The Codex API response was not in the expected format.".to_string(), + )); + } + let new_id_token = string_value(&payload, "id_token"); + Ok(AuthCredentials { + access_token: string_value(&payload, "access_token") + .unwrap_or_else(|| credentials.access_token.clone()), + refresh_token: string_value(&payload, "refresh_token") + .unwrap_or_else(|| credentials.refresh_token.clone()), + id_token: new_id_token + .clone() + .or_else(|| credentials.id_token.clone()), + account_id: credentials + .account_id + .clone() + .or_else(|| account_id_from_id_token(new_id_token.as_deref())), + last_refresh: Some(Utc::now()), + }) + } +} + +impl Default for CodexAccountApi { + fn default() -> Self { + Self::new() + } +} + +// ── URL resolution ────────────────────────────────────────────────────────── + +/// Resolve the usage URL from `config.toml` (`chatgpt_base_url`) or the default. +pub fn resolve_usage_url(codex_home_path: &Path) -> String { + let config_path = codex_home_path.join("config.toml"); + let configured_base = if config_path.exists() { + std::fs::read_to_string(&config_path) + .ok() + .and_then(|raw| parse_chatgpt_base_url(&raw)) + } else { + None + }; + + let mut base = configured_base.unwrap_or_else(|| USAGE_DEFAULT_BASE.to_string()); + while base.ends_with('/') { + base.pop(); + } + if base.starts_with("https://chatgpt.com") && !base.contains("/backend-api") { + base.push_str("/backend-api"); + } + if base.starts_with("https://chat.openai.com") && !base.contains("/backend-api") { + base.push_str("/backend-api"); + } + let path = if base.contains("/backend-api") { + "/wham/usage" + } else { + "/api/codex/usage" + }; + format!("{base}{path}") +} + +/// Extract `chatgpt_base_url` from a Codex `config.toml`. +pub fn parse_chatgpt_base_url(contents: &str) -> Option { + for raw_line in contents.lines() { + let line = raw_line.split('#').next().unwrap_or("").trim(); + if line.is_empty() { + continue; + } + let mut parts = line.splitn(2, '='); + let key = parts.next()?.trim(); + let value = parts.next()?.trim(); + if key != "chatgpt_base_url" { + continue; + } + let value = value + .strip_prefix('"') + .and_then(|v| v.strip_suffix('"')) + .or_else(|| value.strip_prefix('\'').and_then(|v| v.strip_suffix('\''))) + .unwrap_or(value); + return Some(value.to_string()); + } + None +} + +// ── Window normalization (session/weekly) ─────────────────────────────────── + +fn make_window(window: &serde_json::Map) -> Option { + let used_percent = window + .get("used_percent") + .and_then(|v| v.as_f64()) + .unwrap_or(0.0); + let reset_at = window + .get("reset_at") + .and_then(|v| v.as_i64()) + .and_then(|ts| DateTime::::from_timestamp(ts, 0)); + let limit_window_seconds = window + .get("limit_window_seconds") + .and_then(|v| v.as_i64()) + .unwrap_or(0); + Some(UsageWindowSnapshot::new( + used_percent, + reset_at, + limit_window_seconds, + )) +} + +/// Normalize the `rate_limit` object into (primary, secondary) windows with +/// roles assigned and `limit_reached` forced to 100%. +pub fn make_normalized_windows( + rate_limit: Option<&serde_json::Map>, +) -> (Option, Option) { + let Some(rate_limit) = rate_limit else { + return (None, None); + }; + let mut primary = rate_limit + .get("primary_window") + .and_then(|v| v.as_object()) + .and_then(make_window); + let mut secondary = rate_limit + .get("secondary_window") + .and_then(|v| v.as_object()) + .and_then(make_window); + + if rate_limit.get("limit_reached") == Some(&serde_json::Value::Bool(true)) { + if let Some(p) = primary.as_mut() { + p.used_percent = 100.0; + } + if let Some(s) = secondary.as_mut() { + s.used_percent = 100.0; + } + } + + normalize_window_roles(primary, secondary) +} + +/// Put the session window first and the weekly window second. +pub fn normalize_window_roles( + primary: Option, + secondary: Option, +) -> (Option, Option) { + if let (Some(p), Some(s)) = (&primary, &secondary) { + let (pr, sr) = (p.role(), s.role()); + if matches!( + (pr, sr), + (WindowRole::Weekly, WindowRole::Session) | (WindowRole::Weekly, WindowRole::Unknown) + ) { + return (secondary, primary); + } + return (primary, secondary); + } + if let Some(p) = &primary { + if p.role() == WindowRole::Weekly { + return (None, primary); + } + return (primary, None); + } + if let Some(s) = &secondary { + if s.role() == WindowRole::Weekly { + return (None, secondary); + } + return (secondary, None); + } + (None, None) +} + +fn make_credits(credits: &serde_json::Map) -> CreditsBalanceSnapshot { + CreditsBalanceSnapshot { + has_credits: credits + .get("has_credits") + .and_then(|v| v.as_bool()) + .unwrap_or(false), + unlimited: credits + .get("unlimited") + .and_then(|v| v.as_bool()) + .unwrap_or(false), + balance: credits.get("balance").and_then(|v| v.as_f64()), + } +} + +fn extract_error_code(payload: &str) -> String { + let Ok(parsed) = serde_json::from_str::(payload) else { + return String::new(); + }; + let error = parsed.get("error"); + if let Some(error) = error.and_then(|e| e.as_object()) { + return error + .get("code") + .and_then(|c| c.as_str()) + .unwrap_or_default() + .to_string(); + } + if let Some(error) = error.and_then(|e| e.as_str()) { + return error.to_string(); + } + parsed + .get("code") + .and_then(|c| c.as_str()) + .unwrap_or_default() + .to_string() +} + +/// Whether two fetched snapshots are equivalent (CodexControl verification). +pub fn is_equivalent(left: &AccountUsageSnapshot, right: &AccountUsageSnapshot) -> bool { + let email_eq = left.email.as_deref().map(str::to_lowercase) + == right.email.as_deref().map(str::to_lowercase); + email_eq + && left.provider_account_id == right.provider_account_id + && left.plan == right.plan + && left.allowed == right.allowed + && left.limit_reached == right.limit_reached + && windows_equivalent(&left.primary_window, &right.primary_window) + && windows_equivalent(&left.secondary_window, &right.secondary_window) + && credits_equivalent(&left.credits, &right.credits) +} + +fn windows_equivalent( + left: &Option, + right: &Option, +) -> bool { + match (left, right) { + (None, None) => true, + (None, Some(_)) | (Some(_), None) => false, + (Some(l), Some(r)) => { + let reset_matches = match (l.reset_at, r.reset_at) { + (None, None) => true, + (Some(a), Some(b)) => (a - b).num_seconds().abs() <= 1, + _ => false, + }; + l.limit_window_seconds == r.limit_window_seconds + && reset_matches + && (l.used_percent - r.used_percent).abs() < 0.001 + } + } +} + +fn credits_equivalent( + left: &Option, + right: &Option, +) -> bool { + match (left, right) { + (None, None) => true, + (None, Some(_)) | (Some(_), None) => false, + (Some(l), Some(r)) => { + let balance_matches = match (l.balance, r.balance) { + (None, None) => true, + (Some(a), Some(b)) => (a - b).abs() < 0.001, + _ => false, + }; + l.has_credits == r.has_credits && l.unlimited == r.unlimited && balance_matches + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parse_credentials_accepts_api_key() { + let creds = parse_credentials_json(r#"{"OPENAI_API_KEY":"sk-test"}"#).unwrap(); + assert_eq!(creds.access_token, "sk-test"); + assert_eq!(creds.account_id, None); + } + + #[test] + fn parse_credentials_accepts_tokens() { + let creds = parse_credentials_json( + r#"{"tokens":{"access_token":"at","refresh_token":"rt","account_id":"42"},"last_refresh":"2026-01-01T00:00:00Z"}"#, + ) + .unwrap(); + assert_eq!(creds.access_token, "at"); + assert_eq!(creds.refresh_token, "rt"); + assert_eq!(creds.account_id.as_deref(), Some("42")); + } + + #[test] + fn parse_credentials_missing_tokens_errors() { + assert!(parse_credentials_json(r#"{"foo":1}"#).is_err()); + } + + #[test] + fn jwt_payload_decodes() { + let payload = + base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(br#"{"email":"a@b.c"}"#); + let token = format!("eyJhbGciOiJub25lIn0.{payload}."); + let parsed = jwt_payload(&token).unwrap(); + assert_eq!(parsed.get("email").and_then(|v| v.as_str()), Some("a@b.c")); + } + + #[test] + fn normalize_window_roles_orders_session_first() { + let weekly = UsageWindowSnapshot::new(10.0, None, 604_800); + let session = UsageWindowSnapshot::new(10.0, None, 18_000); + let (p, s) = normalize_window_roles(Some(weekly), Some(session)); + assert_eq!(p.unwrap().limit_window_seconds, 18_000); + assert_eq!(s.unwrap().limit_window_seconds, 604_800); + } + + #[test] + fn limit_reached_forces_100() { + let payload = make_rate_limit(); + let (p, s) = make_normalized_windows(Some(&payload)); + assert_eq!(p.as_ref().unwrap().used_percent, 100.0); + assert_eq!(s.as_ref().unwrap().used_percent, 100.0); + } + + fn make_rate_limit() -> serde_json::Map { + serde_json::from_str( + r#"{"allowed":true,"limit_reached":true,"primary_window":{"used_percent":40,"reset_at":0,"limit_window_seconds":18000},"secondary_window":{"used_percent":20,"reset_at":0,"limit_window_seconds":604800}}"#, + ) + .unwrap() + } + + #[test] + fn resolve_usage_url_default() { + let dir = tempfile::tempdir().unwrap(); + let url = resolve_usage_url(dir.path()); + assert_eq!(url, "https://chatgpt.com/backend-api/wham/usage"); + } + + #[test] + fn parse_chatgpt_base_url_parses_quoted() { + let url = parse_chatgpt_base_url( + "# comment\nchatgpt_base_url = \"https://example.com/backend-api\"\n", + ) + .unwrap(); + assert_eq!(url, "https://example.com/backend-api"); + } + + #[test] + fn equivalent_snapshots_match() { + let mk = || AccountUsageSnapshot { + email: Some("a@b.c".to_string()), + provider_account_id: Some("x".to_string()), + plan: Some("pro".to_string()), + allowed: Some(true), + limit_reached: None, + primary_window: Some(UsageWindowSnapshot::new(12.0, Some(Utc::now()), 18_000)), + secondary_window: None, + credits: None, + updated_at: Utc::now(), + }; + assert!(is_equivalent(&mk(), &mk())); + let mut different = mk(); + different.plan = Some("plus".to_string()); + assert!(!is_equivalent(&mk(), &different)); + } + + #[test] + fn account_id_from_id_token_reads_auth() { + let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD + .encode(br#"{"https://api.openai.com/auth":{"chatgpt_account_id":"acct-99"}}"#); + let token = format!("h.{payload}.s"); + assert_eq!( + account_id_from_id_token(Some(&token)).as_deref(), + Some("acct-99") + ); + } +} diff --git a/rust/src/codex_accounts/codex_desktop.rs b/rust/src/codex_accounts/codex_desktop.rs new file mode 100644 index 0000000000..bcda25dfda --- /dev/null +++ b/rust/src/codex_accounts/codex_desktop.rs @@ -0,0 +1,407 @@ +//! Codex Desktop restart control (MSIX). +//! +//! Port of `windows/.../codex_desktop.py` (MIT): builds a hidden PowerShell +//! script that stops the Codex Desktop processes, syncs the MSIX session state +//! (backup/restore of the session entries) and relaunches the app. + +use std::io; +use std::path::{Path, PathBuf}; + +use base64::Engine; + +use super::file_locations::{ + DESKTOP_SESSION_STATE_ENTRIES, app_support_directory, codex_desktop_session_root, + ensure_directories, +}; + +pub const DEFAULT_RESTART_DELAY_SECONDS: f64 = 0.8; + +/// Friendly Codex Desktop control error. +#[derive(Debug, thiserror::Error)] +pub enum CodexDesktopControlError { + #[error("{0}")] + Message(String), + #[error(transparent)] + Io(#[from] io::Error), +} + +pub fn restart_log_path() -> PathBuf { + app_support_directory().join("codex-desktop-restart.log") +} + +pub fn restart_script_path() -> PathBuf { + app_support_directory().join("codex-desktop-restart.ps1") +} + +/// Render the PowerShell restart script. +pub fn build_restart_script( + delay_seconds: f64, + session_root: Option<&Path>, + backup_destination: Option<&Path>, + restore_source: Option<&Path>, +) -> String { + let delay_ms = (delay_seconds.max(0.0) * 1000.0).round() as u64; + let log_path = powershell_literal_path(&restart_log_path()); + let effective_session_root = session_root + .map(Path::to_path_buf) + .or_else(codex_desktop_session_root); + let session_root_literal = powershell_path_or_null(effective_session_root.as_deref()); + let backup_destination_literal = powershell_path_or_null(backup_destination); + let restore_source_literal = powershell_path_or_null(restore_source); + let session_entries_literal = powershell_string_array(DESKTOP_SESSION_STATE_ENTRIES); + + format!( + r#"$ErrorActionPreference = 'Stop' +$logPath = {log_path} +$sessionRoot = {session_root_literal} +$backupDestination = {backup_destination_literal} +$restoreSource = {restore_source_literal} +$sessionEntries = {session_entries_literal} +New-Item -ItemType Directory -Path ([System.IO.Path]::GetDirectoryName($logPath)) -Force | Out-Null +function Write-Log([string]$message) {{ + Add-Content -LiteralPath $logPath -Value ("[{{0}}] {{1}}" -f (Get-Date -Format 'yyyy-MM-dd HH:mm:ss.fff'), $message) +}} +function Clear-SessionEntry([string]$root, [string]$relativePath) {{ + if (-not $root) {{ + return + }} + $targetPath = Join-Path $root $relativePath + if (Test-Path -LiteralPath $targetPath) {{ + Remove-Item -LiteralPath $targetPath -Recurse -Force -ErrorAction Stop + }} +}} +function Copy-SessionEntry([string]$sourceRoot, [string]$destinationRoot, [string]$relativePath) {{ + if (-not $sourceRoot -or -not $destinationRoot) {{ + return + }} + $sourcePath = Join-Path $sourceRoot $relativePath + if (-not (Test-Path -LiteralPath $sourcePath)) {{ + return + }} + $destinationPath = Join-Path $destinationRoot $relativePath + $parentPath = [System.IO.Path]::GetDirectoryName($destinationPath) + if ($parentPath) {{ + New-Item -ItemType Directory -Path $parentPath -Force | Out-Null + }} + $item = Get-Item -LiteralPath $sourcePath -Force + if ($item.PSIsContainer) {{ + Copy-Item -LiteralPath $sourcePath -Destination $destinationPath -Recurse -Force + return + }} + Copy-Item -LiteralPath $sourcePath -Destination $destinationPath -Force +}} +function Sync-DesktopSessionState() {{ + if (-not $sessionRoot) {{ + Write-Log 'Desktop session root was not detected.' + return + }} + Write-Log ("Desktop session root: " + $sessionRoot) + Write-Log ("Backup destination: " + $(if ($backupDestination) {{ $backupDestination }} else {{ '' }})) + Write-Log ("Restore source: " + $(if ($restoreSource) {{ $restoreSource }} else {{ '' }})) + if (-not (Test-Path -LiteralPath $sessionRoot)) {{ + Write-Log ("Desktop session root is missing: " + $sessionRoot) + return + }} + if ($backupDestination) {{ + New-Item -ItemType Directory -Path $backupDestination -Force | Out-Null + foreach ($relativePath in $sessionEntries) {{ + try {{ + Clear-SessionEntry $backupDestination $relativePath + Copy-SessionEntry $sessionRoot $backupDestination $relativePath + Write-Log ("Backed up session entry: " + $relativePath) + }} catch {{ + Write-Log ("Failed to back up session entry " + $relativePath + ": " + $_.Exception.Message) + }} + }} + Write-Log ("Backed up desktop session state to " + $backupDestination) + }} + if ($restoreSource) {{ + if (-not (Test-Path -LiteralPath $restoreSource)) {{ + Write-Log ("Restore source is missing; leaving the current desktop session in place: " + $restoreSource) + return + }} + foreach ($relativePath in $sessionEntries) {{ + try {{ + Clear-SessionEntry $sessionRoot $relativePath + Copy-SessionEntry $restoreSource $sessionRoot $relativePath + Write-Log ("Restored session entry: " + $relativePath) + }} catch {{ + Write-Log ("Failed to restore session entry " + $relativePath + ": " + $_.Exception.Message) + }} + }} + Write-Log ("Restored desktop session state from " + $restoreSource) + }} +}} +Write-Log 'Restart requested.' +$mainProcess = Get-CimInstance Win32_Process | Where-Object {{ + $_.Name -eq 'Codex.exe' -and + $_.ExecutablePath -and + $_.ExecutablePath -notlike '*\resources\codex.exe' -and + $_.CommandLine -notmatch '--type=' +}} | Select-Object -First 1 +$launcherPath = $mainProcess.ExecutablePath +if ($launcherPath) {{ + Write-Log ("Using running launcher path: " + $launcherPath) +}} +if (-not $launcherPath) {{ + $package = Get-AppxPackage | Where-Object {{ + $_.Name -eq 'OpenAI.Codex' -or $_.PackageFamilyName -like 'OpenAI.Codex*' + }} | Sort-Object Version -Descending | Select-Object -First 1 + if ($package -and $package.InstallLocation) {{ + $launcherPath = Join-Path $package.InstallLocation 'app\Codex.exe' + Write-Log ("Using package launcher path: " + $launcherPath) + }} +}} +if (-not $launcherPath) {{ + Write-Log 'Unable to locate the Codex Desktop executable.' + throw 'Unable to locate the Codex Desktop executable.' +}} +Start-Sleep -Milliseconds {delay_ms} +$codexProcesses = Get-CimInstance Win32_Process | Where-Object {{ + $_.Name -ieq 'Codex.exe' -or + $_.ExecutablePath -like '*\OpenAI.Codex_*\app\Codex.exe' -or + $_.ExecutablePath -like '*\OpenAI.Codex_*\app\resources\codex.exe' +}} +Write-Log ("Found " + $codexProcesses.Count + " Codex processes to stop.") +$codexProcesses | ForEach-Object {{ + try {{ + & taskkill.exe /PID $_.ProcessId /F /T | Out-Null + Write-Log ("taskkill succeeded for PID " + $_.ProcessId) + }} catch {{ + Write-Log ("taskkill failed for PID " + $_.ProcessId + ": " + $_.Exception.Message) + }} + try {{ + Stop-Process -Id $_.ProcessId -Force -ErrorAction Stop + Write-Log ("Stop-Process succeeded for PID " + $_.ProcessId) + }} catch {{ + Write-Log ("Stop-Process failed for PID " + $_.ProcessId + ": " + $_.Exception.Message) + }} +}} +$deadline = (Get-Date).AddSeconds(8) +while ((Get-Date) -lt $deadline) {{ + $remaining = Get-CimInstance Win32_Process | Where-Object {{ + $_.Name -ieq 'Codex.exe' -or + $_.ExecutablePath -like '*\OpenAI.Codex_*\app\Codex.exe' -or + $_.ExecutablePath -like '*\OpenAI.Codex_*\app\resources\codex.exe' + }} + if (-not $remaining) {{ + Write-Log 'All Codex processes exited.' + break + }} + Write-Log ("Still waiting for " + $remaining.Count + " Codex processes to exit.") + $remaining | ForEach-Object {{ + try {{ + & taskkill.exe /PID $_.ProcessId /F /T | Out-Null + }} catch {{}} + }} + Start-Sleep -Milliseconds 250 +}} +if (Get-CimInstance Win32_Process | Where-Object {{ + $_.Name -ieq 'Codex.exe' -or + $_.ExecutablePath -like '*\OpenAI.Codex_*\app\Codex.exe' -or + $_.ExecutablePath -like '*\OpenAI.Codex_*\app\resources\codex.exe' +}}) {{ + Write-Log 'Continuing with relaunch after timeout while some Codex processes still appear alive.' +}} +Sync-DesktopSessionState +Start-Sleep -Milliseconds 700 +Start-Process -FilePath $launcherPath +Write-Log 'Codex Desktop relaunched.' +"# + ) + .trim_end() + .to_string() +} + +/// Encode a script as base64 UTF-16LE (PowerShell `-EncodedCommand`). +pub fn encode_powershell_script(script: &str) -> String { + let utf16: Vec = script.encode_utf16().collect(); + let bytes: Vec = utf16 + .into_iter() + .flat_map(|unit| unit.to_le_bytes()) + .collect(); + base64::engine::general_purpose::STANDARD.encode(bytes) +} + +/// Build the hidden PowerShell command line for a script file. +pub fn build_restart_command(script_path: &Path) -> Vec { + let powershell_exe = std::env::var("WINDIR") + .map(|windir| { + PathBuf::from(windir).join("System32\\WindowsPowerShell\\v1.0\\powershell.exe") + }) + .unwrap_or_else(|_| { + PathBuf::from(r"C:\Windows\System32\WindowsPowerShell\v1.0\powershell.exe") + }); + vec![ + powershell_exe.to_string_lossy().to_string(), + "-NoProfile".to_string(), + "-NonInteractive".to_string(), + "-WindowStyle".to_string(), + "Hidden".to_string(), + "-ExecutionPolicy".to_string(), + "Bypass".to_string(), + "-File".to_string(), + script_path.to_string_lossy().to_string(), + ] +} + +/// Write the restart script and launch a hidden PowerShell that runs it. +pub fn restart_codex_desktop( + delay_seconds: f64, + session_root: Option<&Path>, + backup_destination: Option<&Path>, + restore_source: Option<&Path>, +) -> Result<(), CodexDesktopControlError> { + ensure_directories()?; + let script = build_restart_script( + delay_seconds, + session_root, + backup_destination, + restore_source, + ); + fs_write(restart_script_path(), script)?; + launch_hidden_powershell(&restart_script_path()) +} + +#[cfg(windows)] +fn fs_write(path: PathBuf, content: String) -> io::Result<()> { + std::fs::write(path, content) +} + +#[cfg(windows)] +fn launch_hidden_powershell(script_path: &Path) -> Result<(), CodexDesktopControlError> { + use std::os::windows::process::CommandExt; + use std::process::Command; + + let mut command = Command::new(&build_restart_command(script_path)[0]); + command.args(&build_restart_command(script_path)[1..]); + command + .stdin(std::process::Stdio::null()) + .stdout(std::process::Stdio::null()) + .stderr(std::process::Stdio::null()); + // CREATE_NO_WINDOW | CREATE_NEW_PROCESS_GROUP (never DETACHED_PROCESS). + const CREATE_NO_WINDOW: u32 = 0x0800_0000; + const CREATE_NEW_PROCESS_GROUP: u32 = 0x0000_0200; + command.creation_flags(CREATE_NO_WINDOW | CREATE_NEW_PROCESS_GROUP); + command.spawn().map(|_| ()).map_err(|error| { + CodexDesktopControlError::Message(format!("Failed to restart Codex Desktop: {error}")) + }) +} + +#[cfg(not(windows))] +fn fs_write(path: PathBuf, content: String) -> io::Result<()> { + std::fs::write(path, content) +} + +#[cfg(not(windows))] +fn launch_hidden_powershell(_script_path: &Path) -> Result<(), CodexDesktopControlError> { + Err(CodexDesktopControlError::Message( + "Codex Desktop restart is only available on Windows.".to_string(), + )) +} + +fn powershell_literal_path(path: &Path) -> String { + let normalized = path.to_string_lossy().replace('/', "\\"); + format!("'{}'", normalized.replace('\'', "''")) +} + +fn powershell_path_or_null(path: Option<&Path>) -> String { + match path { + Some(path) => powershell_literal_path(path), + None => "$null".to_string(), + } +} + +fn powershell_string_array(values: &[&str]) -> String { + let quoted: Vec = values + .iter() + .map(|value| format!("'{}'", value.replace('\'', "''"))) + .collect(); + format!("@({})", quoted.join(", ")) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn build_restart_script_includes_restart_flow() { + let script = build_restart_script(1.25, None, None, None); + assert!(script.contains("Write-Log")); + assert!(script.contains("Get-CimInstance Win32_Process")); + assert!(script.contains("Get-AppxPackage")); + assert!(script.contains("taskkill.exe /PID")); + assert!(script.contains("Stop-Process -Id $_.ProcessId -Force")); + assert!(script.contains("Start-Process -FilePath $launcherPath")); + assert!(script.contains("Start-Sleep -Milliseconds 1250")); + } + + #[test] + fn build_restart_script_includes_desktop_session_sync() { + let script = build_restart_script( + 0.5, + Some(Path::new( + r"C:\Users\test\AppData\Local\Packages\OpenAI.Codex_test\LocalCache\Roaming\Codex", + )), + Some(Path::new( + r"C:\Users\test\AppData\Roaming\CodexControl\managed-homes\current\desktop-session", + )), + Some(Path::new( + r"C:\Users\test\AppData\Roaming\CodexControl\managed-homes\target\desktop-session", + )), + ); + + assert!(script.contains( + "$sessionRoot = 'C:\\Users\\test\\AppData\\Local\\Packages\\OpenAI.Codex_test\\LocalCache\\Roaming\\Codex'" + )); + assert!(script.contains( + "$backupDestination = 'C:\\Users\\test\\AppData\\Roaming\\CodexControl\\managed-homes\\current\\desktop-session'" + )); + assert!(script.contains( + "$restoreSource = 'C:\\Users\\test\\AppData\\Roaming\\CodexControl\\managed-homes\\target\\desktop-session'" + )); + assert!(script.contains("function Sync-DesktopSessionState()")); + assert!(script.contains("Copy-SessionEntry $sessionRoot $backupDestination $relativePath")); + assert!(script.contains("Copy-SessionEntry $restoreSource $sessionRoot $relativePath")); + assert!(script.contains("Clear-SessionEntry $sessionRoot $relativePath")); + assert!( + script.contains( + "Restore source is missing; leaving the current desktop session in place" + ) + ); + assert!(script.contains("Failed to back up session entry")); + assert!(script.contains("Failed to restore session entry")); + } + + #[test] + fn encode_powershell_script_round_trips_utf16le() { + use base64::Engine; + let script = "Start-Process -FilePath 'Codex.exe'"; + let encoded = encode_powershell_script(script); + let bytes = base64::engine::general_purpose::STANDARD + .decode(encoded) + .unwrap(); + let units: Vec = bytes + .chunks_exact(2) + .map(|chunk| u16::from_le_bytes([chunk[0], chunk[1]])) + .collect(); + let decoded = String::from_utf16(&units).unwrap(); + assert_eq!(decoded, script); + } + + #[test] + fn build_restart_command_invokes_hidden_powershell_file() { + let command = build_restart_command(Path::new(r"C:\temp\restart.ps1")); + assert!(command[0].to_lowercase().ends_with("powershell.exe")); + assert!(command.contains(&"-WindowStyle".to_string())); + assert!(command.contains(&"Hidden".to_string())); + assert_eq!(command[command.len() - 2], "-File"); + assert!( + command + .last() + .unwrap() + .to_lowercase() + .ends_with("temp\\restart.ps1") + ); + } +} diff --git a/rust/src/codex_accounts/file_locations.rs b/rust/src/codex_accounts/file_locations.rs new file mode 100644 index 0000000000..00c8cc29d8 --- /dev/null +++ b/rust/src/codex_accounts/file_locations.rs @@ -0,0 +1,178 @@ +//! Path resolution for Codex account storage and Codex Desktop session state. +//! +//! Mirrors `windows/.../file_locations.py` (MIT), adapted to CodexBar's +//! `%config%/CodexBar` convention. + +use std::path::{Path, PathBuf}; + +/// Entries of the Codex Desktop MSIX session that must be preserved/restored +/// when switching accounts. Mirrors `DESKTOP_SESSION_STATE_ENTRIES`. +pub const DESKTOP_SESSION_STATE_ENTRIES: &[&str] = &[ + "blob_storage", + "DIPS", + "DIPS-wal", + "Local State", + "Local Storage", + "Network", + "Partitions", + "Preferences", + "Session Storage", + "SharedStorage", + "SharedStorage-wal", + "shared_proto_db", +]; + +fn localappdata_directory() -> Option { + if let Ok(path) = std::env::var("LOCALAPPDATA") { + let path = PathBuf::from(path.trim()); + if !path.as_os_str().is_empty() { + return Some(path); + } + } + None +} + +thread_local! { + static APP_SUPPORT_OVERRIDE: std::cell::RefCell> = + const { std::cell::RefCell::new(None) }; + static AMBIENT_HOME_OVERRIDE: std::cell::RefCell> = + const { std::cell::RefCell::new(None) }; + static DESKTOP_SESSION_ROOT_OVERRIDE: std::cell::RefCell> = + const { std::cell::RefCell::new(None) }; +} + +/// Base directory holding the Codex account store (accounts.json, managed-homes). +pub fn app_support_directory() -> PathBuf { + APP_SUPPORT_OVERRIDE + .with(|cell| cell.borrow().clone()) + .unwrap_or_else(|| { + dirs::config_dir() + .map(|dir| dir.join("CodexBar").join("codex-accounts")) + .unwrap_or_else(|| PathBuf::from(".").join("codex-accounts")) + }) +} + +/// Override the app support root (tests / shell). Returns the previous value. +pub fn with_app_support_directory(path: PathBuf) -> Option { + APP_SUPPORT_OVERRIDE.with(|cell| { + let previous = cell.borrow().clone(); + *cell.borrow_mut() = Some(path); + previous + }) +} + +pub fn clear_app_support_directory_override() { + APP_SUPPORT_OVERRIDE.with(|cell| *cell.borrow_mut() = None); +} + +pub fn accounts_file() -> PathBuf { + app_support_directory().join("accounts.json") +} + +pub fn snapshots_file() -> PathBuf { + app_support_directory().join("snapshots.json") +} + +pub fn managed_homes_directory() -> PathBuf { + app_support_directory().join("managed-homes") +} + +pub fn auth_backups_directory() -> PathBuf { + app_support_directory().join("auth-backups") +} + +/// The environment (ambient) Codex home. +pub fn ambient_codex_home() -> PathBuf { + AMBIENT_HOME_OVERRIDE + .with(|cell| cell.borrow().clone()) + .unwrap_or_else(|| { + dirs::home_dir() + .unwrap_or_else(|| PathBuf::from(".")) + .join(".codex") + }) +} + +/// Override the ambient home root (tests / shell). +pub fn with_ambient_codex_home(path: PathBuf) { + AMBIENT_HOME_OVERRIDE.with(|cell| *cell.borrow_mut() = Some(path)); +} + +pub fn clear_ambient_codex_home_override() { + AMBIENT_HOME_OVERRIDE.with(|cell| *cell.borrow_mut() = None); +} + +pub const DESKTOP_SESSION_SNAPSHOT_DIRECTORY_NAME: &str = "desktop-session"; + +/// Ensure required directories exist. +pub fn ensure_directories() -> std::io::Result<()> { + std::fs::create_dir_all(app_support_directory())?; + std::fs::create_dir_all(managed_homes_directory())?; + std::fs::create_dir_all(auth_backups_directory()) +} + +/// Discover the active Codex Desktop MSIX package session root +/// (`%LOCALAPPDATA%\Packages\OpenAI.Codex*\LocalCache\Roaming\Codex`). +pub fn codex_desktop_session_root() -> Option { + if let Some(override_path) = DESKTOP_SESSION_ROOT_OVERRIDE.with(|cell| cell.borrow().clone()) { + return Some(override_path); + } + let packages_root = localappdata_directory()?.join("Packages"); + if !packages_root.exists() { + return None; + } + let entries = std::fs::read_dir(&packages_root).ok()?; + let mut packages: Vec = entries + .filter_map(|e| e.ok().map(|e| e.path())) + .filter(|p| { + p.is_dir() + && p.file_name() + .and_then(|n| n.to_str()) + .is_some_and(|n| n.starts_with("OpenAI.Codex")) + }) + .collect(); + packages.sort(); + for package in packages { + let session_root = package.join("LocalCache").join("Roaming").join("Codex"); + if session_root.exists() { + return Some(session_root); + } + } + None +} + +/// Override the desktop session root (tests / shell). +pub fn with_codex_desktop_session_root(path: PathBuf) { + DESKTOP_SESSION_ROOT_OVERRIDE.with(|cell| *cell.borrow_mut() = Some(path)); +} + +pub fn clear_codex_desktop_session_root_override() { + DESKTOP_SESSION_ROOT_OVERRIDE.with(|cell| *cell.borrow_mut() = None); +} + +/// Per-account desktop-session snapshot directory. +pub fn desktop_session_snapshot_path(home_path: &Path) -> PathBuf { + home_path.join(DESKTOP_SESSION_SNAPSHOT_DIRECTORY_NAME) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn session_state_entries_are_embedded() { + assert!(DESKTOP_SESSION_STATE_ENTRIES.contains(&"Local Storage")); + } + + #[test] + fn desktop_session_snapshot_path_nests_under_home() { + let p = desktop_session_snapshot_path(Path::new("/tmp/acct")); + assert_eq!(p, Path::new("/tmp/acct/desktop-session")); + } + + #[test] + fn app_support_default_resolves_to_config_dir() { + clear_app_support_directory_override(); + let dir = app_support_directory(); + assert!(dir.to_string_lossy().contains("codex-accounts")); + } +} diff --git a/rust/src/codex_accounts/mod.rs b/rust/src/codex_accounts/mod.rs new file mode 100644 index 0000000000..48114a5bf2 --- /dev/null +++ b/rust/src/codex_accounts/mod.rs @@ -0,0 +1,34 @@ +//! Multi-account Codex support for CodexBar. +//! +//! This module is a Rust port of the Windows core of +//! [`ademisler/codexcontrol`](https://github.com/ademisler/codexcontrol) (MIT), +//! which manages multiple Codex accounts through isolated `CODEX_HOME` +//! directories and switches the ambient Codex identity. See `NOTICE`/LICENSE +//! for the upstream MIT attribution. +//! +//! The model is deliberately mirror-shaped: a `CodexAccount` lives in either the +//! ambient home (`~/.codex`) or an app-managed home (`managed-homes/`), +//! quota snapshots are fetched per account, and switching swaps the ambient +//! `auth.json` plus the Codex Desktop MSIX session state. + +pub mod account_manager; +pub mod api; +pub mod codex_desktop; +pub mod file_locations; +pub mod models; +pub mod stores; + +pub use account_manager::{ + CodexAccountManager, CodexAccountManagerError, CodexLoginOutcome, CodexLoginResult, + CodexSwitchResult, ManagedLoginProcess, +}; +pub use api::{AuthBackedIdentity, AuthCredentials, CodexAccountApi, CodexApiError, load_identity}; +pub use codex_desktop::{ + CodexDesktopControlError, build_restart_command, build_restart_script, + encode_powershell_script, restart_codex_desktop, +}; +pub use models::{ + AccountUsageSnapshot, CodexAccount, CodexAccountSource, CreditsBalanceSnapshot, + RemovedAccountIdentity, UsageWindowSnapshot, utc_now, +}; +pub use stores::{AccountStore, SnapshotStore}; diff --git a/rust/src/codex_accounts/models.rs b/rust/src/codex_accounts/models.rs new file mode 100644 index 0000000000..cb79abd214 --- /dev/null +++ b/rust/src/codex_accounts/models.rs @@ -0,0 +1,669 @@ +//! Domain model for Codex accounts and their usage snapshots. +//! +//! Field names intentionally mirror CodexControl's `windows/.../models.py` (MIT) +//! so stored data interops with that project. + +use std::path::{Path, PathBuf}; + +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; +use uuid::Uuid; + +/// `parse_from_rfc3339` requires an offset; append `Z` only when none is present. +pub fn parse_datetime(value: &str) -> Option> { + let text = value.trim(); + if text.is_empty() { + return None; + } + // `parse_from_rfc3339` requires an offset; append `Z` only when none is present. + let text_dt = text.trim(); + let has_offset = text_dt.ends_with(['Z', 'z']) || contains_offset(text_dt); + let normalized = if has_offset { + String::from(text_dt) + } else { + format!("{text_dt}Z") + }; + DateTime::parse_from_rfc3339(&normalized) + .ok() + .map(|dt| dt.with_timezone(&Utc)) +} + +/// Does the string carry an explicit `+HH:MM` / `-HH:MM` UTC offset (not `Z`)? +fn contains_offset(text: &str) -> bool { + let Some(time_start) = text.find('T') else { + return false; + }; + let rest = &text[time_start + 1..]; + let Some(sign) = rest.rfind(['+', '-']) else { + return false; + }; + let tail = &rest[sign + 1..]; + let mut chars = tail.chars(); + let digits = chars.next().is_some_and(|c| c.is_ascii_digit()) + && chars.next().is_some_and(|c| c.is_ascii_digit()); + digits && tail.contains(':') +} + +/// Format a UTC instant the same way CodexControl does (`...Z`). +pub fn format_datetime(value: Option>) -> Option { + value.map(|dt| dt.to_rfc3339_opts(chrono::SecondsFormat::AutoSi, true)) +} + +pub fn utc_now() -> DateTime { + Utc::now() +} + +fn normalize_identifier(value: Option<&str>) -> Option { + value + .map(|v| v.trim().to_lowercase()) + .filter(|v| !v.is_empty()) +} + +/// Where an account's `CODEX_HOME` lives. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub enum CodexAccountSource { + /// The environment's `~/.codex` (the identity the Codex CLI/Desktop uses). + Ambient, + /// An app-owned home directory under `managed-homes/`. + ManagedByApp, +} + +impl CodexAccountSource { + pub fn display_name(self) -> &'static str { + match self { + CodexAccountSource::Ambient => "System", + CodexAccountSource::ManagedByApp => "Managed", + } + } + + /// Whether the app owns (and may delete) this account's files. + pub fn owns_files(self) -> bool { + matches!(self, CodexAccountSource::ManagedByApp) + } + + pub fn from_raw(value: &str) -> Option { + match value { + "ambient" => Some(CodexAccountSource::Ambient), + "managedByApp" | "importedCodexBar" => Some(CodexAccountSource::ManagedByApp), + _ => None, + } + } +} + +/// A stored Codex account. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct CodexAccount { + pub id: Uuid, + pub nickname: Option, + pub email_hint: Option, + pub auth_subject: Option, + pub provider_account_id: Option, + pub codex_home_path: PathBuf, + pub source: CodexAccountSource, + pub created_at: DateTime, + pub updated_at: DateTime, + pub last_authenticated_at: Option>, +} + +impl CodexAccount { + #[allow(clippy::too_many_arguments)] + pub fn new( + id: Uuid, + nickname: Option, + email_hint: Option, + auth_subject: Option, + provider_account_id: Option, + codex_home_path: PathBuf, + source: CodexAccountSource, + created_at: DateTime, + updated_at: DateTime, + last_authenticated_at: Option>, + ) -> Self { + Self { + id, + nickname, + email_hint, + auth_subject, + provider_account_id, + codex_home_path, + source, + created_at, + updated_at, + last_authenticated_at, + } + } + + pub fn display_name(&self) -> String { + if let Some(nickname) = self + .nickname + .as_deref() + .map(str::trim) + .filter(|s| !s.is_empty()) + { + return nickname.to_string(); + } + if let Some(email) = self.email_hint.as_deref().filter(|s| !s.is_empty()) { + return email.to_string(); + } + self.codex_home_path + .file_name() + .map(|name| name.to_string_lossy().to_string()) + .unwrap_or_else(|| self.codex_home_path.display().to_string()) + } + + pub fn normalized_email_hint(&self) -> Option { + normalize_identifier(self.email_hint.as_deref()) + } + + pub fn normalized_auth_subject(&self) -> Option { + normalize_identifier(self.auth_subject.as_deref()) + } + + pub fn normalized_provider_account_id(&self) -> Option { + normalize_identifier(self.provider_account_id.as_deref()) + } + + pub fn standardized_home_path(&self) -> String { + std::path::absolute(&self.codex_home_path) + .unwrap_or_else(|_| self.codex_home_path.clone()) + .to_string_lossy() + .to_lowercase() + } + + fn source_priority(&self) -> u8 { + if self.source.owns_files() { 2 } else { 1 } + } + + fn recency_date(&self) -> DateTime { + self.last_authenticated_at.unwrap_or(self.updated_at) + } + + /// Whether two accounts refer to the same identity. + pub fn matches(&self, other: &CodexAccount) -> bool { + if self.standardized_home_path() == other.standardized_home_path() { + return true; + } + if let (Some(a), Some(b)) = ( + self.normalized_provider_account_id(), + other.normalized_provider_account_id(), + ) && a == b + { + return true; + } + if self.normalized_provider_account_id().is_some() + || other.normalized_provider_account_id().is_some() + { + return false; + } + if let (Some(a), Some(b)) = ( + self.normalized_auth_subject(), + other.normalized_auth_subject(), + ) && a == b + { + return true; + } + if let (Some(a), Some(b)) = (self.normalized_email_hint(), other.normalized_email_hint()) + && a == b + { + return true; + } + false + } + + /// Merge a fresher discovery into this account, preferring managed/recency. + pub fn merge_from(&mut self, other: &CodexAccount) { + if self + .nickname + .as_deref() + .map(str::trim) + .is_none_or(|s| s.is_empty()) + { + self.nickname = other.nickname.clone(); + } + + let prefer_other = other.source_priority() > self.source_priority() + || (other.source_priority() == self.source_priority() + && other.recency_date() >= self.recency_date()); + + let pick = |mine: &mut Option, value: Option<&String>| { + let newer = prefer_other && value.is_some_and(|v| !v.trim().is_empty()); + if newer || mine.is_none() { + *mine = value.cloned(); + } + }; + pick(&mut self.email_hint, other.email_hint.as_ref()); + pick(&mut self.auth_subject, other.auth_subject.as_ref()); + pick( + &mut self.provider_account_id, + other.provider_account_id.as_ref(), + ); + + if prefer_other { + self.source = other.source; + self.codex_home_path = other.codex_home_path.clone(); + } + + self.updated_at = self.updated_at.max(other.updated_at); + self.last_authenticated_at = match (self.last_authenticated_at, other.last_authenticated_at) + { + (Some(a), Some(b)) => Some(a.max(b)), + (a, b) => a.or(b), + }; + } +} + +/// Identity of a previously-removed account, kept to avoid re-adding it. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct RemovedAccountIdentity { + pub id: Uuid, + pub email_hint: Option, + pub auth_subject: Option, + pub provider_account_id: Option, + pub codex_home_path: PathBuf, + pub source: CodexAccountSource, + pub removed_at: DateTime, +} + +impl RemovedAccountIdentity { + pub fn from_account(account: &CodexAccount) -> Self { + Self { + id: Uuid::new_v4(), + email_hint: account.email_hint.clone(), + auth_subject: account.auth_subject.clone(), + provider_account_id: account.provider_account_id.clone(), + codex_home_path: account.codex_home_path.clone(), + source: account.source, + removed_at: utc_now(), + } + } + + pub fn matches(&self, account: &CodexAccount) -> bool { + if self.standardized_home_path() == account.standardized_home_path() { + return true; + } + if let (Some(a), Some(b)) = ( + normalize_identifier(self.provider_account_id.as_deref()), + account.normalized_provider_account_id(), + ) && a == b + { + return true; + } + if self + .provider_account_id + .as_ref() + .map(|v| !v.trim().is_empty()) + .unwrap_or(false) + || account.provider_account_id.as_ref().is_some() + { + return false; + } + if let (Some(a), Some(b)) = ( + normalize_identifier(self.auth_subject.as_deref()), + account.normalized_auth_subject(), + ) && a == b + { + return true; + } + if let (Some(a), Some(b)) = ( + normalize_identifier(self.email_hint.as_deref()), + account.normalized_email_hint(), + ) && a == b + { + return true; + } + false + } + + fn standardized_home_path(&self) -> String { + std::path::absolute(&self.codex_home_path) + .unwrap_or_else(|_| self.codex_home_path.clone()) + .to_string_lossy() + .to_lowercase() + } +} + +/// A single quota window (session or weekly). +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct UsageWindowSnapshot { + pub used_percent: f64, + pub reset_at: Option>, + pub limit_window_seconds: i64, +} + +impl UsageWindowSnapshot { + pub fn new( + used_percent: f64, + reset_at: Option>, + limit_window_seconds: i64, + ) -> Self { + Self { + used_percent, + reset_at, + limit_window_seconds, + } + } + + pub fn remaining_percent(&self) -> f64 { + 100.0_f64.max(self.used_percent) - self.used_percent + } + + pub fn role(&self) -> WindowRole { + match self.limit_window_seconds { + 18_000 => WindowRole::Session, + 604_800 => WindowRole::Weekly, + _ => WindowRole::Unknown, + } + } +} + +/// Normalized role of a window based on its duration. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum WindowRole { + Session, + Weekly, + Unknown, +} + +/// Codex credits balance. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct CreditsBalanceSnapshot { + pub has_credits: bool, + pub unlimited: bool, + pub balance: Option, +} + +impl CreditsBalanceSnapshot { + pub fn display_value(&self) -> String { + if self.unlimited { + return "Unlimited".to_string(); + } + if let Some(balance) = self.balance { + return format!("{balance:.2}") + .trim_end_matches('0') + .trim_end_matches('.') + .to_string(); + } + if self.has_credits { + return "Available".to_string(); + } + "None".to_string() + } +} + +/// A fetched snapshot for one Codex account. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountUsageSnapshot { + pub email: Option, + pub provider_account_id: Option, + pub plan: Option, + pub allowed: Option, + pub limit_reached: Option, + pub primary_window: Option, + pub secondary_window: Option, + pub credits: Option, + pub updated_at: DateTime, +} + +impl AccountUsageSnapshot { + pub fn is_quota_blocked(&self) -> bool { + self.limit_reached == Some(true) || self.allowed == Some(false) + } + + pub fn has_quota_windows(&self) -> bool { + self.primary_window.is_some() || self.secondary_window.is_some() + } + + pub fn has_usable_quota_now(&self) -> bool { + if self.is_quota_blocked() { + return false; + } + let values = [self.primary_window.as_ref(), self.secondary_window.as_ref()] + .into_iter() + .flatten() + .map(|w| w.remaining_percent()); + let mut values = values.peekable(); + values.peek().is_some() && values.any(|v| v > 0.001) + } + + pub fn lowest_remaining_percent(&self) -> f64 { + if self.is_quota_blocked() { + return 0.0; + } + [self.secondary_window.as_ref(), self.primary_window.as_ref()] + .into_iter() + .flatten() + .map(|w| w.remaining_percent()) + .fold(f64::MAX, f64::min) + } + + pub fn next_reset_at(&self) -> Option> { + [self.primary_window.as_ref(), self.secondary_window.as_ref()] + .into_iter() + .flatten() + .filter_map(|w| w.reset_at) + .min() + } +} + +/// Sort weight used to order accounts by practical usefulness. +pub fn account_sort_priority(snapshot: &AccountUsageSnapshot) -> u8 { + if snapshot.has_usable_quota_now() { + 0 + } else if snapshot.next_reset_at().is_some() { + 1 + } else { + 2 + } +} + +fn _path_is_trailing(path: &Path) -> bool { + path.as_os_str() + .to_string_lossy() + .ends_with(std::path::MAIN_SEPARATOR) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn account( + id: &str, + home: &str, + source: CodexAccountSource, + provider_id: Option<&str>, + ) -> CodexAccount { + CodexAccount::new( + Uuid::parse_str(id).unwrap(), + None, + None, + None, + provider_id.map(str::to_string), + PathBuf::from(home), + source, + utc_now(), + utc_now(), + None, + ) + } + + #[test] + fn matches_by_home_path() { + let a = account( + "11111111-1111-1111-1111-111111111111", + "/x/a", + CodexAccountSource::ManagedByApp, + None, + ); + let b = account( + "22222222-2222-2222-2222-222222222222", + "/x/a", + CodexAccountSource::ManagedByApp, + None, + ); + assert!(a.matches(&b)); + } + + #[test] + fn matches_by_provider_account_id() { + let a = account( + "11111111-1111-1111-1111-111111111111", + "/x/a", + CodexAccountSource::ManagedByApp, + Some("acct-1"), + ); + let b = account( + "22222222-2222-2222-2222-222222222222", + "/y/b", + CodexAccountSource::ManagedByApp, + Some("ACCT-1"), + ); + assert!(a.matches(&b)); + } + + #[test] + fn disambiguates_different_provider_ids() { + let a = account( + "11111111-1111-1111-1111-111111111111", + "/x/a", + CodexAccountSource::ManagedByApp, + Some("acct-1"), + ); + let b = account( + "22222222-2222-2222-2222-222222222222", + "/y/b", + CodexAccountSource::ManagedByApp, + Some("acct-2"), + ); + assert!(!a.matches(&b)); + } + + #[test] + fn source_displays_and_ownership() { + assert_eq!(CodexAccountSource::Ambient.display_name(), "System"); + assert_eq!(CodexAccountSource::ManagedByApp.display_name(), "Managed"); + assert!(CodexAccountSource::ManagedByApp.owns_files()); + assert!(!CodexAccountSource::Ambient.owns_files()); + } + + #[test] + fn window_role_classification() { + assert_eq!( + UsageWindowSnapshot::new(0.0, None, 18_000).role(), + WindowRole::Session + ); + assert_eq!( + UsageWindowSnapshot::new(0.0, None, 604_800).role(), + WindowRole::Weekly + ); + assert_eq!( + UsageWindowSnapshot::new(0.0, None, 1234).role(), + WindowRole::Unknown + ); + } + + #[test] + fn blocked_account_has_no_usable_quota() { + let snapshot = AccountUsageSnapshot { + email: None, + provider_account_id: None, + plan: None, + allowed: Some(false), + limit_reached: None, + primary_window: Some(UsageWindowSnapshot::new(10.0, None, 18_000)), + secondary_window: None, + credits: None, + updated_at: utc_now(), + }; + assert!(snapshot.is_quota_blocked()); + assert!(!snapshot.has_usable_quota_now()); + assert_eq!(snapshot.lowest_remaining_percent(), 0.0); + } + + #[test] + fn parse_datetime_accepts_z_and_offset() { + assert!(parse_datetime("2026-01-01T00:00:00Z").is_some()); + assert!(parse_datetime("2026-01-01T00:00:00+00:00").is_some()); + assert!(parse_datetime("").is_none()); + } + + #[test] + fn credits_display_value() { + assert_eq!( + CreditsBalanceSnapshot { + has_credits: true, + unlimited: true, + balance: None + } + .display_value(), + "Unlimited" + ); + assert_eq!( + CreditsBalanceSnapshot { + has_credits: true, + unlimited: false, + balance: Some(12.50) + } + .display_value(), + "12.5" + ); + assert_eq!( + CreditsBalanceSnapshot { + has_credits: true, + unlimited: false, + balance: None + } + .display_value(), + "Available" + ); + assert_eq!( + CreditsBalanceSnapshot { + has_credits: false, + unlimited: false, + balance: None + } + .display_value(), + "None" + ); + } + + #[test] + fn merge_prefers_managed_and_recency() { + let mut managed = account( + "11111111-1111-1111-1111-111111111111", + "/x/managed", + CodexAccountSource::ManagedByApp, + None, + ); + managed.nickname = Some("My acct".to_string()); + let ambient = account( + "22222222-2222-2222-2222-222222222222", + "~/.codex-like/ambient", + CodexAccountSource::Ambient, + None, + ); + managed.merge_from(&ambient); + assert_eq!(managed.source, CodexAccountSource::ManagedByApp); + assert_eq!(managed.display_name(), "My acct"); + } + + #[test] + fn display_name_falls_back_to_home() { + let acct = account( + "11111111-1111-1111-1111-111111111111", + "/x/my-home-dir", + CodexAccountSource::ManagedByApp, + None, + ); + assert!( + acct.display_name().ends_with("my-home-dir") + || acct.display_name().contains("my-home-dir") + ); + let _ = _path_is_trailing(std::path::Path::new("/x/")); + } +} diff --git a/rust/src/codex_accounts/stores.rs b/rust/src/codex_accounts/stores.rs new file mode 100644 index 0000000000..4eace86dd5 --- /dev/null +++ b/rust/src/codex_accounts/stores.rs @@ -0,0 +1,266 @@ +//! JSON persistence for Codex accounts and usage snapshots. +//! +//! Mirrors `windows/.../stores.py` (MIT), using `secure_file` for the +//! secret-bearing accounts file (DPAPI on Windows) and plain JSON for the +//! non-secret snapshot cache. + +use std::collections::HashMap; +use std::io; +use std::path::PathBuf; +use std::str::FromStr; + +use uuid::Uuid; + +use crate::secure_file; + +use super::file_locations::{accounts_file, snapshots_file}; +use super::models::{CodexAccount, RemovedAccountIdentity}; + +#[derive(Debug, serde::Serialize, serde::Deserialize)] +struct AccountsFile { + version: u32, + accounts: Vec, + #[serde(rename = "removedAccounts", default)] + removed_accounts: Vec, +} + +impl Default for AccountsFile { + fn default() -> Self { + Self { + version: Self::CURRENT_VERSION, + accounts: Vec::new(), + removed_accounts: Vec::new(), + } + } +} + +impl AccountsFile { + const CURRENT_VERSION: u32 = 2; +} + +/// Reads/writes the accounts metadata. +pub struct AccountStore { + file_path: PathBuf, +} + +impl AccountStore { + pub fn new() -> Self { + Self { + file_path: accounts_file(), + } + } + + pub fn with_path(path: PathBuf) -> Self { + Self { file_path: path } + } + + pub fn load(&self) -> io::Result<(Vec, Vec)> { + if !self.file_path.exists() { + return Ok((Vec::new(), Vec::new())); + } + let data = secure_file::read_string(&self.file_path)?; + let file: AccountsFile = serde_json::from_str(&data) + .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; + Ok((file.accounts, file.removed_accounts)) + } + + pub fn load_accounts(&self) -> io::Result> { + Ok(self.load()?.0) + } + + pub fn save( + &self, + accounts: &[CodexAccount], + removed_accounts: Option<&[RemovedAccountIdentity]>, + ) -> io::Result<()> { + super::file_locations::ensure_directories()?; + let removed = match removed_accounts { + Some(r) => r.to_vec(), + None => self.load()?.1, + }; + let file = AccountsFile { + version: AccountsFile::CURRENT_VERSION, + accounts: accounts.to_vec(), + removed_accounts: removed, + }; + let data = serde_json::to_vec_pretty(&file).map_err(io::Error::other)?; + let data = String::from_utf8(data).map_err(io::Error::other)?; + secure_file::write_string(&self.file_path, &data) + } + + /// Merge discovered accounts into the stored list, deduping by identity. + pub fn merge( + &self, + existing: &[CodexAccount], + incoming: Vec, + ) -> io::Result> { + let removed = self.load()?.1; + let mut result: Vec = existing + .iter() + .filter(|acct| !removed.iter().any(|r| r.matches(acct))) + .cloned() + .collect(); + for candidate in incoming { + match result.iter_mut().find(|acct| acct.matches(&candidate)) { + Some(existing_account) => existing_account.merge_from(&candidate), + None => result.push(candidate), + } + } + result.sort_by_key(|a| a.display_name().to_lowercase()); + Ok(result) + } +} + +impl Default for AccountStore { + fn default() -> Self { + Self::new() + } +} + +/// Reads/writes the per-account usage snapshot cache. +pub struct SnapshotStore { + file_path: PathBuf, +} + +impl SnapshotStore { + pub fn new() -> Self { + Self { + file_path: snapshots_file(), + } + } + + pub fn with_path(path: PathBuf) -> Self { + Self { file_path: path } + } + + pub fn load(&self) -> io::Result> { + if !self.file_path.exists() { + return Ok(HashMap::new()); + } + let data = std::fs::read_to_string(&self.file_path)?; + let file: serde_json::Value = serde_json::from_str(&data) + .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; + let snapshots = file.get("snapshots"); + let Some(snapshots) = snapshots else { + return Ok(HashMap::new()); + }; + let Some(object) = snapshots.as_object() else { + return Ok(HashMap::new()); + }; + let mut result = HashMap::new(); + for (key, value) in object { + if let Ok(id) = Uuid::from_str(key) + && let Ok(snapshot) = + serde_json::from_value::(value.clone()) + { + result.insert(id, snapshot); + } + } + Ok(result) + } + + pub fn save( + &self, + snapshots: &HashMap, + ) -> io::Result<()> { + super::file_locations::ensure_directories()?; + let mut object = serde_json::Map::new(); + for (id, snapshot) in snapshots { + object.insert( + id.to_string(), + serde_json::to_value(snapshot).map_err(io::Error::other)?, + ); + } + let file = serde_json::json!({ "snapshots": object }); + let data = serde_json::to_vec_pretty(&file).map_err(io::Error::other)?; + std::fs::write(&self.file_path, data) + } +} + +impl Default for SnapshotStore { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::codex_accounts::models::{CodexAccountSource, utc_now}; + + fn store_dir() -> tempfile::TempDir { + let dir = tempfile::tempdir().unwrap(); + super::super::file_locations::with_app_support_directory(dir.path().to_path_buf()); + dir + } + + fn make_account(id: &str, source: CodexAccountSource) -> CodexAccount { + CodexAccount::new( + Uuid::parse_str(id).unwrap(), + Some(format!("acct-{id}")), + Some("person@example.com".to_string()), + None, + None, + PathBuf::from(format!("/tmp/managed/{id}")), + source, + utc_now(), + utc_now(), + None, + ) + } + + #[test] + fn account_store_roundtrips() { + let _guard = store_dir(); + let store = AccountStore::new(); + let acct = make_account( + "11111111-1111-1111-1111-111111111111", + CodexAccountSource::ManagedByApp, + ); + store.save(&[acct], None).unwrap(); + let (loaded, _) = store.load().unwrap(); + assert_eq!(loaded.len(), 1); + assert_eq!( + loaded[0].nickname.as_deref().unwrap(), + "acct-11111111-1111-1111-1111-111111111111" + ); + crate::codex_accounts::file_locations::clear_app_support_directory_override(); + } + + #[test] + fn snapshot_store_roundtrips() { + let _guard = store_dir(); + let store = SnapshotStore::new(); + let id = Uuid::parse_str("11111111-1111-1111-1111-111111111111").unwrap(); + let snapshot = crate::codex_accounts::models::AccountUsageSnapshot { + email: Some("a@b.c".to_string()), + provider_account_id: None, + plan: Some("pro".to_string()), + allowed: Some(true), + limit_reached: None, + primary_window: Some(crate::codex_accounts::models::UsageWindowSnapshot::new( + 12.0, None, 18_000, + )), + secondary_window: None, + credits: None, + updated_at: utc_now(), + }; + let mut map = HashMap::new(); + map.insert(id, snapshot.clone()); + store.save(&map).unwrap(); + let loaded = store.load().unwrap(); + assert_eq!(loaded.len(), 1); + assert_eq!(loaded[&id].plan.as_deref().unwrap(), "pro"); + crate::codex_accounts::file_locations::clear_app_support_directory_override(); + } + + #[test] + fn missing_store_loads_empty() { + let dir = tempfile::tempdir().unwrap(); + crate::codex_accounts::file_locations::with_app_support_directory(dir.path().to_path_buf()); + let store = AccountStore::new(); + let (accounts, removed) = store.load().unwrap(); + assert!(accounts.is_empty() && removed.is_empty()); + crate::codex_accounts::file_locations::clear_app_support_directory_override(); + } +} diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 35e06c8475..1aebe724a2 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -6,6 +6,7 @@ pub mod agent_sessions; pub mod browser; pub mod cli; +pub mod codex_accounts; pub mod codex_workspaces; pub mod core; pub mod cost_scanner; From cc07f94e14d5d96d4c763f77ce4e32b8f9c601c1 Mon Sep 17 00:00:00 2001 From: Finesssee <90105158+Finesssee@users.noreply.github.com> Date: Thu, 6 Aug 2026 13:14:10 +0700 Subject: [PATCH 2/3] docs(adr): record 0003 multi-account Codex coexistence decision --- docs/adr/0003-multi-account-codex.md | 74 ++++++++++++++++++++++++++++ 1 file changed, 74 insertions(+) create mode 100644 docs/adr/0003-multi-account-codex.md diff --git a/docs/adr/0003-multi-account-codex.md b/docs/adr/0003-multi-account-codex.md new file mode 100644 index 0000000000..dfe8a8ad46 --- /dev/null +++ b/docs/adr/0003-multi-account-codex.md @@ -0,0 +1,74 @@ +# ADR 0003: Multi-account Codex coexistence with ambient provider + +Date: 2026-08-06 +Status: Accepted + +## Context + +CodexBar today surfaces exactly one `codex` provider snapshot. That snapshot is +produced by the ambient single-account Codex provider in +`rust/src/providers/codex/`: it reads the identity that happens to be active in +`~/.codex` (`auth.json`), calls `fetch_usage`, and publishes the result through +the standard provider pipeline to the tray, flyout, and settings surfaces. + +A stacked 5-PR series ports the MIT-licensed Windows core of +[`ademisler/codexcontrol`](https://github.com/ademisler/codexcontrol) (Windows +Python modules) into Rust as `rust/src/codex_accounts/`: + +- **PR 1/5 — domain core**: account model (`CodexAccount`, + `CodexAccountSource::Ambient | Managed`), ambient vs app-managed + `CODEX_HOME` discovery (`file_locations`), per-account quota snapshots, + account stores, login runner, and Codex Desktop (MSIX) session switching. +- **PR 2/5 — shell**: Tauri commands exposing discovery/login/switch/remove, + plus managed-account refresh lanes that fetch and cache a snapshot per + managed account in parallel. +- **PR 3/5 — frontend bridge**: TS types, `tauri.ts` wrappers, i18n keys. +- **PR 4/5 — settings panel**: `CodexAccountsSection` in provider settings. +- **PR 5/5 — tray menu**: `CodexAccountsMenu` flyout. + +The new module therefore manages *many* accounts while the existing provider +pipeline still models *one* `codex` provider. Both must work while the port +lands in reviewable slices. + +## Decision + +**COEXISTENCE-NOW.** The ambient single-account provider +(`rust/src/providers/codex/`) remains the source of truth for the single +`codex` provider snapshot published through the existing pipeline today. +`rust/src/codex_accounts/` is added as a *parallel* managed-account domain: + +- PR 2/5 ships managed-account lanes that fetch/cache one snapshot per managed + account without disturbing the ambient provider's snapshot. +- PR 4/5 and PR 5/5 add surfaces (settings panel, tray flyout) that read the + managed-account domain; they present the ambient slot through the same model + (`CodexAccountSource::Ambient`) without changing what the single `codex` + provider row shows. +- No existing provider, surface, or persistence path is rewired in this series; + the ambient provider keeps publishing exactly as before. + +**Replacement intent** (tracked as part of this ADR decision, not executed in +this series): the ambient provider path retires in a follow-up major release +once `codex_accounts` reaches feature parity against it — i.e. the managed +lanes match ambient `fetch_usage` behavior *and* local cost scans are migrated +onto the managed-home model. Retirement is a separate change with its own +review; this series only prepares the ground. + +## Consequences + +- During the transition both worlds publish snapshots: the ambient provider + emits the single `codex` provider snapshot, while the codex_accounts lanes + emit per-managed-account snapshots. They do not overwrite each other. +- **Identity precedence**: the ambient slot's snapshot *is* the `codex` + provider snapshot (authority for credit bars, notifications, ordering) until + the ambient path retires. Managed lanes never claim the provider slot; the + tray/settings account picker is informational for them. +- **Signed/tracked lifetime hook**: the ambient path becomes eligible for + retirement when (a) codex_accounts lane coverage matches ambient + `fetch_usage` behavior (same usage windows, credits, and reset semantics per + account) and (b) local cost scans are migrated to read managed homes. Both + conditions are checked off under this ADR before the retirement PR merges. +- Review cost stays bounded: each PR is a small, independently testable slice; + rollback of any surface PR leaves the domain intact. +- Transitional duplication is accepted: two code paths compute "current Codex + usage" until retirement. The `NOTICE` file carries the upstream + ademisler/codexcontrol MIT attribution. From d3a22c4d56e368a6aa1ff5368a703f7f8fe8ffd6 Mon Sep 17 00:00:00 2001 From: Finesssee <90105158+Finesssee@users.noreply.github.com> Date: Thu, 6 Aug 2026 13:16:10 +0700 Subject: [PATCH 3/3] refactor(codex): split login runner from account manager (1k-line bar) --- rust/src/codex_accounts/account_manager.rs | 231 +------------------- rust/src/codex_accounts/login_runner.rs | 234 +++++++++++++++++++++ rust/src/codex_accounts/mod.rs | 7 +- 3 files changed, 239 insertions(+), 233 deletions(-) create mode 100644 rust/src/codex_accounts/login_runner.rs diff --git a/rust/src/codex_accounts/account_manager.rs b/rust/src/codex_accounts/account_manager.rs index 52ad75e190..636337f12b 100644 --- a/rust/src/codex_accounts/account_manager.rs +++ b/rust/src/codex_accounts/account_manager.rs @@ -8,10 +8,7 @@ use std::fs; use std::path::{Path, PathBuf}; -use std::process::{Child, Command, Stdio}; -use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::{Arc, Mutex}; -use std::time::{Duration, Instant}; +use std::time::Duration; use chrono::{DateTime, Utc}; use thiserror::Error; @@ -22,6 +19,7 @@ use super::file_locations::{ ambient_codex_home, auth_backups_directory, codex_desktop_session_root, desktop_session_snapshot_path, ensure_directories, managed_homes_directory, }; +use super::login_runner::{CodexLoginOutcome, CodexLoginRunner, ManagedLoginProcess}; use super::models::{CodexAccount, CodexAccountSource, utc_now}; /// Friendly account manager error. @@ -39,231 +37,6 @@ impl From for CodexAccountManagerError { } } -/// Outcome of a `codex login` subprocess run. -#[derive(Debug, Clone)] -pub enum CodexLoginOutcome { - MissingBinary, - LaunchFailed(String), - TimedOut(String), - Cancelled, - Failed(String), - Success(String), -} - -impl CodexLoginOutcome { - pub fn as_str(&self) -> &'static str { - match self { - CodexLoginOutcome::MissingBinary => "missing_binary", - CodexLoginOutcome::LaunchFailed(_) => "launch_failed", - CodexLoginOutcome::TimedOut(_) => "timed_out", - CodexLoginOutcome::Cancelled => "cancelled", - CodexLoginOutcome::Failed(_) => "failed", - CodexLoginOutcome::Success(_) => "success", - } - } - - pub fn output(&self) -> &str { - match self { - CodexLoginOutcome::MissingBinary => "", - CodexLoginOutcome::LaunchFailed(output) - | CodexLoginOutcome::TimedOut(output) - | CodexLoginOutcome::Failed(output) - | CodexLoginOutcome::Success(output) => output, - CodexLoginOutcome::Cancelled => "", - } - } -} - -/// Result of a `codex login` subprocess run. -#[derive(Debug, Clone)] -pub struct CodexLoginResult { - pub outcome: CodexLoginOutcome, -} - -/// Handle around an in-flight `codex login` process, for cancellation. -#[derive(Debug, Default, Clone)] -pub struct ManagedLoginProcess { - inner: Arc>>, - cancelled: Arc, -} - -impl ManagedLoginProcess { - fn bind(&self, process: Child) { - *self.inner.lock().expect("login process lock") = Some(process); - self.cancelled.store(false, Ordering::SeqCst); - } - - fn clear(&self) { - *self.inner.lock().expect("login process lock") = None; - } - - pub fn is_cancelled(&self) -> bool { - self.cancelled.load(Ordering::SeqCst) - } - - pub fn cancel(&self) { - self.cancelled.store(true, Ordering::SeqCst); - let mut guard = self.inner.lock().expect("login process lock"); - if let Some(child) = guard.as_mut() { - let _ = child.kill(); - } - } -} - -/// Runs `codex login` inside an isolated `CODEX_HOME`. -pub struct CodexLoginRunner; - -impl CodexLoginRunner { - /// Resolve the `codex` executable, falling back to known install paths. - pub fn locate_codex_binary() -> Option { - if let Ok(found) = which::which("codex") { - return Some(found); - } - path_candidates() - .into_iter() - .find(|candidate| candidate.is_file()) - } - - pub fn run( - home_path: &Path, - timeout: Duration, - handle: Option<&ManagedLoginProcess>, - ) -> CodexLoginResult { - let active_handle = handle.cloned().unwrap_or_default(); - let Some(binary) = Self::locate_codex_binary() else { - return CodexLoginResult { - outcome: CodexLoginOutcome::MissingBinary, - }; - }; - - let mut command = Command::new(binary); - command - .arg("login") - .env("CODEX_HOME", home_path) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()); - let child = match command.spawn() { - Ok(child) => child, - Err(error) => { - return CodexLoginResult { - outcome: CodexLoginOutcome::LaunchFailed(error.to_string()), - }; - } - }; - active_handle.bind(child); - - let output = match wait_for_child(&active_handle, timeout) { - Some(output) => output, - None => { - let output = kill_and_drain(&active_handle); - active_handle.clear(); - return CodexLoginResult { - outcome: CodexLoginOutcome::TimedOut(combine_output(&output)), - }; - } - }; - - active_handle.clear(); - let combined = combine_output(&output); - if active_handle.is_cancelled() { - return CodexLoginResult { - outcome: CodexLoginOutcome::Cancelled, - }; - } - if output.status.success() { - return CodexLoginResult { - outcome: CodexLoginOutcome::Success(combined), - }; - } - CodexLoginResult { - outcome: CodexLoginOutcome::Failed(combined), - } - } -} - -fn path_candidates() -> Vec { - let local_app_data = std::env::var("LOCALAPPDATA") - .map(PathBuf::from) - .unwrap_or_else(|_| { - dirs::home_dir() - .unwrap_or_else(|| PathBuf::from(".")) - .join("AppData") - .join("Local") - }); - let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from(".")); - vec![ - local_app_data - .join("OpenAI") - .join("Codex") - .join("bin") - .join("codex.exe"), - home.join(".bun").join("bin").join("codex.exe"), - local_app_data - .join("Microsoft") - .join("WindowsApps") - .join("codex.exe"), - ] -} - -fn wait_for_child(handle: &ManagedLoginProcess, timeout: Duration) -> Option { - let deadline = Instant::now() + timeout; - loop { - if handle.is_cancelled() { - let output = take_child(handle)?.wait_with_output().ok(); - return output; - } - let polled = { - let mut guard = handle.inner.lock().expect("login process lock"); - match guard.as_mut().map(|child| child.try_wait()) { - Some(Ok(Some(_status))) => take_child(handle)?.wait_with_output().ok(), - Some(Err(_)) => take_child(handle)?.wait_with_output().ok(), - _ => None, - } - }; - if polled.is_some() { - return polled; - } - if Instant::now() >= deadline { - return None; - } - std::thread::sleep(Duration::from_millis(100)); - } -} - -fn take_child(handle: &ManagedLoginProcess) -> Option { - handle.inner.lock().expect("login process lock").take() -} - -fn kill_and_drain(handle: &ManagedLoginProcess) -> std::process::Output { - let mut child = take_child(handle).expect("login process present"); - let _ = child.kill(); - child - .wait_with_output() - .unwrap_or_else(|_| std::process::Output { - status: std::process::ExitStatus::default(), - stdout: Vec::new(), - stderr: Vec::new(), - }) -} - -fn combine_output(output: &std::process::Output) -> String { - let mut parts: Vec = Vec::new(); - for bytes in [&output.stdout, &output.stderr] { - let text = String::from_utf8_lossy(bytes); - let trimmed = text.trim(); - if !trimmed.is_empty() { - parts.push(trimmed.to_string()); - } - } - let merged = parts.join("\n"); - let merged = merged.trim(); - if merged.is_empty() { - "No output captured.".to_string() - } else { - merged.chars().take(4000).collect() - } -} - /// Result of switching the active account. #[derive(Debug, Clone, serde::Serialize)] #[serde(rename_all = "camelCase")] diff --git a/rust/src/codex_accounts/login_runner.rs b/rust/src/codex_accounts/login_runner.rs new file mode 100644 index 0000000000..d96a1f6e1c --- /dev/null +++ b/rust/src/codex_accounts/login_runner.rs @@ -0,0 +1,234 @@ +//! Runs `codex login` inside an isolated `CODEX_HOME`, with cancellation, +//! timeouts, and combined output capture. Split out of `account_manager.rs` +//! (port of the login-running slice of `windows/.../account_manager.py`, MIT). + +use std::path::{Path, PathBuf}; +use std::process::{Child, Command, Stdio}; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +/// Outcome of a `codex login` subprocess run. +#[derive(Debug, Clone)] +pub enum CodexLoginOutcome { + MissingBinary, + LaunchFailed(String), + TimedOut(String), + Cancelled, + Failed(String), + Success(String), +} + +impl CodexLoginOutcome { + pub fn as_str(&self) -> &'static str { + match self { + CodexLoginOutcome::MissingBinary => "missing_binary", + CodexLoginOutcome::LaunchFailed(_) => "launch_failed", + CodexLoginOutcome::TimedOut(_) => "timed_out", + CodexLoginOutcome::Cancelled => "cancelled", + CodexLoginOutcome::Failed(_) => "failed", + CodexLoginOutcome::Success(_) => "success", + } + } + + pub fn output(&self) -> &str { + match self { + CodexLoginOutcome::MissingBinary => "", + CodexLoginOutcome::LaunchFailed(output) + | CodexLoginOutcome::TimedOut(output) + | CodexLoginOutcome::Failed(output) + | CodexLoginOutcome::Success(output) => output, + CodexLoginOutcome::Cancelled => "", + } + } +} + +/// Result of a `codex login` subprocess run. +#[derive(Debug, Clone)] +pub struct CodexLoginResult { + pub outcome: CodexLoginOutcome, +} + +/// Handle around an in-flight `codex login` process, for cancellation. +#[derive(Debug, Default, Clone)] +pub struct ManagedLoginProcess { + inner: Arc>>, + cancelled: Arc, +} + +impl ManagedLoginProcess { + fn bind(&self, process: Child) { + *self.inner.lock().expect("login process lock") = Some(process); + self.cancelled.store(false, Ordering::SeqCst); + } + + fn clear(&self) { + *self.inner.lock().expect("login process lock") = None; + } + + pub fn is_cancelled(&self) -> bool { + self.cancelled.load(Ordering::SeqCst) + } + + pub fn cancel(&self) { + self.cancelled.store(true, Ordering::SeqCst); + let mut guard = self.inner.lock().expect("login process lock"); + if let Some(child) = guard.as_mut() { + let _ = child.kill(); + } + } +} + +/// Runs `codex login` inside an isolated `CODEX_HOME`. +pub struct CodexLoginRunner; + +impl CodexLoginRunner { + /// Resolve the `codex` executable, falling back to known install paths. + pub fn locate_codex_binary() -> Option { + if let Ok(found) = which::which("codex") { + return Some(found); + } + path_candidates() + .into_iter() + .find(|candidate| candidate.is_file()) + } + + pub fn run( + home_path: &Path, + timeout: Duration, + handle: Option<&ManagedLoginProcess>, + ) -> CodexLoginResult { + let active_handle = handle.cloned().unwrap_or_default(); + let Some(binary) = Self::locate_codex_binary() else { + return CodexLoginResult { + outcome: CodexLoginOutcome::MissingBinary, + }; + }; + + let mut command = Command::new(binary); + command + .arg("login") + .env("CODEX_HOME", home_path) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()); + let child = match command.spawn() { + Ok(child) => child, + Err(error) => { + return CodexLoginResult { + outcome: CodexLoginOutcome::LaunchFailed(error.to_string()), + }; + } + }; + active_handle.bind(child); + + let output = match wait_for_child(&active_handle, timeout) { + Some(output) => output, + None => { + let output = kill_and_drain(&active_handle); + active_handle.clear(); + return CodexLoginResult { + outcome: CodexLoginOutcome::TimedOut(combine_output(&output)), + }; + } + }; + + active_handle.clear(); + let combined = combine_output(&output); + if active_handle.is_cancelled() { + return CodexLoginResult { + outcome: CodexLoginOutcome::Cancelled, + }; + } + if output.status.success() { + return CodexLoginResult { + outcome: CodexLoginOutcome::Success(combined), + }; + } + CodexLoginResult { + outcome: CodexLoginOutcome::Failed(combined), + } + } +} + +fn path_candidates() -> Vec { + let local_app_data = std::env::var("LOCALAPPDATA") + .map(PathBuf::from) + .unwrap_or_else(|_| { + dirs::home_dir() + .unwrap_or_else(|| PathBuf::from(".")) + .join("AppData") + .join("Local") + }); + let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from(".")); + vec![ + local_app_data + .join("OpenAI") + .join("Codex") + .join("bin") + .join("codex.exe"), + home.join(".bun").join("bin").join("codex.exe"), + local_app_data + .join("Microsoft") + .join("WindowsApps") + .join("codex.exe"), + ] +} + +fn wait_for_child(handle: &ManagedLoginProcess, timeout: Duration) -> Option { + let deadline = Instant::now() + timeout; + loop { + if handle.is_cancelled() { + let output = take_child(handle)?.wait_with_output().ok(); + return output; + } + let polled = { + let mut guard = handle.inner.lock().expect("login process lock"); + match guard.as_mut().map(|child| child.try_wait()) { + Some(Ok(Some(_status))) => take_child(handle)?.wait_with_output().ok(), + Some(Err(_)) => take_child(handle)?.wait_with_output().ok(), + _ => None, + } + }; + if polled.is_some() { + return polled; + } + if Instant::now() >= deadline { + return None; + } + std::thread::sleep(Duration::from_millis(100)); + } +} + +fn take_child(handle: &ManagedLoginProcess) -> Option { + handle.inner.lock().expect("login process lock").take() +} + +fn kill_and_drain(handle: &ManagedLoginProcess) -> std::process::Output { + let mut child = take_child(handle).expect("login process present"); + let _ = child.kill(); + child + .wait_with_output() + .unwrap_or_else(|_| std::process::Output { + status: std::process::ExitStatus::default(), + stdout: Vec::new(), + stderr: Vec::new(), + }) +} + +fn combine_output(output: &std::process::Output) -> String { + let mut parts: Vec = Vec::new(); + for bytes in [&output.stdout, &output.stderr] { + let text = String::from_utf8_lossy(bytes); + let trimmed = text.trim(); + if !trimmed.is_empty() { + parts.push(trimmed.to_string()); + } + } + let merged = parts.join("\n"); + let merged = merged.trim(); + if merged.is_empty() { + "No output captured.".to_string() + } else { + merged.chars().take(4000).collect() + } +} diff --git a/rust/src/codex_accounts/mod.rs b/rust/src/codex_accounts/mod.rs index 48114a5bf2..494e6e75b3 100644 --- a/rust/src/codex_accounts/mod.rs +++ b/rust/src/codex_accounts/mod.rs @@ -15,18 +15,17 @@ pub mod account_manager; pub mod api; pub mod codex_desktop; pub mod file_locations; +pub mod login_runner; pub mod models; pub mod stores; -pub use account_manager::{ - CodexAccountManager, CodexAccountManagerError, CodexLoginOutcome, CodexLoginResult, - CodexSwitchResult, ManagedLoginProcess, -}; +pub use account_manager::{CodexAccountManager, CodexAccountManagerError, CodexSwitchResult}; pub use api::{AuthBackedIdentity, AuthCredentials, CodexAccountApi, CodexApiError, load_identity}; pub use codex_desktop::{ CodexDesktopControlError, build_restart_command, build_restart_script, encode_powershell_script, restart_codex_desktop, }; +pub use login_runner::{CodexLoginOutcome, CodexLoginResult, ManagedLoginProcess}; pub use models::{ AccountUsageSnapshot, CodexAccount, CodexAccountSource, CreditsBalanceSnapshot, RemovedAccountIdentity, UsageWindowSnapshot, utc_now,