From 634764ece9a895672a58d3b836e4c502aa5ebd15 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Thu, 8 Jan 2026 11:43:56 -0800 Subject: [PATCH] Immutable CodexAuth (#8857) Historically we started with a CodexAuth that knew how to refresh it's own tokens and then added AuthManager that did a different kind of refresh (re-reading from disk). I don't think it makes sense for both `CodexAuth` and `AuthManager` to be mutable and contain behaviors. Move all refresh logic into `AuthManager` and keep `CodexAuth` as a data object. --- .../app-server/src/codex_message_processor.rs | 24 +- codex-rs/backend-client/src/client.rs | 4 +- codex-rs/chatgpt/src/chatgpt_token.rs | 9 +- codex-rs/cli/src/login.rs | 2 +- codex-rs/cloud-tasks/src/lib.rs | 9 +- codex-rs/cloud-tasks/src/util.rs | 4 +- codex-rs/core/src/api_bridge.rs | 4 +- codex-rs/core/src/auth.rs | 196 ++++++-------- codex-rs/core/src/client.rs | 21 +- codex-rs/core/src/codex.rs | 9 +- codex-rs/core/src/models_manager/manager.rs | 4 +- codex-rs/core/tests/suite/auth_refresh.rs | 254 ++++++++++++++---- codex-rs/exec/src/lib.rs | 2 +- codex-rs/tui/src/chatwidget.rs | 13 +- codex-rs/tui/src/lib.rs | 2 +- codex-rs/tui/src/status/helpers.rs | 2 +- codex-rs/tui2/src/chatwidget.rs | 13 +- codex-rs/tui2/src/lib.rs | 2 +- codex-rs/tui2/src/status/helpers.rs | 2 +- 19 files changed, 353 insertions(+), 223 deletions(-) diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index 177410f16..e45b1b788 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -590,7 +590,7 @@ impl CodexMessageProcessor { .await; let payload = AuthStatusChangeNotification { - auth_method: self.auth_manager.auth().map(|auth| auth.mode), + auth_method: self.auth_manager.auth_cached().map(|auth| auth.mode), }; self.outgoing .send_server_notification(ServerNotification::AuthStatusChange(payload)) @@ -620,7 +620,7 @@ impl CodexMessageProcessor { .await; let payload_v2 = AccountUpdatedNotification { - auth_mode: self.auth_manager.auth().map(|auth| auth.mode), + auth_mode: self.auth_manager.auth_cached().map(|auth| auth.mode), }; self.outgoing .send_server_notification(ServerNotification::AccountUpdated(payload_v2)) @@ -712,7 +712,7 @@ impl CodexMessageProcessor { auth_manager.reload(); // Notify clients with the actual current auth mode. - let current_auth_method = auth_manager.auth().map(|a| a.mode); + let current_auth_method = auth_manager.auth_cached().map(|a| a.mode); let payload = AuthStatusChangeNotification { auth_method: current_auth_method, }; @@ -802,7 +802,7 @@ impl CodexMessageProcessor { auth_manager.reload(); // Notify clients with the actual current auth mode. - let current_auth_method = auth_manager.auth().map(|a| a.mode); + let current_auth_method = auth_manager.auth_cached().map(|a| a.mode); let payload_v2 = AccountUpdatedNotification { auth_mode: current_auth_method, }; @@ -914,7 +914,7 @@ impl CodexMessageProcessor { } // Reflect the current auth method after logout (likely None). - Ok(self.auth_manager.auth().map(|auth| auth.mode)) + Ok(self.auth_manager.auth_cached().map(|auth| auth.mode)) } async fn logout_v1(&mut self, request_id: RequestId) { @@ -981,10 +981,10 @@ impl CodexMessageProcessor { requires_openai_auth: Some(false), } } else { - match self.auth_manager.auth() { + match self.auth_manager.auth().await { Some(auth) => { let auth_mode = auth.mode; - let (reported_auth_method, token_opt) = match auth.get_token().await { + let (reported_auth_method, token_opt) = match auth.get_token() { Ok(token) if !token.is_empty() => { let tok = if include_token { Some(token) } else { None }; (Some(auth_mode), tok) @@ -1029,7 +1029,7 @@ impl CodexMessageProcessor { return; } - let account = match self.auth_manager.auth() { + let account = match self.auth_manager.auth_cached() { Some(auth) => Some(match auth.mode { AuthMode::ApiKey => Account::ApiKey {}, AuthMode::ChatGPT => { @@ -1083,7 +1083,7 @@ impl CodexMessageProcessor { } async fn fetch_account_rate_limits(&self) -> Result { - let Some(auth) = self.auth_manager.auth() else { + let Some(auth) = self.auth_manager.auth().await else { return Err(JSONRPCErrorError { code: INVALID_REQUEST_ERROR_CODE, message: "codex account authentication required to read rate limits".to_string(), @@ -1100,7 +1100,6 @@ impl CodexMessageProcessor { } let client = BackendClient::from_auth(self.config.chatgpt_base_url.clone(), &auth) - .await .map_err(|err| JSONRPCErrorError { code: INTERNAL_ERROR_CODE, message: format!("failed to construct backend client: {err}"), @@ -1140,7 +1139,10 @@ impl CodexMessageProcessor { async fn get_user_info(&self, request_id: RequestId) { // Read alleged user email from cached auth (best-effort; not verified). - let alleged_user_email = self.auth_manager.auth().and_then(|a| a.get_account_email()); + let alleged_user_email = self + .auth_manager + .auth_cached() + .and_then(|a| a.get_account_email()); let response = UserInfoResponse { alleged_user_email }; self.outgoing.send_response(request_id, response).await; diff --git a/codex-rs/backend-client/src/client.rs b/codex-rs/backend-client/src/client.rs index 4b5eaa410..fdd4504bb 100644 --- a/codex-rs/backend-client/src/client.rs +++ b/codex-rs/backend-client/src/client.rs @@ -73,8 +73,8 @@ impl Client { }) } - pub async fn from_auth(base_url: impl Into, auth: &CodexAuth) -> Result { - let token = auth.get_token().await.map_err(anyhow::Error::from)?; + pub fn from_auth(base_url: impl Into, auth: &CodexAuth) -> Result { + let token = auth.get_token().map_err(anyhow::Error::from)?; let mut client = Self::new(base_url)? .with_user_agent(get_codex_user_agent()) .with_bearer_token(token); diff --git a/codex-rs/chatgpt/src/chatgpt_token.rs b/codex-rs/chatgpt/src/chatgpt_token.rs index e8879ad21..23f088373 100644 --- a/codex-rs/chatgpt/src/chatgpt_token.rs +++ b/codex-rs/chatgpt/src/chatgpt_token.rs @@ -1,4 +1,4 @@ -use codex_core::CodexAuth; +use codex_core::AuthManager; use std::path::Path; use std::sync::LazyLock; use std::sync::RwLock; @@ -23,9 +23,10 @@ pub async fn init_chatgpt_token_from_auth( codex_home: &Path, auth_credentials_store_mode: AuthCredentialsStoreMode, ) -> std::io::Result<()> { - let auth = CodexAuth::from_auth_storage(codex_home, auth_credentials_store_mode)?; - if let Some(auth) = auth { - let token_data = auth.get_token_data().await?; + let auth_manager = + AuthManager::new(codex_home.to_path_buf(), false, auth_credentials_store_mode); + if let Some(auth) = auth_manager.auth().await { + let token_data = auth.get_token_data()?; set_chatgpt_token_data(token_data); } Ok(()) diff --git a/codex-rs/cli/src/login.rs b/codex-rs/cli/src/login.rs index 8fbf7b04b..3792e2f53 100644 --- a/codex-rs/cli/src/login.rs +++ b/codex-rs/cli/src/login.rs @@ -155,7 +155,7 @@ pub async fn run_login_status(cli_config_overrides: CliConfigOverrides) -> ! { match CodexAuth::from_auth_storage(&config.codex_home, config.cli_auth_credentials_store_mode) { Ok(Some(auth)) => match auth.mode { - AuthMode::ApiKey => match auth.get_token().await { + AuthMode::ApiKey => match auth.get_token() { Ok(api_key) => { eprintln!("Logged in using an API key - {}", safe_format_key(&api_key)); std::process::exit(0); diff --git a/codex-rs/cloud-tasks/src/lib.rs b/codex-rs/cloud-tasks/src/lib.rs index 105f6cfb2..e1bedbc1c 100644 --- a/codex-rs/cloud-tasks/src/lib.rs +++ b/codex-rs/cloud-tasks/src/lib.rs @@ -10,7 +10,6 @@ pub use cli::Cli; use anyhow::anyhow; use chrono::Utc; use codex_cloud_tasks_client::TaskStatus; -use codex_login::AuthManager; use owo_colors::OwoColorize; use owo_colors::Stream; use std::cmp::Ordering; @@ -65,7 +64,11 @@ async fn init_backend(user_agent_suffix: &str) -> anyhow::Result append_error_log(format!("startup: base_url={base_url} path_style={style}")); let auth_manager = util::load_auth_manager().await; - let auth = match auth_manager.as_ref().and_then(AuthManager::auth) { + let auth = match auth_manager.as_ref() { + Some(manager) => manager.auth().await, + None => None, + }; + let auth = match auth { Some(auth) => auth, None => { eprintln!( @@ -79,7 +82,7 @@ async fn init_backend(user_agent_suffix: &str) -> anyhow::Result append_error_log(format!("auth: mode=ChatGPT account_id={acc}")); } - let token = match auth.get_token().await { + let token = match auth.get_token() { Ok(t) if !t.is_empty() => t, _ => { eprintln!( diff --git a/codex-rs/cloud-tasks/src/util.rs b/codex-rs/cloud-tasks/src/util.rs index 9c4ae01cd..cf9236a5b 100644 --- a/codex-rs/cloud-tasks/src/util.rs +++ b/codex-rs/cloud-tasks/src/util.rs @@ -85,8 +85,8 @@ pub async fn build_chatgpt_headers() -> HeaderMap { HeaderValue::from_str(&ua).unwrap_or(HeaderValue::from_static("codex-cli")), ); if let Some(am) = load_auth_manager().await - && let Some(auth) = am.auth() - && let Ok(tok) = auth.get_token().await + && let Some(auth) = am.auth().await + && let Ok(tok) = auth.get_token() && !tok.is_empty() { let v = format!("Bearer {tok}"); diff --git a/codex-rs/core/src/api_bridge.rs b/codex-rs/core/src/api_bridge.rs index 79fd67d65..a19ff5abb 100644 --- a/codex-rs/core/src/api_bridge.rs +++ b/codex-rs/core/src/api_bridge.rs @@ -100,7 +100,7 @@ fn extract_request_id(headers: Option<&HeaderMap>) -> Option { }) } -pub(crate) async fn auth_provider_from_auth( +pub(crate) fn auth_provider_from_auth( auth: Option, provider: &ModelProviderInfo, ) -> crate::error::Result { @@ -119,7 +119,7 @@ pub(crate) async fn auth_provider_from_auth( } if let Some(auth) = auth { - let token = auth.get_token().await?; + let token = auth.get_token()?; Ok(CoreAuthProvider { token: Some(token), account_id: auth.get_account_id(), diff --git a/codex-rs/core/src/auth.rs b/codex-rs/core/src/auth.rs index 71a542912..575bc5119 100644 --- a/codex-rs/core/src/auth.rs +++ b/codex-rs/core/src/auth.rs @@ -8,12 +8,10 @@ use serde::Serialize; use serial_test::serial; use std::env; use std::fmt::Debug; -use std::io::ErrorKind; use std::path::Path; use std::path::PathBuf; use std::sync::Arc; use std::sync::Mutex; -use std::time::Duration; use codex_app_server_protocol::AuthMode; use codex_protocol::config_types::ForcedLoginMethod; @@ -93,40 +91,6 @@ impl From for std::io::Error { } impl CodexAuth { - pub async fn refresh_token(&self) -> Result { - tracing::info!("Refreshing token"); - - let token_data = self.get_current_token_data().ok_or_else(|| { - RefreshTokenError::Transient(std::io::Error::other("Token data is not available.")) - })?; - let token = token_data.refresh_token; - - let refresh_response = try_refresh_token(token, &self.client).await?; - - let updated = update_tokens( - &self.storage, - refresh_response.id_token, - refresh_response.access_token, - refresh_response.refresh_token, - ) - .await - .map_err(RefreshTokenError::from)?; - - if let Ok(mut auth_lock) = self.auth_dot_json.lock() { - *auth_lock = Some(updated.clone()); - } - - let access = match updated.tokens { - Some(t) => t.access_token, - None => { - return Err(RefreshTokenError::other_with_message( - "Token data is not available after refresh.", - )); - } - }; - Ok(access) - } - /// Loads the available auth information from auth storage. pub fn from_auth_storage( codex_home: &Path, @@ -135,62 +99,23 @@ impl CodexAuth { load_auth(codex_home, false, auth_credentials_store_mode) } - pub async fn get_token_data(&self) -> Result { + pub fn get_token_data(&self) -> Result { let auth_dot_json: Option = self.get_current_auth_json(); match auth_dot_json { Some(AuthDotJson { - tokens: Some(mut tokens), - last_refresh: Some(last_refresh), + tokens: Some(tokens), + last_refresh: Some(_), .. - }) => { - if last_refresh < Utc::now() - chrono::Duration::days(TOKEN_REFRESH_INTERVAL) { - let refresh_result = tokio::time::timeout( - Duration::from_secs(60), - try_refresh_token(tokens.refresh_token.clone(), &self.client), - ) - .await; - let refresh_response = match refresh_result { - Ok(Ok(response)) => response, - Ok(Err(err)) => return Err(err.into()), - Err(_) => { - return Err(std::io::Error::new( - ErrorKind::TimedOut, - "timed out while refreshing OpenAI API key", - )); - } - }; - - let updated_auth_dot_json = update_tokens( - &self.storage, - refresh_response.id_token, - refresh_response.access_token, - refresh_response.refresh_token, - ) - .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_lock = self.auth_dot_json.lock().unwrap(); - *auth_lock = Some(updated_auth_dot_json); - } - - Ok(tokens) - } + }) => Ok(tokens), _ => Err(std::io::Error::other("Token data is not available.")), } } - pub async fn get_token(&self) -> Result { + pub fn get_token(&self) -> Result { match self.mode { AuthMode::ApiKey => Ok(self.api_key.clone().unwrap_or_default()), AuthMode::ChatGPT => { - let id_token = self.get_token_data().await?.access_token; + let id_token = self.get_token_data()?.access_token; Ok(id_token) } } @@ -338,7 +263,7 @@ pub fn load_auth_dot_json( storage.load() } -pub async fn enforce_login_restrictions(config: &Config) -> std::io::Result<()> { +pub fn enforce_login_restrictions(config: &Config) -> std::io::Result<()> { let Some(auth) = load_auth( &config.codex_home, true, @@ -376,7 +301,7 @@ pub async fn enforce_login_restrictions(config: &Config) -> std::io::Result<()> return Ok(()); } - let token_data = match auth.get_token_data().await { + let token_data = match auth.get_token_data() { Ok(data) => data, Err(err) => { return logout_with_message( @@ -689,11 +614,22 @@ impl AuthManager { }) } - /// Current cached auth (clone). May be `None` if not logged in or load failed. - pub fn auth(&self) -> Option { + /// Current cached auth (clone) without attempting a refresh. + pub fn auth_cached(&self) -> Option { self.inner.read().ok().and_then(|c| c.auth.clone()) } + /// Current cached auth (clone). May be `None` if not logged in or load failed. + /// Refreshes cached ChatGPT tokens if they are stale before returning. + pub async fn auth(&self) -> Option { + let auth = self.auth_cached()?; + if let Err(err) = self.refresh_if_stale(&auth).await { + tracing::error!("Failed to refresh token: {}", err); + return Some(auth); + } + self.auth_cached() + } + /// Force a reload of the auth information from auth.json. Returns /// whether the auth value changed. pub fn reload(&self) -> bool { @@ -736,24 +672,20 @@ impl AuthManager { /// Attempt to refresh the current auth token (if any). On success, reload /// the auth state from disk so other components observe refreshed token. - /// If the token refresh fails in a permanent (non‑transient) way, logs out - /// to clear invalid auth state. + /// If the token refresh fails, returns the error to the caller. pub async fn refresh_token(&self) -> Result, RefreshTokenError> { - let auth = match self.auth() { - Some(a) => a, + let auth = match self.auth_cached() { + Some(auth) => auth, None => return Ok(None), }; - match auth.refresh_token().await { - Ok(token) => { - // Reload to pick up persisted changes. - self.reload(); - Ok(Some(token)) - } - Err(e) => { - tracing::error!("Failed to refresh token: {}", e); - Err(e) - } - } + tracing::info!("Refreshing token"); + let token_data = auth.get_current_token_data().ok_or_else(|| { + RefreshTokenError::Transient(std::io::Error::other("Token data is not available.")) + })?; + let access = self.refresh_tokens(&auth, token_data.refresh_token).await?; + // Reload to pick up persisted changes. + self.reload(); + Ok(Some(access)) } /// Log out by deleting the on‑disk auth.json (if present). Returns Ok(true) @@ -768,7 +700,56 @@ impl AuthManager { } pub fn get_auth_mode(&self) -> Option { - self.auth().map(|a| a.mode) + self.auth_cached().map(|a| a.mode) + } + + async fn refresh_if_stale(&self, auth: &CodexAuth) -> Result { + if auth.mode != AuthMode::ChatGPT { + return Ok(false); + } + + let auth_dot_json = match auth.get_current_auth_json() { + Some(auth_dot_json) => auth_dot_json, + None => return Ok(false), + }; + let tokens = match auth_dot_json.tokens { + Some(tokens) => tokens, + None => return Ok(false), + }; + let last_refresh = match auth_dot_json.last_refresh { + Some(last_refresh) => last_refresh, + None => return Ok(false), + }; + if last_refresh >= Utc::now() - chrono::Duration::days(TOKEN_REFRESH_INTERVAL) { + return Ok(false); + } + self.refresh_tokens(auth, tokens.refresh_token).await?; + self.reload(); + Ok(true) + } + + async fn refresh_tokens( + &self, + auth: &CodexAuth, + refresh_token: String, + ) -> Result { + let refresh_response = try_refresh_token(refresh_token, &auth.client).await?; + + let updated = update_tokens( + &auth.storage, + refresh_response.id_token, + refresh_response.access_token, + refresh_response.refresh_token, + ) + .await + .map_err(RefreshTokenError::from)?; + + match updated.tokens { + Some(tokens) => Ok(tokens.access_token), + None => Err(RefreshTokenError::other_with_message( + "Token data is not available after refresh.", + )), + } } } @@ -930,7 +911,7 @@ mod tests { 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()); + assert!(auth.get_token_data().is_err()); } #[test] @@ -1058,7 +1039,6 @@ mod tests { let config = build_config(codex_home.path(), Some(ForcedLoginMethod::Chatgpt), None).await; let err = super::enforce_login_restrictions(&config) - .await .expect_err("expected method mismatch to error"); assert!(err.to_string().contains("ChatGPT login is required")); assert!( @@ -1084,7 +1064,6 @@ mod tests { let config = build_config(codex_home.path(), None, Some("org_mine".to_string())).await; let err = super::enforce_login_restrictions(&config) - .await .expect_err("expected workspace mismatch to error"); assert!(err.to_string().contains("workspace org_mine")); assert!( @@ -1109,9 +1088,7 @@ mod tests { let config = build_config(codex_home.path(), None, Some("org_mine".to_string())).await; - super::enforce_login_restrictions(&config) - .await - .expect("matching workspace should succeed"); + super::enforce_login_restrictions(&config).expect("matching workspace should succeed"); assert!( codex_home.path().join("auth.json").exists(), "auth.json should remain when restrictions pass" @@ -1127,9 +1104,7 @@ mod tests { let config = build_config(codex_home.path(), None, Some("org_mine".to_string())).await; - super::enforce_login_restrictions(&config) - .await - .expect("matching workspace should succeed"); + super::enforce_login_restrictions(&config).expect("matching workspace should succeed"); assert!( codex_home.path().join("auth.json").exists(), "auth.json should remain when restrictions pass" @@ -1145,7 +1120,6 @@ mod tests { let config = build_config(codex_home.path(), Some(ForcedLoginMethod::Chatgpt), None).await; let err = super::enforce_login_restrictions(&config) - .await .expect_err("environment API key should not satisfy forced ChatGPT login"); assert!( err.to_string() diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index f1400d4e3..bc4875ddb 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -157,11 +157,14 @@ impl ModelClient { let mut refreshed = false; loop { - let auth = auth_manager.as_ref().and_then(|m| m.auth()); + let auth = match auth_manager.as_ref() { + Some(manager) => manager.auth().await, + None => None, + }; let api_provider = self .provider .to_api_provider(auth.as_ref().map(|a| a.mode))?; - let api_auth = auth_provider_from_auth(auth.clone(), &self.provider).await?; + let api_auth = auth_provider_from_auth(auth.clone(), &self.provider)?; let transport = ReqwestTransport::new(build_reqwest_client()); let (request_telemetry, sse_telemetry) = self.build_streaming_telemetry(); let client = ApiChatClient::new(transport, api_provider, api_auth) @@ -245,11 +248,14 @@ impl ModelClient { let mut refreshed = false; loop { - let auth = auth_manager.as_ref().and_then(|m| m.auth()); + let auth = match auth_manager.as_ref() { + Some(manager) => manager.auth().await, + None => None, + }; let api_provider = self .provider .to_api_provider(auth.as_ref().map(|a| a.mode))?; - let api_auth = auth_provider_from_auth(auth.clone(), &self.provider).await?; + let api_auth = auth_provider_from_auth(auth.clone(), &self.provider)?; let transport = ReqwestTransport::new(build_reqwest_client()); let (request_telemetry, sse_telemetry) = self.build_streaming_telemetry(); let compression = if self @@ -344,11 +350,14 @@ impl ModelClient { return Ok(Vec::new()); } let auth_manager = self.auth_manager.clone(); - let auth = auth_manager.as_ref().and_then(|m| m.auth()); + let auth = match auth_manager.as_ref() { + Some(manager) => manager.auth().await, + None => None, + }; let api_provider = self .provider .to_api_provider(auth.as_ref().map(|a| a.mode))?; - let api_auth = auth_provider_from_auth(auth.clone(), &self.provider).await?; + let api_auth = auth_provider_from_auth(auth.clone(), &self.provider)?; let transport = ReqwestTransport::new(build_reqwest_client()); let request_telemetry = self.build_request_telemetry(); let client = ApiCompactClient::new(transport, api_provider, api_auth) diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index b4fdc7d68..96e90c5cf 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -7,6 +7,7 @@ use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; use crate::AuthManager; +use crate::CodexAuth; use crate::SandboxState; use crate::agent::AgentControl; use crate::agent::AgentStatus; @@ -633,13 +634,15 @@ impl Session { } maybe_push_chat_wire_api_deprecation(&config, &mut post_session_configured_events); + let auth = auth_manager.auth().await; + let auth = auth.as_ref(); let otel_manager = OtelManager::new( conversation_id, session_configuration.model.as_str(), session_configuration.model.as_str(), - auth_manager.auth().and_then(|a| a.get_account_id()), - auth_manager.auth().and_then(|a| a.get_account_email()), - auth_manager.auth().map(|a| a.mode), + auth.and_then(CodexAuth::get_account_id), + auth.and_then(CodexAuth::get_account_email), + auth.map(|a| a.mode), config.otel.log_user_prompt, terminal::user_agent(), session_configuration.session_source.clone(), diff --git a/codex-rs/core/src/models_manager/manager.rs b/codex-rs/core/src/models_manager/manager.rs index 87ff1b76d..0b388d133 100644 --- a/codex-rs/core/src/models_manager/manager.rs +++ b/codex-rs/core/src/models_manager/manager.rs @@ -98,9 +98,9 @@ impl ModelsManager { if !remote_models_feature || self.auth_manager.get_auth_mode() == Some(AuthMode::ApiKey) { return Ok(()); } - let auth = self.auth_manager.auth(); + let auth = self.auth_manager.auth().await; let api_provider = self.provider.to_api_provider(Some(AuthMode::ChatGPT))?; - let api_auth = auth_provider_from_auth(auth.clone(), &self.provider).await?; + let api_auth = auth_provider_from_auth(auth.clone(), &self.provider)?; let transport = ReqwestTransport::new(build_reqwest_client()); let client = ModelsClient::new(transport, api_provider, api_auth); diff --git a/codex-rs/core/tests/suite/auth_refresh.rs b/codex-rs/core/tests/suite/auth_refresh.rs index 6daaf70b5..d63618935 100644 --- a/codex-rs/core/tests/suite/auth_refresh.rs +++ b/codex-rs/core/tests/suite/auth_refresh.rs @@ -3,7 +3,7 @@ use anyhow::Result; use base64::Engine; use chrono::Duration; use chrono::Utc; -use codex_core::CodexAuth; +use codex_core::AuthManager; use codex_core::auth::AuthCredentialsStoreMode; use codex_core::auth::AuthDotJson; use codex_core::auth::REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR; @@ -45,32 +45,149 @@ async fn refresh_token_succeeds_updates_storage() -> Result<()> { .await; let ctx = RefreshTokenTestContext::new(&server)?; - let auth = ctx.auth.clone(); + let initial_last_refresh = Utc::now() - Duration::days(1); + let initial_tokens = build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN); + let initial_auth = AuthDotJson { + openai_api_key: None, + tokens: Some(initial_tokens.clone()), + last_refresh: Some(initial_last_refresh), + }; + ctx.write_auth(&initial_auth)?; - let access = auth + let access = ctx + .auth_manager .refresh_token() .await .context("refresh should succeed")?; - assert_eq!(access, "new-access-token"); + assert_eq!(access, Some("new-access-token".to_string())); + let refreshed_tokens = TokenData { + access_token: "new-access-token".to_string(), + refresh_token: "new-refresh-token".to_string(), + ..initial_tokens.clone() + }; let stored = ctx.load_auth()?; let tokens = stored.tokens.as_ref().context("tokens should exist")?; - assert_eq!(tokens.access_token, "new-access-token"); - assert_eq!(tokens.refresh_token, "new-refresh-token"); + assert_eq!(tokens, &refreshed_tokens); let refreshed_at = stored .last_refresh .as_ref() .context("last_refresh should be recorded")?; assert!( - *refreshed_at >= ctx.initial_last_refresh, + *refreshed_at >= initial_last_refresh, "last_refresh should advance" ); - let cached = auth - .get_token_data() + let cached_auth = ctx + .auth_manager + .auth() .await + .context("auth should be cached")?; + let cached = cached_auth + .get_token_data() .context("token data should be cached")?; - assert_eq!(cached.access_token, "new-access-token"); + assert_eq!(cached, refreshed_tokens); + + server.verify().await; + Ok(()) +} + +#[serial_test::serial(auth_refresh)] +#[tokio::test] +async fn returns_fresh_tokens_as_is() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": "new-access-token", + "refresh_token": "new-refresh-token" + }))) + .mount(&server) + .await; + + let ctx = RefreshTokenTestContext::new(&server)?; + let initial_last_refresh = Utc::now() - Duration::days(1); + let initial_tokens = build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN); + let initial_auth = AuthDotJson { + openai_api_key: None, + tokens: Some(initial_tokens.clone()), + last_refresh: Some(initial_last_refresh), + }; + ctx.write_auth(&initial_auth)?; + + let cached_auth = ctx + .auth_manager + .auth() + .await + .context("auth should be cached")?; + let cached = cached_auth + .get_token_data() + .context("token data should remain cached")?; + assert_eq!(cached, initial_tokens); + + let stored = ctx.load_auth()?; + assert_eq!(stored, initial_auth); + + let requests = server.received_requests().await.unwrap_or_default(); + assert!(requests.is_empty(), "expected no refresh token requests"); + + Ok(()) +} + +#[serial_test::serial(auth_refresh)] +#[tokio::test] +async fn refreshes_token_when_last_refresh_is_stale() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": "new-access-token", + "refresh_token": "new-refresh-token" + }))) + .expect(1) + .mount(&server) + .await; + + let ctx = RefreshTokenTestContext::new(&server)?; + let stale_refresh = Utc::now() - Duration::days(9); + let initial_tokens = build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN); + let initial_auth = AuthDotJson { + openai_api_key: None, + tokens: Some(initial_tokens.clone()), + last_refresh: Some(stale_refresh), + }; + ctx.write_auth(&initial_auth)?; + + let cached_auth = ctx + .auth_manager + .auth() + .await + .context("auth should be cached")?; + let refreshed_tokens = TokenData { + access_token: "new-access-token".to_string(), + refresh_token: "new-refresh-token".to_string(), + ..initial_tokens.clone() + }; + let cached = cached_auth + .get_token_data() + .context("token data should refresh")?; + assert_eq!(cached, refreshed_tokens); + + let stored = ctx.load_auth()?; + let tokens = stored.tokens.as_ref().context("tokens should exist")?; + assert_eq!(tokens, &refreshed_tokens); + let refreshed_at = stored + .last_refresh + .as_ref() + .context("last_refresh should be recorded")?; + assert!( + *refreshed_at >= stale_refresh, + "last_refresh should advance" + ); server.verify().await; Ok(()) @@ -94,9 +211,17 @@ async fn refresh_token_returns_permanent_error_for_expired_refresh_token() -> Re .await; let ctx = RefreshTokenTestContext::new(&server)?; - let auth = ctx.auth.clone(); + let initial_last_refresh = Utc::now() - Duration::days(1); + let initial_tokens = build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN); + let initial_auth = AuthDotJson { + openai_api_key: None, + tokens: Some(initial_tokens.clone()), + last_refresh: Some(initial_last_refresh), + }; + ctx.write_auth(&initial_auth)?; - let err = auth + let err = ctx + .auth_manager .refresh_token() .await .err() @@ -104,16 +229,16 @@ async fn refresh_token_returns_permanent_error_for_expired_refresh_token() -> Re assert_eq!(err.failed_reason(), Some(RefreshTokenFailedReason::Expired)); let stored = ctx.load_auth()?; - let tokens = stored.tokens.as_ref().context("tokens should remain")?; - assert_eq!(tokens.access_token, INITIAL_ACCESS_TOKEN); - assert_eq!(tokens.refresh_token, INITIAL_REFRESH_TOKEN); - assert_eq!( - *stored - .last_refresh - .as_ref() - .context("last_refresh should remain unchanged")?, - ctx.initial_last_refresh, - ); + assert_eq!(stored, initial_auth); + let cached_auth = ctx + .auth_manager + .auth() + .await + .context("auth should remain cached")?; + let cached = cached_auth + .get_token_data() + .context("token data should remain cached")?; + assert_eq!(cached, initial_tokens); server.verify().await; Ok(()) @@ -135,9 +260,17 @@ async fn refresh_token_returns_transient_error_on_server_failure() -> Result<()> .await; let ctx = RefreshTokenTestContext::new(&server)?; - let auth = ctx.auth.clone(); + let initial_last_refresh = Utc::now() - Duration::days(1); + let initial_tokens = build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN); + let initial_auth = AuthDotJson { + openai_api_key: None, + tokens: Some(initial_tokens.clone()), + last_refresh: Some(initial_last_refresh), + }; + ctx.write_auth(&initial_auth)?; - let err = auth + let err = ctx + .auth_manager .refresh_token() .await .err() @@ -146,16 +279,16 @@ async fn refresh_token_returns_transient_error_on_server_failure() -> Result<()> assert_eq!(err.failed_reason(), None); let stored = ctx.load_auth()?; - let tokens = stored.tokens.as_ref().context("tokens should remain")?; - assert_eq!(tokens.access_token, INITIAL_ACCESS_TOKEN); - assert_eq!(tokens.refresh_token, INITIAL_REFRESH_TOKEN); - assert_eq!( - *stored - .last_refresh - .as_ref() - .context("last_refresh should remain unchanged")?, - ctx.initial_last_refresh, - ); + assert_eq!(stored, initial_auth); + let cached_auth = ctx + .auth_manager + .auth() + .await + .context("auth should remain cached")?; + let cached = cached_auth + .get_token_data() + .context("token data should remain cached")?; + assert_eq!(cached, initial_tokens); server.verify().await; Ok(()) @@ -163,44 +296,26 @@ async fn refresh_token_returns_transient_error_on_server_failure() -> Result<()> struct RefreshTokenTestContext { codex_home: TempDir, - auth: CodexAuth, - initial_last_refresh: chrono::DateTime, + auth_manager: AuthManager, _env_guard: EnvGuard, } impl RefreshTokenTestContext { fn new(server: &MockServer) -> Result { let codex_home = TempDir::new()?; - let initial_last_refresh = Utc::now() - Duration::days(1); - let mut id_token = IdTokenInfo::default(); - id_token.raw_jwt = minimal_jwt(); - let tokens = TokenData { - id_token, - access_token: INITIAL_ACCESS_TOKEN.to_string(), - refresh_token: INITIAL_REFRESH_TOKEN.to_string(), - account_id: Some("account-id".to_string()), - }; - let auth_dot_json = AuthDotJson { - openai_api_key: None, - tokens: Some(tokens), - last_refresh: Some(initial_last_refresh), - }; - save_auth( - codex_home.path(), - &auth_dot_json, - AuthCredentialsStoreMode::File, - )?; let endpoint = format!("{}/oauth/token", server.uri()); let env_guard = EnvGuard::set(REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR, endpoint); - let auth = CodexAuth::from_auth_storage(codex_home.path(), AuthCredentialsStoreMode::File)? - .context("auth should load from storage")?; + let auth_manager = AuthManager::new( + codex_home.path().to_path_buf(), + false, + AuthCredentialsStoreMode::File, + ); Ok(Self { codex_home, - auth, - initial_last_refresh, + auth_manager, _env_guard: env_guard, }) } @@ -210,6 +325,16 @@ impl RefreshTokenTestContext { .context("load auth.json")? .context("auth.json should exist") } + + fn write_auth(&self, auth_dot_json: &AuthDotJson) -> Result<()> { + save_auth( + self.codex_home.path(), + auth_dot_json, + AuthCredentialsStoreMode::File, + )?; + self.auth_manager.reload(); + Ok(()) + } } struct EnvGuard { @@ -270,3 +395,14 @@ fn minimal_jwt() -> String { let signature_b64 = b64(b"sig"); format!("{header_b64}.{payload_b64}.{signature_b64}") } + +fn build_tokens(access_token: &str, refresh_token: &str) -> TokenData { + let mut id_token = IdTokenInfo::default(); + id_token.raw_jwt = minimal_jwt(); + TokenData { + id_token, + access_token: access_token.to_string(), + refresh_token: refresh_token.to_string(), + account_id: Some("account-id".to_string()), + } +} diff --git a/codex-rs/exec/src/lib.rs b/codex-rs/exec/src/lib.rs index 89be1ac6c..44ce1d6e2 100644 --- a/codex-rs/exec/src/lib.rs +++ b/codex-rs/exec/src/lib.rs @@ -218,7 +218,7 @@ pub async fn run_main(cli: Cli, codex_linux_sandbox_exe: Option) -> any let config = Config::load_with_cli_overrides_and_harness_overrides(cli_kv_overrides, overrides).await?; - if let Err(err) = enforce_login_restrictions(&config).await { + if let Err(err) = enforce_login_restrictions(&config) { eprintln!("{err}"); std::process::exit(1); } diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index 537adda25..fa0be84a7 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -2310,21 +2310,22 @@ impl ChatWidget { fn prefetch_rate_limits(&mut self) { self.stop_rate_limit_poller(); - let Some(auth) = self.auth_manager.auth() else { - return; - }; - if auth.mode != AuthMode::ChatGPT { + if self.auth_manager.auth_cached().map(|auth| auth.mode) != Some(AuthMode::ChatGPT) { return; } let base_url = self.config.chatgpt_base_url.clone(); let app_event_tx = self.app_event_tx.clone(); + let auth_manager = Arc::clone(&self.auth_manager); let handle = tokio::spawn(async move { let mut interval = tokio::time::interval(Duration::from_secs(60)); loop { - if let Some(snapshot) = fetch_rate_limits(base_url.clone(), auth.clone()).await { + if let Some(auth) = auth_manager.auth().await + && auth.mode == AuthMode::ChatGPT + && let Some(snapshot) = fetch_rate_limits(base_url.clone(), auth).await + { app_event_tx.send(AppEvent::RateLimitSnapshotFetched(snapshot)); } interval.tick().await; @@ -3749,7 +3750,7 @@ fn extract_first_bold(s: &str) -> Option { } async fn fetch_rate_limits(base_url: String, auth: CodexAuth) -> Option { - match BackendClient::from_auth(base_url, &auth).await { + match BackendClient::from_auth(base_url, &auth) { Ok(client) => match client.get_rate_limits().await { Ok(snapshot) => Some(snapshot), Err(err) => { diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index 6f4faaad6..f4e5e771f 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -235,7 +235,7 @@ pub async fn run_main( } #[allow(clippy::print_stderr)] - if let Err(err) = enforce_login_restrictions(&config).await { + if let Err(err) = enforce_login_restrictions(&config) { eprintln!("{err}"); std::process::exit(1); } diff --git a/codex-rs/tui/src/status/helpers.rs b/codex-rs/tui/src/status/helpers.rs index 8ba7ec377..78a956b27 100644 --- a/codex-rs/tui/src/status/helpers.rs +++ b/codex-rs/tui/src/status/helpers.rs @@ -88,7 +88,7 @@ pub(crate) fn compose_account_display( auth_manager: &AuthManager, plan: Option, ) -> Option { - let auth = auth_manager.auth()?; + let auth = auth_manager.auth_cached()?; match auth.mode { AuthMode::ChatGPT => { diff --git a/codex-rs/tui2/src/chatwidget.rs b/codex-rs/tui2/src/chatwidget.rs index bd6506630..4ee3aa5f1 100644 --- a/codex-rs/tui2/src/chatwidget.rs +++ b/codex-rs/tui2/src/chatwidget.rs @@ -2105,21 +2105,22 @@ impl ChatWidget { fn prefetch_rate_limits(&mut self) { self.stop_rate_limit_poller(); - let Some(auth) = self.auth_manager.auth() else { - return; - }; - if auth.mode != AuthMode::ChatGPT { + if self.auth_manager.auth_cached().map(|auth| auth.mode) != Some(AuthMode::ChatGPT) { return; } let base_url = self.config.chatgpt_base_url.clone(); let app_event_tx = self.app_event_tx.clone(); + let auth_manager = Arc::clone(&self.auth_manager); let handle = tokio::spawn(async move { let mut interval = tokio::time::interval(Duration::from_secs(60)); loop { - if let Some(snapshot) = fetch_rate_limits(base_url.clone(), auth.clone()).await { + if let Some(auth) = auth_manager.auth().await + && auth.mode == AuthMode::ChatGPT + && let Some(snapshot) = fetch_rate_limits(base_url.clone(), auth).await + { app_event_tx.send(AppEvent::RateLimitSnapshotFetched(snapshot)); } interval.tick().await; @@ -3502,7 +3503,7 @@ fn extract_first_bold(s: &str) -> Option { } async fn fetch_rate_limits(base_url: String, auth: CodexAuth) -> Option { - match BackendClient::from_auth(base_url, &auth).await { + match BackendClient::from_auth(base_url, &auth) { Ok(client) => match client.get_rate_limits().await { Ok(snapshot) => Some(snapshot), Err(err) => { diff --git a/codex-rs/tui2/src/lib.rs b/codex-rs/tui2/src/lib.rs index 1c161bf62..ac062cf66 100644 --- a/codex-rs/tui2/src/lib.rs +++ b/codex-rs/tui2/src/lib.rs @@ -250,7 +250,7 @@ pub async fn run_main( } #[allow(clippy::print_stderr)] - if let Err(err) = enforce_login_restrictions(&config).await { + if let Err(err) = enforce_login_restrictions(&config) { eprintln!("{err}"); std::process::exit(1); } diff --git a/codex-rs/tui2/src/status/helpers.rs b/codex-rs/tui2/src/status/helpers.rs index 8ba7ec377..78a956b27 100644 --- a/codex-rs/tui2/src/status/helpers.rs +++ b/codex-rs/tui2/src/status/helpers.rs @@ -88,7 +88,7 @@ pub(crate) fn compose_account_display( auth_manager: &AuthManager, plan: Option, ) -> Option { - let auth = auth_manager.auth()?; + let auth = auth_manager.auth_cached()?; match auth.mode { AuthMode::ChatGPT => {