fix(provider): harden Agent Identity OAuth lifecycle

This commit is contained in:
elky
2026-07-23 09:33:00 +08:00
parent e49024d33b
commit 3606290ac8
84 changed files with 8063 additions and 1104 deletions
@@ -69,11 +69,16 @@ impl<'a> AdminAppState<'a> {
)
}
pub(crate) fn masked_catalog_api_key(
pub(crate) fn masked_catalog_api_key_for_provider(
&self,
key: &aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey,
provider_type: &str,
) -> String {
crate::handlers::admin::shared::masked_catalog_api_key(self.app, key)
crate::handlers::admin::shared::masked_catalog_api_key_for_provider(
self.app,
key,
provider_type,
)
}
pub(crate) async fn build_admin_provider_keys_payload(
@@ -227,6 +227,7 @@ impl<'a> AdminAppState<'a> {
.compare_and_update_provider_catalog_key_adaptive_state(
&ProviderCatalogKeyAdaptiveStateUpdate {
key_id: key_id.to_string(),
expected_encrypted_auth_config: None,
expected,
next,
status_snapshot_patch: serde_json::json!({
@@ -257,6 +258,33 @@ impl<'a> AdminAppState<'a> {
) -> Result<
Option<aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey>,
GatewayError,
> {
self.reset_provider_catalog_key_recovery_state_inner(key_id, None)
.await
}
pub(crate) async fn reset_provider_catalog_key_recovery_state_fenced(
&self,
key_id: &str,
expected_encrypted_auth_config: &str,
) -> Result<
Option<aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey>,
GatewayError,
> {
self.reset_provider_catalog_key_recovery_state_inner(
key_id,
Some(expected_encrypted_auth_config),
)
.await
}
async fn reset_provider_catalog_key_recovery_state_inner(
&self,
key_id: &str,
expected_auth_config: Option<&str>,
) -> Result<
Option<aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey>,
GatewayError,
> {
use aether_data_contracts::repository::provider_catalog::ProviderCatalogKeyHealthStateUpdate;
@@ -271,6 +299,11 @@ impl<'a> AdminAppState<'a> {
else {
return Ok(None);
};
if expected_auth_config
.is_some_and(|expected| current.encrypted_auth_config.as_deref() != Some(expected))
{
return Ok(None);
}
if current.health_by_format.as_ref() == Some(&empty)
&& current.circuit_breaker_by_format.as_ref() == Some(&empty)
{
@@ -282,6 +315,7 @@ impl<'a> AdminAppState<'a> {
.compare_and_update_provider_catalog_key_health_state(
&ProviderCatalogKeyHealthStateUpdate {
key_id: key_id.to_string(),
expected_encrypted_auth_config: expected_auth_config.map(ToOwned::to_owned),
expected_health_by_format: current.health_by_format,
expected_circuit_breaker_by_format: current.circuit_breaker_by_format,
health_by_format: Some(empty.clone()),
@@ -299,15 +333,24 @@ impl<'a> AdminAppState<'a> {
"provider key {key_id} health state changed repeatedly while resetting OAuth recovery state"
)));
}
if !self.reset_provider_catalog_key_error_count(key_id).await? {
if expected_auth_config.is_none()
&& !self.reset_provider_catalog_key_error_count(key_id).await?
{
return Ok(None);
}
Ok(self
let current = self
.read_provider_catalog_keys_by_ids(&[key_id.to_string()])
.await?
.into_iter()
.next())
.next();
if current.as_ref().is_some_and(|key| {
expected_auth_config
.is_some_and(|expected| key.encrypted_auth_config.as_deref() != Some(expected))
}) {
return Ok(None);
}
Ok(current)
}
pub(crate) async fn update_provider_catalog_key_status_snapshot(
@@ -90,6 +90,7 @@ impl<'a> AdminAppState<'a> {
provider_id: &str,
provider_type: &str,
pkce_verifier: Option<&str>,
expected_encrypted_auth_config: Option<&str>,
) -> Result<String, GatewayError> {
let nonce = aether_admin::provider::state::generate_provider_oauth_nonce();
let payload = json!({
@@ -98,6 +99,7 @@ impl<'a> AdminAppState<'a> {
"provider_id": provider_id,
"provider_type": provider_type,
"pkce_verifier": pkce_verifier,
"expected_encrypted_auth_config": expected_encrypted_auth_config,
"created_at": aether_admin::provider::state::current_unix_secs(),
});
let key = provider_oauth_state_storage_key(&nonce);
@@ -16,6 +16,10 @@ use serde_json::{json, Map, Value};
use std::collections::{BTreeMap, BTreeSet};
use std::time::{SystemTime, UNIX_EPOCH};
fn provider_skips_automatic_key_cleanup(provider: &StoredProviderCatalogProvider) -> bool {
provider.provider_type.trim().eq_ignore_ascii_case("codex")
}
impl<'a> AdminAppState<'a> {
pub(crate) async fn clear_admin_provider_pool_cooldown(&self, provider_id: &str, key_id: &str) {
crate::handlers::admin::provider::pool::runtime::clear_admin_provider_pool_cooldown(
@@ -368,6 +372,13 @@ impl<'a> AdminAppState<'a> {
) -> Result<usize, GatewayError> {
use aether_admin::provider::pool as admin_provider_pool_pure;
// Codex OAuth credentials can be replaced by a long-lived Agent
// Identity under the same key id. Until deletes support an auth_config
// CAS, automatic cleanup must retain every Codex key.
if provider_skips_automatic_key_cleanup(provider) {
return Ok(0);
}
let banned_keys = self
.list_provider_catalog_keys_by_provider_ids(std::slice::from_ref(&provider.id))
.await?
@@ -408,6 +419,10 @@ impl<'a> AdminAppState<'a> {
) -> Result<usize, GatewayError> {
use aether_admin::provider::pool as admin_provider_pool_pure;
if provider_skips_automatic_key_cleanup(provider) {
return Ok(0);
}
let keys = self
.list_provider_catalog_keys_by_provider_ids(std::slice::from_ref(&provider.id))
.await?;
@@ -849,3 +864,26 @@ impl<'a> AdminAppState<'a> {
.into_response())
}
}
#[cfg(test)]
mod automatic_cleanup_tests {
use super::provider_skips_automatic_key_cleanup;
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider;
fn provider(provider_type: &str) -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
format!("provider-{provider_type}"),
provider_type.to_string(),
None,
provider_type.to_string(),
)
.expect("provider should build")
}
#[test]
fn codex_automatic_cleanup_is_disabled_for_replaceable_agent_credentials() {
assert!(provider_skips_automatic_key_cleanup(&provider("codex")));
assert!(provider_skips_automatic_key_cleanup(&provider("CoDeX")));
assert!(!provider_skips_automatic_key_cleanup(&provider("kiro")));
}
}
@@ -153,6 +153,7 @@ impl<'a> AdminAppState<'a> {
.compare_and_update_provider_catalog_key_adaptive_state(
&ProviderCatalogKeyAdaptiveStateUpdate {
key_id: key.id.clone(),
expected_encrypted_auth_config: None,
expected,
next,
status_snapshot_patch: json!({