From 549846b29ad52f6cb4f8560365a731966054a9b3 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Thu, 31 Jul 2025 10:48:49 -0700 Subject: [PATCH] Add codex login --api-key (#1759) Allow setting the API key via `codex login --api-key` --- codex-rs/Cargo.lock | 1 + codex-rs/chatgpt/src/chatgpt_token.rs | 2 +- codex-rs/cli/src/login.rs | 32 +++- codex-rs/cli/src/main.rs | 10 +- codex-rs/cli/src/proto.rs | 2 +- codex-rs/core/src/codex_wrapper.rs | 2 +- codex-rs/core/tests/client.rs | 8 +- codex-rs/login/Cargo.toml | 3 + codex-rs/login/src/lib.rs | 202 +++++++++++++++++++++----- codex-rs/tui/src/lib.rs | 2 +- 10 files changed, 218 insertions(+), 46 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 120050c2..87d59b21 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -793,6 +793,7 @@ dependencies = [ "reqwest", "serde", "serde_json", + "tempfile", "tokio", ] diff --git a/codex-rs/chatgpt/src/chatgpt_token.rs b/codex-rs/chatgpt/src/chatgpt_token.rs index 55ebc22a..55b6886c 100644 --- a/codex-rs/chatgpt/src/chatgpt_token.rs +++ b/codex-rs/chatgpt/src/chatgpt_token.rs @@ -18,7 +18,7 @@ pub fn set_chatgpt_token_data(value: TokenData) { /// Initialize the ChatGPT token from auth.json file pub async fn init_chatgpt_token_from_auth(codex_home: &Path) -> std::io::Result<()> { - let auth = codex_login::load_auth(codex_home)?; + let auth = codex_login::load_auth(codex_home, true)?; if let Some(auth) = auth { let token_data = auth.get_token_data().await?; set_chatgpt_token_data(token_data); diff --git a/codex-rs/cli/src/login.rs b/codex-rs/cli/src/login.rs index 390c3100..4fa13f0c 100644 --- a/codex-rs/cli/src/login.rs +++ b/codex-rs/cli/src/login.rs @@ -1,8 +1,12 @@ +use std::env; + use codex_common::CliConfigOverrides; use codex_core::config::Config; use codex_core::config::ConfigOverrides; use codex_login::AuthMode; +use codex_login::OPENAI_API_KEY_ENV_VAR; use codex_login::load_auth; +use codex_login::login_with_api_key; use codex_login::login_with_chatgpt; pub async fn run_login_with_chatgpt(cli_config_overrides: CliConfigOverrides) -> ! { @@ -21,14 +25,40 @@ pub async fn run_login_with_chatgpt(cli_config_overrides: CliConfigOverrides) -> } } +pub async fn run_login_with_api_key( + cli_config_overrides: CliConfigOverrides, + api_key: String, +) -> ! { + let config = load_config_or_exit(cli_config_overrides); + + match login_with_api_key(&config.codex_home, &api_key) { + Ok(_) => { + eprintln!("Successfully logged in"); + std::process::exit(0); + } + Err(e) => { + eprintln!("Error logging in: {e}"); + std::process::exit(1); + } + } +} + pub async fn run_login_status(cli_config_overrides: CliConfigOverrides) -> ! { let config = load_config_or_exit(cli_config_overrides); - match load_auth(&config.codex_home) { + match load_auth(&config.codex_home, true) { Ok(Some(auth)) => match auth.mode { AuthMode::ApiKey => { if let Some(api_key) = auth.api_key.as_deref() { eprintln!("Logged in using an API key - {}", safe_format_key(api_key)); + + if let Ok(env_api_key) = env::var(OPENAI_API_KEY_ENV_VAR) { + if env_api_key == api_key { + eprintln!( + " API loaded from OPENAI_API_KEY environment variable or .env file" + ); + } + } } else { eprintln!("Logged in using an API key"); } diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs index c5fd69f9..27f83121 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -8,6 +8,7 @@ use codex_chatgpt::apply_command::run_apply_command; use codex_cli::LandlockCommand; use codex_cli::SeatbeltCommand; use codex_cli::login::run_login_status; +use codex_cli::login::run_login_with_api_key; use codex_cli::login::run_login_with_chatgpt; use codex_cli::proto; use codex_common::CliConfigOverrides; @@ -92,6 +93,9 @@ struct LoginCommand { #[clap(skip)] config_overrides: CliConfigOverrides, + #[arg(long = "api-key", value_name = "API_KEY")] + api_key: Option, + #[command(subcommand)] action: Option, } @@ -133,7 +137,11 @@ async fn cli_main(codex_linux_sandbox_exe: Option) -> anyhow::Result<() run_login_status(login_cli.config_overrides).await; } None => { - run_login_with_chatgpt(login_cli.config_overrides).await; + if let Some(api_key) = login_cli.api_key { + run_login_with_api_key(login_cli.config_overrides, api_key).await; + } else { + run_login_with_chatgpt(login_cli.config_overrides).await; + } } } } diff --git a/codex-rs/cli/src/proto.rs b/codex-rs/cli/src/proto.rs index 291e1680..9f9a94ed 100644 --- a/codex-rs/cli/src/proto.rs +++ b/codex-rs/cli/src/proto.rs @@ -36,7 +36,7 @@ pub async fn run_main(opts: ProtoCli) -> anyhow::Result<()> { .map_err(anyhow::Error::msg)?; let config = Config::load_with_cli_overrides(overrides_vec, ConfigOverrides::default())?; - let auth = load_auth(&config.codex_home)?; + let auth = load_auth(&config.codex_home, true)?; let ctrl_c = notify_on_sigint(); let CodexSpawnOk { codex, .. } = Codex::spawn(config, auth, ctrl_c.clone()).await?; let codex = Arc::new(codex); diff --git a/codex-rs/core/src/codex_wrapper.rs b/codex-rs/core/src/codex_wrapper.rs index 1e26a9eb..eeb4a7b4 100644 --- a/codex-rs/core/src/codex_wrapper.rs +++ b/codex-rs/core/src/codex_wrapper.rs @@ -26,7 +26,7 @@ pub struct CodexConversation { /// that callers can surface the information to the UI. pub async fn init_codex(config: Config) -> anyhow::Result { let ctrl_c = notify_on_sigint(); - let auth = load_auth(&config.codex_home)?; + let auth = load_auth(&config.codex_home, true)?; let CodexSpawnOk { codex, init_id, diff --git a/codex-rs/core/tests/client.rs b/codex-rs/core/tests/client.rs index 67d95cb8..1286928e 100644 --- a/codex-rs/core/tests/client.rs +++ b/codex-rs/core/tests/client.rs @@ -327,14 +327,14 @@ fn auth_from_token(id_token: String) -> CodexAuth { AuthMode::ChatGPT, PathBuf::new(), Some(AuthDotJson { - tokens: TokenData { + openai_api_key: None, + tokens: Some(TokenData { id_token, access_token: "Access Token".to_string(), refresh_token: "test".to_string(), account_id: None, - }, - last_refresh: Utc::now(), - openai_api_key: None, + }), + last_refresh: Some(Utc::now()), }), ) } diff --git a/codex-rs/login/Cargo.toml b/codex-rs/login/Cargo.toml index e10666b0..650291b3 100644 --- a/codex-rs/login/Cargo.toml +++ b/codex-rs/login/Cargo.toml @@ -18,3 +18,6 @@ tokio = { version = "1", features = [ "rt-multi-thread", "signal", ] } + +[dev-dependencies] +tempfile = "3" diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index 47dbbca9..2f0aeb60 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -20,7 +20,7 @@ use tokio::process::Command; const SOURCE_FOR_PYTHON_SERVER: &str = include_str!("./login_with_chatgpt.py"); const CLIENT_ID: &str = "app_EMoamEEZ73f0CkXaXp7hrann"; -const OPENAI_API_KEY_ENV_VAR: &str = "OPENAI_API_KEY"; +pub const OPENAI_API_KEY_ENV_VAR: &str = "OPENAI_API_KEY"; #[derive(Clone, Debug, PartialEq)] pub enum AuthMode { @@ -70,13 +70,16 @@ impl CodexAuth { pub async fn get_token_data(&self) -> Result { #[expect(clippy::unwrap_used)] let auth_dot_json = self.auth_dot_json.lock().unwrap().clone(); - match auth_dot_json { - Some(auth_dot_json) => { - if auth_dot_json.last_refresh < Utc::now() - chrono::Duration::days(28) { + Some(AuthDotJson { + tokens: Some(mut tokens), + last_refresh: Some(last_refresh), + .. + }) => { + if last_refresh < Utc::now() - chrono::Duration::days(28) { let refresh_response = tokio::time::timeout( Duration::from_secs(60), - try_refresh_token(auth_dot_json.tokens.refresh_token.clone()), + try_refresh_token(tokens.refresh_token.clone()), ) .await .map_err(|_| { @@ -92,13 +95,21 @@ impl CodexAuth { ) .await?; + tokens = updated_auth_dot_json + .tokens + .clone() + .ok_or(std::io::Error::other( + "Token data is not available after refresh.", + ))?; + #[expect(clippy::unwrap_used)] - let mut auth_dot_json = self.auth_dot_json.lock().unwrap(); - *auth_dot_json = Some(updated_auth_dot_json); + let mut auth_lock = self.auth_dot_json.lock().unwrap(); + *auth_lock = Some(updated_auth_dot_json); } - Ok(auth_dot_json.tokens.clone()) + + Ok(tokens) } - None => Err(std::io::Error::other("Token data is not available.")), + _ => Err(std::io::Error::other("Token data is not available.")), } } @@ -115,8 +126,8 @@ impl CodexAuth { } // Loads the available auth information from the auth.json or OPENAI_API_KEY environment variable. -pub fn load_auth(codex_home: &Path) -> std::io::Result> { - let auth_file = codex_home.join("auth.json"); +pub fn load_auth(codex_home: &Path, include_env_var: bool) -> std::io::Result> { + let auth_file = get_auth_file(codex_home); let auth_dot_json = try_read_auth_json(&auth_file).ok(); @@ -125,12 +136,21 @@ pub fn load_auth(codex_home: &Path) -> std::io::Result> { .and_then(|a| a.openai_api_key.clone()) .filter(|s| !s.is_empty()); - let openai_api_key = env::var(OPENAI_API_KEY_ENV_VAR) - .ok() - .filter(|s| !s.is_empty()) - .or(auth_json_api_key); + let openai_api_key = if include_env_var { + env::var(OPENAI_API_KEY_ENV_VAR) + .ok() + .filter(|s| !s.is_empty()) + .or(auth_json_api_key) + } else { + auth_json_api_key + }; - if openai_api_key.is_none() && auth_dot_json.is_none() { + let has_tokens = auth_dot_json + .as_ref() + .and_then(|a| a.tokens.as_ref()) + .is_some(); + + if openai_api_key.is_none() && !has_tokens { return Ok(None); } @@ -148,6 +168,10 @@ pub fn load_auth(codex_home: &Path) -> std::io::Result> { })) } +fn get_auth_file(codex_home: &Path) -> PathBuf { + codex_home.join("auth.json") +} + /// Run `python3 -c {{SOURCE_FOR_PYTHON_SERVER}}` with the CODEX_HOME /// environment variable set to the provided `codex_home` path. If the /// subprocess exits 0, read the OPENAI_API_KEY property out of @@ -187,6 +211,15 @@ pub async fn login_with_chatgpt(codex_home: &Path, capture_output: bool) -> std: } } +pub fn login_with_api_key(codex_home: &Path, api_key: &str) -> std::io::Result<()> { + let auth_dot_json = AuthDotJson { + openai_api_key: Some(api_key.to_string()), + tokens: None, + last_refresh: None, + }; + write_auth_json(&get_auth_file(codex_home), &auth_dot_json) +} + /// Attempt to read and refresh the `auth.json` file in the given `CODEX_HOME` directory. /// Returns the full AuthDotJson structure after refreshing if necessary. pub fn try_read_auth_json(auth_file: &Path) -> std::io::Result { @@ -198,35 +231,38 @@ pub fn try_read_auth_json(auth_file: &Path) -> std::io::Result { Ok(auth_dot_json) } -async fn update_tokens( - auth_file: &Path, - id_token: String, - access_token: Option, - refresh_token: Option, -) -> std::io::Result { +fn write_auth_json(auth_file: &Path, auth_dot_json: &AuthDotJson) -> std::io::Result<()> { + let json_data = serde_json::to_string_pretty(auth_dot_json)?; let mut options = OpenOptions::new(); options.truncate(true).write(true).create(true); #[cfg(unix)] { options.mode(0o600); } + let mut file = options.open(auth_file)?; + file.write_all(json_data.as_bytes())?; + file.flush()?; + Ok(()) +} + +async fn update_tokens( + auth_file: &Path, + id_token: String, + access_token: Option, + refresh_token: Option, +) -> std::io::Result { let mut auth_dot_json = try_read_auth_json(auth_file)?; - auth_dot_json.tokens.id_token = id_token.to_string(); + let tokens = auth_dot_json.tokens.get_or_insert_with(TokenData::default); + tokens.id_token = id_token.to_string(); if let Some(access_token) = access_token { - auth_dot_json.tokens.access_token = access_token.to_string(); + tokens.access_token = access_token.to_string(); } if let Some(refresh_token) = refresh_token { - auth_dot_json.tokens.refresh_token = refresh_token.to_string(); - } - auth_dot_json.last_refresh = Utc::now(); - - let json_data = serde_json::to_string_pretty(&auth_dot_json)?; - { - let mut file = options.open(auth_file)?; - file.write_all(json_data.as_bytes())?; - file.flush()?; + tokens.refresh_token = refresh_token.to_string(); } + auth_dot_json.last_refresh = Some(Utc::now()); + write_auth_json(auth_file, &auth_dot_json)?; Ok(auth_dot_json) } @@ -282,12 +318,14 @@ pub struct AuthDotJson { #[serde(rename = "OPENAI_API_KEY")] pub openai_api_key: Option, - pub tokens: TokenData, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub tokens: Option, - pub last_refresh: DateTime, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub last_refresh: Option>, } -#[derive(Deserialize, Serialize, Clone, Debug, PartialEq)] +#[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Default)] pub struct TokenData { /// This is a JWT. pub id_token: String, @@ -299,3 +337,95 @@ pub struct TokenData { pub account_id: Option, } + +#[cfg(test)] +mod tests { + use super::*; + use tempfile::tempdir; + + #[test] + #[expect(clippy::unwrap_used)] + fn writes_api_key_and_loads_auth() { + let dir = tempdir().unwrap(); + login_with_api_key(dir.path(), "sk-test-key").unwrap(); + let auth = load_auth(dir.path(), false).unwrap().unwrap(); + assert_eq!(auth.mode, AuthMode::ApiKey); + assert_eq!(auth.api_key.as_deref(), Some("sk-test-key")); + } + + #[test] + #[expect(clippy::unwrap_used)] + fn loads_from_env_var_if_env_var_exists() { + let dir = tempdir().unwrap(); + + let env_var = std::env::var(OPENAI_API_KEY_ENV_VAR); + + if let Ok(env_var) = env_var { + let auth = load_auth(dir.path(), true).unwrap().unwrap(); + assert_eq!(auth.mode, AuthMode::ApiKey); + assert_eq!(auth.api_key, Some(env_var)); + } + } + + #[tokio::test] + #[expect(clippy::unwrap_used)] + async fn loads_token_data_from_auth_json() { + let dir = tempdir().unwrap(); + let auth_file = dir.path().join("auth.json"); + std::fs::write( + auth_file, + format!( + r#" + {{ + "OPENAI_API_KEY": null, + "tokens": {{ + "id_token": "test-id-token", + "access_token": "test-access-token", + "refresh_token": "test-refresh-token" + }}, + "last_refresh": "{}" + }} + "#, + Utc::now().to_rfc3339() + ), + ) + .unwrap(); + + let auth = load_auth(dir.path(), false).unwrap().unwrap(); + assert_eq!(auth.mode, AuthMode::ChatGPT); + assert_eq!(auth.api_key, None); + assert_eq!( + auth.get_token_data().await.unwrap(), + TokenData { + id_token: "test-id-token".to_string(), + access_token: "test-access-token".to_string(), + refresh_token: "test-refresh-token".to_string(), + account_id: None, + } + ); + } + + #[tokio::test] + #[expect(clippy::unwrap_used)] + async fn loads_api_key_from_auth_json() { + let dir = tempdir().unwrap(); + let auth_file = dir.path().join("auth.json"); + std::fs::write( + auth_file, + r#" + { + "OPENAI_API_KEY": "sk-test-key", + "tokens": null, + "last_refresh": null + } + "#, + ) + .unwrap(); + + let auth = load_auth(dir.path(), false).unwrap().unwrap(); + assert_eq!(auth.mode, AuthMode::ApiKey); + assert_eq!(auth.api_key, Some("sk-test-key".to_string())); + + assert!(auth.get_token_data().await.is_err()); + } +} diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index 6b5fe7f7..7e987f6f 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -226,7 +226,7 @@ fn should_show_login_screen(config: &Config) -> bool { // Reading the OpenAI API key is an async operation because it may need // to refresh the token. Block on it. let codex_home = config.codex_home.clone(); - match load_auth(&codex_home) { + match load_auth(&codex_home, true) { Ok(Some(_)) => false, Ok(None) => true, Err(err) => {