chore: remove unused Bedrock auth lazy loading (#18948)

## Summary

The Bedrock Mantle SigV4 auth provider currently looks like it can
lazily load `AwsAuthContext`, but the provider is only constructed after
`resolve_auth_method` has already loaded that context. Because
`with_context` always pre-populates the `OnceCell`, the
`get_or_try_init` fallback is unused in normal operation and makes the
provider lifecycle harder to reason about.

This change removes that dead lazy-loading path and makes the actual
behavior explicit:

- `BedrockAuthMethod::AwsSdkAuth` carries only the resolved
`AwsAuthContext`.
- `BedrockMantleSigV4AuthProvider` stores the resolved context directly.
- request signing uses the stored context without going through
`OnceCell`.

The existing eager AWS auth resolution behavior is unchanged; this is a
simplification of the provider state, not a behavior change.

## Testing

- `cargo shear`
- `cargo test -p codex-model-provider`
- `just bazel-lock-check`
This commit is contained in:
Celia Chen
2026-04-21 22:01:22 -07:00
committed by GitHub
Unverified
parent 34800d717e
commit 51fdc35945
3 changed files with 12 additions and 35 deletions
-1
View File
@@ -2891,7 +2891,6 @@ dependencies = [
"codex-protocol",
"http 1.4.0",
"pretty_assertions",
"tokio",
]
[[package]]
-1
View File
@@ -21,7 +21,6 @@ codex-login = { workspace = true }
codex-model-provider-info = { workspace = true }
codex-protocol = { workspace = true }
http = { workspace = true }
tokio = { workspace = true, features = ["sync"] }
[dev-dependencies]
pretty_assertions = { workspace = true }
@@ -3,7 +3,6 @@ use std::sync::Arc;
use codex_api::AuthError;
use codex_api::AuthProvider;
use codex_api::SharedAuthProvider;
use codex_aws_auth::AwsAuthConfig;
use codex_aws_auth::AwsAuthContext;
use codex_aws_auth::AwsAuthError;
use codex_aws_auth::AwsRequestToSign;
@@ -14,7 +13,6 @@ use codex_model_provider_info::ModelProviderAwsAuthInfo;
use codex_protocol::error::CodexErr;
use codex_protocol::error::Result;
use http::HeaderMap;
use tokio::sync::OnceCell;
use crate::BearerAuthProvider;
@@ -25,14 +23,8 @@ const AWS_BEARER_TOKEN_BEDROCK_ENV_VAR: &str = "AWS_BEARER_TOKEN_BEDROCK";
const LEGACY_SESSION_ID_HEADER: &str = "session_id";
enum BedrockAuthMethod {
EnvBearerToken {
token: String,
region: String,
},
AwsSdkAuth {
config: AwsAuthConfig,
context: AwsAuthContext,
},
EnvBearerToken { token: String, region: String },
AwsSdkAuth { context: AwsAuthContext },
}
async fn resolve_auth_method(aws: &ModelProviderAwsAuthInfo) -> Result<BedrockAuthMethod> {
@@ -42,10 +34,10 @@ async fn resolve_auth_method(aws: &ModelProviderAwsAuthInfo) -> Result<BedrockAu
}
let config = aws_auth_config(aws);
let context = AwsAuthContext::load(config.clone())
let context = AwsAuthContext::load(config)
.await
.map_err(aws_auth_error_to_codex_error)?;
Ok(BedrockAuthMethod::AwsSdkAuth { config, context })
Ok(BedrockAuthMethod::AwsSdkAuth { context })
}
pub(super) async fn resolve_provider_auth(
@@ -57,9 +49,9 @@ pub(super) async fn resolve_provider_auth(
account_id: None,
is_fedramp_account: false,
})),
BedrockAuthMethod::AwsSdkAuth { config, context } => Ok(Arc::new(
BedrockMantleSigV4AuthProvider::with_context(config, context),
)),
BedrockAuthMethod::AwsSdkAuth { context } => {
Ok(Arc::new(BedrockMantleSigV4AuthProvider::new(context)))
}
}
}
@@ -109,25 +101,12 @@ fn remove_headers_not_preserved_by_bedrock_mantle(headers: &mut HeaderMap) {
/// AWS SigV4 auth provider for Bedrock Mantle OpenAI-compatible requests.
#[derive(Debug)]
struct BedrockMantleSigV4AuthProvider {
config: AwsAuthConfig,
context: OnceCell<AwsAuthContext>,
context: AwsAuthContext,
}
impl BedrockMantleSigV4AuthProvider {
fn with_context(config: AwsAuthConfig, context: AwsAuthContext) -> Self {
let cell = OnceCell::new();
let _ = cell.set(context);
Self {
config,
context: cell,
}
}
async fn context(&self) -> std::result::Result<&AwsAuthContext, AuthError> {
self.context
.get_or_try_init(|| AwsAuthContext::load(self.config.clone()))
.await
.map_err(aws_auth_error_to_auth_error)
fn new(context: AwsAuthContext) -> Self {
Self { context }
}
}
@@ -139,8 +118,8 @@ impl AuthProvider for BedrockMantleSigV4AuthProvider {
let mut request = request;
remove_headers_not_preserved_by_bedrock_mantle(&mut request.headers);
let prepared = request.prepare_body_for_send().map_err(AuthError::Build)?;
let context = self.context().await?;
let signed = context
let signed = self
.context
.sign(AwsRequestToSign {
method: request.method.clone(),
url: request.url.clone(),