model-provider/src/amazon_bedrock/mantle.rs
124 lines
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;use super::auth::BedrockAuthMethod;use super::auth::resolve_auth_method;const BEDROCK_MANTLE_SERVICE_NAME: &str = "bedrock-mantle";const BEDROCK_MANTLE_SUPPORTED_REGIONS: [&str; 12] = [ "us-east-2", "us-east-1", "us-west-2", "ap-southeast-3", "ap-south-1", "ap-northeast-1", "eu-central-1", "eu-west-1", "eu-west-2", "eu-south-1", "eu-north-1", "sa-east-1",];pub(super) fn aws_auth_config(aws: &ModelProviderAwsAuthInfo) -> AwsAuthConfig { AwsAuthConfig { profile: aws.profile.clone(), region: region_from_config(aws), service: BEDROCK_MANTLE_SERVICE_NAME.to_string(), }}pub(super) fn region_from_config(aws: &ModelProviderAwsAuthInfo) -> Option<String> { aws.region .as_deref() .map(str::trim) .filter(|region| !region.is_empty()) .map(str::to_string)}pub(super) fn base_url(region: &str) -> Result<String> { if BEDROCK_MANTLE_SUPPORTED_REGIONS.contains(®ion) { Ok(format!("https://bedrock-mantle.{region}.api.aws/openai/v1")) } else { Err(CodexErr::Fatal(format!( "Amazon Bedrock Mantle does not support region `{region}`" ))) }}pub(super) async fn runtime_base_url( managed_auth: Option<&BedrockApiKeyAuth>, aws: &ModelProviderAwsAuthInfo,) -> Result<String> { let region = resolve_region(managed_auth, aws).await?; base_url(®ion)}async fn resolve_region( managed_auth: Option<&BedrockApiKeyAuth>, aws: &ModelProviderAwsAuthInfo,) -> Result<String> { match resolve_auth_method(managed_auth, aws).await? { BedrockAuthMethod::ManagedBearerToken { region, .. } | BedrockAuthMethod::EnvBearerToken { region, .. } => Ok(region), BedrockAuthMethod::AwsSdkAuth { context } => Ok(context.region().to_string()), }}#[cfg(test)]mod tests { use pretty_assertions::assert_eq; use super::*; #[test] fn base_url_uses_region_endpoint() { assert_eq!( base_url("ap-northeast-1").expect("supported region"), "https://bedrock-mantle.ap-northeast-1.api.aws/openai/v1" ); } #[test] fn base_url_rejects_unsupported_region() { let err = base_url("us-west-1").expect_err("unsupported region"); assert_eq!( err.to_string(), "Fatal error: Amazon Bedrock Mantle does not support region `us-west-1`" ); } #[test] fn aws_auth_config_uses_profile_and_mantle_service() { assert_eq!( aws_auth_config(&ModelProviderAwsAuthInfo { profile: Some("codex-bedrock".to_string()), region: None, }), AwsAuthConfig { profile: Some("codex-bedrock".to_string()), region: None, service: "bedrock-mantle".to_string(), } ); } #[test] fn aws_auth_config_uses_configured_region() { assert_eq!( aws_auth_config(&ModelProviderAwsAuthInfo { profile: None, region: Some(" us-west-2 ".to_string()), }), AwsAuthConfig { profile: None, region: Some("us-west-2".to_string()), service: "bedrock-mantle".to_string(), } ); }}