mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
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:
committed by
GitHub
Unverified
parent
34800d717e
commit
51fdc35945
Generated
-1
@@ -2891,7 +2891,6 @@ dependencies = [
|
||||
"codex-protocol",
|
||||
"http 1.4.0",
|
||||
"pretty_assertions",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
@@ -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(),
|
||||
|
||||
Reference in New Issue
Block a user