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

refactor: 大规模模块拆分与代码精简，新增 ai-pipeline/data-contracts 独立 crate

Browse Source 
- 新增 aether-ai-pipeline 和 aether-data-contracts crate，将 pipeline 逻辑与数据契约从 gateway 中解耦
- 重构 admin handlers：拆分单体模块为 auth/billing/endpoint/features/model/observability/provider/system 等独立子模块
- 合并 chat/cli 重复代码路径：精简 conversion、finalize、planner 中的 sync/chat/cli 分支
- 重构 scheduler/executor/data 层，引入 facade 模式降低模块间耦合
- 移除冗余的 intent 模块，将 plan_fallback/policy/stream_path/sync_path 迁移至 executor
- 前端适配：调整 admin API 调用和 provider 模型测试对话框
This commit is contained in:
fawney19
2026-04-07 02:50:19 +08:00
parent 763ff03a7b
commit 5d96d6673b
732 changed files with 28593 additions and 20666 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
apps/aether-gateway/src/scheduler/affinity.rs Normal file
View File

@@ -0,0 +1,23 @@
use std::time::Duration;
use aether_scheduler_core::{
build_scheduler_affinity_cache_key_for_api_key_id, SchedulerAffinityTarget,
};
use super::state::SchedulerRuntimeState;
pub(crate) const SCHEDULER_AFFINITY_TTL: Duration = Duration::from_secs(300);
pub(crate) fn read_cached_scheduler_affinity_target(
state: &(impl SchedulerRuntimeState + ?Sized),
api_key_id: &str,
api_format: &str,
global_model_name: &str,
) -> Option<SchedulerAffinityTarget> {
let cache_key = build_scheduler_affinity_cache_key_for_api_key_id(
api_key_id,
api_format,
global_model_name,
)?;
state.read_cached_scheduler_affinity_target(&cache_key, SCHEDULER_AFFINITY_TTL)
}

7
apps/aether-gateway/src/scheduler/candidate/affinity.rs
View File

@@ -4,10 +4,11 @@ pub(super) use aether_scheduler_core::{
};
use crate::data::auth::GatewayAuthApiKeySnapshot;
use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
use super::{
GatewayMinimalCandidateSelectionCandidate, SchedulerRuntimeState,
SCHEDULER_AFFINITY_MAX_ENTRIES, SCHEDULER_AFFINITY_TTL,
SchedulerMinimalCandidateSelectionCandidate, SchedulerRuntimeState,
SCHEDULER_AFFINITY_MAX_ENTRIES,
};
pub(super) fn build_scheduler_affinity_cache_key(
@@ -25,7 +26,7 @@ pub(super) fn build_scheduler_affinity_cache_key(
pub(super) fn remember_scheduler_affinity(
affinity_cache_key: Option<&str>,
state: &(impl SchedulerRuntimeState + ?Sized),
candidate: &GatewayMinimalCandidateSelectionCandidate,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) {
let Some(cache_key) = affinity_cache_key else {
return;

242
apps/aether-gateway/src/scheduler/candidate/mod.rs
View File

@@ -1,134 +1,52 @@
use self::affinity::{
build_scheduler_affinity_cache_key, build_scheduler_affinity_cache_key_for_api_key_id,
candidate_affinity_hash, candidate_key, remember_scheduler_affinity,
};
use self::model::{
auth_snapshot_allows_api_format, auth_snapshot_constraints, candidate_model_names,
candidate_supports_required_capability, matches_model_mapping, normalize_api_format,
read_requested_model_rows, resolve_provider_model_name, resolve_requested_global_model_name,
select_provider_model_name,
};
use self::selection::{
is_candidate_selectable, read_provider_concurrent_limits, read_provider_key_rpm_states,
reorder_candidates_by_scheduler_health, should_skip_provider_quota,
};
pub(crate) use self::state::{
SchedulerCandidateSelectionRowSource, SchedulerRuntimeState,
};
use self::affinity::candidate_affinity_hash;
use self::selection::collect_selectable_candidates;
use super::state::SchedulerRuntimeState;
mod affinity;
mod model;
mod runtime;
mod selection;
mod state;
#[cfg(test)]
mod tests;
use std::collections::{BTreeMap, BTreeSet};
use std::time::Duration;
use aether_data::repository::candidate_selection::{
use aether_data_contracts::repository::candidate_selection::{
StoredMinimalCandidateSelectionRow, StoredProviderModelMapping,
};
use aether_data::repository::candidates::StoredRequestCandidate;
use aether_data::repository::provider_catalog::{
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_data::repository::quota::StoredProviderQuotaSnapshot;
use aether_data::DataLayerError;
use aether_data_contracts::repository::quota::StoredProviderQuotaSnapshot;
use aether_scheduler_core::{
auth_api_key_concurrency_limit_reached, build_minimal_candidate_selection,
collect_global_model_names_for_required_capability, collect_selectable_candidates_from_keys,
SchedulerAffinityTarget, SchedulerMinimalCandidateSelectionCandidate,
candidate_model_names, candidate_supports_required_capability, matches_model_mapping,
normalize_api_format, resolve_provider_model_name, select_provider_model_name,
SchedulerMinimalCandidateSelectionCandidate,
};
use aether_wallet::{ProviderBillingType, ProviderQuotaSnapshot};
use regex::Regex;
use sha2::{Digest, Sha256};
use std::collections::BTreeMap;
use crate::data::auth::GatewayAuthApiKeySnapshot;
use crate::{AppState, GatewayError};
use crate::data::candidate_selection::{
read_global_model_names_for_required_capability, MinimalCandidateSelectionRowSource,
};
use crate::GatewayError;
const SCHEDULER_AFFINITY_TTL: Duration = Duration::from_secs(300);
#[cfg_attr(not(test), allow(dead_code))]
const SCHEDULER_AFFINITY_MAX_ENTRIES: usize = 10_000;
pub(crate) use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate as GatewayMinimalCandidateSelectionCandidate;
#[allow(dead_code)]
pub(crate) async fn read_minimal_candidate_selection(
state: &(impl SchedulerCandidateSelectionRowSource + Sync),
api_format: &str,
requested_model_name: &str,
require_streaming: bool,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
) -> Result<Vec<GatewayMinimalCandidateSelectionCandidate>, DataLayerError> {
let normalized_api_format = normalize_api_format(api_format);
if normalized_api_format.is_empty() {
return Ok(Vec::new());
}
if !auth_snapshot_allows_api_format(auth_snapshot, &normalized_api_format) {
return Ok(Vec::new());
}
let Some((resolved_global_model_name, rows)) =
read_requested_model_rows(state, &normalized_api_format, requested_model_name).await?
else {
return Ok(Vec::new());
};
let auth_constraints = auth_snapshot.map(auth_snapshot_constraints);
let affinity_key = auth_snapshot
.map(|snapshot| snapshot.api_key_id.trim())
.filter(|value| !value.is_empty());
build_minimal_candidate_selection(
rows,
&normalized_api_format,
requested_model_name,
resolved_global_model_name.as_str(),
require_streaming,
auth_constraints.as_ref(),
affinity_key,
)
}
#[cfg_attr(not(test), allow(dead_code))]
pub(crate) async fn select_minimal_candidate(
state: &AppState,
api_format: &str,
global_model_name: &str,
require_streaming: bool,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
now_unix_secs: u64,
) -> Result<Option<GatewayMinimalCandidateSelectionCandidate>, GatewayError> {
let affinity_cache_key =
build_scheduler_affinity_cache_key(auth_snapshot, api_format, global_model_name);
let selected = collect_selectable_candidates(
state,
api_format,
global_model_name,
require_streaming,
auth_snapshot,
now_unix_secs,
)
.await?
.into_iter()
.next();
if let Some(candidate) = selected.as_ref() {
remember_scheduler_affinity(affinity_cache_key.as_deref(), state, candidate);
}
Ok(selected)
}
pub(crate) async fn list_selectable_candidates(
state: &AppState,
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
runtime_state: &impl SchedulerRuntimeState,
api_format: &str,
global_model_name: &str,
require_streaming: bool,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
now_unix_secs: u64,
) -> Result<Vec<GatewayMinimalCandidateSelectionCandidate>, GatewayError> {
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
collect_selectable_candidates(
state,
selection_row_source,
runtime_state,
api_format,
global_model_name,
require_streaming,
@@ -139,39 +57,33 @@ pub(crate) async fn list_selectable_candidates(
}
pub(crate) async fn list_selectable_candidates_for_required_capability_without_requested_model(
state: &AppState,
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
runtime_state: &impl SchedulerRuntimeState,
candidate_api_format: &str,
required_capability: &str,
require_streaming: bool,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
now_unix_secs: u64,
) -> Result<Vec<GatewayMinimalCandidateSelectionCandidate>, GatewayError> {
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
let normalized_api_format = normalize_api_format(candidate_api_format);
let required_capability = required_capability.trim();
if normalized_api_format.is_empty() || required_capability.is_empty() {
if normalized_api_format.is_empty() {
return Ok(Vec::new());
}
if !auth_snapshot_allows_api_format(auth_snapshot, &normalized_api_format) {
return Ok(Vec::new());
}
let rows = state
.read_minimal_candidate_selection_rows_for_api_format(&normalized_api_format)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let auth_constraints = auth_snapshot.map(auth_snapshot_constraints);
let model_names = collect_global_model_names_for_required_capability(
rows,
let model_names = read_global_model_names_for_required_capability(
selection_row_source,
&normalized_api_format,
required_capability,
require_streaming,
auth_constraints.as_ref(),
);
auth_snapshot,
)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
for global_model_name in model_names {
let candidates = list_selectable_candidates(
state,
selection_row_source,
runtime_state,
&normalized_api_format,
&global_model_name,
require_streaming,
@@ -192,93 +104,3 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
Ok(Vec::new())
}
pub(crate) fn read_cached_scheduler_affinity_target(
state: &AppState,
api_key_id: &str,
api_format: &str,
global_model_name: &str,
) -> Option<SchedulerAffinityTarget> {
let cache_key = build_scheduler_affinity_cache_key_for_api_key_id(
api_key_id,
api_format,
global_model_name,
)?;
state.read_cached_scheduler_affinity_target(&cache_key, SCHEDULER_AFFINITY_TTL)
}
async fn collect_selectable_candidates(
state: &AppState,
api_format: &str,
global_model_name: &str,
require_streaming: bool,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
now_unix_secs: u64,
) -> Result<Vec<GatewayMinimalCandidateSelectionCandidate>, GatewayError> {
let mut candidates = state
.read_minimal_candidate_selection(
api_format,
global_model_name,
require_streaming,
auth_snapshot,
)
.await?;
let recent_candidates = state.read_recent_request_candidates(128).await?;
let provider_concurrent_limits = read_provider_concurrent_limits(state, &candidates).await?;
let provider_key_rpm_states = read_provider_key_rpm_states(state, &candidates).await?;
reorder_candidates_by_scheduler_health(
&mut candidates,
&provider_key_rpm_states,
auth_snapshot,
);
let affinity_cache_key =
build_scheduler_affinity_cache_key(auth_snapshot, api_format, global_model_name);
let cached_affinity_target = affinity_cache_key.as_deref().and_then(|cache_key| {
state.read_cached_scheduler_affinity_target(cache_key, SCHEDULER_AFFINITY_TTL)
});
if let Some((api_key_id, limit)) = auth_snapshot.and_then(|snapshot| {
usize::try_from(snapshot.api_key_concurrent_limit?)
.ok()
.and_then(|limit| {
if limit == 0 {
return None;
}
Some((snapshot.api_key_id.as_str(), limit))
})
}) {
if auth_api_key_concurrency_limit_reached(
&recent_candidates,
now_unix_secs,
api_key_id,
limit,
) {
return Ok(Vec::new());
}
}
let mut selected_keys = BTreeSet::new();
for candidate in &candidates {
if !is_candidate_selectable(
candidate,
&recent_candidates,
&provider_concurrent_limits,
&provider_key_rpm_states,
now_unix_secs,
cached_affinity_target.as_ref(),
state,
)
.await?
{
continue;
}
selected_keys.insert(candidate_key(candidate));
}
Ok(collect_selectable_candidates_from_keys(
candidates,
&selected_keys,
cached_affinity_target.as_ref(),
))
}

96
apps/aether-gateway/src/scheduler/candidate/model.rs
View File

@@ -1,96 +0,0 @@
use aether_data::repository::candidate_selection::StoredMinimalCandidateSelectionRow;
use aether_data::DataLayerError;
pub(super) use aether_scheduler_core::{
auth_constraints_allow_api_format, auth_constraints_allow_model,
auth_constraints_allow_provider, candidate_model_names, candidate_supports_required_capability,
extract_global_priority_for_format, matches_model_mapping, normalize_api_format,
resolve_provider_model_name, resolve_requested_global_model_name,
row_supports_required_capability, select_provider_model_name, SchedulerAuthConstraints,
};
use crate::data::auth::GatewayAuthApiKeySnapshot;
use super::state::SchedulerCandidateSelectionRowSource;
use super::GatewayMinimalCandidateSelectionCandidate;
pub(super) fn auth_snapshot_allows_provider(
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
provider_id: &str,
provider_name: &str,
) -> bool {
auth_constraints_allow_provider(
auth_snapshot.map(auth_snapshot_constraints).as_ref(),
provider_id,
provider_name,
)
}
pub(super) fn auth_snapshot_allows_api_format(
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
api_format: &str,
) -> bool {
auth_constraints_allow_api_format(
auth_snapshot.map(auth_snapshot_constraints).as_ref(),
api_format,
)
}
pub(super) fn auth_snapshot_allows_model(
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
requested_model_name: &str,
resolved_global_model_name: &str,
) -> bool {
auth_constraints_allow_model(
auth_snapshot.map(auth_snapshot_constraints).as_ref(),
requested_model_name,
resolved_global_model_name,
)
}
pub(super) async fn read_requested_model_rows(
state: &(impl SchedulerCandidateSelectionRowSource + Sync),
api_format: &str,
requested_model_name: &str,
) -> Result<Option<(String, Vec<StoredMinimalCandidateSelectionRow>)>, DataLayerError> {
let exact_rows = state
.read_minimal_candidate_selection_rows_for_api_format_and_global_model(
api_format,
requested_model_name,
)
.await?;
if !exact_rows.is_empty() {
return Ok(Some((requested_model_name.to_string(), exact_rows)));
}
let rows = state
.read_minimal_candidate_selection_rows_for_api_format(api_format)
.await?;
let Some(resolved_global_model_name) =
resolve_requested_global_model_name(&rows, requested_model_name, api_format)
else {
return Ok(None);
};
Ok(Some((
resolved_global_model_name.clone(),
rows.into_iter()
.filter(|row| row.global_model_name == resolved_global_model_name)
.collect(),
)))
}
pub(super) fn auth_snapshot_constraints(
snapshot: &GatewayAuthApiKeySnapshot,
) -> SchedulerAuthConstraints {
SchedulerAuthConstraints {
allowed_providers: snapshot
.effective_allowed_providers()
.map(|items| items.to_vec()),
allowed_api_formats: snapshot
.effective_allowed_api_formats()
.map(|items| items.to_vec()),
allowed_models: snapshot
.effective_allowed_models()
.map(|items| items.to_vec()),
}
}

182
apps/aether-gateway/src/scheduler/candidate/runtime.rs Normal file
View File

@@ -0,0 +1,182 @@
use std::collections::{BTreeMap, BTreeSet};
use aether_data_contracts::repository::candidates::StoredRequestCandidate;
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use aether_scheduler_core::{
auth_api_key_concurrency_limit_reached, build_provider_concurrent_limit_map,
candidate_is_selectable_with_runtime_state, SchedulerAffinityTarget,
};
use crate::data::auth::GatewayAuthApiKeySnapshot;
use crate::GatewayError;
use super::{SchedulerMinimalCandidateSelectionCandidate, SchedulerRuntimeState};
pub(super) use aether_scheduler_core::should_skip_provider_quota;
pub(super) struct CandidateRuntimeSelectionSnapshot {
pub(super) recent_candidates: Vec<StoredRequestCandidate>,
pub(super) provider_concurrent_limits: BTreeMap<String, usize>,
pub(super) provider_key_rpm_states: BTreeMap<String, StoredProviderCatalogKey>,
provider_quota_blocks_requests: BTreeMap<String, bool>,
provider_key_rpm_reset_ats: BTreeMap<String, Option<u64>>,
}
pub(super) async fn read_candidate_runtime_selection_snapshot(
state: &(impl SchedulerRuntimeState + ?Sized),
candidates: &[SchedulerMinimalCandidateSelectionCandidate],
now_unix_secs: u64,
) -> Result<CandidateRuntimeSelectionSnapshot, GatewayError> {
let recent_candidates = state.read_recent_request_candidates(128).await?;
let provider_concurrent_limits = read_provider_concurrent_limits(state, candidates).await?;
let provider_key_rpm_states = read_provider_key_rpm_states(state, candidates).await?;
let provider_quota_blocks_requests =
read_provider_quota_block_map(state, candidates, now_unix_secs).await?;
let provider_key_rpm_reset_ats =
read_provider_key_rpm_reset_at_map(state, candidates, now_unix_secs);
Ok(CandidateRuntimeSelectionSnapshot {
recent_candidates,
provider_concurrent_limits,
provider_key_rpm_states,
provider_quota_blocks_requests,
provider_key_rpm_reset_ats,
})
}
pub(super) fn auth_snapshot_concurrency_limit_reached(
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
snapshot: &CandidateRuntimeSelectionSnapshot,
now_unix_secs: u64,
) -> bool {
auth_snapshot
.and_then(|snapshot| {
usize::try_from(snapshot.api_key_concurrent_limit?)
.ok()
.and_then(|limit| {
if limit == 0 {
return None;
}
Some((snapshot.api_key_id.as_str(), limit))
})
})
.is_some_and(|(api_key_id, limit)| {
auth_api_key_concurrency_limit_reached(
&snapshot.recent_candidates,
now_unix_secs,
api_key_id,
limit,
)
})
}
pub(super) fn is_candidate_selectable(
candidate: &SchedulerMinimalCandidateSelectionCandidate,
snapshot: &CandidateRuntimeSelectionSnapshot,
now_unix_secs: u64,
cached_affinity_target: Option<&SchedulerAffinityTarget>,
) -> bool {
let provider_quota_blocks_requests = snapshot
.provider_quota_blocks_requests
.get(candidate.provider_id.as_str())
.copied()
.unwrap_or(false);
let rpm_reset_at = snapshot
.provider_key_rpm_reset_ats
.get(candidate.key_id.as_str())
.copied()
.flatten();
candidate_is_selectable_with_runtime_state(
candidate,
&snapshot.recent_candidates,
&snapshot.provider_concurrent_limits,
&snapshot.provider_key_rpm_states,
now_unix_secs,
cached_affinity_target,
provider_quota_blocks_requests,
rpm_reset_at,
)
}
pub(super) async fn read_provider_concurrent_limits(
state: &(impl SchedulerRuntimeState + ?Sized),
candidates: &[SchedulerMinimalCandidateSelectionCandidate],
) -> Result<BTreeMap<String, usize>, GatewayError> {
let provider_ids = candidates
.iter()
.map(|candidate| candidate.provider_id.clone())
.collect::<BTreeSet<_>>()
.into_iter()
.collect::<Vec<_>>();
if provider_ids.is_empty() {
return Ok(BTreeMap::new());
}
let providers = state
.read_provider_catalog_providers_by_ids(&provider_ids)
.await?;
Ok(build_provider_concurrent_limit_map(providers))
}
pub(super) async fn read_provider_key_rpm_states(
state: &(impl SchedulerRuntimeState + ?Sized),
candidates: &[SchedulerMinimalCandidateSelectionCandidate],
) -> Result<BTreeMap<String, StoredProviderCatalogKey>, GatewayError> {
let key_ids = candidates
.iter()
.map(|candidate| candidate.key_id.clone())
.collect::<BTreeSet<_>>()
.into_iter()
.collect::<Vec<_>>();
if key_ids.is_empty() {
return Ok(BTreeMap::new());
}
let keys = state.read_provider_catalog_keys_by_ids(&key_ids).await?;
Ok(keys
.into_iter()
.map(|key| (key.id.clone(), key))
.collect::<BTreeMap<_, _>>())
}
async fn read_provider_quota_block_map(
state: &(impl SchedulerRuntimeState + ?Sized),
candidates: &[SchedulerMinimalCandidateSelectionCandidate],
now_unix_secs: u64,
) -> Result<BTreeMap<String, bool>, GatewayError> {
let provider_ids = candidates
.iter()
.map(|candidate| candidate.provider_id.clone())
.collect::<BTreeSet<_>>()
.into_iter()
.collect::<Vec<_>>();
let mut quota_blocks = BTreeMap::new();
for provider_id in provider_ids {
let blocks_requests = state
.read_provider_quota_snapshot(&provider_id)
.await?
.as_ref()
.is_some_and(|quota| should_skip_provider_quota(quota, now_unix_secs));
quota_blocks.insert(provider_id, blocks_requests);
}
Ok(quota_blocks)
}
fn read_provider_key_rpm_reset_at_map(
state: &(impl SchedulerRuntimeState + ?Sized),
candidates: &[SchedulerMinimalCandidateSelectionCandidate],
now_unix_secs: u64,
) -> BTreeMap<String, Option<u64>> {
candidates
.iter()
.map(|candidate| {
(
candidate.key_id.clone(),
state.provider_key_rpm_reset_at(candidate.key_id.as_str(), now_unix_secs),
)
})
.collect::<BTreeMap<_, _>>()
}

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

@@ -1,20 +1,29 @@
use std::collections::{BTreeMap, BTreeSet};
use aether_data::repository::candidates::StoredRequestCandidate;
use aether_data::repository::provider_catalog::StoredProviderCatalogKey;
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use aether_scheduler_core::{
build_provider_concurrent_limit_map, candidate_is_selectable_with_runtime_state,
collect_selectable_candidates_from_keys,
reorder_candidates_by_scheduler_health as reorder_candidates_by_scheduler_health_in_core,
SchedulerAffinityTarget,
};
use crate::data::auth::GatewayAuthApiKeySnapshot;
use crate::data::candidate_selection::{
read_minimal_candidate_selection, MinimalCandidateSelectionRowSource,
};
use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
use crate::GatewayError;
use super::{GatewayMinimalCandidateSelectionCandidate, SchedulerRuntimeState};
use super::affinity::{
build_scheduler_affinity_cache_key, candidate_key, remember_scheduler_affinity,
};
use super::runtime::{
auth_snapshot_concurrency_limit_reached, is_candidate_selectable,
read_candidate_runtime_selection_snapshot,
};
use super::{SchedulerMinimalCandidateSelectionCandidate, SchedulerRuntimeState};
pub(super) fn reorder_candidates_by_scheduler_health(
candidates: &mut [GatewayMinimalCandidateSelectionCandidate],
candidates: &mut [SchedulerMinimalCandidateSelectionCandidate],
provider_key_rpm_states: &BTreeMap<String, StoredProviderCatalogKey>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
) {
@@ -28,73 +37,89 @@ pub(super) fn reorder_candidates_by_scheduler_health(
);
}
pub(super) use aether_scheduler_core::should_skip_provider_quota;
pub(super) async fn is_candidate_selectable(
candidate: &GatewayMinimalCandidateSelectionCandidate,
recent_candidates: &[StoredRequestCandidate],
provider_concurrent_limits: &BTreeMap<String, usize>,
provider_key_rpm_states: &BTreeMap<String, StoredProviderCatalogKey>,
#[cfg_attr(not(test), allow(dead_code))]
pub(super) async fn select_minimal_candidate(
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
runtime_state: &impl SchedulerRuntimeState,
api_format: &str,
global_model_name: &str,
require_streaming: bool,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
now_unix_secs: u64,
cached_affinity_target: Option<&SchedulerAffinityTarget>,
state: &(impl SchedulerRuntimeState + ?Sized),
) -> Result<bool, GatewayError> {
let provider_quota_blocks_requests = state
.read_provider_quota_snapshot(&candidate.provider_id)
.await?
.as_ref()
.is_some_and(|quota| should_skip_provider_quota(quota, now_unix_secs));
let rpm_reset_at = state.provider_key_rpm_reset_at(candidate.key_id.as_str(), now_unix_secs);
Ok(candidate_is_selectable_with_runtime_state(
candidate,
recent_candidates,
provider_concurrent_limits,
provider_key_rpm_states,
) -> Result<Option<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
let affinity_cache_key =
build_scheduler_affinity_cache_key(auth_snapshot, api_format, global_model_name);
let selected = collect_selectable_candidates(
selection_row_source,
runtime_state,
api_format,
global_model_name,
require_streaming,
auth_snapshot,
now_unix_secs,
cached_affinity_target,
provider_quota_blocks_requests,
rpm_reset_at,
)
.await?
.into_iter()
.next();
if let Some(candidate) = selected.as_ref() {
remember_scheduler_affinity(affinity_cache_key.as_deref(), runtime_state, candidate);
}
Ok(selected)
}
pub(super) async fn collect_selectable_candidates(
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
runtime_state: &impl SchedulerRuntimeState,
api_format: &str,
global_model_name: &str,
require_streaming: bool,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
now_unix_secs: u64,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
let mut candidates = read_minimal_candidate_selection(
selection_row_source,
api_format,
global_model_name,
require_streaming,
auth_snapshot,
)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let runtime_snapshot =
read_candidate_runtime_selection_snapshot(runtime_state, &candidates, now_unix_secs)
.await?;
reorder_candidates_by_scheduler_health(
&mut candidates,
&runtime_snapshot.provider_key_rpm_states,
auth_snapshot,
);
let affinity_cache_key =
build_scheduler_affinity_cache_key(auth_snapshot, api_format, global_model_name);
let cached_affinity_target = affinity_cache_key.as_deref().and_then(|cache_key| {
runtime_state.read_cached_scheduler_affinity_target(cache_key, SCHEDULER_AFFINITY_TTL)
});
if auth_snapshot_concurrency_limit_reached(auth_snapshot, &runtime_snapshot, now_unix_secs) {
return Ok(Vec::new());
}
let mut selected_keys = BTreeSet::new();
for candidate in &candidates {
if !is_candidate_selectable(
candidate,
&runtime_snapshot,
now_unix_secs,
cached_affinity_target.as_ref(),
) {
continue;
}
selected_keys.insert(candidate_key(candidate));
}
Ok(collect_selectable_candidates_from_keys(
candidates,
&selected_keys,
cached_affinity_target.as_ref(),
))
}
pub(super) async fn read_provider_concurrent_limits(
state: &(impl SchedulerRuntimeState + ?Sized),
candidates: &[GatewayMinimalCandidateSelectionCandidate],
) -> Result<BTreeMap<String, usize>, GatewayError> {
let provider_ids = candidates
.iter()
.map(|candidate| candidate.provider_id.clone())
.collect::<BTreeSet<_>>()
.into_iter()
.collect::<Vec<_>>();
if provider_ids.is_empty() {
return Ok(BTreeMap::new());
}
let providers = state
.read_provider_catalog_providers_by_ids(&provider_ids)
.await?;
Ok(build_provider_concurrent_limit_map(providers))
}
pub(super) async fn read_provider_key_rpm_states(
state: &(impl SchedulerRuntimeState + ?Sized),
candidates: &[GatewayMinimalCandidateSelectionCandidate],
) -> Result<BTreeMap<String, StoredProviderCatalogKey>, GatewayError> {
let key_ids = candidates
.iter()
.map(|candidate| candidate.key_id.clone())
.collect::<BTreeSet<_>>()
.into_iter()
.collect::<Vec<_>>();
if key_ids.is_empty() {
return Ok(BTreeMap::new());
}
let keys = state.read_provider_catalog_keys_by_ids(&key_ids).await?;
Ok(keys
.into_iter()
.map(|key| (key.id.clone(), key))
.collect::<BTreeMap<_, _>>())
}

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

File diff suppressed because it is too large Load Diff

246
apps/aether-gateway/src/scheduler/candidate/tests/affinity.rs Normal file
View File

@@ -0,0 +1,246 @@
use std::sync::Arc;
use std::time::Duration;
use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository;
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data::repository::quota::InMemoryProviderQuotaRepository;
use aether_data_contracts::repository::candidate_selection::StoredProviderModelMapping;
use aether_data_contracts::repository::candidates::{
RequestCandidateStatus, StoredRequestCandidate,
};
use crate::cache::SchedulerAffinityTarget;
use crate::data::candidate_selection::read_minimal_candidate_selection;
use crate::data::GatewayDataState;
use crate::AppState;
use super::super::affinity::{build_scheduler_affinity_cache_key, candidate_affinity_hash};
use super::super::selection::select_minimal_candidate as select_candidate;
use super::support::{sample_auth_snapshot, sample_key, sample_provider, sample_row};
#[tokio::test]
async fn same_priority_candidates_are_distributed_by_affinity_key() {
let mut first = sample_row();
first.provider_id = "provider-a".to_string();
first.provider_name = "openai-a".to_string();
first.endpoint_id = "endpoint-a".to_string();
first.key_id = "key-a".to_string();
first.key_name = "alpha".to_string();
first.provider_priority = 10;
first.key_internal_priority = 10;
first.key_global_priority_by_format = Some(serde_json::json!({"openai:chat": 1}));
first.model_provider_model_name = "gpt-4.1-a".to_string();
first.model_provider_model_mappings = Some(vec![StoredProviderModelMapping {
name: "gpt-4.1-a".to_string(),
priority: 1,
api_formats: Some(vec!["openai:chat".to_string()]),
}]);
let mut second = sample_row();
second.provider_id = "provider-b".to_string();
second.provider_name = "openai-b".to_string();
second.endpoint_id = "endpoint-b".to_string();
second.key_id = "key-b".to_string();
second.key_name = "beta".to_string();
second.provider_priority = 10;
second.key_internal_priority = 10;
second.key_global_priority_by_format = Some(serde_json::json!({"openai:chat": 1}));
second.model_provider_model_name = "gpt-4.1-b".to_string();
second.model_provider_model_mappings = Some(vec![StoredProviderModelMapping {
name: "gpt-4.1-b".to_string(),
priority: 1,
api_formats: Some(vec!["openai:chat".to_string()]),
}]);
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
first, second,
]));
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
let state = GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas);
let auth_snapshot = sample_auth_snapshot("affinity-key-1");
let selection = read_minimal_candidate_selection(
&state,
"openai:chat",
"gpt-4.1",
false,
Some(&auth_snapshot),
)
.await
.expect("selection should succeed");
assert_eq!(selection.len(), 2);
let left_hash = candidate_affinity_hash(&auth_snapshot.api_key_id, &selection[0]);
let right_hash = candidate_affinity_hash(&auth_snapshot.api_key_id, &selection[1]);
assert!(
left_hash <= right_hash,
"same-priority candidates should be ordered by affinity hash"
);
}
#[tokio::test]
async fn reuses_cached_scheduler_affinity_candidate_before_sorted_fallback() {
let mut first = sample_row();
first.provider_id = "provider-a".to_string();
first.provider_name = "openai-a".to_string();
first.endpoint_id = "endpoint-a".to_string();
first.key_id = "key-a".to_string();
first.key_name = "alpha".to_string();
first.key_global_priority_by_format = Some(serde_json::json!({"openai:chat": 1}));
let mut second = sample_row();
second.provider_id = "provider-b".to_string();
second.provider_name = "openai-b".to_string();
second.endpoint_id = "endpoint-b".to_string();
second.key_id = "key-b".to_string();
second.key_name = "beta".to_string();
second.key_global_priority_by_format = Some(serde_json::json!({"openai:chat": 2}));
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
first, second,
]));
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(
GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas),
);
let auth_snapshot = sample_auth_snapshot("affinity-key-1");
let cache_key =
build_scheduler_affinity_cache_key(Some(&auth_snapshot), "openai:chat", "gpt-4.1")
.expect("cache key should build");
state.scheduler_affinity_cache.insert(
cache_key,
SchedulerAffinityTarget {
provider_id: "provider-b".to_string(),
endpoint_id: "endpoint-b".to_string(),
key_id: "key-b".to_string(),
},
Duration::from_secs(300),
100,
);
let selected = select_candidate(
state.data.as_ref(),
&state,
"openai:chat",
"gpt-4.1",
false,
Some(&auth_snapshot),
100,
)
.await
.expect("selection should succeed")
.expect("candidate should exist");
assert_eq!(selected.provider_id, "provider-b");
assert_eq!(selected.key_id, "key-b");
}
#[tokio::test]
async fn cached_affinity_candidate_can_use_reserved_provider_key_rpm_capacity() {
let mut first = sample_row();
first.provider_id = "provider-a".to_string();
first.provider_name = "openai-a".to_string();
first.endpoint_id = "endpoint-a".to_string();
first.key_id = "key-a".to_string();
first.key_name = "alpha".to_string();
first.key_global_priority_by_format = Some(serde_json::json!({"openai:chat": 1}));
let mut second = sample_row();
second.provider_id = "provider-b".to_string();
second.provider_name = "openai-b".to_string();
second.endpoint_id = "endpoint-b".to_string();
second.key_id = "key-b".to_string();
second.key_name = "beta".to_string();
second.key_global_priority_by_format = Some(serde_json::json!({"openai:chat": 2}));
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
first, second,
]));
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![
sample_provider("provider-a", None),
sample_provider("provider-b", None),
],
Vec::new(),
vec![
sample_key("key-a", "provider-a", Some(10)),
sample_key("key-b", "provider-b", Some(10)),
],
));
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
StoredRequestCandidate::new(
"cand-1".to_string(),
"req-1".to_string(),
None,
Some("api-key-cached-user".to_string()),
None,
None,
0,
0,
Some("provider-a".to_string()),
Some("endpoint-a".to_string()),
Some("key-a".to_string()),
RequestCandidateStatus::Success,
None,
false,
Some(200),
None,
None,
Some(10),
Some(9),
None,
None,
95,
Some(95),
Some(96),
)
.expect("candidate should build"),
]));
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 auth_snapshot = sample_auth_snapshot("api-key-cached-user");
let cache_key =
build_scheduler_affinity_cache_key(Some(&auth_snapshot), "openai:chat", "gpt-4.1")
.expect("cache key should build");
state.scheduler_affinity_cache.insert(
cache_key,
SchedulerAffinityTarget {
provider_id: "provider-a".to_string(),
endpoint_id: "endpoint-a".to_string(),
key_id: "key-a".to_string(),
},
Duration::from_secs(300),
100,
);
let selected = select_candidate(
state.data.as_ref(),
&state,
"openai:chat",
"gpt-4.1",
false,
Some(&auth_snapshot),
100,
)
.await
.expect("selection should succeed")
.expect("candidate should exist");
assert_eq!(selected.provider_id, "provider-a");
assert_eq!(selected.key_id, "key-a");
}

4
apps/aether-gateway/src/scheduler/candidate/tests/mod.rs Normal file
View File

@@ -0,0 +1,4 @@
mod affinity;
mod model;
mod selection;
mod support;

180
apps/aether-gateway/src/scheduler/candidate/tests/model.rs Normal file
View File

@@ -0,0 +1,180 @@
use std::sync::Arc;
use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository;
use aether_data::repository::quota::InMemoryProviderQuotaRepository;
use aether_data_contracts::repository::candidate_selection::StoredProviderModelMapping;
use aether_scheduler_core::{
resolve_requested_global_model_name, SchedulerMinimalCandidateSelectionCandidate,
};
use crate::data::candidate_selection::read_minimal_candidate_selection;
use crate::data::GatewayDataState;
use super::super::{
candidate_model_names, matches_model_mapping, resolve_provider_model_name,
select_provider_model_name,
};
use super::support::{sample_auth_snapshot, sample_row};
#[test]
fn selects_provider_model_name_with_api_format_scope() {
let row = sample_row();
assert_eq!(
select_provider_model_name(&row, "openai:chat"),
"gpt-4.1-canary"
);
}
#[test]
fn candidate_model_names_keep_base_and_scoped_mappings() {
let row = sample_row();
let names = candidate_model_names(&row, "openai:chat");
assert!(names.contains("gpt-4.1-upstream"));
assert!(names.contains("gpt-4.1-canary"));
assert!(!names.contains("gpt-4.1-responses"));
}
#[test]
fn resolves_mapping_matched_model_from_key_allowed_models() {
let mut row = sample_row();
row.key_allowed_models = Some(vec!["gpt-4.1-canary".to_string()]);
let resolved = resolve_provider_model_name(&row, "gpt-4.1", "openai:chat")
.expect("candidate should resolve");
assert_eq!(resolved.0, "gpt-4.1-canary");
assert_eq!(resolved.1, Some("gpt-4.1-canary".to_string()));
}
#[test]
fn resolves_mapping_matched_model_from_global_regex_mapping() {
let mut row = sample_row();
row.key_allowed_models = Some(vec!["gpt-4.1-variant".to_string()]);
let resolved = resolve_provider_model_name(&row, "gpt-4.1", "openai:chat")
.expect("candidate should resolve");
assert_eq!(resolved.0, "gpt-4.1-variant");
assert_eq!(resolved.1, Some("gpt-4.1-variant".to_string()));
}
#[test]
fn invalid_regex_mapping_is_treated_as_non_match() {
assert!(!matches_model_mapping("(", "gpt-4.1-variant"));
}
#[test]
fn resolves_requested_global_model_from_provider_model_alias() {
let mut row = sample_row();
row.global_model_name = "gpt-5".to_string();
row.model_provider_model_name = "gpt-5.2".to_string();
row.model_provider_model_mappings = Some(vec![StoredProviderModelMapping {
name: "gpt-5.2".to_string(),
priority: 1,
api_formats: Some(vec!["openai:chat".to_string()]),
}]);
let resolved = resolve_requested_global_model_name(&[row], "gpt-5.2", "openai:chat");
assert_eq!(resolved.as_deref(), Some("gpt-5"));
}
#[test]
fn resolves_requested_global_model_from_global_regex_mapping() {
let mut row = sample_row();
row.global_model_name = "gpt-5".to_string();
row.global_model_mappings = Some(vec!["gpt-5(?:\\.\\d+)?".to_string()]);
let resolved = resolve_requested_global_model_name(&[row], "gpt-5.2", "openai:chat");
assert_eq!(resolved.as_deref(), Some("gpt-5"));
}
#[test]
fn scheduler_candidate_is_serializable() {
let candidate = SchedulerMinimalCandidateSelectionCandidate {
provider_id: "provider-1".to_string(),
provider_name: "OpenAI".to_string(),
provider_type: "custom".to_string(),
provider_priority: 10,
endpoint_id: "endpoint-1".to_string(),
endpoint_api_format: "openai:chat".to_string(),
key_id: "key-1".to_string(),
key_name: "prod".to_string(),
key_auth_type: "api_key".to_string(),
key_internal_priority: 50,
key_global_priority_for_format: Some(2),
key_capabilities: Some(serde_json::json!({"cache_1h": true})),
model_id: "model-1".to_string(),
global_model_id: "global-model-1".to_string(),
global_model_name: "gpt-4.1".to_string(),
selected_provider_model_name: "gpt-4.1-canary".to_string(),
mapping_matched_model: Some("gpt-4.1-canary".to_string()),
};
let json = serde_json::to_value(candidate).expect("candidate should serialize");
assert_eq!(json["provider_name"], "OpenAI");
}
#[tokio::test]
async fn read_minimal_candidate_selection_resolves_provider_model_alias() {
let mut row = sample_row();
row.global_model_name = "gpt-5".to_string();
row.model_provider_model_name = "gpt-5.2".to_string();
row.model_provider_model_mappings = Some(vec![StoredProviderModelMapping {
name: "gpt-5.2".to_string(),
priority: 1,
api_formats: Some(vec!["openai:chat".to_string()]),
}]);
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
row,
]));
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
let state = GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas);
let selection = read_minimal_candidate_selection(&state, "openai:chat", "gpt-5.2", false, None)
.await
.expect("selection should succeed");
assert_eq!(selection.len(), 1);
assert_eq!(selection[0].global_model_name, "gpt-5");
assert_eq!(selection[0].selected_provider_model_name, "gpt-5.2");
}
#[tokio::test]
async fn read_minimal_candidate_selection_allows_resolved_global_model_in_auth_snapshot() {
let mut row = sample_row();
row.global_model_name = "gpt-5".to_string();
row.global_model_mappings = Some(vec!["gpt-5(?:\\.\\d+)?".to_string()]);
row.model_provider_model_name = "gpt-5-upstream".to_string();
row.model_provider_model_mappings = Some(vec![StoredProviderModelMapping {
name: "gpt-5-upstream".to_string(),
priority: 1,
api_formats: Some(vec!["openai:chat".to_string()]),
}]);
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
row,
]));
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
let state = GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas);
let mut auth_snapshot = sample_auth_snapshot("api-key-1");
auth_snapshot.user_allowed_models = Some(vec!["gpt-5".to_string()]);
auth_snapshot.api_key_allowed_models = Some(vec!["gpt-5".to_string()]);
let selection = read_minimal_candidate_selection(
&state,
"openai:chat",
"gpt-5.2",
false,
Some(&auth_snapshot),
)
.await
.expect("selection should succeed");
assert_eq!(selection.len(), 1);
assert_eq!(selection[0].global_model_name, "gpt-5");
}

702
apps/aether-gateway/src/scheduler/candidate/tests/selection.rs Normal file
View File

@@ -0,0 +1,702 @@
use std::sync::Arc;
use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository;
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data::repository::quota::InMemoryProviderQuotaRepository;
use aether_data_contracts::repository::candidate_selection::StoredProviderModelMapping;
use aether_data_contracts::repository::candidates::{
RequestCandidateStatus, StoredRequestCandidate,
};
use aether_data_contracts::repository::quota::StoredProviderQuotaSnapshot;
use crate::data::GatewayDataState;
use crate::AppState;
use super::super::runtime::should_skip_provider_quota;
use super::super::selection::select_minimal_candidate as select_candidate;
use super::support::{sample_auth_snapshot, sample_key, sample_provider, sample_row};
#[test]
fn skips_inactive_or_exhausted_monthly_quota_provider() {
let inactive = StoredProviderQuotaSnapshot::new(
"provider-1".to_string(),
"monthly_quota".to_string(),
Some(10.0),
1.0,
Some(30),
Some(1_000),
None,
false,
)
.expect("quota should build");
assert!(should_skip_provider_quota(&inactive, 2_000));
let exhausted = StoredProviderQuotaSnapshot::new(
"provider-1".to_string(),
"monthly_quota".to_string(),
Some(10.0),
10.0,
Some(30),
Some(1_000),
None,
true,
)
.expect("quota should build");
assert!(should_skip_provider_quota(&exhausted, 2_000));
let payg = StoredProviderQuotaSnapshot::new(
"provider-1".to_string(),
"pay_as_you_go".to_string(),
None,
10.0,
None,
None,
None,
true,
)
.expect("quota should build");
assert!(!should_skip_provider_quota(&payg, 2_000));
}
#[tokio::test]
async fn selects_next_candidate_when_first_provider_quota_is_exhausted() {
let mut first = sample_row();
first.provider_id = "provider-1".to_string();
first.provider_name = "openai-primary".to_string();
first.endpoint_id = "endpoint-1".to_string();
first.key_id = "key-1".to_string();
first.key_name = "primary".to_string();
first.model_provider_model_name = "gpt-4.1-primary".to_string();
first.model_provider_model_mappings = Some(vec![StoredProviderModelMapping {
name: "gpt-4.1-primary".to_string(),
priority: 1,
api_formats: Some(vec!["openai:chat".to_string()]),
}]);
first.key_global_priority_by_format = Some(serde_json::json!({"openai:chat": 1}));
let mut second = sample_row();
second.provider_id = "provider-2".to_string();
second.provider_name = "openai-secondary".to_string();
second.endpoint_id = "endpoint-2".to_string();
second.key_id = "key-2".to_string();
second.key_name = "secondary".to_string();
second.model_provider_model_name = "gpt-4.1-secondary".to_string();
second.model_provider_model_mappings = Some(vec![StoredProviderModelMapping {
name: "gpt-4.1-secondary".to_string(),
priority: 1,
api_formats: Some(vec!["openai:chat".to_string()]),
}]);
second.key_global_priority_by_format = Some(serde_json::json!({"openai:chat": 2}));
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
first, second,
]));
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![
StoredProviderQuotaSnapshot::new(
"provider-1".to_string(),
"monthly_quota".to_string(),
Some(10.0),
10.0,
Some(30),
Some(1_000),
None,
true,
)
.expect("quota should build"),
]));
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(
GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas),
);
let selected = select_candidate(
state.data.as_ref(),
&state,
"openai:chat",
"gpt-4.1",
false,
None,
2_000,
)
.await
.expect("selection should succeed")
.expect("candidate should exist");
assert_eq!(selected.provider_id, "provider-2");
assert_eq!(selected.selected_provider_model_name, "gpt-4.1-secondary");
}
#[tokio::test]
async fn cooled_down_when_recent_failures_are_recorded_for_same_key() {
let mut first = sample_row();
first.provider_id = "provider-1".to_string();
first.provider_name = "openai-primary".to_string();
first.endpoint_id = "endpoint-1".to_string();
first.key_id = "key-1".to_string();
first.key_name = "primary".to_string();
first.model_provider_model_name = "gpt-4.1-primary".to_string();
first.model_provider_model_mappings = Some(vec![StoredProviderModelMapping {
name: "gpt-4.1-primary".to_string(),
priority: 1,
api_formats: Some(vec!["openai:chat".to_string()]),
}]);
first.key_global_priority_by_format = Some(serde_json::json!({"openai:chat": 1}));
let mut second = sample_row();
second.provider_id = "provider-2".to_string();
second.provider_name = "openai-secondary".to_string();
second.endpoint_id = "endpoint-2".to_string();
second.key_id = "key-2".to_string();
second.key_name = "secondary".to_string();
second.model_provider_model_name = "gpt-4.1-secondary".to_string();
second.model_provider_model_mappings = Some(vec![StoredProviderModelMapping {
name: "gpt-4.1-secondary".to_string(),
priority: 1,
api_formats: Some(vec!["openai:chat".to_string()]),
}]);
second.key_global_priority_by_format = Some(serde_json::json!({"openai:chat": 2}));
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
first, second,
]));
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
StoredRequestCandidate::new(
"cand-1".to_string(),
"req-1".to_string(),
None,
None,
None,
None,
0,
0,
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("key-1".to_string()),
RequestCandidateStatus::Failed,
None,
false,
Some(502),
None,
Some("upstream".to_string()),
Some(100),
None,
None,
None,
95,
Some(95),
Some(95),
)
.expect("candidate should build"),
StoredRequestCandidate::new(
"cand-2".to_string(),
"req-2".to_string(),
None,
None,
None,
None,
0,
0,
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("key-1".to_string()),
RequestCandidateStatus::Cancelled,
None,
false,
Some(499),
None,
Some("cancelled".to_string()),
Some(80),
None,
None,
None,
98,
Some(98),
Some(98),
)
.expect("candidate should build"),
]));
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(
GatewayDataState::with_candidate_selection_quota_and_request_candidates_for_tests(
candidates,
quotas,
request_candidates,
),
);
let selected = select_candidate(
state.data.as_ref(),
&state,
"openai:chat",
"gpt-4.1",
false,
None,
100,
)
.await
.expect("selection should succeed")
.expect("candidate should exist");
assert_eq!(selected.provider_id, "provider-2");
assert_eq!(selected.selected_provider_model_name, "gpt-4.1-secondary");
}
#[tokio::test]
async fn selects_next_candidate_when_first_provider_concurrent_limit_is_reached() {
let mut first = sample_row();
first.provider_id = "provider-a".to_string();
first.provider_name = "openai-a".to_string();
first.endpoint_id = "endpoint-a".to_string();
first.key_id = "key-a".to_string();
first.key_name = "alpha".to_string();
first.key_global_priority_by_format = Some(serde_json::json!({"openai:chat": 1}));
let mut second = sample_row();
second.provider_id = "provider-b".to_string();
second.provider_name = "openai-b".to_string();
second.endpoint_id = "endpoint-b".to_string();
second.key_id = "key-b".to_string();
second.key_name = "beta".to_string();
second.key_global_priority_by_format = Some(serde_json::json!({"openai:chat": 2}));
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
first, second,
]));
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![
sample_provider("provider-a", Some(1)),
sample_provider("provider-b", None),
],
Vec::new(),
Vec::new(),
));
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
StoredRequestCandidate::new(
"cand-1".to_string(),
"req-1".to_string(),
None,
None,
None,
None,
0,
0,
Some("provider-a".to_string()),
Some("endpoint-a".to_string()),
Some("key-a".to_string()),
RequestCandidateStatus::Streaming,
None,
false,
None,
None,
None,
None,
None,
None,
None,
95,
Some(95),
None,
)
.expect("candidate should build"),
]));
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 = select_candidate(
state.data.as_ref(),
&state,
"openai:chat",
"gpt-4.1",
false,
None,
100,
)
.await
.expect("selection should succeed")
.expect("candidate should exist");
assert_eq!(selected.provider_id, "provider-b");
assert_eq!(selected.key_id, "key-b");
}
#[tokio::test]
async fn returns_none_when_auth_api_key_concurrent_limit_is_reached() {
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_row(),
]));
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-1", None)],
Vec::new(),
Vec::new(),
));
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
StoredRequestCandidate::new(
"cand-1".to_string(),
"req-1".to_string(),
Some("user-1".to_string()),
Some("api-key-1".to_string()),
None,
None,
0,
0,
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("key-1".to_string()),
RequestCandidateStatus::Pending,
None,
false,
None,
None,
None,
None,
None,
None,
None,
95,
Some(95),
None,
)
.expect("candidate should build"),
]));
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 mut auth_snapshot = sample_auth_snapshot("api-key-1");
auth_snapshot.api_key_concurrent_limit = Some(1);
let selected = select_candidate(
state.data.as_ref(),
&state,
"openai:chat",
"gpt-4.1",
false,
Some(&auth_snapshot),
100,
)
.await
.expect("selection should succeed");
assert!(selected.is_none());
}
#[tokio::test]
async fn selects_next_candidate_when_first_provider_key_rpm_slots_are_reserved_for_new_user() {
let mut first = sample_row();
first.provider_id = "provider-a".to_string();
first.provider_name = "openai-a".to_string();
first.endpoint_id = "endpoint-a".to_string();
first.key_id = "key-a".to_string();
first.key_name = "alpha".to_string();
first.key_global_priority_by_format = Some(serde_json::json!({"openai:chat": 1}));
let mut second = sample_row();
second.provider_id = "provider-b".to_string();
second.provider_name = "openai-b".to_string();
second.endpoint_id = "endpoint-b".to_string();
second.key_id = "key-b".to_string();
second.key_name = "beta".to_string();
second.key_global_priority_by_format = Some(serde_json::json!({"openai:chat": 2}));
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
first, second,
]));
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![
sample_provider("provider-a", None),
sample_provider("provider-b", None),
],
Vec::new(),
vec![
sample_key("key-a", "provider-a", Some(10)),
sample_key("key-b", "provider-b", Some(10)),
],
));
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
StoredRequestCandidate::new(
"cand-1".to_string(),
"req-1".to_string(),
None,
Some("api-key-new-user".to_string()),
None,
None,
0,
0,
Some("provider-a".to_string()),
Some("endpoint-a".to_string()),
Some("key-a".to_string()),
RequestCandidateStatus::Success,
None,
false,
Some(200),
None,
None,
Some(10),
Some(9),
None,
None,
95,
Some(95),
Some(96),
)
.expect("candidate should build"),
]));
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 auth_snapshot = sample_auth_snapshot("api-key-new-user");
let selected = select_candidate(
state.data.as_ref(),
&state,
"openai:chat",
"gpt-4.1",
false,
Some(&auth_snapshot),
100,
)
.await
.expect("selection should succeed")
.expect("candidate should exist");
assert_eq!(selected.provider_id, "provider-b");
assert_eq!(selected.key_id, "key-b");
}
#[tokio::test]
async fn selects_next_candidate_when_first_provider_key_circuit_is_open() {
let mut first = sample_row();
first.provider_id = "provider-a".to_string();
first.provider_name = "openai-a".to_string();
first.endpoint_id = "endpoint-a".to_string();
first.key_id = "key-a".to_string();
first.key_name = "alpha".to_string();
first.key_global_priority_by_format = Some(serde_json::json!({"openai:chat": 1}));
let mut second = sample_row();
second.provider_id = "provider-b".to_string();
second.provider_name = "openai-b".to_string();
second.endpoint_id = "endpoint-b".to_string();
second.key_id = "key-b".to_string();
second.key_name = "beta".to_string();
second.key_global_priority_by_format = Some(serde_json::json!({"openai:chat": 2}));
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
first, second,
]));
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![
sample_provider("provider-a", None),
sample_provider("provider-b", None),
],
Vec::new(),
vec![
sample_key("key-a", "provider-a", Some(10)).with_health_fields(
Some(serde_json::json!({"openai:chat": {"health_score": 0.2}})),
Some(serde_json::json!({"openai:chat": {"open": true}})),
),
sample_key("key-b", "provider-b", Some(10)),
],
));
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 = select_candidate(
state.data.as_ref(),
&state,
"openai:chat",
"gpt-4.1",
false,
None,
100,
)
.await
.expect("selection should succeed")
.expect("candidate should exist");
assert_eq!(selected.provider_id, "provider-b");
assert_eq!(selected.key_id, "key-b");
}
#[tokio::test]
async fn same_priority_candidates_prefer_healthier_provider_key_before_id_order() {
let mut first = sample_row();
first.provider_id = "provider-a".to_string();
first.provider_name = "openai-a".to_string();
first.endpoint_id = "endpoint-a".to_string();
first.key_id = "key-a".to_string();
first.key_name = "alpha".to_string();
first.key_global_priority_by_format = Some(serde_json::json!({"openai:chat": 1}));
let mut second = sample_row();
second.provider_id = "provider-b".to_string();
second.provider_name = "openai-b".to_string();
second.endpoint_id = "endpoint-b".to_string();
second.key_id = "key-b".to_string();
second.key_name = "beta".to_string();
second.key_global_priority_by_format = Some(serde_json::json!({"openai:chat": 1}));
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
first, second,
]));
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![
sample_provider("provider-a", None),
sample_provider("provider-b", None),
],
Vec::new(),
vec![
sample_key("key-a", "provider-a", Some(10)).with_health_fields(
Some(serde_json::json!({"openai:chat": {"health_score": 0.30}})),
None,
),
sample_key("key-b", "provider-b", Some(10)).with_health_fields(
Some(serde_json::json!({"openai:chat": {"health_score": 0.95}})),
None,
),
],
));
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 = select_candidate(
state.data.as_ref(),
&state,
"openai:chat",
"gpt-4.1",
false,
None,
100,
)
.await
.expect("selection should succeed")
.expect("candidate should exist");
assert_eq!(selected.provider_id, "provider-b");
assert_eq!(selected.key_id, "key-b");
}
#[tokio::test]
async fn same_priority_candidates_use_aggregate_health_score_when_api_format_specific_health_is_missing(
) {
let mut first = sample_row();
first.provider_id = "provider-a".to_string();
first.provider_name = "openai-a".to_string();
first.endpoint_id = "endpoint-a".to_string();
first.key_id = "key-a".to_string();
first.key_name = "alpha".to_string();
first.key_global_priority_by_format = Some(serde_json::json!({"openai:chat": 1}));
let mut second = sample_row();
second.provider_id = "provider-b".to_string();
second.provider_name = "openai-b".to_string();
second.endpoint_id = "endpoint-b".to_string();
second.key_id = "key-b".to_string();
second.key_name = "beta".to_string();
second.key_global_priority_by_format = Some(serde_json::json!({"openai:chat": 1}));
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
first, second,
]));
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![
sample_provider("provider-a", None),
sample_provider("provider-b", None),
],
Vec::new(),
vec![
sample_key("key-a", "provider-a", Some(10)).with_health_fields(
Some(serde_json::json!({
"openai:responses": {"health_score": 0.40},
"claude:chat": {"health_score": 0.55}
})),
None,
),
sample_key("key-b", "provider-b", Some(10)).with_health_fields(
Some(serde_json::json!({
"openai:responses": {"health_score": 0.90},
"claude:chat": {"health_score": 0.92}
})),
None,
),
],
));
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 = select_candidate(
state.data.as_ref(),
&state,
"openai:chat",
"gpt-4.1",
false,
None,
100,
)
.await
.expect("selection should succeed")
.expect("candidate should exist");
assert_eq!(selected.provider_id, "provider-b");
assert_eq!(selected.key_id, "key-b");
}

122
apps/aether-gateway/src/scheduler/candidate/tests/support.rs Normal file
View File

@@ -0,0 +1,122 @@
use aether_data_contracts::repository::candidate_selection::{
StoredMinimalCandidateSelectionRow, StoredProviderModelMapping,
};
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use crate::data::auth::GatewayAuthApiKeySnapshot;
pub(super) fn sample_row() -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: "provider-1".to_string(),
provider_name: "OpenAI".to_string(),
provider_type: "custom".to_string(),
provider_priority: 10,
provider_is_active: true,
endpoint_id: "endpoint-1".to_string(),
endpoint_api_format: "openai:chat".to_string(),
endpoint_api_family: Some("openai".to_string()),
endpoint_kind: Some("chat".to_string()),
endpoint_is_active: true,
key_id: "key-1".to_string(),
key_name: "prod".to_string(),
key_auth_type: "api_key".to_string(),
key_is_active: true,
key_api_formats: Some(vec!["openai:chat".to_string()]),
key_allowed_models: None,
key_capabilities: Some(serde_json::json!({"cache_1h": true})),
key_internal_priority: 50,
key_global_priority_by_format: Some(serde_json::json!({"openai:chat": 2})),
model_id: "model-1".to_string(),
global_model_id: "global-model-1".to_string(),
global_model_name: "gpt-4.1".to_string(),
global_model_mappings: Some(vec!["gpt-4\\.1-.*".to_string()]),
global_model_supports_streaming: Some(true),
model_provider_model_name: "gpt-4.1-upstream".to_string(),
model_provider_model_mappings: Some(vec![
StoredProviderModelMapping {
name: "gpt-4.1-canary".to_string(),
priority: 1,
api_formats: Some(vec!["openai:chat".to_string()]),
},
StoredProviderModelMapping {
name: "gpt-4.1-responses".to_string(),
priority: 1,
api_formats: Some(vec!["openai:responses".to_string()]),
},
]),
model_supports_streaming: None,
model_is_active: true,
model_is_available: true,
}
}
pub(super) fn sample_provider(
id: &str,
concurrent_limit: Option<i32>,
) -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
id.to_string(),
format!("provider-{id}"),
Some("https://example.com".to_string()),
"custom".to_string(),
)
.expect("provider should build")
.with_transport_fields(
true,
false,
false,
concurrent_limit,
None,
None,
None,
None,
None,
)
}
pub(super) fn sample_key(
id: &str,
provider_id: &str,
rpm_limit: Option<u32>,
) -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
id.to_string(),
provider_id.to_string(),
format!("key-{id}"),
"api_key".to_string(),
None,
true,
)
.expect("key should build")
.with_rate_limit_fields(rpm_limit, None, None, None, None, None, Some(20), Some(20))
}
pub(super) fn sample_auth_snapshot(api_key_id: &str) -> GatewayAuthApiKeySnapshot {
GatewayAuthApiKeySnapshot {
user_id: "user-1".to_string(),
username: "alice".to_string(),
email: None,
user_role: "user".to_string(),
user_auth_source: "local".to_string(),
user_is_active: true,
user_is_deleted: false,
user_rate_limit: None,
user_allowed_providers: None,
user_allowed_api_formats: None,
user_allowed_models: None,
api_key_id: api_key_id.to_string(),
api_key_name: Some("default".to_string()),
api_key_is_active: true,
api_key_is_locked: false,
api_key_is_standalone: false,
api_key_rate_limit: None,
api_key_concurrent_limit: None,
api_key_expires_at_unix_secs: None,
api_key_allowed_providers: None,
api_key_allowed_api_formats: None,
api_key_allowed_models: None,
currently_usable: true,
}
}

316
apps/aether-gateway/src/scheduler/failover.rs
View File

@@ -1,316 +0,0 @@
use aether_contracts::ExecutionResult;
fn is_local_candidate_attempt(report_context: Option<&serde_json::Value>) -> bool {
report_context
.and_then(serde_json::Value::as_object)
.and_then(|context| context.get("candidate_index"))
.and_then(serde_json::Value::as_u64)
.is_some()
}
fn is_retryable_local_upstream_status(status_code: u16) -> bool {
status_code == 429 || status_code >= 500
}
pub(crate) fn should_retry_next_local_candidate_sync(
plan_kind: &str,
report_context: Option<&serde_json::Value>,
result: &ExecutionResult,
) -> bool {
is_local_candidate_attempt(report_context)
&& plan_kind == "openai_chat_sync"
&& is_retryable_local_upstream_status(result.status_code)
}
pub(crate) fn should_fallback_to_control_sync(
plan_kind: &str,
result: &ExecutionResult,
body_json: Option<&serde_json::Value>,
has_body_bytes: bool,
explicit_finalize: bool,
mapped_error_finalize: bool,
) -> bool {
if explicit_finalize
&& matches!(
plan_kind,
"openai_video_delete_sync" | "openai_video_cancel_sync" | "gemini_video_cancel_sync"
)
{
return false;
}
if !matches!(
plan_kind,
"openai_video_create_sync"
| "openai_video_remix_sync"
| "gemini_video_create_sync"
| "openai_chat_sync"
| "openai_cli_sync"
| "openai_compact_sync"
| "claude_chat_sync"
| "gemini_chat_sync"
| "claude_cli_sync"
| "gemini_cli_sync"
) {
return false;
}
if explicit_finalize {
return result.status_code < 400 && body_json.is_none() && !has_body_bytes;
}
if mapped_error_finalize {
return false;
}
if result.status_code >= 400 {
return true;
}
let Some(body_json) = body_json else {
return true;
};
body_json.get("error").is_some()
}
pub(crate) fn should_finalize_sync_response(report_kind: Option<&str>) -> bool {
report_kind.is_some_and(|kind| kind.ends_with("_finalize"))
}
pub(crate) fn resolve_core_sync_error_finalize_report_kind(
plan_kind: &str,
result: &ExecutionResult,
body_json: Option<&serde_json::Value>,
) -> Option<String> {
let has_embedded_error = body_json.is_some_and(|value| value.get("error").is_some());
if result.status_code < 400 && !has_embedded_error {
return None;
}
let report_kind = match plan_kind {
"openai_chat_sync" => "openai_chat_sync_finalize",
"openai_cli_sync" => "openai_cli_sync_finalize",
"openai_compact_sync" => "openai_compact_sync_finalize",
"claude_chat_sync" => "claude_chat_sync_finalize",
"gemini_chat_sync" => "gemini_chat_sync_finalize",
"claude_cli_sync" => "claude_cli_sync_finalize",
"gemini_cli_sync" => "gemini_cli_sync_finalize",
_ => return None,
};
Some(report_kind.to_string())
}
pub(crate) fn should_retry_next_local_candidate_stream(
plan_kind: &str,
report_context: Option<&serde_json::Value>,
status_code: u16,
) -> bool {
is_local_candidate_attempt(report_context)
&& plan_kind == "openai_chat_stream"
&& is_retryable_local_upstream_status(status_code)
}
pub(crate) fn should_fallback_to_control_stream(
plan_kind: &str,
status_code: u16,
mapped_error_finalize: bool,
) -> bool {
if mapped_error_finalize {
return false;
}
matches!(
plan_kind,
"openai_chat_stream"
| "claude_chat_stream"
| "gemini_chat_stream"
| "openai_cli_stream"
| "openai_compact_stream"
| "claude_cli_stream"
| "gemini_cli_stream"
) && status_code >= 400
}
pub(crate) fn resolve_core_stream_error_finalize_report_kind(
plan_kind: &str,
status_code: u16,
) -> Option<String> {
if status_code < 400 {
return None;
}
let report_kind = match plan_kind {
"openai_chat_stream" => "openai_chat_sync_finalize",
"claude_chat_stream" => "claude_chat_sync_finalize",
"gemini_chat_stream" => "gemini_chat_sync_finalize",
"openai_cli_stream" => "openai_cli_sync_finalize",
"openai_compact_stream" => "openai_compact_sync_finalize",
"claude_cli_stream" => "claude_cli_sync_finalize",
"gemini_cli_stream" => "gemini_cli_sync_finalize",
_ => return None,
};
Some(report_kind.to_string())
}
pub(crate) fn resolve_core_stream_direct_finalize_report_kind(plan_kind: &str) -> Option<String> {
let report_kind = match plan_kind {
"openai_chat_stream" => "openai_chat_sync_finalize",
"claude_chat_stream" => "claude_chat_sync_finalize",
"gemini_chat_stream" => "gemini_chat_sync_finalize",
"openai_cli_stream" => "openai_cli_sync_finalize",
"openai_compact_stream" => "openai_compact_sync_finalize",
"claude_cli_stream" => "claude_cli_sync_finalize",
"gemini_cli_stream" => "gemini_cli_sync_finalize",
_ => return None,
};
Some(report_kind.to_string())
}
#[cfg(test)]
mod tests {
use aether_contracts::ExecutionResult;
use super::{
resolve_core_stream_error_finalize_report_kind,
resolve_core_sync_error_finalize_report_kind, should_fallback_to_control_stream,
should_fallback_to_control_sync, should_retry_next_local_candidate_stream,
should_retry_next_local_candidate_sync,
};
#[test]
fn sync_failover_marks_chat_errors() {
let result = ExecutionResult {
request_id: "req-1".to_string(),
candidate_id: None,
status_code: 502,
headers: Default::default(),
body: None,
telemetry: None,
error: None,
};
assert!(should_fallback_to_control_sync(
"openai_chat_sync",
&result,
None,
false,
false,
false,
));
assert_eq!(
resolve_core_sync_error_finalize_report_kind("openai_chat_sync", &result, None),
Some("openai_chat_sync_finalize".to_string())
);
}
#[test]
fn stream_failover_marks_chat_errors() {
assert!(should_fallback_to_control_stream(
"openai_chat_stream",
502,
false,
));
assert_eq!(
resolve_core_stream_error_finalize_report_kind("openai_chat_stream", 502),
Some("openai_chat_sync_finalize".to_string())
);
}
#[test]
fn sync_retry_next_candidate_is_local_openai_chat_only() {
let result = ExecutionResult {
request_id: "req-1".to_string(),
candidate_id: None,
status_code: 502,
headers: Default::default(),
body: None,
telemetry: None,
error: None,
};
let local_report_context = serde_json::json!({
"candidate_index": 0,
"retry_index": 0,
});
assert!(should_retry_next_local_candidate_sync(
"openai_chat_sync",
Some(&local_report_context),
&result,
));
assert!(!should_retry_next_local_candidate_sync(
"openai_chat_sync",
None,
&result,
));
assert!(!should_retry_next_local_candidate_sync(
"claude_chat_sync",
None,
&result,
));
}
#[test]
fn sync_retry_next_candidate_treats_rate_limit_as_retryable() {
let result = ExecutionResult {
request_id: "req-1".to_string(),
candidate_id: None,
status_code: 429,
headers: Default::default(),
body: None,
telemetry: None,
error: None,
};
let local_report_context = serde_json::json!({
"candidate_index": 0,
"retry_index": 0,
});
assert!(should_retry_next_local_candidate_sync(
"openai_chat_sync",
Some(&local_report_context),
&result,
));
}
#[test]
fn stream_retry_next_candidate_is_local_openai_chat_only() {
let local_report_context = serde_json::json!({
"candidate_index": 0,
"retry_index": 0,
});
assert!(should_retry_next_local_candidate_stream(
"openai_chat_stream",
Some(&local_report_context),
502,
));
assert!(!should_retry_next_local_candidate_stream(
"openai_chat_stream",
None,
502,
));
assert!(!should_retry_next_local_candidate_stream(
"claude_chat_stream",
None,
502,
));
}
#[test]
fn stream_retry_next_candidate_treats_rate_limit_as_retryable() {
let local_report_context = serde_json::json!({
"candidate_index": 0,
"retry_index": 0,
});
assert!(should_retry_next_local_candidate_stream(
"openai_chat_stream",
Some(&local_report_context),
429,
));
}
}

9
apps/aether-gateway/src/scheduler/health.rs
View File

@@ -1,9 +0,0 @@
pub(crate) use aether_scheduler_core::{
aggregate_provider_key_health_score, count_recent_active_requests_for_api_key,
count_recent_active_requests_for_provider, count_recent_rpm_requests_for_provider_key,
count_recent_rpm_requests_for_provider_key_since, effective_provider_key_health_score,
effective_provider_key_rpm_limit, is_candidate_in_recent_failure_cooldown,
is_provider_key_circuit_open, provider_key_health_bucket, provider_key_health_score,
provider_key_rpm_allows_request, provider_key_rpm_allows_request_since,
ProviderKeyHealthBucket, PROVIDER_KEY_RPM_WINDOW_SECS,
};

41
apps/aether-gateway/src/scheduler/mod.rs
View File

@@ -1,38 +1,3 @@
mod candidate;
mod failover;
mod health;
mod request_candidate_state;
mod request_candidates;
mod route;
pub(crate) use candidate::{
list_selectable_candidates,
list_selectable_candidates_for_required_capability_without_requested_model,
read_cached_scheduler_affinity_target, read_minimal_candidate_selection,
GatewayMinimalCandidateSelectionCandidate,
};
pub(crate) use candidate::{
SchedulerCandidateSelectionRowSource, SchedulerRuntimeState,
};
pub(crate) use failover::{
resolve_core_stream_direct_finalize_report_kind,
resolve_core_stream_error_finalize_report_kind, resolve_core_sync_error_finalize_report_kind,
should_fallback_to_control_stream, should_fallback_to_control_sync,
should_finalize_sync_response, should_retry_next_local_candidate_stream,
should_retry_next_local_candidate_sync,
};
pub(crate) use health::{
count_recent_rpm_requests_for_provider_key, count_recent_rpm_requests_for_provider_key_since,
is_provider_key_circuit_open, provider_key_health_score, provider_key_rpm_allows_request_since,
PROVIDER_KEY_RPM_WINDOW_SECS,
};
pub(crate) use request_candidate_state::SchedulerRequestCandidateRuntimeState;
pub(crate) use request_candidates::{
current_unix_secs, ensure_execution_request_candidate_slot, execution_error_details,
record_local_request_candidate_status, record_report_request_candidate_status,
};
pub(crate) use route::{
is_matching_stream_request, resolve_execution_runtime_stream_plan_kind,
resolve_execution_runtime_sync_plan_kind, supports_stream_scheduler_decision_kind,
supports_sync_scheduler_decision_kind,
};
pub(crate) mod affinity;
pub(crate) mod candidate;
pub(crate) mod state;

19
apps/aether-gateway/src/scheduler/request_candidate_state.rs
View File

@@ -1,19 +0,0 @@
use aether_data::repository::candidates::{StoredRequestCandidate, UpsertRequestCandidateRecord};
use async_trait::async_trait;
use crate::GatewayError;
#[async_trait]
pub(crate) trait SchedulerRequestCandidateRuntimeState {
fn has_request_candidate_data_writer(&self) -> bool;
async fn read_request_candidates_by_request_id(
&self,
request_id: &str,
) -> Result<Vec<StoredRequestCandidate>, GatewayError>;
async fn upsert_request_candidate(
&self,
candidate: UpsertRequestCandidateRecord,
) -> Result<Option<StoredRequestCandidate>, GatewayError>;
}

377
apps/aether-gateway/src/scheduler/request_candidates.rs
View File

@@ -1,377 +0,0 @@
use std::time::{SystemTime, UNIX_EPOCH};
use aether_contracts::ExecutionPlan;
use aether_data::repository::candidates::RequestCandidateStatus;
use aether_scheduler_core::{
build_execution_request_candidate_seed, build_local_request_candidate_status_record,
build_report_request_candidate_status_record,
finalize_execution_request_candidate_report_context, parse_request_candidate_report_context,
resolve_report_request_candidate_slot as resolve_report_request_candidate_slot_from_candidates,
SchedulerResolvedReportRequestCandidateSlot,
};
use serde_json::Value;
use tracing::warn;
use uuid::Uuid;
use super::request_candidate_state::SchedulerRequestCandidateRuntimeState;
pub(crate) use aether_scheduler_core::execution_error_details;
pub(crate) async fn record_local_request_candidate_status(
state: &(impl SchedulerRequestCandidateRuntimeState + ?Sized),
plan: &ExecutionPlan,
report_context: Option<&Value>,
status: RequestCandidateStatus,
status_code: Option<u16>,
error_type: Option<String>,
error_message: Option<String>,
latency_ms: Option<u64>,
started_at_unix_secs: Option<u64>,
finished_at_unix_secs: Option<u64>,
) {
let Some(record) = build_local_request_candidate_status_record(
plan,
report_context,
status,
status_code,
error_type,
error_message,
latency_ms,
started_at_unix_secs,
finished_at_unix_secs,
) else {
return;
};
let candidate_id = record.id.clone();
if let Err(err) = state.upsert_request_candidate(record).await {
warn!(
event_name = "request_candidate_status_persist_failed",
log_type = "event",
request_id = %plan.request_id,
candidate_id = %candidate_id,
error = ?err,
"gateway failed to persist request candidate status update"
);
}
}
pub(crate) async fn record_report_request_candidate_status(
state: &(impl SchedulerRequestCandidateRuntimeState + ?Sized),
report_context: Option<&Value>,
status: RequestCandidateStatus,
status_code: Option<u16>,
error_type: Option<String>,
error_message: Option<String>,
latency_ms: Option<u64>,
started_at_unix_secs: Option<u64>,
finished_at_unix_secs: Option<u64>,
) {
let Some(slot) = resolve_report_request_candidate_slot(state, report_context).await else {
return;
};
let request_id = slot.request_id.clone();
let candidate_index = slot.candidate_index;
let retry_index = slot.retry_index;
let record = build_report_request_candidate_status_record(
slot,
status,
status_code,
error_type,
error_message,
latency_ms,
started_at_unix_secs,
finished_at_unix_secs,
current_unix_secs(),
);
if let Err(err) = state.upsert_request_candidate(record).await {
warn!(
event_name = "request_candidate_report_status_persist_failed",
log_type = "event",
request_id = %request_id,
candidate_index,
retry_index,
error = ?err,
"gateway failed to persist report-driven request candidate status update"
);
}
}
pub(crate) async fn ensure_execution_request_candidate_slot(
state: &(impl SchedulerRequestCandidateRuntimeState + ?Sized),
plan: &mut ExecutionPlan,
report_context: &mut Option<Value>,
) {
if !state.has_request_candidate_data_writer() {
return;
}
if plan
.candidate_id
.as_deref()
.map(str::trim)
.is_some_and(|value| !value.is_empty())
{
return;
}
let seed = build_execution_request_candidate_seed(
plan,
report_context.as_ref(),
current_unix_secs(),
Uuid::new_v4().to_string(),
);
let generated_candidate_id = seed.upsert_record.id.clone();
let candidate_id = match state.upsert_request_candidate(seed.upsert_record).await {
Ok(Some(stored)) => stored.id,
Ok(None) => generated_candidate_id,
Err(err) => {
warn!(
event_name = "request_candidate_slot_seed_failed",
log_type = "event",
request_id = %plan.request_id,
error = ?err,
"gateway failed to seed execution request candidate slot"
);
return;
}
};
plan.candidate_id = Some(candidate_id.clone());
*report_context = Some(finalize_execution_request_candidate_report_context(
seed.report_context,
&candidate_id,
));
}
pub(crate) fn current_unix_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
}
async fn resolve_report_request_candidate_slot(
state: &(impl SchedulerRequestCandidateRuntimeState + ?Sized),
report_context: Option<&Value>,
) -> Option<SchedulerResolvedReportRequestCandidateSlot> {
let metadata = parse_request_candidate_report_context(report_context)?;
let request_id = metadata.request_id.clone()?;
let existing_candidates = state
.read_request_candidates_by_request_id(request_id.as_str())
.await
.ok()
.unwrap_or_default();
resolve_report_request_candidate_slot_from_candidates(
&existing_candidates,
metadata,
current_unix_secs(),
Uuid::new_v4().to_string(),
)
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use std::sync::Arc;
use aether_contracts::{ExecutionPlan, RequestBody};
use aether_data::repository::candidates::{
InMemoryRequestCandidateRepository, RequestCandidateReadRepository, RequestCandidateStatus,
StoredRequestCandidate,
};
use aether_data::repository::usage::InMemoryUsageReadRepository;
use serde_json::json;
use super::{ensure_execution_request_candidate_slot, record_report_request_candidate_status};
use crate::data::GatewayDataState;
use crate::AppState;
fn build_test_state(repository: Arc<InMemoryRequestCandidateRepository>) -> AppState {
AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(
GatewayDataState::with_request_candidate_and_usage_repository_for_tests(
repository,
Arc::new(InMemoryUsageReadRepository::default()),
),
)
}
fn sample_plan() -> ExecutionPlan {
ExecutionPlan {
request_id: "req-request-candidate-seed-123".to_string(),
candidate_id: None,
provider_name: Some("openai".to_string()),
provider_id: "provider-request-candidate-seed-123".to_string(),
endpoint_id: "endpoint-request-candidate-seed-123".to_string(),
key_id: "key-request-candidate-seed-123".to_string(),
method: "POST".to_string(),
url: "https://api.openai.example/v1/chat/completions".to_string(),
headers: BTreeMap::new(),
content_type: Some("application/json".to_string()),
content_encoding: None,
body: RequestBody::from_json(json!({"model": "gpt-5", "messages": []})),
stream: false,
client_api_format: "openai:chat".to_string(),
provider_api_format: "openai:chat".to_string(),
model_name: Some("gpt-5".to_string()),
proxy: None,
tls_profile: None,
timeouts: None,
}
}
#[tokio::test]
async fn seeds_execution_request_candidate_slot_for_plan_without_candidate_id() {
let repository = Arc::new(InMemoryRequestCandidateRepository::default());
let state = build_test_state(Arc::clone(&repository));
let mut plan = sample_plan();
let mut report_context = Some(json!({
"request_id": "req-request-candidate-seed-123",
"client_api_format": "openai:chat"
}));
ensure_execution_request_candidate_slot(&state, &mut plan, &mut report_context).await;
let candidate_id = plan
.candidate_id
.clone()
.expect("candidate id should be seeded");
let report_context = report_context.expect("report context should be populated");
assert_eq!(
report_context
.get("candidate_id")
.and_then(|value| value.as_str()),
Some(candidate_id.as_str())
);
assert_eq!(
report_context
.get("candidate_index")
.and_then(|value| value.as_u64()),
Some(0)
);
assert_eq!(
report_context
.get("provider_id")
.and_then(|value| value.as_str()),
Some("provider-request-candidate-seed-123")
);
let stored = repository
.list_by_request_id("req-request-candidate-seed-123")
.await
.expect("request candidates should read");
assert_eq!(stored.len(), 1);
assert_eq!(stored[0].id, candidate_id);
assert_eq!(stored[0].status, RequestCandidateStatus::Pending);
assert_eq!(
stored[0].provider_id.as_deref(),
Some("provider-request-candidate-seed-123")
);
assert_eq!(
stored[0].endpoint_id.as_deref(),
Some("endpoint-request-candidate-seed-123")
);
assert_eq!(
stored[0].key_id.as_deref(),
Some("key-request-candidate-seed-123")
);
}
#[tokio::test]
async fn does_not_reseed_execution_request_candidate_slot_when_plan_already_has_candidate_id() {
let repository = Arc::new(InMemoryRequestCandidateRepository::default());
let state = build_test_state(Arc::clone(&repository));
let mut plan = sample_plan();
plan.candidate_id = Some("cand-existing-123".to_string());
let mut report_context = Some(json!({
"request_id": "req-request-candidate-seed-123"
}));
ensure_execution_request_candidate_slot(&state, &mut plan, &mut report_context).await;
assert_eq!(plan.candidate_id.as_deref(), Some("cand-existing-123"));
let stored = repository
.list_by_request_id("req-request-candidate-seed-123")
.await
.expect("request candidates should read");
assert!(stored.is_empty());
assert_eq!(
report_context
.as_ref()
.and_then(|value| value.get("candidate_id"))
.and_then(|value| value.as_str()),
None
);
}
#[tokio::test]
async fn records_report_request_candidate_status_for_existing_slot() {
let repository = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
StoredRequestCandidate::new(
"cand-report-123".to_string(),
"req-report-123".to_string(),
Some("user-1".to_string()),
Some("api-key-1".to_string()),
None,
None,
0,
0,
Some("provider-report-123".to_string()),
Some("endpoint-report-123".to_string()),
Some("key-report-123".to_string()),
RequestCandidateStatus::Pending,
None,
false,
None,
None,
None,
None,
None,
None,
None,
100,
Some(100),
None,
)
.expect("request candidate should build"),
]));
let state = build_test_state(Arc::clone(&repository));
let report_context = json!({
"request_id": "req-report-123",
"candidate_id": "cand-report-123",
"candidate_index": 0,
"retry_index": 0,
"provider_id": "provider-report-123",
"endpoint_id": "endpoint-report-123",
"key_id": "key-report-123"
});
record_report_request_candidate_status(
&state,
Some(&report_context),
RequestCandidateStatus::Success,
Some(200),
None,
None,
Some(25),
Some(101),
Some(102),
)
.await;
let stored = repository
.list_by_request_id("req-report-123")
.await
.expect("request candidates should read");
assert_eq!(stored.len(), 1);
assert_eq!(stored[0].id, "cand-report-123");
assert_eq!(stored[0].status, RequestCandidateStatus::Success);
assert_eq!(stored[0].status_code, Some(200));
assert_eq!(stored[0].latency_ms, Some(25));
assert_eq!(stored[0].started_at_unix_secs, Some(101));
assert_eq!(stored[0].finished_at_unix_secs, Some(102));
}
}

355
apps/aether-gateway/src/scheduler/route.rs
View File

@@ -1,355 +0,0 @@
use crate::control::GatewayControlDecision;
pub(crate) fn resolve_execution_runtime_stream_plan_kind(
parts: &http::request::Parts,
decision: &GatewayControlDecision,
) -> Option<&'static str> {
if decision.route_class.as_deref() != Some("ai_public") {
return None;
}
if decision.route_family.as_deref() == Some("gemini")
&& decision.route_kind.as_deref() == Some("files")
&& parts.method == http::Method::GET
&& parts.uri.path().ends_with(":download")
{
return Some("gemini_files_download");
}
if decision.route_family.as_deref() == Some("openai")
&& decision.route_kind.as_deref() == Some("chat")
&& parts.method == http::Method::POST
&& parts.uri.path() == "/v1/chat/completions"
{
return Some("openai_chat_stream");
}
if decision.route_family.as_deref() == Some("claude")
&& decision.route_kind.as_deref() == Some("chat")
&& parts.method == http::Method::POST
&& parts.uri.path() == "/v1/messages"
{
return Some("claude_chat_stream");
}
if decision.route_family.as_deref() == Some("claude")
&& decision.route_kind.as_deref() == Some("cli")
&& parts.method == http::Method::POST
&& parts.uri.path() == "/v1/messages"
{
return Some("claude_cli_stream");
}
if decision.route_family.as_deref() == Some("gemini")
&& decision.route_kind.as_deref() == Some("chat")
&& parts.method == http::Method::POST
&& parts.uri.path().ends_with(":streamGenerateContent")
{
return Some("gemini_chat_stream");
}
if decision.route_family.as_deref() == Some("gemini")
&& decision.route_kind.as_deref() == Some("cli")
&& parts.method == http::Method::POST
&& parts.uri.path().ends_with(":streamGenerateContent")
{
return Some("gemini_cli_stream");
}
if decision.route_family.as_deref() == Some("openai")
&& decision.route_kind.as_deref() == Some("cli")
&& parts.method == http::Method::POST
&& parts.uri.path() == "/v1/responses"
{
return Some("openai_cli_stream");
}
if decision.route_family.as_deref() == Some("openai")
&& decision.route_kind.as_deref() == Some("compact")
&& parts.method == http::Method::POST
&& parts.uri.path() == "/v1/responses/compact"
{
return Some("openai_compact_stream");
}
if decision.route_family.as_deref() == Some("openai")
&& decision.route_kind.as_deref() == Some("video")
&& parts.method == http::Method::GET
&& parts.uri.path().ends_with("/content")
{
return Some("openai_video_content");
}
None
}
pub(crate) fn resolve_execution_runtime_sync_plan_kind(
parts: &http::request::Parts,
decision: &GatewayControlDecision,
) -> Option<&'static str> {
if decision.route_class.as_deref() != Some("ai_public") {
return None;
}
if decision.route_family.as_deref() == Some("openai")
&& decision.route_kind.as_deref() == Some("video")
&& parts.method == http::Method::POST
&& parts.uri.path().starts_with("/v1/videos/")
&& parts.uri.path().ends_with("/cancel")
{
return Some("openai_video_cancel_sync");
}
if decision.route_family.as_deref() == Some("openai")
&& decision.route_kind.as_deref() == Some("video")
&& parts.method == http::Method::POST
&& parts.uri.path().starts_with("/v1/videos/")
&& parts.uri.path().ends_with("/remix")
{
return Some("openai_video_remix_sync");
}
if decision.route_family.as_deref() == Some("openai")
&& decision.route_kind.as_deref() == Some("video")
&& parts.method == http::Method::POST
&& parts.uri.path() == "/v1/videos"
{
return Some("openai_video_create_sync");
}
if decision.route_family.as_deref() == Some("openai")
&& decision.route_kind.as_deref() == Some("video")
&& parts.method == http::Method::DELETE
&& parts.uri.path().starts_with("/v1/videos/")
{
return Some("openai_video_delete_sync");
}
if decision.route_family.as_deref() == Some("gemini")
&& decision.route_kind.as_deref() == Some("video")
&& parts.method == http::Method::POST
&& parts.uri.path().ends_with(":cancel")
{
return Some("gemini_video_cancel_sync");
}
if decision.route_family.as_deref() == Some("gemini")
&& decision.route_kind.as_deref() == Some("video")
&& parts.method == http::Method::POST
&& parts.uri.path().ends_with(":predictLongRunning")
{
return Some("gemini_video_create_sync");
}
if decision.route_family.as_deref() == Some("openai")
&& decision.route_kind.as_deref() == Some("chat")
&& parts.method == http::Method::POST
&& parts.uri.path() == "/v1/chat/completions"
{
return Some("openai_chat_sync");
}
if decision.route_family.as_deref() == Some("openai")
&& decision.route_kind.as_deref() == Some("cli")
&& parts.method == http::Method::POST
&& parts.uri.path() == "/v1/responses"
{
return Some("openai_cli_sync");
}
if decision.route_family.as_deref() == Some("openai")
&& decision.route_kind.as_deref() == Some("compact")
&& parts.method == http::Method::POST
&& parts.uri.path() == "/v1/responses/compact"
{
return Some("openai_compact_sync");
}
if decision.route_family.as_deref() == Some("claude")
&& decision.route_kind.as_deref() == Some("chat")
&& parts.method == http::Method::POST
&& parts.uri.path() == "/v1/messages"
{
return Some("claude_chat_sync");
}
if decision.route_family.as_deref() == Some("claude")
&& decision.route_kind.as_deref() == Some("cli")
&& parts.method == http::Method::POST
&& parts.uri.path() == "/v1/messages"
{
return Some("claude_cli_sync");
}
if decision.route_family.as_deref() == Some("gemini")
&& decision.route_kind.as_deref() == Some("chat")
&& parts.method == http::Method::POST
&& parts.uri.path().ends_with(":generateContent")
{
return Some("gemini_chat_sync");
}
if decision.route_family.as_deref() == Some("gemini")
&& decision.route_kind.as_deref() == Some("cli")
&& parts.method == http::Method::POST
&& parts.uri.path().ends_with(":generateContent")
{
return Some("gemini_cli_sync");
}
if decision.route_family.as_deref() == Some("gemini")
&& decision.route_kind.as_deref() == Some("files")
{
if parts.method == http::Method::POST && parts.uri.path() == "/upload/v1beta/files" {
return Some("gemini_files_upload");
}
if parts.method == http::Method::GET && parts.uri.path() == "/v1beta/files" {
return Some("gemini_files_list");
}
if parts.method == http::Method::GET
&& parts.uri.path().starts_with("/v1beta/files/")
&& !parts.uri.path().ends_with(":download")
{
return Some("gemini_files_get");
}
if parts.method == http::Method::DELETE
&& parts.uri.path().starts_with("/v1beta/files/")
&& !parts.uri.path().ends_with(":download")
{
return Some("gemini_files_delete");
}
}
None
}
pub(crate) fn is_matching_stream_request(
plan_kind: &str,
parts: &http::request::Parts,
body_json: &serde_json::Value,
) -> bool {
match plan_kind {
"openai_chat_stream"
| "claude_chat_stream"
| "openai_cli_stream"
| "openai_compact_stream"
| "claude_cli_stream" => body_json
.get("stream")
.and_then(|value| value.as_bool())
.unwrap_or(false),
"gemini_chat_stream" | "gemini_cli_stream" => {
parts.uri.path().ends_with(":streamGenerateContent")
}
_ => true,
}
}
pub(crate) fn supports_sync_scheduler_decision_kind(plan_kind: &str) -> bool {
matches!(
plan_kind,
"openai_chat_sync"
| "openai_cli_sync"
| "openai_compact_sync"
| "claude_chat_sync"
| "claude_cli_sync"
| "gemini_chat_sync"
| "gemini_cli_sync"
| "gemini_files_upload"
| "openai_video_create_sync"
| "openai_video_remix_sync"
| "openai_video_cancel_sync"
| "openai_video_delete_sync"
| "gemini_video_create_sync"
| "gemini_video_cancel_sync"
| "gemini_files_get"
| "gemini_files_list"
| "gemini_files_delete"
)
}
pub(crate) fn supports_stream_scheduler_decision_kind(plan_kind: &str) -> bool {
matches!(
plan_kind,
"openai_chat_stream"
| "claude_chat_stream"
| "gemini_chat_stream"
| "openai_cli_stream"
| "openai_compact_stream"
| "claude_cli_stream"
| "gemini_cli_stream"
| "gemini_files_download"
| "openai_video_content"
)
}
#[cfg(test)]
mod tests {
use axum::http::{Method, Request};
use super::{
is_matching_stream_request, resolve_execution_runtime_stream_plan_kind,
resolve_execution_runtime_sync_plan_kind, supports_stream_scheduler_decision_kind,
supports_sync_scheduler_decision_kind,
};
use crate::control::GatewayControlDecision;
fn sample_decision(route_family: &str, route_kind: &str) -> GatewayControlDecision {
GatewayControlDecision {
public_path: "/".to_string(),
public_query_string: None,
route_class: Some("ai_public".to_string()),
route_family: Some(route_family.to_string()),
route_kind: Some(route_kind.to_string()),
auth_context: None,
admin_principal: None,
auth_endpoint_signature: None,
execution_runtime_candidate: true,
local_auth_rejection: None,
}
}
#[test]
fn resolves_openai_chat_plan_kinds() {
let request = Request::builder()
.method(Method::POST)
.uri("/v1/chat/completions")
.body(())
.expect("request should build");
let (parts, _) = request.into_parts();
let decision = sample_decision("openai", "chat");
assert_eq!(
resolve_execution_runtime_sync_plan_kind(&parts, &decision),
Some("openai_chat_sync")
);
assert_eq!(
resolve_execution_runtime_stream_plan_kind(&parts, &decision),
Some("openai_chat_stream")
);
}
#[test]
fn stream_matching_requires_openai_stream_flag() {
let request = Request::builder()
.method(Method::POST)
.uri("/v1/chat/completions")
.body(())
.expect("request should build");
let (parts, _) = request.into_parts();
assert!(!is_matching_stream_request(
"openai_chat_stream",
&parts,
&serde_json::json!({"stream": false}),
));
assert!(is_matching_stream_request(
"openai_chat_stream",
&parts,
&serde_json::json!({"stream": true}),
));
assert!(supports_sync_scheduler_decision_kind("openai_chat_sync"));
assert!(supports_stream_scheduler_decision_kind(
"openai_chat_stream"
));
}
}

32
apps/aether-gateway/src/scheduler/candidate/state.rs → apps/aether-gateway/src/scheduler/state.rs
View File

@@ -1,33 +1,15 @@
use std::time::Duration;
use aether_data::repository::candidate_selection::StoredMinimalCandidateSelectionRow;
use aether_data::repository::candidates::StoredRequestCandidate;
use aether_data::repository::provider_catalog::{
use aether_data_contracts::repository::candidates::StoredRequestCandidate;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_data::repository::quota::StoredProviderQuotaSnapshot;
use aether_data::DataLayerError;
use aether_data_contracts::repository::quota::StoredProviderQuotaSnapshot;
use aether_scheduler_core::SchedulerAffinityTarget;
use async_trait::async_trait;
use crate::GatewayError;
use super::GatewayMinimalCandidateSelectionCandidate;
#[async_trait]
pub(crate) trait SchedulerCandidateSelectionRowSource {
async fn read_minimal_candidate_selection_rows_for_api_format_and_global_model(
&self,
api_format: &str,
global_model_name: &str,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError>;
async fn read_minimal_candidate_selection_rows_for_api_format(
&self,
api_format: &str,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError>;
}
#[async_trait]
pub(crate) trait SchedulerRuntimeState {
async fn read_provider_quota_snapshot(
@@ -50,14 +32,6 @@ pub(crate) trait SchedulerRuntimeState {
limit: usize,
) -> Result<Vec<StoredRequestCandidate>, GatewayError>;
async fn read_minimal_candidate_selection(
&self,
api_format: &str,
global_model_name: &str,
require_streaming: bool,
auth_snapshot: Option<&crate::data::auth::GatewayAuthApiKeySnapshot>,
) -> Result<Vec<GatewayMinimalCandidateSelectionCandidate>, GatewayError>;
fn provider_key_rpm_reset_at(&self, key_id: &str, now_unix_secs: u64) -> Option<u64>;
fn read_cached_scheduler_affinity_target(