openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-30 19:18:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

Centralize provider request preparation auth checks (#58121)

Browse Source 
## What changed

Move the model client's request-preparation loop into the shared `prepare_provider_request` helper and add `RequestAuthChange` to classify auth state changes. For managed auth, retry preparation when credentials change and reject account changes, including when the provider starts or stops exposing an auth manager during setup.

Capture the auth owner and credential generations together and pass them into the preparation callback so the resulting client setup records the state used for that attempt.

## Testing

Add a test verifying that an explicit provider bearer token survives a workspace switch during setup without retrying preparation.

GitOrigin-RevId: ed8710d1c344b48ebf6182e25e6269214dd2d3fe
This commit is contained in:
sgarg-oai
2026-10-29 16:39:38 +00:00
committed by copyberry
parent 6174d05a6a
commit 53d3d7f64e
7 changed files with 184 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/client.rs
+41 -56
View File
@@ -1278,63 +1278,48 @@ impl ModelClient {
routing: ClientRouting,
provider: &SharedModelProvider,
) -> Result<CurrentClientSetup> {
// Capture before resolving credentials so an account switch during setup cannot label
// an old connection with the new owner's revision.
let auth_owner_generation = self.auth_owner_generation();
let auth_changes = provider
.auth_manager()
.map(|manager| manager.auth_change_receiver());
loop {
let revision = auth_changes.as_ref().map(|changes| *changes.borrow());
let auth = provider.auth().await;
let (api_provider, redirect_policy) = match routing {
ClientRouting::Workspace => {
let resolved = provider
.responses_api_provider(&self.state.workspace_routing)
.await?;
(resolved.provider, resolved.redirect_policy)
}
ClientRouting::ConfiguredProvider => (
provider.api_provider().await?,
ClientRedirectPolicy::Default,
),
};
let resolved_auth = provider
.api_auth_for_scope(ProviderAuthScope {
agent_identity_policy: self.agent_identity_policy,
session_source: self.state.session_source.clone(),
agent_identity_session_fallback: self
.state
.agent_identity_session_fallback
.clone(),
let routing = &routing;
codex_model_provider::prepare_provider_request(
provider.as_ref(),
provider.info().auth.is_none(),
|state| async move {
let auth = provider.auth().await;
let (api_provider, redirect_policy) = match routing {
ClientRouting::Workspace => {
let resolved = provider
.responses_api_provider(&self.state.workspace_routing)
.await?;
(resolved.provider, resolved.redirect_policy)
}
ClientRouting::ConfiguredProvider => (
provider.api_provider().await?,
ClientRedirectPolicy::Default,
),
};
let resolved_auth = provider
.api_auth_for_scope(ProviderAuthScope {
agent_identity_policy: self.agent_identity_policy,
session_source: self.state.session_source.clone(),
agent_identity_session_fallback: self
.state
.agent_identity_session_fallback
.clone(),
})
.await?;
Ok(CurrentClientSetup {
auth,
auth_owner_generation: state.map(|state| state.owner_generation),
auth_revision: state.map(|state| state.generation),
api_provider,
http_client_factory: self.http_client_factory.clone(),
redirect_policy,
api_auth: resolved_auth.auth,
agent_identity_telemetry: resolved_auth.agent_identity_telemetry,
runtime_id: None,
})
.await?;
// Command-backed bearer refreshes cannot change workspace routing. Keep the captured
// revisions so a refresh during setup still invalidates cached WebSocket state.
if provider.info().auth.is_none() {
if self.auth_owner_generation() != auth_owner_generation {
return Err(std::io::Error::other(
"account changed while preparing model request",
)
.into());
}
// Rebuild the entire bundle after a credential-only refresh, including routing.
if auth_changes.as_ref().map(|changes| *changes.borrow()) != revision {
continue;
}
}
return Ok(CurrentClientSetup {
auth,
auth_owner_generation,
auth_revision: revision,
api_provider,
http_client_factory: self.http_client_factory.clone(),
redirect_policy,
api_auth: resolved_auth.auth,
agent_identity_telemetry: resolved_auth.agent_identity_telemetry,
runtime_id: None,
});
}
},
)
.await
}
fn responses_headers(&self, auth: Option<&CodexAuth>, model: &str) -> ApiHeaderMap {
codex-rs/login/src/auth/change_state.rs
+21
View File
@@ -19,6 +19,27 @@ pub struct AuthChangeState {
pub login_generation: u64,
}
/// Whether request preparation can use the credentials it selected.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum RequestAuthChange {
Unchanged,
CredentialsChanged,
OwnerChanged,
}
impl AuthChangeState {
/// Compare a completed attempt with its original owner and starting credential revision.
pub fn compare_for_request(&self, expected_owner: u64, prepared: Self) -> RequestAuthChange {
if self.owner_generation != expected_owner {
RequestAuthChange::OwnerChanged
} else if self.generation != prepared.generation {
RequestAuthChange::CredentialsChanged
} else {
RequestAuthChange::Unchanged
}
}
}
pub(super) fn same_owner(previous: Option<&CodexAuth>, current: Option<&CodexAuth>) -> bool {
let (Some(previous), Some(current)) = (previous, current) else {
return false;
codex-rs/login/src/auth/mod.rs
+1
View File
@@ -24,6 +24,7 @@ pub use bedrock_access_keys::login_with_bedrock_access_keys;
pub use bedrock_api_key::BedrockApiKeyAuth;
pub use bedrock_api_key::login_with_bedrock_api_key;
pub use change_state::AuthChangeState;
pub use change_state::RequestAuthChange;
pub use error::RefreshTokenFailedError;
pub use error::RefreshTokenFailedReason;
pub use manager::*;
codex-rs/login/src/lib.rs
+1
View File
@@ -58,6 +58,7 @@ pub use auth::OPENAI_API_KEY_ENV_VAR;
pub use auth::REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR;
pub use auth::REVOKE_TOKEN_URL_OVERRIDE_ENV_VAR;
pub use auth::RefreshTokenError;
pub use auth::RequestAuthChange;
pub use auth::UnauthorizedRecovery;
pub use auth::default_client;
pub use auth::enforce_login_restrictions;
codex-rs/model-provider/src/lib.rs
+2
View File
@@ -8,6 +8,7 @@ mod gcp;
mod models_endpoint;
mod models_identity;
mod provider;
mod request_preparation;
mod runtime_auth;
mod shared_state;
pub mod test_support;
@@ -50,6 +51,7 @@ pub use provider::create_model_provider;
pub use provider::provider_uses_first_party_auth_path;
pub use provider::resolve_api_provider;
pub use request_preparation::prepare_provider_request;
pub use runtime_auth::select_runtime_auth;
pub use runtime_auth::validate_runtime_destination;
pub use runtime_auth::validate_runtime_provider;
codex-rs/model-provider/src/request_preparation.rs
+56
View File
@@ -0,0 +1,56 @@
//! Keeps request preparation bound to the credential revision it started with.
use std::future::Future;
use std::io;
use codex_login::AuthChangeState;
use codex_login::RequestAuthChange;
use codex_protocol::error::Result;
use crate::ModelProvider;
/// Prepare a request without dispatching it. When the request uses managed auth,
/// repeat preparation after a credential refresh and fail if the owner changes.
pub async fn prepare_provider_request<T, Fut>(
provider: &dyn ModelProvider,
check_auth: bool,
mut prepare: impl FnMut(Option<AuthChangeState>) -> Fut,
) -> Result<T>
where
Fut: Future<Output = Result<T>>,
{
let auth_state = || {
provider
.auth_manager()
.map(|manager| *manager.auth_change_state_receiver().borrow())
};
let mut before = auth_state();
loop {
let prepared = prepare(before).await?;
if !check_auth {
return Ok(prepared);
}
// A managed provider can begin or cease exposing its auth manager during setup.
let after = auth_state();
let change = match (before, after) {
(Some(before), Some(after)) => {
after.compare_for_request(before.owner_generation, before)
}
(None, None) => RequestAuthChange::Unchanged,
_ => RequestAuthChange::OwnerChanged,
};
match change {
RequestAuthChange::Unchanged => return Ok(prepared),
RequestAuthChange::CredentialsChanged => before = after,
RequestAuthChange::OwnerChanged => {
return Err(
io::Error::other("account changed while preparing model request").into(),
);
}
}
}
}
#[cfg(test)]
#[path = "request_preparation_tests.rs"]
mod tests;
codex-rs/model-provider/src/request_preparation_tests.rs
+62
View File
@@ -0,0 +1,62 @@
use super::prepare_provider_request;
use crate::create_model_provider;
use crate::provider_uses_first_party_auth_path;
use base64::Engine;
use codex_login::AuthManager;
use codex_login::CodexAuth;
use codex_model_provider_info::ModelProviderInfo;
use pretty_assertions::assert_eq;
use serde_json::json;
use tempfile::TempDir;
#[tokio::test]
async fn explicit_provider_token_survives_a_workspace_switch_during_setup() {
let home = TempDir::new().unwrap();
let claims = json!({"https://api.openai.com/auth": {"chatgpt_user_id": "user-a"}});
let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(claims.to_string());
let token = format!("header.{payload}.signature");
let initial = CodexAuth::from_external_chatgpt_tokens(
&token,
"workspace-a",
/*chatgpt_plan_type*/ None,
)
.unwrap();
let manager = AuthManager::from_auth_for_testing_with_home(initial, home.path().to_path_buf());
let mut info = ModelProviderInfo::create_openai_provider(/*base_url*/ None);
info.experimental_bearer_token = Some("provider-token".into());
let provider = create_model_provider(info, Some(manager.clone()));
let mut attempts = 0;
let auth = prepare_provider_request(
provider.as_ref(),
provider_uses_first_party_auth_path(provider.info()),
|_| {
attempts += 1;
async {
let (_, auth) = provider.api_provider_and_auth().await?;
codex_login::auth::login_with_chatgpt_auth_tokens(
home.path(),
&token,
"workspace-b",
/*chatgpt_plan_type*/ None,
)?;
manager.reload().await;
Ok(auth)
}
},
)
.await
.unwrap();
assert_eq!(
manager.auth_cached().unwrap().get_account_id().as_deref(),
Some("workspace-b")
);
assert_eq!(attempts, 1);
let mut headers = http::HeaderMap::new();
auth.add_auth_headers(&mut headers).unwrap();
assert_eq!(
headers[http::header::AUTHORIZATION],
"Bearer provider-token"
);
}