From c0be3ad57aeeb9450a1ab21785bd8984e2681357 Mon Sep 17 00:00:00 2001 From: celia-oai Date: Mon, 6 Jul 2026 17:25:32 -0700 Subject: [PATCH] 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 | 47 +++-- .../src/request_processors/bedrock_auth.rs | 64 +++++++ codex-rs/app-server/tests/suite/v2/account.rs | 170 ++++++++++++++++++ codex-rs/login/src/auth/bedrock_api_key.rs | 13 ++ .../login/src/auth/bedrock_api_key_tests.rs | 33 ++++ 7 files changed, 312 insertions(+), 18 deletions(-) diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index c530590237aa..9cfa457d3996 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -2053,6 +2053,8 @@ Codex stores the key and region in its configured auth store and writes `model_p { "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 119ea6179061..5b205f07d15d 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -36,7 +36,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 8527722046d9..b1e8f6ccc060 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 chrono::DateTime; @@ -776,7 +778,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; @@ -785,6 +787,27 @@ impl AccountRequestProcessor { } } + if self + .auth_manager + .has_stored_bedrock_api_key_auth() + .map_err(|err| internal_error(format!("failed to read stored auth: {err}")))? + { + let user_model_provider = user_model_provider_state(&self.config_manager).await?; + self.auth_manager + .remove_stored_bedrock_api_key_auth() + .map_err(|err| internal_error(format!("logout failed: {err}")))?; + self.auth_manager + .reload_bedrock_api_key_auth() + .map_err(|err| { + internal_error(format!("failed to reload Amazon Bedrock auth: {err}")) + })?; + clear_user_model_provider_if_bedrock(&self.config_manager, user_model_provider).await?; + return Ok(()); + } + if self.config.model_provider.is_amazon_bedrock() { + return Ok(()); + } + match self.auth_manager.logout_with_revoke().await { Ok(_) => {} Err(err) => { @@ -799,26 +822,16 @@ 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)) + 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().await) + } 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 e387e4ceb922..fb6d82924144 100644 --- a/codex-rs/app-server/tests/suite/v2/account.rs +++ b/codex-rs/app-server/tests/suite/v2/account.rs @@ -1083,6 +1083,10 @@ async fn login_amazon_bedrock_persists_changes_until_restart() -> Result<()> { responses::ev_response_created("resp-after-login"), responses::ev_completed("resp-after-login"), ]), + responses::sse(vec![ + responses::ev_response_created("resp-after-logout"), + responses::ev_completed("resp-after-logout"), + ]), ], ) .await; @@ -1114,6 +1118,11 @@ async fn login_amazon_bedrock_persists_changes_until_restart() -> Result<()> { "model_provider".to_string(), toml::Value::String("amazon-bedrock".to_string()), ); + let mut expected_config_after_logout = expected_config.clone(); + expected_config_after_logout + .as_table_mut() + .expect("config should be a table") + .remove("model_provider"); let request_id = mcp .send_login_account_amazon_bedrock_request(" managed-bedrock-api-key ", " us-west-2 ") .await?; @@ -1169,6 +1178,31 @@ async fn login_amazon_bedrock_persists_changes_until_restart() -> Result<()> { } ); run_turn(&mut mcp, &thread.id).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_after_logout + ); + assert_account_updated(&mut mcp, Some(AuthMode::ApiKey)).await?; + assert_eq!( + read_account(&mut mcp).await?, + GetAccountResponse { + account: Some(Account::ApiKey {}), + requires_openai_auth: true, + } + ); + run_turn(&mut mcp, &thread.id).await?; assert_eq!( responses_mock .requests() @@ -1178,12 +1212,148 @@ async fn login_amazon_bedrock_persists_changes_until_restart() -> Result<()> { vec![ Some("Bearer sk-test-key".to_string()), Some("Bearer sk-test-key".to_string()), + Some("Bearer sk-test-key".to_string()), ] ); 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, /*auth_mode*/ None).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()?; diff --git a/codex-rs/login/src/auth/bedrock_api_key.rs b/codex-rs/login/src/auth/bedrock_api_key.rs index 57c02e058ebd..32286b2ee8f4 100644 --- a/codex-rs/login/src/auth/bedrock_api_key.rs +++ b/codex-rs/login/src/auth/bedrock_api_key.rs @@ -7,6 +7,7 @@ use serde::Serialize; use super::manager::AuthManager; use super::manager::load_auth_dot_json; +use super::manager::logout; use super::manager::save_auth; use super::storage::AuthDotJson; use super::storage::AuthKeyringBackendKind; @@ -115,6 +116,18 @@ impl AuthManager { )?; Ok(self.set_cached_bedrock_api_key_auth(bedrock_api_key)) } + + /// Removes stored managed Bedrock auth without changing cached general Codex auth. + pub fn remove_stored_bedrock_api_key_auth(&self) -> std::io::Result { + if !self.has_stored_bedrock_api_key_auth()? { + return Ok(false); + } + logout( + &self.codex_home, + self.auth_credentials_store_mode, + self.keyring_backend_kind, + ) + } } #[cfg(test)] diff --git a/codex-rs/login/src/auth/bedrock_api_key_tests.rs b/codex-rs/login/src/auth/bedrock_api_key_tests.rs index a00fbe68e2f4..d40e6fca470e 100644 --- a/codex-rs/login/src/auth/bedrock_api_key_tests.rs +++ b/codex-rs/login/src/auth/bedrock_api_key_tests.rs @@ -102,6 +102,39 @@ async fn managed_bedrock_persistence_preserves_cached_openai_auth() -> anyhow::R Ok(()) } +#[tokio::test] +#[serial(codex_auth_env)] +async fn managed_bedrock_removal_preserves_cached_openai_auth() -> anyhow::Result<()> { + let codex_home = tempdir()?; + let storage = FileAuthStorage::new(codex_home.path().to_path_buf()); + storage.save(&api_key_auth())?; + let auth_manager = AuthManager::new( + codex_home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, + /*chatgpt_base_url*/ None, + AuthKeyringBackendKind::default(), + /*auth_route_config*/ None, + ) + .await; + auth_manager.persist_bedrock_api_key_auth("bedrock-api-key-test", "us-east-1")?; + assert!(auth_manager.reload_bedrock_api_key_auth()?); + + assert!(auth_manager.remove_stored_bedrock_api_key_auth()?); + assert!(auth_manager.reload_bedrock_api_key_auth()?); + + assert_eq!(storage.load()?, None); + assert_eq!( + auth_manager + .auth_cached() + .and_then(|auth| auth.api_key().map(str::to_string)), + Some("sk-test-key".to_string()) + ); + assert_eq!(auth_manager.bedrock_api_key_auth_cached(), None); + Ok(()) +} + #[test] fn bedrock_auth_debug_redacts_api_key() { assert_eq!(