mirror of
https://github.com/openai/codex.git
synced 2026-10-21 17:48:44 +00:00
Propagate authentication header failures before sending requests (#55607)
## Why Synchronous auth header construction could silently return no credentials when managed authentication became unavailable. Callers need an error to avoid sending requests with missing or partial authentication. ## What changed - Make `AuthProvider::add_auth_headers`, `to_auth_headers`, and auth header telemetry fallible, and propagate failures through request and WebSocket setup. - Return a transient error when managed authentication is unavailable, and preserve network policy denials through auth error conversion. - Retry transient MCP authentication errors while treating build errors and policy denials as terminal. - Cancel local MCP event streams before resolving authentication so cleanup still works when credentials become unavailable. ## Testing Add regression coverage for rejected partial authentication, async header resolution, terminal auth errors without sending requests, managed account changes, MCP retry classification, and local stream cancellation after authentication becomes unavailable. GitOrigin-RevId: 39d981e232cba1e302de502a2fa6239101114ed1
This commit is contained in:
@@ -1,5 +1,6 @@
|
||||
use codex_client::Request;
|
||||
use codex_client::TransportError;
|
||||
use codex_http_client::NetworkPolicyDenied;
|
||||
use http::HeaderMap;
|
||||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
@@ -8,6 +9,8 @@ use std::sync::Arc;
|
||||
/// Error returned while applying authentication to an outbound request.
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum AuthError {
|
||||
#[error(transparent)]
|
||||
Policy(#[from] NetworkPolicyDenied),
|
||||
#[error("request auth build error: {0}")]
|
||||
Build(String),
|
||||
#[error("transient auth error: {0}")]
|
||||
@@ -17,6 +20,7 @@ pub enum AuthError {
|
||||
impl From<AuthError> for TransportError {
|
||||
fn from(error: AuthError) -> Self {
|
||||
match error {
|
||||
AuthError::Policy(denied) => TransportError::Policy(denied),
|
||||
AuthError::Build(message) => TransportError::Build(message),
|
||||
AuthError::Transient(message) => TransportError::Network(message),
|
||||
}
|
||||
@@ -31,14 +35,16 @@ pub trait AuthProvider: Send + Sync {
|
||||
/// Adds any auth headers that are available without request body access.
|
||||
///
|
||||
/// Implementations should be cheap and non-blocking. This method is also
|
||||
/// used by telemetry and non-HTTP request paths.
|
||||
fn add_auth_headers(&self, headers: &mut HeaderMap);
|
||||
/// used by telemetry and non-HTTP request paths. Providers that require async
|
||||
/// work must return an error here and override `resolve_auth_headers`.
|
||||
/// Callers must not send the request if this returns an error.
|
||||
fn add_auth_headers(&self, headers: &mut HeaderMap) -> Result<(), AuthError>;
|
||||
|
||||
/// Returns any auth headers that are available without request body access.
|
||||
fn to_auth_headers(&self) -> HeaderMap {
|
||||
fn to_auth_headers(&self) -> Result<HeaderMap, AuthError> {
|
||||
let mut headers = HeaderMap::new();
|
||||
self.add_auth_headers(&mut headers);
|
||||
headers
|
||||
self.add_auth_headers(&mut headers)?;
|
||||
Ok(headers)
|
||||
}
|
||||
|
||||
/// Resolves auth headers for an outbound request.
|
||||
@@ -47,7 +53,7 @@ pub trait AuthProvider: Send + Sync {
|
||||
/// credentials before returning. Header-only providers with static credentials can rely on the
|
||||
/// default implementation.
|
||||
fn resolve_auth_headers(&self) -> AuthHeadersFuture<'_> {
|
||||
Box::pin(async { Ok(self.to_auth_headers()) })
|
||||
Box::pin(async { self.to_auth_headers() })
|
||||
}
|
||||
|
||||
/// Applies auth to a complete outbound request and returns the request to send.
|
||||
@@ -91,14 +97,14 @@ pub struct AuthHeaderTelemetry {
|
||||
pub name: Option<&'static str>,
|
||||
}
|
||||
|
||||
pub fn auth_header_telemetry(auth: &dyn AuthProvider) -> AuthHeaderTelemetry {
|
||||
pub fn auth_header_telemetry(auth: &dyn AuthProvider) -> Result<AuthHeaderTelemetry, AuthError> {
|
||||
let mut headers = HeaderMap::new();
|
||||
auth.add_auth_headers(&mut headers);
|
||||
auth.add_auth_headers(&mut headers)?;
|
||||
let name = headers
|
||||
.contains_key(http::header::AUTHORIZATION)
|
||||
.then_some("authorization");
|
||||
AuthHeaderTelemetry {
|
||||
Ok(AuthHeaderTelemetry {
|
||||
attached: name.is_some(),
|
||||
name,
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
@@ -139,7 +139,9 @@ mod tests {
|
||||
struct DummyAuth;
|
||||
|
||||
impl AuthProvider for DummyAuth {
|
||||
fn add_auth_headers(&self, _headers: &mut HeaderMap) {}
|
||||
fn add_auth_headers(&self, _headers: &mut HeaderMap) -> Result<(), crate::AuthError> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
|
||||
@@ -102,7 +102,9 @@ mod tests {
|
||||
struct DummyAuth;
|
||||
|
||||
impl AuthProvider for DummyAuth {
|
||||
fn add_auth_headers(&self, _headers: &mut HeaderMap) {}
|
||||
fn add_auth_headers(&self, _headers: &mut HeaderMap) -> Result<(), crate::AuthError> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
|
||||
@@ -184,7 +184,9 @@ mod tests {
|
||||
struct DummyAuth;
|
||||
|
||||
impl AuthProvider for DummyAuth {
|
||||
fn add_auth_headers(&self, _headers: &mut HeaderMap) {}
|
||||
fn add_auth_headers(&self, _headers: &mut HeaderMap) -> Result<(), crate::AuthError> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
@@ -193,8 +195,9 @@ mod tests {
|
||||
}
|
||||
|
||||
impl AuthProvider for RetryAuth {
|
||||
fn add_auth_headers(&self, headers: &mut HeaderMap) {
|
||||
fn add_auth_headers(&self, headers: &mut HeaderMap) -> Result<(), AuthError> {
|
||||
headers.insert(http::header::AUTHORIZATION, "Bearer test".parse().unwrap());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn apply_auth(&self, mut request: Request) -> AuthProviderFuture<'_> {
|
||||
@@ -204,7 +207,7 @@ mod tests {
|
||||
if limits.len() == 1 {
|
||||
return Err(AuthError::Transient("retry authentication".to_string()));
|
||||
}
|
||||
self.add_auth_headers(&mut request.headers);
|
||||
self.add_auth_headers(&mut request.headers)?;
|
||||
Ok(request)
|
||||
})
|
||||
}
|
||||
|
||||
@@ -360,11 +360,12 @@ mod tests {
|
||||
struct DummyAuth;
|
||||
|
||||
impl AuthProvider for DummyAuth {
|
||||
fn add_auth_headers(&self, headers: &mut HeaderMap) {
|
||||
fn add_auth_headers(&self, headers: &mut HeaderMap) -> Result<(), crate::AuthError> {
|
||||
headers.insert(
|
||||
http::header::AUTHORIZATION,
|
||||
HeaderValue::from_static("Bearer test-token"),
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -466,7 +466,7 @@ impl ResponsesWebsocketClient {
|
||||
|
||||
let mut headers =
|
||||
merge_request_headers(&self.provider.headers, extra_headers, default_headers);
|
||||
self.auth.add_auth_headers(&mut headers);
|
||||
self.auth.add_auth_headers(&mut headers)?;
|
||||
|
||||
let connector = WebSocketConnector::new(http_client_factory)
|
||||
.map_err(|err| ApiError::Stream(format!("failed to configure websocket TLS: {err}")))?;
|
||||
|
||||
@@ -35,7 +35,7 @@ impl ResponsesWebsocketClient {
|
||||
.map_err(|err| ApiError::Stream(format!("failed to build websocket URL: {err}")))?;
|
||||
let mut headers =
|
||||
merge_request_headers(&self.provider.headers, extra_headers, default_headers);
|
||||
self.auth.add_auth_headers(&mut headers);
|
||||
self.auth.add_auth_headers(&mut headers)?;
|
||||
let (stream, _status, server_reasoning_included, server_model) =
|
||||
connect_websocket(ws_url, headers, connector, turn_state).await?;
|
||||
Ok(ResponsesWebsocketConnection::new(
|
||||
|
||||
@@ -83,7 +83,9 @@ mod tests {
|
||||
struct DummyAuth;
|
||||
|
||||
impl AuthProvider for DummyAuth {
|
||||
fn add_auth_headers(&self, _headers: &mut HeaderMap) {}
|
||||
fn add_auth_headers(&self, _headers: &mut HeaderMap) -> Result<(), crate::AuthError> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
|
||||
@@ -66,3 +66,9 @@ pub(crate) fn parse_flex_unavailable(error: &Value) -> Option<ApiError> {
|
||||
(error.get("code").and_then(Value::as_str) == Some("flex_unavailable"))
|
||||
.then_some(ApiError::FlexUnavailable)
|
||||
}
|
||||
|
||||
impl From<crate::AuthError> for ApiError {
|
||||
fn from(error: crate::AuthError) -> Self {
|
||||
Self::Transport(error.into())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -46,6 +46,8 @@ pub struct UploadedOpenAiFile {
|
||||
pub enum OpenAiFileError {
|
||||
#[error("failed to open OpenAI file upload contents: {0}")]
|
||||
ReadContents(#[source] std::io::Error),
|
||||
#[error(transparent)]
|
||||
Auth(#[from] crate::AuthError),
|
||||
#[error(
|
||||
"file `{file_name}` is too large: {size_bytes} bytes exceeds the limit of {limit_bytes} bytes"
|
||||
)]
|
||||
@@ -152,7 +154,7 @@ where
|
||||
"file_size": file_size_bytes,
|
||||
"use_case": OPENAI_FILE_USE_CASE,
|
||||
});
|
||||
let request = authorized_request(client_pool, auth, Method::POST, &create_url);
|
||||
let request = authorized_request(client_pool, auth, Method::POST, &create_url)?;
|
||||
let create_request = match hosted_upload {
|
||||
Some(context) => serde_json::json!({
|
||||
"file_name": file_name.as_str(),
|
||||
@@ -363,7 +365,7 @@ where
|
||||
};
|
||||
let finalize_started_at = Instant::now();
|
||||
loop {
|
||||
let finalize_response = authorized_request(client_pool, auth, Method::POST, &finalize_url)
|
||||
let finalize_response = authorized_request(client_pool, auth, Method::POST, &finalize_url)?
|
||||
.json(&finalize_request)
|
||||
.send()
|
||||
.await
|
||||
@@ -435,14 +437,14 @@ fn authorized_request(
|
||||
auth: &dyn AuthProvider,
|
||||
method: Method,
|
||||
url: &str,
|
||||
) -> RequestBuilder {
|
||||
) -> Result<RequestBuilder, OpenAiFileError> {
|
||||
let mut headers = http::HeaderMap::new();
|
||||
auth.add_auth_headers(&mut headers);
|
||||
auth.add_auth_headers(&mut headers)?;
|
||||
|
||||
client_pool
|
||||
Ok(client_pool
|
||||
.request(method, url)
|
||||
.timeout(OPENAI_FILE_REQUEST_TIMEOUT)
|
||||
.headers(headers)
|
||||
.headers(headers))
|
||||
}
|
||||
|
||||
fn upload_response_header(response: &HttpResponse, header: &str) -> String {
|
||||
@@ -511,12 +513,13 @@ mod tests {
|
||||
}
|
||||
|
||||
impl AuthProvider for ChatGptTestAuth {
|
||||
fn add_auth_headers(&self, headers: &mut http::HeaderMap) {
|
||||
fn add_auth_headers(&self, headers: &mut http::HeaderMap) -> Result<(), crate::AuthError> {
|
||||
headers.insert(
|
||||
http::header::AUTHORIZATION,
|
||||
HeaderValue::from_static("Bearer token"),
|
||||
);
|
||||
headers.insert("ChatGPT-Account-ID", HeaderValue::from_static("account_id"));
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -105,7 +105,9 @@ impl HttpTransport for RecordingTransport {
|
||||
struct NoAuth;
|
||||
|
||||
impl AuthProvider for NoAuth {
|
||||
fn add_auth_headers(&self, _headers: &mut HeaderMap) {}
|
||||
fn add_auth_headers(&self, _headers: &mut HeaderMap) -> Result<(), AuthError> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
@@ -124,7 +126,7 @@ impl StaticAuth {
|
||||
}
|
||||
|
||||
impl AuthProvider for StaticAuth {
|
||||
fn add_auth_headers(&self, headers: &mut HeaderMap) {
|
||||
fn add_auth_headers(&self, headers: &mut HeaderMap) -> Result<(), AuthError> {
|
||||
let token = &self.token;
|
||||
if let Ok(header) = HeaderValue::from_str(&format!("Bearer {token}")) {
|
||||
headers.insert(http::header::AUTHORIZATION, header);
|
||||
@@ -132,6 +134,7 @@ impl AuthProvider for StaticAuth {
|
||||
if let Ok(header) = HeaderValue::from_str(&self.account_id) {
|
||||
headers.insert("ChatGPT-Account-ID", header);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -199,19 +202,10 @@ struct FailsOnceAuth {
|
||||
}
|
||||
|
||||
impl FailsOnceAuth {
|
||||
fn transient() -> Self {
|
||||
fn new(error: AuthError) -> Self {
|
||||
Self {
|
||||
attempts: Arc::new(Mutex::new(0)),
|
||||
error: Arc::new(AuthError::Transient(
|
||||
"sts temporarily unavailable".to_string(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn build() -> Self {
|
||||
Self {
|
||||
attempts: Arc::new(Mutex::new(0)),
|
||||
error: Arc::new(AuthError::Build("invalid auth configuration".to_string())),
|
||||
error: Arc::new(error),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -231,6 +225,7 @@ impl FailsOnceAuth {
|
||||
|
||||
if *attempts == 1 {
|
||||
return match self.error.as_ref() {
|
||||
AuthError::Policy(denied) => Err(AuthError::Policy(*denied)),
|
||||
AuthError::Build(message) => Err(AuthError::Build(message.clone())),
|
||||
AuthError::Transient(message) => Err(AuthError::Transient(message.clone())),
|
||||
};
|
||||
@@ -241,7 +236,9 @@ impl FailsOnceAuth {
|
||||
}
|
||||
|
||||
impl AuthProvider for FailsOnceAuth {
|
||||
fn add_auth_headers(&self, _headers: &mut HeaderMap) {}
|
||||
fn add_auth_headers(&self, _headers: &mut HeaderMap) -> Result<(), AuthError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn apply_auth(&self, request: Request) -> codex_api::AuthProviderFuture<'_> {
|
||||
Box::pin(FailsOnceAuth::apply_auth(self, request))
|
||||
@@ -545,7 +542,7 @@ async fn streaming_client_retries_on_transport_error() -> Result<()> {
|
||||
async fn streaming_client_retries_on_transient_auth_error() -> Result<()> {
|
||||
let state = RecordingState::default();
|
||||
let transport = RecordingTransport::new(state.clone());
|
||||
let auth = FailsOnceAuth::transient();
|
||||
let auth = FailsOnceAuth::new(AuthError::Transient("sts temporarily unavailable".into()));
|
||||
|
||||
let mut provider = provider("openai");
|
||||
provider.retry.max_attempts = 2;
|
||||
@@ -567,36 +564,39 @@ async fn streaming_client_retries_on_transient_auth_error() -> Result<()> {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn streaming_client_does_not_retry_auth_build_error() -> Result<()> {
|
||||
let state = RecordingState::default();
|
||||
let transport = RecordingTransport::new(state.clone());
|
||||
let auth = FailsOnceAuth::build();
|
||||
|
||||
let mut provider = provider("openai");
|
||||
provider.retry.max_attempts = 2;
|
||||
|
||||
let client = ResponsesClient::new(transport, provider, Arc::new(auth.clone()));
|
||||
let body = serde_json::json!({ "model": "gpt-test" });
|
||||
let result = client
|
||||
.stream(
|
||||
body,
|
||||
HeaderMap::new(),
|
||||
Compression::None,
|
||||
/*turn_state*/ None,
|
||||
)
|
||||
.await;
|
||||
let err = result
|
||||
.err()
|
||||
.expect("auth build errors should fail without retry");
|
||||
|
||||
assert!(matches!(
|
||||
err,
|
||||
ApiError::Transport(TransportError::Build(message))
|
||||
if message == "invalid auth configuration"
|
||||
));
|
||||
assert_eq!(auth.attempts(), 1);
|
||||
assert_eq!(state.take_stream_requests().len(), 0);
|
||||
Ok(())
|
||||
async fn streaming_client_does_not_retry_terminal_auth_errors() {
|
||||
for error in [
|
||||
AuthError::Build("invalid auth configuration".into()),
|
||||
AuthError::Policy(codex_http_client::NetworkPolicyDenied::Revoked),
|
||||
] {
|
||||
let state = RecordingState::default();
|
||||
let transport = RecordingTransport::new(state.clone());
|
||||
let auth = FailsOnceAuth::new(error);
|
||||
let mut provider = provider("openai");
|
||||
provider.retry.max_attempts = 2;
|
||||
let client = ResponsesClient::new(transport, provider, Arc::new(auth.clone()));
|
||||
let error = client
|
||||
.stream(
|
||||
serde_json::json!({ "model": "gpt-test" }),
|
||||
HeaderMap::new(),
|
||||
Compression::None,
|
||||
/*turn_state*/ None,
|
||||
)
|
||||
.await
|
||||
.err()
|
||||
.expect("terminal auth errors must fail without retry");
|
||||
match (error, auth.error.as_ref()) {
|
||||
(ApiError::Transport(TransportError::Build(actual)), AuthError::Build(expected)) => {
|
||||
assert_eq!(&actual, expected);
|
||||
}
|
||||
(ApiError::Transport(TransportError::Policy(actual)), AuthError::Policy(expected)) => {
|
||||
assert_eq!(&actual, expected);
|
||||
}
|
||||
(actual, expected) => panic!("expected {expected:?}, got {actual:?}"),
|
||||
}
|
||||
assert_eq!(auth.attempts(), 1);
|
||||
assert!(state.take_stream_requests().is_empty());
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -688,3 +688,84 @@ async fn azure_store_sends_ids_and_headers() -> Result<()> {
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
struct RejectedHeaderAuth;
|
||||
|
||||
impl AuthProvider for RejectedHeaderAuth {
|
||||
fn add_auth_headers(&self, headers: &mut HeaderMap) -> Result<(), AuthError> {
|
||||
headers.insert(
|
||||
http::header::AUTHORIZATION,
|
||||
HeaderValue::from_static("Bearer partial-auth"),
|
||||
);
|
||||
Err(AuthError::Build("header authentication failed".to_string()))
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn header_auth_error_prevents_sending_partial_authentication() {
|
||||
let state = RecordingState::default();
|
||||
let client = ResponsesClient::new(
|
||||
RecordingTransport::new(state.clone()),
|
||||
provider("openai"),
|
||||
Arc::new(RejectedHeaderAuth),
|
||||
);
|
||||
let result = client
|
||||
.stream(
|
||||
serde_json::json!({"echo": true}),
|
||||
HeaderMap::new(),
|
||||
Compression::None,
|
||||
/*turn_state*/ None,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert!(matches!(
|
||||
result,
|
||||
Err(ApiError::Transport(TransportError::Build(_)))
|
||||
));
|
||||
assert!(state.take_stream_requests().is_empty());
|
||||
}
|
||||
|
||||
struct AsyncHeaderAuth;
|
||||
|
||||
impl AuthProvider for AsyncHeaderAuth {
|
||||
fn add_auth_headers(&self, _headers: &mut HeaderMap) -> Result<(), codex_api::AuthError> {
|
||||
Err(codex_api::AuthError::Build(
|
||||
"authentication requires async header resolution".to_string(),
|
||||
))
|
||||
}
|
||||
|
||||
fn resolve_auth_headers(&self) -> codex_api::AuthHeadersFuture<'_> {
|
||||
Box::pin(async {
|
||||
Ok(HeaderMap::from_iter([(
|
||||
http::header::AUTHORIZATION,
|
||||
HeaderValue::from_static("Bearer async-auth"),
|
||||
)]))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn async_auth_provider_rejects_sync_resolution_and_authenticates_requests() -> Result<()> {
|
||||
assert!(AsyncHeaderAuth.to_auth_headers().is_err());
|
||||
let state = RecordingState::default();
|
||||
let client = ResponsesClient::new(
|
||||
RecordingTransport::new(state.clone()),
|
||||
provider("openai"),
|
||||
Arc::new(AsyncHeaderAuth),
|
||||
);
|
||||
let _stream = client
|
||||
.stream(
|
||||
serde_json::json!({"echo": true}),
|
||||
HeaderMap::new(),
|
||||
Compression::None,
|
||||
/*turn_state*/ None,
|
||||
)
|
||||
.await?;
|
||||
let requests = state.take_stream_requests();
|
||||
assert_eq!(requests.len(), 1);
|
||||
assert_eq!(
|
||||
requests[0].headers.get(http::header::AUTHORIZATION),
|
||||
Some(&HeaderValue::from_static("Bearer async-auth"))
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -28,7 +28,9 @@ use wiremock::matchers::path;
|
||||
struct DummyAuth;
|
||||
|
||||
impl AuthProvider for DummyAuth {
|
||||
fn add_auth_headers(&self, _headers: &mut HeaderMap) {}
|
||||
fn add_auth_headers(&self, _headers: &mut HeaderMap) -> Result<(), codex_api::AuthError> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn provider(base_url: &str) -> Provider {
|
||||
|
||||
@@ -53,7 +53,9 @@ impl HttpTransport for FixtureSseTransport {
|
||||
struct NoAuth;
|
||||
|
||||
impl AuthProvider for NoAuth {
|
||||
fn add_auth_headers(&self, _headers: &mut HeaderMap) {}
|
||||
fn add_auth_headers(&self, _headers: &mut HeaderMap) -> Result<(), codex_api::AuthError> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn provider(name: &str) -> Provider {
|
||||
|
||||
Reference in New Issue
Block a user