Refine OAuth auto cleanup conditions

(cherry picked from commit 78826eae47ab18ace6931dd190a41070416b5e10)
2026-09-01 17:00:21 +08:00 · 2026-05-16 11:07:07 +08:00
parent f59cf1090d
commit f72ab383c9
8 changed files with 322 additions and 16 deletions
--- a/apps/aether-gateway/src/handlers/admin/provider/oauth/errors.rs
+++ b/apps/aether-gateway/src/handlers/admin/provider/oauth/errors.rs
@@ -1,5 +1,5 @@
 use crate::handlers::admin::provider::shared::payloads::{
-    OAUTH_ACCOUNT_BLOCK_PREFIX, OAUTH_EXPIRED_PREFIX,
+    OAUTH_ACCOUNT_BLOCK_PREFIX, OAUTH_EXPIRED_PREFIX, OAUTH_REFRESH_FAILED_PREFIX,
 };
 use axum::{
    body::Body,
@@ -147,6 +147,14 @@ pub(crate) fn merge_provider_oauth_refresh_failure_reason(
        return Some(refresh_reason.to_string());
    }
    if current_reason.starts_with(OAUTH_EXPIRED_PREFIX) {
+        if refresh_reason.starts_with(OAUTH_REFRESH_FAILED_PREFIX)
+            && !current_reason
+                .lines()
+                .map(str::trim)
+                .any(|line| line.starts_with(OAUTH_REFRESH_FAILED_PREFIX))
+        {
+            return Some(format!("{current_reason}\n{refresh_reason}"));
+        }
        return Some(current_reason.to_string());
    }
    if oauth_invalid_reason_is_account_level_block(Some(current_reason)) {
@@ -190,13 +198,16 @@ mod tests {
    }

    #[test]
-    fn refresh_failure_does_not_replace_access_token_expired_marker() {
+    fn refresh_failure_is_appended_to_access_token_expired_marker() {
        assert_eq!(
            merge_provider_oauth_refresh_failure_reason(
                Some("[OAUTH_EXPIRED] access token invalid"),
                "[REFRESH_FAILED] Token 续期失败 (401): refresh_token 无效",
            ),
-            Some("[OAUTH_EXPIRED] access token invalid".to_string()),
+            Some(
+                "[OAUTH_EXPIRED] access token invalid\n[REFRESH_FAILED] Token 续期失败 (401): refresh_token 无效"
+                    .to_string()
+            ),
        );
    }
 }
--- a/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/codex/mod.rs
+++ b/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/codex/mod.rs
@@ -15,7 +15,7 @@ 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_structured_reason,
+    quota_refresh_success_invalid_state, should_auto_remove_oauth_invalid_key,
    ProviderQuotaExecutionOutcome,
 };
 use crate::handlers::admin::request::AdminAppState;
@@ -261,8 +261,22 @@ 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_structured_reason(oauth_invalid_reason.as_deref());
+            && should_auto_remove_oauth_invalid_key(
+                &auto_remove_key,
+                oauth_invalid_reason.as_deref(),
+                now_unix_secs,
+            );
        if auto_removed {
            if state.delete_provider_catalog_key(&key.id).await? {
                auto_removed_count += 1;
--- a/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs
+++ b/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs
@@ -48,8 +48,16 @@ pub(super) fn provider_auto_remove_banned_keys(config: Option<&serde_json::Value
    admin_provider_quota_pure::provider_auto_remove_banned_keys(config)
 }

-pub(super) fn should_auto_remove_structured_reason(reason: Option<&str>) -> bool {
-    admin_provider_quota_pure::should_auto_remove_structured_reason(reason)
+pub(super) fn should_auto_remove_oauth_invalid_key(
+    key: &StoredProviderCatalogKey,
+    candidate_reason: Option<&str>,
+    now_unix_secs: u64,
+) -> bool {
+    admin_provider_quota_pure::should_auto_remove_oauth_invalid_key(
+        key,
+        candidate_reason,
+        now_unix_secs,
+    )
 }

 pub(crate) fn normalize_string_id_list(values: Option<Vec<String>>) -> Option<Vec<String>> {
--- a/apps/aether-gateway/src/state/oauth.rs
+++ b/apps/aether-gateway/src/state/oauth.rs
@@ -290,6 +290,14 @@ fn merge_local_oauth_refresh_failure_reason(
        return Some(refresh_reason.to_string());
    }
    if current_reason.starts_with(OAUTH_EXPIRED_PREFIX) {
+        if refresh_reason.starts_with(OAUTH_REFRESH_FAILED_PREFIX)
+            && !current_reason
+                .lines()
+                .map(str::trim)
+                .any(|line| line.starts_with(OAUTH_REFRESH_FAILED_PREFIX))
+        {
+            return Some(format!("{current_reason}\n{refresh_reason}"));
+        }
        return Some(current_reason.to_string());
    }
    if oauth_invalid_reason_is_account_block(Some(current_reason)) {
@@ -1717,13 +1725,16 @@ mod tests {
    }

    #[test]
-    fn local_refresh_failure_does_not_replace_access_token_expired_marker() {
+    fn local_refresh_failure_is_appended_to_access_token_expired_marker() {
        assert_eq!(
            super::merge_local_oauth_refresh_failure_reason(
                Some("[OAUTH_EXPIRED] access token invalid"),
                "[REFRESH_FAILED] Token 续期失败 (401): refresh_token 无效",
            ),
-            Some("[OAUTH_EXPIRED] access token invalid".to_string()),
+            Some(
+                "[OAUTH_EXPIRED] access token invalid\n[REFRESH_FAILED] Token 续期失败 (401): refresh_token 无效"
+                    .to_string()
+            ),
        );
    }

--- a/apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs
+++ b/apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs
@@ -270,6 +270,126 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_codex_with_trusted_a
    upstream_handle.abort();
 }

+#[tokio::test]
+async fn gateway_auto_removes_codex_key_when_refresh_and_access_tokens_are_invalid() {
+    let upstream = Router::new().route(
+        "/api/admin/endpoints/providers/provider-codex/refresh-quota",
+        any(move |_request: Request| async move {
+            (StatusCode::OK, Body::from("unexpected upstream hit"))
+        }),
+    );
+
+    let execution_runtime = Router::new().route(
+        "/v1/execute/sync",
+        any(move |request: Request| async move {
+            let plan: aether_contracts::ExecutionPlan = serde_json::from_slice(
+                &to_bytes(request.into_body(), usize::MAX)
+                    .await
+                    .expect("body should read"),
+            )
+            .expect("plan should parse");
+            let result = aether_contracts::ExecutionResult {
+                request_id: plan.request_id,
+                candidate_id: None,
+                status_code: 401,
+                headers: BTreeMap::new(),
+                body: Some(aether_contracts::ResponseBody {
+                    json_body: Some(json!({
+                        "error": {
+                            "message": "session expired"
+                        }
+                    })),
+                    body_bytes_b64: None,
+                }),
+                telemetry: None,
+                error: None,
+            };
+            (StatusCode::OK, Json(result))
+        }),
+    );
+
+    let mut provider = StoredProviderCatalogProvider::new(
+        "provider-codex".to_string(),
+        "codex".to_string(),
+        Some("https://example.com".to_string()),
+        "codex".to_string(),
+    )
+    .expect("provider should build");
+    provider.config = Some(json!({
+        "pool_advanced": {
+            "auto_remove_banned_keys": true
+        }
+    }));
+
+    let mut key = sample_key(
+        "key-codex-expired",
+        "provider-codex",
+        "openai:responses",
+        "stale-access-token",
+    );
+    key.auth_type = "oauth".to_string();
+    key.expires_at_unix_secs = Some(1);
+    key.oauth_invalid_at_unix_secs = Some(1);
+    key.oauth_invalid_reason = Some(
+        "[REFRESH_FAILED] Token 续期失败 (401): refresh_token 无效、已过期或已撤销，请重新登录授权"
+            .to_string(),
+    );
+
+    let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
+        vec![provider],
+        vec![sample_endpoint(
+            "endpoint-codex-cli",
+            "provider-codex",
+            "openai:responses",
+            "https://chatgpt.com/backend-api",
+        )],
+        vec![key],
+    ));
+
+    let (_upstream_url, upstream_handle) = start_server(upstream).await;
+    let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
+    let gateway = build_router_with_state(
+        build_state_with_execution_runtime_override(execution_runtime_url.clone())
+            .with_data_state_for_tests(
+                GatewayDataState::with_provider_catalog_repository_for_tests(
+                    provider_catalog_repository.clone(),
+                )
+                .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
+            ),
+    );
+    let (gateway_url, gateway_handle) = start_server(gateway).await;
+
+    let response = reqwest::Client::new()
+        .post(format!(
+            "{gateway_url}/api/admin/endpoints/providers/provider-codex/refresh-quota"
+        ))
+        .header(GATEWAY_HEADER, "rust-phase3b")
+        .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
+        .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
+        .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
+        .send()
+        .await
+        .expect("request should succeed");
+
+    assert_eq!(response.status(), StatusCode::OK);
+    let payload: serde_json::Value = response.json().await.expect("json body should parse");
+    assert_eq!(payload["success"], 0);
+    assert_eq!(payload["failed"], 1);
+    assert_eq!(payload["auto_removed"], 1);
+    assert_eq!(payload["results"][0]["status"], "auth_invalid");
+    assert_eq!(payload["results"][0]["auto_removed"], true);
+
+    let reloaded = provider_catalog_repository
+        .list_keys_by_ids(&["key-codex-expired".to_string()])
+        .await
+        .expect("keys should read");
+    assert!(reloaded.is_empty());
+
+    gateway_handle.abort();
+    execution_runtime_handle.abort();
+    upstream_handle.abort();
+}
+
 #[tokio::test]
 async fn gateway_refreshes_admin_provider_quota_locally_for_requested_codex_keys_only() {
    let upstream_hits = Arc::new(Mutex::new(0usize));
--- a/crates/aether-admin/src/provider/quota.rs
+++ b/crates/aether-admin/src/provider/quota.rs
@@ -26,6 +26,57 @@ pub fn should_auto_remove_structured_reason(reason: Option<&str>) -> bool {
    ))
 }

+fn oauth_reason_has_tag(reason: Option<&str>, tag: &str) -> bool {
+    reason
+        .map(str::trim)
+        .filter(|value| !value.is_empty())
+        .is_some_and(|reason| {
+            reason
+                .lines()
+                .map(str::trim)
+                .any(|line| line.starts_with(tag))
+        })
+}
+
+fn oauth_access_token_expired(key: &StoredProviderCatalogKey, now_unix_secs: u64) -> bool {
+    let now_unix_secs = if now_unix_secs == 0 {
+        std::time::SystemTime::now()
+            .duration_since(std::time::UNIX_EPOCH)
+            .ok()
+            .map(|duration| duration.as_secs())
+            .unwrap_or(0)
+    } else {
+        now_unix_secs
+    };
+    key.expires_at_unix_secs
+        .is_none_or(|expires_at| expires_at == 0 || expires_at <= now_unix_secs)
+}
+
+pub fn should_auto_remove_oauth_invalid_key(
+    key: &StoredProviderCatalogKey,
+    candidate_reason: Option<&str>,
+    now_unix_secs: u64,
+) -> bool {
+    if should_auto_remove_structured_reason(candidate_reason)
+        || should_auto_remove_structured_reason(key.oauth_invalid_reason.as_deref())
+    {
+        return true;
+    }
+
+    let refresh_token_failed = oauth_reason_has_tag(candidate_reason, OAUTH_REFRESH_FAILED_PREFIX)
+        || oauth_reason_has_tag(
+            key.oauth_invalid_reason.as_deref(),
+            OAUTH_REFRESH_FAILED_PREFIX,
+        );
+    if !refresh_token_failed {
+        return false;
+    }
+
+    oauth_reason_has_tag(candidate_reason, OAUTH_EXPIRED_PREFIX)
+        || oauth_reason_has_tag(key.oauth_invalid_reason.as_deref(), OAUTH_EXPIRED_PREFIX)
+        || oauth_access_token_expired(key, now_unix_secs)
+}
+
 pub fn normalize_string_id_list(values: Option<Vec<String>>) -> Option<Vec<String>> {
    let mut out = Vec::new();
    let mut seen = std::collections::BTreeSet::new();
@@ -435,8 +486,19 @@ fn codex_merge_invalid_reason(current: &str, candidate_reason: &str) -> String {
        return current.to_string();
    }
    if current.starts_with(OAUTH_EXPIRED_PREFIX)
-        && (candidate_reason.starts_with(OAUTH_REQUEST_FAILED_PREFIX)
-            || candidate_reason.starts_with(OAUTH_REFRESH_FAILED_PREFIX))
+        && candidate_reason.starts_with(OAUTH_REFRESH_FAILED_PREFIX)
+    {
+        if current
+            .lines()
+            .map(str::trim)
+            .any(|line| line.starts_with(OAUTH_REFRESH_FAILED_PREFIX))
+        {
+            return current.to_string();
+        }
+        return format!("{current}\n{candidate_reason}");
+    }
+    if current.starts_with(OAUTH_EXPIRED_PREFIX)
+        && candidate_reason.starts_with(OAUTH_REQUEST_FAILED_PREFIX)
    {
        return current.to_string();
    }
@@ -944,7 +1006,7 @@ mod tests {
    }

    #[test]
-    fn codex_invalid_state_keeps_oauth_expired_over_refresh_failure() {
+    fn codex_invalid_state_appends_refresh_failure_to_oauth_expired() {
        let mut key = StoredProviderCatalogKey::new(
            "key-1".to_string(),
            "provider-1".to_string(),
@@ -964,10 +1026,28 @@ mod tests {
                200,
            ),
            (
-                Some(100),
-                Some(format!("{OAUTH_EXPIRED_PREFIX}session expired"))
+                Some(200),
+                Some(format!(
+                    "{OAUTH_EXPIRED_PREFIX}session expired\n{OAUTH_REFRESH_FAILED_PREFIX}Token 续期失败"
+                ))
            )
        );
+    }
+
+    #[test]
+    fn codex_invalid_state_keeps_oauth_expired_over_request_failure() {
+        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_at_unix_secs = Some(100);
+        key.oauth_invalid_reason = Some(format!("{OAUTH_EXPIRED_PREFIX}session expired"));
+
        assert_eq!(
            codex_build_invalid_state(
                &key,
@@ -1010,6 +1090,68 @@ mod tests {
        );
    }

+    #[test]
+    fn auto_remove_refresh_failed_after_access_token_expiry() {
+        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.expires_at_unix_secs = Some(1_000);
+        key.oauth_invalid_reason = Some(format!("{OAUTH_REFRESH_FAILED_PREFIX}Token 续期失败"));
+
+        assert!(!super::should_auto_remove_oauth_invalid_key(
+            &key, None, 999
+        ));
+        assert!(super::should_auto_remove_oauth_invalid_key(
+            &key, None, 1_000
+        ));
+    }
+
+    #[test]
+    fn auto_remove_combined_refresh_and_access_token_failure() {
+        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.expires_at_unix_secs = Some(2_000);
+        key.oauth_invalid_reason = Some(format!("{OAUTH_REFRESH_FAILED_PREFIX}Token 续期失败"));
+
+        assert!(super::should_auto_remove_oauth_invalid_key(
+            &key,
+            Some("[OAUTH_EXPIRED] access token invalid"),
+            1_000,
+        ));
+    }
+
+    #[test]
+    fn does_not_auto_remove_access_token_failure_without_refresh_failure() {
+        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.expires_at_unix_secs = Some(1_000);
+        key.oauth_invalid_reason = Some(format!("{OAUTH_EXPIRED_PREFIX}session expired"));
+
+        assert!(!super::should_auto_remove_oauth_invalid_key(
+            &key, None, 1_001
+        ));
+    }
+
    #[test]
    fn parses_codex_spark_quota_from_additional_rate_limits() {
        let parsed = parse_codex_wham_usage_response(
--- a/frontend/src/features/pool/utils/tests/poolAdvancedDialog.spec.ts
+++ b/frontend/src/features/pool/utils/tests/poolAdvancedDialog.spec.ts
@@ -38,7 +38,7 @@ describe('poolAdvancedDialog', () => {
      {
        key: 'auto_remove_banned_keys',
        label: '异常自动清除',
-        description: '仅在检测到不可恢复的账号异常时自动从号池移除，不处理纯 Token 失效。',
+        description: '检测到不可恢复账号异常，或 RT 与 AT 均失效时自动从号池移除。',
      },
      {
        key: 'skip_exhausted_accounts',
--- a/frontend/src/features/pool/utils/poolAdvancedDialog.ts
+++ b/frontend/src/features/pool/utils/poolAdvancedDialog.ts
@@ -45,7 +45,7 @@ export function buildPoolHealthToggleCards(): PoolHealthToggleCard[] {
    {
      key: 'auto_remove_banned_keys',
      label: '异常自动清除',
-      description: '仅在检测到不可恢复的账号异常时自动从号池移除，不处理纯 Token 失效。',
+      description: '检测到不可恢复账号异常，或 RT 与 AT 均失效时自动从号池移除。',
    },
    {
      key: 'skip_exhausted_accounts',