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

fix: tighten OAuth auto cleanup signals

Browse Source
This commit is contained in:
fawney19
2026-05-20 10:34:22 +08:00
parent fbda210b84
commit 61bcbe826a
13 changed files with 922 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/refresh/execution.rs
+48 -1
View File
@@ -1,7 +1,10 @@
use super::super::super::errors::{
merge_provider_oauth_refresh_failure_reason, normalize_provider_oauth_refresh_error_message,
};
use super::super::super::quota::shared::persist_provider_quota_refresh_state;
use super::super::super::quota::shared::{
persist_provider_quota_refresh_state, provider_auto_remove_banned_keys,
should_auto_remove_oauth_invalid_key,
};
use super::super::super::runtime::refresh_provider_oauth_account_state_after_update;
use super::helpers::{self, RefreshDispatch, RefreshRequestContext, RefreshSuccessContext};
use super::response;
@@ -76,6 +79,50 @@ pub(super) async fn execute_admin_provider_oauth_refresh(
None,
)
.await?;
if provider_auto_remove_banned_keys(provider.config.as_ref()) {
let now_unix_secs = helpers::unix_now_secs();
let latest_key = state
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&key_id))
.await?
.into_iter()
.next();
if latest_key.as_ref().is_some_and(|latest_key| {
should_auto_remove_oauth_invalid_key(
latest_key,
None,
false,
now_unix_secs,
)
}) {
state
.clear_admin_provider_pool_cooldown(&provider.id, &key_id)
.await;
state
.reset_admin_provider_pool_cost(&provider.id, &key_id)
.await;
if state.delete_provider_catalog_key(&key_id).await? {
let deleted_key_ids = [key_id.clone()];
state
.cleanup_deleted_provider_catalog_refs(
&provider.id,
&[],
&deleted_key_ids,
)
.await?;
tracing::info!(
trace_id = %trace_id,
key_id = %key_id,
provider_id = %provider.id,
provider_type = %provider_type,
event_name = "auto_removed_oauth_refresh_failed",
"gateway manual provider oauth refresh auto-removed unusable key"
);
return Ok(RefreshDispatch::Respond(
response::oauth_refresh_auto_removed_response(&error_reason),
));
}
}
}
}
}
return Ok(RefreshDispatch::Respond(
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/refresh/response.rs
+9
View File
@@ -24,6 +24,15 @@ pub(super) fn oauth_refresh_failed_bad_request_response(
)
}
pub(super) fn oauth_refresh_auto_removed_response(error_reason: impl AsRef<str>) -> Response<Body> {
Json(json!({
"status": "auto_removed",
"message": "已自动删除",
"detail": format!("Token 刷新失败且 OAuth 凭证已不可用，已自动删除：{}", error_reason.as_ref()),
}))
.into_response()
}
pub(super) fn oauth_refresh_failed_service_unavailable_response(
error_reason: impl Into<String>,
) -> Response<Body> {
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/antigravity.rs
+11 -4
View File
@@ -1,7 +1,8 @@
use super::shared::{
build_provider_quota_execution_plan, build_quota_snapshot_payload, coerce_json_f64,
coerce_json_string, default_provider_quota_execution_timeouts, execute_provider_quota_plan,
extract_execution_error_message, persist_provider_quota_refresh_state,
extract_execution_error_message, oauth_refresh_auto_removed_result,
persist_provider_quota_refresh_state, quota_key_auto_removed,
quota_refresh_success_invalid_state, ProviderQuotaExecutionOutcome,
};
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
@@ -65,6 +66,7 @@ pub(crate) async fn refresh_antigravity_provider_quota_locally(
let mut results = Vec::new();
let mut success_count = 0usize;
let mut failed_count = 0usize;
let mut auto_removed_count = 0usize;
for key in keys {
let transport = match state
@@ -87,6 +89,11 @@ pub(crate) async fn refresh_antigravity_provider_quota_locally(
let authorization = match state.resolve_local_oauth_header_auth(&transport).await? {
Some(auth) => auth,
_ => {
if quota_key_auto_removed(state, &key.id).await? {
auto_removed_count += 1;
results.push(oauth_refresh_auto_removed_result(&key));
continue;
}
failed_count += 1;
results.push(json!({
"key_id": key.id,
@@ -247,9 +254,9 @@ pub(crate) async fn refresh_antigravity_provider_quota_locally(
Ok(Some(json!({
"success": success_count,
"failed": failed_count,
"total": success_count + failed_count,
"total": results.len(),
"results": results,
"message": format!("已处理 {} 个 Key", success_count + failed_count),
"auto_removed": 0,
"message": format!("已处理 {} 个 Key", results.len()),
"auto_removed": auto_removed_count,
})))
}
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/chatgpt_web.rs
+11 -5
View File
@@ -1,8 +1,8 @@
use super::shared::{
build_quota_snapshot_payload, default_provider_quota_execution_timeouts,
execute_provider_quota_plan, extract_execution_error_message,
persist_provider_quota_refresh_state, quota_refresh_success_invalid_state,
ProviderQuotaExecutionOutcome,
oauth_refresh_auto_removed_result, persist_provider_quota_refresh_state,
quota_key_auto_removed, quota_refresh_success_invalid_state, ProviderQuotaExecutionOutcome,
};
use crate::handlers::admin::provider::shared::payloads::{
OAUTH_ACCOUNT_BLOCK_PREFIX, OAUTH_EXPIRED_PREFIX,
@@ -113,6 +113,7 @@ pub(crate) async fn refresh_chatgpt_web_provider_quota_locally(
let mut results = Vec::new();
let mut success_count = 0usize;
let mut failed_count = 0usize;
let mut auto_removed_count = 0usize;
for key in keys {
let transport = match state
@@ -135,6 +136,11 @@ pub(crate) async fn refresh_chatgpt_web_provider_quota_locally(
let authorization = match resolve_chatgpt_web_quota_auth(state, &transport).await? {
Some(auth) => auth,
None => {
if quota_key_auto_removed(state, &key.id).await? {
auto_removed_count += 1;
results.push(oauth_refresh_auto_removed_result(&key));
continue;
}
failed_count += 1;
results.push(json!({
"key_id": key.id,
@@ -291,9 +297,9 @@ pub(crate) async fn refresh_chatgpt_web_provider_quota_locally(
Ok(Some(json!({
"success": success_count,
"failed": failed_count,
"total": success_count + failed_count,
"total": results.len(),
"results": results,
"message": format!("已处理 {} 个 Key", success_count + failed_count),
"auto_removed": 0,
"message": format!("已处理 {} 个 Key", results.len()),
"auto_removed": auto_removed_count,
})))
}
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/codex/mod.rs
+25 -24
View File
@@ -14,9 +14,9 @@ use self::parse::{
use self::plan::{build_codex_quota_request_spec, execute_codex_quota_plan};
use super::shared::{
build_quota_snapshot_payload, extract_execution_error_message,
persist_provider_quota_refresh_state, provider_auto_remove_banned_keys,
quota_refresh_success_invalid_state, should_auto_remove_oauth_invalid_key,
ProviderQuotaExecutionOutcome,
oauth_refresh_auto_removed_result, persist_provider_quota_refresh_state,
provider_auto_remove_banned_keys, quota_key_auto_removed, quota_refresh_success_invalid_state,
should_auto_remove_structured_reason, ProviderQuotaExecutionOutcome,
};
use crate::handlers::admin::request::AdminAppState;
use crate::provider_key_auth::provider_key_is_oauth_managed;
@@ -75,12 +75,27 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
}
};
let resolved_oauth_auth =
if provider_key_is_oauth_managed(&key, provider.provider_type.as_str()) {
state.resolve_local_oauth_header_auth(&transport).await?
} else {
None
};
let is_oauth_managed = provider_key_is_oauth_managed(&key, provider.provider_type.as_str());
let resolved_oauth_auth = if is_oauth_managed {
state.resolve_local_oauth_header_auth(&transport).await?
} else {
None
};
if is_oauth_managed && quota_key_auto_removed(state, &key.id).await? {
auto_removed_count += 1;
results.push(oauth_refresh_auto_removed_result(&key));
continue;
}
if is_oauth_managed && resolved_oauth_auth.is_none() {
failed_count += 1;
results.push(json!({
"key_id": key.id,
"key_name": key.name,
"status": "error",
"message": "缺少 Codex OAuth 认证信息，请先重新授权/刷新 Token",
}));
continue;
}
let request_spec = match build_codex_quota_request_spec(&transport, resolved_oauth_auth) {
Ok(request_spec) => request_spec,
@@ -261,22 +276,8 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
}
}
let auto_remove_key = if auto_remove_abnormal_keys {
state
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&key.id))
.await?
.into_iter()
.next()
.unwrap_or_else(|| key.clone())
} else {
key.clone()
};
let auto_removed = auto_remove_abnormal_keys
&& should_auto_remove_oauth_invalid_key(
&auto_remove_key,
oauth_invalid_reason.as_deref(),
now_unix_secs,
);
&& should_auto_remove_structured_reason(oauth_invalid_reason.as_deref());
if auto_removed {
if state.delete_provider_catalog_key(&key.id).await? {
auto_removed_count += 1;
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/kiro/mod.rs
+13 -4
View File
@@ -5,7 +5,8 @@ use self::parse::parse_kiro_usage_response;
use self::plan::execute_kiro_quota_plan;
use super::shared::{
build_quota_snapshot_payload, extract_execution_error_message,
persist_provider_quota_refresh_state, quota_refresh_success_invalid_state,
oauth_refresh_auto_removed_result, persist_provider_quota_refresh_state,
persist_quota_oauth_refresh_failure_state, quota_refresh_success_invalid_state,
ProviderQuotaExecutionOutcome,
};
use crate::handlers::admin::request::{AdminAppState, AdminLocalOAuthRefreshError};
@@ -95,6 +96,7 @@ pub(crate) async fn refresh_kiro_provider_quota_locally(
let mut results = Vec::new();
let mut success_count = 0usize;
let mut failed_count = 0usize;
let mut auto_removed_count = 0usize;
for key in keys {
let transport = match state
@@ -145,6 +147,13 @@ pub(crate) async fn refresh_kiro_provider_quota_locally(
}
},
Err(err) => {
if persist_quota_oauth_refresh_failure_state(state, &transport, &err).await?
|| super::shared::quota_key_auto_removed(state, &key.id).await?
{
auto_removed_count += 1;
results.push(oauth_refresh_auto_removed_result(&key));
continue;
}
failed_count += 1;
let mut payload = serde_json::Map::new();
payload.insert("key_id".to_string(), json!(key.id));
@@ -328,10 +337,10 @@ pub(crate) async fn refresh_kiro_provider_quota_locally(
Ok(Some(json!({
"success": success_count,
"failed": failed_count,
"total": success_count + failed_count,
"total": results.len(),
"results": results,
"message": format!("已处理 {} 个 Key", success_count + failed_count),
"auto_removed": 0,
"message": format!("已处理 {} 个 Key", results.len()),
"auto_removed": auto_removed_count,
})))
}
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs
+58 -3
View File
@@ -1,7 +1,9 @@
use crate::handlers::admin::provider::shared::payloads::{
OAUTH_ACCOUNT_BLOCK_PREFIX, OAUTH_REFRESH_FAILED_PREFIX,
};
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
use crate::handlers::admin::request::{
AdminAppState, AdminGatewayProviderTransportSnapshot, AdminLocalOAuthRefreshError,
};
use crate::handlers::shared::{
sync_provider_key_oauth_status_snapshot, sync_provider_key_quota_status_snapshot,
};
@@ -44,22 +46,75 @@ pub(super) fn default_provider_quota_execution_timeouts(
}
}
pub(super) fn provider_auto_remove_banned_keys(config: Option<&serde_json::Value>) -> bool {
pub(crate) fn provider_auto_remove_banned_keys(config: Option<&serde_json::Value>) -> bool {
admin_provider_quota_pure::provider_auto_remove_banned_keys(config)
}
pub(super) fn should_auto_remove_oauth_invalid_key(
pub(super) fn should_auto_remove_structured_reason(reason: Option<&str>) -> bool {
admin_provider_quota_pure::should_auto_remove_structured_reason(reason)
}
pub(crate) fn should_auto_remove_oauth_invalid_key(
key: &StoredProviderCatalogKey,
candidate_reason: Option<&str>,
access_token_invalid_proven: bool,
now_unix_secs: u64,
) -> bool {
admin_provider_quota_pure::should_auto_remove_oauth_invalid_key(
key,
candidate_reason,
access_token_invalid_proven,
now_unix_secs,
)
}
pub(crate) async fn persist_quota_oauth_refresh_failure_state(
state: &AdminAppState<'_>,
transport: &AdminGatewayProviderTransportSnapshot,
err: &AdminLocalOAuthRefreshError,
) -> Result<bool, GatewayError> {
let AdminLocalOAuthRefreshError::HttpStatus {
status_code,
body_excerpt,
..
} = err
else {
return Ok(false);
};
if !matches!(*status_code, 400 | 401 | 403) {
return Ok(false);
}
state
.app()
.persist_local_oauth_refresh_failure_state(transport, *status_code, body_excerpt, false)
.await
}
pub(crate) async fn quota_key_auto_removed(
state: &AdminAppState<'_>,
key_id: &str,
) -> Result<bool, GatewayError> {
if key_id.trim().is_empty() {
return Ok(false);
}
Ok(state
.read_provider_catalog_keys_by_ids(&[key_id.to_string()])
.await?
.is_empty())
}
pub(crate) fn oauth_refresh_auto_removed_result(
key: &StoredProviderCatalogKey,
) -> serde_json::Value {
serde_json::json!({
"key_id": key.id,
"key_name": key.name,
"status": "auto_removed",
"message": "OAuth refresh 失败且凭证已不可用，已自动删除",
"auto_removed": true,
})
}
pub(crate) fn normalize_string_id_list(values: Option<Vec<String>>) -> Option<Vec<String>> {
admin_provider_quota_pure::normalize_string_id_list(values)
}