feat(providers): add credential-fenced OAuth key cleanup

This commit is contained in:
ZheFox
2026-07-28 11:32:11 +08:00
parent 531cf11025
commit f0f87b56a3
17 changed files with 702 additions and 142 deletions
@@ -20,14 +20,17 @@ use self::plan::{
use super::shared::{
build_quota_snapshot_payload, extract_execution_error_message,
oauth_refresh_auto_removed_result, persist_fenced_provider_quota_refresh_state,
persist_provider_quota_refresh_state, quota_key_auto_removed,
quota_refresh_success_invalid_state, ProviderQuotaExecutionOutcome,
persist_provider_quota_refresh_state, provider_auto_remove_banned_keys,
provider_auto_remove_quota_exhausted_keys, quota_key_auto_removed,
quota_refresh_success_invalid_state, should_auto_remove_oauth_invalid_key,
ProviderQuotaExecutionOutcome,
};
use crate::handlers::admin::request::AdminAppState;
use crate::provider_key_auth::provider_key_is_oauth_managed;
use crate::GatewayError;
use aether_contracts::ProxySnapshot;
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogKeyOAuthCredentialCasDelete, ProviderCatalogKeyOAuthCredentialFence,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use axum::http::StatusCode;
@@ -695,20 +698,68 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
}));
continue;
}
// Codex quota responses never auto-delete keys. Without a repository
// conditional delete, any read-then-delete sequence could remove a
// replacement Agent Identity installed while the response was in flight.
let auto_removed_hard_banned = false;
let credential_cas_delete = quota_auth_config_fence.as_ref().map(|auth_config| {
ProviderCatalogKeyOAuthCredentialCasDelete {
key_id: key.id.clone(),
expected_encrypted_auth_config: Some(auth_config.clone()),
expected_credential: ProviderCatalogKeyOAuthCredentialFence {
encrypted_api_key: key.encrypted_api_key.clone(),
auth_type: key.auth_type.clone(),
provider_id: key.provider_id.clone(),
provider_type: provider.provider_type.clone(),
},
}
});
let should_auto_remove_hard_banned =
provider_auto_remove_banned_keys(provider.config.as_ref())
&& should_auto_remove_oauth_invalid_key(
&key,
oauth_invalid_reason.as_deref(),
matches!(status_code, Some(401 | 403)),
now_unix_secs,
);
let auto_removed_hard_banned = if should_auto_remove_hard_banned {
match credential_cas_delete.as_ref() {
Some(delete) => {
state
.compare_and_delete_provider_catalog_key_oauth_credential(delete)
.await?
}
None => false,
}
} else {
false
};
if auto_removed_hard_banned {
auto_removed_count += 1;
auto_removed_hard_banned_count += 1;
}
let auto_removed_quota_exhausted = false;
let auto_removed_quota_exhausted = if !auto_removed_hard_banned
&& status == "quota_exhausted"
&& provider_auto_remove_quota_exhausted_keys(provider.config.as_ref())
{
match credential_cas_delete.as_ref() {
Some(delete) => {
state
.compare_and_delete_provider_catalog_key_oauth_credential(delete)
.await?
}
None => false,
}
} else {
false
};
if auto_removed_quota_exhausted {
auto_removed_count += 1;
status = "quota_exhausted".to_string();
}
let auto_removed = auto_removed_hard_banned || auto_removed_quota_exhausted;
if auto_removed {
let deleted_key_ids = [key.id.clone()];
state
.cleanup_deleted_provider_catalog_refs(&provider.id, false, &[], &deleted_key_ids)
.await?;
}
let refresh_fixed =
status == "success" && had_oauth_refresh_issue && oauth_invalid_reason.is_none();
if refresh_fixed {