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

fix provider pool quota status handling

Browse Source
This commit is contained in:
fawney19
2026-05-07 11:28:44 +08:00
parent 01c9f9be49
commit fd44906bb7
37 changed files with 1844 additions and 500 deletions
Download Patch File Download Diff File Expand all files Collapse all files

52
apps/aether-gateway/src/scheduler/candidate/runtime.rs
View File

@@ -10,6 +10,7 @@ use aether_scheduler_core::{
candidate_is_selectable_with_runtime_state, candidate_runtime_skip_reason_with_state,
CandidateRuntimeSelectabilityInput,
};
use std::time::{SystemTime, UNIX_EPOCH};
use crate::data::auth::GatewayAuthApiKeySnapshot;
use crate::GatewayError;
@@ -343,7 +344,7 @@ fn read_key_oauth_invalid_map(
fn key_requires_oauth_reauth(
key: &StoredProviderCatalogKey,
provider_type: &str,
_now_unix_secs: u64,
now_unix_secs: u64,
) -> bool {
if !key.auth_type.trim().eq_ignore_ascii_case("oauth") {
return false;
@@ -355,34 +356,69 @@ fn key_requires_oauth_reauth(
.map(str::trim)
.unwrap_or_default();
if !invalid_reason.is_empty() {
return oauth_invalid_reason_is_hard_account_block(key, provider_type, invalid_reason);
return oauth_invalid_reason_blocks_scheduling(
key,
provider_type,
invalid_reason,
now_unix_secs,
);
}
false
}
fn oauth_invalid_reason_is_hard_account_block(
fn oauth_invalid_reason_blocks_scheduling(
key: &StoredProviderCatalogKey,
provider_type: &str,
invalid_reason: &str,
now_unix_secs: u64,
) -> bool {
if provider_type.trim().eq_ignore_ascii_case("kiro")
&& invalid_reason.trim().starts_with("[REFRESH_FAILED] ")
{
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),
key.upstream_metadata.as_ref(),
Some(invalid_reason),
Some(trimmed_reason),
);
account_state.blocked
if account_state.blocked
&& !account_state.recoverable
&& account_state
.code
.as_deref()
.is_some_and(oauth_account_state_code_is_hard_block)
{
return true;
}
if oauth_invalid_reason_has_tag(trimmed_reason, "[REFRESH_FAILED]") {
return oauth_access_token_expired(key, now_unix_secs);
}
false
}
fn oauth_invalid_reason_has_tag(reason: &str, tag: &str) -> bool {
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 {
SystemTime::now()
.duration_since(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)
}
fn oauth_account_state_code_is_hard_block(code: &str) -> bool {

198
apps/aether-gateway/src/scheduler/candidate/tests/selection.rs
View File

@@ -1692,6 +1692,7 @@ async fn keeps_refresh_failed_oauth_candidate_selectable_before_local_auth_resol
{
let mut key = sample_key("key-codex", "provider-codex", Some(10));
key.auth_type = "oauth".to_string();
key.expires_at_unix_secs = Some(1_710_000_200);
key.oauth_invalid_at_unix_secs = Some(1_710_000_000);
key.oauth_invalid_reason = Some(
"[REFRESH_FAILED] Token 续期失败 (401): refresh_token 已被使用并轮换，请重新登录授权"
@@ -1733,6 +1734,136 @@ async fn keeps_refresh_failed_oauth_candidate_selectable_before_local_auth_resol
assert!(skipped.is_empty());
}
#[tokio::test]
async fn skips_refresh_failed_oauth_candidate_after_access_token_expiry() {
let mut row = sample_row();
row.provider_id = "provider-codex".to_string();
row.provider_name = "codex".to_string();
row.provider_type = "codex".to_string();
row.endpoint_id = "endpoint-codex".to_string();
row.endpoint_api_format = "openai:responses".to_string();
row.key_id = "key-codex".to_string();
row.key_name = "codex-refresh-failed-expired".to_string();
row.key_auth_type = "oauth".to_string();
row.key_api_formats = Some(vec!["openai:responses".to_string()]);
row.key_global_priority_by_format = Some(serde_json::json!({"openai:responses": 1}));
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
row,
]));
let mut provider = sample_provider("provider-codex", None);
provider.provider_type = "codex".to_string();
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
Vec::new(),
vec![{
let mut key = sample_key("key-codex", "provider-codex", Some(10));
key.auth_type = "oauth".to_string();
key.expires_at_unix_secs = Some(1_710_000_000);
key.oauth_invalid_at_unix_secs = Some(1_710_000_000);
key.oauth_invalid_reason = Some(
"[REFRESH_FAILED] Token 续期失败 (401): refresh_token 已被使用并轮换，请重新登录授权"
.to_string(),
);
key
}],
));
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![]));
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(
GatewayDataState::with_candidate_selection_provider_catalog_quota_and_request_candidates_for_tests(
candidates,
provider_catalog,
quotas,
request_candidates,
),
);
let (selected, skipped) = collect_selectable_candidates_with_skip_reasons(
state.data.as_ref(),
&state,
"openai:responses",
"gpt-4.1",
false,
None,
1_710_000_100,
)
.await
.expect("selection should succeed");
assert!(selected.is_empty());
assert_eq!(skipped.len(), 1);
assert_eq!(skipped[0].candidate.key_id, "key-codex");
assert_eq!(skipped[0].skip_reason, "oauth_invalid");
}
#[tokio::test]
async fn skips_oauth_candidate_with_account_block_even_when_refresh_failed_is_present() {
let mut row = sample_row();
row.provider_id = "provider-codex".to_string();
row.provider_name = "codex".to_string();
row.provider_type = "codex".to_string();
row.endpoint_id = "endpoint-codex".to_string();
row.endpoint_api_format = "openai:responses".to_string();
row.key_id = "key-codex".to_string();
row.key_name = "codex-account-blocked-refresh-failed".to_string();
row.key_auth_type = "oauth".to_string();
row.key_api_formats = Some(vec!["openai:responses".to_string()]);
row.key_global_priority_by_format = Some(serde_json::json!({"openai:responses": 1}));
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
row,
]));
let mut provider = sample_provider("provider-codex", None);
provider.provider_type = "codex".to_string();
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
Vec::new(),
vec![{
let mut key = sample_key("key-codex", "provider-codex", Some(10));
key.auth_type = "oauth".to_string();
key.expires_at_unix_secs = Some(1_710_000_200);
key.oauth_invalid_at_unix_secs = Some(1_710_000_000);
key.oauth_invalid_reason = Some(
"[REFRESH_FAILED] Token 续期失败 (401): refresh_token 已被使用并轮换，请重新登录授权\n[ACCOUNT_BLOCK] account has been deactivated"
.to_string(),
);
key
}],
));
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![]));
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(
GatewayDataState::with_candidate_selection_provider_catalog_quota_and_request_candidates_for_tests(
candidates,
provider_catalog,
quotas,
request_candidates,
),
);
let (selected, skipped) = collect_selectable_candidates_with_skip_reasons(
state.data.as_ref(),
&state,
"openai:responses",
"gpt-4.1",
false,
None,
1_710_000_100,
)
.await
.expect("selection should succeed");
assert!(selected.is_empty());
assert_eq!(skipped.len(), 1);
assert_eq!(skipped[0].candidate.key_id, "key-codex");
assert_eq!(skipped[0].skip_reason, "oauth_invalid");
}
#[tokio::test]
async fn keeps_request_failed_oauth_candidate_selectable() {
let mut row = sample_row();
@@ -1971,7 +2102,7 @@ async fn keeps_refreshable_kiro_candidate_selectable_when_oauth_token_expired()
}
#[tokio::test]
async fn skips_kiro_candidate_after_refresh_token_failure() {
async fn keeps_kiro_candidate_selectable_after_refresh_token_failure_until_access_token_expiry() {
let mut row = sample_row();
row.provider_id = "provider-kiro".to_string();
row.provider_name = "kiro".to_string();
@@ -1995,6 +2126,71 @@ async fn skips_kiro_candidate_after_refresh_token_failure() {
let mut key = sample_key("key-kiro", "provider-kiro", Some(10));
key.auth_type = "oauth".to_string();
key.encrypted_auth_config = Some("encrypted-refreshable-session".to_string());
key.expires_at_unix_secs = Some(1_710_000_200);
key.oauth_invalid_at_unix_secs = Some(1_710_000_000);
key.oauth_invalid_reason = Some(
"[REFRESH_FAILED] Token 续期失败 (401): refresh_token 无效、已过期或已撤销，请重新登录授权"
.to_string(),
);
key
}],
));
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![]));
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(
GatewayDataState::with_candidate_selection_provider_catalog_quota_and_request_candidates_for_tests(
candidates,
provider_catalog,
quotas,
request_candidates,
),
);
let (selected, skipped) = collect_selectable_candidates_with_skip_reasons(
state.data.as_ref(),
&state,
"claude:messages",
"gpt-4.1",
false,
None,
1_710_000_100,
)
.await
.expect("selection should succeed");
assert_eq!(selected.len(), 1);
assert_eq!(selected[0].provider_id, "provider-kiro");
assert!(skipped.is_empty());
}
#[tokio::test]
async fn skips_kiro_candidate_after_refresh_token_failure_and_access_token_expiry() {
let mut row = sample_row();
row.provider_id = "provider-kiro".to_string();
row.provider_name = "kiro".to_string();
row.provider_type = "kiro".to_string();
row.endpoint_id = "endpoint-kiro".to_string();
row.endpoint_api_format = "claude:messages".to_string();
row.key_id = "key-kiro".to_string();
row.key_name = "kiro-refresh-failed-expired".to_string();
row.key_auth_type = "oauth".to_string();
row.key_api_formats = Some(vec!["claude:messages".to_string()]);
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
row,
]));
let mut provider = sample_provider("provider-kiro", None);
provider.provider_type = "kiro".to_string();
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
Vec::new(),
vec![{
let mut key = sample_key("key-kiro", "provider-kiro", Some(10));
key.auth_type = "oauth".to_string();
key.encrypted_auth_config = Some("encrypted-refreshable-session".to_string());
key.expires_at_unix_secs = Some(1_710_000_000);
key.oauth_invalid_at_unix_secs = Some(1_710_000_000);
key.oauth_invalid_reason = Some(
"[REFRESH_FAILED] Token 续期失败 (401): refresh_token 无效、已过期或已撤销，请重新登录授权"