diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json index 2f2dde5afd3d..b79ab0a35411 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json @@ -5997,20 +5997,16 @@ }, { "properties": { - "credentialSource": { - "allOf": [ - { - "$ref": "#/definitions/v2/AmazonBedrockCredentialSource" - } - ], - "default": "awsManaged" - }, "type": { "enum": [ "amazonBedrock" ], "title": "AmazonBedrockAccountType", "type": "string" + }, + "usesCodexManagedCredentials": { + "default": false, + "type": "boolean" } }, "required": [ @@ -6324,13 +6320,6 @@ "AgentPath": { "type": "string" }, - "AmazonBedrockCredentialSource": { - "enum": [ - "codexManaged", - "awsManaged" - ], - "type": "string" - }, "AnalyticsConfig": { "additionalProperties": true, "properties": { diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json index 93e9ea5edc5f..413d55a506d5 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json @@ -52,20 +52,16 @@ }, { "properties": { - "credentialSource": { - "allOf": [ - { - "$ref": "#/definitions/AmazonBedrockCredentialSource" - } - ], - "default": "awsManaged" - }, "type": { "enum": [ "amazonBedrock" ], "title": "AmazonBedrockAccountType", "type": "string" + }, + "usesCodexManagedCredentials": { + "default": false, + "type": "boolean" } }, "required": [ @@ -379,13 +375,6 @@ "AgentPath": { "type": "string" }, - "AmazonBedrockCredentialSource": { - "enum": [ - "codexManaged", - "awsManaged" - ], - "type": "string" - }, "AnalyticsConfig": { "additionalProperties": true, "properties": { diff --git a/codex-rs/app-server-protocol/schema/json/v2/GetAccountResponse.json b/codex-rs/app-server-protocol/schema/json/v2/GetAccountResponse.json index ca641799edb6..89d8827d738a 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/GetAccountResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/GetAccountResponse.json @@ -48,20 +48,16 @@ }, { "properties": { - "credentialSource": { - "allOf": [ - { - "$ref": "#/definitions/AmazonBedrockCredentialSource" - } - ], - "default": "awsManaged" - }, "type": { "enum": [ "amazonBedrock" ], "title": "AmazonBedrockAccountType", "type": "string" + }, + "usesCodexManagedCredentials": { + "default": false, + "type": "boolean" } }, "required": [ @@ -72,13 +68,6 @@ } ] }, - "AmazonBedrockCredentialSource": { - "enum": [ - "codexManaged", - "awsManaged" - ], - "type": "string" - }, "PlanType": { "enum": [ "free", diff --git a/codex-rs/app-server-protocol/schema/typescript/AmazonBedrockCredentialSource.ts b/codex-rs/app-server-protocol/schema/typescript/AmazonBedrockCredentialSource.ts deleted file mode 100644 index 7822a2f8b427..000000000000 --- a/codex-rs/app-server-protocol/schema/typescript/AmazonBedrockCredentialSource.ts +++ /dev/null @@ -1,5 +0,0 @@ -// GENERATED CODE! DO NOT MODIFY BY HAND! - -// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. - -export type AmazonBedrockCredentialSource = "codexManaged" | "awsManaged"; diff --git a/codex-rs/app-server-protocol/schema/typescript/index.ts b/codex-rs/app-server-protocol/schema/typescript/index.ts index 30f145590c16..77f4af946632 100644 --- a/codex-rs/app-server-protocol/schema/typescript/index.ts +++ b/codex-rs/app-server-protocol/schema/typescript/index.ts @@ -3,7 +3,6 @@ export type { AbsolutePathBuf } from "./AbsolutePathBuf"; export type { AgentMessageInputContent } from "./AgentMessageInputContent"; export type { AgentPath } from "./AgentPath"; -export type { AmazonBedrockCredentialSource } from "./AmazonBedrockCredentialSource"; export type { ApplyPatchApprovalParams } from "./ApplyPatchApprovalParams"; export type { ApplyPatchApprovalResponse } from "./ApplyPatchApprovalResponse"; export type { AuthMode } from "./AuthMode"; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/Account.ts b/codex-rs/app-server-protocol/schema/typescript/v2/Account.ts index 1b7953e5df18..1f1ad851c757 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/Account.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/Account.ts @@ -1,7 +1,6 @@ // GENERATED CODE! DO NOT MODIFY BY HAND! // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. -import type { AmazonBedrockCredentialSource } from "../AmazonBedrockCredentialSource"; import type { PlanType } from "../PlanType"; -export type Account = { "type": "apiKey", } | { "type": "chatgpt", email: string | null, planType: PlanType, } | { "type": "amazonBedrock", credentialSource: AmazonBedrockCredentialSource, }; +export type Account = { "type": "apiKey", } | { "type": "chatgpt", email: string | null, planType: PlanType, } | { "type": "amazonBedrock", usesCodexManagedCredentials: boolean, }; diff --git a/codex-rs/app-server-protocol/src/protocol/common.rs b/codex-rs/app-server-protocol/src/protocol/common.rs index 6199e779bfd5..ab7609912c63 100644 --- a/codex-rs/app-server-protocol/src/protocol/common.rs +++ b/codex-rs/app-server-protocol/src/protocol/common.rs @@ -1755,7 +1755,6 @@ mod tests { use super::*; use anyhow::Result; use codex_protocol::ThreadId; - use codex_protocol::account::AmazonBedrockCredentialSource; use codex_protocol::account::PlanType; use codex_protocol::config_types::MultiAgentMode; use codex_protocol::models::BUILT_IN_PERMISSION_PROFILE_READ_ONLY; @@ -2995,35 +2994,35 @@ mod tests { ); let codex_managed_bedrock = v2::Account::AmazonBedrock { - credential_source: AmazonBedrockCredentialSource::CodexManaged, + uses_codex_managed_credentials: true, }; assert_eq!( json!({ "type": "amazonBedrock", - "credentialSource": "codexManaged", + "usesCodexManagedCredentials": true, }), serde_json::to_value(&codex_managed_bedrock)?, ); - let aws_managed_bedrock = v2::Account::AmazonBedrock { - credential_source: AmazonBedrockCredentialSource::AwsManaged, + let externally_managed_bedrock = v2::Account::AmazonBedrock { + uses_codex_managed_credentials: false, }; assert_eq!( json!({ "type": "amazonBedrock", - "credentialSource": "awsManaged", + "usesCodexManagedCredentials": false, }), - serde_json::to_value(&aws_managed_bedrock)?, + serde_json::to_value(&externally_managed_bedrock)?, ); Ok(()) } #[test] - fn account_defaults_legacy_bedrock_credential_source() -> Result<()> { + fn account_defaults_legacy_bedrock_managed_credentials_flag() -> Result<()> { assert_eq!( v2::Account::AmazonBedrock { - credential_source: AmazonBedrockCredentialSource::AwsManaged, + uses_codex_managed_credentials: false, }, serde_json::from_value(json!({ "type": "amazonBedrock", diff --git a/codex-rs/app-server-protocol/src/protocol/v2/account.rs b/codex-rs/app-server-protocol/src/protocol/v2/account.rs index 9a0f796d207d..1ef94f9d818b 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/account.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/account.rs @@ -1,6 +1,5 @@ use crate::protocol::common::AuthMode; use codex_experimental_api_macros::ExperimentalApi; -use codex_protocol::account::AmazonBedrockCredentialSource; use codex_protocol::account::PlanType; use codex_protocol::account::ProviderAccount; use codex_protocol::protocol::CreditsSnapshot as CoreCreditsSnapshot; @@ -34,8 +33,8 @@ pub enum Account { #[serde(rename = "amazonBedrock", rename_all = "camelCase")] #[ts(rename = "amazonBedrock", rename_all = "camelCase")] AmazonBedrock { - #[serde(default = "default_bedrock_credential_source")] - credential_source: AmazonBedrockCredentialSource, + #[serde(default)] + uses_codex_managed_credentials: bool, }, } @@ -45,18 +44,16 @@ fn nullable_string_schema( generator.subschema_for::>() } -fn default_bedrock_credential_source() -> AmazonBedrockCredentialSource { - AmazonBedrockCredentialSource::AwsManaged -} - impl From for Account { fn from(account: ProviderAccount) -> Self { match account { ProviderAccount::ApiKey => Self::ApiKey {}, ProviderAccount::Chatgpt { email, plan_type } => Self::Chatgpt { email, plan_type }, - ProviderAccount::AmazonBedrock { credential_source } => { - Self::AmazonBedrock { credential_source } - } + ProviderAccount::AmazonBedrock { + uses_codex_managed_credentials, + } => Self::AmazonBedrock { + uses_codex_managed_credentials, + }, } } } diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 38527a786826..f00c9d051b66 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -2010,13 +2010,8 @@ Request: Response examples: ```json -{ "id": 1, "result": { "account": null, "requiresOpenaiAuth": false } } // No OpenAI auth needed (e.g., OSS/local models) -{ "id": 1, "result": { "account": null, "requiresOpenaiAuth": true } } // OpenAI auth required (typical for OpenAI-hosted models) -{ "id": 1, "result": { "account": { "type": "apiKey" }, "requiresOpenaiAuth": true } } { "id": 1, "result": { "account": { "type": "chatgpt", "email": "user@example.com", "planType": "pro" }, "requiresOpenaiAuth": true } } -{ "id": 1, "result": { "account": { "type": "chatgpt", "email": null, "planType": "enterprise" }, "requiresOpenaiAuth": true } } -{ "id": 1, "result": { "account": { "type": "amazonBedrock", "credentialSource": "codexManaged" }, "requiresOpenaiAuth": false } } -{ "id": 1, "result": { "account": { "type": "amazonBedrock", "credentialSource": "awsManaged" }, "requiresOpenaiAuth": false } } +{ "id": 1, "result": { "account": { "type": "amazonBedrock", "usesCodexManagedCredentials": false }, "requiresOpenaiAuth": false } } ``` Field notes: @@ -2024,7 +2019,7 @@ Field notes: - `refreshToken` (bool): set `true` to force a token refresh. - `email` is `null` when the ChatGPT account does not have an email address. - `requiresOpenaiAuth` reflects the active provider; when `false`, Codex can run without OpenAI credentials. -- Amazon Bedrock reports `credentialSource: "codexManaged"` when it uses a Bedrock API key managed by Codex. Otherwise it reports `credentialSource: "awsManaged"` for the external AWS credential path. This identifies the selected credential source; it does not validate that the AWS credential chain can resolve credentials. +- Amazon Bedrock reports `usesCodexManagedCredentials: true` when it uses a Bedrock API key managed by Codex. It reports `false` for external credential paths, including the AWS credential chain and configured command auth. This identifies whether Codex-managed credentials are selected; it does not validate that the credential source can resolve credentials. ### 2) Log in with an API key diff --git a/codex-rs/app-server/tests/suite/v2/account.rs b/codex-rs/app-server/tests/suite/v2/account.rs index 09a2bf7d95b9..79dd824e0ce2 100644 --- a/codex-rs/app-server/tests/suite/v2/account.rs +++ b/codex-rs/app-server/tests/suite/v2/account.rs @@ -47,7 +47,6 @@ use codex_login::auth::BedrockApiKeyAuth; use codex_login::load_auth_dot_json; use codex_login::login_with_api_key; use codex_login::login_with_bedrock_api_key; -use codex_protocol::account::AmazonBedrockCredentialSource; use codex_protocol::account::PlanType as AccountPlanType; use codex_protocol::auth::AuthMode as DomainAuthMode; use core_test_support::responses; @@ -1555,7 +1554,7 @@ async fn login_managed_bedrock_updates_active_bedrock_account() -> Result<()> { read_account(&mut mcp).await?, GetAccountResponse { account: Some(Account::AmazonBedrock { - credential_source: AmazonBedrockCredentialSource::CodexManaged, + uses_codex_managed_credentials: true, }), requires_openai_auth: false, } @@ -2583,7 +2582,7 @@ region = "us-west-2" let expected = GetAccountResponse { account: Some(Account::AmazonBedrock { - credential_source: AmazonBedrockCredentialSource::AwsManaged, + uses_codex_managed_credentials: false, }), requires_openai_auth: false, }; @@ -2591,6 +2590,45 @@ region = "us-west-2" Ok(()) } +#[tokio::test] +async fn get_account_with_user_managed_bedrock_provider() -> Result<()> { + let codex_home = TempDir::new()?; + create_config_toml( + codex_home.path(), + CreateConfigTomlParams { + model_provider_id: Some("amazon-bedrock".to_string()), + extra_provider_config: Some( + r#"[model_providers.amazon-bedrock] +base_url = "https://bedrock.example.com/v1" + +[model_providers.amazon-bedrock.auth] +command = "print-token" +"# + .to_string(), + ), + ..Default::default() + }, + )?; + + let mut mcp = TestAppServer::builder() + .with_codex_home(codex_home.path()) + .without_auto_env() + .build() + .await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + assert_eq!( + read_account(&mut mcp).await?, + GetAccountResponse { + account: Some(Account::AmazonBedrock { + uses_codex_managed_credentials: false, + }), + requires_openai_auth: false, + } + ); + Ok(()) +} + #[tokio::test] async fn account_reads_use_startup_config_when_config_reload_fails() -> Result<()> { let codex_home = TempDir::new()?; @@ -2622,7 +2660,7 @@ region = "us-west-2" read_account(&mut mcp).await?, GetAccountResponse { account: Some(Account::AmazonBedrock { - credential_source: AmazonBedrockCredentialSource::AwsManaged, + uses_codex_managed_credentials: false, }), requires_openai_auth: false, } @@ -2692,7 +2730,7 @@ async fn get_account_with_managed_bedrock_provider() -> Result<()> { received, GetAccountResponse { account: Some(Account::AmazonBedrock { - credential_source: AmazonBedrockCredentialSource::CodexManaged, + uses_codex_managed_credentials: true, }), requires_openai_auth: false, } diff --git a/codex-rs/config/src/config_toml.rs b/codex-rs/config/src/config_toml.rs index 3ac2fa3c3a4b..62b5606bd3a4 100644 --- a/codex-rs/config/src/config_toml.rs +++ b/codex-rs/config/src/config_toml.rs @@ -923,18 +923,17 @@ 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" - )); + if key != AMAZON_BEDROCK_PROVIDER_ID { + 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() @@ -1021,4 +1020,21 @@ mod tests { assert!(message.contains("TOML list of strings")); assert!(message.contains("comma-separated strings are not supported")); } + + #[test] + fn amazon_bedrock_auth_command_must_not_be_empty() { + let err = toml::from_str::( + r#" +[model_providers.amazon-bedrock.auth] +command = " " +"#, + ) + .expect_err("empty Amazon Bedrock auth command should be rejected"); + + assert!( + err.to_string().contains( + "model_providers.amazon-bedrock: provider auth.command must not be empty" + ) + ); + } } diff --git a/codex-rs/core/src/config/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index 7d878e84aa3c..23980fa9ce67 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -75,6 +75,7 @@ use codex_model_provider_info::OLLAMA_OSS_PROVIDER_ID; use codex_model_provider_info::WireApi; use codex_models_manager::bundled_models_response; use codex_network_proxy::NetworkMode; +use codex_protocol::config_types::ModelProviderAuthInfo; use codex_protocol::config_types::SERVICE_TIER_DEFAULT_REQUEST_VALUE; use codex_protocol::config_types::ServiceTier; use codex_protocol::models::ActivePermissionProfile; @@ -844,6 +845,53 @@ region = "us-west-2" ); } +#[tokio::test] +async fn load_config_applies_amazon_bedrock_transport_overrides() { + let cfg = toml::from_str::( + r#" +model_provider = "amazon-bedrock" + +[model_providers.amazon-bedrock] +base_url = "https://bedrock.example.com/v1" +http_headers = { "X-Custom-Header" = "value" } + +[model_providers.amazon-bedrock.auth] +command = "print-token" +"#, + ) + .expect("Amazon Bedrock transport overrides should deserialize"); + + let config = Config::load_from_base_config_with_overrides( + cfg, + ConfigOverrides::default(), + tempdir().expect("tempdir").abs(), + ) + .await + .expect("load config"); + + let mut expected_provider = built_in_model_providers(/*openai_base_url*/ None) + .remove("amazon-bedrock") + .expect("Amazon Bedrock provider should be built in"); + expected_provider.base_url = Some("https://bedrock.example.com/v1".to_string()); + expected_provider.auth = Some(ModelProviderAuthInfo { + command: "print-token".to_string(), + args: Vec::new(), + timeout_ms: std::num::NonZeroU64::new(5_000).expect("timeout should be non-zero"), + refresh_interval_ms: 300_000, + cwd: std::env::current_dir() + .expect("current directory should be available") + .try_into() + .expect("current directory should be absolute"), + }); + expected_provider + .http_headers + .get_or_insert_default() + .insert("X-Custom-Header".to_string(), "value".to_string()); + + assert_eq!(config.model_provider_id, "amazon-bedrock"); + assert_eq!(config.model_provider, expected_provider); +} + #[tokio::test] async fn load_config_rejects_unsupported_amazon_bedrock_overrides() { let cfg = toml::from_str::( @@ -852,13 +900,8 @@ 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" -region = "us-west-2" "#, ) .expect("Amazon Bedrock unsupported overrides should deserialize"); @@ -873,7 +916,7 @@ region = "us-west-2" assert_eq!(err.kind(), std::io::ErrorKind::InvalidData); assert!(err.to_string().contains( - "model_providers.amazon-bedrock only supports changing `aws.profile` and `aws.region`; other non-default provider fields are not supported" + "model_providers.amazon-bedrock only supports changing `base_url`, `auth`, `http_headers`, `aws.profile`, and `aws.region`; other non-default provider fields are not supported" )); } diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index c73ccf4932ad..827708516d4a 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -14,6 +14,7 @@ use codex_login::AuthManager; use codex_login::CodexAuth; use codex_login::auth::AgentIdentityAuthPolicy; use codex_login::default_client::originator; +use codex_model_provider_info::AMAZON_BEDROCK_PROVIDER_ID; use codex_model_provider_info::ModelProviderInfo; use codex_model_provider_info::WireApi; use codex_model_provider_info::built_in_model_providers; @@ -1282,11 +1283,49 @@ async fn provider_auth_command_refreshes_after_401() { send_provider_auth_request(&server, auth_fixture.auth()).await; } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn amazon_bedrock_proxy_uses_command_auth_and_custom_headers() { + skip_if_no_network!(); + + let server = MockServer::start().await; + let response = mount_sse_once( + &server, + sse(vec![ev_response_created("resp1"), ev_completed("resp1")]), + ) + .await; + let auth_fixture = ProviderAuthCommandFixture::new(&["command-token"]).unwrap(); + let mut provider = built_in_model_providers(/*openai_base_url*/ None) + .remove(AMAZON_BEDROCK_PROVIDER_ID) + .expect("Amazon Bedrock provider should be built in"); + provider.base_url = Some(format!("{}/v1", server.uri())); + provider.auth = Some(auth_fixture.auth()); + provider.aws = None; + provider + .http_headers + .get_or_insert_default() + .insert("x-some-header".to_string(), "foo".to_string()); + + send_request_with_provider(provider).await; + + let request = response.single_request(); + assert_eq!(request.path(), "/v1/responses"); + assert_eq!( + request.header("authorization"), + Some("Bearer command-token".to_string()) + ); + assert_eq!(request.header("x-amz-date"), None); + assert_eq!(request.header("x-some-header"), Some("foo".to_string())); + assert_eq!( + request.header("x-amzn-mantle-client-agent"), + Some("codex".to_string()) + ); + assert_eq!(request.body_json()["store"], false); +} + /// Issues one streamed Responses request through a provider configured with command-backed auth. /// /// The caller owns the server-side assertions, so this helper only validates that the request /// reaches `Completed` without surfacing an auth or transport error to the client. -#[expect(clippy::unwrap_used)] async fn send_provider_auth_request(server: &MockServer, auth: ModelProviderAuthInfo) { let provider = ModelProviderInfo { name: "corp".into(), @@ -1308,6 +1347,11 @@ async fn send_provider_auth_request(server: &MockServer, auth: ModelProviderAuth supports_websockets: false, }; + send_request_with_provider(provider).await; +} + +#[expect(clippy::unwrap_used)] +async fn send_request_with_provider(provider: ModelProviderInfo) { let codex_home = TempDir::new().unwrap(); let mut config = load_default_config_for_test(&codex_home).await; config.model_provider_id = provider.name.clone(); diff --git a/codex-rs/model-provider-info/src/lib.rs b/codex-rs/model-provider-info/src/lib.rs index 61cd717d588e..803d858789c9 100644 --- a/codex-rs/model-provider-info/src/lib.rs +++ b/codex-rs/model-provider-info/src/lib.rs @@ -368,7 +368,10 @@ impl ModelProviderInfo { ) -> ModelProviderInfo { ModelProviderInfo { name: AMAZON_BEDROCK_PROVIDER_NAME.into(), - base_url: Some(AMAZON_BEDROCK_DEFAULT_BASE_URL.into()), + // The runtime provider derives the regional Mantle endpoint when + // this is unset. A configured value is therefore unambiguously an + // endpoint override. + base_url: None, env_key: None, env_key_instructions: None, experimental_bearer_token: None, @@ -459,30 +462,36 @@ pub fn built_in_model_providers( /// /// 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` and `aws.region`. +/// user to customize its 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 { + 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 \ -`aws.profile` and `aws.region`; other non-default provider fields are not supported" +`base_url`, `auth`, `http_headers`, `aws.profile`, and `aws.region`; other non-default \ +provider fields are not supported" )); } - if let Some(aws_override) = aws_override - && let Some(built_in_provider) = model_providers.get_mut(AMAZON_BEDROCK_PROVIDER_ID) - && let Some(built_in_aws) = built_in_provider.aws.as_mut() - { - if let Some(profile) = aws_override.profile { - built_in_aws.profile = Some(profile); + if let Some(built_in_provider) = model_providers.get_mut(AMAZON_BEDROCK_PROVIDER_ID) { + built_in_provider.base_url = base_url_override; + built_in_provider.auth = auth_override; + if let Some(aws_override) = aws_override { + built_in_provider.aws = Some(aws_override); } - if let Some(region) = aws_override.region { - built_in_aws.region = Some(region); + if let Some(http_headers_override) = http_headers_override { + built_in_provider + .http_headers + .get_or_insert_default() + .extend(http_headers_override); } } } else { 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 9dbb23ecc19d..cc7f6851f298 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 @@ -288,7 +288,7 @@ fn test_create_amazon_bedrock_provider() { ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None), ModelProviderInfo { name: "Amazon Bedrock".to_string(), - base_url: Some("https://bedrock-mantle.us-east-1.api.aws/openai/v1".to_string()), + base_url: None, env_key: None, env_key_instructions: None, experimental_bearer_token: None, @@ -314,6 +314,19 @@ fn test_create_amazon_bedrock_provider() { ); } +fn provider_auth_for_test() -> ModelProviderAuthInfo { + ModelProviderAuthInfo { + command: "token-fetcher".to_string(), + args: vec!["fetch".to_string()], + timeout_ms: NonZeroU64::new(5_000).expect("timeout should be non-zero"), + refresh_interval_ms: 300_000, + cwd: std::env::current_dir() + .expect("current directory should be available") + .try_into() + .expect("current directory should be absolute"), + } +} + #[test] fn test_amazon_bedrock_provider_adds_mantle_client_agent_header() { let api_provider = ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None) @@ -394,6 +407,49 @@ fn test_merge_configured_model_providers_applies_amazon_bedrock_profile_override ); } +#[test] +fn test_merge_configured_model_providers_applies_amazon_bedrock_transport_overrides() { + let auth = provider_auth_for_test(); + let configured_model_providers = std::collections::HashMap::from([( + AMAZON_BEDROCK_PROVIDER_ID.to_string(), + ModelProviderInfo { + base_url: Some("https://proxy.example.com/v1".to_string()), + auth: Some(auth.clone()), + aws: Some(ModelProviderAwsAuthInfo { + profile: Some("codex-bedrock".to_string()), + region: Some("us-west-2".to_string()), + }), + http_headers: Some(maplit::hashmap! { + "x-example-header".to_string() => "value".to_string(), + }), + ..ModelProviderInfo::default() + }, + )]); + + let mut expected = built_in_model_providers(/*openai_base_url*/ None); + let expected_provider = expected + .get_mut(AMAZON_BEDROCK_PROVIDER_ID) + .expect("Amazon Bedrock provider should be built in"); + expected_provider.base_url = Some("https://proxy.example.com/v1".to_string()); + expected_provider.auth = Some(auth); + expected_provider.aws = Some(ModelProviderAwsAuthInfo { + profile: Some("codex-bedrock".to_string()), + region: Some("us-west-2".to_string()), + }); + expected_provider + .http_headers + .get_or_insert_default() + .insert("x-example-header".to_string(), "value".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([( @@ -414,7 +470,7 @@ fn test_merge_configured_model_providers_rejects_amazon_bedrock_non_default_fiel configured_model_providers, ), Err( - "model_providers.amazon-bedrock only supports changing `aws.profile` and `aws.region`; other non-default provider fields are not supported" + "model_providers.amazon-bedrock only supports changing `base_url`, `auth`, `http_headers`, `aws.profile`, and `aws.region`; other non-default provider fields are not supported" .to_string() ) ); diff --git a/codex-rs/model-provider/src/amazon_bedrock/auth.rs b/codex-rs/model-provider/src/amazon_bedrock/auth.rs index b9b318d9b1c7..20701670358d 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/auth.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/auth.rs @@ -80,7 +80,7 @@ fn non_empty_env_var_from( .filter(|value| !value.is_empty()) } -fn bearer_token_region( +pub(super) fn bearer_token_region( aws: &ModelProviderAwsAuthInfo, env_var: impl Fn(&'static str) -> std::result::Result + Copy, ) -> Result { diff --git a/codex-rs/model-provider/src/amazon_bedrock/mantle.rs b/codex-rs/model-provider/src/amazon_bedrock/mantle.rs index e31ce96edc01..d9f9ac48848d 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/mantle.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/mantle.rs @@ -54,7 +54,7 @@ pub(super) fn base_url(region: &str) -> Result { } } -pub(super) async fn runtime_base_url( +pub(super) async fn bedrock_mantle_runtime_base_url( managed_auth: Option<&BedrockApiKeyAuth>, aws: &ModelProviderAwsAuthInfo, ) -> Result { diff --git a/codex-rs/model-provider/src/amazon_bedrock/mod.rs b/codex-rs/model-provider/src/amazon_bedrock/mod.rs index c013907703cd..d807da1e957a 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/mod.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/mod.rs @@ -17,22 +17,23 @@ use codex_model_provider_info::ModelProviderAwsAuthInfo; use codex_model_provider_info::ModelProviderInfo; use codex_models_manager::manager::SharedModelsManager; use codex_models_manager::manager::StaticModelsManager; -use codex_protocol::account::AmazonBedrockCredentialSource; use codex_protocol::account::ProviderAccount; use codex_protocol::error::CodexErr; use codex_protocol::error::Result; use codex_protocol::openai_models::ModelsResponse; +use crate::auth::auth_manager_for_provider; +use crate::auth::resolve_provider_auth as resolve_configured_provider_auth; use crate::provider::ModelProvider; use crate::provider::ModelProviderFuture; use crate::provider::ProviderAccountResult; use crate::provider::ProviderAccountState; use crate::provider::ProviderCapabilities; -use auth::resolve_provider_auth; +use auth::resolve_provider_auth as resolve_bedrock_provider_auth; pub(crate) use catalog::static_model_catalog; use catalog::with_default_only_service_tier; +use mantle::bedrock_mantle_runtime_base_url; pub use mantle::is_supported_amazon_bedrock_region; -use mantle::runtime_base_url; /// Runtime provider for Amazon Bedrock's OpenAI-compatible Mantle endpoint. #[derive(Clone, Debug)] @@ -47,6 +48,7 @@ impl AmazonBedrockModelProvider { provider_info: ModelProviderInfo, auth_manager: Option>, ) -> Self { + let auth_manager = auth_manager_for_provider(auth_manager, &provider_info); let aws = provider_info .aws .clone() @@ -77,27 +79,39 @@ impl AmazonBedrockModelProvider { } async fn auth(&self) -> Option { - self.managed_auth().map(CodexAuth::BedrockApiKey) + if self.info.has_command_auth() { + match self.auth_manager.as_ref() { + Some(auth_manager) => auth_manager.auth().await, + None => None, + } + } else { + self.managed_auth().map(CodexAuth::BedrockApiKey) + } } async fn api_provider(&self) -> Result { - let managed_auth = self.managed_auth(); let mut api_provider_info = self.info.clone(); - api_provider_info.base_url = - Some(runtime_base_url(managed_auth.as_ref(), &self.aws).await?); + api_provider_info.base_url = self.runtime_base_url().await?; api_provider_info.to_api_provider(/*auth_mode*/ None) } async fn runtime_base_url(&self) -> Result> { + if let Some(base_url) = self.info.base_url.clone() { + return Ok(Some(base_url)); + } let managed_auth = self.managed_auth(); Ok(Some( - runtime_base_url(managed_auth.as_ref(), &self.aws).await?, + bedrock_mantle_runtime_base_url(managed_auth.as_ref(), &self.aws).await?, )) } async fn api_auth(&self) -> Result { + if self.info.has_command_auth() { + let auth = self.auth().await; + return resolve_configured_provider_auth(auth.as_ref(), &self.info); + } let managed_auth = self.managed_auth(); - resolve_provider_auth(managed_auth.as_ref(), &self.aws).await + resolve_bedrock_provider_auth(managed_auth.as_ref(), &self.aws).await } } @@ -127,8 +141,11 @@ impl ModelProvider for AmazonBedrockModelProvider { } fn auth_manager(&self) -> Option> { - self.managed_auth() - .and_then(|_| self.auth_manager.as_ref().cloned()) + if self.info.has_command_auth() || self.managed_auth().is_some() { + self.auth_manager.clone() + } else { + None + } } fn auth(&self) -> ModelProviderFuture<'_, Option> { @@ -136,13 +153,10 @@ impl ModelProvider for AmazonBedrockModelProvider { } fn account_state(&self) -> ProviderAccountResult { - let credential_source = if self.managed_auth().is_some() { - AmazonBedrockCredentialSource::CodexManaged - } else { - AmazonBedrockCredentialSource::AwsManaged - }; Ok(ProviderAccountState { - account: Some(ProviderAccount::AmazonBedrock { credential_source }), + account: Some(ProviderAccount::AmazonBedrock { + uses_codex_managed_credentials: self.managed_auth().is_some(), + }), requires_openai_auth: false, }) } @@ -191,11 +205,30 @@ mod error_tests; #[cfg(test)] mod tests { + use std::num::NonZeroU64; + + use codex_protocol::config_types::ModelProviderAuthInfo; use http::HeaderValue; use pretty_assertions::assert_eq; use super::*; + fn command_auth_provider(base_url: Option<&str>) -> ModelProviderInfo { + let mut provider = ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None); + provider.base_url = base_url.map(str::to_string); + provider.auth = Some(ModelProviderAuthInfo { + command: "token-fetcher".to_string(), + args: vec!["fetch".to_string()], + timeout_ms: NonZeroU64::new(5_000).expect("timeout should be non-zero"), + refresh_interval_ms: 300_000, + cwd: std::env::current_dir() + .expect("current directory should be available") + .try_into() + .expect("current directory should be absolute"), + }); + provider + } + #[test] fn api_provider_for_bedrock_bearer_token_uses_configured_region_endpoint() { let region = "eu-central-1"; @@ -212,6 +245,39 @@ mod tests { ); } + #[tokio::test] + async fn command_auth_uses_configured_base_url_without_resolving_aws() { + let mut provider_info = command_auth_provider(Some("https://proxy.example.com/v1")); + provider_info.aws = Some(ModelProviderAwsAuthInfo { + profile: Some("aws-profile-that-should-not-be-loaded".to_string()), + region: Some("us-west-2".to_string()), + }); + let provider = AmazonBedrockModelProvider::new(provider_info, /*auth_manager*/ None); + + assert_eq!( + provider + .runtime_base_url() + .await + .expect("configured base URL should resolve"), + Some("https://proxy.example.com/v1".to_string()) + ); + assert!( + provider + .auth_manager() + .expect("command auth manager should be exposed") + .has_external_auth() + ); + assert_eq!( + provider.account_state(), + Ok(ProviderAccountState { + account: Some(ProviderAccount::AmazonBedrock { + uses_codex_managed_credentials: false, + }), + requires_openai_auth: false, + }) + ); + } + #[tokio::test] async fn managed_auth_takes_precedence_over_aws_auth() { let managed_auth = BedrockApiKeyAuth { @@ -242,7 +308,7 @@ mod tests { provider.account_state(), Ok(ProviderAccountState { account: Some(ProviderAccount::AmazonBedrock { - credential_source: AmazonBedrockCredentialSource::CodexManaged, + uses_codex_managed_credentials: true, }), requires_openai_auth: false, }) @@ -280,7 +346,7 @@ mod tests { provider.account_state(), Ok(ProviderAccountState { account: Some(ProviderAccount::AmazonBedrock { - credential_source: AmazonBedrockCredentialSource::AwsManaged, + uses_codex_managed_credentials: false, }), requires_openai_auth: false, }) diff --git a/codex-rs/model-provider/src/provider.rs b/codex-rs/model-provider/src/provider.rs index f0ad92281b9f..756bd641375c 100644 --- a/codex-rs/model-provider/src/provider.rs +++ b/codex-rs/model-provider/src/provider.rs @@ -676,8 +676,7 @@ mod tests { provider.account_state(), Ok(ProviderAccountState { account: Some(ProviderAccount::AmazonBedrock { - credential_source: - codex_protocol::account::AmazonBedrockCredentialSource::AwsManaged, + uses_codex_managed_credentials: false, }), requires_openai_auth: false, }) diff --git a/codex-rs/protocol/src/account.rs b/codex-rs/protocol/src/account.rs index dfcac51d30bc..aedf5474be47 100644 --- a/codex-rs/protocol/src/account.rs +++ b/codex-rs/protocol/src/account.rs @@ -39,18 +39,10 @@ pub enum ProviderAccount { plan_type: PlanType, }, AmazonBedrock { - credential_source: AmazonBedrockCredentialSource, + uses_codex_managed_credentials: bool, }, } -#[derive(Serialize, Deserialize, Copy, Clone, Debug, PartialEq, Eq, JsonSchema, TS)] -#[serde(rename_all = "camelCase")] -#[ts(rename_all = "camelCase")] -pub enum AmazonBedrockCredentialSource { - CodexManaged, - AwsManaged, -} - impl PlanType { pub fn is_team_like(self) -> bool { matches!(self, Self::Team | Self::SelfServeBusinessUsageBased)