diff --git a/codex-rs/app-server-test-client/README.md b/codex-rs/app-server-test-client/README.md index b1afc778723a..9f73b1db6f5c 100644 --- a/codex-rs/app-server-test-client/README.md +++ b/codex-rs/app-server-test-client/README.md @@ -22,6 +22,30 @@ cargo run -p codex-app-server-test-client -- model-list When Codex asks a question, choose a numbered option (or `o` for a free-form answer when offered) and the client will send the response and continue streaming the same turn. +## Testing Codex-managed Amazon Bedrock login + +`test-login --amazon-bedrock` initializes the experimental app-server API, sends an +`account/login/start` request with an Amazon Bedrock API key, and waits for the +`account/login/completed` and `account/updated` notifications. Login replaces the current primary +credential and sets `model_provider = "amazon-bedrock"`, so use an isolated `CODEX_HOME` when +testing. + +```bash +export CODEX_HOME="$(mktemp -d)" +printf 'cli_auth_credentials_store = "file"\n' > "$CODEX_HOME/config.toml" + +cargo build -p codex-cli --bin codex +cargo run -p codex-app-server-test-client -- \ + --codex-bin ./target/debug/codex \ + test-login \ + --amazon-bedrock \ + --api-key "" \ + --region us-west-2 +``` + +The test client redacts `apiKey` from its outbound request log. After login, start a fresh Codex +process with the same `CODEX_HOME` to verify that it uses the persisted managed credential. + ## Testing Plugin Analytics The `plugin-analytics-smoke` command exercises `plugin/installed`, plugin diff --git a/codex-rs/app-server-test-client/src/lib.rs b/codex-rs/app-server-test-client/src/lib.rs index 0e32365d12fe..10d66dbc01d5 100644 --- a/codex-rs/app-server-test-client/src/lib.rs +++ b/codex-rs/app-server-test-client/src/lib.rs @@ -230,11 +230,20 @@ enum CliCommand { #[arg(long)] abort_on: Option, }, - /// Trigger the ChatGPT login flow and wait for completion. + /// Trigger a ChatGPT or Amazon Bedrock login flow. TestLogin { /// Use the device-code login flow instead of the browser callback flow. - #[arg(long, default_value_t = false)] + #[arg(long, default_value_t = false, conflicts_with = "amazon_bedrock")] device_code: bool, + /// Use a Codex-managed Amazon Bedrock API key. + #[arg(long, default_value_t = false, conflicts_with = "device_code")] + amazon_bedrock: bool, + /// Amazon Bedrock API key. + #[arg(long, value_name = "API_KEY")] + api_key: Option, + /// AWS Region for the Amazon Bedrock Mantle endpoint. + #[arg(long, value_name = "REGION")] + region: Option, }, /// Fetch the current account rate limits from the Codex app-server. GetAccountRateLimits, @@ -313,6 +322,12 @@ enum CliCommand { }, } +enum TestLoginMode { + ChatgptBrowser, + ChatgptDeviceCode, + AmazonBedrock { api_key: String, region: String }, +} + pub async fn run() -> Result<()> { let Cli { codex_bin, @@ -415,10 +430,24 @@ pub async fn run() -> Result<()> { ) .await } - CliCommand::TestLogin { device_code } => { + CliCommand::TestLogin { + device_code, + amazon_bedrock, + api_key, + region, + } => { ensure_dynamic_tools_unused(&dynamic_tools, "test-login")?; let endpoint = resolve_endpoint(codex_bin, url)?; - test_login(&endpoint, &config_overrides, device_code).await + let mode = if amazon_bedrock { + let api_key = api_key.context("--api-key is required with --amazon-bedrock")?; + let region = region.context("--region is required with --amazon-bedrock")?; + TestLoginMode::AmazonBedrock { api_key, region } + } else if device_code { + TestLoginMode::ChatgptDeviceCode + } else { + TestLoginMode::ChatgptBrowser + }; + test_login(&endpoint, &config_overrides, mode).await } CliCommand::GetAccountRateLimits => { ensure_dynamic_tools_unused(&dynamic_tools, "get-account-rate-limits")?; @@ -1128,16 +1157,45 @@ async fn send_follow_up_v2( async fn test_login( endpoint: &Endpoint, config_overrides: &[String], - device_code: bool, + mode: TestLoginMode, ) -> Result<()> { with_client("test-login", endpoint, config_overrides, |client| { let initialize = client.initialize()?; println!("< initialize response: {initialize:?}"); - let login_response = if device_code { - client.login_account_chatgpt_device_code()? - } else { - client.login_account_chatgpt()? + let login_response = match mode { + TestLoginMode::ChatgptBrowser => client.login_account_chatgpt()?, + TestLoginMode::ChatgptDeviceCode => client.login_account_chatgpt_device_code()?, + TestLoginMode::AmazonBedrock { api_key, region } => { + let request_id = client.request_id(); + let login_response: LoginAccountResponse = client.send_request( + ClientRequest::LoginAccount { + request_id: request_id.clone(), + params: codex_app_server_protocol::LoginAccountParams::AmazonBedrock { + api_key, + region, + }, + }, + request_id, + "account/login/start", + )?; + println!("< account/login/start response: {login_response:?}"); + + let completion = + client.wait_for_account_login_completion(/*expected_login_id*/ None)?; + println!("< account/login/completed notification: {completion:?}"); + + loop { + let notification = client.next_notification()?; + if let Ok(ServerNotification::AccountUpdated(account_updated)) = + ServerNotification::try_from(notification) + { + println!("< account/updated notification: {account_updated:?}"); + break; + } + } + return Ok(()); + } }; println!("< account/login/start response: {login_response:?}"); let login_id = match login_response { @@ -1158,7 +1216,7 @@ async fn test_login( _ => bail!("expected chatgpt login response"), }; - let completion = client.wait_for_account_login_completion(&login_id)?; + let completion = client.wait_for_account_login_completion(Some(&login_id))?; println!("< account/login/completed notification: {completion:?}"); if completion.success { @@ -1799,7 +1857,7 @@ impl CodexClient { fn wait_for_account_login_completion( &mut self, - expected_login_id: &str, + expected_login_id: Option<&str>, ) -> Result { loop { let notification = self.next_notification()?; @@ -1807,7 +1865,7 @@ impl CodexClient { if let Ok(server_notification) = ServerNotification::try_from(notification) { match server_notification { ServerNotification::AccountLoginCompleted(completion) => { - if completion.login_id.as_deref() == Some(expected_login_id) { + if completion.login_id.as_deref() == expected_login_id { return Ok(completion); } @@ -1956,7 +2014,13 @@ impl CodexClient { .context("client request was not a valid JSON-RPC request")?; request.trace = current_span_w3c_trace_context(); let request_json = serde_json::to_string(&request)?; - let request_pretty = serde_json::to_string_pretty(&request)?; + let mut request_for_logging = serde_json::to_value(&request)?; + if request.method == "account/login/start" + && let Some(api_key) = request_for_logging.pointer_mut("/params/apiKey") + { + *api_key = Value::String("".to_string()); + } + let request_pretty = serde_json::to_string_pretty(&request_for_logging)?; print_multiline_with_prefix("> ", &request_pretty); self.write_payload(&request_json) } diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index a7a1f6ce8196..b58e430d7e99 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -387,6 +387,7 @@ use codex_login::ServerOptions as LoginServerOptions; use codex_login::ShutdownHandle; use codex_login::complete_device_code_login; use codex_login::login_with_api_key; +use codex_login::login_with_bedrock_api_key; use codex_login::oauth_client_id; use codex_login::request_device_code; use codex_login::run_login_server; @@ -497,6 +498,7 @@ use codex_app_server_protocol::ServerRequest; mod account_processor; mod apps_processor; +mod bedrock_auth; mod catalog_processor; mod command_exec_processor; mod config_processor; diff --git a/codex-rs/app-server/src/request_processors/account_processor.rs b/codex-rs/app-server/src/request_processors/account_processor.rs index d98977ce580b..f0e1ebfc3da2 100644 --- a/codex-rs/app-server/src/request_processors/account_processor.rs +++ b/codex-rs/app-server/src/request_processors/account_processor.rs @@ -1,7 +1,9 @@ +use super::bedrock_auth::set_user_model_provider_to_bedrock; use super::*; use crate::auth_mode::auth_mode_to_api; use crate::external_auth::ExternalAuthBridge; use chrono::DateTime; +use codex_model_provider::is_supported_amazon_bedrock_region; mod rate_limit_resets; @@ -191,6 +193,13 @@ impl AccountRequestProcessor { } } + async fn load_latest_config(&self) -> Result { + self.config_manager + .load_latest_config(/*fallback_cwd*/ None) + .await + .map_err(|err| internal_error(format!("failed to reload config: {err}"))) + } + async fn maybe_refresh_plugin_caches_for_current_config( config_manager: &ConfigManager, thread_manager: &Arc, @@ -293,8 +302,9 @@ impl AccountRequestProcessor { ) .await; } - LoginAccountParams::AmazonBedrock { .. } => { - return Err(invalid_request("Amazon Bedrock login is not implemented")); + LoginAccountParams::AmazonBedrock { api_key, region } => { + self.login_amazon_bedrock_v2(request_id, api_key, region) + .await; } } Ok(()) @@ -359,6 +369,65 @@ impl AccountRequestProcessor { } } + async fn login_amazon_bedrock_v2( + &self, + request_id: ConnectionRequestId, + api_key: String, + region: String, + ) { + let result = async { + if self.auth_manager.is_external_chatgpt_auth_active() { + return Err(self.external_auth_active_error()); + } + if matches!( + self.config.forced_login_method, + Some(ForcedLoginMethod::Chatgpt) + ) { + return Err(invalid_request( + "Amazon Bedrock login is disabled. Use ChatGPT login instead.", + )); + } + + let api_key = api_key.trim(); + if api_key.is_empty() { + return Err(invalid_request("Amazon Bedrock API key must not be empty.")); + } + let region = region.trim(); + if !is_supported_amazon_bedrock_region(region) { + return Err(invalid_request(format!( + "Amazon Bedrock Mantle does not support region `{region}`" + ))); + } + + { + let mut guard = self.active_login.lock().await; + if let Some(active) = guard.take() { + drop(active); + } + } + + set_user_model_provider_to_bedrock(&self.config_manager).await?; + login_with_bedrock_api_key( + &self.config.codex_home, + api_key, + region, + self.config.cli_auth_credentials_store_mode, + self.config.auth_keyring_backend_kind(), + ) + .map_err(|err| internal_error(format!("failed to save Amazon Bedrock auth: {err}")))?; + self.auth_manager.reload().await; + Ok(LoginAccountResponse::AmazonBedrock {}) + } + .await; + let logged_in = result.is_ok(); + self.outgoing.send_result(request_id, result).await; + + if logged_in { + self.send_login_success_notifications(/*login_id*/ None) + .await; + } + } + // Build options for a ChatGPT login attempt; performs validation. async fn login_chatgpt_common( &self, @@ -837,7 +906,8 @@ impl AccountRequestProcessor { // Determine whether auth is required based on the active model provider. // If a custom provider is configured with `requires_openai_auth == false`, // then no auth step is required; otherwise, default to requiring auth. - let requires_openai_auth = self.config.model_provider.requires_openai_auth; + let config = self.load_latest_config().await?; + let requires_openai_auth = config.model_provider.requires_openai_auth; let response = if !requires_openai_auth { GetAuthStatusResponse { @@ -905,10 +975,9 @@ impl AccountRequestProcessor { self.refresh_token_if_requested(do_refresh).await; - let provider = create_model_provider( - self.config.model_provider.clone(), - Some(self.auth_manager.clone()), - ); + let config = self.load_latest_config().await?; + let provider = + create_model_provider(config.model_provider, Some(self.auth_manager.clone())); let account_state = match provider.account_state() { Ok(account_state) => account_state, Err(err) => return Err(invalid_request(err.to_string())), diff --git a/codex-rs/app-server/src/request_processors/bedrock_auth.rs b/codex-rs/app-server/src/request_processors/bedrock_auth.rs new file mode 100644 index 000000000000..88576e5ffe9c --- /dev/null +++ b/codex-rs/app-server/src/request_processors/bedrock_auth.rs @@ -0,0 +1,35 @@ +use super::config_processor::map_error as map_config_error; +use crate::config_manager::ConfigManager; +use codex_app_server_protocol::ConfigValueWriteParams; +use codex_app_server_protocol::JSONRPCErrorError; +use codex_app_server_protocol::MergeStrategy; +use codex_model_provider::AMAZON_BEDROCK_PROVIDER_ID; + +pub(super) async fn set_user_model_provider_to_bedrock( + config_manager: &ConfigManager, +) -> Result<(), JSONRPCErrorError> { + write_user_model_provider( + config_manager, + serde_json::json!(AMAZON_BEDROCK_PROVIDER_ID), + /*expected_version*/ None, + ) + .await +} + +async fn write_user_model_provider( + config_manager: &ConfigManager, + value: serde_json::Value, + expected_version: Option, +) -> Result<(), JSONRPCErrorError> { + config_manager + .write_value(ConfigValueWriteParams { + key_path: "model_provider".to_string(), + value, + merge_strategy: MergeStrategy::Replace, + file_path: None, + expected_version, + }) + .await + .map(|_| ()) + .map_err(map_config_error) +} diff --git a/codex-rs/app-server/src/request_processors/config_processor.rs b/codex-rs/app-server/src/request_processors/config_processor.rs index b6dbb30b0834..0eeb1aaca98a 100644 --- a/codex-rs/app-server/src/request_processors/config_processor.rs +++ b/codex-rs/app-server/src/request_processors/config_processor.rs @@ -555,7 +555,7 @@ fn map_network_unix_socket_permission_to_api( } } -fn map_error(err: ConfigManagerError) -> JSONRPCErrorError { +pub(super) fn map_error(err: ConfigManagerError) -> JSONRPCErrorError { if let Some(code) = err.write_error_code() { return config_write_error(code, err.to_string()); } diff --git a/codex-rs/app-server/tests/common/test_app_server.rs b/codex-rs/app-server/tests/common/test_app_server.rs index 1a74e1df5ad4..dcd62a8ae9d5 100644 --- a/codex-rs/app-server/tests/common/test_app_server.rs +++ b/codex-rs/app-server/tests/common/test_app_server.rs @@ -1325,6 +1325,20 @@ impl TestAppServer { self.send_login_account_request(params).await } + /// Send an `account/login/start` JSON-RPC request for managed Amazon Bedrock login. + pub async fn send_login_account_amazon_bedrock_request( + &mut self, + api_key: &str, + region: &str, + ) -> anyhow::Result { + let params = serde_json::json!({ + "type": "amazonBedrock", + "apiKey": api_key, + "region": region, + }); + self.send_request("account/login/start", Some(params)).await + } + /// Send an `account/login/start` JSON-RPC request for ChatGPT login. pub async fn send_login_account_chatgpt_request(&mut self) -> anyhow::Result { let params = serde_json::json!({ diff --git a/codex-rs/app-server/tests/suite/v2/account.rs b/codex-rs/app-server/tests/suite/v2/account.rs index a8e8c0903d46..2ffec32bb9b5 100644 --- a/codex-rs/app-server/tests/suite/v2/account.rs +++ b/codex-rs/app-server/tests/suite/v2/account.rs @@ -5,24 +5,30 @@ use app_test_support::to_response; use app_test_support::ChatGptAuthFixture; use app_test_support::ChatGptIdTokenClaims; +use app_test_support::DEFAULT_CLIENT_NAME; use app_test_support::encode_id_token; use app_test_support::write_chatgpt_auth; use app_test_support::write_models_cache; use chrono::Duration as ChronoDuration; use chrono::Utc; use codex_app_server_protocol::Account; +use codex_app_server_protocol::AccountLoginCompletedNotification; +use codex_app_server_protocol::AccountUpdatedNotification; use codex_app_server_protocol::AuthMode; use codex_app_server_protocol::CancelLoginAccountParams; use codex_app_server_protocol::CancelLoginAccountResponse; use codex_app_server_protocol::CancelLoginAccountStatus; use codex_app_server_protocol::ChatgptAuthTokensRefreshReason; use codex_app_server_protocol::ChatgptAuthTokensRefreshResponse; +use codex_app_server_protocol::ClientInfo; use codex_app_server_protocol::GetAccountParams; use codex_app_server_protocol::GetAccountResponse; use codex_app_server_protocol::GetAuthStatusParams; use codex_app_server_protocol::GetAuthStatusResponse; +use codex_app_server_protocol::InitializeCapabilities; use codex_app_server_protocol::JSONRPCError; use codex_app_server_protocol::JSONRPCErrorError; +use codex_app_server_protocol::JSONRPCMessage; use codex_app_server_protocol::JSONRPCNotification; use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::LoginAccountResponse; @@ -33,13 +39,17 @@ use codex_app_server_protocol::ServerRequest; use codex_app_server_protocol::TurnCompletedNotification; use codex_app_server_protocol::TurnStatus; use codex_config::types::AuthCredentialsStoreMode; +use codex_login::AuthDotJson; use codex_login::AuthKeyringBackendKind; use codex_login::CLIENT_ID_OVERRIDE_ENV_VAR; use codex_login::REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR; +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; use pretty_assertions::assert_eq; use serde_json::json; @@ -142,6 +152,56 @@ shell_snapshot = false std::fs::write(config_toml, contents) } +fn read_config_toml(codex_home: &Path) -> Result { + Ok(toml::from_str(&std::fs::read_to_string( + codex_home.join("config.toml"), + )?)?) +} + +fn load_file_auth(codex_home: &Path) -> Result> { + Ok(load_auth_dot_json( + codex_home, + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?) +} + +async fn read_account(mcp: &mut TestAppServer) -> Result { + let request_id = mcp + .send_get_account_request(GetAccountParams { + refresh_token: false, + }) + .await?; + let response = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + to_response(response) +} + +async fn assert_account_updated( + mcp: &mut TestAppServer, + auth_mode: Option, +) -> Result<()> { + let notification = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("account/updated"), + ) + .await??; + let ServerNotification::AccountUpdated(payload) = notification.try_into()? else { + bail!("unexpected notification") + }; + assert_eq!( + payload, + AccountUpdatedNotification { + auth_mode, + plan_type: None, + } + ); + Ok(()) +} + async fn mock_device_code_usercode(server: &MockServer, interval_seconds: u64) { Mock::given(method("POST")) .and(path("/api/accounts/deviceauth/usercode")) @@ -1011,6 +1071,300 @@ async fn login_account_api_key_succeeds_and_notifies() -> Result<()> { Ok(()) } +#[tokio::test] +async fn login_amazon_bedrock_replaces_primary_auth_and_persists_provider() -> Result<()> { + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?; + login_with_api_key( + codex_home.path(), + "sk-test-key", + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + let mut mcp = TestAppServer::builder() + .with_codex_home(codex_home.path()) + .without_auto_env() + .with_env_overrides(&[("OPENAI_API_KEY", None)]) + .build() + .await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + let mut expected_config = read_config_toml(codex_home.path())?; + expected_config + .as_table_mut() + .expect("config should be a table") + .insert( + "model_provider".to_string(), + toml::Value::String("amazon-bedrock".to_string()), + ); + let request_id = mcp + .send_login_account_amazon_bedrock_request(" managed-bedrock-api-key ", " us-west-2 ") + .await?; + let response: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + assert_eq!( + to_response::(response)?, + LoginAccountResponse::AmazonBedrock {} + ); + + assert_eq!( + load_file_auth(codex_home.path())?, + Some(AuthDotJson { + auth_mode: Some(DomainAuthMode::BedrockApiKey), + openai_api_key: None, + tokens: None, + last_refresh: None, + agent_identity: None, + personal_access_token: None, + bedrock_api_key: Some(BedrockApiKeyAuth { + api_key: "managed-bedrock-api-key".to_string(), + region: "us-west-2".to_string(), + }), + }) + ); + assert_eq!(read_config_toml(codex_home.path())?, expected_config); + + let notification = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("account/login/completed"), + ) + .await??; + let ServerNotification::AccountLoginCompleted(payload) = notification.try_into()? else { + bail!("unexpected notification") + }; + assert_eq!( + payload, + AccountLoginCompletedNotification { + login_id: None, + success: true, + error: None, + } + ); + assert_account_updated(&mut mcp, Some(AuthMode::BedrockApiKey)).await?; + + Ok(()) +} + +#[tokio::test] +async fn managed_bedrock_login_requires_experimental_api() -> Result<()> { + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?; + let mut mcp = TestAppServer::builder() + .with_codex_home(codex_home.path()) + .without_auto_env() + .build() + .await?; + let initialized = mcp + .initialize_with_capabilities( + ClientInfo { + name: DEFAULT_CLIENT_NAME.to_string(), + title: None, + version: "0.1.0".to_string(), + }, + Some(InitializeCapabilities { + experimental_api: false, + ..Default::default() + }), + ) + .await?; + assert!(matches!(initialized, JSONRPCMessage::Response(_))); + + let request_id = mcp + .send_login_account_amazon_bedrock_request("managed-bedrock-api-key", "us-west-2") + .await?; + let error = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_error_message(RequestId::Integer(request_id)), + ) + .await??; + assert_eq!( + error.error.message, + "account/login/start.amazonBedrock requires experimentalApi capability" + ); + assert_eq!(load_file_auth(codex_home.path())?, None); + Ok(()) +} + +#[tokio::test] +async fn login_managed_bedrock_updates_active_bedrock_account() -> Result<()> { + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?; + + let mut mcp = TestAppServer::builder() + .with_codex_home(codex_home.path()) + .without_auto_env() + .with_env_overrides(&[("OPENAI_API_KEY", None)]) + .build() + .await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + let request_id = mcp + .send_login_account_amazon_bedrock_request("managed-bedrock-api-key", "us-west-2") + .await?; + let response: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + assert_eq!( + to_response::(response)?, + LoginAccountResponse::AmazonBedrock {} + ); + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("account/login/completed"), + ) + .await??; + assert_account_updated(&mut mcp, Some(AuthMode::BedrockApiKey)).await?; + assert_eq!( + read_account(&mut mcp).await?, + GetAccountResponse { + account: Some(Account::AmazonBedrock { + credential_source: AmazonBedrockCredentialSource::CodexManaged, + }), + requires_openai_auth: false, + } + ); + + assert!(codex_home.path().join("auth.json").exists()); + Ok(()) +} + +#[tokio::test] +async fn login_account_amazon_bedrock_rejects_invalid_credentials_without_changes() -> Result<()> { + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?; + + let mut mcp = TestAppServer::builder() + .with_codex_home(codex_home.path()) + .without_auto_env() + .with_env_overrides(&[("OPENAI_API_KEY", None)]) + .build() + .await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + let expected_config = read_config_toml(codex_home.path())?; + + let request_id = mcp + .send_login_account_amazon_bedrock_request(" ", "us-west-2") + .await?; + let error = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_error_message(RequestId::Integer(request_id)), + ) + .await??; + assert_eq!( + error.error.message, + "Amazon Bedrock API key must not be empty." + ); + + let request_id = mcp + .send_login_account_amazon_bedrock_request("managed-bedrock-api-key", "us-west-1") + .await?; + let error = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_error_message(RequestId::Integer(request_id)), + ) + .await??; + assert_eq!( + error.error.message, + "Amazon Bedrock Mantle does not support region `us-west-1`" + ); + assert_eq!(load_file_auth(codex_home.path())?, None); + assert_eq!(read_config_toml(codex_home.path())?, expected_config); + + Ok(()) +} + +#[tokio::test] +async fn login_account_amazon_bedrock_rejected_when_forced_chatgpt() -> Result<()> { + let codex_home = TempDir::new()?; + create_config_toml( + codex_home.path(), + CreateConfigTomlParams { + forced_method: Some("chatgpt".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??; + let request_id = mcp + .send_login_account_amazon_bedrock_request("managed-bedrock-api-key", "us-west-2") + .await?; + let error = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_error_message(RequestId::Integer(request_id)), + ) + .await??; + + assert_eq!( + error.error.message, + "Amazon Bedrock login is disabled. Use ChatGPT login instead." + ); + assert_eq!(load_file_auth(codex_home.path())?, None); + Ok(()) +} + +#[tokio::test] +async fn login_account_amazon_bedrock_rejected_with_external_chatgpt_auth() -> Result<()> { + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?; + let access_token = encode_id_token( + &ChatGptIdTokenClaims::new() + .email("embedded@example.com") + .plan_type("pro") + .chatgpt_account_id(WORKSPACE_ID_EMBEDDED), + )?; + + let mut mcp = TestAppServer::builder() + .with_codex_home(codex_home.path()) + .without_auto_env() + .build() + .await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + let set_id = mcp + .send_chatgpt_auth_tokens_login_request( + access_token, + WORKSPACE_ID_EMBEDDED.to_string(), + Some("pro".to_string()), + ) + .await?; + let set_response = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(set_id)), + ) + .await??; + assert_eq!( + to_response::(set_response)?, + LoginAccountResponse::ChatgptAuthTokens {} + ); + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("account/updated"), + ) + .await??; + + let request_id = mcp + .send_login_account_amazon_bedrock_request("managed-bedrock-api-key", "us-west-2") + .await?; + let error = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_error_message(RequestId::Integer(request_id)), + ) + .await??; + assert_eq!( + error.error.message, + "External auth is active. Use account/login/start (chatgptAuthTokens) to update it or account/logout to clear it." + ); + assert_eq!(load_file_auth(codex_home.path())?, None); + Ok(()) +} + #[tokio::test] async fn login_account_api_key_rejected_when_forced_chatgpt() -> Result<()> { let codex_home = TempDir::new()?; diff --git a/codex-rs/model-provider/src/amazon_bedrock/mantle.rs b/codex-rs/model-provider/src/amazon_bedrock/mantle.rs index 9727802d1196..e31ce96edc01 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/mantle.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/mantle.rs @@ -39,8 +39,13 @@ pub(super) fn region_from_config(aws: &ModelProviderAwsAuthInfo) -> Option bool { + BEDROCK_MANTLE_SUPPORTED_REGIONS.contains(®ion) +} + pub(super) fn base_url(region: &str) -> Result { - if BEDROCK_MANTLE_SUPPORTED_REGIONS.contains(®ion) { + if is_supported_amazon_bedrock_region(region) { Ok(format!("https://bedrock-mantle.{region}.api.aws/openai/v1")) } else { Err(CodexErr::Fatal(format!( diff --git a/codex-rs/model-provider/src/amazon_bedrock/mod.rs b/codex-rs/model-provider/src/amazon_bedrock/mod.rs index 9fdf937bcb0c..2a7230e08153 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/mod.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/mod.rs @@ -31,6 +31,7 @@ use crate::provider::ProviderCapabilities; use auth::resolve_provider_auth; pub(crate) use catalog::static_model_catalog; use catalog::with_default_only_service_tier; +pub use mantle::is_supported_amazon_bedrock_region; use mantle::runtime_base_url; /// Runtime provider for Amazon Bedrock's OpenAI-compatible Mantle endpoint. diff --git a/codex-rs/model-provider/src/lib.rs b/codex-rs/model-provider/src/lib.rs index 15967c4fe033..2a23343dfc36 100644 --- a/codex-rs/model-provider/src/lib.rs +++ b/codex-rs/model-provider/src/lib.rs @@ -4,6 +4,7 @@ mod bearer_auth_provider; mod models_endpoint; mod provider; +pub use amazon_bedrock::is_supported_amazon_bedrock_region; pub use auth::AgentIdentitySessionFallback; pub use auth::ProviderAuthScope; pub use auth::ResolvedProviderAuth; @@ -11,6 +12,7 @@ pub use auth::auth_provider_from_auth; 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::CHATGPT_CODEX_BASE_URL; pub use codex_protocol::account::ProviderAccount; pub use provider::ModelProvider;