From cefcfe43b93719f97cde6901881896ba45b637d0 Mon Sep 17 00:00:00 2001 From: Celia Chen Date: Mon, 20 Apr 2026 17:54:05 -0700 Subject: [PATCH] feat: add a built-in Amazon Bedrock model provider (#18744) ## Why Codex needs a first-class `amazon-bedrock` model provider so users can select Bedrock without copying a full provider definition into `config.toml`. The provider has Codex-owned defaults for the pieces that should stay consistent across users: the display `name`, Bedrock `base_url`, and `wire_api`. At the same time, users still need a way to choose the AWS credential profile used by their local environment. This change makes `amazon-bedrock` a partially modifiable built-in provider: code owns the provider identity and endpoint defaults, while user config can set `model_providers.amazon-bedrock.aws.profile`. For example: ```toml model_provider = "amazon-bedrock" [model_providers.amazon-bedrock.aws] profile = "codex-bedrock" ``` ## What Changed - Added `amazon-bedrock` to the built-in model provider map with: - `name = "Amazon Bedrock"` - `base_url = "https://bedrock-mantle.us-east-1.api.aws/v1"` - `wire_api = "responses"` - Added AWS provider auth config with a profile-only shape: `model_providers..aws.profile`. - Kept AWS auth config restricted to `amazon-bedrock`; custom providers that set `aws` are rejected. - Allowed `model_providers.amazon-bedrock` through reserved-provider validation so it can act as a partial override. - During config loading, only `aws.profile` is copied from the user-provided `amazon-bedrock` entry onto the built-in provider. Other Bedrock provider fields remain hard-coded by the built-in definition. - Updated the generated config schema for the new provider AWS profile config. --- .../app-server/src/codex_message_processor.rs | 1 + codex-rs/config/src/config_toml.rs | 22 +- codex-rs/config/src/thread_config.rs | 1 + codex-rs/core/config.schema.json | 23 ++- codex-rs/core/src/compact_tests.rs | 1 + codex-rs/core/src/config/config_tests.rs | 103 ++++++++++ codex-rs/core/src/config/mod.rs | 9 +- codex-rs/core/tests/responses_headers.rs | 3 + codex-rs/core/tests/suite/client.rs | 4 + .../core/tests/suite/client_websockets.rs | 1 + .../suite/stream_error_allows_next_turn.rs | 1 + .../core/tests/suite/stream_no_completed.rs | 1 + codex-rs/login/src/auth_env_telemetry.rs | 1 + codex-rs/model-provider-info/src/lib.rs | 111 +++++++++- .../src/model_provider_info_tests.rs | 189 ++++++++++++++++++ codex-rs/models-manager/src/manager_tests.rs | 1 + 16 files changed, 461 insertions(+), 11 deletions(-) diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index 6e0c7ce91..52a3d2fcb 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -10595,6 +10595,7 @@ mod tests { env_key_instructions: None, experimental_bearer_token: None, auth: None, + aws: None, wire_api: WireApi::Responses, query_params: None, http_headers: None, diff --git a/codex-rs/config/src/config_toml.rs b/codex-rs/config/src/config_toml.rs index 35e213bc8..0c8e14c16 100644 --- a/codex-rs/config/src/config_toml.rs +++ b/codex-rs/config/src/config_toml.rs @@ -29,6 +29,7 @@ use crate::types::WindowsToml; use codex_app_server_protocol::Tools; use codex_app_server_protocol::UserSavedConfig; use codex_features::FeaturesToml; +use codex_model_provider_info::AMAZON_BEDROCK_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; @@ -56,7 +57,8 @@ use serde::Deserialize; use serde::Deserializer; use serde::Serialize; -const RESERVED_MODEL_PROVIDER_IDS: [&str; 3] = [ +const RESERVED_MODEL_PROVIDER_IDS: [&str; 4] = [ + AMAZON_BEDROCK_PROVIDER_ID, OPENAI_PROVIDER_ID, OLLAMA_OSS_PROVIDER_ID, LMSTUDIO_OSS_PROVIDER_ID, @@ -780,7 +782,10 @@ pub fn validate_reserved_model_provider_ids( ) -> Result<(), String> { let mut conflicts = model_providers .keys() - .filter(|key| RESERVED_MODEL_PROVIDER_IDS.contains(&key.as_str())) + .filter(|key| { + key.as_str() != AMAZON_BEDROCK_PROVIDER_ID + && RESERVED_MODEL_PROVIDER_IDS.contains(&key.as_str()) + }) .map(|key| format!("`{key}`")) .collect::>(); conflicts.sort_unstable(); @@ -800,6 +805,19 @@ 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 { + continue; + } + if provider.aws.is_some() { + return Err(format!( + "model_providers.{key}: provider aws is only supported for `{AMAZON_BEDROCK_PROVIDER_ID}`" + )); + } + if provider.name.trim().is_empty() { + return Err(format!( + "model_providers.{key}: provider name must not be empty" + )); + } provider .validate() .map_err(|message| format!("model_providers.{key}: {message}"))?; diff --git a/codex-rs/config/src/thread_config.rs b/codex-rs/config/src/thread_config.rs index bfdfa144d..e1e598708 100644 --- a/codex-rs/config/src/thread_config.rs +++ b/codex-rs/config/src/thread_config.rs @@ -298,6 +298,7 @@ mod tests { env_key_instructions: None, experimental_bearer_token: None, auth: None, + aws: None, wire_api: WireApi::Responses, query_params: None, http_headers: None, diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 996e9428e..b0faea2b5 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -1027,6 +1027,17 @@ ], "type": "object" }, + "ModelProviderAwsAuthInfo": { + "additionalProperties": false, + "description": "AWS SigV4 auth configuration for a model provider.", + "properties": { + "profile": { + "description": "AWS profile name to use. When unset, the AWS SDK default chain decides.", + "type": "string" + } + }, + "type": "object" + }, "ModelProviderInfo": { "additionalProperties": false, "description": "Serializable representation of a provider definition.", @@ -1039,6 +1050,14 @@ ], "description": "Command-backed bearer-token configuration for this provider." }, + "aws": { + "allOf": [ + { + "$ref": "#/definitions/ModelProviderAwsAuthInfo" + } + ], + "description": "AWS SigV4 auth configuration for this provider." + }, "base_url": { "description": "Base URL for the provider's OpenAI-compatible API.", "type": "string" @@ -1070,6 +1089,7 @@ "type": "object" }, "name": { + "default": "", "description": "Friendly display name.", "type": "string" }, @@ -1124,9 +1144,6 @@ "description": "Which wire protocol this provider expects." } }, - "required": [ - "name" - ], "type": "object" }, "MultiAgentV2ConfigToml": { diff --git a/codex-rs/core/src/compact_tests.rs b/codex-rs/core/src/compact_tests.rs index ed2fd8487..fbdfdb051 100644 --- a/codex-rs/core/src/compact_tests.rs +++ b/codex-rs/core/src/compact_tests.rs @@ -198,6 +198,7 @@ fn should_use_remote_compact_task_for_azure_provider() { env_key_instructions: None, experimental_bearer_token: None, auth: None, + aws: None, wire_api: WireApi::Responses, query_params: None, http_headers: None, diff --git a/codex-rs/core/src/config/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index 9a6338875..fb48e7316 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -369,6 +369,108 @@ command = "print-token" ); } +#[test] +fn rejects_provider_aws_for_custom_provider() { + let err = toml::from_str::( + r#" +[model_providers.custom] +name = "Custom Provider" + +[model_providers.custom.aws] +profile = "codex-bedrock" +"#, + ) + .unwrap_err(); + + assert!( + err.to_string().contains( + "model_providers.custom: provider aws is only supported for `amazon-bedrock`" + ) + ); +} + +#[test] +fn accepts_amazon_bedrock_aws_profile_override() { + let cfg = toml::from_str::( + r#" +[model_providers.amazon-bedrock.aws] +profile = "codex-bedrock" +"#, + ) + .expect("Amazon Bedrock AWS profile override should deserialize"); + + assert_eq!( + cfg.model_providers + .get("amazon-bedrock") + .and_then(|provider| provider.aws.as_ref()) + .and_then(|aws| aws.profile.as_deref()), + Some("codex-bedrock") + ); +} + +#[tokio::test] +async fn load_config_applies_amazon_bedrock_aws_profile_override() { + let cfg = toml::from_str::( + r#" +model_provider = "amazon-bedrock" + +[model_providers.amazon-bedrock.aws] +profile = "codex-bedrock" +"#, + ) + .expect("Amazon Bedrock AWS profile override should deserialize"); + + let config = Config::load_from_base_config_with_overrides( + cfg, + ConfigOverrides::default(), + tempdir().expect("tempdir").abs(), + ) + .await + .expect("load config"); + + assert_eq!(config.model_provider_id, "amazon-bedrock"); + assert_eq!( + config + .model_provider + .aws + .as_ref() + .and_then(|aws| aws.profile.as_deref()), + Some("codex-bedrock") + ); +} + +#[tokio::test] +async fn load_config_rejects_unsupported_amazon_bedrock_overrides() { + let cfg = toml::from_str::( + r#" +model_provider = "amazon-bedrock" + +[model_providers.amazon-bedrock] +name = "Custom Bedrock" +base_url = "https://bedrock.example.com/v1" +requires_openai_auth = true +supports_websockets = true + +[model_providers.amazon-bedrock.aws] +profile = "codex-bedrock" +"#, + ) + .expect("Amazon Bedrock unsupported overrides should deserialize"); + + let err = Config::load_from_base_config_with_overrides( + cfg, + ConfigOverrides::default(), + tempdir().expect("tempdir").abs(), + ) + .await + .unwrap_err(); + + assert_eq!(err.kind(), std::io::ErrorKind::InvalidData); + assert!(err.to_string().contains( + "model_providers.amazon-bedrock only supports changing `aws.profile`; other non-default provider fields are not supported" + )); +} + #[test] fn config_toml_deserializes_model_availability_nux() { let toml = r#" @@ -4755,6 +4857,7 @@ model_verbosity = "high" env_key_instructions: None, experimental_bearer_token: None, auth: None, + aws: None, query_params: None, http_headers: None, env_http_headers: None, diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 1dd46d94c..2a4171acf 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -65,6 +65,7 @@ use codex_model_provider_info::LEGACY_OLLAMA_CHAT_PROVIDER_ID; use codex_model_provider_info::ModelProviderInfo; use codex_model_provider_info::OLLAMA_CHAT_PROVIDER_REMOVED_ERROR; use codex_model_provider_info::built_in_model_providers; +use codex_model_provider_info::merge_configured_model_providers; use codex_models_manager::ModelsManagerConfig; use codex_protocol::config_types::AltScreenMode; use codex_protocol::config_types::ForcedLoginMethod; @@ -1821,11 +1822,9 @@ impl Config { .clone() .filter(|value| !value.is_empty()); - let mut model_providers = built_in_model_providers(openai_base_url); - // Merge user-defined providers into the built-in list. - for (key, provider) in cfg.model_providers.into_iter() { - model_providers.entry(key).or_insert(provider); - } + let model_providers = + merge_configured_model_providers(built_in_model_providers(openai_base_url), cfg.model_providers) + .map_err(|message| std::io::Error::new(std::io::ErrorKind::InvalidData, message))?; let model_provider_id = model_provider .or(config_profile.model_provider) diff --git a/codex-rs/core/tests/responses_headers.rs b/codex-rs/core/tests/responses_headers.rs index db4dc794b..85b23af31 100644 --- a/codex-rs/core/tests/responses_headers.rs +++ b/codex-rs/core/tests/responses_headers.rs @@ -57,6 +57,7 @@ async fn responses_stream_includes_subagent_header_on_review() { env_key_instructions: None, experimental_bearer_token: None, auth: None, + aws: None, wire_api: WireApi::Responses, query_params: None, http_headers: None, @@ -181,6 +182,7 @@ async fn responses_stream_includes_subagent_header_on_other() { env_key_instructions: None, experimental_bearer_token: None, auth: None, + aws: None, wire_api: WireApi::Responses, query_params: None, http_headers: None, @@ -290,6 +292,7 @@ async fn responses_respects_model_info_overrides_from_config() { env_key_instructions: None, experimental_bearer_token: None, auth: None, + aws: None, wire_api: WireApi::Responses, query_params: None, http_headers: None, diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index 13a6073fa..2086367b2 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -841,6 +841,7 @@ async fn send_provider_auth_request(server: &MockServer, auth: ModelProviderAuth env_key_instructions: None, experimental_bearer_token: None, auth: Some(auth), + aws: None, wire_api: WireApi::Responses, query_params: None, http_headers: None, @@ -2132,6 +2133,7 @@ async fn azure_responses_request_includes_store_and_reasoning_ids() { env_key_instructions: None, experimental_bearer_token: None, auth: None, + aws: None, wire_api: WireApi::Responses, query_params: None, http_headers: None, @@ -2750,6 +2752,7 @@ async fn azure_overrides_assign_properties_used_for_responses_url() { env_key: Some(EXISTING_ENV_VAR_WITH_NON_EMPTY_VALUE.to_string()), experimental_bearer_token: None, auth: None, + aws: None, query_params: Some(std::collections::HashMap::from([( "api-version".to_string(), "2025-04-01-preview".to_string(), @@ -2841,6 +2844,7 @@ async fn env_var_overrides_loaded_auth() { env_key_instructions: None, experimental_bearer_token: None, auth: None, + aws: None, wire_api: WireApi::Responses, http_headers: Some(std::collections::HashMap::from([( "Custom-Header".to_string(), diff --git a/codex-rs/core/tests/suite/client_websockets.rs b/codex-rs/core/tests/suite/client_websockets.rs index 883a244ad..f6cfd0d91 100755 --- a/codex-rs/core/tests/suite/client_websockets.rs +++ b/codex-rs/core/tests/suite/client_websockets.rs @@ -1736,6 +1736,7 @@ fn websocket_provider_with_connect_timeout( env_key_instructions: None, experimental_bearer_token: None, auth: None, + aws: None, wire_api: WireApi::Responses, query_params: None, http_headers: None, diff --git a/codex-rs/core/tests/suite/stream_error_allows_next_turn.rs b/codex-rs/core/tests/suite/stream_error_allows_next_turn.rs index 254ac7b12..950306e97 100644 --- a/codex-rs/core/tests/suite/stream_error_allows_next_turn.rs +++ b/codex-rs/core/tests/suite/stream_error_allows_next_turn.rs @@ -70,6 +70,7 @@ async fn continue_after_stream_error() { env_key_instructions: None, experimental_bearer_token: None, auth: None, + aws: None, wire_api: WireApi::Responses, query_params: None, http_headers: None, diff --git a/codex-rs/core/tests/suite/stream_no_completed.rs b/codex-rs/core/tests/suite/stream_no_completed.rs index 149ba1c53..2dd73e0f6 100644 --- a/codex-rs/core/tests/suite/stream_no_completed.rs +++ b/codex-rs/core/tests/suite/stream_no_completed.rs @@ -54,6 +54,7 @@ async fn retries_on_early_close() { env_key_instructions: None, experimental_bearer_token: None, auth: None, + aws: None, wire_api: WireApi::Responses, query_params: None, http_headers: None, diff --git a/codex-rs/login/src/auth_env_telemetry.rs b/codex-rs/login/src/auth_env_telemetry.rs index f8824268a..3cec5a4ce 100644 --- a/codex-rs/login/src/auth_env_telemetry.rs +++ b/codex-rs/login/src/auth_env_telemetry.rs @@ -65,6 +65,7 @@ mod tests { env_key_instructions: None, experimental_bearer_token: None, auth: None, + aws: None, wire_api: WireApi::Responses, query_params: None, http_headers: None, diff --git a/codex-rs/model-provider-info/src/lib.rs b/codex-rs/model-provider-info/src/lib.rs index 80fb70560..97b1e166d 100644 --- a/codex-rs/model-provider-info/src/lib.rs +++ b/codex-rs/model-provider-info/src/lib.rs @@ -34,6 +34,9 @@ const MAX_REQUEST_MAX_RETRIES: u64 = 100; const OPENAI_PROVIDER_NAME: &str = "OpenAI"; pub const OPENAI_PROVIDER_ID: &str = "openai"; +const AMAZON_BEDROCK_PROVIDER_NAME: &str = "Amazon Bedrock"; +pub const AMAZON_BEDROCK_PROVIDER_ID: &str = "amazon-bedrock"; +pub const AMAZON_BEDROCK_DEFAULT_BASE_URL: &str = "https://bedrock-mantle.us-east-1.api.aws/v1"; const CHAT_WIRE_API_REMOVED_ERROR: &str = "`wire_api = \"chat\"` is no longer supported.\nHow to fix: set `wire_api = \"responses\"` in your provider config.\nMore info: https://github.com/openai/codex/discussions/7782"; pub const LEGACY_OLLAMA_CHAT_PROVIDER_ID: &str = "ollama-chat"; pub const OLLAMA_CHAT_PROVIDER_REMOVED_ERROR: &str = "`ollama-chat` is no longer supported.\nHow to fix: replace `ollama-chat` with `ollama` in `model_provider`, `oss_provider`, or `--local-provider`.\nMore info: https://github.com/openai/codex/discussions/7782"; @@ -71,10 +74,11 @@ impl<'de> Deserialize<'de> for WireApi { } /// Serializable representation of a provider definition. -#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema)] +#[derive(Debug, Clone, Default, Deserialize, Serialize, PartialEq, JsonSchema)] #[schemars(deny_unknown_fields)] pub struct ModelProviderInfo { /// Friendly display name. + #[serde(default)] pub name: String, /// Base URL for the provider's OpenAI-compatible API. pub base_url: Option, @@ -90,6 +94,8 @@ pub struct ModelProviderInfo { pub experimental_bearer_token: Option, /// Command-backed bearer-token configuration for this provider. pub auth: Option, + /// AWS SigV4 auth configuration for this provider. + pub aws: Option, /// Which wire protocol this provider expects. #[serde(default)] pub wire_api: WireApi, @@ -124,8 +130,46 @@ pub struct ModelProviderInfo { pub supports_websockets: bool, } +/// AWS SigV4 auth configuration for a model provider. +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, JsonSchema)] +#[schemars(deny_unknown_fields)] +pub struct ModelProviderAwsAuthInfo { + /// AWS profile name to use. When unset, the AWS SDK default chain decides. + pub profile: Option, +} + impl ModelProviderInfo { pub fn validate(&self) -> std::result::Result<(), String> { + if self.aws.is_some() { + if self.supports_websockets { + // TODO(celia-oai): Support AWS SigV4 signing for WebSocket + // upgrade requests before allowing AWS-authenticated providers + // to enable Responses-over-WebSocket. + return Err("provider aws cannot be combined with supports_websockets".to_string()); + } + + let mut conflicts = Vec::new(); + if self.env_key.is_some() { + conflicts.push("env_key"); + } + if self.experimental_bearer_token.is_some() { + conflicts.push("experimental_bearer_token"); + } + if self.auth.is_some() { + conflicts.push("auth"); + } + if self.requires_openai_auth { + conflicts.push("requires_openai_auth"); + } + + if !conflicts.is_empty() { + return Err(format!( + "provider aws cannot be combined with {}", + conflicts.join(", ") + )); + } + } + let Some(auth) = self.auth.as_ref() else { return Ok(()); }; @@ -269,6 +313,7 @@ impl ModelProviderInfo { env_key_instructions: None, experimental_bearer_token: None, auth: None, + aws: None, wire_api: WireApi::Responses, query_params: None, http_headers: Some( @@ -297,10 +342,38 @@ impl ModelProviderInfo { } } + pub fn create_amazon_bedrock_provider( + aws: Option, + ) -> ModelProviderInfo { + ModelProviderInfo { + name: AMAZON_BEDROCK_PROVIDER_NAME.into(), + base_url: Some(AMAZON_BEDROCK_DEFAULT_BASE_URL.into()), + env_key: None, + env_key_instructions: None, + experimental_bearer_token: None, + auth: None, + aws: Some(aws.unwrap_or(ModelProviderAwsAuthInfo { profile: None })), + wire_api: WireApi::Responses, + query_params: None, + http_headers: None, + env_http_headers: None, + request_max_retries: None, + stream_max_retries: None, + stream_idle_timeout_ms: None, + websocket_connect_timeout_ms: None, + requires_openai_auth: false, + supports_websockets: false, + } + } + pub fn is_openai(&self) -> bool { self.name == OPENAI_PROVIDER_NAME } + pub fn is_amazon_bedrock(&self) -> bool { + self.name == AMAZON_BEDROCK_PROVIDER_NAME + } + pub fn supports_remote_compaction(&self) -> bool { self.is_openai() || is_azure_responses_provider(&self.name, self.base_url.as_deref()) } @@ -322,6 +395,7 @@ pub fn built_in_model_providers( ) -> HashMap { 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); // 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 @@ -329,6 +403,7 @@ pub fn built_in_model_providers( // `model_providers` in config.toml to add their own providers. [ (OPENAI_PROVIDER_ID, openai_provider), + (AMAZON_BEDROCK_PROVIDER_ID, amazon_bedrock_provider), ( OLLAMA_OSS_PROVIDER_ID, create_oss_provider(DEFAULT_OLLAMA_PORT, WireApi::Responses), @@ -343,6 +418,39 @@ pub fn built_in_model_providers( .collect() } +/// 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 set `aws.profile`. +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 { + let aws_override = provider.aws.take(); + if provider != ModelProviderInfo::default() { + return Err(format!( + "model_providers.{AMAZON_BEDROCK_PROVIDER_ID} only supports changing \ +`aws.profile`; other non-default provider fields are not supported" + )); + } + + if let Some(profile) = aws_override.and_then(|aws| aws.profile) + && let Some(built_in) = model_providers.get_mut(AMAZON_BEDROCK_PROVIDER_ID) + && let Some(aws) = built_in.aws.as_mut() + { + aws.profile = Some(profile); + } + } else { + model_providers.entry(key).or_insert(provider); + } + } + + Ok(model_providers) +} + pub fn create_oss_provider(default_provider_port: u16, wire_api: WireApi) -> ModelProviderInfo { // These CODEX_OSS_ environment variables are experimental: we may // switch to reading values from config.toml instead. @@ -370,6 +478,7 @@ pub fn create_oss_provider_with_base_url(base_url: &str, wire_api: WireApi) -> M env_key_instructions: None, experimental_bearer_token: None, auth: None, + aws: None, wire_api, query_params: None, http_headers: None, 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 6cd1ef236..20440de32 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 @@ -18,6 +18,7 @@ base_url = "http://localhost:11434/v1" env_key_instructions: None, experimental_bearer_token: None, auth: None, + aws: None, wire_api: WireApi::Responses, query_params: None, http_headers: None, @@ -49,6 +50,7 @@ query_params = { api-version = "2025-04-01-preview" } env_key_instructions: None, experimental_bearer_token: None, auth: None, + aws: None, wire_api: WireApi::Responses, query_params: Some(maplit::hashmap! { "api-version".to_string() => "2025-04-01-preview".to_string(), @@ -83,6 +85,7 @@ env_http_headers = { "X-Example-Env-Header" = "EXAMPLE_ENV_VAR" } env_key_instructions: None, experimental_bearer_token: None, auth: None, + aws: None, wire_api: WireApi::Responses, query_params: None, http_headers: Some(maplit::hashmap! { @@ -145,6 +148,7 @@ fn test_supports_remote_compaction_for_azure_name() { env_key_instructions: None, experimental_bearer_token: None, auth: None, + aws: None, wire_api: WireApi::Responses, query_params: None, http_headers: None, @@ -169,6 +173,7 @@ fn test_supports_remote_compaction_for_non_openai_non_azure_provider() { env_key_instructions: None, experimental_bearer_token: None, auth: None, + aws: None, wire_api: WireApi::Responses, query_params: None, http_headers: None, @@ -212,6 +217,190 @@ args = ["--format=text"] ); } +#[test] +fn test_deserialize_provider_aws_config() { + let provider_toml = r#" +name = "Amazon Bedrock" +base_url = "https://bedrock.example.com/v1" + +[aws] +profile = "codex-bedrock" + "#; + + let provider: ModelProviderInfo = toml::from_str(provider_toml).unwrap(); + + assert_eq!( + provider.aws, + Some(ModelProviderAwsAuthInfo { + profile: Some("codex-bedrock".to_string()), + }) + ); +} + +#[test] +fn test_create_amazon_bedrock_provider() { + assert_eq!( + ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None), + ModelProviderInfo { + name: "Amazon Bedrock".to_string(), + base_url: Some("https://bedrock-mantle.us-east-1.api.aws/v1".to_string()), + env_key: None, + env_key_instructions: None, + experimental_bearer_token: None, + auth: None, + aws: Some(ModelProviderAwsAuthInfo { profile: None }), + wire_api: WireApi::Responses, + query_params: None, + http_headers: None, + env_http_headers: None, + request_max_retries: None, + stream_max_retries: None, + stream_idle_timeout_ms: None, + websocket_connect_timeout_ms: None, + requires_openai_auth: false, + supports_websockets: false, + } + ); +} + +#[test] +fn test_built_in_model_providers_include_amazon_bedrock() { + let providers = built_in_model_providers(/*openai_base_url*/ None); + + assert_eq!( + providers + .get(AMAZON_BEDROCK_PROVIDER_ID) + .map(ModelProviderInfo::is_amazon_bedrock), + Some(true) + ); +} + +#[test] +fn test_merge_configured_model_providers_adds_custom_provider() { + let custom_provider = ModelProviderInfo { + name: "Custom".to_string(), + base_url: Some("https://example.com/v1".to_string()), + ..ModelProviderInfo::default() + }; + let configured_model_providers = + std::collections::HashMap::from([("custom".to_string(), custom_provider.clone())]); + + let mut expected = built_in_model_providers(/*openai_base_url*/ None); + expected.insert("custom".to_string(), custom_provider); + + 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_profile_override() { + let configured_model_providers = std::collections::HashMap::from([( + AMAZON_BEDROCK_PROVIDER_ID.to_string(), + ModelProviderInfo { + aws: Some(ModelProviderAwsAuthInfo { + profile: Some("codex-bedrock".to_string()), + }), + ..ModelProviderInfo::default() + }, + )]); + + let mut expected = built_in_model_providers(/*openai_base_url*/ None); + expected + .get_mut(AMAZON_BEDROCK_PROVIDER_ID) + .expect("Amazon Bedrock provider should be built in") + .aws = Some(ModelProviderAwsAuthInfo { + profile: Some("codex-bedrock".to_string()), + }); + + 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_rejects_amazon_bedrock_non_default_fields() { + let configured_model_providers = std::collections::HashMap::from([( + AMAZON_BEDROCK_PROVIDER_ID.to_string(), + ModelProviderInfo { + name: "Custom Bedrock".to_string(), + aws: Some(ModelProviderAwsAuthInfo { + profile: Some("codex-bedrock".to_string()), + }), + ..ModelProviderInfo::default() + }, + )]); + + assert_eq!( + merge_configured_model_providers( + built_in_model_providers(/*openai_base_url*/ None), + configured_model_providers, + ), + Err( + "model_providers.amazon-bedrock only supports changing `aws.profile`; other non-default provider fields are not supported" + .to_string() + ) + ); +} + +#[test] +fn test_merge_configured_model_providers_allows_amazon_bedrock_default_fields() { + let configured_model_providers = std::collections::HashMap::from([( + AMAZON_BEDROCK_PROVIDER_ID.to_string(), + ModelProviderInfo { + aws: Some(ModelProviderAwsAuthInfo { profile: None }), + wire_api: WireApi::Responses, + ..ModelProviderInfo::default() + }, + )]); + + assert_eq!( + merge_configured_model_providers( + built_in_model_providers(/*openai_base_url*/ None), + configured_model_providers, + ), + Ok(built_in_model_providers(/*openai_base_url*/ None)) + ); +} + +#[test] +fn test_validate_provider_aws_rejects_conflicting_auth() { + let provider = ModelProviderInfo { + aws: Some(ModelProviderAwsAuthInfo { profile: None }), + env_key: Some("AWS_BEARER_TOKEN_BEDROCK".to_string()), + supports_websockets: false, + ..ModelProviderInfo::create_openai_provider(/*base_url*/ None) + }; + + assert_eq!( + provider.validate(), + Err("provider aws cannot be combined with env_key, requires_openai_auth".to_string()) + ); +} + +#[test] +fn test_validate_provider_aws_rejects_websockets() { + let provider = ModelProviderInfo { + aws: Some(ModelProviderAwsAuthInfo { profile: None }), + requires_openai_auth: false, + supports_websockets: true, + ..ModelProviderInfo::create_openai_provider(/*base_url*/ None) + }; + + assert_eq!( + provider.validate(), + Err("provider aws cannot be combined with supports_websockets".to_string()) + ); +} + #[test] fn test_deserialize_provider_auth_config_allows_zero_refresh_interval() { let base_dir = tempdir().unwrap(); diff --git a/codex-rs/models-manager/src/manager_tests.rs b/codex-rs/models-manager/src/manager_tests.rs index a70f46567..d004dd894 100644 --- a/codex-rs/models-manager/src/manager_tests.rs +++ b/codex-rs/models-manager/src/manager_tests.rs @@ -94,6 +94,7 @@ fn provider_for(base_url: String) -> ModelProviderInfo { env_key_instructions: None, experimental_bearer_token: None, auth: None, + aws: None, wire_api: WireApi::Responses, query_params: None, http_headers: None,