diff --git a/rust/src/providers/codex/api.rs b/rust/src/providers/codex/api.rs index 6ce382deab..004c5dbdf2 100755 --- a/rust/src/providers/codex/api.rs +++ b/rust/src/providers/codex/api.rs @@ -21,6 +21,11 @@ const DEFAULT_BASE_URL: &str = "https://chatgpt.com/backend-api"; const USAGE_PATH: &str = "/wham/usage"; const RESET_CREDITS_PATH: &str = "/wham/rate-limit-reset-credits"; const CREDENTIAL_CACHE_TTL: Duration = Duration::from_secs(5); +/// Upstream 0.69.0 #4088: the Codex CLI owns `auth.json` and may be publishing a +/// replacement while we read it. A failed or stale read is repeated up to this many +/// times, `CREDENTIAL_READ_RETRY_DELAY` apart, before the error is reported. +const CREDENTIAL_READ_RETRIES: u32 = 2; +const CREDENTIAL_READ_RETRY_DELAY: Duration = Duration::from_millis(50); const EXTERNAL_OAUTH_REFRESH_WINDOW: chrono::TimeDelta = chrono::Duration::minutes(5); static CREDENTIAL_CACHE: OnceLock>> = OnceLock::new(); @@ -113,7 +118,7 @@ impl CodexApi { pub async fn fetch_usage( &self, ) -> Result<(UsageSnapshot, Option, Option), ProviderError> { - let creds = self.load_credentials()?; + let creds = self.load_credentials().await?; let base_url = self.resolve_base_url(); let auth_path = self.get_auth_path(); let scope = weekly_reset::scope_key(creds.account_id.as_deref(), &auth_path); @@ -306,37 +311,54 @@ impl CodexApi { decode_reset_credits(&response.bytes().await?) } - fn load_credentials(&self) -> Result { - let auth_path = self.get_auth_path(); - - if !auth_path.exists() { - // Upstream 0.50.0 #2679: when the CLI targets Amazon Bedrock or - // another custom backend without ChatGPT auth, sign-in guidance - // is wrong — rate limits simply are not available there. - if self.uses_custom_backend() { - return Err(ProviderError::NotInstalled( - "Codex uses a custom backend (chatgpt_base_url / model_provider) without \ - ChatGPT auth. ChatGPT rate limits are unavailable for this setup." - .to_string(), - )); + /// Load credentials, tolerating a brief owner publication of `auth.json`. + /// + /// Upstream 0.69.0 #4088 (`CodexOAuthFetchStrategy.loadCredentials` on the + /// usage path, `retryStale: true`): every failed read is repeated. That covers + /// a missing (`NotInstalled`), unreadable (`Other`), malformed or incomplete + /// (`Parse`) file, and a credential the gate rejects as stale (`AuthRequired`, + /// such as a token inside its renewal window), because the CLI may be + /// publishing its renewal. This only rereads the file: no token is redeemed, + /// nothing is written, and the credential cache semantics are unchanged. After + /// the last read the error keeps its category, so unchanged stale credentials + /// still need their owner's renewal. + async fn load_credentials(&self) -> Result { + Self::reread_during_owner_publication(|| self.load_credentials_once()).await + } + + /// The bounded reread behind [`Self::load_credentials`]: `read` runs up to + /// `1 + CREDENTIAL_READ_RETRIES` times, `CREDENTIAL_READ_RETRY_DELAY` apart, + /// until it succeeds, and the last result is returned unchanged. Dropping the + /// returned future cancels the pending delay and any further read (upstream + /// checks task cancellation before each read). + async fn reread_during_owner_publication( + mut read: impl FnMut() -> Result, + ) -> Result { + let mut retries_remaining = CREDENTIAL_READ_RETRIES; + loop { + match read() { + Err(_) if retries_remaining > 0 => { + retries_remaining -= 1; + tokio::time::sleep(CREDENTIAL_READ_RETRY_DELAY).await; + } + result => return result, } - return Err(ProviderError::NotInstalled( - "Codex auth.json not found. Run `codex login` in a terminal to sign in." - .to_string(), - )); } + } - let modified = std::fs::metadata(&auth_path) - .ok() - .and_then(|metadata| metadata.modified().ok()); + fn load_credentials_once(&self) -> Result { + let auth_path = self.get_auth_path(); + + let metadata = + std::fs::metadata(&auth_path).map_err(|error| self.credential_file_error(error))?; + let modified = metadata.modified().ok(); if let Some(cached) = Self::cached_credentials(&auth_path, modified) { Self::enforce_external_oauth_gate(&cached)?; return Ok(cached); } - let content = std::fs::read_to_string(&auth_path).map_err(|e| { - ProviderError::Other(format!("Failed to read Codex credentials: {}", e)) - })?; + let content = std::fs::read_to_string(&auth_path) + .map_err(|error| self.credential_file_error(error))?; let credentials = Self::parse_credentials_json(&content)?; Self::enforce_external_oauth_gate(&credentials)?; @@ -344,6 +366,31 @@ impl CodexApi { Ok(credentials) } + fn missing_credentials_error(&self) -> ProviderError { + // Upstream 0.50.0 #2679: when the CLI targets Amazon Bedrock or + // another custom backend without ChatGPT auth, sign-in guidance + // is wrong — rate limits simply are not available there. + if self.uses_custom_backend() { + return ProviderError::NotInstalled( + "Codex uses a custom backend (chatgpt_base_url / model_provider) without \ + ChatGPT auth. ChatGPT rate limits are unavailable for this setup." + .to_string(), + ); + } + + ProviderError::NotInstalled( + "Codex auth.json not found. Run `codex login` in a terminal to sign in.".to_string(), + ) + } + + fn credential_file_error(&self, error: std::io::Error) -> ProviderError { + if error.kind() == std::io::ErrorKind::NotFound { + return self.missing_credentials_error(); + } + + ProviderError::Other(format!("Failed to read Codex credentials: {error}")) + } + fn parse_credentials_json(content: &str) -> Result { let json: serde_json::Value = serde_json::from_str(content) .map_err(|e| ProviderError::Parse(format!("Invalid Codex credentials JSON: {}", e)))?; @@ -414,11 +461,14 @@ impl CodexApi { /// request may use it. When the access token is a JWT, its native expiry /// is the validity authority; opaque tokens are sent to the server. fn enforce_external_oauth_gate(credentials: &CodexCredentials) -> Result<(), ProviderError> { - Self::enforce_external_oauth_gate_at( - credentials, - crate::settings::Settings::load().codex_external_oauth_sources_allowed, - Utc::now(), - ) + if !credentials.is_external_oauth { + return Ok(()); + } + // The opt-in only matters without refresh provenance. Skip the settings + // load otherwise: credential reads repeat while the owner publishes. + let external_sources_allowed = credentials.last_refresh.is_some() + || crate::settings::Settings::load().codex_external_oauth_sources_allowed; + Self::enforce_external_oauth_gate_at(credentials, external_sources_allowed, Utc::now()) } fn enforce_external_oauth_gate_at( @@ -1350,6 +1400,9 @@ fn capitalize(s: &str) -> String { } } +#[cfg(test)] +mod credential_retry_tests; + #[cfg(test)] mod tests { use super::*; diff --git a/rust/src/providers/codex/api/credential_retry_tests.rs b/rust/src/providers/codex/api/credential_retry_tests.rs new file mode 100644 index 0000000000..7d9aca1163 --- /dev/null +++ b/rust/src/providers/codex/api/credential_retry_tests.rs @@ -0,0 +1,210 @@ +//! Upstream 0.69.0 #4088: bounded reread of `auth.json` while its owner +//! publishes it (`CodexOAuthExpiryPipelineTests`). +//! +//! The loop tests drive [`CodexApi::reread_during_owner_publication`] with the +//! real file read, parser and gate, and count reads the way upstream's +//! injected reader does. The clock is paused, so the 50 ms retry delays are +//! exact and the tests never race a writer thread. + +use super::*; +use std::cell::Cell; +use std::path::Path; +use std::time::Duration; + +/// Upstream's fresh fixture expiry (2100-01-01T00:00:00Z). +const FRESH_EXPIRY: i64 = 4_102_444_800; + +fn api_for(home: &Path) -> CodexApi { + CodexApi::new().with_codex_home(home) +} + +/// A native OAuth `auth.json` whose JWT access token expires at `exp` +/// (seconds since the epoch). It carries `last_refresh`, so the gate decides +/// on expiry alone and never reads the host settings. +fn oauth_auth_json(exp: i64) -> String { + let payload = + base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(format!(r#"{{"exp":{exp}}}"#)); + format!( + r#"{{"last_refresh":"2000-01-01T00:00:00Z","tokens":{{"access_token":"h.{payload}.s","refresh_token":"r","account_id":"acct_pub"}}}}"# + ) +} + +fn access_token_of(auth_json: &str) -> String { + let json: serde_json::Value = serde_json::from_str(auth_json).expect("fixture json"); + json["tokens"]["access_token"] + .as_str() + .expect("fixture access token") + .to_string() +} + +/// Replace `auth.json` atomically, the way the Codex CLI publishes it. +fn publish(home: &Path, contents: &str) { + let staging = home.join("auth.json.tmp"); + std::fs::write(&staging, contents).expect("stage auth.json"); + std::fs::rename(&staging, home.join("auth.json")).expect("publish auth.json"); +} + +/// Leave `auth.json` in the named state of an unusable or in-progress file. +fn write_state(home: &Path, state: &str) { + let auth = home.join("auth.json"); + match state { + "missing" => {} + // A directory named auth.json exists but cannot be read as a file. + "unreadable" => std::fs::create_dir(&auth).expect("auth.json directory"), + "partial" => std::fs::write(&auth, r#"{"tokens":"#).expect("partial auth.json"), + "incomplete" => std::fs::write(&auth, r#"{"tokens":{}}"#).expect("incomplete auth.json"), + "expired" => std::fs::write(&auth, oauth_auth_json(1)).expect("expired auth.json"), + "near-expiry" => std::fs::write(&auth, oauth_auth_json(Utc::now().timestamp() + 120)) + .expect("near-expiry auth.json"), + other => panic!("unknown auth.json state {other}"), + } +} + +fn auth_bytes(home: &Path) -> Option> { + std::fs::read(home.join("auth.json")).ok() +} + +/// Upstream `OAuth fetch retries an owner publication in progress`: the first +/// read sees the publication in progress (missing, torn, incomplete, expired +/// or inside the renewal window), the second read sees the owner's +/// replacement, keeps its workspace, and nothing is written back. +#[tokio::test(start_paused = true)] +async fn usage_read_retries_an_owner_publication_in_progress() { + for publication in ["missing", "partial", "incomplete", "expired", "near-expiry"] { + let dir = tempfile::tempdir().expect("codex home"); + let api = api_for(dir.path()); + let fresh = oauth_auth_json(FRESH_EXPIRY); + write_state(dir.path(), publication); + let reads = Cell::new(0); + + let credentials = CodexApi::reread_during_owner_publication(|| { + reads.set(reads.get() + 1); + let read = api.load_credentials_once(); + if reads.get() == 1 { + publish(dir.path(), &fresh); + } + read + }) + .await + .unwrap_or_else(|error| panic!("{publication}: {error:?}")); + + assert_eq!( + credentials.access_token, + access_token_of(&fresh), + "{publication}" + ); + assert_eq!( + credentials.account_id.as_deref(), + Some("acct_pub"), + "{publication}" + ); + assert_eq!(reads.get(), 2, "{publication}"); + assert_eq!( + auth_bytes(dir.path()), + Some(fresh.into_bytes()), + "{publication}" + ); + } +} + +/// Upstream `OAuth read retries are bounded and preserve the final error`: +/// three reads, then the last failure keeps its category (unchanged stale +/// credentials still need their owner's renewal) and nothing is written. +#[tokio::test(start_paused = true)] +async fn usage_read_retries_are_bounded_and_keep_the_final_error() { + for failure in ["missing", "partial", "incomplete", "expired", "unreadable"] { + let dir = tempfile::tempdir().expect("codex home"); + let api = api_for(dir.path()); + write_state(dir.path(), failure); + let before = auth_bytes(dir.path()); + let reads = Cell::new(0); + + let error = CodexApi::reread_during_owner_publication(|| { + reads.set(reads.get() + 1); + api.load_credentials_once() + }) + .await + .err() + .unwrap_or_else(|| panic!("{failure}: expected a credential error")); + + let category_kept = match failure { + "missing" => matches!(error, ProviderError::NotInstalled(_)), + "partial" | "incomplete" => matches!(error, ProviderError::Parse(_)), + "expired" => matches!(error, ProviderError::AuthRequired), + "unreadable" => matches!(error, ProviderError::Other(_)), + _ => false, + }; + assert!(category_kept, "{failure}: {error:?}"); + assert_eq!(reads.get(), 3, "{failure}"); + assert_eq!(auth_bytes(dir.path()), before, "{failure}"); + } +} + +/// `load_credentials` waits exactly two retry delays before it reports a +/// missing file or unchanged stale credentials. +#[tokio::test(start_paused = true)] +async fn load_credentials_rereads_before_reporting_the_final_error() { + for failure in ["missing", "expired"] { + let dir = tempfile::tempdir().expect("codex home"); + write_state(dir.path(), failure); + let started = tokio::time::Instant::now(); + + let error = api_for(dir.path()) + .load_credentials() + .await + .err() + .expect("credential error"); + + let waited = started.elapsed(); + assert!( + waited >= 2 * CREDENTIAL_READ_RETRY_DELAY && waited < 3 * CREDENTIAL_READ_RETRY_DELAY, + "{failure}: waited {waited:?}" + ); + if failure == "missing" { + assert!(matches!(error, ProviderError::NotInstalled(_)), "{error:?}"); + } else { + assert!(matches!(error, ProviderError::AuthRequired), "{error:?}"); + } + } +} + +#[tokio::test(start_paused = true)] +async fn load_credentials_returns_usable_credentials_without_waiting() { + let dir = tempfile::tempdir().expect("codex home"); + let fresh = oauth_auth_json(FRESH_EXPIRY); + publish(dir.path(), &fresh); + let started = tokio::time::Instant::now(); + + let credentials = api_for(dir.path()) + .load_credentials() + .await + .expect("fresh credentials"); + + assert_eq!(started.elapsed(), Duration::ZERO); + assert_eq!(credentials.access_token, access_token_of(&fresh)); +} + +/// Upstream `cancelled OAuth fetch does not read credentials`: a load that is +/// never polled does not read, and dropping a load during its retry delay +/// stops further reads. +#[tokio::test(start_paused = true)] +async fn dropping_the_load_cancels_the_retry_delay() { + let dir = tempfile::tempdir().expect("codex home"); + let api = api_for(dir.path()); + let reads = Cell::new(0); + let read = || { + reads.set(reads.get() + 1); + api.load_credentials_once() + }; + + drop(CodexApi::reread_during_owner_publication(read)); + assert_eq!(reads.get(), 0, "an unpolled load must not read"); + + let outcome = tokio::time::timeout( + Duration::from_millis(20), + CodexApi::reread_during_owner_publication(read), + ) + .await; + assert!(outcome.is_err(), "the retry delay must be cancellable"); + assert_eq!(reads.get(), 1, "no read after the load was dropped"); +}