mirror of
https://github.com/openai/codex.git
synced 2026-09-05 15:18:41 +00:00
Retry provider auth commands after initial failures (#40523)
## Why A provider auth command can fail during initial credential resolution, leaving no cached authentication. A subsequent `401` should still give the configured provider one bounded opportunity to recover. ## What changed - Treat configured external authentication with an empty cache as eligible for unauthorized recovery. - Allow authority refresh to invoke the external provider when no prior auth is cached, while preserving the existing requirements for managers without an external provider. ## Testing Added unit and client coverage for a provider command that fails initially and succeeds during `401` recovery. GitOrigin-RevId: 7cab9a03f90ce29b412be8fbd97b297acdad4482
This commit is contained in:
@@ -745,6 +745,9 @@ impl ProviderAuthCommandFixture {
|
||||
std::fs::write(
|
||||
&script_path,
|
||||
r#"#!/bin/sh
|
||||
if [ -f fail-until-401 ]; then
|
||||
exit 1
|
||||
fi
|
||||
first_line=$(sed -n '1p' tokens.txt)
|
||||
printf '%s\n' "$first_line"
|
||||
tail -n +2 tokens.txt > tokens.next
|
||||
@@ -767,6 +770,7 @@ mv tokens.next tokens.txt
|
||||
&script_path,
|
||||
r#"@echo off
|
||||
setlocal EnableExtensions DisableDelayedExpansion
|
||||
if exist fail-until-401 exit /b 1
|
||||
|
||||
set "first_line="
|
||||
<tokens.txt set /p first_line=
|
||||
@@ -1420,6 +1424,43 @@ async fn provider_auth_command_refreshes_after_401() {
|
||||
send_provider_auth_request(&server, auth_fixture.auth()).await;
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn provider_auth_command_recovers_after_initial_resolution_failure() {
|
||||
skip_if_no_network!();
|
||||
|
||||
let server = MockServer::start().await;
|
||||
let auth_fixture = ProviderAuthCommandFixture::new(&["recovered-token"]).unwrap();
|
||||
let failure_marker = auth_fixture.tempdir.path().join("fail-until-401");
|
||||
std::fs::write(&failure_marker, "").unwrap();
|
||||
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/v1/responses"))
|
||||
.and(|request: &wiremock::Request| !request.headers.contains_key("authorization"))
|
||||
.respond_with(move |_request: &wiremock::Request| {
|
||||
std::fs::remove_file(&failure_marker).unwrap();
|
||||
ResponseTemplate::new(401).set_body_string("unauthorized")
|
||||
})
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/v1/responses"))
|
||||
.and(header("authorization", "Bearer recovered-token"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.insert_header("content-type", "text/event-stream")
|
||||
.set_body_raw(
|
||||
sse(vec![ev_response_created("resp1"), ev_completed("resp1")]),
|
||||
"text/event-stream",
|
||||
),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
send_provider_auth_request(&server, auth_fixture.auth()).await;
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn amazon_bedrock_proxy_uses_command_auth_and_custom_headers() {
|
||||
skip_if_no_network!();
|
||||
|
||||
@@ -1280,6 +1280,44 @@ async fn external_bearer_only_auth_manager_returns_none_when_command_fails() {
|
||||
assert_eq!(manager.auth().await, None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unauthorized_recovery_retries_provider_command_after_initial_failure() {
|
||||
let script = ProviderAuthScript::new(&["provider-token"]).unwrap();
|
||||
std::fs::write(script.tempdir.path().join("fail-once"), "").unwrap();
|
||||
let manager = AuthManager::external_bearer_only(script.auth_config());
|
||||
let mut recovery = manager.unauthorized_recovery();
|
||||
|
||||
assert_eq!(manager.auth().await, None);
|
||||
assert_eq!(manager.auth_cached(), None);
|
||||
assert!(recovery.has_next());
|
||||
assert_eq!(recovery.unavailable_reason(), "ready");
|
||||
|
||||
let result = recovery
|
||||
.next()
|
||||
.await
|
||||
.expect("external refresh should succeed");
|
||||
|
||||
assert_eq!(result.auth_state_changed(), Some(true));
|
||||
assert_eq!(
|
||||
manager.auth_cached(),
|
||||
Some(CodexAuth::from_api_key("provider-token"))
|
||||
);
|
||||
assert!(!recovery.has_next());
|
||||
assert_eq!(recovery.unavailable_reason(), "recovery_exhausted");
|
||||
recovery.next().await.expect_err("recovery is bounded");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unauthorized_recovery_without_an_external_provider_still_requires_refreshable_auth() {
|
||||
for auth in [None, Some(CodexAuth::from_api_key("static-token"))] {
|
||||
let manager = AuthManager::from_optional_auth_for_testing(auth);
|
||||
let recovery = manager.unauthorized_recovery();
|
||||
|
||||
assert!(!recovery.has_next());
|
||||
assert_eq!(recovery.unavailable_reason(), "not_chatgpt_auth");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unauthorized_recovery_uses_external_refresh_for_bearer_manager() {
|
||||
let script = ProviderAuthScript::new(&["provider-token", "refreshed-provider-token"]).unwrap();
|
||||
@@ -1619,6 +1657,10 @@ impl ProviderAuthScript {
|
||||
std::fs::write(
|
||||
&script_path,
|
||||
r#"#!/bin/sh
|
||||
if [ -f fail-once ]; then
|
||||
rm fail-once
|
||||
exit 1
|
||||
fi
|
||||
first_line=$(sed -n '1p' tokens.txt)
|
||||
printf '%s\n' "$first_line"
|
||||
tail -n +2 tokens.txt > tokens.next
|
||||
@@ -1641,6 +1683,10 @@ mv tokens.next tokens.txt
|
||||
&script_path,
|
||||
r#"@echo off
|
||||
setlocal EnableExtensions DisableDelayedExpansion
|
||||
if exist fail-once (
|
||||
del fail-once
|
||||
exit /b 1
|
||||
)
|
||||
set "first_line="
|
||||
<tokens.txt set /p "first_line="
|
||||
if not defined first_line exit /b 1
|
||||
|
||||
@@ -1890,7 +1890,7 @@ impl UnauthorizedRecovery {
|
||||
}
|
||||
|
||||
pub fn has_next(&self) -> bool {
|
||||
if self.manager.has_external_api_key_auth() {
|
||||
if self.manager.has_refreshable_external_auth() {
|
||||
return !matches!(self.step, UnauthorizedRecoveryStep::Done);
|
||||
}
|
||||
|
||||
@@ -1911,7 +1911,7 @@ impl UnauthorizedRecovery {
|
||||
}
|
||||
|
||||
pub fn unavailable_reason(&self) -> &'static str {
|
||||
if self.manager.has_external_api_key_auth() {
|
||||
if self.manager.has_refreshable_external_auth() {
|
||||
return if matches!(self.step, UnauthorizedRecoveryStep::Done) {
|
||||
"recovery_exhausted"
|
||||
} else {
|
||||
@@ -2741,12 +2741,12 @@ impl AuthManager {
|
||||
.and_then(|external_auth| external_auth.clone())
|
||||
}
|
||||
|
||||
fn has_external_api_key_auth(&self) -> bool {
|
||||
fn has_refreshable_external_auth(&self) -> bool {
|
||||
self.has_external_auth()
|
||||
&& self
|
||||
.auth_cached()
|
||||
.as_ref()
|
||||
.is_some_and(CodexAuth::is_api_key_auth)
|
||||
.is_none_or(|auth| auth.is_api_key_auth() || auth.supports_unauthorized_recovery())
|
||||
}
|
||||
|
||||
async fn resolve_external_auth(
|
||||
@@ -2817,40 +2817,44 @@ impl AuthManager {
|
||||
async fn refresh_token_from_authority_impl(&self) -> Result<(), RefreshTokenError> {
|
||||
tracing::info!("Refreshing token");
|
||||
|
||||
let auth = match self.auth_cached() {
|
||||
Some(auth) => auth,
|
||||
None => return Ok(()),
|
||||
};
|
||||
if let Some(error) = self.refresh_failure_for_auth(&auth) {
|
||||
let attempted_auth = self.auth_cached();
|
||||
if let Some(error) = attempted_auth
|
||||
.as_ref()
|
||||
.and_then(|auth| self.refresh_failure_for_auth(auth))
|
||||
{
|
||||
return Err(RefreshTokenError::Permanent(error));
|
||||
}
|
||||
|
||||
let attempted_auth = auth.clone();
|
||||
let result = if self.has_external_auth() {
|
||||
self.refresh_external_auth(ExternalAuthRefreshReason::Unauthorized)
|
||||
.await
|
||||
} else {
|
||||
match auth {
|
||||
CodexAuth::Chatgpt(chatgpt_auth) => {
|
||||
match attempted_auth.as_ref() {
|
||||
Some(CodexAuth::Chatgpt(chatgpt_auth)) => {
|
||||
let token_data = chatgpt_auth.current_token_data().ok_or_else(|| {
|
||||
RefreshTokenError::Transient(std::io::Error::other(
|
||||
"Token data is not available.",
|
||||
))
|
||||
})?;
|
||||
self.refresh_and_persist_chatgpt_token(&chatgpt_auth, token_data.refresh_token)
|
||||
self.refresh_and_persist_chatgpt_token(chatgpt_auth, token_data.refresh_token)
|
||||
.await
|
||||
}
|
||||
CodexAuth::ApiKey(_)
|
||||
| CodexAuth::ChatgptAuthTokens(_)
|
||||
| CodexAuth::Headers(_)
|
||||
| CodexAuth::AgentIdentity(_)
|
||||
| CodexAuth::PersonalAccessToken(_)
|
||||
| CodexAuth::BedrockApiKey(_)
|
||||
| CodexAuth::BedrockAccessKeys(_) => Ok(()),
|
||||
Some(
|
||||
CodexAuth::ApiKey(_)
|
||||
| CodexAuth::ChatgptAuthTokens(_)
|
||||
| CodexAuth::Headers(_)
|
||||
| CodexAuth::AgentIdentity(_)
|
||||
| CodexAuth::PersonalAccessToken(_)
|
||||
| CodexAuth::BedrockApiKey(_)
|
||||
| CodexAuth::BedrockAccessKeys(_),
|
||||
)
|
||||
| None => Ok(()),
|
||||
}
|
||||
};
|
||||
if let Err(RefreshTokenError::Permanent(error)) = &result {
|
||||
self.record_permanent_refresh_failure_if_unchanged(&attempted_auth, error);
|
||||
if let Some(attempted_auth) = attempted_auth.as_ref()
|
||||
&& let Err(RefreshTokenError::Permanent(error)) = &result
|
||||
{
|
||||
self.record_permanent_refresh_failure_if_unchanged(attempted_auth, error);
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user