i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-12 22:20:19 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: auto-remove invalid oauth pool keys

Browse Source
This commit is contained in:
beilo
2026-05-16 15:08:16 +08:00
parent c916e76bd2
commit 2ca4b486ec
5 changed files with 88 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/handlers/admin/request/provider/tasks.rs
+27 -12
View File
@@ -1,4 +1,5 @@
use super::*;
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider;
use axum::{
body::Body,
http,
@@ -193,8 +194,6 @@ impl<'a> AdminAppState<'a> {
&self,
provider_id: &str,
) -> Result<Response<Body>, GatewayError> {
use aether_admin::provider::pool as admin_provider_pool_pure;
let Some(provider) = self
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&provider_id.to_string()))
.await?
@@ -208,6 +207,30 @@ impl<'a> AdminAppState<'a> {
.into_response());
};
let affected = self
.cleanup_known_banned_provider_catalog_keys(&provider)
.await?;
if affected == 0 {
return Ok(Json(
aether_admin::provider::pool::build_admin_pool_cleanup_empty_payload(
"未发现可清理的异常账号",
),
)
.into_response());
}
Ok(Json(
aether_admin::provider::pool::build_admin_pool_cleanup_result_payload(affected),
)
.into_response())
}
pub(crate) async fn cleanup_known_banned_provider_catalog_keys(
&self,
provider: &StoredProviderCatalogProvider,
) -> Result<usize, GatewayError> {
use aether_admin::provider::pool as admin_provider_pool_pure;
let banned_keys = self
.list_provider_catalog_keys_by_provider_ids(std::slice::from_ref(&provider.id))
.await?
@@ -215,12 +238,7 @@ impl<'a> AdminAppState<'a> {
.filter(admin_provider_pool_pure::admin_pool_key_is_known_banned)
.collect::<Vec<_>>();
if banned_keys.is_empty() {
return Ok(Json(
admin_provider_pool_pure::build_admin_pool_cleanup_empty_payload(
"未发现可清理的异常账号",
),
)
.into_response());
return Ok(0);
}
let deleted_key_ids = banned_keys
@@ -243,10 +261,7 @@ impl<'a> AdminAppState<'a> {
self.cleanup_deleted_provider_catalog_refs(&provider.id, &[], &deleted_key_ids)
.await?;
Ok(
Json(admin_provider_pool_pure::build_admin_pool_cleanup_result_payload(affected))
.into_response(),
)
Ok(affected)
}
pub(crate) async fn build_admin_pool_batch_action_response(
apps/aether-gateway/src/maintenance/runtime/pool_quota_probe.rs
+17 -1
View File
@@ -1202,6 +1202,23 @@ async fn perform_pool_quota_probe_for_provider(
now_ts: u64,
) -> Result<PoolQuotaProbeRunSummary, GatewayError> {
let mut summary = PoolQuotaProbeRunSummary::empty();
let provider_short_id = provider.id.chars().take(8).collect::<String>();
if aether_admin::provider::quota::provider_auto_remove_banned_keys(provider.config.as_ref()) {
let auto_removed = admin_state
.cleanup_known_banned_provider_catalog_keys(provider)
.await?;
if auto_removed > 0 {
summary.auto_removed += auto_removed;
info!(
provider_id = %provider_short_id,
provider_type,
auto_removed,
"gateway pool quota probe auto-cleaned known abnormal provider keys"
);
}
}
let Some(endpoint) = endpoint_for_probe_with_reconcile(
state,
admin_state,
@@ -1283,7 +1300,6 @@ async fn perform_pool_quota_probe_for_provider(
record_score_probe_in_progress_for_key(state, &provider.id, key_id, now_ts).await;
}
let provider_short_id = provider.id.chars().take(8).collect::<String>();
let probe_concurrency = pool_config.probe_concurrency.clamp(1, 64) as usize;
let probe_concurrency = probe_concurrency.min(config.global_concurrency).max(1);
let probe_results = stream::iter(keys.into_iter().map(|key| {
apps/aether-gateway/src/tests/control/admin/pool.rs
+11 -3
View File
@@ -3433,6 +3433,14 @@ async fn gateway_cleans_up_admin_pool_banned_keys_locally_with_trusted_admin_pri
);
banned_key.name = "banned".to_string();
banned_key.oauth_invalid_reason = Some("account_banned".to_string());
let mut oauth_expired_key = sample_key(
"key-openai-oauth-expired",
"provider-openai",
"openai:chat",
"sk-oauth-expired",
);
oauth_expired_key.name = "oauth-expired".to_string();
oauth_expired_key.oauth_invalid_reason = Some("[OAUTH_EXPIRED] token invalidated".to_string());
let mut healthy_key = sample_key(
"key-openai-healthy",
"provider-openai",
@@ -3444,7 +3452,7 @@ async fn gateway_cleans_up_admin_pool_banned_keys_locally_with_trusted_admin_pri
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
Vec::new(),
vec![banned_key, healthy_key],
vec![banned_key, oauth_expired_key, healthy_key],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
@@ -3474,8 +3482,8 @@ async fn gateway_cleans_up_admin_pool_banned_keys_locally_with_trusted_admin_pri
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["affected"], 1);
assert_eq!(payload["message"], "已清理 1 个异常账号");
assert_eq!(payload["affected"], 2);
assert_eq!(payload["message"], "已清理 2 个异常账号");
let remaining_keys = provider_catalog_repository
.list_keys_by_provider_ids(&["provider-openai".to_string()])