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

core: unify request auth and 401 recovery plumbing

Browse Source
This commit is contained in:
Michael Bolin
2026-03-30 13:49:13 -07:00
parent 258ba436f1
commit d4cf2faae8
5 changed files with 282 additions and 156 deletions
Download Patch File Download Diff File Expand all files Collapse all files

218
codex-rs/core/src/client.rs
View File

@@ -31,9 +31,7 @@ use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use crate::api_bridge::CoreAuthProvider;
use crate::api_bridge::auth_provider_from_auth;
use crate::api_bridge::map_api_error;
use crate::auth::UnauthorizedRecovery;
use crate::auth_env_telemetry::AuthEnvTelemetry;
use crate::auth_env_telemetry::collect_auth_env_telemetry;
use codex_api::CompactClient as ApiCompactClient;
@@ -104,6 +102,12 @@ use crate::error::Result;
use crate::flags::CODEX_RS_SSE_FIXTURE;
use crate::model_provider_info::ModelProviderInfo;
use crate::model_provider_info::WireApi;
use crate::request_auth::RequestUnauthorizedRecovery;
use crate::request_auth::ResolvedRequestAuth;
use crate::request_auth::UnauthorizedRecoveryError;
use crate::request_auth::UnauthorizedRecoveryExecution;
use crate::request_auth::UnauthorizedRecoveryOutcome;
use crate::request_auth::resolve_request_auth;
use crate::response_debug_context::extract_response_debug_context;
use crate::response_debug_context::extract_response_debug_context_from_api_error;
use crate::response_debug_context::telemetry_api_error_message;
@@ -144,16 +148,6 @@ struct ModelClientState {
cached_websocket_session: StdMutex<WebsocketSession>,
}
/// Resolved API client setup for a single request attempt.
///
/// Keeping this as a single bundle ensures prewarm and normal request paths
/// share the same auth/provider setup flow.
struct CurrentClientSetup {
auth: Option<CodexAuth>,
api_provider: codex_api::Provider,
api_auth: CoreAuthProvider,
}
#[derive(Clone, Copy)]
struct RequestRouteTelemetry {
endpoint: &'static str,
@@ -523,21 +517,8 @@ impl ModelClient {
///
/// This centralizes setup used by both prewarm and normal request paths so they stay in
/// lockstep when auth/provider resolution changes.
async fn current_client_setup(&self) -> Result<CurrentClientSetup> {
let auth = match self.state.auth_manager.as_ref() {
Some(manager) => manager.auth().await,
None => None,
};
let api_provider = self
.state
.provider
.to_api_provider(auth.as_ref().map(CodexAuth::auth_mode))?;
let api_auth = auth_provider_from_auth(auth.clone(), &self.state.provider)?;
Ok(CurrentClientSetup {
auth,
api_provider,
api_auth,
})
async fn current_client_setup(&self) -> Result<ResolvedRequestAuth> {
resolve_request_auth(self.state.auth_manager.as_ref(), &self.state.provider).await
}
/// Opens a websocket connection using the same header and telemetry wiring as normal turns.
@@ -1016,11 +997,11 @@ impl ModelClientSession {
return Ok(stream);
}
let auth_manager = self.client.state.auth_manager.clone();
let mut auth_recovery = auth_manager
.as_ref()
.map(super::auth::AuthManager::unauthorized_recovery);
let mut unauthorized_recovery =
RequestUnauthorizedRecovery::new(self.client.state.auth_manager.as_ref());
let mut pending_retry = PendingUnauthorizedRetry::default();
// Only loop after a successful auth-recovery step. Each retry must rebuild
// the client and request headers before issuing the same streaming request again.
loop {
let client_setup = self.client.current_client_setup().await?;
let transport = ReqwestTransport::new(build_reqwest_client());
@@ -1065,7 +1046,7 @@ impl ModelClientSession {
pending_retry = PendingUnauthorizedRetry::from_recovery(
handle_unauthorized(
unauthorized_transport,
&mut auth_recovery,
&mut unauthorized_recovery,
session_telemetry,
)
.await?,
@@ -1104,12 +1085,11 @@ impl ModelClientSession {
warmup: bool,
request_trace: Option<W3cTraceContext>,
) -> Result<WebsocketStreamOutcome> {
let auth_manager = self.client.state.auth_manager.clone();
let mut auth_recovery = auth_manager
.as_ref()
.map(super::auth::AuthManager::unauthorized_recovery);
let mut unauthorized_recovery =
RequestUnauthorizedRecovery::new(self.client.state.auth_manager.as_ref());
let mut pending_retry = PendingUnauthorizedRetry::default();
// Only loop after a successful auth-recovery step. WebSocket auth is attached
// during connect, so a recovered token requires a fresh connection attempt.
loop {
let client_setup = self.client.current_client_setup().await?;
let request_auth_context = AuthRequestTelemetryContext::new(
@@ -1162,14 +1142,13 @@ impl ModelClientSession {
Err(ApiError::Transport(
unauthorized_transport @ TransportError::Http { status, .. },
)) if status == StatusCode::UNAUTHORIZED => {
pending_retry = PendingUnauthorizedRetry::from_recovery(
handle_unauthorized(
unauthorized_transport,
&mut auth_recovery,
session_telemetry,
)
.await?,
);
let recovery = handle_unauthorized(
unauthorized_transport,
&mut unauthorized_recovery,
session_telemetry,
)
.await?;
pending_retry = PendingUnauthorizedRetry::from_recovery(recovery);
continue;
}
Err(err) => return Err(map_api_error(err)),
@@ -1484,16 +1463,6 @@ where
(ResponseStream { rx_event }, rx_last_response)
}
/// Handles a 401 response by optionally refreshing ChatGPT tokens once.
///
/// When refresh succeeds, the caller should retry the API call; otherwise
/// the mapped `CodexErr` is returned to the caller.
#[derive(Clone, Copy, Debug)]
struct UnauthorizedRecoveryExecution {
mode: &'static str,
phase: &'static str,
}
#[derive(Clone, Copy, Debug, Default)]
struct PendingUnauthorizedRetry {
retry_after_unauthorized: bool,
@@ -1551,45 +1520,71 @@ struct WebsocketConnectParams<'a> {
request_route_telemetry: RequestRouteTelemetry,
}
/// Handles a `401 Unauthorized` from the transport used by the request loops above.
///
/// The helper centralizes three coupled concerns:
/// - ask `RequestUnauthorizedRecovery` whether another recovery step is available
/// - record the matching telemetry / feedback tags for the outcome of that step
/// - return the recovery execution to the caller so it can rebuild auth state and retry,
/// or map the failure into the `CodexErr` that should terminate the loop
async fn handle_unauthorized(
transport: TransportError,
auth_recovery: &mut Option<UnauthorizedRecovery>,
unauthorized_recovery: &mut RequestUnauthorizedRecovery,
session_telemetry: &SessionTelemetry,
) -> Result<UnauthorizedRecoveryExecution> {
let debug = extract_response_debug_context(&transport);
if let Some(recovery) = auth_recovery
&& recovery.has_next()
{
let mode = recovery.mode_name();
let phase = recovery.step_name();
return match recovery.next().await {
Ok(step_result) => {
match unauthorized_recovery.next().await {
Ok(UnauthorizedRecoveryOutcome::Recovered(recovery)) => {
session_telemetry.record_auth_recovery(
recovery.mode,
recovery.phase,
"recovery_succeeded",
debug.request_id.as_deref(),
debug.cf_ray.as_deref(),
debug.auth_error.as_deref(),
debug.auth_error_code.as_deref(),
/*recovery_reason*/ None,
recovery.auth_state_changed,
);
emit_feedback_auth_recovery_tags(
recovery.mode,
recovery.phase,
"recovery_succeeded",
debug.request_id.as_deref(),
debug.cf_ray.as_deref(),
debug.auth_error.as_deref(),
debug.auth_error_code.as_deref(),
);
Ok(recovery)
}
Ok(UnauthorizedRecoveryOutcome::Unavailable(unavailable)) => {
session_telemetry.record_auth_recovery(
unavailable.mode,
unavailable.phase,
"recovery_not_run",
debug.request_id.as_deref(),
debug.cf_ray.as_deref(),
debug.auth_error.as_deref(),
debug.auth_error_code.as_deref(),
Some(unavailable.reason),
/*auth_state_changed*/ None,
);
emit_feedback_auth_recovery_tags(
unavailable.mode,
unavailable.phase,
"recovery_not_run",
debug.request_id.as_deref(),
debug.cf_ray.as_deref(),
debug.auth_error.as_deref(),
debug.auth_error_code.as_deref(),
);
Err(map_api_error(ApiError::Transport(transport)))
}
Err(UnauthorizedRecoveryError::Chatgpt { execution, error }) => match error {
RefreshTokenError::Permanent(failed) => {
session_telemetry.record_auth_recovery(
mode,
phase,
"recovery_succeeded",
debug.request_id.as_deref(),
debug.cf_ray.as_deref(),
debug.auth_error.as_deref(),
debug.auth_error_code.as_deref(),
/*recovery_reason*/ None,
step_result.auth_state_changed(),
);
emit_feedback_auth_recovery_tags(
mode,
phase,
"recovery_succeeded",
debug.request_id.as_deref(),
debug.cf_ray.as_deref(),
debug.auth_error.as_deref(),
debug.auth_error_code.as_deref(),
);
Ok(UnauthorizedRecoveryExecution { mode, phase })
}
Err(RefreshTokenError::Permanent(failed)) => {
session_telemetry.record_auth_recovery(
mode,
phase,
execution.mode,
execution.phase,
"recovery_failed_permanent",
debug.request_id.as_deref(),
debug.cf_ray.as_deref(),
@@ -1599,8 +1594,8 @@ async fn handle_unauthorized(
/*auth_state_changed*/ None,
);
emit_feedback_auth_recovery_tags(
mode,
phase,
execution.mode,
execution.phase,
"recovery_failed_permanent",
debug.request_id.as_deref(),
debug.cf_ray.as_deref(),
@@ -1609,10 +1604,10 @@ async fn handle_unauthorized(
);
Err(CodexErr::RefreshTokenFailed(failed))
}
Err(RefreshTokenError::Transient(other)) => {
RefreshTokenError::Transient(other) => {
session_telemetry.record_auth_recovery(
mode,
phase,
execution.mode,
execution.phase,
"recovery_failed_transient",
debug.request_id.as_deref(),
debug.cf_ray.as_deref(),
@@ -1622,8 +1617,8 @@ async fn handle_unauthorized(
/*auth_state_changed*/ None,
);
emit_feedback_auth_recovery_tags(
mode,
phase,
execution.mode,
execution.phase,
"recovery_failed_transient",
debug.request_id.as_deref(),
debug.cf_ray.as_deref(),
@@ -1632,39 +1627,8 @@ async fn handle_unauthorized(
);
Err(CodexErr::Io(other))
}
};
},
}
let (mode, phase, recovery_reason) = match auth_recovery.as_ref() {
Some(recovery) => (
recovery.mode_name(),
recovery.step_name(),
Some(recovery.unavailable_reason()),
),
None => ("none", "none", Some("auth_manager_missing")),
};
session_telemetry.record_auth_recovery(
mode,
phase,
"recovery_not_run",
debug.request_id.as_deref(),
debug.cf_ray.as_deref(),
debug.auth_error.as_deref(),
debug.auth_error_code.as_deref(),
recovery_reason,
/*auth_state_changed*/ None,
);
emit_feedback_auth_recovery_tags(
mode,
phase,
"recovery_not_run",
debug.request_id.as_deref(),
debug.cf_ray.as_deref(),
debug.auth_error.as_deref(),
debug.auth_error_code.as_deref(),
);
Err(map_api_error(ApiError::Transport(transport)))
}
fn api_error_http_status(error: &ApiError) -> Option<u16> {

1
codex-rs/core/src/client_tests.rs
View File

@@ -110,6 +110,7 @@ fn auth_request_telemetry_context_tracks_attached_auth_and_retry_phase() {
PendingUnauthorizedRetry::from_recovery(UnauthorizedRecoveryExecution {
mode: "managed",
phase: "refresh_token",
auth_state_changed: None,
}),
);

1
codex-rs/core/src/lib.rs
View File

@@ -52,6 +52,7 @@ pub mod models_manager;
mod network_policy_decision;
pub mod network_proxy_loader;
mod original_image_detail;
mod request_auth;
pub use mcp_connection_manager::MCP_SANDBOX_STATE_CAPABILITY;
pub use mcp_connection_manager::MCP_SANDBOX_STATE_METHOD;
pub use mcp_connection_manager::SandboxState;

81
codex-rs/core/src/models_manager/manager.rs
View File

@@ -1,9 +1,7 @@
use super::cache::ModelsCacheManager;
use crate::api_bridge::auth_provider_from_auth;
use crate::api_bridge::map_api_error;
use crate::auth::AuthManager;
use crate::auth::AuthMode;
use crate::auth::CodexAuth;
use crate::auth_env_telemetry::AuthEnvTelemetry;
use crate::auth_env_telemetry::collect_auth_env_telemetry;
use crate::config::Config;
@@ -14,6 +12,9 @@ use crate::model_provider_info::ModelProviderInfo;
use crate::models_manager::collaboration_mode_presets::CollaborationModesConfig;
use crate::models_manager::collaboration_mode_presets::builtin_collaboration_mode_presets;
use crate::models_manager::model_info;
use crate::request_auth::RequestUnauthorizedRecovery;
use crate::request_auth::UnauthorizedRecoveryOutcome;
use crate::request_auth::resolve_request_auth;
use crate::response_debug_context::extract_response_debug_context;
use crate::response_debug_context::telemetry_transport_error_message;
use crate::util::FeedbackRequestTags;
@@ -22,6 +23,7 @@ use codex_api::ModelsClient;
use codex_api::RequestTelemetry;
use codex_api::ReqwestTransport;
use codex_api::TransportError;
use codex_api::error::ApiError;
use codex_otel::TelemetryAuthMode;
use codex_protocol::config_types::CollaborationModeMask;
use codex_protocol::openai_models::ModelInfo;
@@ -431,39 +433,60 @@ impl ModelsManager {
async fn fetch_and_update_models(&self) -> CoreResult<()> {
let _timer =
codex_otel::start_global_timer("codex.remote_models.fetch_update.duration_ms", &[]);
let auth = self.auth_manager.auth().await;
let auth_mode = auth.as_ref().map(CodexAuth::auth_mode);
let api_provider = self.provider.to_api_provider(auth_mode)?;
let api_auth = auth_provider_from_auth(auth.clone(), &self.provider)?;
let auth_env = collect_auth_env_telemetry(
&self.provider,
self.auth_manager.codex_api_key_env_enabled(),
);
let transport = ReqwestTransport::new(build_reqwest_client());
let request_telemetry: Arc<dyn RequestTelemetry> = Arc::new(ModelsRequestTelemetry {
auth_mode: auth_mode.map(|mode| TelemetryAuthMode::from(mode).to_string()),
auth_header_attached: api_auth.auth_header_attached(),
auth_header_name: api_auth.auth_header_name(),
auth_env,
});
let client = ModelsClient::new(transport, api_provider, api_auth)
.with_telemetry(Some(request_telemetry));
let client_version = crate::models_manager::client_version_to_whole();
let (models, etag) = timeout(
MODELS_REFRESH_TIMEOUT,
client.list_models(&client_version, HeaderMap::new()),
)
.await
.map_err(|_| CodexErr::Timeout)?
.map_err(map_api_error)?;
let mut unauthorized_recovery = RequestUnauthorizedRecovery::new(Some(&self.auth_manager));
self.apply_remote_models(models.clone()).await;
*self.etag.write().await = etag.clone();
self.cache_manager
.persist_cache(&models, etag, client_version)
.await;
Ok(())
// Only loop after a successful auth-recovery step so `/models` retries with
// the same freshly resolved auth state as normal request paths.
loop {
let request_auth =
resolve_request_auth(Some(&self.auth_manager), &self.provider).await?;
let transport = ReqwestTransport::new(build_reqwest_client());
let request_telemetry: Arc<dyn RequestTelemetry> = Arc::new(ModelsRequestTelemetry {
auth_mode: request_auth
.auth_mode
.map(|mode| TelemetryAuthMode::from(mode).to_string()),
auth_header_attached: request_auth.api_auth.auth_header_attached(),
auth_header_name: request_auth.api_auth.auth_header_name(),
auth_env: auth_env.clone(),
});
let client =
ModelsClient::new(transport, request_auth.api_provider, request_auth.api_auth)
.with_telemetry(Some(request_telemetry));
match timeout(
MODELS_REFRESH_TIMEOUT,
client.list_models(&client_version, HeaderMap::new()),
)
.await
.map_err(|_| CodexErr::Timeout)?
{
Ok((models, etag)) => {
self.apply_remote_models(models.clone()).await;
*self.etag.write().await = etag.clone();
self.cache_manager
.persist_cache(&models, etag, client_version)
.await;
return Ok(());
}
Err(ApiError::Transport(
unauthorized_transport @ TransportError::Http { status, .. },
)) if status == http::StatusCode::UNAUTHORIZED => {
match unauthorized_recovery.next().await {
Ok(UnauthorizedRecoveryOutcome::Recovered(_)) => continue,
Ok(UnauthorizedRecoveryOutcome::Unavailable(_)) => {
return Err(map_api_error(ApiError::Transport(unauthorized_transport)));
}
Err(error) => return Err(error.into_codex_err()),
}
}
Err(err) => return Err(map_api_error(err)),
}
}
}
async fn get_etag(&self) -> Option<String> {

137
codex-rs/core/src/request_auth.rs Normal file
View File

@@ -0,0 +1,137 @@
use std::sync::Arc;
use crate::api_bridge::CoreAuthProvider;
use crate::api_bridge::auth_provider_from_auth;
use crate::auth::AuthManager;
use crate::auth::AuthMode;
use crate::auth::CodexAuth;
use crate::auth::RefreshTokenError;
use crate::auth::UnauthorizedRecovery;
use crate::error::CodexErr;
use crate::error::Result;
use crate::model_provider_info::ModelProviderInfo;
#[derive(Clone)]
pub(crate) struct ResolvedRequestAuth {
pub(crate) auth: Option<CodexAuth>,
pub(crate) auth_mode: Option<AuthMode>,
pub(crate) api_provider: codex_api::Provider,
pub(crate) api_auth: CoreAuthProvider,
}
pub(crate) async fn resolve_request_auth(
auth_manager: Option<&Arc<AuthManager>>,
provider: &ModelProviderInfo,
) -> Result<ResolvedRequestAuth> {
let auth = match auth_manager {
Some(manager) => manager.auth().await,
None => None,
};
let auth_mode = auth.as_ref().map(CodexAuth::auth_mode);
let api_provider = provider.to_api_provider(auth_mode)?;
let api_auth = auth_provider_from_auth(auth.clone(), provider)?;
Ok(ResolvedRequestAuth {
auth,
auth_mode,
api_provider,
api_auth,
})
}
#[derive(Clone, Copy, Debug)]
pub(crate) struct UnauthorizedRecoveryExecution {
pub(crate) mode: &'static str,
pub(crate) phase: &'static str,
pub(crate) auth_state_changed: Option<bool>,
}
#[derive(Clone, Copy, Debug)]
pub(crate) struct UnauthorizedRecoveryUnavailable {
pub(crate) mode: &'static str,
pub(crate) phase: &'static str,
pub(crate) reason: &'static str,
}
#[derive(Debug)]
pub(crate) enum UnauthorizedRecoveryOutcome {
Recovered(UnauthorizedRecoveryExecution),
Unavailable(UnauthorizedRecoveryUnavailable),
}
#[derive(Debug)]
pub(crate) enum UnauthorizedRecoveryError {
Chatgpt {
execution: UnauthorizedRecoveryExecution,
error: RefreshTokenError,
},
}
impl UnauthorizedRecoveryError {
pub(crate) fn into_codex_err(self) -> CodexErr {
match self {
Self::Chatgpt { error, .. } => match error {
RefreshTokenError::Permanent(failed) => CodexErr::RefreshTokenFailed(failed),
RefreshTokenError::Transient(error) => CodexErr::Io(error),
},
}
}
}
/// Tracks which auth-recovery steps are still available after a request fails with `401`.
///
/// The request loops reuse the same instance across retries so this type can enforce the
/// recovery ordering and make sure each one-shot step only runs once.
pub(crate) struct RequestUnauthorizedRecovery {
auth_recovery: Option<UnauthorizedRecovery>,
}
impl RequestUnauthorizedRecovery {
pub(crate) fn new(auth_manager: Option<&Arc<AuthManager>>) -> Self {
Self {
auth_recovery: auth_manager.map(AuthManager::unauthorized_recovery),
}
}
/// Runs the next available auth-recovery step, if any.
///
/// Provider auth refresh is attempted first because it is specific to the configured
/// model provider and should repair the currently attached bearer token before falling
/// back to ChatGPT-managed recovery. Once all recovery steps have been consumed, this
/// returns `UnauthorizedRecoveryOutcome::Unavailable` so the caller can stop retrying.
pub(crate) async fn next(
&mut self,
) -> std::result::Result<UnauthorizedRecoveryOutcome, UnauthorizedRecoveryError> {
if let Some(recovery) = self.auth_recovery.as_mut()
&& recovery.has_next()
{
let execution = UnauthorizedRecoveryExecution {
mode: recovery.mode_name(),
phase: recovery.step_name(),
auth_state_changed: None,
};
return match recovery.next().await {
Ok(step_result) => Ok(UnauthorizedRecoveryOutcome::Recovered(
UnauthorizedRecoveryExecution {
auth_state_changed: step_result.auth_state_changed(),
..execution
},
)),
Err(error) => Err(UnauthorizedRecoveryError::Chatgpt { execution, error }),
};
}
let unavailable = match self.auth_recovery.as_ref() {
Some(recovery) => UnauthorizedRecoveryUnavailable {
mode: recovery.mode_name(),
phase: recovery.step_name(),
reason: recovery.unavailable_reason(),
},
None => UnauthorizedRecoveryUnavailable {
mode: "none",
phase: "none",
reason: "auth_manager_missing",
},
};
Ok(UnauthorizedRecoveryOutcome::Unavailable(unavailable))
}
}