mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-08 10:27:46 +08:00
feat(providers): add credential-fenced OAuth key cleanup
This commit is contained in:
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user