i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-04 16:37:46 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/codex/mod.rs
+58 -7
View File
@@ -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 {
apps/aether-gateway/src/handlers/admin/request/provider/catalog.rs
+9
View File
@@ -389,6 +389,15 @@ impl<'a> AdminAppState<'a> {
self.app.delete_provider_catalog_key(key_id).await
}
pub(crate) async fn compare_and_delete_provider_catalog_key_oauth_credential(
&self,
delete: &aether_data_contracts::repository::provider_catalog::ProviderCatalogKeyOAuthCredentialCasDelete,
) -> Result<bool, GatewayError> {
self.app
.compare_and_delete_provider_catalog_key_oauth_credential(delete)
.await
}
pub(crate) async fn create_provider_catalog_provider(
&self,
provider: &aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider,
apps/aether-gateway/src/handlers/admin/request/provider/tasks.rs
+120 -52
View File
@@ -16,7 +16,7 @@ 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 {
fn provider_requires_credential_cas_cleanup(provider: &StoredProviderCatalogProvider) -> bool {
provider.provider_type.trim().eq_ignore_ascii_case("codex")
}
@@ -372,13 +372,6 @@ 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?
@@ -389,27 +382,25 @@ impl<'a> AdminAppState<'a> {
return Ok(0);
}
let deleted_key_ids = banned_keys
.iter()
.map(|key| key.id.clone())
.collect::<Vec<_>>();
let mut deleted_key_ids = Vec::new();
for key in &banned_keys {
self.clear_admin_provider_pool_cooldown(&provider.id, &key.id)
.await;
self.reset_admin_provider_pool_cost(&provider.id, &key.id)
.await;
}
let mut affected = 0usize;
for key_id in &deleted_key_ids {
if self.delete_provider_catalog_key(key_id).await? {
affected += 1;
if self
.delete_provider_catalog_key_for_automatic_cleanup(provider, key)
.await?
{
self.clear_admin_provider_pool_cooldown(&provider.id, &key.id)
.await;
self.reset_admin_provider_pool_cost(&provider.id, &key.id)
.await;
deleted_key_ids.push(key.id.clone());
}
}
self.cleanup_deleted_provider_catalog_refs(&provider.id, false, &[], &deleted_key_ids)
.await?;
if !deleted_key_ids.is_empty() {
self.cleanup_deleted_provider_catalog_refs(&provider.id, false, &[], &deleted_key_ids)
.await?;
}
Ok(affected)
Ok(deleted_key_ids.len())
}
pub(crate) async fn cleanup_quota_exhausted_provider_catalog_keys(
@@ -419,10 +410,6 @@ 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?;
@@ -487,27 +474,48 @@ impl<'a> AdminAppState<'a> {
return Ok(0);
}
let deleted_key_ids = exhausted_keys
.iter()
.map(|key| key.id.clone())
.collect::<Vec<_>>();
let mut deleted_key_ids = Vec::new();
for key in exhausted_keys {
self.clear_admin_provider_pool_cooldown(&provider.id, &key.id)
.await;
self.reset_admin_provider_pool_cost(&provider.id, &key.id)
.await;
}
let mut affected = 0usize;
for key_id in &deleted_key_ids {
if self.delete_provider_catalog_key(key_id).await? {
affected += 1;
if self
.delete_provider_catalog_key_for_automatic_cleanup(provider, key)
.await?
{
self.clear_admin_provider_pool_cooldown(&provider.id, &key.id)
.await;
self.reset_admin_provider_pool_cost(&provider.id, &key.id)
.await;
deleted_key_ids.push(key.id.clone());
}
}
self.cleanup_deleted_provider_catalog_refs(&provider.id, false, &[], &deleted_key_ids)
.await?;
if !deleted_key_ids.is_empty() {
self.cleanup_deleted_provider_catalog_refs(&provider.id, false, &[], &deleted_key_ids)
.await?;
}
Ok(affected)
Ok(deleted_key_ids.len())
}
async fn delete_provider_catalog_key_for_automatic_cleanup(
&self,
provider: &StoredProviderCatalogProvider,
key: &StoredProviderCatalogKey,
) -> Result<bool, GatewayError> {
if !provider_requires_credential_cas_cleanup(provider) {
return self.delete_provider_catalog_key(&key.id).await;
}
self.compare_and_delete_provider_catalog_key_oauth_credential(
&aether_data_contracts::repository::provider_catalog::ProviderCatalogKeyOAuthCredentialCasDelete {
key_id: key.id.clone(),
expected_encrypted_auth_config: key.encrypted_auth_config.clone(),
expected_credential: aether_data_contracts::repository::provider_catalog::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(),
},
},
)
.await
}
pub(crate) async fn cleanup_provider_catalog_key_if_current<F>(
@@ -867,8 +875,12 @@ impl<'a> AdminAppState<'a> {
#[cfg(test)]
mod automatic_cleanup_tests {
use super::provider_skips_automatic_key_cleanup;
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider;
use super::{provider_requires_credential_cas_cleanup, AdminAppState};
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogReadRepository, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use std::sync::Arc;
fn provider(provider_type: &str) -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
@@ -881,9 +893,65 @@ mod automatic_cleanup_tests {
}
#[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")));
fn codex_automatic_cleanup_uses_credential_cas() {
assert!(provider_requires_credential_cas_cleanup(&provider("codex")));
assert!(provider_requires_credential_cas_cleanup(&provider("CoDeX")));
assert!(!provider_requires_credential_cas_cleanup(&provider("kiro")));
}
#[tokio::test]
async fn cleanup_removes_existing_hard_invalid_codex_key() {
let provider = provider("codex");
let mut key = StoredProviderCatalogKey::new(
"key-codex-invalid".to_string(),
provider.id.clone(),
"Invalid Codex OAuth".to_string(),
"oauth".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
None,
"encrypted-api-key".to_string(),
Some("encrypted-auth-config".to_string()),
None,
None,
None,
None,
None,
None,
)
.expect("key transport should build");
key.oauth_invalid_at_unix_secs = Some(1);
key.oauth_invalid_reason = Some(
"[OAUTH_EXPIRED] Your authentication token has been invalidated. Please try signing in again."
.to_string(),
);
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider.clone()],
vec![],
vec![key],
));
let state = crate::AppState::new()
.expect("state should build")
.with_data_state_for_tests(
crate::data::GatewayDataState::with_provider_catalog_repository_for_tests(
repository.clone(),
),
);
let admin_state = AdminAppState::new(&state);
let affected = admin_state
.cleanup_known_banned_provider_catalog_keys(&provider)
.await
.expect("cleanup should succeed");
assert_eq!(affected, 1);
assert!(repository
.list_keys_by_ids(&["key-codex-invalid".to_string()])
.await
.expect("keys should read")
.is_empty());
}
}