Distinguish expired OAuth token status

This commit is contained in:
elky
2026-06-12 19:43:59 +08:00
parent 308cc88ef7
commit 68038c182b
16 changed files with 332 additions and 58 deletions
@@ -4,7 +4,9 @@ use std::sync::{
Arc, LazyLock,
};
use aether_admin::provider::pool as admin_provider_pool_pure;
use aether_admin::provider::{
pool as admin_provider_pool_pure, status as admin_provider_status_pure,
};
use aether_data_contracts::repository::candidate_selection::{
StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateOrder,
StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery,
@@ -1217,9 +1219,15 @@ fn pool_key_requires_reauth_for_scheduling(
.map(str::trim)
.unwrap_or_default();
if !invalid_reason.is_empty() {
if pool_oauth_reason_has_tag(invalid_reason, "[OAUTH_EXPIRED]")
|| pool_oauth_reason_has_tag(invalid_reason, "[ACCOUNT_BLOCK]")
{
let account_state = admin_provider_status_pure::resolve_pool_account_state(
None,
key.upstream_metadata.as_ref(),
Some(invalid_reason),
);
if account_state.blocked && !account_state.recoverable {
return true;
}
if pool_oauth_reason_has_tag(invalid_reason, "[ACCOUNT_BLOCK]") {
return true;
}
if pool_oauth_reason_has_tag(invalid_reason, "[REQUEST_FAILED]") {
@@ -1230,6 +1238,9 @@ fn pool_key_requires_reauth_for_scheduling(
.expires_at_unix_secs
.is_none_or(|expires_at| expires_at == 0 || expires_at <= now_unix_secs);
}
if pool_oauth_reason_has_tag(invalid_reason, "[OAUTH_EXPIRED]") {
return false;
}
return true;
}
@@ -3265,8 +3276,7 @@ mod tests {
let mut key_a_invalid = sample_codex_pool_key("provider-a", "key-a-invalid");
key_a_invalid.oauth_invalid_at_unix_secs = Some(1_710_000_000);
key_a_invalid.oauth_invalid_reason =
Some("[OAUTH_EXPIRED] Codex Token 无效或已过期 (401)".to_string());
key_a_invalid.oauth_invalid_reason = Some("[OAUTH_EXPIRED] token invalidated".to_string());
let exhausted_status_snapshot = json!({
"quota": {
"provider_type": "codex",
@@ -3370,8 +3380,7 @@ mod tests {
let mut key_a_invalid = sample_codex_pool_key("provider-a", "key-a-invalid");
key_a_invalid.oauth_invalid_at_unix_secs = Some(1_710_000_000);
key_a_invalid.oauth_invalid_reason =
Some("[OAUTH_EXPIRED] Codex Token 无效或已过期 (401)".to_string());
key_a_invalid.oauth_invalid_reason = Some("[OAUTH_EXPIRED] token invalidated".to_string());
key_a_invalid.status_snapshot = Some(json!({
"quota": {
"provider_type": "codex",
@@ -3463,6 +3472,9 @@ mod tests {
key.oauth_invalid_reason = Some("[REQUEST_FAILED] 账号状态检查失败".to_string());
key.oauth_invalid_at_unix_secs = Some(100);
assert!(!pool_key_requires_reauth_for_scheduling(&key, 300));
key.oauth_invalid_reason = Some("[OAUTH_EXPIRED] session expired".to_string());
assert!(!pool_key_requires_reauth_for_scheduling(&key, 300));
}
#[test]
@@ -3471,6 +3483,10 @@ mod tests {
key.oauth_invalid_reason = Some("[ACCOUNT_BLOCK] account has been deactivated".to_string());
assert!(pool_key_requires_reauth_for_scheduling(&key, 100));
key.oauth_invalid_reason = Some("[OAUTH_EXPIRED] token invalidated".to_string());
key.oauth_invalid_at_unix_secs = None;
assert!(pool_key_requires_reauth_for_scheduling(&key, 100));
key.oauth_invalid_reason = Some("Kiro Token 无效或已过期".to_string());
key.oauth_invalid_at_unix_secs = None;
assert!(pool_key_requires_reauth_for_scheduling(&key, 100));
@@ -21,7 +21,10 @@ fn oauth_invalid_reason_is_account_level_block(reason: Option<&str>) -> bool {
snapshot.blocked
&& !matches!(
snapshot.code.trim().to_ascii_lowercase().as_str(),
"oauth_token_invalid" | "oauth_expired" | "oauth_refresh_failed"
"oauth_token_invalid"
| "oauth_token_expired"
| "oauth_expired"
| "oauth_refresh_failed"
)
}
@@ -13,6 +13,10 @@ pub(super) fn codex_looks_like_token_invalidated(message: Option<&str>) -> bool
admin_provider_quota_pure::codex_looks_like_token_invalidated(message)
}
pub(super) fn codex_looks_like_token_expired(message: Option<&str>) -> bool {
admin_provider_quota_pure::codex_looks_like_token_expired(message)
}
pub(super) fn codex_looks_like_workspace_deactivated(message: Option<&str>) -> bool {
admin_provider_quota_pure::codex_looks_like_workspace_deactivated(message)
}
@@ -3,7 +3,7 @@ mod parse;
mod plan;
use self::invalid::{
codex_build_invalid_state, codex_looks_like_token_invalidated,
codex_build_invalid_state, codex_looks_like_token_expired, codex_looks_like_token_invalidated,
codex_looks_like_workspace_deactivated, codex_soft_request_failure_reason,
codex_structured_invalid_reason,
};
@@ -275,6 +275,7 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
}
403 => {
let candidate_reason = if codex_looks_like_token_invalidated(err_msg.as_deref())
|| codex_looks_like_token_expired(err_msg.as_deref())
{
codex_structured_invalid_reason(403, err_msg.as_deref())
} else {
@@ -304,6 +304,7 @@ fn admin_pool_trimmed_string(value: Option<&Value>) -> Option<String> {
fn admin_pool_account_code_status_filter(code: &str) -> Option<&'static str> {
match code.trim().to_ascii_lowercase().as_str() {
"oauth_token_invalid" => Some("invalid"),
"oauth_token_expired" => Some("expired"),
"account_banned" | "account_suspended" => Some("account_banned"),
"account_disabled" => Some("account_disabled"),
"workspace_deactivated" => Some("workspace_deactivated"),
@@ -265,14 +265,16 @@ fn build_provider_key_oauth_status_snapshot(key: &StoredProviderCatalogKey) -> V
if let Some(reason) =
tagged_oauth_invalid_reason(invalid_reason.as_deref(), OAUTH_EXPIRED_PREFIX)
{
let (code, label) =
admin_provider_status_pure::oauth_token_snapshot_status_parts(reason.as_str());
return json!({
"code": "invalid",
"label": "已失效",
"code": code,
"label": label,
"reason": reason,
"expires_at": expires_at_unix_secs,
"invalid_at": invalid_at_unix_secs,
"source": "oauth_invalid",
"requires_reauth": true,
"requires_reauth": code == "invalid",
"expiring_soon": false,
});
}
@@ -2031,7 +2031,7 @@ mod tests {
}
#[tokio::test]
async fn oauth_invalidation_marks_codex_key_invalid() {
async fn oauth_invalidation_marks_codex_key_expired() {
let state = codex_state();
let plan = sample_codex_plan();
@@ -2069,7 +2069,7 @@ mod tests {
.and_then(|value| value.get("oauth"))
.and_then(|value| value.get("code"))
.and_then(Value::as_str),
Some("invalid")
Some("expired")
);
}
@@ -376,9 +376,6 @@ fn oauth_invalid_reason_blocks_scheduling(
now_unix_secs: u64,
) -> bool {
let trimmed_reason = invalid_reason.trim();
if oauth_invalid_reason_has_tag(trimmed_reason, "[OAUTH_EXPIRED]") {
return true;
}
let account_state = admin_provider_status_pure::resolve_pool_account_state(
Some(provider_type),
@@ -433,6 +430,7 @@ fn oauth_account_state_code_is_hard_block(code: &str) -> bool {
| "account_forbidden"
| "account_blocked"
| "account_verification"
| "oauth_token_invalid"
)
}
@@ -2064,7 +2064,7 @@ async fn keeps_codex_candidate_selectable_when_oauth_token_is_expired() {
let mut key = sample_key("key-codex", "provider-codex", Some(10));
key.auth_type = "oauth".to_string();
key.oauth_invalid_at_unix_secs = Some(1_710_000_000);
key.oauth_invalid_reason = Some("Codex Token 无效或已过期".to_string());
key.oauth_invalid_reason = Some("[OAUTH_EXPIRED] session expired".to_string());
key
}],
));
+9 -4
View File
@@ -171,7 +171,10 @@ fn oauth_invalid_reason_is_account_block(reason: Option<&str>) -> bool {
snapshot.blocked
&& !matches!(
snapshot.code.trim().to_ascii_lowercase().as_str(),
"oauth_token_invalid" | "oauth_expired" | "oauth_refresh_failed"
"oauth_token_invalid"
| "oauth_token_expired"
| "oauth_expired"
| "oauth_refresh_failed"
)
}
@@ -356,14 +359,16 @@ fn build_oauth_status_snapshot_value(key: &StoredProviderCatalogKey) -> Value {
let invalid_reason = trimmed_reason(key.oauth_invalid_reason.as_deref());
if let Some(reason) = tagged_reason(invalid_reason.as_deref(), OAUTH_EXPIRED_PREFIX) {
let (code, label) =
aether_admin::provider::status::oauth_token_snapshot_status_parts(reason.as_str());
return json!({
"code": "invalid",
"label": "已失效",
"code": code,
"label": label,
"reason": reason,
"expires_at": expires_at_unix_secs,
"invalid_at": invalid_at_unix_secs,
"source": "oauth_invalid",
"requires_reauth": true,
"requires_reauth": code == "invalid",
"expiring_soon": false,
});
}
@@ -5403,7 +5403,7 @@ async fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_p
assert!(stored_key.oauth_invalid_at_unix_secs.is_some());
assert_eq!(
stored_key.oauth_invalid_reason.as_deref(),
Some("[OAUTH_EXPIRED] Codex Token 无效或已过期 (401)")
Some("[OAUTH_EXPIRED] Codex Token 已过期 (401)")
);
} else if account_state_recheck_attempted
&& payload["account_state_recheck_error"] == "wham/usage API 返回状态码 403"
@@ -5449,7 +5449,40 @@ async fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_p
oauth_snapshot.get("expires_at"),
auth_config.get("expires_at")
);
if stored_key.oauth_invalid_reason.is_some() {
if stored_key
.oauth_invalid_reason
.as_deref()
.is_some_and(|reason| reason.starts_with("[OAUTH_EXPIRED]"))
{
assert_eq!(
oauth_snapshot
.get("code")
.and_then(serde_json::Value::as_str),
Some("expired")
);
assert_eq!(
oauth_snapshot
.get("label")
.and_then(serde_json::Value::as_str),
Some("已过期")
);
assert_eq!(
oauth_snapshot.get("reason"),
Some(&json!("Codex Token 已过期 (401)"))
);
assert_eq!(
oauth_snapshot
.get("requires_reauth")
.and_then(serde_json::Value::as_bool),
Some(false)
);
assert_eq!(
oauth_snapshot
.get("expiring_soon")
.and_then(serde_json::Value::as_bool),
Some(false)
);
} else if stored_key.oauth_invalid_reason.is_some() {
assert_eq!(
oauth_snapshot
.get("code")
@@ -5463,8 +5496,11 @@ async fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_p
Some("已失效")
);
assert_eq!(
oauth_snapshot.get("reason"),
Some(&json!("Codex Token 无效或已过期 (401)"))
oauth_snapshot
.get("reason")
.and_then(serde_json::Value::as_str)
.is_some_and(|reason| !reason.trim().is_empty()),
true
);
assert_eq!(
oauth_snapshot
@@ -3526,6 +3526,15 @@ 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_invalidated_key = sample_key(
"key-openai-oauth-invalidated",
"provider-openai",
"openai:chat",
"sk-oauth-invalidated",
);
oauth_invalidated_key.name = "oauth-invalidated".to_string();
oauth_invalidated_key.oauth_invalid_reason =
Some("[OAUTH_EXPIRED] token invalidated".to_string());
let mut oauth_expired_key = sample_key(
"key-openai-oauth-expired",
"provider-openai",
@@ -3533,7 +3542,7 @@ async fn gateway_cleans_up_admin_pool_banned_keys_locally_with_trusted_admin_pri
"sk-oauth-expired",
);
oauth_expired_key.name = "oauth-expired".to_string();
oauth_expired_key.oauth_invalid_reason = Some("[OAUTH_EXPIRED] token invalidated".to_string());
oauth_expired_key.oauth_invalid_reason = Some("[OAUTH_EXPIRED] session expired".to_string());
let mut healthy_key = sample_key(
"key-openai-healthy",
"provider-openai",
@@ -3545,7 +3554,12 @@ 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, oauth_expired_key, healthy_key],
vec![
banned_key,
oauth_invalidated_key,
oauth_expired_key,
healthy_key,
],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
@@ -3575,8 +3589,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()])
@@ -3586,6 +3600,9 @@ async fn gateway_cleans_up_admin_pool_banned_keys_locally_with_trusted_admin_pri
assert!(remaining_keys
.iter()
.any(|key| key.id == "key-openai-oauth-expired"));
assert!(!remaining_keys
.iter()
.any(|key| key.id == "key-openai-oauth-invalidated"));
assert!(remaining_keys
.iter()
.any(|key| key.id == "key-openai-healthy"));