i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-06 01:17:46 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge origin/main into fix/gemini-cli-v1internal

Browse Source
This commit is contained in:
Mas0nShi
2026-05-28 11:58:00 +08:00
parent 0d80db8a8d 14ad6e9b75
commit 9ad9858ac2
410 changed files with 38026 additions and 6621 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/dispatch/pool_scheduler.rs
+697 -52
View File
@@ -1,5 +1,8 @@
use std::collections::{btree_map::Entry, BTreeMap, BTreeSet, VecDeque};
use std::sync::atomic::{AtomicU64, Ordering as AtomicOrdering};
use std::sync::{
atomic::{AtomicU64, Ordering as AtomicOrdering},
Arc, LazyLock,
};
use aether_admin::provider::pool as admin_provider_pool_pure;
use aether_data_contracts::repository::candidate_selection::{
@@ -19,7 +22,8 @@ use aether_pool_core::{
};
use aether_provider_pool::ProviderPoolService;
use aether_routing_core::{RankingOverlay, ResolvedRoutingPolicy};
use tracing::warn;
use tokio::sync::Semaphore;
use tracing::{debug, warn};
use crate::ai_serving::{
candidate_auth_channel_skip_reason, candidate_common_transport_skip_reason,
@@ -43,8 +47,13 @@ use crate::maintenance::spawn_pool_quota_probe_replenish_for_request;
use crate::orchestration::LocalExecutionCandidateMetadata;
static LOAD_BALANCE_SEQUENCE: AtomicU64 = AtomicU64::new(0);
static POOL_SCORE_SCHEDULE_INTEREST_SEMAPHORE: LazyLock<Arc<Semaphore>> =
LazyLock::new(|| Arc::new(Semaphore::new(POOL_SCORE_SCHEDULE_INTEREST_CONCURRENCY)));
const POOL_ACTIVE_PROBE_SEALED_SKIP_REASON: &str = "pool_active_probe_sealed";
const ROUTING_PROFILE_DISALLOWED_KEY_SKIP_REASON: &str = "routing_profile_disallowed_key";
const POOL_SCORE_SCHEDULE_INTEREST_CONCURRENCY: usize = 4;
const POOL_SCORE_SCHEDULE_INTEREST_MAX_PER_BATCH: usize = 16;
const POOL_SCORE_SCHEDULE_INTEREST_MIN_INTERVAL_SECS: u64 = 60;
type PoolCatalogKeyContext = PoolMemberSignals;
@@ -334,9 +343,11 @@ pub(crate) struct PoolKeyCursor<'a> {
pool_key_order: StoredPoolKeyCandidateOrder,
next_offset: u32,
scanned_keys: u32,
budget_scanned_keys: u32,
window_size: u32,
page_size: u32,
max_scanned_keys: u32,
absolute_max_scanned_keys: u32,
score_top_n: u32,
score_phase_loaded: bool,
skip_reason_counts: BTreeMap<&'static str, u32>,
@@ -384,13 +395,14 @@ impl<'a> PoolKeyCursor<'a> {
.map(|config| config.score_top_n)
.unwrap_or(u64::from(aether_dispatch_core::DEFAULT_POOL_PAGE_SIZE))
.clamp(1, u64::from(u32::MAX)) as u32;
let max_scanned_keys = pool_config
let configured_max_scanned_keys = pool_config
.as_ref()
.map(|config| config.score_fallback_scan_limit)
.unwrap_or(u64::from(aether_dispatch_core::DEFAULT_POOL_MAX_SCAN))
.clamp(1, u64::from(u32::MAX)) as u32;
let window_config = crate::dispatch::pool::default_pool_window_config().normalized();
let max_scanned_keys = max_scanned_keys.min(window_config.max_scan);
let max_scanned_keys = configured_max_scanned_keys.min(window_config.max_scan);
let absolute_max_scanned_keys = configured_max_scanned_keys.max(max_scanned_keys);
Self {
state,
group,
@@ -403,9 +415,11 @@ impl<'a> PoolKeyCursor<'a> {
pool_key_order,
next_offset: 0,
scanned_keys: 0,
budget_scanned_keys: 0,
window_size: window_config.window_size,
page_size: window_config.page_size,
max_scanned_keys: max_scanned_keys.max(window_config.window_size),
absolute_max_scanned_keys: absolute_max_scanned_keys.max(window_config.window_size),
score_top_n,
score_phase_loaded: false,
skip_reason_counts: BTreeMap::new(),
@@ -477,6 +491,7 @@ impl<'a> PoolKeyCursor<'a> {
extra_data: Some(serde_json::json!({
"pool_group_exhaustion": {
"scanned_keys": self.scanned_keys,
"budget_scanned_keys": self.budget_scanned_keys,
"skip_reason_counts": skip_reason_counts,
}
})),
@@ -495,6 +510,9 @@ impl<'a> PoolKeyCursor<'a> {
endpoint_id = %self.group.candidate.endpoint_id,
model_id = %self.group.candidate.model_id,
scanned_keys = self.scanned_keys,
budget_scanned_keys = self.budget_scanned_keys,
max_scanned_keys = self.max_scanned_keys,
absolute_max_scanned_keys = self.absolute_max_scanned_keys,
skip_reason_counts = ?self.skip_reason_counts,
"gateway pool scheduler exhausted pool group without a schedulable key"
);
@@ -539,13 +557,16 @@ impl<'a> PoolKeyCursor<'a> {
}
}
if self.scanned_keys >= self.max_scanned_keys {
if self.budget_scanned_keys >= self.max_scanned_keys
|| self.scanned_keys >= self.absolute_max_scanned_keys
{
return None;
}
let limit = self
.page_size
.min(self.max_scanned_keys - self.scanned_keys);
.min(self.max_scanned_keys - self.budget_scanned_keys)
.min(self.absolute_max_scanned_keys - self.scanned_keys);
let query = StoredPoolKeyCandidateRowsQuery {
api_format: self.group.candidate.endpoint_api_format.clone(),
provider_id: self.group.candidate.provider_id.clone(),
@@ -584,11 +605,21 @@ impl<'a> PoolKeyCursor<'a> {
}
self.scanned_keys += rows.len() as u32;
self.budget_scanned_keys += rows.len() as u32;
self.next_offset = self.next_offset.saturating_add(rows.len() as u32);
Some(self.build_page_eligible_candidates(rows).await)
}
async fn next_score_candidates(&mut self) -> Option<Vec<EligibleLocalExecutionCandidate>> {
if self.scanned_keys >= self.absolute_max_scanned_keys {
return None;
}
let limit = self
.score_top_n
.min(self.absolute_max_scanned_keys - self.scanned_keys);
if limit == 0 {
return None;
}
let scope = provider_key_pool_score_scope();
let query = ListRankedPoolMembersQuery {
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
@@ -599,7 +630,7 @@ impl<'a> PoolKeyCursor<'a> {
hard_states: vec![PoolMemberHardState::Available, PoolMemberHardState::Unknown],
probe_statuses: None,
offset: 0,
limit: self.score_top_n as usize,
limit: limit as usize,
};
let scores = match self.state.app().data.list_ranked_pool_members(&query).await {
Ok(scores) => scores,
@@ -621,7 +652,7 @@ impl<'a> PoolKeyCursor<'a> {
return None;
}
self.record_score_schedule_interest(&scores).await;
self.spawn_score_schedule_interest_recording(&scores);
let key_ids = scores
.iter()
@@ -658,6 +689,7 @@ impl<'a> PoolKeyCursor<'a> {
}
};
self.scanned_keys = self.scanned_keys.saturating_add(scores.len() as u32);
self.budget_scanned_keys = self.budget_scanned_keys.saturating_add(scores.len() as u32);
Some(self.build_page_eligible_candidates(rows).await)
}
@@ -849,60 +881,97 @@ impl<'a> PoolKeyCursor<'a> {
);
}
async fn record_score_schedule_interest(&self, scores: &[StoredPoolMemberScore]) {
if scores.is_empty() {
fn spawn_score_schedule_interest_recording(&self, scores: &[StoredPoolMemberScore]) {
if scores.is_empty() || !self.state.app().data.has_pool_score_writer() {
return;
}
let scheduled_at = current_unix_ms() / 1000;
let mut failed = 0usize;
for score in scores {
let identity = PoolMemberIdentity {
pool_kind: score.pool_kind.clone(),
pool_id: score.pool_id.clone(),
member_kind: score.member_kind.clone(),
member_id: score.member_id.clone(),
};
let scope = PoolScoreScope {
capability: score.capability.clone(),
scope_kind: score.scope_kind.clone(),
scope_id: score.scope_id.clone(),
};
let result = self
.state
.app()
.data
.record_pool_member_schedule_feedback(PoolMemberScheduleFeedback {
identity,
scope: Some(scope),
scheduled_at,
succeeded: None,
hard_state: None,
score_delta: None,
score_reason_patch: Some(serde_json::json!({
"last_schedule_interest": {
"provider_id": self.group.candidate.provider_id.as_str(),
"endpoint_id": self.group.candidate.endpoint_id.as_str(),
"model_id": self.group.candidate.model_id.as_str()
}
})),
let provider_id = self.group.candidate.provider_id.clone();
let endpoint_id = self.group.candidate.endpoint_id.clone();
let model_id = self.group.candidate.model_id.clone();
let feedback = scores
.iter()
.filter(|score| {
score.last_scheduled_at.is_none_or(|last_scheduled_at| {
scheduled_at.saturating_sub(last_scheduled_at)
>= POOL_SCORE_SCHEDULE_INTEREST_MIN_INTERVAL_SECS
})
.await;
if result.is_err() {
failed += 1;
}
})
.take(POOL_SCORE_SCHEDULE_INTEREST_MAX_PER_BATCH)
.map(|score| PoolMemberScheduleFeedback {
identity: PoolMemberIdentity {
pool_kind: score.pool_kind.clone(),
pool_id: score.pool_id.clone(),
member_kind: score.member_kind.clone(),
member_id: score.member_id.clone(),
},
scope: Some(PoolScoreScope {
capability: score.capability.clone(),
scope_kind: score.scope_kind.clone(),
scope_id: score.scope_id.clone(),
}),
scheduled_at,
succeeded: None,
hard_state: None,
score_delta: None,
score_reason_patch: Some(serde_json::json!({
"last_schedule_interest": {
"provider_id": provider_id.as_str(),
"endpoint_id": endpoint_id.as_str(),
"model_id": model_id.as_str()
}
})),
})
.collect::<Vec<_>>();
let score_count = feedback.len();
if feedback.is_empty() {
return;
}
if failed > 0 {
warn!(
event_name = "pool_group_score_interest_update_failed",
let Ok(permit) = POOL_SCORE_SCHEDULE_INTEREST_SEMAPHORE
.clone()
.try_acquire_owned()
else {
debug!(
event_name = "pool_group_score_interest_dropped",
log_type = "event",
provider_id = %self.group.candidate.provider_id,
endpoint_id = %self.group.candidate.endpoint_id,
model_id = %self.group.candidate.model_id,
failed_count = failed,
score_count = scores.len(),
"gateway pool scheduler failed to record some pool score schedule interests"
"gateway pool scheduler dropped score schedule interest because the background writer is saturated"
);
}
return;
};
let app = self.state.app().clone();
tokio::spawn(async move {
let _permit = permit;
let mut failed = 0usize;
for feedback in feedback {
let result = app
.data
.record_pool_member_schedule_feedback(feedback)
.await;
if result.is_err() {
failed += 1;
}
}
if failed > 0 {
warn!(
event_name = "pool_group_score_interest_update_failed",
log_type = "event",
provider_id = %provider_id,
endpoint_id = %endpoint_id,
model_id = %model_id,
failed_count = failed,
score_count,
"gateway pool scheduler failed to record some pool score schedule interests"
);
}
});
}
async fn build_page_eligible_candidates(
@@ -964,9 +1033,20 @@ impl<'a> PoolKeyCursor<'a> {
for skipped_candidate in skipped_candidates {
self.record_skip_reason(skipped_candidate.skip_reason);
}
let prefiltered_count = skipped_candidates
.iter()
.filter(|candidate| pool_skip_reason_releases_scan_budget(candidate.skip_reason))
.count();
self.budget_scanned_keys = self
.budget_scanned_keys
.saturating_sub(u32::try_from(prefiltered_count).unwrap_or(u32::MAX));
}
}
fn pool_skip_reason_releases_scan_budget(skip_reason: &str) -> bool {
skip_reason == POOL_ACCOUNT_EXHAUSTED_SKIP_REASON
}
fn pool_candidate_transport_policy_facts(
candidate: &aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate,
) -> CandidateTransportPolicyFacts<'_> {
@@ -1098,12 +1178,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 +1576,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 +3001,324 @@ mod tests {
}));
}
#[tokio::test]
async fn pool_key_cursor_does_not_spend_effective_scan_budget_on_exhausted_accounts() {
let provider_config = Some(json!({
"pool_advanced": {
"skip_exhausted_accounts": true
}
}));
let (provider, endpoint, mut keys, rows) = large_pool_fixture(700, provider_config.clone());
for key in keys.iter_mut().take(600) {
key.status_snapshot = Some(json!({
"quota": {
"provider_type": "openai",
"exhausted": true,
"usage_ratio": 1.0,
"windows": [
{
"code": "daily",
"used_ratio": 1.0,
"remaining_ratio": 0.0
}
]
}
}));
}
let data_state =
GatewayDataState::with_provider_catalog_and_minimal_candidate_selection_for_tests(
Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
keys,
)),
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);
let group = sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"pool-group",
10,
provider_config,
);
let mut cursor = PoolKeyCursor::new(PlannerAppState::new(&app), group, None, None, None);
assert_eq!(
cursor.max_scanned_keys,
aether_dispatch_core::DEFAULT_POOL_MAX_SCAN
);
assert!(
cursor.absolute_max_scanned_keys > cursor.max_scanned_keys,
"pool config scan limit should be retained as the absolute cap"
);
let candidate = cursor
.next_key()
.await
.expect("cursor should scan past exhausted accounts within the absolute cap");
let key_index = candidate
.candidate
.key_id
.strip_prefix("key-")
.and_then(|value| value.parse::<usize>().ok())
.expect("fixture key id should contain a numeric suffix");
assert!(
key_index >= 600,
"cursor should not return one of the exhausted leading keys"
);
assert_eq!(candidate.orchestration.pool_key_index, Some(0));
assert_eq!(cursor.scanned_keys, 640);
assert_eq!(cursor.budget_scanned_keys, 40);
assert_eq!(
cursor
.skip_reason_counts
.get(aether_pool_core::POOL_ACCOUNT_EXHAUSTED_SKIP_REASON),
Some(&600)
);
}
#[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 +3766,210 @@ 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,
upstream_metadata: 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 {