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
This commit is contained in:
Celia Chen
2026-08-14 01:23:59 +00:00
committed by copyberry
parent 3711943d11
commit d5e256ceb2
15 changed files with 643 additions and 55 deletions
@@ -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(())
}
@@ -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;
@@ -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::<ConfigToml>(
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::<ConfigToml>(
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`"
));
}
+17 -5
View File
@@ -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::<Vec<_>>();
@@ -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(()),
+61
View File
@@ -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!();
+33 -5
View File
@@ -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<ModelProviderAwsAuthInfo>,
) -> 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<String, ModelProviderInfo>,
configured_model_providers: HashMap<String, ModelProviderInfo>,
) -> Result<HashMap<String, ModelProviderInfo>, 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 {
@@ -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<_>>(),
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();
@@ -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<BedrockAuthMethod> {
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<SharedAuthProvider> {
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<String> {
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<String, std::env::VarError>,
@@ -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<Request, AuthError> {
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))
}
}
@@ -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<String> {
let region = resolve_region(managed_auth, aws).await?;
let region = resolve_region(managed_auth, aws, BedrockEndpoint::Mantle).await?;
base_url(&region)
}
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;
+125 -16
View File
@@ -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<Arc<AuthManager>>,
}
@@ -50,6 +64,11 @@ impl AmazonBedrockModelProvider {
auth_manager: Option<Arc<AuthManager>>,
) -> 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<SharedAuthProvider> {
@@ -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<Arc<AuthManager>> {
@@ -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_provider(/*aws*/ None),
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
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!(
(
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,
)
);
}
}
@@ -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<String> {
let region = resolve_region(managed_auth, aws, BedrockEndpoint::Runtime).await?;
Ok(base_url(&region))
}
#[cfg(test)]
#[path = "runtime_tests.rs"]
mod tests;
@@ -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;
@@ -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<_>>(),
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<_>>(),
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),
]
);
}
@@ -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(),
}
);
}
+1
View File
@@ -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;