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

Merge remote-tracking branch 'origin/pr/475'

Browse Source 
# Conflicts:
#	apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/refresh/execution.rs
#	apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/refresh/response.rs
#	apps/aether-gateway/src/handlers/admin/provider/oauth/errors.rs
#	apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs
#	apps/aether-gateway/src/state/oauth.rs
#	apps/aether-gateway/src/tests/control/admin/oauth.rs
#	crates/aether-admin/src/provider/quota.rs
This commit is contained in:
fawney19
2026-05-23 21:26:22 +08:00
parent 5643b2c901 75aa3dc0cc
commit 8087a98c9d
11 changed files with 277 additions and 73 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
+50 -40
View File
@@ -81,48 +81,42 @@ pub(super) async fn execute_admin_provider_oauth_refresh(
.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,
let auto_removed = state
.cleanup_provider_catalog_key_if_current(
&provider,
&key_id,
|latest_key| {
should_auto_remove_oauth_invalid_key(
latest_key,
Some(&failure_reason),
false,
now_unix_secs,
)
.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),
));
}
},
)
.await?;
if auto_removed {
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),
));
}
}
tracing::info!(
trace_id = %trace_id,
key_id = %key_id,
provider_id = %provider.id,
provider_type = %provider_type,
event_name = "refresh_failed_retained",
"gateway manual provider oauth refresh failure retained key"
);
}
}
return Ok(RefreshDispatch::Respond(
@@ -171,9 +165,25 @@ pub(super) async fn execute_admin_provider_oauth_refresh(
};
if !helpers::key_is_account_blocked(&key, OAUTH_ACCOUNT_BLOCK_PREFIX) {
let _ = state
let previous_oauth_refresh_issue =
key.oauth_invalid_reason.as_deref().is_some_and(|reason| {
reason.lines().map(str::trim).any(|line| {
line.starts_with("[OAUTH_EXPIRED]") || line.starts_with("[REFRESH_FAILED]")
})
});
let cleared = state
.clear_provider_catalog_key_oauth_invalid_marker(&key_id)
.await?;
if cleared && previous_oauth_refresh_issue {
tracing::info!(
trace_id = %trace_id,
key_id = %key_id,
provider_id = %provider.id,
provider_type = %provider_type,
event_name = "refresh_fixed",
"gateway manual provider oauth refresh cleared oauth invalid marker"
);
}
}
let refreshed_key = state
apps/aether-gateway/src/handlers/admin/provider/oauth/provisioning.rs
-1
View File
@@ -204,7 +204,6 @@ pub(crate) async fn update_existing_provider_oauth_catalog_key(
updated.encrypted_api_key = Some(encrypted_api_key);
updated.encrypted_auth_config = Some(encrypted_auth_config);
updated.api_formats = provider_oauth_catalog_key_api_formats(provider_type, api_formats);
updated.is_active = true;
updated.expires_at_unix_secs = expires_at_unix_secs;
updated.oauth_invalid_at_unix_secs = None;
updated.oauth_invalid_reason = None;
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/codex/mod.rs
+51 -8
View File
@@ -44,6 +44,15 @@ fn merge_codex_quota_metadata(
serde_json::Value::Object(merged)
}
fn codex_oauth_refresh_issue_reason(reason: Option<&str>) -> bool {
reason.is_some_and(|reason| {
reason
.lines()
.map(str::trim)
.any(|line| line.starts_with("[OAUTH_EXPIRED]") || line.starts_with("[REFRESH_FAILED]"))
})
}
pub(crate) async fn refresh_codex_provider_quota_locally(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
@@ -56,8 +65,13 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
let mut success_count = 0usize;
let mut failed_count = 0usize;
let mut auto_removed_count = 0usize;
let mut refresh_fixed_count = 0usize;
let mut refresh_failed_retained_count = 0usize;
let mut auto_removed_hard_banned_count = 0usize;
for key in keys {
let had_oauth_refresh_issue =
codex_oauth_refresh_issue_reason(key.oauth_invalid_reason.as_deref());
let transport = match state
.read_provider_transport_snapshot(&provider.id, &endpoint.id, &key.id)
.await?
@@ -276,13 +290,9 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
}
}
let auto_removed = auto_remove_abnormal_keys
let auto_remove_candidate = auto_remove_abnormal_keys
&& 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;
}
} else if !persist_provider_quota_refresh_state(
let persisted = persist_provider_quota_refresh_state(
state,
&key.id,
metadata_update.as_ref(),
@@ -290,8 +300,8 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
oauth_invalid_reason.clone(),
None,
)
.await?
{
.await?;
if !persisted {
failed_count += 1;
results.push(json!({
"key_id": key.id,
@@ -301,6 +311,29 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
}));
continue;
}
let auto_removed = if auto_remove_candidate {
state
.cleanup_provider_catalog_key_if_current(provider, &key.id, |latest_key| {
should_auto_remove_structured_reason(latest_key.oauth_invalid_reason.as_deref())
})
.await?
} else {
false
};
if auto_removed {
auto_removed_count += 1;
auto_removed_hard_banned_count += 1;
}
let refresh_fixed =
status == "success" && had_oauth_refresh_issue && oauth_invalid_reason.is_none();
if refresh_fixed {
refresh_fixed_count += 1;
}
let refresh_failed_retained =
status != "success" && oauth_invalid_reason.is_some() && !auto_removed;
if refresh_failed_retained {
refresh_failed_retained_count += 1;
}
if status == "success" {
success_count += 1;
@@ -336,6 +369,13 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
}
if auto_removed {
payload.insert("auto_removed".to_string(), json!(true));
payload.insert("auto_removed_hard_banned".to_string(), json!(true));
}
if refresh_fixed {
payload.insert("refresh_fixed".to_string(), json!(true));
}
if refresh_failed_retained {
payload.insert("refresh_failed_retained".to_string(), json!(true));
}
results.push(serde_json::Value::Object(payload));
}
@@ -346,5 +386,8 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
"total": results.len(),
"results": results,
"auto_removed": auto_removed_count,
"refresh_fixed": refresh_fixed_count,
"refresh_failed_retained": refresh_failed_retained_count,
"auto_removed_hard_banned": auto_removed_hard_banned_count,
})))
}
apps/aether-gateway/src/handlers/admin/request/provider/tasks.rs
+64 -12
View File
@@ -1,4 +1,7 @@
use super::*;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use axum::{
body::Body,
http,
@@ -193,8 +196,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 +209,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 +240,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 +263,42 @@ 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 cleanup_provider_catalog_key_if_current<F>(
&self,
provider: &StoredProviderCatalogProvider,
key_id: &str,
should_delete: F,
) -> Result<bool, GatewayError>
where
F: FnOnce(&StoredProviderCatalogKey) -> bool,
{
let key_ids = [key_id.to_string()];
let Some(key) = self
.read_provider_catalog_keys_by_ids(&key_ids)
.await?
.into_iter()
.next()
else {
return Ok(false);
};
if key.provider_id != provider.id || !should_delete(&key) {
return Ok(false);
}
self.clear_admin_provider_pool_cooldown(&provider.id, &key.id)
.await;
self.reset_admin_provider_pool_cost(&provider.id, &key.id)
.await;
let deleted = self.delete_provider_catalog_key(&key.id).await?;
if deleted {
let deleted_key_ids = [key.id.clone()];
self.cleanup_deleted_provider_catalog_refs(&provider.id, &[], &deleted_key_ids)
.await?;
}
Ok(deleted)
}
pub(crate) async fn build_admin_pool_batch_action_response(
apps/aether-gateway/src/maintenance/runtime/account_self_check.rs
+25
View File
@@ -98,6 +98,10 @@ enum AccountSelfCheckOutcome {
status_code: Option<u16>,
message: String,
},
AutoRemoved {
status_code: Option<u16>,
message: String,
},
Failed {
status_code: Option<u16>,
message: String,
@@ -112,6 +116,7 @@ impl AccountSelfCheckOutcome {
match self {
Self::Success { .. } => "success",
Self::Blocked { .. } => "blocked",
Self::AutoRemoved { .. } => "auto_removed",
Self::Failed { .. } => "failed",
Self::Skipped { .. } => "skipped",
}
@@ -121,6 +126,7 @@ impl AccountSelfCheckOutcome {
match self {
Self::Success { status_code, .. }
| Self::Blocked { status_code, .. }
| Self::AutoRemoved { status_code, .. }
| Self::Failed { status_code, .. } => *status_code,
Self::Skipped { .. } => None,
}
@@ -130,6 +136,7 @@ impl AccountSelfCheckOutcome {
match self {
Self::Success { message, .. } => message.as_deref(),
Self::Blocked { message, .. }
| Self::AutoRemoved { message, .. }
| Self::Failed { message, .. }
| Self::Skipped { message, .. } => Some(message.as_str()),
}
@@ -392,6 +399,10 @@ fn quota_payload_result_for_key(key_id: &str, payload: Option<Value>) -> Account
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
let auto_removed = item
.get("auto_removed")
.and_then(Value::as_bool)
.unwrap_or(false);
if status == "success" {
return AccountSelfCheckOutcome::Success {
@@ -399,6 +410,12 @@ fn quota_payload_result_for_key(key_id: &str, payload: Option<Value>) -> Account
message,
};
}
if auto_removed {
return AccountSelfCheckOutcome::AutoRemoved {
status_code,
message: message.unwrap_or_else(|| "已自动删除".to_string()),
};
}
if quota_result_status_is_blocked(&status, status_code, message.as_deref()) {
return AccountSelfCheckOutcome::Blocked {
status_code,
@@ -499,6 +516,11 @@ async fn record_score_probe_result_for_key(
Some(PoolMemberHardState::Banned),
PoolMemberProbeStatus::Failed,
),
AccountSelfCheckOutcome::AutoRemoved { .. } => (
false,
Some(PoolMemberHardState::Banned),
PoolMemberProbeStatus::Failed,
),
AccountSelfCheckOutcome::Failed { .. } => (
false,
Some(PoolMemberHardState::Cooldown),
@@ -565,6 +587,9 @@ fn update_summary_from_outcome(
AccountSelfCheckOutcome::Blocked { .. } => {
summary.blocked = summary.blocked.saturating_add(1);
}
AccountSelfCheckOutcome::AutoRemoved { .. } => {
summary.auto_removed = summary.auto_removed.saturating_add(1);
}
AccountSelfCheckOutcome::Failed { .. } => {
summary.failed = summary.failed.saturating_add(1);
}
apps/aether-gateway/src/maintenance/runtime/pool_quota_probe.rs
+18 -1
View File
@@ -1227,6 +1227,24 @@ 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!(
event_name = "auto_removed_hard_banned",
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,
@@ -1308,7 +1326,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/state/oauth.rs
-1
View File
@@ -1166,7 +1166,6 @@ impl AppState {
latest_key.encrypted_api_key = Some(encrypted_api_key);
latest_key.encrypted_auth_config = encrypted_auth_config;
latest_key.is_active = true;
latest_key.expires_at_unix_secs = entry.expires_at_unix_secs;
let (oauth_invalid_at_unix_secs, oauth_invalid_reason) =
local_oauth_refresh_success_invalid_state(&latest_key);
apps/aether-gateway/src/tests/control/admin/oauth.rs
+5 -6
View File
@@ -6144,10 +6144,10 @@ async fn gateway_auto_removes_manual_oauth_refresh_failure_after_access_token_ex
StatusCode::UNAUTHORIZED,
Json(json!({
"error": {
"message": "Your refresh token has already been used to generate a new access token. Please try signing in again.",
"message": "Could not validate your refresh token. Please try signing in again.",
"type": "invalid_request_error",
"param": serde_json::Value::Null,
"code": "refresh_token_reused"
"code": "refresh_token_expired"
}
})),
)
@@ -6178,19 +6178,18 @@ async fn gateway_auto_removes_manual_oauth_refresh_failure_after_access_token_ex
"openai:responses",
"https://chatgpt.com/backend-api/codex",
);
let mut key = sample_key(
"key-codex-oauth-refresh-expired",
"provider-codex",
"openai:responses",
"stale-codex-access-token",
"expired-codex-access-token",
);
key.auth_type = "oauth".to_string();
key.expires_at_unix_secs = Some(1);
key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"codex","refresh_token":"used-refresh-token","email":"alice@example.com","account_id":"acct-codex-123","plan_type":"plus","expires_at":1}"#,
r#"{"provider_type":"codex","refresh_token":"expired-refresh-token","email":"alice@example.com","account_id":"acct-codex-123","plan_type":"plus","expires_at":1}"#,
)
.expect("auth config ciphertext should build"),
);
@@ -6200,7 +6199,6 @@ async fn gateway_auto_removes_manual_oauth_refresh_failure_after_access_token_ex
vec![endpoint],
vec![key],
));
let (token_url, token_handle) = start_server(token_server).await;
let oauth_refresh =
crate::provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![
@@ -6240,6 +6238,7 @@ async fn gateway_auto_removes_manual_oauth_refresh_failure_after_access_token_ex
.await
.expect("refresh payload should parse");
assert_eq!(refresh_payload["status"], json!("auto_removed"));
assert_eq!(refresh_payload["message"], json!("已自动删除"));
assert_eq!(*token_hits.lock().expect("mutex should lock"), 1);
let keys = provider_catalog_repository
apps/aether-gateway/src/tests/control/admin/pool.rs
+16 -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;
@@ -3481,8 +3489,13 @@ async fn gateway_cleans_up_admin_pool_banned_keys_locally_with_trusted_admin_pri
.list_keys_by_provider_ids(&["provider-openai".to_string()])
.await
.expect("remaining keys should load");
assert_eq!(remaining_keys.len(), 1);
assert_eq!(remaining_keys[0].id, "key-openai-healthy");
assert_eq!(remaining_keys.len(), 2);
assert!(remaining_keys
.iter()
.any(|key| key.id == "key-openai-oauth-expired"));
assert!(remaining_keys
.iter()
.any(|key| key.id == "key-openai-healthy"));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
crates/aether-admin/src/provider/quota.rs
+35 -1
View File
@@ -199,7 +199,7 @@ pub fn quota_refresh_success_invalid_state(
.as_deref()
.map(str::trim)
.unwrap_or_default();
if current_reason.starts_with(OAUTH_REFRESH_FAILED_PREFIX) {
if current_reason.starts_with(OAUTH_ACCOUNT_BLOCK_PREFIX) {
return (
key.oauth_invalid_at_unix_secs,
(!current_reason.is_empty()).then_some(current_reason.to_string()),
@@ -1454,6 +1454,7 @@ mod tests {
parse_chatgpt_web_conversation_init_response, parse_codex_backend_me_response,
parse_codex_wham_usage_response, parse_windsurf_model_configs_response,
parse_windsurf_rate_limit_response, parse_windsurf_user_status_response,
quota_refresh_success_invalid_state, should_auto_remove_structured_reason,
OAUTH_ACCOUNT_BLOCK_PREFIX, OAUTH_EXPIRED_PREFIX, OAUTH_REFRESH_FAILED_PREFIX,
OAUTH_REQUEST_FAILED_PREFIX,
};
@@ -1576,6 +1577,13 @@ mod tests {
);
}
#[test]
fn auto_remove_structured_reason_keeps_oauth_expired_token_invalid() {
assert!(!should_auto_remove_structured_reason(Some(
"[OAUTH_EXPIRED] token invalidated"
)));
}
#[test]
fn auto_remove_refresh_failed_after_access_token_expiry() {
let mut key = StoredProviderCatalogKey::new(
@@ -1726,6 +1734,32 @@ mod tests {
));
}
#[test]
fn quota_refresh_success_clears_refresh_failed_marker() {
let mut key = StoredProviderCatalogKey::new(
"key-1".to_string(),
"provider-1".to_string(),
"key-1".to_string(),
"oauth".to_string(),
None,
true,
)
.expect("key should build");
key.oauth_invalid_reason = Some("[REFRESH_FAILED] Token 续期失败".to_string());
assert_eq!(quota_refresh_success_invalid_state(&key), (None, None));
}
#[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
+13
View File
@@ -644,6 +644,19 @@ mod tests {
assert!(!snapshot.recoverable);
}
#[test]
fn oauth_expired_state_is_not_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(