From d5e256ceb210dacbab9d8dcc7eafda58f90f6e51 Mon Sep 17 00:00:00 2001 From: Celia Chen Date: Fri, 14 Aug 2026 01:18:49 +0000 Subject: [PATCH] Add an Amazon Bedrock Runtime provider (#38470) ## What changed - Add the built-in `amazon-bedrock-runtime` provider for regional `bedrock-runtime` OpenAI-compatible endpoints. - Use endpoint-specific SigV4 service configuration while preserving bearer token auth and per-provider AWS profile, region, and transport overrides. - Provide global and US cross-region GPT-5.6 model variants, prefer global routing for fallback and background tasks, and disable unsupported web search. ## Testing - Cover provider configuration and capabilities, endpoint and auth selection, runtime catalog ordering, model fallback, and web search omission. GitOrigin-RevId: d158ab74a5768e1b4c61063365a0c06735313322 --- .../v2/model_provider_capabilities_read.rs | 30 ++++ .../app-server/tests/suite/v2/thread_start.rs | 31 ++++ codex-rs/config/src/bedrock_runtime_tests.rs | 46 ++++++ codex-rs/config/src/config_toml.rs | 22 ++- codex-rs/core/tests/suite/web_search.rs | 61 ++++++++ codex-rs/model-provider-info/src/lib.rs | 38 ++++- .../src/model_provider_info_tests.rs | 102 ++++++++++++- .../model-provider/src/amazon_bedrock/auth.rs | 44 ++++-- .../src/amazon_bedrock/mantle.rs | 17 +-- .../model-provider/src/amazon_bedrock/mod.rs | 139 ++++++++++++++++-- .../src/amazon_bedrock/runtime.rs | 34 +++++ .../src/amazon_bedrock/runtime_catalog.rs | 41 ++++++ .../amazon_bedrock/runtime_catalog_tests.rs | 61 ++++++++ .../src/amazon_bedrock/runtime_tests.rs | 29 ++++ codex-rs/model-provider/src/lib.rs | 1 + 15 files changed, 642 insertions(+), 54 deletions(-) create mode 100644 codex-rs/config/src/bedrock_runtime_tests.rs create mode 100644 codex-rs/model-provider/src/amazon_bedrock/runtime.rs create mode 100644 codex-rs/model-provider/src/amazon_bedrock/runtime_catalog.rs create mode 100644 codex-rs/model-provider/src/amazon_bedrock/runtime_catalog_tests.rs create mode 100644 codex-rs/model-provider/src/amazon_bedrock/runtime_tests.rs diff --git a/codex-rs/app-server/tests/suite/v2/model_provider_capabilities_read.rs b/codex-rs/app-server/tests/suite/v2/model_provider_capabilities_read.rs index e05b7e2e3acc..8a497bfbbd5b 100644 --- a/codex-rs/app-server/tests/suite/v2/model_provider_capabilities_read.rs +++ b/codex-rs/app-server/tests/suite/v2/model_provider_capabilities_read.rs @@ -62,3 +62,33 @@ async fn read_amazon_bedrock_provider_capabilities() -> Result<()> { assert_eq!(received, expected); Ok(()) } + +#[tokio::test] +async fn read_amazon_bedrock_runtime_provider_capabilities() -> Result<()> { + let codex_home = TempDir::new()?; + std::fs::write( + codex_home.path().join("config.toml"), + r#"model_provider = "amazon-bedrock-runtime" +"#, + )?; + let mut mcp = TestAppServer::builder() + .with_codex_home(codex_home.path()) + .build_initialized_with_timeout(DEFAULT_TIMEOUT) + .await?; + + let request_id = mcp + .send_model_provider_capabilities_read_request(ModelProviderCapabilitiesReadParams {}) + .await?; + let received: ModelProviderCapabilitiesReadResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; + + assert_eq!( + received, + ModelProviderCapabilitiesReadResponse { + namespace_tools: true, + image_generation: false, + web_search: false, + } + ); + Ok(()) +} diff --git a/codex-rs/app-server/tests/suite/v2/thread_start.rs b/codex-rs/app-server/tests/suite/v2/thread_start.rs index c4a12be6841b..5d3135b1a227 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_start.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_start.rs @@ -267,6 +267,37 @@ async fn thread_start_provider_model_fallback_uses_bedrock_static_catalog() -> R Ok(()) } +#[tokio::test] +async fn thread_start_bedrock_runtime_prefers_global_cross_region_models() -> Result<()> { + let codex_home = TempDir::new()?; + std::fs::write( + codex_home.path().join("config.toml"), + r#"model_provider = "amazon-bedrock-runtime" +"#, + )?; + let mut mcp = TestAppServer::builder() + .with_codex_home(codex_home.path()) + .build_initialized() + .await?; + + for model in ["global.openai.gpt-5.6-sol", "us.openai.gpt-5.6-sol"] { + let response = + start_thread_with_model(&mut mcp, model, /*allow_provider_model_fallback*/ true) + .await?; + assert_eq!(response.model, model); + } + + let response = start_thread_with_model( + &mut mcp, + "openai.gpt-5.6-sol", + /*allow_provider_model_fallback*/ true, + ) + .await?; + assert_eq!(response.model, "global.openai.gpt-5.6-sol"); + + Ok(()) +} + #[tokio::test] async fn thread_start_provider_model_fallback_ignores_dynamic_catalog() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; diff --git a/codex-rs/config/src/bedrock_runtime_tests.rs b/codex-rs/config/src/bedrock_runtime_tests.rs new file mode 100644 index 000000000000..8bba3aca394d --- /dev/null +++ b/codex-rs/config/src/bedrock_runtime_tests.rs @@ -0,0 +1,46 @@ +use codex_model_provider_info::AMAZON_BEDROCK_RUNTIME_PROVIDER_ID; +use codex_model_provider_info::ModelProviderAwsAuthInfo; +use pretty_assertions::assert_eq; + +use super::ConfigToml; + +#[test] +fn runtime_provider_accepts_aws_profile_and_region_overrides() { + let config = toml::from_str::( + r#" +[model_providers.amazon-bedrock-runtime.aws] +profile = "runtime-profile" +region = "us-west-2" +"#, + ) + .expect("Bedrock Runtime AWS overrides should deserialize"); + + assert_eq!( + config + .model_providers + .get(AMAZON_BEDROCK_RUNTIME_PROVIDER_ID) + .and_then(|provider| provider.aws.clone()), + Some(ModelProviderAwsAuthInfo { + profile: Some("runtime-profile".to_string()), + region: Some("us-west-2".to_string()), + }) + ); +} + +#[test] +fn custom_provider_still_rejects_aws_auth() { + let error = toml::from_str::( + r#" +[model_providers.custom] +name = "Custom" + +[model_providers.custom.aws] +region = "us-west-2" +"#, + ) + .expect_err("custom providers must not accept AWS auth"); + + assert!(error.to_string().contains( + "provider aws is only supported for `amazon-bedrock` or `amazon-bedrock-runtime`" + )); +} diff --git a/codex-rs/config/src/config_toml.rs b/codex-rs/config/src/config_toml.rs index a4ecf0a3e69b..20d64d21babc 100644 --- a/codex-rs/config/src/config_toml.rs +++ b/codex-rs/config/src/config_toml.rs @@ -30,6 +30,7 @@ use crate::types::UriBasedFileOpener; use crate::types::WindowsToml; use codex_features::FeaturesToml; use codex_model_provider_info::AMAZON_BEDROCK_PROVIDER_ID; +use codex_model_provider_info::AMAZON_BEDROCK_RUNTIME_PROVIDER_ID; use codex_model_provider_info::LEGACY_OLLAMA_CHAT_PROVIDER_ID; use codex_model_provider_info::LMSTUDIO_OSS_PROVIDER_ID; use codex_model_provider_info::ModelProviderInfo; @@ -59,8 +60,9 @@ use serde::Serialize; use serde::de::Error as SerdeError; use serde_json::Value as JsonValue; -const RESERVED_MODEL_PROVIDER_IDS: [&str; 4] = [ +const RESERVED_MODEL_PROVIDER_IDS: [&str; 5] = [ AMAZON_BEDROCK_PROVIDER_ID, + AMAZON_BEDROCK_RUNTIME_PROVIDER_ID, OPENAI_PROVIDER_ID, OLLAMA_OSS_PROVIDER_ID, LMSTUDIO_OSS_PROVIDER_ID, @@ -883,8 +885,10 @@ pub fn validate_reserved_model_provider_ids( let mut conflicts = model_providers .keys() .filter(|key| { - key.as_str() != AMAZON_BEDROCK_PROVIDER_ID - && RESERVED_MODEL_PROVIDER_IDS.contains(&key.as_str()) + !matches!( + key.as_str(), + AMAZON_BEDROCK_PROVIDER_ID | AMAZON_BEDROCK_RUNTIME_PROVIDER_ID + ) && RESERVED_MODEL_PROVIDER_IDS.contains(&key.as_str()) }) .map(|key| format!("`{key}`")) .collect::>(); @@ -905,10 +909,14 @@ pub fn validate_model_providers( ) -> Result<(), String> { validate_reserved_model_provider_ids(model_providers)?; for (key, provider) in model_providers { - if key != AMAZON_BEDROCK_PROVIDER_ID { + if !matches!( + key.as_str(), + AMAZON_BEDROCK_PROVIDER_ID | AMAZON_BEDROCK_RUNTIME_PROVIDER_ID + ) { if provider.aws.is_some() { return Err(format!( - "model_providers.{key}: provider aws is only supported for `{AMAZON_BEDROCK_PROVIDER_ID}`" + "model_providers.{key}: provider aws is only supported for \ +`{AMAZON_BEDROCK_PROVIDER_ID}` or `{AMAZON_BEDROCK_RUNTIME_PROVIDER_ID}`" )); } if provider.name.trim().is_empty() { @@ -935,6 +943,10 @@ where Ok(model_providers) } +#[cfg(test)] +#[path = "bedrock_runtime_tests.rs"] +mod bedrock_runtime_tests; + pub fn validate_oss_provider(provider: &str) -> std::io::Result<()> { match provider { LMSTUDIO_OSS_PROVIDER_ID | OLLAMA_OSS_PROVIDER_ID => Ok(()), diff --git a/codex-rs/core/tests/suite/web_search.rs b/codex-rs/core/tests/suite/web_search.rs index 8b0f9cf208e9..d21417bae077 100644 --- a/codex-rs/core/tests/suite/web_search.rs +++ b/codex-rs/core/tests/suite/web_search.rs @@ -6,6 +6,7 @@ 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::AMAZON_BEDROCK_PROVIDER_ID; +use codex_model_provider_info::AMAZON_BEDROCK_RUNTIME_PROVIDER_ID; use codex_model_provider_info::ModelProviderInfo; use codex_protocol::config_types::WebSearchMode; use codex_protocol::models::PermissionProfile; @@ -147,6 +148,66 @@ async fn amazon_bedrock_web_search_uses_text_only_hosted_tools() { } } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn amazon_bedrock_runtime_preserves_cross_region_models_without_web_search() { + skip_if_no_network!(); + + for model in ["us.openai.gpt-5.6-sol", "global.openai.gpt-5.6-sol"] { + let server = start_mock_server().await; + let response = responses::mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_completed("resp-1"), + ]), + ) + .await; + let auth = CodexAuth::BedrockApiKey(BedrockApiKeyAuth { + api_key: "dummy".to_string(), + region: "us-east-1".to_string(), + }); + let mut builder = test_codex() + .with_auth(auth) + .with_model(model) + .with_config(|config| { + let base_url = config.model_provider.base_url.clone(); + config.model_provider_id = AMAZON_BEDROCK_RUNTIME_PROVIDER_ID.to_string(); + config.model_provider = + ModelProviderInfo::create_amazon_bedrock_runtime_provider(/*aws*/ None); + config.model_provider.base_url = base_url; + config + .web_search_mode + .set(WebSearchMode::Cached) + .expect("test web search mode should satisfy constraints"); + }); + let test = builder + .build_with_auto_env(&server) + .await + .expect("create test Bedrock Runtime conversation"); + + test.submit_turn_with_permission_profile( + "hello Bedrock Runtime", + PermissionProfile::Disabled, + ) + .await + .expect("submit Bedrock Runtime turn"); + + let body = response.single_request().body_json(); + let web_search_tool = body + .get("tools") + .and_then(Value::as_array) + .and_then(|tools| { + tools + .iter() + .find(|tool| tool.get("type").and_then(Value::as_str) == Some("web_search")) + }); + assert_eq!( + (body["model"].as_str(), web_search_tool), + (Some(model), None) + ); + } +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn amazon_bedrock_web_search_is_disabled_when_managed_requirements_prohibit_cached_search() { skip_if_no_network!(); diff --git a/codex-rs/model-provider-info/src/lib.rs b/codex-rs/model-provider-info/src/lib.rs index 95d9401fa52f..14f55c2042d0 100644 --- a/codex-rs/model-provider-info/src/lib.rs +++ b/codex-rs/model-provider-info/src/lib.rs @@ -37,11 +37,16 @@ pub const OPENAI_PROVIDER_ID: &str = "openai"; pub const CHATGPT_CODEX_BASE_URL: &str = "https://chatgpt.com/backend-api/codex"; const AMAZON_BEDROCK_PROVIDER_NAME: &str = "Amazon Bedrock"; pub const AMAZON_BEDROCK_PROVIDER_ID: &str = "amazon-bedrock"; +const AMAZON_BEDROCK_RUNTIME_PROVIDER_NAME: &str = "Amazon Bedrock Runtime"; +pub const AMAZON_BEDROCK_RUNTIME_PROVIDER_ID: &str = "amazon-bedrock-runtime"; pub const AMAZON_BEDROCK_GPT_5_5_MODEL_ID: &str = "openai.gpt-5.5"; pub const AMAZON_BEDROCK_GPT_5_4_MODEL_ID: &str = "openai.gpt-5.4"; pub const AMAZON_BEDROCK_GPT_5_6_SOL_MODEL_ID: &str = "openai.gpt-5.6-sol"; pub const AMAZON_BEDROCK_GPT_5_6_TERRA_MODEL_ID: &str = "openai.gpt-5.6-terra"; pub const AMAZON_BEDROCK_GPT_5_6_LUNA_MODEL_ID: &str = "openai.gpt-5.6-luna"; +pub const AMAZON_BEDROCK_RUNTIME_GLOBAL_GPT_5_6_TERRA_MODEL_ID: &str = + "global.openai.gpt-5.6-terra"; +pub const AMAZON_BEDROCK_RUNTIME_GLOBAL_GPT_5_6_LUNA_MODEL_ID: &str = "global.openai.gpt-5.6-luna"; pub const AMAZON_BEDROCK_DEFAULT_BASE_URL: &str = "https://bedrock-mantle.us-east-1.api.aws/openai/v1"; const AMAZON_BEDROCK_MANTLE_CLIENT_AGENT_HEADER: &str = "x-amzn-mantle-client-agent"; @@ -400,6 +405,15 @@ impl ModelProviderInfo { } } + pub fn create_amazon_bedrock_runtime_provider( + aws: Option, + ) -> ModelProviderInfo { + let mut provider = Self::create_amazon_bedrock_provider(aws); + provider.name = AMAZON_BEDROCK_RUNTIME_PROVIDER_NAME.into(); + provider.http_headers = None; + provider + } + pub fn is_openai(&self) -> bool { self.name == OPENAI_PROVIDER_NAME } @@ -416,6 +430,11 @@ impl ModelProviderInfo { pub fn is_amazon_bedrock(&self) -> bool { self.name == AMAZON_BEDROCK_PROVIDER_NAME + || self.name == AMAZON_BEDROCK_RUNTIME_PROVIDER_NAME + } + + pub fn is_amazon_bedrock_runtime(&self) -> bool { + self.name == AMAZON_BEDROCK_RUNTIME_PROVIDER_NAME } pub fn has_command_auth(&self) -> bool { @@ -436,6 +455,8 @@ pub fn built_in_model_providers( use ModelProviderInfo as P; let openai_provider = P::create_openai_provider(openai_base_url); let amazon_bedrock_provider = P::create_amazon_bedrock_provider(/*aws*/ None); + let amazon_bedrock_runtime_provider = + P::create_amazon_bedrock_runtime_provider(/*aws*/ None); // We do not want to be in the business of adjucating which third-party // providers are bundled with Codex CLI, so we only include the OpenAI and @@ -444,6 +465,10 @@ pub fn built_in_model_providers( [ (OPENAI_PROVIDER_ID, openai_provider), (AMAZON_BEDROCK_PROVIDER_ID, amazon_bedrock_provider), + ( + AMAZON_BEDROCK_RUNTIME_PROVIDER_ID, + amazon_bedrock_runtime_provider, + ), ( OLLAMA_OSS_PROVIDER_ID, create_oss_provider(DEFAULT_OLLAMA_PORT, WireApi::Responses), @@ -461,27 +486,30 @@ pub fn built_in_model_providers( /// Merge configured providers into the built-in provider catalog. /// /// Configured providers extend the built-in set. Built-in providers are not -/// generally overridable, but the built-in Amazon Bedrock provider allows the -/// user to customize its endpoint, authentication, headers, and AWS settings. +/// generally overridable, but built-in Amazon Bedrock providers allow the user +/// to customize their endpoint, authentication, headers, and AWS settings. pub fn merge_configured_model_providers( mut model_providers: HashMap, configured_model_providers: HashMap, ) -> Result, String> { for (key, mut provider) in configured_model_providers { - if key == AMAZON_BEDROCK_PROVIDER_ID { + if matches!( + key.as_str(), + AMAZON_BEDROCK_PROVIDER_ID | AMAZON_BEDROCK_RUNTIME_PROVIDER_ID + ) { let base_url_override = provider.base_url.take(); let auth_override = provider.auth.take(); let aws_override = provider.aws.take(); let http_headers_override = provider.http_headers.take(); if provider != ModelProviderInfo::default() { return Err(format!( - "model_providers.{AMAZON_BEDROCK_PROVIDER_ID} only supports changing \ + "model_providers.{key} only supports changing \ `base_url`, `auth`, `http_headers`, `aws.profile`, and `aws.region`; other non-default \ provider fields are not supported" )); } - if let Some(built_in_provider) = model_providers.get_mut(AMAZON_BEDROCK_PROVIDER_ID) { + if let Some(built_in_provider) = model_providers.get_mut(&key) { built_in_provider.base_url = base_url_override; built_in_provider.auth = auth_override; if let Some(aws_override) = aws_override { diff --git a/codex-rs/model-provider-info/src/model_provider_info_tests.rs b/codex-rs/model-provider-info/src/model_provider_info_tests.rs index 39ceba81998b..26950e41ac54 100644 --- a/codex-rs/model-provider-info/src/model_provider_info_tests.rs +++ b/codex-rs/model-provider-info/src/model_provider_info_tests.rs @@ -262,6 +262,45 @@ fn test_create_amazon_bedrock_provider() { ); } +#[test] +fn test_create_amazon_bedrock_runtime_provider() { + let mut expected = ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None); + expected.name = "Amazon Bedrock Runtime".to_string(); + expected.http_headers = None; + + assert_eq!( + ModelProviderInfo::create_amazon_bedrock_runtime_provider(/*aws*/ None), + expected + ); +} + +#[test] +fn test_create_amazon_bedrock_runtime_provider_with_aws_configuration() { + let provider = + ModelProviderInfo::create_amazon_bedrock_runtime_provider(Some(ModelProviderAwsAuthInfo { + profile: Some("runtime-profile".to_string()), + region: Some("us-west-2".to_string()), + })); + + assert_eq!( + ( + provider.name.as_str(), + provider.aws, + provider.http_headers, + provider.supports_standalone_web_search, + ), + ( + "Amazon Bedrock Runtime", + Some(ModelProviderAwsAuthInfo { + profile: Some("runtime-profile".to_string()), + region: Some("us-west-2".to_string()), + }), + None, + false, + ) + ); +} + fn provider_auth_for_test() -> ModelProviderAuthInfo { ModelProviderAuthInfo { command: "token-fetcher".to_string(), @@ -291,14 +330,39 @@ fn test_amazon_bedrock_provider_adds_mantle_client_agent_header() { } #[test] -fn test_built_in_model_providers_include_amazon_bedrock() { +fn test_built_in_model_providers_include_amazon_bedrock_endpoints() { let providers = built_in_model_providers(/*openai_base_url*/ None); assert_eq!( - providers + [ + AMAZON_BEDROCK_PROVIDER_ID, + AMAZON_BEDROCK_RUNTIME_PROVIDER_ID + ] + .into_iter() + .map(|provider_id| { + providers + .get(provider_id) + .map(ModelProviderInfo::is_amazon_bedrock) + }) + .collect::>(), + vec![Some(true), Some(true)] + ); +} + +#[test] +fn test_built_in_model_providers_include_amazon_bedrock_runtime() { + let providers = built_in_model_providers(/*openai_base_url*/ None); + let runtime = providers + .get(AMAZON_BEDROCK_RUNTIME_PROVIDER_ID) + .expect("Amazon Bedrock Runtime provider should be built in"); + + assert!(runtime.is_amazon_bedrock()); + assert!(runtime.is_amazon_bedrock_runtime()); + assert!( + !providers .get(AMAZON_BEDROCK_PROVIDER_ID) - .map(ModelProviderInfo::is_amazon_bedrock), - Some(true) + .expect("Amazon Bedrock provider should be built in") + .is_amazon_bedrock_runtime() ); } @@ -355,6 +419,36 @@ fn test_merge_configured_model_providers_applies_amazon_bedrock_profile_override ); } +#[test] +fn test_merge_configured_model_providers_applies_runtime_overrides_independently() { + let runtime_aws = ModelProviderAwsAuthInfo { + profile: Some("runtime-profile".to_string()), + region: Some("eu-west-1".to_string()), + }; + let configured_model_providers = std::collections::HashMap::from([( + AMAZON_BEDROCK_RUNTIME_PROVIDER_ID.to_string(), + ModelProviderInfo { + base_url: Some("https://runtime.example.com/openai/v1".to_string()), + aws: Some(runtime_aws.clone()), + ..ModelProviderInfo::default() + }, + )]); + let mut expected = built_in_model_providers(/*openai_base_url*/ None); + let expected_runtime = expected + .get_mut(AMAZON_BEDROCK_RUNTIME_PROVIDER_ID) + .expect("Amazon Bedrock Runtime provider should be built in"); + expected_runtime.base_url = Some("https://runtime.example.com/openai/v1".to_string()); + expected_runtime.aws = Some(runtime_aws); + + assert_eq!( + merge_configured_model_providers( + built_in_model_providers(/*openai_base_url*/ None), + configured_model_providers, + ), + Ok(expected) + ); +} + #[test] fn test_merge_configured_model_providers_applies_amazon_bedrock_transport_overrides() { let auth = provider_auth_for_test(); diff --git a/codex-rs/model-provider/src/amazon_bedrock/auth.rs b/codex-rs/model-provider/src/amazon_bedrock/auth.rs index 20701670358d..1420f146543b 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/auth.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/auth.rs @@ -17,8 +17,10 @@ use http::HeaderMap; use crate::BearerAuthProvider; +use super::BedrockEndpoint; use super::mantle::aws_auth_config; use super::mantle::region_from_config; +use super::runtime; const AWS_BEARER_TOKEN_BEDROCK_ENV_VAR: &str = "AWS_BEARER_TOKEN_BEDROCK"; const AWS_REGION_ENV_VAR: &str = "AWS_REGION"; @@ -33,6 +35,7 @@ pub(super) enum BedrockAuthMethod { pub(super) async fn resolve_auth_method( managed_auth: Option<&BedrockApiKeyAuth>, aws: &ModelProviderAwsAuthInfo, + endpoint: BedrockEndpoint, ) -> Result { if let Some(managed_auth) = managed_auth { return Ok(BedrockAuthMethod::ManagedBearerToken { @@ -46,7 +49,10 @@ pub(super) async fn resolve_auth_method( return Ok(BedrockAuthMethod::EnvBearerToken { token, region }); } - let config = aws_auth_config(aws); + let config = match endpoint { + BedrockEndpoint::Mantle => aws_auth_config(aws), + BedrockEndpoint::Runtime => runtime::aws_auth_config(aws), + }; let context = AwsAuthContext::load(config) .await .map_err(aws_auth_error_to_codex_error)?; @@ -56,8 +62,9 @@ pub(super) async fn resolve_auth_method( pub(super) async fn resolve_provider_auth( managed_auth: Option<&BedrockApiKeyAuth>, aws: &ModelProviderAwsAuthInfo, + endpoint: BedrockEndpoint, ) -> Result { - match resolve_auth_method(managed_auth, aws).await? { + match resolve_auth_method(managed_auth, aws, endpoint).await? { BedrockAuthMethod::ManagedBearerToken { token, .. } | BedrockAuthMethod::EnvBearerToken { token, .. } => Ok(Arc::new(BearerAuthProvider { token: Some(token), @@ -65,11 +72,23 @@ pub(super) async fn resolve_provider_auth( is_fedramp_account: false, })), BedrockAuthMethod::AwsSdkAuth { context } => { - Ok(Arc::new(BedrockMantleSigV4AuthProvider::new(context))) + Ok(Arc::new(BedrockSigV4AuthProvider::new(context, endpoint))) } } } +pub(super) async fn resolve_region( + managed_auth: Option<&BedrockApiKeyAuth>, + aws: &ModelProviderAwsAuthInfo, + endpoint: BedrockEndpoint, +) -> Result { + match resolve_auth_method(managed_auth, aws, endpoint).await? { + BedrockAuthMethod::ManagedBearerToken { region, .. } + | BedrockAuthMethod::EnvBearerToken { region, .. } => Ok(region), + BedrockAuthMethod::AwsSdkAuth { context } => Ok(context.region().to_string()), + } +} + fn non_empty_env_var_from( name: &'static str, env_var: impl Fn(&'static str) -> std::result::Result, @@ -123,20 +142,23 @@ fn remove_headers_not_preserved_by_bedrock_mantle(headers: &mut HeaderMap) { } } -/// AWS SigV4 auth provider for Bedrock Mantle OpenAI-compatible requests. +/// AWS SigV4 auth provider for Bedrock OpenAI-compatible requests. #[derive(Debug)] -struct BedrockMantleSigV4AuthProvider { +struct BedrockSigV4AuthProvider { context: AwsAuthContext, + endpoint: BedrockEndpoint, } -impl BedrockMantleSigV4AuthProvider { - fn new(context: AwsAuthContext) -> Self { - Self { context } +impl BedrockSigV4AuthProvider { + fn new(context: AwsAuthContext, endpoint: BedrockEndpoint) -> Self { + Self { context, endpoint } } async fn apply_auth(&self, request: Request) -> std::result::Result { let mut request = request; - remove_headers_not_preserved_by_bedrock_mantle(&mut request.headers); + if self.endpoint == BedrockEndpoint::Mantle { + remove_headers_not_preserved_by_bedrock_mantle(&mut request.headers); + } let prepared = request.prepare_body_for_send().map_err(AuthError::Build)?; let signed = self .context @@ -157,11 +179,11 @@ impl BedrockMantleSigV4AuthProvider { } } -impl AuthProvider for BedrockMantleSigV4AuthProvider { +impl AuthProvider for BedrockSigV4AuthProvider { fn add_auth_headers(&self, _headers: &mut HeaderMap) {} fn apply_auth(&self, request: Request) -> codex_api::AuthProviderFuture<'_> { - Box::pin(BedrockMantleSigV4AuthProvider::apply_auth(self, request)) + Box::pin(BedrockSigV4AuthProvider::apply_auth(self, request)) } } diff --git a/codex-rs/model-provider/src/amazon_bedrock/mantle.rs b/codex-rs/model-provider/src/amazon_bedrock/mantle.rs index d9f9ac48848d..a75f124cb895 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/mantle.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/mantle.rs @@ -4,8 +4,8 @@ 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; +use super::BedrockEndpoint; +use super::auth::resolve_region; const BEDROCK_MANTLE_SERVICE_NAME: &str = "bedrock-mantle"; const BEDROCK_MANTLE_SUPPORTED_REGIONS: [&str; 12] = [ @@ -58,21 +58,10 @@ pub(super) async fn bedrock_mantle_runtime_base_url( managed_auth: Option<&BedrockApiKeyAuth>, aws: &ModelProviderAwsAuthInfo, ) -> Result { - let region = resolve_region(managed_auth, aws).await?; + let region = resolve_region(managed_auth, aws, BedrockEndpoint::Mantle).await?; base_url(®ion) } -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()), - } -} - #[cfg(test)] mod tests { use pretty_assertions::assert_eq; diff --git a/codex-rs/model-provider/src/amazon_bedrock/mod.rs b/codex-rs/model-provider/src/amazon_bedrock/mod.rs index 6e3140d0e2c6..f7707d0ab84c 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/mod.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/mod.rs @@ -2,6 +2,8 @@ mod auth; mod catalog; mod error; mod mantle; +mod runtime; +mod runtime_catalog; use std::path::PathBuf; use std::sync::Arc; @@ -12,7 +14,10 @@ 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::AMAZON_BEDROCK_GPT_5_6_LUNA_MODEL_ID; +use codex_model_provider_info::AMAZON_BEDROCK_GPT_5_6_TERRA_MODEL_ID; +use codex_model_provider_info::AMAZON_BEDROCK_RUNTIME_GLOBAL_GPT_5_6_LUNA_MODEL_ID; +use codex_model_provider_info::AMAZON_BEDROCK_RUNTIME_GLOBAL_GPT_5_6_TERRA_MODEL_ID; use codex_model_provider_info::ModelProviderAwsAuthInfo; use codex_model_provider_info::ModelProviderInfo; use codex_models_manager::manager::SharedModelsManager; @@ -35,12 +40,21 @@ use catalog::normalize_bedrock_catalog; pub(crate) use catalog::static_model_catalog; use mantle::bedrock_mantle_runtime_base_url; pub use mantle::is_supported_amazon_bedrock_region; +use runtime::bedrock_runtime_base_url; +use runtime_catalog::static_runtime_model_catalog; -/// Runtime provider for Amazon Bedrock's OpenAI-compatible Mantle endpoint. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) enum BedrockEndpoint { + Mantle, + Runtime, +} + +/// Runtime provider for Amazon Bedrock's OpenAI-compatible endpoints. #[derive(Clone, Debug)] pub(crate) struct AmazonBedrockModelProvider { pub(crate) info: ModelProviderInfo, pub(crate) aws: ModelProviderAwsAuthInfo, + endpoint: BedrockEndpoint, auth_manager: Option>, } @@ -50,6 +64,11 @@ impl AmazonBedrockModelProvider { auth_manager: Option>, ) -> Self { let auth_manager = auth_manager_for_provider(auth_manager, &provider_info); + let endpoint = if provider_info.is_amazon_bedrock_runtime() { + BedrockEndpoint::Runtime + } else { + BedrockEndpoint::Mantle + }; let aws = provider_info .aws .clone() @@ -60,6 +79,7 @@ impl AmazonBedrockModelProvider { Self { info: provider_info, aws, + endpoint, auth_manager, } } @@ -101,9 +121,15 @@ impl AmazonBedrockModelProvider { return Ok(Some(base_url)); } let managed_auth = self.managed_auth(); - Ok(Some( - bedrock_mantle_runtime_base_url(managed_auth.as_ref(), &self.aws).await?, - )) + let base_url = match self.endpoint { + BedrockEndpoint::Mantle => { + bedrock_mantle_runtime_base_url(managed_auth.as_ref(), &self.aws).await? + } + BedrockEndpoint::Runtime => { + bedrock_runtime_base_url(managed_auth.as_ref(), &self.aws).await? + } + }; + Ok(Some(base_url)) } async fn api_auth(&self) -> Result { @@ -112,7 +138,14 @@ impl AmazonBedrockModelProvider { return resolve_configured_provider_auth(auth.as_ref(), &self.info); } let managed_auth = self.managed_auth(); - resolve_bedrock_provider_auth(managed_auth.as_ref(), &self.aws).await + resolve_bedrock_provider_auth(managed_auth.as_ref(), &self.aws, self.endpoint).await + } + + fn default_model_catalog(&self) -> ModelsResponse { + match self.endpoint { + BedrockEndpoint::Mantle => static_model_catalog(), + BedrockEndpoint::Runtime => static_runtime_model_catalog(), + } } } @@ -125,22 +158,31 @@ impl ModelProvider for AmazonBedrockModelProvider { ProviderCapabilities { namespace_tools: true, image_generation: false, - web_search: true, + web_search: self.endpoint == BedrockEndpoint::Mantle, external_web_access: false, remote_compaction: RemoteCompactionSupport::V1, } } fn approval_review_preferred_model(&self) -> &'static str { - AMAZON_BEDROCK_GPT_5_4_MODEL_ID + match self.endpoint { + BedrockEndpoint::Mantle => AMAZON_BEDROCK_GPT_5_6_LUNA_MODEL_ID, + BedrockEndpoint::Runtime => AMAZON_BEDROCK_RUNTIME_GLOBAL_GPT_5_6_LUNA_MODEL_ID, + } } fn memory_extraction_preferred_model(&self) -> &'static str { - AMAZON_BEDROCK_GPT_5_4_MODEL_ID + match self.endpoint { + BedrockEndpoint::Mantle => AMAZON_BEDROCK_GPT_5_6_LUNA_MODEL_ID, + BedrockEndpoint::Runtime => AMAZON_BEDROCK_RUNTIME_GLOBAL_GPT_5_6_LUNA_MODEL_ID, + } } fn memory_consolidation_preferred_model(&self) -> &'static str { - AMAZON_BEDROCK_GPT_5_4_MODEL_ID + match self.endpoint { + BedrockEndpoint::Mantle => AMAZON_BEDROCK_GPT_5_6_TERRA_MODEL_ID, + BedrockEndpoint::Runtime => AMAZON_BEDROCK_RUNTIME_GLOBAL_GPT_5_6_TERRA_MODEL_ID, + } } fn auth_manager(&self) -> Option> { @@ -187,7 +229,8 @@ impl ModelProvider for AmazonBedrockModelProvider { ) -> SharedModelsManager { Arc::new(StaticModelsManager::new( /*auth_manager*/ None, - config_model_catalog.map_or_else(static_model_catalog, normalize_bedrock_catalog), + config_model_catalog + .map_or_else(|| self.default_model_catalog(), normalize_bedrock_catalog), )) } @@ -197,7 +240,8 @@ impl ModelProvider for AmazonBedrockModelProvider { ) -> SharedModelsManager { Arc::new(StaticModelsManager::new( /*auth_manager*/ None, - config_model_catalog.map_or_else(static_model_catalog, normalize_bedrock_catalog), + config_model_catalog + .map_or_else(|| self.default_model_catalog(), normalize_bedrock_catalog), )) } } @@ -376,15 +420,80 @@ mod tests { } #[test] - fn approval_review_preferred_model_uses_bedrock_gpt_5_4() { + fn runtime_capabilities_disable_web_search_and_support_v1_remote_compaction() { + let provider = AmazonBedrockModelProvider::new( + ModelProviderInfo::create_amazon_bedrock_runtime_provider(/*aws*/ None), + /*auth_manager*/ None, + ); + + assert_eq!( + provider.capabilities(), + ProviderCapabilities { + namespace_tools: true, + image_generation: false, + web_search: false, + external_web_access: false, + remote_compaction: RemoteCompactionSupport::V1, + } + ); + } + + #[tokio::test] + async fn runtime_managed_auth_resolves_runtime_endpoint() { + let managed_auth = BedrockApiKeyAuth { + api_key: "managed-bedrock-api-key".to_string(), + region: "eu-west-1".to_string(), + }; + let auth_manager = + AuthManager::from_auth_for_testing(CodexAuth::BedrockApiKey(managed_auth)); let provider = AmazonBedrockModelProvider::new( + ModelProviderInfo::create_amazon_bedrock_runtime_provider(/*aws*/ None), + Some(auth_manager), + ); + + assert_eq!( + provider + .runtime_base_url() + .await + .expect("managed Bedrock Runtime region should resolve"), + Some("https://bedrock-runtime.eu-west-1.amazonaws.com/openai/v1".to_string()) + ); + } + + #[test] + fn preferred_background_models_match_bedrock_endpoint() { + let mantle_provider = AmazonBedrockModelProvider::new( ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None), /*auth_manager*/ None, ); + let runtime_provider = AmazonBedrockModelProvider::new( + ModelProviderInfo::create_amazon_bedrock_runtime_provider(/*aws*/ None), + /*auth_manager*/ None, + ); assert_eq!( - provider.approval_review_preferred_model(), - AMAZON_BEDROCK_GPT_5_4_MODEL_ID + ( + mantle_provider.approval_review_preferred_model(), + mantle_provider.memory_extraction_preferred_model(), + mantle_provider.memory_consolidation_preferred_model(), + ), + ( + AMAZON_BEDROCK_GPT_5_6_LUNA_MODEL_ID, + AMAZON_BEDROCK_GPT_5_6_LUNA_MODEL_ID, + AMAZON_BEDROCK_GPT_5_6_TERRA_MODEL_ID, + ) + ); + assert_eq!( + ( + runtime_provider.approval_review_preferred_model(), + runtime_provider.memory_extraction_preferred_model(), + runtime_provider.memory_consolidation_preferred_model(), + ), + ( + AMAZON_BEDROCK_RUNTIME_GLOBAL_GPT_5_6_LUNA_MODEL_ID, + AMAZON_BEDROCK_RUNTIME_GLOBAL_GPT_5_6_LUNA_MODEL_ID, + AMAZON_BEDROCK_RUNTIME_GLOBAL_GPT_5_6_TERRA_MODEL_ID, + ) ); } } diff --git a/codex-rs/model-provider/src/amazon_bedrock/runtime.rs b/codex-rs/model-provider/src/amazon_bedrock/runtime.rs new file mode 100644 index 000000000000..b1de3402bce7 --- /dev/null +++ b/codex-rs/model-provider/src/amazon_bedrock/runtime.rs @@ -0,0 +1,34 @@ +use codex_aws_auth::AwsAuthConfig; +use codex_login::auth::BedrockApiKeyAuth; +use codex_model_provider_info::ModelProviderAwsAuthInfo; +use codex_protocol::error::Result; + +use super::BedrockEndpoint; +use super::auth::resolve_region; +use super::mantle::region_from_config; + +const BEDROCK_RUNTIME_SERVICE_NAME: &str = "bedrock"; + +pub(super) fn aws_auth_config(aws: &ModelProviderAwsAuthInfo) -> AwsAuthConfig { + AwsAuthConfig { + profile: aws.profile.clone(), + region: region_from_config(aws), + service: BEDROCK_RUNTIME_SERVICE_NAME.to_string(), + } +} + +pub(super) fn base_url(region: &str) -> String { + format!("https://bedrock-runtime.{region}.amazonaws.com/openai/v1") +} + +pub(super) async fn bedrock_runtime_base_url( + managed_auth: Option<&BedrockApiKeyAuth>, + aws: &ModelProviderAwsAuthInfo, +) -> Result { + let region = resolve_region(managed_auth, aws, BedrockEndpoint::Runtime).await?; + Ok(base_url(®ion)) +} + +#[cfg(test)] +#[path = "runtime_tests.rs"] +mod tests; diff --git a/codex-rs/model-provider/src/amazon_bedrock/runtime_catalog.rs b/codex-rs/model-provider/src/amazon_bedrock/runtime_catalog.rs new file mode 100644 index 000000000000..0f2766efa049 --- /dev/null +++ b/codex-rs/model-provider/src/amazon_bedrock/runtime_catalog.rs @@ -0,0 +1,41 @@ +use codex_model_provider_info::AMAZON_BEDROCK_GPT_5_6_LUNA_MODEL_ID; +use codex_model_provider_info::AMAZON_BEDROCK_GPT_5_6_SOL_MODEL_ID; +use codex_model_provider_info::AMAZON_BEDROCK_GPT_5_6_TERRA_MODEL_ID; +use codex_protocol::openai_models::ModelsResponse; + +use super::catalog::static_model_catalog; + +const ROUTING_VARIANTS: [(&str, &str, i32); 2] = + [("global.", "Global", 0), ("us.", "US cross-region", 1)]; + +pub(super) fn static_runtime_model_catalog() -> ModelsResponse { + let models = static_model_catalog() + .models + .into_iter() + .filter(|model| { + matches!( + model.slug.as_str(), + AMAZON_BEDROCK_GPT_5_6_SOL_MODEL_ID + | AMAZON_BEDROCK_GPT_5_6_TERRA_MODEL_ID + | AMAZON_BEDROCK_GPT_5_6_LUNA_MODEL_ID + ) + }) + .flat_map(|model| { + ROUTING_VARIANTS + .into_iter() + .map(move |(prefix, routing_label, routing_priority)| { + let mut variant = model.clone(); + variant.slug = format!("{prefix}{}", model.slug); + variant.display_name = format!("{} ({routing_label})", model.display_name); + variant.priority = model.priority * 2 + routing_priority; + variant.supports_search_tool = false; + variant + }) + }) + .collect(); + ModelsResponse { models } +} + +#[cfg(test)] +#[path = "runtime_catalog_tests.rs"] +mod tests; diff --git a/codex-rs/model-provider/src/amazon_bedrock/runtime_catalog_tests.rs b/codex-rs/model-provider/src/amazon_bedrock/runtime_catalog_tests.rs new file mode 100644 index 000000000000..f03eb06187c0 --- /dev/null +++ b/codex-rs/model-provider/src/amazon_bedrock/runtime_catalog_tests.rs @@ -0,0 +1,61 @@ +use pretty_assertions::assert_eq; + +use super::static_runtime_model_catalog; + +#[test] +fn runtime_catalog_includes_supported_cross_region_models_in_priority_order() { + let catalog = static_runtime_model_catalog(); + + assert_eq!( + catalog + .models + .iter() + .map(|model| ( + model.slug.as_str(), + model.display_name.as_str(), + model.priority, + )) + .collect::>(), + vec![ + ("global.openai.gpt-5.6-sol", "GPT-5.6 Sol (Global)", 0), + ("us.openai.gpt-5.6-sol", "GPT-5.6 Sol (US cross-region)", 1), + ("global.openai.gpt-5.6-terra", "GPT-5.6 Terra (Global)", 2), + ( + "us.openai.gpt-5.6-terra", + "GPT-5.6 Terra (US cross-region)", + 3, + ), + ("global.openai.gpt-5.6-luna", "GPT-5.6 Luna (Global)", 4), + ( + "us.openai.gpt-5.6-luna", + "GPT-5.6 Luna (US cross-region)", + 5 + ), + ] + ); +} + +#[test] +fn runtime_catalog_disables_web_search_without_overriding_review_models() { + let catalog = static_runtime_model_catalog(); + + assert_eq!( + catalog + .models + .iter() + .map(|model| ( + model.slug.as_str(), + model.auto_review_model_override.as_deref(), + model.supports_search_tool, + )) + .collect::>(), + vec![ + ("global.openai.gpt-5.6-sol", None, false), + ("us.openai.gpt-5.6-sol", None, false), + ("global.openai.gpt-5.6-terra", None, false), + ("us.openai.gpt-5.6-terra", None, false), + ("global.openai.gpt-5.6-luna", None, false), + ("us.openai.gpt-5.6-luna", None, false), + ] + ); +} diff --git a/codex-rs/model-provider/src/amazon_bedrock/runtime_tests.rs b/codex-rs/model-provider/src/amazon_bedrock/runtime_tests.rs new file mode 100644 index 000000000000..6ba3439481e6 --- /dev/null +++ b/codex-rs/model-provider/src/amazon_bedrock/runtime_tests.rs @@ -0,0 +1,29 @@ +use codex_aws_auth::AwsAuthConfig; +use codex_model_provider_info::ModelProviderAwsAuthInfo; +use pretty_assertions::assert_eq; + +use super::aws_auth_config; +use super::base_url; + +#[test] +fn base_url_uses_regional_runtime_endpoint() { + assert_eq!( + base_url("us-west-2"), + "https://bedrock-runtime.us-west-2.amazonaws.com/openai/v1" + ); +} + +#[test] +fn aws_auth_config_uses_bedrock_service() { + assert_eq!( + aws_auth_config(&ModelProviderAwsAuthInfo { + profile: Some("runtime-profile".to_string()), + region: Some(" eu-west-1 ".to_string()), + }), + AwsAuthConfig { + profile: Some("runtime-profile".to_string()), + region: Some("eu-west-1".to_string()), + service: "bedrock".to_string(), + } + ); +} diff --git a/codex-rs/model-provider/src/lib.rs b/codex-rs/model-provider/src/lib.rs index 6da0053d9d46..b1bd6d21622a 100644 --- a/codex-rs/model-provider/src/lib.rs +++ b/codex-rs/model-provider/src/lib.rs @@ -14,6 +14,7 @@ pub use auth::unauthenticated_auth_provider; pub use bearer_auth_provider::BearerAuthProvider; pub use bearer_auth_provider::BearerAuthProvider as CoreAuthProvider; pub use codex_model_provider_info::AMAZON_BEDROCK_PROVIDER_ID; +pub use codex_model_provider_info::AMAZON_BEDROCK_RUNTIME_PROVIDER_ID; pub use codex_model_provider_info::CHATGPT_CODEX_BASE_URL; pub use codex_protocol::account::ProviderAccount; pub use provider::ModelProvider;