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:
mpc-oai
2026-08-25 02:30:24 +00:00
committed by copyberry
parent c1db22a3cd
commit a7b86b6201
3 changed files with 113 additions and 22 deletions

View File

@@ -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!();

View File

@@ -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

View File

@@ -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
}