修复热池调度已失效 OAuth 账号

This commit is contained in:
Novick Yuan
2026-05-23 03:17:09 +08:00
parent c7641dad0a
commit c0252387b4
@@ -1098,12 +1098,54 @@ fn build_pool_catalog_key_context(
let mut signals =
provider_pool_service.member_signals(provider_type, key, auth_config.as_ref());
signals.account_blocked |= admin_provider_pool_pure::admin_pool_key_is_known_banned(key);
signals.account_blocked |=
pool_key_requires_reauth_for_scheduling(key, current_unix_ms().saturating_div(1000));
signals.health_score = health_score;
signals.latency_avg_ms = latency_avg_ms;
signals.catalog_lru_score = Some(key.last_used_at_unix_secs.unwrap_or(0) as f64);
signals
}
fn pool_key_requires_reauth_for_scheduling(
key: &StoredProviderCatalogKey,
now_unix_secs: u64,
) -> bool {
if !key.auth_type.trim().eq_ignore_ascii_case("oauth") {
return false;
}
let invalid_reason = key
.oauth_invalid_reason
.as_deref()
.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]")
{
return true;
}
if pool_oauth_reason_has_tag(invalid_reason, "[REQUEST_FAILED]") {
return false;
}
if pool_oauth_reason_has_tag(invalid_reason, "[REFRESH_FAILED]") {
return key
.expires_at_unix_secs
.is_none_or(|expires_at| expires_at == 0 || expires_at <= now_unix_secs);
}
return true;
}
key.oauth_invalid_at_unix_secs.is_some()
}
fn pool_oauth_reason_has_tag(reason: &str, tag: &str) -> bool {
reason
.lines()
.map(str::trim)
.any(|line| line.starts_with(tag))
}
fn apply_local_execution_pool_scheduler_with_runtime_map(
candidates: Vec<EligibleLocalExecutionCandidate>,
runtime_by_provider: &BTreeMap<String, AdminProviderPoolRuntimeState>,
@@ -1454,10 +1496,11 @@ fn apply_pool_orchestration(
#[cfg(test)]
mod tests {
use super::{
admin_provider_pool_quota_probe_active_members_key,
admin_provider_pool_quota_probe_active_members_key, apply_local_execution_pool_scheduler,
apply_local_execution_pool_scheduler_with_runtime_map,
apply_local_execution_pool_scheduler_with_runtime_map_outcome,
build_pool_catalog_key_context, pool_config_for_candidate,
pool_key_requires_reauth_for_scheduling,
prune_unschedulable_active_probe_members_for_request,
remove_active_probe_members_for_request, should_trigger_active_probe_burst_for_request,
PoolCatalogKeyContext, PoolKeyCursor, POOL_ACTIVE_PROBE_SEALED_SKIP_REASON,
@@ -2878,6 +2921,243 @@ mod tests {
}));
}
#[tokio::test]
async fn pool_scheduler_skips_invalid_and_exhausted_high_priority_hot_pool_before_fallback_provider(
) {
let provider_config = Some(json!({
"pool_advanced": {
"probing_enabled": true,
"skip_exhausted_accounts": true,
"scheduling_presets": [
{"preset": "single_account", "enabled": true}
]
}
}));
let provider_a = sample_codex_pool_provider("provider-a", 0, provider_config.clone());
let provider_b = sample_codex_pool_provider("provider-b", 10, provider_config.clone());
let endpoint_a = sample_codex_pool_endpoint("provider-a", "endpoint-a");
let endpoint_b = sample_codex_pool_endpoint("provider-b", "endpoint-b");
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());
let exhausted_status_snapshot = json!({
"quota": {
"provider_type": "codex",
"exhausted": true,
"usage_ratio": 1.0,
"windows": [
{
"code": "daily",
"used_ratio": 1.0,
"remaining_ratio": 0.0
}
]
}
});
key_a_invalid.status_snapshot = Some(exhausted_status_snapshot.clone());
let mut key_a_exhausted = sample_codex_pool_key("provider-a", "key-a-exhausted");
key_a_exhausted.status_snapshot = Some(exhausted_status_snapshot);
let key_b_ready = sample_codex_pool_key("provider-b", "key-b-ready");
let rows = vec![
sample_codex_pool_row("provider-a", "endpoint-a", "key-a-invalid", 0),
sample_codex_pool_row("provider-a", "endpoint-a", "key-a-exhausted", 0),
sample_codex_pool_row("provider-b", "endpoint-b", "key-b-ready", 10),
];
let data_state =
GatewayDataState::with_provider_catalog_and_minimal_candidate_selection_for_tests(
Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider_a, provider_b],
vec![endpoint_a, endpoint_b],
vec![key_a_invalid, key_a_exhausted, key_b_ready],
)),
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(rows)),
)
.with_encryption_key_for_tests(aether_crypto::DEVELOPMENT_ENCRYPTION_KEY);
let app = AppState::new()
.expect("state should build")
.with_data_state_for_tests(data_state);
app.runtime_state
.set_add(
&admin_provider_pool_quota_probe_active_members_key("provider-a"),
"key-a-invalid",
)
.await
.expect("provider-a hot member should insert");
app.runtime_state
.set_add(
&admin_provider_pool_quota_probe_active_members_key("provider-b"),
"key-b-ready",
)
.await
.expect("provider-b hot member should insert");
let group_a =
sample_codex_pool_group("provider-a", "endpoint-a", 0, provider_config.clone());
let group_b = sample_codex_pool_group("provider-b", "endpoint-b", 10, provider_config);
let (scheduled, skipped) = apply_local_execution_pool_scheduler(
PlannerAppState::new(&app),
vec![group_a, group_b],
None,
Some("gpt-5"),
None,
)
.await;
assert_eq!(
scheduled
.iter()
.map(|item| item.candidate.key_id.as_str())
.collect::<Vec<_>>(),
vec!["key-b-ready"]
);
let skipped_pairs = skipped
.iter()
.map(|item| (item.candidate.key_id.as_str(), item.skip_reason))
.collect::<Vec<_>>();
assert!(skipped_pairs.contains(&("key-a-invalid", "pool_account_blocked")));
assert!(skipped_pairs.contains(&(
"key-a-exhausted",
aether_pool_core::POOL_ACCOUNT_EXHAUSTED_SKIP_REASON
)));
}
#[tokio::test]
async fn pool_scheduler_skips_invalid_high_priority_hot_pool_account_even_with_remaining_quota()
{
let provider_config = Some(json!({
"pool_advanced": {
"probing_enabled": true,
"skip_exhausted_accounts": true,
"scheduling_presets": [
{"preset": "single_account", "enabled": true}
]
}
}));
let provider_a = sample_codex_pool_provider("provider-a", 0, provider_config.clone());
let provider_b = sample_codex_pool_provider("provider-b", 10, provider_config.clone());
let endpoint_a = sample_codex_pool_endpoint("provider-a", "endpoint-a");
let endpoint_b = sample_codex_pool_endpoint("provider-b", "endpoint-b");
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.status_snapshot = Some(json!({
"quota": {
"provider_type": "codex",
"exhausted": false,
"usage_ratio": 0.25,
"windows": [
{
"code": "daily",
"used_ratio": 0.25,
"remaining_ratio": 0.75
}
]
}
}));
let key_b_ready = sample_codex_pool_key("provider-b", "key-b-ready");
let rows = vec![
sample_codex_pool_row("provider-a", "endpoint-a", "key-a-invalid", 0),
sample_codex_pool_row("provider-b", "endpoint-b", "key-b-ready", 10),
];
let data_state =
GatewayDataState::with_provider_catalog_and_minimal_candidate_selection_for_tests(
Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider_a, provider_b],
vec![endpoint_a, endpoint_b],
vec![key_a_invalid, key_b_ready],
)),
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(rows)),
)
.with_encryption_key_for_tests(aether_crypto::DEVELOPMENT_ENCRYPTION_KEY);
let app = AppState::new()
.expect("state should build")
.with_data_state_for_tests(data_state);
app.runtime_state
.set_add(
&admin_provider_pool_quota_probe_active_members_key("provider-a"),
"key-a-invalid",
)
.await
.expect("provider-a hot member should insert");
app.runtime_state
.set_add(
&admin_provider_pool_quota_probe_active_members_key("provider-b"),
"key-b-ready",
)
.await
.expect("provider-b hot member should insert");
let group_a =
sample_codex_pool_group("provider-a", "endpoint-a", 0, provider_config.clone());
let group_b = sample_codex_pool_group("provider-b", "endpoint-b", 10, provider_config);
let (scheduled, skipped) = apply_local_execution_pool_scheduler(
PlannerAppState::new(&app),
vec![group_a, group_b],
None,
Some("gpt-5"),
None,
)
.await;
assert_eq!(
scheduled
.iter()
.map(|item| item.candidate.key_id.as_str())
.collect::<Vec<_>>(),
vec!["key-b-ready"]
);
let skipped_pairs = skipped
.iter()
.map(|item| (item.candidate.key_id.as_str(), item.skip_reason))
.collect::<Vec<_>>();
assert!(skipped_pairs.contains(&("key-a-invalid", "pool_account_blocked")));
}
#[test]
fn pool_key_reauth_scheduling_keeps_recoverable_oauth_markers_usable() {
let mut key = sample_codex_pool_key("provider-a", "key-refresh-failed");
key.expires_at_unix_secs = Some(200);
key.oauth_invalid_reason = Some(
"[REFRESH_FAILED] Token 续期失败 (401): refresh_token 已被使用并轮换,请重新登录授权"
.to_string(),
);
assert!(!pool_key_requires_reauth_for_scheduling(&key, 100));
assert!(pool_key_requires_reauth_for_scheduling(&key, 200));
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));
}
#[test]
fn pool_key_reauth_scheduling_blocks_invalid_oauth_markers_without_affecting_non_oauth_keys() {
let mut key = sample_codex_pool_key("provider-a", "key-invalid");
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("Kiro Token 无效或已过期".to_string());
key.oauth_invalid_at_unix_secs = None;
assert!(pool_key_requires_reauth_for_scheduling(&key, 100));
key.oauth_invalid_reason = None;
key.oauth_invalid_at_unix_secs = Some(100);
assert!(pool_key_requires_reauth_for_scheduling(&key, 100));
key.auth_type = "api_key".to_string();
assert!(!pool_key_requires_reauth_for_scheduling(&key, 100));
}
#[tokio::test]
async fn pool_key_cursor_simulates_large_lru_pool_with_lazy_pages_and_dynamic_skips() {
const KEY_COUNT: usize = 2048;
@@ -3325,6 +3605,209 @@ mod tests {
(provider, endpoint, keys, rows)
}
fn sample_codex_pool_provider(
provider_id: &str,
provider_priority: i32,
provider_config: Option<serde_json::Value>,
) -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
provider_id.to_string(),
provider_id.to_string(),
Some("https://example.com".to_string()),
"codex".to_string(),
)
.expect("provider should build")
.with_routing_fields(provider_priority)
.with_transport_fields(
true,
false,
false,
None,
None,
None,
None,
None,
provider_config,
)
}
fn sample_codex_pool_endpoint(
provider_id: &str,
endpoint_id: &str,
) -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
endpoint_id.to_string(),
provider_id.to_string(),
"openai:responses".to_string(),
Some("openai".to_string()),
Some("responses".to_string()),
true,
)
.expect("endpoint should build")
.with_health_score(1.0)
.with_transport_fields(
"https://example.com/v1/responses".to_string(),
None,
None,
None,
None,
None,
None,
None,
)
.expect("endpoint transport should build")
}
fn sample_codex_pool_key(provider_id: &str, key_id: &str) -> StoredProviderCatalogKey {
let mut key = StoredProviderCatalogKey::new(
key_id.to_string(),
provider_id.to_string(),
key_id.to_string(),
"oauth".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(json!(["openai:responses"])),
Some(format!("secret-{key_id}")),
None,
None,
Some(json!({"openai:responses": 1})),
None,
Some(4_102_444_800),
None,
None,
)
.expect("key transport should build");
key.internal_priority = 10;
key
}
fn sample_codex_pool_row(
provider_id: &str,
endpoint_id: &str,
key_id: &str,
provider_priority: i32,
) -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: provider_id.to_string(),
provider_name: provider_id.to_string(),
provider_type: "codex".to_string(),
provider_priority,
provider_is_active: true,
endpoint_id: endpoint_id.to_string(),
endpoint_api_format: "openai:responses".to_string(),
endpoint_api_family: Some("openai".to_string()),
endpoint_kind: Some("responses".to_string()),
endpoint_is_active: true,
key_id: key_id.to_string(),
key_name: key_id.to_string(),
key_auth_type: "oauth".to_string(),
key_is_active: true,
key_api_formats: Some(vec!["openai:responses".to_string()]),
key_allowed_models: None,
key_capabilities: None,
key_internal_priority: 10,
key_global_priority_by_format: Some(json!({"openai:responses": 1})),
model_id: "model-1".to_string(),
global_model_id: "global-model-1".to_string(),
global_model_name: "gpt-5".to_string(),
global_model_mappings: None,
global_model_supports_streaming: Some(true),
model_provider_model_name: "gpt-5".to_string(),
model_provider_model_mappings: None,
model_supports_streaming: Some(true),
model_is_active: true,
model_is_available: true,
}
}
fn sample_codex_pool_group(
provider_id: &str,
endpoint_id: &str,
provider_priority: i32,
provider_config: Option<serde_json::Value>,
) -> EligibleLocalExecutionCandidate {
EligibleLocalExecutionCandidate {
kind: LocalExecutionCandidateKind::PoolGroup,
candidate: SchedulerMinimalCandidateSelectionCandidate {
provider_id: provider_id.to_string(),
provider_name: provider_id.to_string(),
provider_type: "codex".to_string(),
provider_priority,
endpoint_id: endpoint_id.to_string(),
endpoint_api_format: "openai:responses".to_string(),
key_id: format!("{provider_id}-pool-group"),
key_name: format!("{provider_id}-pool-group"),
key_auth_type: "oauth".to_string(),
key_internal_priority: 10,
key_global_priority_for_format: Some(1),
key_capabilities: None,
model_id: "model-1".to_string(),
global_model_id: "global-model-1".to_string(),
global_model_name: "gpt-5".to_string(),
selected_provider_model_name: "gpt-5".to_string(),
mapping_matched_model: None,
},
provider_api_format: "openai:responses".to_string(),
orchestration: LocalExecutionCandidateMetadata::default(),
ranking: None,
transport: Arc::new(crate::ai_serving::GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider {
id: provider_id.to_string(),
name: provider_id.to_string(),
provider_type: "codex".to_string(),
website: None,
is_active: true,
keep_priority_on_conversion: false,
enable_format_conversion: false,
concurrent_limit: None,
max_retries: None,
proxy: None,
request_timeout_secs: None,
stream_first_byte_timeout_secs: None,
config: provider_config,
},
endpoint: GatewayProviderTransportEndpoint {
id: endpoint_id.to_string(),
provider_id: provider_id.to_string(),
api_format: "openai:responses".to_string(),
api_family: Some("openai".to_string()),
endpoint_kind: Some("responses".to_string()),
is_active: true,
base_url: "https://example.com/v1/responses".to_string(),
header_rules: None,
body_rules: None,
max_retries: None,
custom_path: None,
config: None,
format_acceptance_config: None,
proxy: None,
},
key: GatewayProviderTransportKey {
id: format!("{provider_id}-pool-group"),
provider_id: provider_id.to_string(),
name: format!("{provider_id}-pool-group"),
auth_type: "oauth".to_string(),
is_active: true,
api_formats: Some(vec!["openai:responses".to_string()]),
auth_type_by_format: None,
allow_auth_channel_mismatch_formats: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,
global_priority_by_format: None,
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
decrypted_api_key: "secret".to_string(),
decrypted_auth_config: None,
},
}),
}
}
fn routing_policy_with_allowed_keys<const N: usize>(
key_ids: [&str; N],
) -> ResolvedRoutingPolicy {