mirror of
https://github.com/openai/codex.git
synced 2026-10-31 19:28:49 +00:00
Fix remote executor authentication timeout recovery (#58355)
## Why Registration can time out while resolving credentials, before any request is sent. Treating this as a transport timeout prevents a safe retry. Credential resolution timeouts can also incorrectly quarantine credentials when no OAuth refresh exchange has started. ## What changed - Classify registration deadlines reached before the first send as `RequestNotSent`, allowing the existing registration retry loop to recover. Preserve transport timeout handling once a request has been sent. - Track unsettled ChatGPT refreshes in `AuthManager` from the start of the OAuth exchange through credential persistence. Cancellation while waiting for the refresh guard or reading credentials does not mark them as unsettled. - Have the executor auth provider withhold credentials with an unsettled refresh. Preserve uncertainty after a cancelled exchange even if a later refresh fails, while allowing replacement credentials and successfully completed refreshes. ## Testing Add regression coverage for registration and connection deadlines across auth stages, retrying registration after an unsent credential timeout, recovery after slow replacement PAT metadata, and cancellation before and during OAuth refresh. GitOrigin-RevId: f4d7ec2db7f79b8d313f1ee52dbd441094856669
This commit is contained in:
@@ -180,16 +180,12 @@ impl EnvironmentRegistryClient {
|
||||
security_profile: NOISE_RELAY_SECURITY_PROFILE.to_string(),
|
||||
executor_public_key: executor_public_key.clone(),
|
||||
};
|
||||
let (response, executor_credential) = timeout_at(
|
||||
deadline,
|
||||
self.send_authenticated(Request::new(Method::POST, url).with_json(&body)),
|
||||
)
|
||||
.await
|
||||
.unwrap_or_else(|_| {
|
||||
Err(ExecServerError::EnvironmentRegistryRequest(
|
||||
codex_http_client::RouteAwareRequestError::Timeout,
|
||||
))
|
||||
})?;
|
||||
let (response, executor_credential) = self
|
||||
.send_authenticated(
|
||||
Request::new(Method::POST, url).with_json(&body),
|
||||
Some(deadline),
|
||||
)
|
||||
.await?;
|
||||
let status = response.status();
|
||||
// Read diagnostics within the same attempt budget, preserving a known error status.
|
||||
let response: EnvironmentRegistryRegistrationResponse = timeout_at(deadline, async {
|
||||
@@ -264,7 +260,10 @@ impl EnvironmentRegistryClient {
|
||||
let response: EnvironmentRegistryConnectResponse =
|
||||
timeout_at(tokio::time::Instant::now() + self.connect_timeout, async {
|
||||
let (response, _) = self
|
||||
.send_authenticated(Request::new(Method::POST, url).with_json(&body))
|
||||
.send_authenticated(
|
||||
Request::new(Method::POST, url).with_json(&body),
|
||||
/*deadline*/ None,
|
||||
)
|
||||
.await?;
|
||||
self.parse_json_response(response).await
|
||||
})
|
||||
@@ -421,7 +420,10 @@ impl HarnessKeyValidator for RegistryHarnessKeyValidator {
|
||||
})?,
|
||||
);
|
||||
}
|
||||
let (response, executor_credential) = self.client.send_authenticated(request).await?;
|
||||
let (response, executor_credential) = self
|
||||
.client
|
||||
.send_authenticated(request, /*deadline*/ None)
|
||||
.await?;
|
||||
let status = response.status();
|
||||
if !status.is_success() {
|
||||
// The request contains the short-lived authorization. Do not include
|
||||
|
||||
@@ -1,10 +1,13 @@
|
||||
//! Sends registry requests with current credentials and at most one retry after a 401.
|
||||
//! Transport failures and other HTTP errors are never replayed here.
|
||||
//! A deadline before the first transport send remains a retryable unsent request.
|
||||
|
||||
use codex_http_client::HttpResponse;
|
||||
use codex_http_client::Request;
|
||||
use http::HeaderValue;
|
||||
use http::StatusCode;
|
||||
use tokio::time::Instant;
|
||||
use tokio::time::timeout_at;
|
||||
|
||||
use super::EnvironmentRegistryClient;
|
||||
use super::RegistryAuthMode;
|
||||
@@ -24,52 +27,70 @@ impl EnvironmentRegistryClient {
|
||||
pub(super) async fn send_authenticated(
|
||||
&self,
|
||||
request: Request,
|
||||
deadline: Option<Instant>,
|
||||
) -> Result<(HttpResponse, Option<HeaderValue>), ExecServerError> {
|
||||
let request = request
|
||||
.into_prepared()
|
||||
.map_err(ExecServerError::EnvironmentRegistryConfig)?;
|
||||
let mut retried = false;
|
||||
loop {
|
||||
let (authenticated, credential) =
|
||||
if self.auth_mode == RegistryAuthMode::ExecutorCredential {
|
||||
let (headers, credential) = self.resolve_auth().await?;
|
||||
let mut authenticated = request.clone();
|
||||
authenticated.headers.extend(headers);
|
||||
(authenticated, credential)
|
||||
} else {
|
||||
(
|
||||
self.auth_provider
|
||||
.apply_auth(request.clone())
|
||||
.await
|
||||
.map_err(registry_auth_error)?,
|
||||
None,
|
||||
)
|
||||
};
|
||||
let prepared = authenticated
|
||||
.prepare_body_for_send()
|
||||
.map_err(ExecServerError::EnvironmentRegistryConfig)?;
|
||||
let rejected_headers = prepared.headers.clone();
|
||||
let mut outgoing = self
|
||||
.http_for_auth()
|
||||
.request(authenticated.method, authenticated.url)
|
||||
.headers(prepared.headers)
|
||||
.headers(current_trace_context_headers())
|
||||
.body(prepared.body.unwrap_or_default());
|
||||
if let Some(timeout) = authenticated.timeout {
|
||||
outgoing = outgoing.timeout(timeout);
|
||||
let mut sent = false;
|
||||
let send = async {
|
||||
let mut retried = false;
|
||||
loop {
|
||||
let (authenticated, credential) =
|
||||
if self.auth_mode == RegistryAuthMode::ExecutorCredential {
|
||||
let (headers, credential) = self.resolve_auth().await?;
|
||||
let mut authenticated = request.clone();
|
||||
authenticated.headers.extend(headers);
|
||||
(authenticated, credential)
|
||||
} else {
|
||||
(
|
||||
self.auth_provider
|
||||
.apply_auth(request.clone())
|
||||
.await
|
||||
.map_err(registry_auth_error)?,
|
||||
None,
|
||||
)
|
||||
};
|
||||
let prepared = authenticated
|
||||
.prepare_body_for_send()
|
||||
.map_err(ExecServerError::EnvironmentRegistryConfig)?;
|
||||
let rejected_headers = prepared.headers.clone();
|
||||
let mut outgoing = self
|
||||
.http_for_auth()
|
||||
.request(authenticated.method, authenticated.url)
|
||||
.headers(prepared.headers)
|
||||
.headers(current_trace_context_headers())
|
||||
.body(prepared.body.unwrap_or_default());
|
||||
if let Some(timeout) = authenticated.timeout {
|
||||
outgoing = outgoing.timeout(timeout);
|
||||
}
|
||||
sent = true;
|
||||
let response = outgoing.send().await?;
|
||||
if !retried
|
||||
&& response.status() == StatusCode::UNAUTHORIZED
|
||||
&& self
|
||||
.auth_provider
|
||||
.refresh_after_unauthorized(rejected_headers)
|
||||
.await
|
||||
{
|
||||
retried = true;
|
||||
continue;
|
||||
}
|
||||
return Ok((response, credential));
|
||||
}
|
||||
let response = outgoing.send().await?;
|
||||
if !retried
|
||||
&& response.status() == StatusCode::UNAUTHORIZED
|
||||
&& self
|
||||
.auth_provider
|
||||
.refresh_after_unauthorized(rejected_headers)
|
||||
.await
|
||||
{
|
||||
retried = true;
|
||||
continue;
|
||||
};
|
||||
if let Some(deadline) = deadline {
|
||||
match timeout_at(deadline, send).await {
|
||||
Ok(result) => result,
|
||||
Err(_) if !sent => Err(ExecServerError::RequestNotSent(
|
||||
"environment registry credential resolution timed out".to_string(),
|
||||
)),
|
||||
Err(_) => Err(ExecServerError::EnvironmentRegistryRequest(
|
||||
codex_http_client::RouteAwareRequestError::Timeout,
|
||||
)),
|
||||
}
|
||||
return Ok((response, credential));
|
||||
} else {
|
||||
send.await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -84,32 +84,48 @@ impl AuthProvider for RotatingAuth {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn connect_deadline_covers_initial_auth_refresh_and_retry() -> anyhow::Result<()> {
|
||||
for stage in [AuthStage::Initial, AuthStage::Refresh, AuthStage::Retry] {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(wiremock::matchers::any())
|
||||
.respond_with(ResponseTemplate::new(401))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let provider = Arc::new(RotatingAuth {
|
||||
stall: Some(stage),
|
||||
..Default::default()
|
||||
});
|
||||
let mut client = EnvironmentRegistryClient::new(server.uri(), provider.clone())?;
|
||||
client.connect_timeout = Duration::from_secs(60);
|
||||
let key = NoiseChannelIdentity::generate()?.public_key();
|
||||
let task = tokio::spawn(async move { client.connect_environment("env", key).await });
|
||||
provider.reached.notified().await;
|
||||
// Advance only after real socket I/O reached the selected auth stage.
|
||||
tokio::time::pause();
|
||||
tokio::time::advance(Duration::from_secs(61)).await;
|
||||
let result = tokio::time::timeout(Duration::from_secs(1), task)
|
||||
.await
|
||||
.expect("connect must finish at its original deadline")?;
|
||||
assert!(
|
||||
matches!(result, Err(ExecServerError::EnvironmentRegistryRequest(error)) if error.is_timeout())
|
||||
);
|
||||
tokio::time::resume();
|
||||
async fn registry_deadline_covers_initial_auth_refresh_and_retry() -> anyhow::Result<()> {
|
||||
for operation in ["connect", "register"] {
|
||||
for stage in [AuthStage::Initial, AuthStage::Refresh, AuthStage::Retry] {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(wiremock::matchers::any())
|
||||
.respond_with(ResponseTemplate::new(401))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let provider = Arc::new(RotatingAuth {
|
||||
stall: Some(stage),
|
||||
..Default::default()
|
||||
});
|
||||
let mut client = EnvironmentRegistryClient::new(server.uri(), provider.clone())?;
|
||||
client.connect_timeout = Duration::from_secs(60);
|
||||
let key = NoiseChannelIdentity::generate()?.public_key();
|
||||
let task = tokio::spawn(async move {
|
||||
if operation == "register" {
|
||||
client.register_environment("env", &key).await.map(|_| ())
|
||||
} else {
|
||||
client.connect_environment("env", key).await.map(|_| ())
|
||||
}
|
||||
});
|
||||
provider.reached.notified().await;
|
||||
// Advance only after real socket I/O reached the selected auth stage.
|
||||
tokio::time::pause();
|
||||
tokio::time::advance(Duration::from_secs(61)).await;
|
||||
let result = tokio::time::timeout(Duration::from_secs(1), task)
|
||||
.await
|
||||
.expect("registry operation must finish at its original deadline")?;
|
||||
if operation == "register" && stage == AuthStage::Initial {
|
||||
assert!(matches!(result, Err(ExecServerError::RequestNotSent(_))));
|
||||
} else {
|
||||
assert!(
|
||||
matches!(result, Err(ExecServerError::EnvironmentRegistryRequest(error)) if error.is_timeout())
|
||||
);
|
||||
}
|
||||
assert_eq!(
|
||||
server.received_requests().await.unwrap().len(),
|
||||
if stage == AuthStage::Initial { 0 } else { 1 }
|
||||
);
|
||||
tokio::time::resume();
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -192,7 +208,10 @@ async fn only_401_retries_and_only_once() -> anyhow::Result<()> {
|
||||
let provider = Arc::new(RotatingAuth::default());
|
||||
let client = EnvironmentRegistryClient::new(server.uri(), provider.clone())?;
|
||||
let (response, _) = client
|
||||
.send_authenticated(Request::new(Method::POST, server.uri()))
|
||||
.send_authenticated(
|
||||
Request::new(Method::POST, server.uri()),
|
||||
/*deadline*/ None,
|
||||
)
|
||||
.await?;
|
||||
assert_eq!(response.status().as_u16(), status);
|
||||
assert_eq!(provider.refreshes.load(Ordering::SeqCst), refreshes);
|
||||
|
||||
@@ -41,6 +41,7 @@ const SUCCESS_BODY: &str = r#"{"environment_id":"registration-retry-test","url":
|
||||
struct RegistrationAuthProvider {
|
||||
calls: AtomicUsize,
|
||||
fail_first: bool,
|
||||
stall_first: bool,
|
||||
}
|
||||
|
||||
impl AuthProvider for RegistrationAuthProvider {
|
||||
@@ -51,6 +52,9 @@ impl AuthProvider for RegistrationAuthProvider {
|
||||
fn apply_auth_headers(&self, headers: HeaderMap) -> codex_api::AuthHeadersFuture<'_> {
|
||||
if self.calls.fetch_add(1, Ordering::Relaxed) == 0 && self.fail_first {
|
||||
return Box::pin(async {
|
||||
if self.stall_first {
|
||||
std::future::pending::<()>().await;
|
||||
}
|
||||
Err(codex_api::AuthError::Transient(
|
||||
"metadata unavailable".to_string(),
|
||||
))
|
||||
@@ -62,30 +66,34 @@ impl AuthProvider for RegistrationAuthProvider {
|
||||
|
||||
#[tokio::test]
|
||||
async fn transient_credential_failure_retries_before_sending_registration() -> anyhow::Result<()> {
|
||||
let server = wiremock::MockServer::start().await;
|
||||
wiremock::Mock::given(wiremock::matchers::any())
|
||||
.respond_with(
|
||||
wiremock::ResponseTemplate::new(200).set_body_raw(SUCCESS_BODY, "application/json"),
|
||||
for stall_first in [false, true] {
|
||||
let server = wiremock::MockServer::start().await;
|
||||
wiremock::Mock::given(wiremock::matchers::any())
|
||||
.respond_with(
|
||||
wiremock::ResponseTemplate::new(200).set_body_raw(SUCCESS_BODY, "application/json"),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let provider = Arc::new(RegistrationAuthProvider {
|
||||
fail_first: true,
|
||||
stall_first,
|
||||
..Default::default()
|
||||
});
|
||||
let mut client = EnvironmentRegistryClient::new(server.uri(), provider.clone())?;
|
||||
client.connect_timeout = Duration::from_secs(1);
|
||||
let response = timeout(
|
||||
Duration::from_secs(10),
|
||||
client.register_environment_with_retry(
|
||||
ENVIRONMENT_ID,
|
||||
&NoiseChannelIdentity::generate()?.public_key(),
|
||||
),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let provider = Arc::new(RegistrationAuthProvider {
|
||||
fail_first: true,
|
||||
..Default::default()
|
||||
});
|
||||
let client = EnvironmentRegistryClient::new(server.uri(), provider.clone())?;
|
||||
let response = timeout(
|
||||
Duration::from_secs(10),
|
||||
client.register_environment_with_retry(
|
||||
ENVIRONMENT_ID,
|
||||
&NoiseChannelIdentity::generate()?.public_key(),
|
||||
),
|
||||
)
|
||||
.await??;
|
||||
assert_eq!(response.environment_id, ENVIRONMENT_ID);
|
||||
assert_eq!(provider.calls.load(Ordering::Relaxed), 2);
|
||||
server.verify().await;
|
||||
.await??;
|
||||
assert_eq!(response.environment_id, ENVIRONMENT_ID);
|
||||
assert_eq!(provider.calls.load(Ordering::Relaxed), 2);
|
||||
server.verify().await;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user