agent/codex
2
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

Immutable CodexAuth (#8857)

Browse Source 
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.
This commit is contained in:
pakrym-oai
2026-01-08 11:43:56 -08:00
committed by GitHub
Unverified
parent 5bc3e325a6
commit 634764ece9
19 changed files with 353 additions and 223 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server/src/codex_message_processor.rs
+13 -11
View File
@@ -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<CoreRateLimitSnapshot, JSONRPCErrorError> {
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;
codex-rs/backend-client/src/client.rs
+2 -2
View File
@@ -73,8 +73,8 @@ impl Client {
})
}
pub async fn from_auth(base_url: impl Into<String>, auth: &CodexAuth) -> Result<Self> {
let token = auth.get_token().await.map_err(anyhow::Error::from)?;
pub fn from_auth(base_url: impl Into<String>, auth: &CodexAuth) -> Result<Self> {
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);
codex-rs/chatgpt/src/chatgpt_token.rs
+5 -4
View File
@@ -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(())
codex-rs/cli/src/login.rs
+1 -1
View File
@@ -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);
codex-rs/cloud-tasks/src/lib.rs
+6 -3
View File
@@ -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<BackendContext>
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<BackendContext>
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!(
codex-rs/cloud-tasks/src/util.rs
+2 -2
View File
@@ -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}");
codex-rs/core/src/api_bridge.rs
+2 -2
View File
@@ -100,7 +100,7 @@ fn extract_request_id(headers: Option<&HeaderMap>) -> Option<String> {
})
}
pub(crate) async fn auth_provider_from_auth(
pub(crate) fn auth_provider_from_auth(
auth: Option<CodexAuth>,
provider: &ModelProviderInfo,
) -> crate::error::Result<CoreAuthProvider> {
@@ -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(),
codex-rs/core/src/auth.rs
+85 -111
View File
@@ -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<RefreshTokenError> for std::io::Error {
}
impl CodexAuth {
pub async fn refresh_token(&self) -> Result<String, RefreshTokenError> {
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<TokenData, std::io::Error> {
pub fn get_token_data(&self) -> Result<TokenData, std::io::Error> {
let auth_dot_json: Option<AuthDotJson> = 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<String, std::io::Error> {
pub fn get_token(&self) -> Result<String, std::io::Error> {
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<CodexAuth> {
/// Current cached auth (clone) without attempting a refresh.
pub fn auth_cached(&self) -> Option<CodexAuth> {
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<CodexAuth> {
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<Option<String>, 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<AuthMode> {
self.auth().map(|a| a.mode)
self.auth_cached().map(|a| a.mode)
}
async fn refresh_if_stale(&self, auth: &CodexAuth) -> Result<bool, RefreshTokenError> {
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<String, RefreshTokenError> {
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()
codex-rs/core/src/client.rs
+15 -6
View File
@@ -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)
codex-rs/core/src/codex.rs
+6 -3
View File
@@ -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(),
codex-rs/core/src/models_manager/manager.rs
+2 -2
View File
@@ -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);
codex-rs/core/tests/suite/auth_refresh.rs
+195 -59
View File
@@ -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<Utc>,
auth_manager: AuthManager,
_env_guard: EnvGuard,
}
impl RefreshTokenTestContext {
fn new(server: &MockServer) -> Result<Self> {
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()),
}
}
codex-rs/exec/src/lib.rs
+1 -1
View File
@@ -218,7 +218,7 @@ pub async fn run_main(cli: Cli, codex_linux_sandbox_exe: Option<PathBuf>) -> 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);
}
codex-rs/tui/src/chatwidget.rs
+7 -6
View File
@@ -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<String> {
}
async fn fetch_rate_limits(base_url: String, auth: CodexAuth) -> Option<RateLimitSnapshot> {
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) => {
codex-rs/tui/src/lib.rs
+1 -1
View File
@@ -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);
}
codex-rs/tui/src/status/helpers.rs
+1 -1
View File
@@ -88,7 +88,7 @@ pub(crate) fn compose_account_display(
auth_manager: &AuthManager,
plan: Option<PlanType>,
) -> Option<StatusAccountDisplay> {
let auth = auth_manager.auth()?;
let auth = auth_manager.auth_cached()?;
match auth.mode {
AuthMode::ChatGPT => {
codex-rs/tui2/src/chatwidget.rs
+7 -6
View File
@@ -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<String> {
}
async fn fetch_rate_limits(base_url: String, auth: CodexAuth) -> Option<RateLimitSnapshot> {
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) => {
codex-rs/tui2/src/lib.rs
+1 -1
View File
@@ -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);
}
codex-rs/tui2/src/status/helpers.rs
+1 -1
View File
@@ -88,7 +88,7 @@ pub(crate) fn compose_account_display(
auth_manager: &AuthManager,
plan: Option<PlanType>,
) -> Option<StatusAccountDisplay> {
let auth = auth_manager.auth()?;
let auth = auth_manager.auth_cached()?;
match auth.mode {
AuthMode::ChatGPT => {