mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-08 02:17:46 +08:00
Distinguish expired OAuth token status
This commit is contained in:
@@ -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
|
||||
}],
|
||||
));
|
||||
|
||||
@@ -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"));
|
||||
|
||||
Reference in New Issue
Block a user