i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-11 21:50: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()])
crates/aether-admin/src/provider/quota.rs
+19 -2
View File
@@ -914,8 +914,8 @@ mod tests {
use super::{
codex_build_invalid_state, codex_runtime_invalid_reason,
parse_chatgpt_web_conversation_init_response, parse_codex_wham_usage_response,
OAUTH_ACCOUNT_BLOCK_PREFIX, OAUTH_EXPIRED_PREFIX, OAUTH_REFRESH_FAILED_PREFIX,
OAUTH_REQUEST_FAILED_PREFIX,
should_auto_remove_structured_reason, OAUTH_ACCOUNT_BLOCK_PREFIX, OAUTH_EXPIRED_PREFIX,
OAUTH_REFRESH_FAILED_PREFIX, OAUTH_REQUEST_FAILED_PREFIX,
};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use serde_json::json;
@@ -1010,6 +1010,23 @@ mod tests {
);
}
#[test]
fn auto_remove_structured_reason_removes_oauth_expired_token_invalid() {
assert!(should_auto_remove_structured_reason(Some(
"[OAUTH_EXPIRED] token invalidated"
)));
}
#[test]
fn auto_remove_structured_reason_keeps_request_and_refresh_failures() {
assert!(!should_auto_remove_structured_reason(Some(
"[REQUEST_FAILED] 账号状态检查失败"
)));
assert!(!should_auto_remove_structured_reason(Some(
"[REFRESH_FAILED] Token 续期失败 (401): refresh_token 已失效"
)));
}
#[test]
fn parses_codex_spark_quota_from_additional_rate_limits() {
let parsed = parse_codex_wham_usage_response(
crates/aether-admin/src/provider/status.rs
+14
View File
@@ -42,6 +42,7 @@ const AUTO_REMOVABLE_ACCOUNT_STATE_CODES: &[&str] = &[
"account_disabled",
"workspace_deactivated",
"account_forbidden",
"oauth_token_invalid",
];
#[derive(Debug, Clone, Default, PartialEq, Eq)]
@@ -602,6 +603,19 @@ mod tests {
assert!(!snapshot.recoverable);
}
#[test]
fn oauth_expired_state_is_auto_removed() {
let state = resolve_pool_account_state(
Some("codex"),
None,
Some("[OAUTH_EXPIRED] token invalidated"),
);
assert!(state.blocked);
assert_eq!(state.code.as_deref(), Some("oauth_token_invalid"));
assert!(should_auto_remove_account_state(&state));
}
#[test]
fn account_snapshot_detects_account_block_and_verification() {
let snapshot = resolve_account_status_snapshot(