openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-14 11:57:03 +00:00
Code Issues Packages Projects Releases Wiki Activity

Mitigate token refresh storms

Browse Source
This commit is contained in:
Eric Traut
2026-03-10 15:25:40 -06:00
parent e4edafe1a8
commit a7f756819d
4 changed files with 269 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
codex-rs/app-server/src/codex_message_processor.rs
View File

@@ -1325,20 +1325,25 @@ impl CodexMessageProcessor {
}
}
async fn refresh_token_if_requested(&self, do_refresh: bool) {
async fn refresh_token_if_requested(&self, do_refresh: bool) -> bool {
if self.auth_manager.is_external_auth_active() {
return;
return false;
}
if do_refresh && let Err(err) = self.auth_manager.refresh_token().await {
tracing::warn!("failed to refresh token while getting account: {err}");
let failed_reason = err.failed_reason();
if failed_reason.is_none() {
tracing::warn!("failed to refresh token while getting account: {err}");
}
return failed_reason.is_some();
}
false
}
async fn get_auth_status(&self, request_id: ConnectionRequestId, params: GetAuthStatusParams) {
let include_token = params.include_token.unwrap_or(false);
let do_refresh = params.refresh_token.unwrap_or(false);
self.refresh_token_if_requested(do_refresh).await;
let refresh_failed_permanently = self.refresh_token_if_requested(do_refresh).await;
// Determine whether auth is required based on the active model provider.
// If a custom provider is configured with `requires_openai_auth == false`,
@@ -1352,18 +1357,25 @@ impl CodexMessageProcessor {
requires_openai_auth: Some(false),
}
} else {
let refresh_failure = self.auth_manager.refresh_failure();
match self.auth_manager.auth().await {
Some(auth) => {
let auth_mode = auth.api_auth_mode();
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)
}
Ok(_) => (None, None),
Err(err) => {
tracing::warn!("failed to get token for auth status: {err}");
(None, None)
let (reported_auth_method, token_opt) = if include_token
&& (refresh_failure.is_some() || refresh_failed_permanently)
{
(Some(auth_mode), None)
} else {
match auth.get_token() {
Ok(token) if !token.is_empty() => {
let tok = if include_token { Some(token) } else { None };
(Some(auth_mode), tok)
}
Ok(_) => (None, None),
Err(err) => {
tracing::warn!("failed to get token for auth status: {err}");
(None, None)
}
}
};
GetAuthStatusResponse {

90
codex-rs/app-server/tests/suite/auth.rs
View File

@@ -1,6 +1,8 @@
use anyhow::Result;
use app_test_support::ChatGptAuthFixture;
use app_test_support::McpProcess;
use app_test_support::to_response;
use app_test_support::write_chatgpt_auth;
use codex_app_server_protocol::AuthMode;
use codex_app_server_protocol::GetAuthStatusParams;
use codex_app_server_protocol::GetAuthStatusResponse;
@@ -8,10 +10,17 @@ use codex_app_server_protocol::JSONRPCError;
use codex_app_server_protocol::JSONRPCResponse;
use codex_app_server_protocol::LoginAccountResponse;
use codex_app_server_protocol::RequestId;
use codex_core::auth::AuthCredentialsStoreMode;
use codex_core::auth::REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR;
use pretty_assertions::assert_eq;
use std::path::Path;
use tempfile::TempDir;
use tokio::time::timeout;
use wiremock::Mock;
use wiremock::MockServer;
use wiremock::ResponseTemplate;
use wiremock::matchers::method;
use wiremock::matchers::path;
const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
@@ -207,6 +216,87 @@ async fn get_auth_status_with_api_key_no_include_token() -> Result<()> {
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn get_auth_status_omits_token_after_permanent_refresh_failure() -> Result<()> {
let codex_home = TempDir::new()?;
create_config_toml(codex_home.path())?;
write_chatgpt_auth(
codex_home.path(),
ChatGptAuthFixture::new("stale-access-token")
.refresh_token("stale-refresh-token")
.account_id("acct_123")
.email("user@example.com")
.plan_type("pro"),
AuthCredentialsStoreMode::File,
)?;
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/oauth/token"))
.respond_with(ResponseTemplate::new(401).set_body_json(serde_json::json!({
"error": {
"code": "refresh_token_reused"
}
})))
.expect(1)
.mount(&server)
.await;
let refresh_url = format!("{}/oauth/token", server.uri());
let mut mcp = McpProcess::new_with_env(
codex_home.path(),
&[
("OPENAI_API_KEY", None),
(
REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR,
Some(refresh_url.as_str()),
),
],
)
.await?;
timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??;
let request_id = mcp
.send_get_auth_status_request(GetAuthStatusParams {
include_token: Some(true),
refresh_token: Some(true),
})
.await?;
let resp: JSONRPCResponse = timeout(
DEFAULT_READ_TIMEOUT,
mcp.read_stream_until_response_message(RequestId::Integer(request_id)),
)
.await??;
let status: GetAuthStatusResponse = to_response(resp)?;
assert_eq!(
status,
GetAuthStatusResponse {
auth_method: Some(AuthMode::Chatgpt),
auth_token: None,
requires_openai_auth: Some(true),
}
);
let second_request_id = mcp
.send_get_auth_status_request(GetAuthStatusParams {
include_token: Some(true),
refresh_token: Some(true),
})
.await?;
let second_resp: JSONRPCResponse = timeout(
DEFAULT_READ_TIMEOUT,
mcp.read_stream_until_response_message(RequestId::Integer(second_request_id)),
)
.await??;
let second_status: GetAuthStatusResponse = to_response(second_resp)?;
assert_eq!(second_status, status);
server.verify().await;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn login_api_key_rejected_when_forced_chatgpt() -> Result<()> {
let codex_home = TempDir::new()?;

91
codex-rs/core/src/auth.rs
View File

@@ -815,6 +815,13 @@ struct CachedAuth {
auth: Option<CodexAuth>,
/// Callback used to refresh external auth by asking the parent app for new tokens.
external_refresher: Option<Arc<dyn ExternalAuthRefresher>>,
poisoned_managed_auth: Option<PoisonedManagedAuth>,
}
#[derive(Clone, Debug)]
struct PoisonedManagedAuth {
auth_dot_json: AuthDotJson,
error: RefreshTokenFailedError,
}
impl Debug for CachedAuth {
@@ -828,6 +835,13 @@ impl Debug for CachedAuth {
"external_refresher",
&self.external_refresher.as_ref().map(|_| "present"),
)
.field(
"poisoned_managed_auth",
&self
.poisoned_managed_auth
.as_ref()
.map(|poisoned| poisoned.error.reason),
)
.finish()
}
}
@@ -998,6 +1012,7 @@ impl AuthManager {
inner: RwLock::new(CachedAuth {
auth: managed_auth,
external_refresher: None,
poisoned_managed_auth: None,
}),
enable_codex_api_key_env,
auth_credentials_store_mode,
@@ -1010,6 +1025,7 @@ impl AuthManager {
let cached = CachedAuth {
auth: Some(auth),
external_refresher: None,
poisoned_managed_auth: None,
};
Arc::new(Self {
@@ -1029,6 +1045,7 @@ impl AuthManager {
let cached = CachedAuth {
auth: Some(auth),
external_refresher: None,
poisoned_managed_auth: None,
};
Arc::new(Self {
codex_home,
@@ -1044,6 +1061,11 @@ impl AuthManager {
self.inner.read().ok().and_then(|c| c.auth.clone())
}
pub fn refresh_failure(&self) -> Option<RefreshTokenFailedError> {
let auth = self.auth_cached()?;
self.refresh_failure_for_auth(&auth)
}
/// 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> {
@@ -1118,6 +1140,40 @@ impl AuthManager {
}
}
fn refresh_failure_for_auth(&self, auth: &CodexAuth) -> Option<RefreshTokenFailedError> {
let auth_dot_json = auth.get_current_auth_json()?;
self.inner
.read()
.ok()
.and_then(|cached| cached.poisoned_managed_auth.clone())
.filter(|poisoned| poisoned.auth_dot_json == auth_dot_json)
.map(|poisoned| poisoned.error)
}
fn poison_managed_auth_if_unchanged(
&self,
attempted_auth: &CodexAuth,
error: &RefreshTokenFailedError,
) {
let Some(attempted_auth_dot_json) = attempted_auth.get_current_auth_json() else {
return;
};
if let Ok(mut guard) = self.inner.write() {
let current_auth_matches = guard
.auth
.as_ref()
.and_then(CodexAuth::get_current_auth_json)
.is_some_and(|current| current == attempted_auth_dot_json);
if current_auth_matches {
guard.poisoned_managed_auth = Some(PoisonedManagedAuth {
auth_dot_json: attempted_auth_dot_json,
error: error.clone(),
});
}
}
}
fn load_auth_from_storage(&self) -> Option<CodexAuth> {
load_auth(
&self.codex_home,
@@ -1132,6 +1188,19 @@ impl AuthManager {
if let Ok(mut guard) = self.inner.write() {
let previous = guard.auth.as_ref();
let changed = !AuthManager::auths_equal(previous, new_auth.as_ref());
let poisoned_auth_still_matches = guard
.poisoned_managed_auth
.as_ref()
.and_then(|poisoned| {
new_auth
.as_ref()
.and_then(CodexAuth::get_current_auth_json)
.map(|current| current == poisoned.auth_dot_json)
})
.unwrap_or(false);
if !poisoned_auth_still_matches {
guard.poisoned_managed_auth = None;
}
tracing::info!("Reloaded auth, changed: {changed}");
guard.auth = new_auth;
changed
@@ -1197,6 +1266,11 @@ impl AuthManager {
/// token is the same as the cached, then ask the token authority to refresh.
pub async fn refresh_token(&self) -> Result<(), RefreshTokenError> {
let auth_before_reload = self.auth_cached();
if let Some(auth_before_reload) = auth_before_reload.as_ref()
&& let Some(error) = self.refresh_failure_for_auth(auth_before_reload)
{
return Err(RefreshTokenError::Permanent(error));
}
let expected_account_id = auth_before_reload
.as_ref()
.and_then(CodexAuth::get_account_id);
@@ -1227,7 +1301,12 @@ impl AuthManager {
Some(auth) => auth,
None => return Ok(()),
};
match auth {
if let Some(error) = self.refresh_failure_for_auth(&auth) {
return Err(RefreshTokenError::Permanent(error));
}
let attempted_auth = auth.clone();
let result = match auth {
CodexAuth::ChatgptAuthTokens(_) => {
self.refresh_external_auth(ExternalAuthRefreshReason::Unauthorized)
.await
@@ -1239,11 +1318,14 @@ impl AuthManager {
))
})?;
self.refresh_and_persist_chatgpt_token(&chatgpt_auth, token_data.refresh_token)
.await?;
Ok(())
.await
}
CodexAuth::ApiKey(_) => Ok(()),
};
if let Err(RefreshTokenError::Permanent(error)) = &result {
self.poison_managed_auth_if_unchanged(&attempted_auth, error);
}
result
}
/// Log out by deleting the on‑disk auth.json (if present). Returns Ok(true)
@@ -1266,6 +1348,9 @@ impl AuthManager {
}
async fn refresh_if_stale(&self, auth: &CodexAuth) -> Result<bool, RefreshTokenError> {
if let Some(error) = self.refresh_failure_for_auth(auth) {
return Err(RefreshTokenError::Permanent(error));
}
let chatgpt_auth = match auth {
CodexAuth::Chatgpt(chatgpt_auth) => chatgpt_auth,
_ => return Ok(false),

66
codex-rs/core/tests/suite/auth_refresh.rs
View File

@@ -433,6 +433,72 @@ async fn refresh_token_returns_permanent_error_for_expired_refresh_token() -> Re
Ok(())
}
#[serial_test::serial(auth_refresh)]
#[tokio::test]
async fn refresh_token_does_not_retry_after_permanent_failure() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/oauth/token"))
.respond_with(ResponseTemplate::new(401).set_body_json(json!({
"error": {
"code": "refresh_token_reused"
}
})))
.expect(1)
.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 {
auth_mode: Some(AuthMode::Chatgpt),
openai_api_key: None,
tokens: Some(initial_tokens.clone()),
last_refresh: Some(initial_last_refresh),
};
ctx.write_auth(&initial_auth)?;
let first_err = ctx
.auth_manager
.refresh_token()
.await
.err()
.context("first refresh should fail")?;
assert_eq!(
first_err.failed_reason(),
Some(RefreshTokenFailedReason::Exhausted)
);
let second_err = ctx
.auth_manager
.refresh_token()
.await
.err()
.context("second refresh should fail without retrying")?;
assert_eq!(
second_err.failed_reason(),
Some(RefreshTokenFailedReason::Exhausted)
);
let stored = ctx.load_auth()?;
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(())
}
#[serial_test::serial(auth_refresh)]
#[tokio::test]
async fn refresh_token_returns_transient_error_on_server_failure() -> Result<()> {