diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 120050c227..87d59b21be 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 55ebc22a08..55b6886c59 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/debug_sandbox.rs b/codex-rs/cli/src/debug_sandbox.rs index 905b746168..7f0983cbc6 100644 --- a/codex-rs/cli/src/debug_sandbox.rs +++ b/codex-rs/cli/src/debug_sandbox.rs @@ -4,10 +4,10 @@ use codex_common::CliConfigOverrides; use codex_core::config::Config; use codex_core::config::ConfigOverrides; use codex_core::config_types::SandboxMode; -use codex_core::exec::StdioPolicy; use codex_core::exec::spawn_command_under_linux_sandbox; -use codex_core::exec::spawn_command_under_seatbelt; use codex_core::exec_env::create_env; +use codex_core::seatbelt::spawn_command_under_seatbelt; +use codex_core::spawn::StdioPolicy; use crate::LandlockCommand; use crate::SeatbeltCommand; diff --git a/codex-rs/cli/src/login.rs b/codex-rs/cli/src/login.rs index 390c310030..4fa13f0cc6 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 c5fd69f9cd..27f8312193 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 291e1680f1..9f9a94ed4d 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 1e26a9ebed..eeb4a7b470 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/src/exec.rs b/codex-rs/core/src/exec.rs index 230c4ec134..06416e6768 100644 --- a/codex-rs/core/src/exec.rs +++ b/codex-rs/core/src/exec.rs @@ -6,7 +6,6 @@ use std::io; use std::path::Path; use std::path::PathBuf; use std::process::ExitStatus; -use std::process::Stdio; use std::sync::Arc; use std::time::Duration; use std::time::Instant; @@ -15,14 +14,15 @@ use tokio::io::AsyncRead; use tokio::io::AsyncReadExt; use tokio::io::BufReader; use tokio::process::Child; -use tokio::process::Command; use tokio::sync::Notify; -use tracing::trace; use crate::error::CodexErr; use crate::error::Result; use crate::error::SandboxErr; use crate::protocol::SandboxPolicy; +use crate::seatbelt::spawn_command_under_seatbelt; +use crate::spawn::StdioPolicy; +use crate::spawn::spawn_child_async; // Maximum we send for each stream, which is either: // - 10KiB OR @@ -37,24 +37,6 @@ const DEFAULT_TIMEOUT_MS: u64 = 10_000; const SIGKILL_CODE: i32 = 9; const TIMEOUT_CODE: i32 = 64; -const MACOS_SEATBELT_BASE_POLICY: &str = include_str!("seatbelt_base_policy.sbpl"); - -/// When working with `sandbox-exec`, only consider `sandbox-exec` in `/usr/bin` -/// to defend against an attacker trying to inject a malicious version on the -/// PATH. If /usr/bin/sandbox-exec has been tampered with, then the attacker -/// already has root access. -const MACOS_PATH_TO_SEATBELT_EXECUTABLE: &str = "/usr/bin/sandbox-exec"; - -/// Experimental environment variable that will be set to some non-empty value -/// if both of the following are true: -/// -/// 1. The process was spawned by Codex as part of a shell tool call. -/// 2. SandboxPolicy.has_full_network_access() was false for the tool call. -/// -/// We may try to have just one environment variable for all sandboxing -/// attributes, so this may change in the future. -pub const CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR: &str = "CODEX_SANDBOX_NETWORK_DISABLED"; - #[derive(Debug, Clone)] pub struct ExecParams { pub command: Vec, @@ -168,27 +150,6 @@ pub async fn process_exec_tool_call( } } -pub async fn spawn_command_under_seatbelt( - command: Vec, - sandbox_policy: &SandboxPolicy, - cwd: PathBuf, - stdio_policy: StdioPolicy, - env: HashMap, -) -> std::io::Result { - let args = create_seatbelt_command_args(command, sandbox_policy, &cwd); - let arg0 = None; - spawn_child_async( - PathBuf::from(MACOS_PATH_TO_SEATBELT_EXECUTABLE), - args, - arg0, - cwd, - sandbox_policy, - stdio_policy, - env, - ) - .await -} - /// Spawn a shell tool command under the Linux Landlock+seccomp sandbox helper /// (codex-linux-sandbox). /// @@ -248,65 +209,6 @@ fn create_linux_sandbox_command_args( linux_cmd } -fn create_seatbelt_command_args( - command: Vec, - sandbox_policy: &SandboxPolicy, - cwd: &Path, -) -> Vec { - let (file_write_policy, extra_cli_args) = { - if sandbox_policy.has_full_disk_write_access() { - // Allegedly, this is more permissive than `(allow file-write*)`. - ( - r#"(allow file-write* (regex #"^/"))"#.to_string(), - Vec::::new(), - ) - } else { - let writable_roots = sandbox_policy.get_writable_roots_with_cwd(cwd); - let (writable_folder_policies, cli_args): (Vec, Vec) = writable_roots - .iter() - .enumerate() - .map(|(index, root)| { - let param_name = format!("WRITABLE_ROOT_{index}"); - let policy: String = format!("(subpath (param \"{param_name}\"))"); - let cli_arg = format!("-D{param_name}={}", root.to_string_lossy()); - (policy, cli_arg) - }) - .unzip(); - if writable_folder_policies.is_empty() { - ("".to_string(), Vec::::new()) - } else { - let file_write_policy = format!( - "(allow file-write*\n{}\n)", - writable_folder_policies.join(" ") - ); - (file_write_policy, cli_args) - } - } - }; - - let file_read_policy = if sandbox_policy.has_full_disk_read_access() { - "; allow read-only file operations\n(allow file-read*)" - } else { - "" - }; - - // TODO(mbolin): apply_patch calls must also honor the SandboxPolicy. - let network_policy = if sandbox_policy.has_full_network_access() { - "(allow network-outbound)\n(allow network-inbound)\n(allow system-socket)" - } else { - "" - }; - - let full_policy = format!( - "{MACOS_SEATBELT_BASE_POLICY}\n{file_read_policy}\n{file_write_policy}\n{network_policy}" - ); - let mut seatbelt_args: Vec = vec!["-p".to_string(), full_policy]; - seatbelt_args.extend(extra_cli_args); - seatbelt_args.push("--".to_string()); - seatbelt_args.extend(command); - seatbelt_args -} - #[derive(Debug)] pub struct RawExecToolCallOutput { pub exit_status: ExitStatus, @@ -352,90 +254,6 @@ async fn exec( consume_truncated_output(child, ctrl_c, timeout_ms).await } -#[derive(Debug, Clone, Copy)] -pub enum StdioPolicy { - RedirectForShellTool, - Inherit, -} - -/// Spawns the appropriate child process for the ExecParams and SandboxPolicy, -/// ensuring the args and environment variables used to create the `Command` -/// (and `Child`) honor the configuration. -/// -/// For now, we take `SandboxPolicy` as a parameter to spawn_child() because -/// we need to determine whether to set the -/// `CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR` environment variable. -async fn spawn_child_async( - program: PathBuf, - args: Vec, - #[cfg_attr(not(unix), allow(unused_variables))] arg0: Option<&str>, - cwd: PathBuf, - sandbox_policy: &SandboxPolicy, - stdio_policy: StdioPolicy, - env: HashMap, -) -> std::io::Result { - trace!( - "spawn_child_async: {program:?} {args:?} {arg0:?} {cwd:?} {sandbox_policy:?} {stdio_policy:?} {env:?}" - ); - - let mut cmd = Command::new(&program); - #[cfg(unix)] - cmd.arg0(arg0.map_or_else(|| program.to_string_lossy().to_string(), String::from)); - cmd.args(args); - cmd.current_dir(cwd); - cmd.env_clear(); - cmd.envs(env); - - if !sandbox_policy.has_full_network_access() { - cmd.env(CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR, "1"); - } - - // If this Codex process dies (including being killed via SIGKILL), we want - // any child processes that were spawned as part of a `"shell"` tool call - // to also be terminated. - - // This relies on prctl(2), so it only works on Linux. - #[cfg(target_os = "linux")] - unsafe { - cmd.pre_exec(|| { - // This prctl call effectively requests, "deliver SIGTERM when my - // current parent dies." - if libc::prctl(libc::PR_SET_PDEATHSIG, libc::SIGTERM) == -1 { - return Err(io::Error::last_os_error()); - } - - // Though if there was a race condition and this pre_exec() block is - // run _after_ the parent (i.e., the Codex process) has already - // exited, then the parent is the _init_ process (which will never - // die), so we should just terminate the child process now. - if libc::getppid() == 1 { - libc::raise(libc::SIGTERM); - } - Ok(()) - }); - } - - match stdio_policy { - StdioPolicy::RedirectForShellTool => { - // Do not create a file descriptor for stdin because otherwise some - // commands may hang forever waiting for input. For example, ripgrep has - // a heuristic where it may try to read from stdin as explained here: - // https://github.com/BurntSushi/ripgrep/blob/e2362d4d5185d02fa857bf381e7bd52e66fafc73/crates/core/flags/hiargs.rs#L1101-L1103 - cmd.stdin(Stdio::null()); - - cmd.stdout(Stdio::piped()).stderr(Stdio::piped()); - } - StdioPolicy::Inherit => { - // Inherit stdin, stdout, and stderr from the parent process. - cmd.stdin(Stdio::inherit()) - .stdout(Stdio::inherit()) - .stderr(Stdio::inherit()); - } - } - - cmd.kill_on_drop(true).spawn() -} - /// Consumes the output of a child process, truncating it so it is suitable for /// use as the output of a `shell` tool call. Also enforces specified timeout. pub(crate) async fn consume_truncated_output( diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 054abd742a..1b5a2a6a20 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -39,9 +39,12 @@ mod project_doc; pub mod protocol; mod rollout; mod safety; +pub mod seatbelt; pub mod shell; +pub mod spawn; mod user_notification; pub mod util; pub use apply_patch::CODEX_APPLY_PATCH_ARG1; pub use client_common::model_supports_reasoning_summaries; +pub use spawn::CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR; diff --git a/codex-rs/core/src/seatbelt.rs b/codex-rs/core/src/seatbelt.rs new file mode 100644 index 0000000000..be2acb1bdc --- /dev/null +++ b/codex-rs/core/src/seatbelt.rs @@ -0,0 +1,96 @@ +use std::collections::HashMap; +use std::path::Path; +use std::path::PathBuf; +use tokio::process::Child; + +use crate::protocol::SandboxPolicy; +use crate::spawn::StdioPolicy; +use crate::spawn::spawn_child_async; + +const MACOS_SEATBELT_BASE_POLICY: &str = include_str!("seatbelt_base_policy.sbpl"); + +/// When working with `sandbox-exec`, only consider `sandbox-exec` in `/usr/bin` +/// to defend against an attacker trying to inject a malicious version on the +/// PATH. If /usr/bin/sandbox-exec has been tampered with, then the attacker +/// already has root access. +const MACOS_PATH_TO_SEATBELT_EXECUTABLE: &str = "/usr/bin/sandbox-exec"; + +pub async fn spawn_command_under_seatbelt( + command: Vec, + sandbox_policy: &SandboxPolicy, + cwd: PathBuf, + stdio_policy: StdioPolicy, + env: HashMap, +) -> std::io::Result { + let args = create_seatbelt_command_args(command, sandbox_policy, &cwd); + let arg0 = None; + spawn_child_async( + PathBuf::from(MACOS_PATH_TO_SEATBELT_EXECUTABLE), + args, + arg0, + cwd, + sandbox_policy, + stdio_policy, + env, + ) + .await +} + +fn create_seatbelt_command_args( + command: Vec, + sandbox_policy: &SandboxPolicy, + cwd: &Path, +) -> Vec { + let (file_write_policy, extra_cli_args) = { + if sandbox_policy.has_full_disk_write_access() { + // Allegedly, this is more permissive than `(allow file-write*)`. + ( + r#"(allow file-write* (regex #"^/"))"#.to_string(), + Vec::::new(), + ) + } else { + let writable_roots = sandbox_policy.get_writable_roots_with_cwd(cwd); + let (writable_folder_policies, cli_args): (Vec, Vec) = writable_roots + .iter() + .enumerate() + .map(|(index, root)| { + let param_name = format!("WRITABLE_ROOT_{index}"); + let policy: String = format!("(subpath (param \"{param_name}\"))"); + let cli_arg = format!("-D{param_name}={}", root.to_string_lossy()); + (policy, cli_arg) + }) + .unzip(); + if writable_folder_policies.is_empty() { + ("".to_string(), Vec::::new()) + } else { + let file_write_policy = format!( + "(allow file-write*\n{}\n)", + writable_folder_policies.join(" ") + ); + (file_write_policy, cli_args) + } + } + }; + + let file_read_policy = if sandbox_policy.has_full_disk_read_access() { + "; allow read-only file operations\n(allow file-read*)" + } else { + "" + }; + + // TODO(mbolin): apply_patch calls must also honor the SandboxPolicy. + let network_policy = if sandbox_policy.has_full_network_access() { + "(allow network-outbound)\n(allow network-inbound)\n(allow system-socket)" + } else { + "" + }; + + let full_policy = format!( + "{MACOS_SEATBELT_BASE_POLICY}\n{file_read_policy}\n{file_write_policy}\n{network_policy}" + ); + let mut seatbelt_args: Vec = vec!["-p".to_string(), full_policy]; + seatbelt_args.extend(extra_cli_args); + seatbelt_args.push("--".to_string()); + seatbelt_args.extend(command); + seatbelt_args +} diff --git a/codex-rs/core/src/spawn.rs b/codex-rs/core/src/spawn.rs new file mode 100644 index 0000000000..5dab353642 --- /dev/null +++ b/codex-rs/core/src/spawn.rs @@ -0,0 +1,102 @@ +use std::collections::HashMap; +use std::path::PathBuf; +use std::process::Stdio; +use tokio::process::Child; +use tokio::process::Command; +use tracing::trace; + +use crate::protocol::SandboxPolicy; + +/// Experimental environment variable that will be set to some non-empty value +/// if both of the following are true: +/// +/// 1. The process was spawned by Codex as part of a shell tool call. +/// 2. SandboxPolicy.has_full_network_access() was false for the tool call. +/// +/// We may try to have just one environment variable for all sandboxing +/// attributes, so this may change in the future. +pub const CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR: &str = "CODEX_SANDBOX_NETWORK_DISABLED"; + +#[derive(Debug, Clone, Copy)] +pub enum StdioPolicy { + RedirectForShellTool, + Inherit, +} + +/// Spawns the appropriate child process for the ExecParams and SandboxPolicy, +/// ensuring the args and environment variables used to create the `Command` +/// (and `Child`) honor the configuration. +/// +/// For now, we take `SandboxPolicy` as a parameter to spawn_child() because +/// we need to determine whether to set the +/// `CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR` environment variable. +pub(crate) async fn spawn_child_async( + program: PathBuf, + args: Vec, + #[cfg_attr(not(unix), allow(unused_variables))] arg0: Option<&str>, + cwd: PathBuf, + sandbox_policy: &SandboxPolicy, + stdio_policy: StdioPolicy, + env: HashMap, +) -> std::io::Result { + trace!( + "spawn_child_async: {program:?} {args:?} {arg0:?} {cwd:?} {sandbox_policy:?} {stdio_policy:?} {env:?}" + ); + + let mut cmd = Command::new(&program); + #[cfg(unix)] + cmd.arg0(arg0.map_or_else(|| program.to_string_lossy().to_string(), String::from)); + cmd.args(args); + cmd.current_dir(cwd); + cmd.env_clear(); + cmd.envs(env); + + if !sandbox_policy.has_full_network_access() { + cmd.env(CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR, "1"); + } + + // If this Codex process dies (including being killed via SIGKILL), we want + // any child processes that were spawned as part of a `"shell"` tool call + // to also be terminated. + + // This relies on prctl(2), so it only works on Linux. + #[cfg(target_os = "linux")] + unsafe { + cmd.pre_exec(|| { + // This prctl call effectively requests, "deliver SIGTERM when my + // current parent dies." + if libc::prctl(libc::PR_SET_PDEATHSIG, libc::SIGTERM) == -1 { + return Err(io::Error::last_os_error()); + } + + // Though if there was a race condition and this pre_exec() block is + // run _after_ the parent (i.e., the Codex process) has already + // exited, then the parent is the _init_ process (which will never + // die), so we should just terminate the child process now. + if libc::getppid() == 1 { + libc::raise(libc::SIGTERM); + } + Ok(()) + }); + } + + match stdio_policy { + StdioPolicy::RedirectForShellTool => { + // Do not create a file descriptor for stdin because otherwise some + // commands may hang forever waiting for input. For example, ripgrep has + // a heuristic where it may try to read from stdin as explained here: + // https://github.com/BurntSushi/ripgrep/blob/e2362d4d5185d02fa857bf381e7bd52e66fafc73/crates/core/flags/hiargs.rs#L1101-L1103 + cmd.stdin(Stdio::null()); + + cmd.stdout(Stdio::piped()).stderr(Stdio::piped()); + } + StdioPolicy::Inherit => { + // Inherit stdin, stdout, and stderr from the parent process. + cmd.stdin(Stdio::inherit()) + .stdout(Stdio::inherit()) + .stderr(Stdio::inherit()); + } + } + + cmd.kill_on_drop(true).spawn() +} diff --git a/codex-rs/core/tests/cli_stream.rs b/codex-rs/core/tests/cli_stream.rs index ee0377fc10..be45240f85 100644 --- a/codex-rs/core/tests/cli_stream.rs +++ b/codex-rs/core/tests/cli_stream.rs @@ -1,7 +1,7 @@ #![expect(clippy::unwrap_used)] use assert_cmd::Command as AssertCommand; -use codex_core::exec::CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR; +use codex_core::CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR; use std::time::Duration; use std::time::Instant; use tempfile::TempDir; diff --git a/codex-rs/core/tests/client.rs b/codex-rs/core/tests/client.rs index 67d95cb8f6..5de2552495 100644 --- a/codex-rs/core/tests/client.rs +++ b/codex-rs/core/tests/client.rs @@ -1,11 +1,11 @@ use std::path::PathBuf; use chrono::Utc; +use codex_core::CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR; use codex_core::Codex; use codex_core::CodexSpawnOk; use codex_core::ModelProviderInfo; use codex_core::built_in_model_providers; -use codex_core::exec::CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR; use codex_core::protocol::EventMsg; use codex_core::protocol::InputItem; use codex_core::protocol::Op; @@ -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/core/tests/stream_no_completed.rs b/codex-rs/core/tests/stream_no_completed.rs index d2fc035569..2efd31db9f 100644 --- a/codex-rs/core/tests/stream_no_completed.rs +++ b/codex-rs/core/tests/stream_no_completed.rs @@ -3,10 +3,10 @@ use std::time::Duration; +use codex_core::CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR; use codex_core::Codex; use codex_core::CodexSpawnOk; use codex_core::ModelProviderInfo; -use codex_core::exec::CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR; use codex_core::protocol::EventMsg; use codex_core::protocol::InputItem; use codex_core::protocol::Op; diff --git a/codex-rs/login/Cargo.toml b/codex-rs/login/Cargo.toml index e10666b092..650291b3bc 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 47dbbca9fb..2f0aeb60bd 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/mcp-server/tests/codex_tool.rs b/codex-rs/mcp-server/tests/codex_tool.rs index 0f06483f24..cf8abcdcef 100644 --- a/codex-rs/mcp-server/tests/codex_tool.rs +++ b/codex-rs/mcp-server/tests/codex_tool.rs @@ -3,7 +3,7 @@ use std::env; use std::path::Path; use std::path::PathBuf; -use codex_core::exec::CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR; +use codex_core::CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR; use codex_core::protocol::FileChange; use codex_core::protocol::ReviewDecision; use codex_mcp_server::CodexToolCallParam; diff --git a/codex-rs/mcp-server/tests/interrupt.rs b/codex-rs/mcp-server/tests/interrupt.rs index 313bc7afab..dc2474df0b 100644 --- a/codex-rs/mcp-server/tests/interrupt.rs +++ b/codex-rs/mcp-server/tests/interrupt.rs @@ -3,7 +3,7 @@ use std::path::Path; -use codex_core::exec::CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR; +use codex_core::CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR; use codex_mcp_server::CodexToolCallParam; use mcp_types::JSONRPCResponse; use mcp_types::RequestId; diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index 6b5fe7f7ae..7e987f6ff4 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) => {