Centralize provider request preparation auth checks (#58121)

## 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
+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 {