Simplify managed Bedrock provider cleanup

This commit is contained in:
celia-oai
2026-07-09 17:07:04 -07:00
parent 9f6bee570e
commit 20d4ae6146
5 changed files with 91 additions and 70 deletions
@@ -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,
@@ -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");
@@ -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(())
@@ -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<String>,
version: Option<String>,
}
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<UserModelProviderState, JSONRPCErrorError> {
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<String>,
) -> 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)
}
@@ -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,