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 authored and copyberry committed 2026-08-14 01:23:59 +00:00
1 parent 3711943d11
commit d5e256ceb2
15 files changed
+643 -55

No files matched your search

+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();