From eed07ed50b0a63fbfea7264f073c6fc9eec5dab6 Mon Sep 17 00:00:00 2001 From: celia-oai Date: Mon, 6 Jul 2026 17:09:34 -0700 Subject: [PATCH 01/12] app-server: add managed Bedrock login --- codex-rs/app-server/src/request_processors.rs | 2 + .../request_processors/account_processor.rs | 65 +++ .../src/request_processors/bedrock_auth.rs | 35 ++ .../request_processors/config_processor.rs | 2 +- .../tests/common/test_app_server.rs | 14 + codex-rs/app-server/tests/suite/v2/account.rs | 436 ++++++++++++++++++ .../src/amazon_bedrock/mantle.rs | 7 +- .../model-provider/src/amazon_bedrock/mod.rs | 1 + codex-rs/model-provider/src/lib.rs | 2 + 9 files changed, 562 insertions(+), 2 deletions(-) create mode 100644 codex-rs/app-server/src/request_processors/bedrock_auth.rs 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 f4acc5a8e45a..243e45632a3a 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; @@ -293,6 +295,10 @@ impl AccountRequestProcessor { ) .await; } + LoginAccountParams::AmazonBedrock { api_key, region } => { + self.login_amazon_bedrock_v2(request_id, api_key, region) + .await; + } } Ok(()) } @@ -356,6 +362,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); + } + } + + 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}")))?; + set_user_model_provider_to_bedrock(&self.config_manager).await?; + 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, 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..be9d6c5a301b 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,70 @@ 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(), + )?) +} + +fn aws_managed_bedrock_config() -> CreateConfigTomlParams { + CreateConfigTomlParams { + model_provider_id: Some("amazon-bedrock".to_string()), + extra_provider_config: Some( + r#"[model_providers.amazon-bedrock.aws] +profile = "codex-bedrock" +region = "us-west-2" +"# + .to_string(), + ), + ..Default::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 +1085,368 @@ 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::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).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::new(codex_home.path()).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(), aws_managed_bedrock_config())?; + + let mut mcp = + TestAppServer::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).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::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).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_persists_under_provider_override() -> Result<()> { + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?; + + let mut mcp = TestAppServer::new_with_args( + codex_home.path(), + &["-c", "model_provider=\"mock_provider\""], + ) + .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); + 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: None, + requires_openai_auth: false, + } + ); + 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::new(codex_home.path()).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_allowed_when_forced_api() -> Result<()> { + let codex_home = TempDir::new()?; + create_config_toml( + codex_home.path(), + CreateConfigTomlParams { + forced_method: Some("api".to_string()), + ..Default::default() + }, + )?; + + let mut mcp = TestAppServer::new(codex_home.path()).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 = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + + assert_eq!( + to_response::(response)?, + LoginAccountResponse::AmazonBedrock {} + ); + 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::new(codex_home.path()).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; From 912387141f1109a8a4813417971d3caeea8bd8fa Mon Sep 17 00:00:00 2001 From: celia-oai Date: Wed, 8 Jul 2026 14:25:27 -0700 Subject: [PATCH 02/12] Remove redundant managed Bedrock login tests --- codex-rs/app-server/tests/suite/v2/account.rs | 92 ------------------- 1 file changed, 92 deletions(-) diff --git a/codex-rs/app-server/tests/suite/v2/account.rs b/codex-rs/app-server/tests/suite/v2/account.rs index be9d6c5a301b..d12352f79643 100644 --- a/codex-rs/app-server/tests/suite/v2/account.rs +++ b/codex-rs/app-server/tests/suite/v2/account.rs @@ -1274,69 +1274,6 @@ async fn login_account_amazon_bedrock_rejects_invalid_credentials_without_change Ok(()) } -#[tokio::test] -async fn login_account_amazon_bedrock_persists_under_provider_override() -> Result<()> { - let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?; - - let mut mcp = TestAppServer::new_with_args( - codex_home.path(), - &["-c", "model_provider=\"mock_provider\""], - ) - .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); - 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: None, - requires_openai_auth: false, - } - ); - Ok(()) -} - #[tokio::test] async fn login_account_amazon_bedrock_rejected_when_forced_chatgpt() -> Result<()> { let codex_home = TempDir::new()?; @@ -1367,35 +1304,6 @@ async fn login_account_amazon_bedrock_rejected_when_forced_chatgpt() -> Result<( Ok(()) } -#[tokio::test] -async fn login_account_amazon_bedrock_allowed_when_forced_api() -> Result<()> { - let codex_home = TempDir::new()?; - create_config_toml( - codex_home.path(), - CreateConfigTomlParams { - forced_method: Some("api".to_string()), - ..Default::default() - }, - )?; - - let mut mcp = TestAppServer::new(codex_home.path()).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 = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - - assert_eq!( - to_response::(response)?, - LoginAccountResponse::AmazonBedrock {} - ); - Ok(()) -} - #[tokio::test] async fn login_account_amazon_bedrock_rejected_with_external_chatgpt_auth() -> Result<()> { let codex_home = TempDir::new()?; From 622cba673dcc233d47f7f1f17fe9fc429945638c Mon Sep 17 00:00:00 2001 From: celia-oai Date: Wed, 8 Jul 2026 15:40:18 -0700 Subject: [PATCH 03/12] app-server: reload config for account reads --- .../src/request_processors/account_processor.rs | 17 ++++++++++++----- codex-rs/app-server/tests/suite/v2/account.rs | 16 +--------------- 2 files changed, 13 insertions(+), 20 deletions(-) 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 243e45632a3a..8a7243ffdd5b 100644 --- a/codex-rs/app-server/src/request_processors/account_processor.rs +++ b/codex-rs/app-server/src/request_processors/account_processor.rs @@ -193,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, @@ -899,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 { @@ -967,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/tests/suite/v2/account.rs b/codex-rs/app-server/tests/suite/v2/account.rs index d12352f79643..db39c20d73d5 100644 --- a/codex-rs/app-server/tests/suite/v2/account.rs +++ b/codex-rs/app-server/tests/suite/v2/account.rs @@ -166,20 +166,6 @@ fn load_file_auth(codex_home: &Path) -> Result> { )?) } -fn aws_managed_bedrock_config() -> CreateConfigTomlParams { - CreateConfigTomlParams { - model_provider_id: Some("amazon-bedrock".to_string()), - extra_provider_config: Some( - r#"[model_providers.amazon-bedrock.aws] -profile = "codex-bedrock" -region = "us-west-2" -"# - .to_string(), - ), - ..Default::default() - } -} - async fn read_account(mcp: &mut TestAppServer) -> Result { let request_id = mcp .send_get_account_request(GetAccountParams { @@ -1196,7 +1182,7 @@ async fn managed_bedrock_login_requires_experimental_api() -> Result<()> { #[tokio::test] async fn login_managed_bedrock_updates_active_bedrock_account() -> Result<()> { let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), aws_managed_bedrock_config())?; + create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?; let mut mcp = TestAppServer::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; From 89b772f721c891d8e0127eba7067a7d8121625a5 Mon Sep 17 00:00:00 2001 From: celia-oai Date: Wed, 8 Jul 2026 16:10:06 -0700 Subject: [PATCH 04/12] app-server-test-client: add managed Bedrock login --- codex-rs/app-server-test-client/README.md | 24 +++++ codex-rs/app-server-test-client/src/lib.rs | 106 ++++++++++++++++++++- 2 files changed, 126 insertions(+), 4 deletions(-) diff --git a/codex-rs/app-server-test-client/README.md b/codex-rs/app-server-test-client/README.md index b1afc778723a..c49c49f9e422 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-amazon-bedrock-login` initializes the experimental app-server API, logs in with an Amazon +Bedrock API key, waits for the login completion notification, and verifies that `account/read` +reports a Codex-managed Amazon Bedrock account. 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" +export AWS_BEARER_TOKEN_BEDROCK="" + +cargo build -p codex-cli --bin codex +cargo run -p codex-app-server-test-client -- \ + --codex-bin ./target/debug/codex \ + test-amazon-bedrock-login \ + --region us-west-2 +``` + +The key can instead be passed with `--api-key`, but the environment variable avoids exposing it in +shell history. The test client redacts `apiKey` from its outbound request log. After login, unset +`AWS_BEARER_TOKEN_BEDROCK` and 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..b6b485fa1023 100644 --- a/codex-rs/app-server-test-client/src/lib.rs +++ b/codex-rs/app-server-test-client/src/lib.rs @@ -25,6 +25,7 @@ use anyhow::bail; use clap::ArgAction; use clap::Parser; use clap::Subcommand; +use codex_app_server_protocol::Account; use codex_app_server_protocol::AccountLoginCompletedNotification; use codex_app_server_protocol::AskForApproval; use codex_app_server_protocol::ClientInfo; @@ -37,7 +38,9 @@ use codex_app_server_protocol::DynamicToolSpec; use codex_app_server_protocol::FileChangeApprovalDecision; use codex_app_server_protocol::FileChangeRequestApprovalParams; use codex_app_server_protocol::FileChangeRequestApprovalResponse; +use codex_app_server_protocol::GetAccountParams; use codex_app_server_protocol::GetAccountRateLimitsResponse; +use codex_app_server_protocol::GetAccountResponse; use codex_app_server_protocol::InitializeCapabilities; use codex_app_server_protocol::InitializeParams; use codex_app_server_protocol::InitializeResponse; @@ -70,6 +73,7 @@ use codex_app_server_protocol::UserInput as V2UserInput; use codex_core::config::Config; use codex_otel::OtelProvider; use codex_otel::current_span_w3c_trace_context; +use codex_protocol::account::AmazonBedrockCredentialSource; use codex_protocol::dynamic_tools::normalize_dynamic_tool_specs; use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::protocol::W3cTraceContext; @@ -236,6 +240,15 @@ enum CliCommand { #[arg(long, default_value_t = false)] device_code: bool, }, + /// Log in with a Codex-managed Amazon Bedrock API key and verify the account state. + TestAmazonBedrockLogin { + /// Amazon Bedrock API key. Defaults to AWS_BEARER_TOKEN_BEDROCK. + #[arg(long, env = "AWS_BEARER_TOKEN_BEDROCK", hide_env_values = true)] + api_key: String, + /// AWS Region for the Amazon Bedrock Mantle endpoint. + #[arg(long, env = "AWS_REGION")] + region: String, + }, /// Fetch the current account rate limits from the Codex app-server. GetAccountRateLimits, /// List the available models from the Codex app-server. @@ -420,6 +433,11 @@ pub async fn run() -> Result<()> { let endpoint = resolve_endpoint(codex_bin, url)?; test_login(&endpoint, &config_overrides, device_code).await } + CliCommand::TestAmazonBedrockLogin { api_key, region } => { + ensure_dynamic_tools_unused(&dynamic_tools, "test-amazon-bedrock-login")?; + let endpoint = resolve_endpoint(codex_bin, url)?; + test_amazon_bedrock_login(&endpoint, &config_overrides, api_key, region).await + } CliCommand::GetAccountRateLimits => { ensure_dynamic_tools_unused(&dynamic_tools, "get-account-rate-limits")?; let endpoint = resolve_endpoint(codex_bin, url)?; @@ -1158,7 +1176,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 { @@ -1177,6 +1195,80 @@ async fn test_login( .await } +async fn test_amazon_bedrock_login( + endpoint: &Endpoint, + config_overrides: &[String], + api_key: String, + region: String, +) -> Result<()> { + with_client( + "test-amazon-bedrock-login", + endpoint, + config_overrides, + |client| { + let initialize = + client.initialize_with_experimental_api(/*experimental_api*/ true)?; + println!("< initialize response: {initialize:?}"); + + 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:?}"); + if login_response != (LoginAccountResponse::AmazonBedrock {}) { + bail!("expected Amazon Bedrock login response, got {login_response:?}"); + } + + let completion = + client.wait_for_account_login_completion(/*expected_login_id*/ None)?; + println!("< account/login/completed notification: {completion:?}"); + if !completion.success { + bail!( + "Amazon Bedrock login failed: {}", + completion + .error + .as_deref() + .unwrap_or("unknown error from account/login/completed") + ); + } + + let request_id = client.request_id(); + let account: GetAccountResponse = client.send_request( + ClientRequest::GetAccount { + request_id: request_id.clone(), + params: GetAccountParams { + refresh_token: false, + }, + }, + request_id, + "account/read", + )?; + println!("< account/read response: {account:?}"); + let expected_account = GetAccountResponse { + account: Some(Account::AmazonBedrock { + credential_source: AmazonBedrockCredentialSource::CodexManaged, + }), + requires_openai_auth: false, + }; + if account != expected_account { + bail!("expected managed Amazon Bedrock account, got {account:?}"); + } + + println!("Amazon Bedrock login succeeded."); + Ok(()) + }, + ) + .await +} + async fn get_account_rate_limits(endpoint: &Endpoint, config_overrides: &[String]) -> Result<()> { with_client( "get-account-rate-limits", @@ -1799,7 +1891,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 +1899,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 +2048,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) } From a551b191cf492246c769cb997c4127b72d44732d Mon Sep 17 00:00:00 2001 From: celia-oai Date: Wed, 8 Jul 2026 16:20:49 -0700 Subject: [PATCH 05/12] app-server-test-client: fold Bedrock into test-login --- codex-rs/app-server-test-client/README.md | 11 +- codex-rs/app-server-test-client/src/lib.rs | 155 ++++++++------------- 2 files changed, 61 insertions(+), 105 deletions(-) diff --git a/codex-rs/app-server-test-client/README.md b/codex-rs/app-server-test-client/README.md index c49c49f9e422..7c5d0a3e0720 100644 --- a/codex-rs/app-server-test-client/README.md +++ b/codex-rs/app-server-test-client/README.md @@ -24,10 +24,10 @@ and the client will send the response and continue streaming the same turn. ## Testing Codex-managed Amazon Bedrock login -`test-amazon-bedrock-login` initializes the experimental app-server API, logs in with an Amazon -Bedrock API key, waits for the login completion notification, and verifies that `account/read` -reports a Codex-managed Amazon Bedrock account. Login replaces the current primary credential and -sets `model_provider = "amazon-bedrock"`, so use an isolated `CODEX_HOME` when testing. +`test-login --amazon-bedrock` initializes the experimental app-server API and sends an +`account/login/start` request with an Amazon Bedrock API key. 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)" @@ -37,7 +37,8 @@ export AWS_BEARER_TOKEN_BEDROCK="" cargo build -p codex-cli --bin codex cargo run -p codex-app-server-test-client -- \ --codex-bin ./target/debug/codex \ - test-amazon-bedrock-login \ + test-login \ + --amazon-bedrock \ --region us-west-2 ``` diff --git a/codex-rs/app-server-test-client/src/lib.rs b/codex-rs/app-server-test-client/src/lib.rs index b6b485fa1023..a3fd631eedb2 100644 --- a/codex-rs/app-server-test-client/src/lib.rs +++ b/codex-rs/app-server-test-client/src/lib.rs @@ -25,7 +25,6 @@ use anyhow::bail; use clap::ArgAction; use clap::Parser; use clap::Subcommand; -use codex_app_server_protocol::Account; use codex_app_server_protocol::AccountLoginCompletedNotification; use codex_app_server_protocol::AskForApproval; use codex_app_server_protocol::ClientInfo; @@ -38,9 +37,7 @@ use codex_app_server_protocol::DynamicToolSpec; use codex_app_server_protocol::FileChangeApprovalDecision; use codex_app_server_protocol::FileChangeRequestApprovalParams; use codex_app_server_protocol::FileChangeRequestApprovalResponse; -use codex_app_server_protocol::GetAccountParams; use codex_app_server_protocol::GetAccountRateLimitsResponse; -use codex_app_server_protocol::GetAccountResponse; use codex_app_server_protocol::InitializeCapabilities; use codex_app_server_protocol::InitializeParams; use codex_app_server_protocol::InitializeResponse; @@ -73,7 +70,6 @@ use codex_app_server_protocol::UserInput as V2UserInput; use codex_core::config::Config; use codex_otel::OtelProvider; use codex_otel::current_span_w3c_trace_context; -use codex_protocol::account::AmazonBedrockCredentialSource; use codex_protocol::dynamic_tools::normalize_dynamic_tool_specs; use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::protocol::W3cTraceContext; @@ -234,20 +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, - }, - /// Log in with a Codex-managed Amazon Bedrock API key and verify the account state. - TestAmazonBedrockLogin { + /// 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. Defaults to AWS_BEARER_TOKEN_BEDROCK. #[arg(long, env = "AWS_BEARER_TOKEN_BEDROCK", hide_env_values = true)] - api_key: String, + api_key: Option, /// AWS Region for the Amazon Bedrock Mantle endpoint. #[arg(long, env = "AWS_REGION")] - region: String, + region: Option, }, /// Fetch the current account rate limits from the Codex app-server. GetAccountRateLimits, @@ -326,6 +322,12 @@ enum CliCommand { }, } +enum TestLoginMode { + ChatgptBrowser, + ChatgptDeviceCode, + AmazonBedrock { api_key: String, region: String }, +} + pub async fn run() -> Result<()> { let Cli { codex_bin, @@ -428,15 +430,27 @@ 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 - } - CliCommand::TestAmazonBedrockLogin { api_key, region } => { - ensure_dynamic_tools_unused(&dynamic_tools, "test-amazon-bedrock-login")?; - let endpoint = resolve_endpoint(codex_bin, url)?; - test_amazon_bedrock_login(&endpoint, &config_overrides, api_key, region).await + let mode = if amazon_bedrock { + let api_key = api_key.context( + "--api-key or AWS_BEARER_TOKEN_BEDROCK is required with --amazon-bedrock", + )?; + let region = + region.context("--region or AWS_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")?; @@ -1146,16 +1160,31 @@ 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:?}"); + return Ok(()); + } }; println!("< account/login/start response: {login_response:?}"); let login_id = match login_response { @@ -1176,7 +1205,7 @@ async fn test_login( _ => bail!("expected chatgpt login response"), }; - let completion = client.wait_for_account_login_completion(Some(&login_id))?; + let completion = client.wait_for_account_login_completion(&login_id)?; println!("< account/login/completed notification: {completion:?}"); if completion.success { @@ -1195,80 +1224,6 @@ async fn test_login( .await } -async fn test_amazon_bedrock_login( - endpoint: &Endpoint, - config_overrides: &[String], - api_key: String, - region: String, -) -> Result<()> { - with_client( - "test-amazon-bedrock-login", - endpoint, - config_overrides, - |client| { - let initialize = - client.initialize_with_experimental_api(/*experimental_api*/ true)?; - println!("< initialize response: {initialize:?}"); - - 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:?}"); - if login_response != (LoginAccountResponse::AmazonBedrock {}) { - bail!("expected Amazon Bedrock login response, got {login_response:?}"); - } - - let completion = - client.wait_for_account_login_completion(/*expected_login_id*/ None)?; - println!("< account/login/completed notification: {completion:?}"); - if !completion.success { - bail!( - "Amazon Bedrock login failed: {}", - completion - .error - .as_deref() - .unwrap_or("unknown error from account/login/completed") - ); - } - - let request_id = client.request_id(); - let account: GetAccountResponse = client.send_request( - ClientRequest::GetAccount { - request_id: request_id.clone(), - params: GetAccountParams { - refresh_token: false, - }, - }, - request_id, - "account/read", - )?; - println!("< account/read response: {account:?}"); - let expected_account = GetAccountResponse { - account: Some(Account::AmazonBedrock { - credential_source: AmazonBedrockCredentialSource::CodexManaged, - }), - requires_openai_auth: false, - }; - if account != expected_account { - bail!("expected managed Amazon Bedrock account, got {account:?}"); - } - - println!("Amazon Bedrock login succeeded."); - Ok(()) - }, - ) - .await -} - async fn get_account_rate_limits(endpoint: &Endpoint, config_overrides: &[String]) -> Result<()> { with_client( "get-account-rate-limits", @@ -1891,7 +1846,7 @@ impl CodexClient { fn wait_for_account_login_completion( &mut self, - expected_login_id: Option<&str>, + expected_login_id: &str, ) -> Result { loop { let notification = self.next_notification()?; @@ -1899,7 +1854,7 @@ impl CodexClient { if let Ok(server_notification) = ServerNotification::try_from(notification) { match server_notification { ServerNotification::AccountLoginCompleted(completion) => { - if completion.login_id.as_deref() == expected_login_id { + if completion.login_id.as_deref() == Some(expected_login_id) { return Ok(completion); } From ca8dd0616d8c0315539a1c2a5231e4ed4420d820 Mon Sep 17 00:00:00 2001 From: celia-oai Date: Wed, 8 Jul 2026 16:24:04 -0700 Subject: [PATCH 06/12] app-server-test-client: require Bedrock login arguments --- codex-rs/app-server-test-client/README.md | 8 +++----- codex-rs/app-server-test-client/src/lib.rs | 13 +++++-------- 2 files changed, 8 insertions(+), 13 deletions(-) diff --git a/codex-rs/app-server-test-client/README.md b/codex-rs/app-server-test-client/README.md index 7c5d0a3e0720..872eeb683989 100644 --- a/codex-rs/app-server-test-client/README.md +++ b/codex-rs/app-server-test-client/README.md @@ -32,20 +32,18 @@ testing. ```bash export CODEX_HOME="$(mktemp -d)" printf 'cli_auth_credentials_store = "file"\n' > "$CODEX_HOME/config.toml" -export AWS_BEARER_TOKEN_BEDROCK="" 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 key can instead be passed with `--api-key`, but the environment variable avoids exposing it in -shell history. The test client redacts `apiKey` from its outbound request log. After login, unset -`AWS_BEARER_TOKEN_BEDROCK` and start a fresh Codex process with the same `CODEX_HOME` to verify that -it uses the persisted managed credential. +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 diff --git a/codex-rs/app-server-test-client/src/lib.rs b/codex-rs/app-server-test-client/src/lib.rs index a3fd631eedb2..b3796fda0768 100644 --- a/codex-rs/app-server-test-client/src/lib.rs +++ b/codex-rs/app-server-test-client/src/lib.rs @@ -238,11 +238,11 @@ enum CliCommand { /// 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. Defaults to AWS_BEARER_TOKEN_BEDROCK. - #[arg(long, env = "AWS_BEARER_TOKEN_BEDROCK", hide_env_values = true)] + /// Amazon Bedrock API key. + #[arg(long, value_name = "API_KEY")] api_key: Option, /// AWS Region for the Amazon Bedrock Mantle endpoint. - #[arg(long, env = "AWS_REGION")] + #[arg(long, value_name = "REGION")] region: Option, }, /// Fetch the current account rate limits from the Codex app-server. @@ -439,11 +439,8 @@ pub async fn run() -> Result<()> { ensure_dynamic_tools_unused(&dynamic_tools, "test-login")?; let endpoint = resolve_endpoint(codex_bin, url)?; let mode = if amazon_bedrock { - let api_key = api_key.context( - "--api-key or AWS_BEARER_TOKEN_BEDROCK is required with --amazon-bedrock", - )?; - let region = - region.context("--region or AWS_REGION is required with --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 From 67029393c3684e1156394423a3c89a46149806ba Mon Sep 17 00:00:00 2001 From: celia-oai Date: Wed, 8 Jul 2026 16:30:35 -0700 Subject: [PATCH 07/12] app-server-test-client: wait for Bedrock login notifications --- codex-rs/app-server-test-client/README.md | 5 +++-- codex-rs/app-server-test-client/src/lib.rs | 20 +++++++++++++++++--- 2 files changed, 20 insertions(+), 5 deletions(-) diff --git a/codex-rs/app-server-test-client/README.md b/codex-rs/app-server-test-client/README.md index 872eeb683989..9f73b1db6f5c 100644 --- a/codex-rs/app-server-test-client/README.md +++ b/codex-rs/app-server-test-client/README.md @@ -24,8 +24,9 @@ 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 and sends an -`account/login/start` request with an Amazon Bedrock API key. Login replaces the current primary +`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. diff --git a/codex-rs/app-server-test-client/src/lib.rs b/codex-rs/app-server-test-client/src/lib.rs index b3796fda0768..10d66dbc01d5 100644 --- a/codex-rs/app-server-test-client/src/lib.rs +++ b/codex-rs/app-server-test-client/src/lib.rs @@ -1180,6 +1180,20 @@ async fn test_login( "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(()); } }; @@ -1202,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 { @@ -1843,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()?; @@ -1851,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); } From 476ee1b40b69f2fdd4ce032b312af13429257dd1 Mon Sep 17 00:00:00 2001 From: celia-oai Date: Mon, 6 Jul 2026 17:25:32 -0700 Subject: [PATCH 08/12] app-server: add managed Bedrock logout --- codex-rs/app-server/README.md | 2 + codex-rs/app-server/src/request_processors.rs | 1 - .../request_processors/account_processor.rs | 43 +++--- .../src/request_processors/bedrock_auth.rs | 64 +++++++++ codex-rs/app-server/tests/suite/v2/account.rs | 135 ++++++++++++++++++ 5 files changed, 227 insertions(+), 18 deletions(-) diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 6280cf520409..a4f75d1541e9 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -2060,6 +2060,8 @@ Codex stores the key and region as the primary Codex auth, replacing any previou { "method": "account/updated", "params": { "authMode": null, "planType": null } } ``` +When Codex-managed Amazon Bedrock auth is stored, logout removes that credential and clears the user-level `model_provider` only when it is still `"amazon-bedrock"`. A concurrent user config change is preserved. For AWS-managed Bedrock auth, logout is a no-op because the AWS SDK credential chain is managed outside Codex. + ### 7) Rate limits (ChatGPT) ```json diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index b58e430d7e99..7556c009b90d 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -35,7 +35,6 @@ use codex_app_server_protocol::AppTemplateUnavailableReason; use codex_app_server_protocol::AppsListParams; use codex_app_server_protocol::AppsListResponse; use codex_app_server_protocol::AskForApproval; -use codex_app_server_protocol::AuthMode; use codex_app_server_protocol::CancelLoginAccountParams; use codex_app_server_protocol::CancelLoginAccountResponse; use codex_app_server_protocol::CancelLoginAccountStatus; 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 8a7243ffdd5b..f4d0c9459f51 100644 --- a/codex-rs/app-server/src/request_processors/account_processor.rs +++ b/codex-rs/app-server/src/request_processors/account_processor.rs @@ -1,4 +1,6 @@ +use super::bedrock_auth::clear_user_model_provider_if_bedrock; use super::bedrock_auth::set_user_model_provider_to_bedrock; +use super::bedrock_auth::user_model_provider_state; use super::*; use crate::auth_mode::auth_mode_to_api; use crate::external_auth::ExternalAuthBridge; @@ -824,7 +826,7 @@ impl AccountRequestProcessor { } } - async fn logout_common(&self) -> std::result::Result, JSONRPCErrorError> { + async fn logout_common(&self) -> Result<(), JSONRPCErrorError> { // Cancel any active login attempt. { let mut guard = self.active_login.lock().await; @@ -833,6 +835,19 @@ impl AccountRequestProcessor { } } + let managed_bedrock_auth = matches!( + self.auth_manager.auth_cached(), + Some(CodexAuth::BedrockApiKey(_)) + ); + let user_model_provider = if managed_bedrock_auth { + Some(user_model_provider_state(&self.config_manager).await?) + } else { + None + }; + if self.config.model_provider.is_amazon_bedrock() && !managed_bedrock_auth { + return Ok(()); + } + match self.auth_manager.logout_with_revoke().await { Ok(_) => {} Err(err) => { @@ -847,26 +862,20 @@ impl AccountRequestProcessor { ) .await; - // Reflect the current auth method after logout (likely None). - Ok(self - .auth_manager - .auth_cached() - .as_ref() - .map(CodexAuth::api_auth_mode) - .map(auth_mode_to_api)) + if let Some(user_model_provider) = user_model_provider { + clear_user_model_provider_if_bedrock(&self.config_manager, user_model_provider).await?; + } + + Ok(()) } async fn logout_v2(&self, request_id: ConnectionRequestId) -> Result<(), JSONRPCErrorError> { let result = self.logout_common().await; - let account_updated = - result - .as_ref() - .ok() - .cloned() - .map(|auth_mode| AccountUpdatedNotification { - auth_mode, - plan_type: None, - }); + let account_updated = if result.is_ok() { + Some(self.current_account_updated_notification()) + } else { + None + }; self.outgoing .send_result(request_id, result.map(|_| LogoutAccountResponse {})) .await; diff --git a/codex-rs/app-server/src/request_processors/bedrock_auth.rs b/codex-rs/app-server/src/request_processors/bedrock_auth.rs index 88576e5ffe9c..333d905d95f6 100644 --- a/codex-rs/app-server/src/request_processors/bedrock_auth.rs +++ b/codex-rs/app-server/src/request_processors/bedrock_auth.rs @@ -1,10 +1,18 @@ use super::config_processor::map_error as map_config_error; use crate::config_manager::ConfigManager; +use crate::error_code::internal_error; +use codex_app_server_protocol::ConfigLayerSource; +use codex_app_server_protocol::ConfigReadParams; 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) struct UserModelProviderState { + model_provider: Option, + version: Option, +} + pub(super) async fn set_user_model_provider_to_bedrock( config_manager: &ConfigManager, ) -> Result<(), JSONRPCErrorError> { @@ -16,6 +24,62 @@ pub(super) async fn set_user_model_provider_to_bedrock( .await } +pub(super) async fn clear_user_model_provider_if_bedrock( + config_manager: &ConfigManager, + user_model_provider: UserModelProviderState, +) -> Result<(), JSONRPCErrorError> { + if user_model_provider.model_provider.as_deref() == Some(AMAZON_BEDROCK_PROVIDER_ID) { + write_user_model_provider( + config_manager, + serde_json::Value::Null, + user_model_provider.version, + ) + .await?; + } + + Ok(()) +} + +pub(super) async fn user_model_provider_state( + config_manager: &ConfigManager, +) -> Result { + let user_config_path = config_manager + .user_config_path() + .map_err(|err| internal_error(format!("failed to resolve user config path: {err}")))?; + let response = config_manager + .read(ConfigReadParams { + include_layers: true, + cwd: None, + }) + .await + .map_err(|err| internal_error(format!("failed to read user config: {err}")))?; + let layer = response + .layers + .unwrap_or_default() + .into_iter() + .find(|layer| { + matches!( + &layer.name, + ConfigLayerSource::User { file, .. } if file == &user_config_path + ) + }); + let Some(layer) = layer else { + return Ok(UserModelProviderState { + model_provider: None, + version: None, + }); + }; + let model_provider = layer + .config + .get("model_provider") + .and_then(serde_json::Value::as_str) + .map(str::to_string); + Ok(UserModelProviderState { + model_provider, + version: Some(layer.version), + }) +} + async fn write_user_model_provider( config_manager: &ConfigManager, value: serde_json::Value, diff --git a/codex-rs/app-server/tests/suite/v2/account.rs b/codex-rs/app-server/tests/suite/v2/account.rs index db39c20d73d5..aca640141fe6 100644 --- a/codex-rs/app-server/tests/suite/v2/account.rs +++ b/codex-rs/app-server/tests/suite/v2/account.rs @@ -1143,6 +1143,141 @@ async fn login_amazon_bedrock_replaces_primary_auth_and_persists_provider() -> R Ok(()) } +#[tokio::test] +async fn logout_managed_bedrock_restores_aws_managed_account() -> Result<()> { + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), aws_managed_bedrock_config())?; + let mut expected_config = read_config_toml(codex_home.path())?; + expected_config + .as_table_mut() + .expect("config should be a table") + .remove("model_provider"); + + let mut mcp = + TestAppServer::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).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 = 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?; + + let request_id = mcp.send_logout_account_request().await?; + let response = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + assert_eq!( + to_response::(response)?, + LogoutAccountResponse {} + ); + assert_eq!(load_file_auth(codex_home.path())?, None); + assert_eq!(read_config_toml(codex_home.path())?, expected_config); + assert_account_updated(&mut mcp, /*auth_mode*/ None).await?; + assert_eq!( + read_account(&mut mcp).await?, + GetAccountResponse { + account: Some(Account::AmazonBedrock { + credential_source: AmazonBedrockCredentialSource::AwsManaged, + }), + requires_openai_auth: false, + } + ); + Ok(()) +} + +#[tokio::test] +async fn logout_aws_managed_bedrock_preserves_openai_auth_and_config() -> Result<()> { + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), aws_managed_bedrock_config())?; + login_with_api_key( + codex_home.path(), + "sk-test-key", + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + let expected_auth = load_file_auth(codex_home.path())?; + let expected_config = read_config_toml(codex_home.path())?; + + let mut mcp = + TestAppServer::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + let request_id = mcp.send_logout_account_request().await?; + let response = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + assert_eq!( + to_response::(response)?, + LogoutAccountResponse {} + ); + assert_eq!(load_file_auth(codex_home.path())?, expected_auth); + assert_eq!(read_config_toml(codex_home.path())?, expected_config); + assert_account_updated(&mut mcp, Some(AuthMode::ApiKey)).await?; + Ok(()) +} + +#[tokio::test] +async fn logout_managed_bedrock_preserves_changed_provider_without_experimental_api() -> Result<()> +{ + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?; + login_with_bedrock_api_key( + codex_home.path(), + "managed-bedrock-api-key", + "us-west-2", + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + let expected_config = read_config_toml(codex_home.path())?; + + let mut mcp = TestAppServer::new(codex_home.path()).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_logout_account_request().await?; + let response = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + assert_eq!( + to_response::(response)?, + LogoutAccountResponse {} + ); + assert_eq!(load_file_auth(codex_home.path())?, None); + assert_eq!(read_config_toml(codex_home.path())?, expected_config); + assert_account_updated(&mut mcp, /*auth_mode*/ None).await?; + Ok(()) +} + #[tokio::test] async fn managed_bedrock_login_requires_experimental_api() -> Result<()> { let codex_home = TempDir::new()?; From 340776f912476fc37d54aac374468b3538fdf55e Mon Sep 17 00:00:00 2001 From: celia-oai Date: Wed, 8 Jul 2026 14:50:25 -0700 Subject: [PATCH 09/12] Simplify managed Bedrock provider cleanup --- .../app-server/src/config_manager_service.rs | 37 ++++++++ .../src/config_manager_service_tests.rs | 34 +++++++ .../request_processors/account_processor.rs | 10 +- .../src/request_processors/bedrock_auth.rs | 93 +++---------------- codex-rs/app-server/tests/suite/v2/account.rs | 6 +- 5 files changed, 91 insertions(+), 89 deletions(-) diff --git a/codex-rs/app-server/src/config_manager_service.rs b/codex-rs/app-server/src/config_manager_service.rs index 94980bf8e012..5640c9cb7406 100644 --- a/codex-rs/app-server/src/config_manager_service.rs +++ b/codex-rs/app-server/src/config_manager_service.rs @@ -181,6 +181,43 @@ impl ConfigManager { .await } + /// Clears a value from the active user config only when its current raw value matches. + pub(crate) async fn clear_user_value_if_matches( + &self, + key_path: &str, + expected_value: JsonValue, + ) -> Result<(), ConfigManagerError> { + let layers = self + .load_thread_agnostic_config() + .await + .map_err(|err| ConfigManagerError::io("failed to load configuration", err))?; + let Some(user_layer) = layers.get_active_user_layer() else { + return Ok(()); + }; + let segments = parse_key_path(key_path).map_err(|message| { + ConfigManagerError::write(ConfigWriteErrorCode::ConfigValidationError, message) + })?; + let expected_value = parse_value(expected_value).map_err(|message| { + ConfigManagerError::write(ConfigWriteErrorCode::ConfigValidationError, message) + })?; + if value_at_path(&user_layer.config, &segments) != expected_value.as_ref() { + return Ok(()); + } + let expected_version = Some(user_layer.version.clone()); + + self.apply_edits( + /*file_path*/ None, + expected_version, + vec![( + key_path.to_string(), + JsonValue::Null, + MergeStrategy::Replace, + )], + ) + .await?; + Ok(()) + } + pub(crate) async fn batch_write( &self, params: ConfigBatchWriteParams, diff --git a/codex-rs/app-server/src/config_manager_service_tests.rs b/codex-rs/app-server/src/config_manager_service_tests.rs index 1ec79a749b3c..3407c1b423ff 100644 --- a/codex-rs/app-server/src/config_manager_service_tests.rs +++ b/codex-rs/app-server/src/config_manager_service_tests.rs @@ -130,6 +130,40 @@ async fn clear_missing_nested_config_is_noop() -> Result<()> { Ok(()) } +#[tokio::test] +async fn clear_user_value_if_matches_clears_matching_value() -> Result<()> { + let tmp = tempdir().expect("tempdir"); + let path = tmp.path().join(CONFIG_TOML_FILE); + std::fs::write(&path, "model = \"gpt-5.2\"\napproval_policy = \"never\"\n")?; + + let service = ConfigManager::without_managed_config_for_tests(tmp.path().to_path_buf()); + service + .clear_user_value_if_matches("model", serde_json::json!("gpt-5.2")) + .await?; + + assert_eq!( + std::fs::read_to_string(&path)?, + "approval_policy = \"never\"\n" + ); + Ok(()) +} + +#[tokio::test] +async fn clear_user_value_if_matches_preserves_non_matching_value() -> Result<()> { + let tmp = tempdir().expect("tempdir"); + let path = tmp.path().join(CONFIG_TOML_FILE); + let original = "model = \"gpt-5.2\"\napproval_policy = \"never\"\n"; + std::fs::write(&path, original)?; + + let service = ConfigManager::without_managed_config_for_tests(tmp.path().to_path_buf()); + service + .clear_user_value_if_matches("model", serde_json::json!("gpt-5.3")) + .await?; + + assert_eq!(std::fs::read_to_string(&path)?, original); + Ok(()) +} + #[tokio::test] async fn write_value_rejects_legacy_profile_selector() -> Result<()> { let tmp = tempdir().expect("tempdir"); 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 f4d0c9459f51..1b8bce00f794 100644 --- a/codex-rs/app-server/src/request_processors/account_processor.rs +++ b/codex-rs/app-server/src/request_processors/account_processor.rs @@ -1,6 +1,5 @@ use super::bedrock_auth::clear_user_model_provider_if_bedrock; use super::bedrock_auth::set_user_model_provider_to_bedrock; -use super::bedrock_auth::user_model_provider_state; use super::*; use crate::auth_mode::auth_mode_to_api; use crate::external_auth::ExternalAuthBridge; @@ -839,11 +838,6 @@ impl AccountRequestProcessor { self.auth_manager.auth_cached(), Some(CodexAuth::BedrockApiKey(_)) ); - let user_model_provider = if managed_bedrock_auth { - Some(user_model_provider_state(&self.config_manager).await?) - } else { - None - }; if self.config.model_provider.is_amazon_bedrock() && !managed_bedrock_auth { return Ok(()); } @@ -862,8 +856,8 @@ impl AccountRequestProcessor { ) .await; - if let Some(user_model_provider) = user_model_provider { - clear_user_model_provider_if_bedrock(&self.config_manager, user_model_provider).await?; + if managed_bedrock_auth { + clear_user_model_provider_if_bedrock(&self.config_manager).await?; } Ok(()) diff --git a/codex-rs/app-server/src/request_processors/bedrock_auth.rs b/codex-rs/app-server/src/request_processors/bedrock_auth.rs index 333d905d95f6..4a023d989a65 100644 --- a/codex-rs/app-server/src/request_processors/bedrock_auth.rs +++ b/codex-rs/app-server/src/request_processors/bedrock_auth.rs @@ -1,99 +1,34 @@ use super::config_processor::map_error as map_config_error; use crate::config_manager::ConfigManager; -use crate::error_code::internal_error; -use codex_app_server_protocol::ConfigLayerSource; -use codex_app_server_protocol::ConfigReadParams; 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) struct UserModelProviderState { - model_provider: Option, - version: Option, -} - 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 -} - -pub(super) async fn clear_user_model_provider_if_bedrock( - config_manager: &ConfigManager, - user_model_provider: UserModelProviderState, -) -> Result<(), JSONRPCErrorError> { - if user_model_provider.model_provider.as_deref() == Some(AMAZON_BEDROCK_PROVIDER_ID) { - write_user_model_provider( - config_manager, - serde_json::Value::Null, - user_model_provider.version, - ) - .await?; - } - - Ok(()) -} - -pub(super) async fn user_model_provider_state( - config_manager: &ConfigManager, -) -> Result { - let user_config_path = config_manager - .user_config_path() - .map_err(|err| internal_error(format!("failed to resolve user config path: {err}")))?; - let response = config_manager - .read(ConfigReadParams { - include_layers: true, - cwd: None, - }) - .await - .map_err(|err| internal_error(format!("failed to read user config: {err}")))?; - let layer = response - .layers - .unwrap_or_default() - .into_iter() - .find(|layer| { - matches!( - &layer.name, - ConfigLayerSource::User { file, .. } if file == &user_config_path - ) - }); - let Some(layer) = layer else { - return Ok(UserModelProviderState { - model_provider: None, - version: None, - }); - }; - let model_provider = layer - .config - .get("model_provider") - .and_then(serde_json::Value::as_str) - .map(str::to_string); - Ok(UserModelProviderState { - model_provider, - version: Some(layer.version), - }) -} - -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, + value: serde_json::json!(AMAZON_BEDROCK_PROVIDER_ID), merge_strategy: MergeStrategy::Replace, file_path: None, - expected_version, + expected_version: None, }) .await .map(|_| ()) .map_err(map_config_error) } + +pub(super) async fn clear_user_model_provider_if_bedrock( + config_manager: &ConfigManager, +) -> Result<(), JSONRPCErrorError> { + config_manager + .clear_user_value_if_matches( + "model_provider", + serde_json::json!(AMAZON_BEDROCK_PROVIDER_ID), + ) + .await + .map_err(map_config_error) +} diff --git a/codex-rs/app-server/tests/suite/v2/account.rs b/codex-rs/app-server/tests/suite/v2/account.rs index aca640141fe6..c97c0a42ec95 100644 --- a/codex-rs/app-server/tests/suite/v2/account.rs +++ b/codex-rs/app-server/tests/suite/v2/account.rs @@ -1236,7 +1236,7 @@ async fn logout_aws_managed_bedrock_preserves_openai_auth_and_config() -> Result async fn logout_managed_bedrock_preserves_changed_provider_without_experimental_api() -> Result<()> { let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?; + create_config_toml(codex_home.path(), aws_managed_bedrock_config())?; login_with_bedrock_api_key( codex_home.path(), "managed-bedrock-api-key", @@ -1244,7 +1244,6 @@ async fn logout_managed_bedrock_preserves_changed_provider_without_experimental_ AuthCredentialsStoreMode::File, AuthKeyringBackendKind::default(), )?; - let expected_config = read_config_toml(codex_home.path())?; let mut mcp = TestAppServer::new(codex_home.path()).await?; let initialized = mcp @@ -1262,6 +1261,9 @@ async fn logout_managed_bedrock_preserves_changed_provider_without_experimental_ .await?; assert!(matches!(initialized, JSONRPCMessage::Response(_))); + create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?; + let expected_config = read_config_toml(codex_home.path())?; + let request_id = mcp.send_logout_account_request().await?; let response = timeout( DEFAULT_READ_TIMEOUT, From fb4c66620a6e2eea079e6d7db5bb72843c9c0a24 Mon Sep 17 00:00:00 2001 From: celia-oai Date: Wed, 8 Jul 2026 15:20:59 -0700 Subject: [PATCH 10/12] changes --- codex-rs/app-server/src/request_processors.rs | 1 + .../request_processors/account_processor.rs | 31 ++++++++++++++----- 2 files changed, 24 insertions(+), 8 deletions(-) diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index 7556c009b90d..b58e430d7e99 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -35,6 +35,7 @@ use codex_app_server_protocol::AppTemplateUnavailableReason; use codex_app_server_protocol::AppsListParams; use codex_app_server_protocol::AppsListResponse; use codex_app_server_protocol::AskForApproval; +use codex_app_server_protocol::AuthMode; use codex_app_server_protocol::CancelLoginAccountParams; use codex_app_server_protocol::CancelLoginAccountResponse; use codex_app_server_protocol::CancelLoginAccountStatus; 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 1b8bce00f794..1bd13531c2ad 100644 --- a/codex-rs/app-server/src/request_processors/account_processor.rs +++ b/codex-rs/app-server/src/request_processors/account_processor.rs @@ -825,7 +825,7 @@ impl AccountRequestProcessor { } } - async fn logout_common(&self) -> Result<(), JSONRPCErrorError> { + async fn logout_common(&self) -> std::result::Result, JSONRPCErrorError> { // Cancel any active login attempt. { let mut guard = self.active_login.lock().await; @@ -839,7 +839,12 @@ impl AccountRequestProcessor { Some(CodexAuth::BedrockApiKey(_)) ); if self.config.model_provider.is_amazon_bedrock() && !managed_bedrock_auth { - return Ok(()); + return Ok(self + .auth_manager + .auth_cached() + .as_ref() + .map(CodexAuth::api_auth_mode) + .map(auth_mode_to_api)); } match self.auth_manager.logout_with_revoke().await { @@ -860,16 +865,26 @@ impl AccountRequestProcessor { clear_user_model_provider_if_bedrock(&self.config_manager).await?; } - Ok(()) + // Reflect the current auth method after logout (likely None). + Ok(self + .auth_manager + .auth_cached() + .as_ref() + .map(CodexAuth::api_auth_mode) + .map(auth_mode_to_api)) } async fn logout_v2(&self, request_id: ConnectionRequestId) -> Result<(), JSONRPCErrorError> { let result = self.logout_common().await; - let account_updated = if result.is_ok() { - Some(self.current_account_updated_notification()) - } else { - None - }; + let account_updated = + result + .as_ref() + .ok() + .cloned() + .map(|auth_mode| AccountUpdatedNotification { + auth_mode, + plan_type: None, + }); self.outgoing .send_result(request_id, result.map(|_| LogoutAccountResponse {})) .await; From 3b8549af35d92cd7a707b5bf779951ecc6f51535 Mon Sep 17 00:00:00 2001 From: celia-oai Date: Wed, 8 Jul 2026 15:43:00 -0700 Subject: [PATCH 11/12] app-server: reload config after Bedrock logout --- .../request_processors/account_processor.rs | 11 ++--- codex-rs/app-server/tests/suite/v2/account.rs | 40 ++++++++++++++++--- 2 files changed, 40 insertions(+), 11 deletions(-) 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 1bd13531c2ad..584aef36bbc3 100644 --- a/codex-rs/app-server/src/request_processors/account_processor.rs +++ b/codex-rs/app-server/src/request_processors/account_processor.rs @@ -838,7 +838,8 @@ impl AccountRequestProcessor { self.auth_manager.auth_cached(), Some(CodexAuth::BedrockApiKey(_)) ); - if self.config.model_provider.is_amazon_bedrock() && !managed_bedrock_auth { + let config = self.load_latest_config().await?; + if config.model_provider.is_amazon_bedrock() && !managed_bedrock_auth { return Ok(self .auth_manager .auth_cached() @@ -854,6 +855,10 @@ impl AccountRequestProcessor { } } + if managed_bedrock_auth { + clear_user_model_provider_if_bedrock(&self.config_manager).await?; + } + Self::maybe_refresh_plugin_caches_for_current_config( &self.config_manager, &self.thread_manager, @@ -861,10 +866,6 @@ impl AccountRequestProcessor { ) .await; - if managed_bedrock_auth { - clear_user_model_provider_if_bedrock(&self.config_manager).await?; - } - // Reflect the current auth method after logout (likely None). Ok(self .auth_manager diff --git a/codex-rs/app-server/tests/suite/v2/account.rs b/codex-rs/app-server/tests/suite/v2/account.rs index c97c0a42ec95..32294fbd89b3 100644 --- a/codex-rs/app-server/tests/suite/v2/account.rs +++ b/codex-rs/app-server/tests/suite/v2/account.rs @@ -166,6 +166,20 @@ fn load_file_auth(codex_home: &Path) -> Result> { )?) } +fn aws_managed_bedrock_config() -> CreateConfigTomlParams { + CreateConfigTomlParams { + model_provider_id: Some("amazon-bedrock".to_string()), + extra_provider_config: Some( + r#"[model_providers.amazon-bedrock.aws] +profile = "codex-bedrock" +region = "us-west-2" +"# + .to_string(), + ), + ..Default::default() + } +} + async fn read_account(mcp: &mut TestAppServer) -> Result { let request_id = mcp .send_get_account_request(GetAccountParams { @@ -1144,9 +1158,9 @@ async fn login_amazon_bedrock_replaces_primary_auth_and_persists_provider() -> R } #[tokio::test] -async fn logout_managed_bedrock_restores_aws_managed_account() -> Result<()> { +async fn logout_managed_bedrock_restores_default_account() -> Result<()> { let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), aws_managed_bedrock_config())?; + create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?; let mut expected_config = read_config_toml(codex_home.path())?; expected_config .as_table_mut() @@ -1174,6 +1188,15 @@ async fn logout_managed_bedrock_restores_aws_managed_account() -> Result<()> { ) .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, + } + ); let request_id = mcp.send_logout_account_request().await?; let response = timeout( @@ -1191,10 +1214,8 @@ async fn logout_managed_bedrock_restores_aws_managed_account() -> Result<()> { assert_eq!( read_account(&mut mcp).await?, GetAccountResponse { - account: Some(Account::AmazonBedrock { - credential_source: AmazonBedrockCredentialSource::AwsManaged, - }), - requires_openai_auth: false, + account: None, + requires_openai_auth: true, } ); Ok(()) @@ -1277,6 +1298,13 @@ async fn logout_managed_bedrock_preserves_changed_provider_without_experimental_ assert_eq!(load_file_auth(codex_home.path())?, None); assert_eq!(read_config_toml(codex_home.path())?, expected_config); assert_account_updated(&mut mcp, /*auth_mode*/ None).await?; + assert_eq!( + read_account(&mut mcp).await?, + GetAccountResponse { + account: None, + requires_openai_auth: false, + } + ); Ok(()) } From 3df78d4b91ec0c2c2aff742151f84744208d0ae5 Mon Sep 17 00:00:00 2001 From: celia-oai Date: Wed, 8 Jul 2026 16:42:26 -0700 Subject: [PATCH 12/12] app-server-test-client: add logout command --- codex-rs/app-server-test-client/README.md | 12 +++++++ codex-rs/app-server-test-client/src/lib.rs | 39 ++++++++++++++++++++++ 2 files changed, 51 insertions(+) diff --git a/codex-rs/app-server-test-client/README.md b/codex-rs/app-server-test-client/README.md index 9f73b1db6f5c..9d559d80b4a2 100644 --- a/codex-rs/app-server-test-client/README.md +++ b/codex-rs/app-server-test-client/README.md @@ -46,6 +46,18 @@ cargo run -p codex-app-server-test-client -- \ 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 logout + +`test-logout` initializes the app-server, sends an `account/logout` request, and waits for the +resulting `account/updated` notification. It uses the active `CODEX_HOME`, so point it at an +isolated directory when testing credential cleanup. + +```bash +cargo run -p codex-app-server-test-client -- \ + --codex-bin ./target/debug/codex \ + test-logout +``` + ## 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 10d66dbc01d5..21becb302fbe 100644 --- a/codex-rs/app-server-test-client/src/lib.rs +++ b/codex-rs/app-server-test-client/src/lib.rs @@ -46,6 +46,7 @@ use codex_app_server_protocol::JSONRPCNotification; use codex_app_server_protocol::JSONRPCRequest; use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::LoginAccountResponse; +use codex_app_server_protocol::LogoutAccountResponse; use codex_app_server_protocol::ModelListParams; use codex_app_server_protocol::ModelListResponse; use codex_app_server_protocol::RequestId; @@ -245,6 +246,8 @@ enum CliCommand { #[arg(long, value_name = "REGION")] region: Option, }, + /// Log out of the current account and wait for the account update. + TestLogout, /// Fetch the current account rate limits from the Codex app-server. GetAccountRateLimits, /// List the available models from the Codex app-server. @@ -449,6 +452,11 @@ pub async fn run() -> Result<()> { }; test_login(&endpoint, &config_overrides, mode).await } + CliCommand::TestLogout => { + ensure_dynamic_tools_unused(&dynamic_tools, "test-logout")?; + let endpoint = resolve_endpoint(codex_bin, url)?; + test_logout(&endpoint, &config_overrides).await + } CliCommand::GetAccountRateLimits => { ensure_dynamic_tools_unused(&dynamic_tools, "get-account-rate-limits")?; let endpoint = resolve_endpoint(codex_bin, url)?; @@ -1253,6 +1261,27 @@ async fn get_account_rate_limits(endpoint: &Endpoint, config_overrides: &[String .await } +async fn test_logout(endpoint: &Endpoint, config_overrides: &[String]) -> Result<()> { + with_client("test-logout", endpoint, config_overrides, |client| { + let initialize = client.initialize()?; + println!("< initialize response: {initialize:?}"); + + let response = client.logout_account()?; + println!("< account/logout response: {response:?}"); + + loop { + let notification = client.next_notification()?; + if let Ok(ServerNotification::AccountUpdated(account_updated)) = + ServerNotification::try_from(notification) + { + println!("< account/updated notification: {account_updated:?}"); + return Ok(()); + } + } + }) + .await +} + async fn model_list(endpoint: &Endpoint, config_overrides: &[String]) -> Result<()> { with_client("model-list", endpoint, config_overrides, |client| { let initialize = client.initialize()?; @@ -1809,6 +1838,16 @@ impl CodexClient { self.send_request(request, request_id, "account/rateLimits/read") } + fn logout_account(&mut self) -> Result { + let request_id = self.request_id(); + let request = ClientRequest::LogoutAccount { + request_id: request_id.clone(), + params: None, + }; + + self.send_request(request, request_id, "account/logout") + } + fn model_list(&mut self, params: ModelListParams) -> Result { let request_id = self.request_id(); let request = ClientRequest::ModelList {