From 20d4ae6146f64ccf9caf91bf0909322bb3c9f527 Mon Sep 17 00:00:00 2001 From: celia-oai Date: Wed, 8 Jul 2026 14:50:25 -0700 Subject: [PATCH] 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 | 74 ++++--------------- codex-rs/app-server/tests/suite/v2/account.rs | 6 +- 5 files changed, 91 insertions(+), 70 deletions(-) diff --git a/codex-rs/app-server/src/config_manager_service.rs b/codex-rs/app-server/src/config_manager_service.rs index 94980bf8e0..5640c9cb74 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 1ec79a749b..3407c1b423 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 caabb0d7c9..25ca80b3b8 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; @@ -846,11 +845,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(()); } @@ -869,8 +863,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 fb41fb95bd..9deef08557 100644 --- a/codex-rs/app-server/src/request_processors/bedrock_auth.rs +++ b/codex-rs/app-server/src/request_processors/bedrock_auth.rs @@ -10,11 +10,6 @@ use codex_config::ConfigLayerSource; use codex_config::format_config_layer_source; 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> { @@ -50,68 +45,27 @@ pub(super) async fn set_user_model_provider_to_bedrock( ))); } - 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 layers = config_manager - .load_config_layers(/*cwd*/ None) - .await - .map_err(|err| internal_error(format!("failed to load configuration layers: {err}")))?; - let Some(layer) = layers.get_active_user_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.clone()), - }) -} - -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 4aff323498..20826afa1f 100644 --- a/codex-rs/app-server/tests/suite/v2/account.rs +++ b/codex-rs/app-server/tests/suite/v2/account.rs @@ -1360,7 +1360,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", @@ -1368,7 +1368,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 @@ -1386,6 +1385,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,