From 5379cb3b27a4258a942069c5ed2904c8ed1a94e6 Mon Sep 17 00:00:00 2001 From: celia-oai Date: Fri, 29 May 2026 20:59:32 -0700 Subject: [PATCH] changes --- .../model-provider/src/amazon_bedrock/auth.rs | 13 +++- .../src/amazon_bedrock/mantle.rs | 6 +- .../model-provider/src/amazon_bedrock/mod.rs | 67 +++++++++++++------ .../src/amazon_bedrock/provider_auth.rs | 2 + 4 files changed, 60 insertions(+), 28 deletions(-) diff --git a/codex-rs/model-provider/src/amazon_bedrock/auth.rs b/codex-rs/model-provider/src/amazon_bedrock/auth.rs index 11e20a1cbe..75f7de1409 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/auth.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/auth.rs @@ -19,6 +19,7 @@ use crate::BearerAuthProvider; use super::mantle::aws_auth_config; use super::mantle::region_from_config; use super::provider_auth::AmazonBedrockAuth; +use super::provider_auth::StoredAmazonBedrockAuth; const AWS_BEARER_TOKEN_BEDROCK_ENV_VAR: &str = "AWS_BEARER_TOKEN_BEDROCK"; const AWS_REGION_ENV_VAR: &str = "AWS_REGION"; @@ -31,9 +32,14 @@ pub(super) enum BedrockAuthMethod { } pub(super) async fn resolve_auth_method( - stored_auth: Option<&AmazonBedrockAuth>, + stored_auth: &StoredAmazonBedrockAuth, aws: &ModelProviderAwsAuthInfo, ) -> Result { + let stored_auth = stored_auth + .as_ref() + .map(Option::as_ref) + .map_err(|err| CodexErr::Fatal(err.clone()))?; + if let Some(auth) = stored_auth { return Ok(BedrockAuthMethod::StoredBearerToken { token: auth.bearer_token.clone(), @@ -54,7 +60,7 @@ pub(super) async fn resolve_auth_method( } pub(super) async fn resolve_provider_auth( - stored_auth: Option<&AmazonBedrockAuth>, + stored_auth: &StoredAmazonBedrockAuth, aws: &ModelProviderAwsAuthInfo, ) -> Result { match resolve_auth_method(stored_auth, aws).await? { @@ -212,9 +218,10 @@ mod tests { bearer_token: "stored-bedrock-key".to_string(), region: "eu-west-1".to_string(), }; + let stored_auth: StoredAmazonBedrockAuth = Ok(Some(auth)); let method = resolve_auth_method( - Some(&auth), + &stored_auth, &ModelProviderAwsAuthInfo { profile: None, region: Some("us-east-1".to_string()), diff --git a/codex-rs/model-provider/src/amazon_bedrock/mantle.rs b/codex-rs/model-provider/src/amazon_bedrock/mantle.rs index 1e22eab30e..bec61645a0 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/mantle.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/mantle.rs @@ -5,7 +5,7 @@ use codex_protocol::error::Result; use super::auth::BedrockAuthMethod; use super::auth::resolve_auth_method; -use super::provider_auth::AmazonBedrockAuth; +use super::provider_auth::StoredAmazonBedrockAuth; const BEDROCK_MANTLE_SERVICE_NAME: &str = "bedrock-mantle"; const BEDROCK_MANTLE_SUPPORTED_REGIONS: [&str; 12] = [ @@ -54,7 +54,7 @@ pub fn is_supported_region(region: &str) -> bool { } pub(super) async fn runtime_base_url( - stored_auth: Option<&AmazonBedrockAuth>, + stored_auth: &StoredAmazonBedrockAuth, aws: &ModelProviderAwsAuthInfo, ) -> Result { let region = resolve_region(stored_auth, aws).await?; @@ -62,7 +62,7 @@ pub(super) async fn runtime_base_url( } async fn resolve_region( - stored_auth: Option<&AmazonBedrockAuth>, + stored_auth: &StoredAmazonBedrockAuth, aws: &ModelProviderAwsAuthInfo, ) -> Result { match resolve_auth_method(stored_auth, aws).await? { diff --git a/codex-rs/model-provider/src/amazon_bedrock/mod.rs b/codex-rs/model-provider/src/amazon_bedrock/mod.rs index fa02c122c5..828003664f 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/mod.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/mod.rs @@ -28,7 +28,7 @@ pub(crate) use catalog::static_model_catalog; use catalog::with_default_only_service_tier; pub use mantle::is_supported_region; use mantle::runtime_base_url; -pub use provider_auth::AmazonBedrockAuth; +use provider_auth::StoredAmazonBedrockAuth; pub use provider_auth::delete_amazon_bedrock_auth; pub use provider_auth::load_amazon_bedrock_auth; pub use provider_auth::save_amazon_bedrock_auth; @@ -38,7 +38,7 @@ pub use provider_auth::save_amazon_bedrock_auth; pub(crate) struct AmazonBedrockModelProvider { pub(crate) info: ModelProviderInfo, pub(crate) aws: ModelProviderAwsAuthInfo, - codex_home: Option, + stored_auth: StoredAmazonBedrockAuth, } impl AmazonBedrockModelProvider { @@ -50,23 +50,17 @@ impl AmazonBedrockModelProvider { profile: None, region: None, }); + let stored_auth = match codex_home.as_deref() { + Some(codex_home) => load_amazon_bedrock_auth(codex_home) + .map_err(|err| format!("failed to load Amazon Bedrock auth: {err}")), + None => Ok(None), + }; Self { info: provider_info, aws, - codex_home, + stored_auth, } } - - fn stored_auth(&self) -> Result> { - let Some(codex_home) = self.codex_home.as_deref() else { - return Ok(None); - }; - load_amazon_bedrock_auth(codex_home).map_err(|err| { - codex_protocol::error::CodexErr::Fatal(format!( - "failed to load Amazon Bedrock auth: {err}" - )) - }) - } } #[async_trait::async_trait] @@ -104,21 +98,16 @@ impl ModelProvider for AmazonBedrockModelProvider { async fn api_provider(&self) -> Result { let mut api_provider_info = self.info.clone(); - let stored_auth = self.stored_auth()?; - api_provider_info.base_url = Some(runtime_base_url(stored_auth.as_ref(), &self.aws).await?); + api_provider_info.base_url = Some(runtime_base_url(&self.stored_auth, &self.aws).await?); api_provider_info.to_api_provider(/*auth_mode*/ None) } async fn runtime_base_url(&self) -> Result> { - let stored_auth = self.stored_auth()?; - Ok(Some( - runtime_base_url(stored_auth.as_ref(), &self.aws).await?, - )) + Ok(Some(runtime_base_url(&self.stored_auth, &self.aws).await?)) } async fn api_auth(&self) -> Result { - let stored_auth = self.stored_auth()?; - resolve_provider_auth(stored_auth.as_ref(), &self.aws).await + resolve_provider_auth(&self.stored_auth, &self.aws).await } fn models_manager( @@ -135,6 +124,9 @@ impl ModelProvider for AmazonBedrockModelProvider { #[cfg(test)] mod tests { + use std::time::SystemTime; + use std::time::UNIX_EPOCH; + use pretty_assertions::assert_eq; use super::*; @@ -184,4 +176,35 @@ mod tests { AMAZON_BEDROCK_GPT_5_4_MODEL_ID ); } + + #[tokio::test] + async fn stored_auth_is_loaded_when_provider_is_created() { + let codex_home = std::env::temp_dir().join(format!( + "codex-bedrock-provider-auth-init-{}-{}", + std::process::id(), + SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("time") + .as_nanos() + )); + let _ = std::fs::remove_dir_all(&codex_home); + save_amazon_bedrock_auth(&codex_home, "bedrock-key", "eu-west-2").expect("save auth"); + + let provider = AmazonBedrockModelProvider::new( + ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None), + Some(codex_home.clone()), + ); + + delete_amazon_bedrock_auth(&codex_home).expect("delete auth"); + + assert_eq!( + provider + .runtime_base_url() + .await + .expect("runtime base URL should use initialized auth"), + Some("https://bedrock-mantle.eu-west-2.api.aws/openai/v1".to_string()) + ); + + let _ = std::fs::remove_dir_all(&codex_home); + } } diff --git a/codex-rs/model-provider/src/amazon_bedrock/provider_auth.rs b/codex-rs/model-provider/src/amazon_bedrock/provider_auth.rs index c72199e8cf..1397fb6c83 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/provider_auth.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/provider_auth.rs @@ -13,6 +13,8 @@ pub struct AmazonBedrockAuth { pub region: String, } +pub(super) type StoredAmazonBedrockAuth = std::result::Result, String>; + pub fn save_amazon_bedrock_auth( codex_home: &Path, bearer_token: &str,