diff --git a/codex-rs/model-provider/src/amazon_bedrock/auth.rs b/codex-rs/model-provider/src/amazon_bedrock/auth.rs index ecfd2dd53..841b45bb4 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/auth.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/auth.rs @@ -9,6 +9,7 @@ use codex_aws_auth::AwsRequestToSign; use codex_client::Request; use codex_client::RequestBody; use codex_client::RequestCompression; +use codex_login::auth::BedrockApiKeyAuth; use codex_model_provider_info::ModelProviderAwsAuthInfo; use codex_protocol::error::CodexErr; use codex_protocol::error::Result; @@ -24,13 +25,22 @@ const AWS_REGION_ENV_VAR: &str = "AWS_REGION"; const AWS_DEFAULT_REGION_ENV_VAR: &str = "AWS_DEFAULT_REGION"; pub(super) enum BedrockAuthMethod { + ManagedBearerToken { token: String, region: String }, EnvBearerToken { token: String, region: String }, AwsSdkAuth { context: AwsAuthContext }, } pub(super) async fn resolve_auth_method( + managed_auth: Option<&BedrockApiKeyAuth>, aws: &ModelProviderAwsAuthInfo, ) -> Result { + if let Some(managed_auth) = managed_auth { + return Ok(BedrockAuthMethod::ManagedBearerToken { + token: managed_auth.api_key.clone(), + region: managed_auth.region.clone(), + }); + } + if let Some(token) = non_empty_env_var_from(AWS_BEARER_TOKEN_BEDROCK_ENV_VAR, std::env::var) { let region = bearer_token_region(aws, std::env::var)?; return Ok(BedrockAuthMethod::EnvBearerToken { token, region }); @@ -44,10 +54,12 @@ pub(super) async fn resolve_auth_method( } pub(super) async fn resolve_provider_auth( + managed_auth: Option<&BedrockApiKeyAuth>, aws: &ModelProviderAwsAuthInfo, ) -> Result { - match resolve_auth_method(aws).await? { - BedrockAuthMethod::EnvBearerToken { token, .. } => Ok(Arc::new(BearerAuthProvider { + match resolve_auth_method(managed_auth, aws).await? { + BedrockAuthMethod::ManagedBearerToken { token, .. } + | BedrockAuthMethod::EnvBearerToken { token, .. } => Ok(Arc::new(BearerAuthProvider { token: Some(token), account_id: None, is_fedramp_account: false, diff --git a/codex-rs/model-provider/src/amazon_bedrock/mantle.rs b/codex-rs/model-provider/src/amazon_bedrock/mantle.rs index 7881845e4..9727802d1 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/mantle.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/mantle.rs @@ -1,4 +1,5 @@ use codex_aws_auth::AwsAuthConfig; +use codex_login::auth::BedrockApiKeyAuth; use codex_model_provider_info::ModelProviderAwsAuthInfo; use codex_protocol::error::CodexErr; use codex_protocol::error::Result; @@ -48,14 +49,21 @@ pub(super) fn base_url(region: &str) -> Result { } } -pub(super) async fn runtime_base_url(aws: &ModelProviderAwsAuthInfo) -> Result { - let region = resolve_region(aws).await?; +pub(super) async fn runtime_base_url( + managed_auth: Option<&BedrockApiKeyAuth>, + aws: &ModelProviderAwsAuthInfo, +) -> Result { + let region = resolve_region(managed_auth, aws).await?; base_url(®ion) } -async fn resolve_region(aws: &ModelProviderAwsAuthInfo) -> Result { - match resolve_auth_method(aws).await? { - BedrockAuthMethod::EnvBearerToken { region, .. } => Ok(region), +async fn resolve_region( + managed_auth: Option<&BedrockApiKeyAuth>, + aws: &ModelProviderAwsAuthInfo, +) -> Result { + match resolve_auth_method(managed_auth, aws).await? { + BedrockAuthMethod::ManagedBearerToken { region, .. } + | BedrockAuthMethod::EnvBearerToken { region, .. } => Ok(region), BedrockAuthMethod::AwsSdkAuth { context } => Ok(context.region().to_string()), } } diff --git a/codex-rs/model-provider/src/amazon_bedrock/mod.rs b/codex-rs/model-provider/src/amazon_bedrock/mod.rs index 2a2af7043..872be005b 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/mod.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/mod.rs @@ -9,6 +9,7 @@ use codex_api::Provider; use codex_api::SharedAuthProvider; use codex_login::AuthManager; use codex_login::CodexAuth; +use codex_login::auth::BedrockApiKeyAuth; use codex_model_provider_info::AMAZON_BEDROCK_GPT_5_4_MODEL_ID; use codex_model_provider_info::ModelProviderAwsAuthInfo; use codex_model_provider_info::ModelProviderInfo; @@ -32,10 +33,14 @@ use mantle::runtime_base_url; pub(crate) struct AmazonBedrockModelProvider { pub(crate) info: ModelProviderInfo, pub(crate) aws: ModelProviderAwsAuthInfo, + auth_manager: Option>, } impl AmazonBedrockModelProvider { - pub(crate) fn new(provider_info: ModelProviderInfo) -> Self { + pub(crate) fn new( + provider_info: ModelProviderInfo, + auth_manager: Option>, + ) -> Self { let aws = provider_info .aws .clone() @@ -46,8 +51,23 @@ impl AmazonBedrockModelProvider { Self { info: provider_info, aws, + auth_manager, } } + + fn managed_auth(&self) -> Option { + self.auth_manager + .as_ref() + .and_then(|auth_manager| auth_manager.auth_cached()) + .and_then(|auth| match auth { + CodexAuth::BedrockApiKey(auth) => Some(auth), + CodexAuth::ApiKey(_) + | CodexAuth::Chatgpt(_) + | CodexAuth::ChatgptAuthTokens(_) + | CodexAuth::AgentIdentity(_) + | CodexAuth::PersonalAccessToken(_) => None, + }) + } } #[async_trait::async_trait] @@ -77,11 +97,12 @@ impl ModelProvider for AmazonBedrockModelProvider { } fn auth_manager(&self) -> Option> { - None + self.managed_auth() + .and_then(|_| self.auth_manager.as_ref().cloned()) } async fn auth(&self) -> Option { - None + self.managed_auth().map(CodexAuth::BedrockApiKey) } fn account_state(&self) -> ProviderAccountResult { @@ -92,17 +113,23 @@ impl ModelProvider for AmazonBedrockModelProvider { } async fn api_provider(&self) -> Result { + let managed_auth = self.managed_auth(); let mut api_provider_info = self.info.clone(); - api_provider_info.base_url = Some(runtime_base_url(&self.aws).await?); + api_provider_info.base_url = + Some(runtime_base_url(managed_auth.as_ref(), &self.aws).await?); api_provider_info.to_api_provider(/*auth_mode*/ None) } async fn runtime_base_url(&self) -> Result> { - Ok(Some(runtime_base_url(&self.aws).await?)) + let managed_auth = self.managed_auth(); + Ok(Some( + runtime_base_url(managed_auth.as_ref(), &self.aws).await?, + )) } async fn api_auth(&self) -> Result { - resolve_provider_auth(&self.aws).await + let managed_auth = self.managed_auth(); + resolve_provider_auth(managed_auth.as_ref(), &self.aws).await } fn models_manager( @@ -119,6 +146,7 @@ impl ModelProvider for AmazonBedrockModelProvider { #[cfg(test)] mod tests { + use http::HeaderValue; use pretty_assertions::assert_eq; use super::*; @@ -139,10 +167,68 @@ mod tests { ); } + #[tokio::test] + async fn managed_auth_takes_precedence_over_aws_auth() { + let managed_auth = BedrockApiKeyAuth { + api_key: "managed-bedrock-api-key".to_string(), + region: "us-east-1".to_string(), + }; + let auth_manager = + AuthManager::from_auth_for_testing(CodexAuth::BedrockApiKey(managed_auth.clone())); + let provider = AmazonBedrockModelProvider::new( + ModelProviderInfo::create_amazon_bedrock_provider(Some(ModelProviderAwsAuthInfo { + profile: Some("aws-profile-that-should-not-be-loaded".to_string()), + region: Some("us-west-2".to_string()), + })), + Some(auth_manager.clone()), + ); + + assert!(Arc::ptr_eq( + &provider + .auth_manager() + .expect("managed Bedrock auth manager should be exposed"), + &auth_manager, + )); + assert_eq!( + provider.auth().await, + Some(CodexAuth::BedrockApiKey(managed_auth)) + ); + assert_eq!( + provider + .runtime_base_url() + .await + .expect("managed Bedrock region should resolve"), + Some("https://bedrock-mantle.us-east-1.api.aws/openai/v1".to_string()) + ); + assert_eq!( + provider + .api_auth() + .await + .expect("managed Bedrock auth should resolve") + .to_auth_headers() + .get(http::header::AUTHORIZATION), + Some(&HeaderValue::from_static("Bearer managed-bedrock-api-key")) + ); + } + + #[tokio::test] + async fn openai_auth_is_not_exposed_to_bedrock() { + let provider = AmazonBedrockModelProvider::new( + ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None), + Some(AuthManager::from_auth_for_testing(CodexAuth::from_api_key( + "openai-api-key", + ))), + ); + + assert!(provider.auth_manager().is_none()); + assert_eq!(provider.auth().await, None); + } + #[test] fn capabilities_disable_unsupported_hosted_tools() { let provider = AmazonBedrockModelProvider::new( ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None), + /*auth_manager*/ None, ); assert_eq!( @@ -159,6 +245,7 @@ mod tests { fn approval_review_preferred_model_uses_bedrock_gpt_5_4() { let provider = AmazonBedrockModelProvider::new( ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None), + /*auth_manager*/ None, ); assert_eq!( diff --git a/codex-rs/model-provider/src/provider.rs b/codex-rs/model-provider/src/provider.rs index 95427465e..6f6602d01 100644 --- a/codex-rs/model-provider/src/provider.rs +++ b/codex-rs/model-provider/src/provider.rs @@ -179,7 +179,7 @@ pub fn create_model_provider( auth_manager: Option>, ) -> SharedModelProvider { if provider_info.is_amazon_bedrock() { - Arc::new(AmazonBedrockModelProvider::new(provider_info)) + Arc::new(AmazonBedrockModelProvider::new(provider_info, auth_manager)) } else { Arc::new(ConfiguredModelProvider::new(provider_info, auth_manager)) } @@ -460,6 +460,17 @@ mod tests { assert!(provider.auth_manager().is_none()); } + #[tokio::test] + async fn create_model_provider_uses_managed_auth_for_amazon_bedrock_provider() { + let auth = bedrock_api_key_auth(); + let provider = create_model_provider( + ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None), + Some(AuthManager::from_auth_for_testing(auth.clone())), + ); + + assert_eq!(provider.auth().await, Some(auth)); + } + #[test] fn openai_provider_returns_unauthenticated_openai_account_state() { let provider = create_model_provider(