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

refactor: 拆分 gateway 单体为独立 crate，新增 systemd 部署方案

Browse Source 
将 gateway 内部的 model-fetch、provider-transport、scheduler-core、
usage-runtime、video-tasks-core 模块提取为独立 crate；重构 gateway
内部模块结构（state/router/cache/data/query 等）；移除大量遗留模块
文件；新增 systemd 二进制部署骨架及相关文档；更新前端 usage 相关
API 和组件。
This commit is contained in:
fawney19
2026-04-05 20:23:16 +08:00
parent cbc811f6ce
commit 763ff03a7b
777 changed files with 42659 additions and 21469 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1,16 +1,17 @@
use sha2::{Digest, Sha256};
pub(super) use aether_scheduler_core::{
build_scheduler_affinity_cache_key_for_api_key_id, candidate_affinity_hash, candidate_key,
compare_affinity_order, matches_affinity_target, SchedulerAffinityTarget,
};
use crate::gateway::gateway_cache::SchedulerAffinityTarget;
use crate::gateway::gateway_data::StoredGatewayAuthApiKeySnapshot;
use crate::gateway::AppState;
use crate::data::auth::GatewayAuthApiKeySnapshot;
use super::{
normalize_api_format, GatewayMinimalCandidateSelectionCandidate,
GatewayMinimalCandidateSelectionCandidate, SchedulerRuntimeState,
SCHEDULER_AFFINITY_MAX_ENTRIES, SCHEDULER_AFFINITY_TTL,
};
pub(super) fn build_scheduler_affinity_cache_key(
auth_snapshot: Option<&StoredGatewayAuthApiKeySnapshot>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
api_format: &str,
global_model_name: &str,
) -> Option<String> {
@@ -20,87 +21,18 @@ pub(super) fn build_scheduler_affinity_cache_key(
build_scheduler_affinity_cache_key_for_api_key_id(api_key_id, api_format, global_model_name)
}
pub(super) fn build_scheduler_affinity_cache_key_for_api_key_id(
api_key_id: &str,
api_format: &str,
global_model_name: &str,
) -> Option<String> {
let api_key_id = api_key_id.trim();
if api_key_id.is_empty() {
return None;
}
let api_format = normalize_api_format(api_format);
let global_model_name = global_model_name.trim();
if api_format.is_empty() || global_model_name.is_empty() {
return None;
}
Some(format!(
"scheduler_affinity:{api_key_id}:{api_format}:{global_model_name}"
))
}
pub(super) fn compare_affinity_order(
left: &GatewayMinimalCandidateSelectionCandidate,
right: &GatewayMinimalCandidateSelectionCandidate,
affinity_key: Option<&str>,
) -> std::cmp::Ordering {
let Some(affinity_key) = affinity_key else {
return std::cmp::Ordering::Equal;
};
candidate_affinity_hash(affinity_key, left).cmp(&candidate_affinity_hash(affinity_key, right))
}
pub(super) fn candidate_affinity_hash(
affinity_key: &str,
candidate: &GatewayMinimalCandidateSelectionCandidate,
) -> u64 {
let mut hasher = Sha256::new();
hasher.update(affinity_key.as_bytes());
hasher.update(b":");
hasher.update(candidate.provider_id.as_bytes());
hasher.update(b":");
hasher.update(candidate.endpoint_id.as_bytes());
hasher.update(b":");
hasher.update(candidate.key_id.as_bytes());
let digest = hasher.finalize();
u64::from_be_bytes([
digest[0], digest[1], digest[2], digest[3], digest[4], digest[5], digest[6], digest[7],
])
}
pub(super) fn matches_affinity_target(
candidate: &GatewayMinimalCandidateSelectionCandidate,
target: &SchedulerAffinityTarget,
) -> bool {
candidate.provider_id == target.provider_id
&& candidate.endpoint_id == target.endpoint_id
&& candidate.key_id == target.key_id
}
pub(super) fn candidate_key(
candidate: &GatewayMinimalCandidateSelectionCandidate,
) -> (String, String, String) {
(
candidate.provider_id.clone(),
candidate.endpoint_id.clone(),
candidate.key_id.clone(),
)
}
#[cfg_attr(not(test), allow(dead_code))]
pub(super) fn remember_scheduler_affinity(
affinity_cache_key: Option<&str>,
state: &AppState,
state: &(impl SchedulerRuntimeState + ?Sized),
candidate: &GatewayMinimalCandidateSelectionCandidate,
) {
let Some(cache_key) = affinity_cache_key else {
return;
};
state.scheduler_affinity_cache.insert(
cache_key.to_string(),
state.remember_scheduler_affinity_target(
cache_key,
SchedulerAffinityTarget {
provider_id: candidate.provider_id.clone(),
endpoint_id: candidate.endpoint_id.clone(),

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

@@ -1,23 +1,25 @@
use self::affinity::{
build_scheduler_affinity_cache_key, build_scheduler_affinity_cache_key_for_api_key_id,
candidate_affinity_hash, candidate_key, compare_affinity_order, matches_affinity_target,
remember_scheduler_affinity,
candidate_affinity_hash, candidate_key, remember_scheduler_affinity,
};
use self::model::{
auth_snapshot_allows_api_format, auth_snapshot_allows_model, auth_snapshot_allows_provider,
candidate_model_names, candidate_supports_required_capability,
extract_global_priority_for_format, matches_model_mapping, normalize_api_format,
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,
row_supports_required_capability, select_provider_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,
};
mod affinity;
mod model;
mod selection;
mod state;
#[cfg(test)]
mod tests;
@@ -34,54 +36,31 @@ use aether_data::repository::provider_catalog::{
};
use aether_data::repository::quota::StoredProviderQuotaSnapshot;
use aether_data::DataLayerError;
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,
};
use aether_wallet::{ProviderBillingType, ProviderQuotaSnapshot};
use regex::Regex;
use sha2::{Digest, Sha256};
use crate::gateway::gateway_cache::SchedulerAffinityTarget;
use crate::gateway::gateway_data::{GatewayDataState, StoredGatewayAuthApiKeySnapshot};
use crate::gateway::{AppState, GatewayError};
use super::health::{
count_recent_active_requests_for_api_key, count_recent_active_requests_for_provider,
effective_provider_key_health_score, is_candidate_in_recent_failure_cooldown,
is_provider_key_circuit_open, provider_key_health_bucket, provider_key_health_score,
provider_key_rpm_allows_request_since,
};
use crate::data::auth::GatewayAuthApiKeySnapshot;
use crate::{AppState, 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;
#[allow(dead_code)]
#[derive(Debug, Clone, PartialEq, serde::Serialize)]
pub(crate) struct GatewayMinimalCandidateSelectionCandidate {
pub(crate) provider_id: String,
pub(crate) provider_name: String,
pub(crate) provider_type: String,
pub(crate) provider_priority: i32,
pub(crate) endpoint_id: String,
pub(crate) endpoint_api_format: String,
pub(crate) key_id: String,
pub(crate) key_name: String,
pub(crate) key_auth_type: String,
pub(crate) key_internal_priority: i32,
pub(crate) key_global_priority_for_format: Option<i32>,
pub(crate) key_capabilities: Option<serde_json::Value>,
pub(crate) model_id: String,
pub(crate) global_model_id: String,
pub(crate) global_model_name: String,
pub(crate) selected_provider_model_name: String,
pub(crate) mapping_matched_model: Option<String>,
}
pub(crate) use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate as GatewayMinimalCandidateSelectionCandidate;
#[allow(dead_code)]
pub(crate) async fn read_minimal_candidate_selection(
state: &GatewayDataState,
state: &(impl SchedulerCandidateSelectionRowSource + Sync),
api_format: &str,
requested_model_name: &str,
require_streaming: bool,
auth_snapshot: Option<&StoredGatewayAuthApiKeySnapshot>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
) -> Result<Vec<GatewayMinimalCandidateSelectionCandidate>, DataLayerError> {
let normalized_api_format = normalize_api_format(api_format);
if normalized_api_format.is_empty() {
@@ -97,73 +76,19 @@ pub(crate) async fn read_minimal_candidate_selection(
else {
return Ok(Vec::new());
};
if !auth_snapshot_allows_model(
auth_snapshot,
requested_model_name,
resolved_global_model_name.as_str(),
) {
return Ok(Vec::new());
}
let mut candidates = Vec::new();
for row in rows {
if !auth_snapshot_allows_provider(auth_snapshot, &row.provider_id, &row.provider_name) {
continue;
}
if require_streaming && !row.supports_streaming() {
continue;
}
let Some((selected_provider_model_name, mapping_matched_model)) =
resolve_provider_model_name(&row, requested_model_name, &normalized_api_format)
else {
continue;
};
candidates.push(GatewayMinimalCandidateSelectionCandidate {
provider_id: row.provider_id,
provider_name: row.provider_name,
provider_type: row.provider_type,
provider_priority: row.provider_priority,
endpoint_id: row.endpoint_id,
endpoint_api_format: row.endpoint_api_format,
key_id: row.key_id,
key_name: row.key_name,
key_auth_type: row.key_auth_type,
key_internal_priority: row.key_internal_priority,
key_global_priority_for_format: extract_global_priority_for_format(
row.key_global_priority_by_format.as_ref(),
&normalized_api_format,
)?,
key_capabilities: row.key_capabilities,
model_id: row.model_id,
global_model_id: row.global_model_id,
global_model_name: row.global_model_name,
selected_provider_model_name,
mapping_matched_model,
});
}
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());
candidates.sort_by(|left, right| {
left.key_global_priority_for_format
.unwrap_or(i32::MAX)
.cmp(&right.key_global_priority_for_format.unwrap_or(i32::MAX))
.then_with(|| compare_affinity_order(left, right, affinity_key))
.then(left.provider_priority.cmp(&right.provider_priority))
.then(left.key_internal_priority.cmp(&right.key_internal_priority))
.then(left.provider_id.cmp(&right.provider_id))
.then(left.endpoint_id.cmp(&right.endpoint_id))
.then(left.key_id.cmp(&right.key_id))
.then(
left.selected_provider_model_name
.cmp(&right.selected_provider_model_name),
)
});
Ok(candidates)
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))]
@@ -172,7 +97,7 @@ pub(crate) async fn select_minimal_candidate(
api_format: &str,
global_model_name: &str,
require_streaming: bool,
auth_snapshot: Option<&StoredGatewayAuthApiKeySnapshot>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
now_unix_secs: u64,
) -> Result<Option<GatewayMinimalCandidateSelectionCandidate>, GatewayError> {
let affinity_cache_key =
@@ -199,7 +124,7 @@ pub(crate) async fn list_selectable_candidates(
api_format: &str,
global_model_name: &str,
require_streaming: bool,
auth_snapshot: Option<&StoredGatewayAuthApiKeySnapshot>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
now_unix_secs: u64,
) -> Result<Vec<GatewayMinimalCandidateSelectionCandidate>, GatewayError> {
collect_selectable_candidates(
@@ -218,7 +143,7 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
candidate_api_format: &str,
required_capability: &str,
require_streaming: bool,
auth_snapshot: Option<&StoredGatewayAuthApiKeySnapshot>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
now_unix_secs: u64,
) -> Result<Vec<GatewayMinimalCandidateSelectionCandidate>, GatewayError> {
let normalized_api_format = normalize_api_format(candidate_api_format);
@@ -232,28 +157,17 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
}
let rows = state
.list_minimal_candidate_selection_rows_for_api_format(&normalized_api_format)
.await?;
let mut model_names = BTreeSet::new();
for row in rows {
if !auth_snapshot_allows_provider(auth_snapshot, &row.provider_id, &row.provider_name) {
continue;
}
if !row_supports_required_capability(&row, required_capability) {
continue;
}
if require_streaming && !row.supports_streaming() {
continue;
}
if !auth_snapshot_allows_model(
auth_snapshot,
&row.global_model_name,
&row.global_model_name,
) {
continue;
}
model_names.insert(row.global_model_name);
}
.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,
&normalized_api_format,
required_capability,
require_streaming,
auth_constraints.as_ref(),
);
for global_model_name in model_names {
let candidates = list_selectable_candidates(
@@ -290,9 +204,7 @@ pub(crate) fn read_cached_scheduler_affinity_target(
api_format,
global_model_name,
)?;
state
.scheduler_affinity_cache
.get_fresh(&cache_key, SCHEDULER_AFFINITY_TTL)
state.read_cached_scheduler_affinity_target(&cache_key, SCHEDULER_AFFINITY_TTL)
}
async fn collect_selectable_candidates(
@@ -300,7 +212,7 @@ async fn collect_selectable_candidates(
api_format: &str,
global_model_name: &str,
require_streaming: bool,
auth_snapshot: Option<&StoredGatewayAuthApiKeySnapshot>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
now_unix_secs: u64,
) -> Result<Vec<GatewayMinimalCandidateSelectionCandidate>, GatewayError> {
let mut candidates = state
@@ -322,9 +234,7 @@ async fn collect_selectable_candidates(
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
.scheduler_affinity_cache
.get_fresh(cache_key, SCHEDULER_AFFINITY_TTL)
state.read_cached_scheduler_affinity_target(cache_key, SCHEDULER_AFFINITY_TTL)
});
if let Some((api_key_id, limit)) = auth_snapshot.and_then(|snapshot| {
@@ -337,45 +247,21 @@ async fn collect_selectable_candidates(
Some((snapshot.api_key_id.as_str(), limit))
})
}) {
let active_requests =
count_recent_active_requests_for_api_key(&recent_candidates, api_key_id, now_unix_secs);
if active_requests >= limit {
if auth_api_key_concurrency_limit_reached(
&recent_candidates,
now_unix_secs,
api_key_id,
limit,
) {
return Ok(Vec::new());
}
}
let mut selected = Vec::new();
let mut selected_keys = BTreeSet::new();
if let Some(target) = cached_affinity_target.as_ref() {
if let Some(candidate) = candidates
.iter()
.find(|candidate| matches_affinity_target(candidate, target))
.cloned()
{
if is_candidate_selectable(
&candidate,
&recent_candidates,
&provider_concurrent_limits,
&provider_key_rpm_states,
now_unix_secs,
cached_affinity_target.as_ref(),
state,
)
.await?
{
selected_keys.insert(candidate_key(&candidate));
selected.push(candidate);
}
}
}
for candidate in candidates {
if selected_keys.contains(&candidate_key(&candidate)) {
continue;
}
for candidate in &candidates {
if !is_candidate_selectable(
&candidate,
candidate,
&recent_candidates,
&provider_concurrent_limits,
&provider_key_rpm_states,
@@ -387,9 +273,12 @@ async fn collect_selectable_candidates(
{
continue;
}
selected_keys.insert(candidate_key(&candidate));
selected.push(candidate);
selected_keys.insert(candidate_key(candidate));
}
Ok(selected)
Ok(collect_selectable_candidates_from_keys(
candidates,
&selected_keys,
cached_affinity_target.as_ref(),
))
}

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

@@ -1,77 +1,69 @@
use std::collections::BTreeSet;
use aether_data::repository::candidate_selection::{
StoredMinimalCandidateSelectionRow, StoredProviderModelMapping,
};
use aether_data::repository::candidate_selection::StoredMinimalCandidateSelectionRow;
use aether_data::DataLayerError;
use regex::Regex;
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::gateway::gateway_data::{GatewayDataState, StoredGatewayAuthApiKeySnapshot};
use crate::data::auth::GatewayAuthApiKeySnapshot;
use super::state::SchedulerCandidateSelectionRowSource;
use super::GatewayMinimalCandidateSelectionCandidate;
pub(super) fn auth_snapshot_allows_provider(
auth_snapshot: Option<&StoredGatewayAuthApiKeySnapshot>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
provider_id: &str,
provider_name: &str,
) -> bool {
let Some(allowed) =
auth_snapshot.and_then(StoredGatewayAuthApiKeySnapshot::effective_allowed_providers)
else {
return true;
};
allowed.iter().any(|value| {
value.trim().eq_ignore_ascii_case(provider_id.trim())
|| value.trim().eq_ignore_ascii_case(provider_name.trim())
})
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<&StoredGatewayAuthApiKeySnapshot>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
api_format: &str,
) -> bool {
let Some(allowed) =
auth_snapshot.and_then(StoredGatewayAuthApiKeySnapshot::effective_allowed_api_formats)
else {
return true;
};
allowed
.iter()
.any(|value| normalize_api_format(value) == api_format)
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<&StoredGatewayAuthApiKeySnapshot>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
requested_model_name: &str,
resolved_global_model_name: &str,
) -> bool {
let Some(allowed) =
auth_snapshot.and_then(StoredGatewayAuthApiKeySnapshot::effective_allowed_models)
else {
return true;
};
allowed
.iter()
.any(|value| value == requested_model_name || value == resolved_global_model_name)
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: &GatewayDataState,
state: &(impl SchedulerCandidateSelectionRowSource + Sync),
api_format: &str,
requested_model_name: &str,
) -> Result<Option<(String, Vec<StoredMinimalCandidateSelectionRow>)>, DataLayerError> {
let exact_rows = state
.list_minimal_candidate_selection_rows(api_format, requested_model_name)
.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
.list_minimal_candidate_selection_rows_for_api_format(api_format)
.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)
@@ -87,262 +79,18 @@ pub(super) async fn read_requested_model_rows(
)))
}
pub(super) fn resolve_requested_global_model_name(
rows: &[StoredMinimalCandidateSelectionRow],
requested_model_name: &str,
api_format: &str,
) -> Option<String> {
resolve_global_model_name_by(rows, |row| {
row.model_provider_model_name == requested_model_name
})
.or_else(|| {
resolve_global_model_name_by(rows, |row| {
row.model_provider_model_mappings
.as_ref()
.is_some_and(|mappings| {
mappings.iter().any(|mapping| {
mapping_scope_matches(mapping, api_format)
&& mapping.name == requested_model_name
})
})
})
})
.or_else(|| {
resolve_global_model_name_by(rows, |row| {
row.global_model_mappings.as_ref().is_some_and(|patterns| {
patterns
.iter()
.any(|pattern| matches_model_mapping(pattern, requested_model_name))
})
})
})
}
fn resolve_global_model_name_by<F>(
rows: &[StoredMinimalCandidateSelectionRow],
matches: F,
) -> Option<String>
where
F: Fn(&StoredMinimalCandidateSelectionRow) -> bool,
{
let mut matches = rows
.iter()
.filter(|row| matches(row))
.map(|row| row.global_model_name.trim())
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.collect::<BTreeSet<_>>()
.into_iter();
matches.next()
}
pub(super) fn resolve_provider_model_name(
row: &StoredMinimalCandidateSelectionRow,
requested_model_name: &str,
api_format: &str,
) -> Option<(String, Option<String>)> {
let selected_provider_model_name = select_provider_model_name(row, api_format);
let Some(key_allowed_models) = row.key_allowed_models.as_ref() else {
return Some((selected_provider_model_name, None));
};
if key_allowed_models.is_empty() {
return None;
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()),
}
if key_allowed_models
.iter()
.any(|value| value == requested_model_name)
{
return Some((selected_provider_model_name, None));
}
let candidate_models = candidate_model_names(row, api_format);
let mut sorted_allowed_models = key_allowed_models
.iter()
.map(|value| value.trim())
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.collect::<Vec<_>>();
sorted_allowed_models.sort();
for allowed_model in &sorted_allowed_models {
if candidate_models.contains(allowed_model.as_str()) {
return Some((allowed_model.clone(), Some(allowed_model.clone())));
}
}
let Some(global_model_mappings) = row.global_model_mappings.as_ref() else {
return None;
};
for allowed_model in sorted_allowed_models {
for pattern in global_model_mappings {
if matches_model_mapping(pattern, &allowed_model) {
return Some((allowed_model.clone(), Some(allowed_model)));
}
}
}
None
}
pub(super) fn select_provider_model_name(
row: &StoredMinimalCandidateSelectionRow,
api_format: &str,
) -> String {
let Some(mappings) = row.model_provider_model_mappings.as_ref() else {
return row.model_provider_model_name.clone();
};
let mut scoped = mappings
.iter()
.filter(|mapping| mapping_scope_matches(mapping, api_format))
.collect::<Vec<_>>();
if scoped.is_empty() {
return row.model_provider_model_name.clone();
}
scoped.sort_by(|left, right| {
left.priority
.cmp(&right.priority)
.then(left.name.cmp(&right.name))
});
let top_priority = scoped[0].priority;
scoped
.into_iter()
.find(|mapping| mapping.priority == top_priority)
.map(|mapping| mapping.name.clone())
.unwrap_or_else(|| row.model_provider_model_name.clone())
}
pub(super) fn candidate_model_names(
row: &StoredMinimalCandidateSelectionRow,
api_format: &str,
) -> BTreeSet<String> {
let mut names = BTreeSet::from([row.model_provider_model_name.clone()]);
if let Some(mappings) = row.model_provider_model_mappings.as_ref() {
for mapping in mappings {
if mapping_scope_matches(mapping, api_format) {
names.insert(mapping.name.clone());
}
}
}
names
}
fn mapping_scope_matches(mapping: &StoredProviderModelMapping, api_format: &str) -> bool {
let Some(api_formats) = mapping.api_formats.as_ref() else {
return true;
};
api_formats
.iter()
.any(|value| normalize_api_format(value) == api_format)
}
pub(super) fn row_supports_required_capability(
row: &StoredMinimalCandidateSelectionRow,
required_capability: &str,
) -> bool {
capabilities_support_required_capability(row.key_capabilities.as_ref(), required_capability)
}
pub(super) fn candidate_supports_required_capability(
candidate: &GatewayMinimalCandidateSelectionCandidate,
required_capability: &str,
) -> bool {
capabilities_support_required_capability(
candidate.key_capabilities.as_ref(),
required_capability,
)
}
fn capabilities_support_required_capability(
capabilities: Option<&serde_json::Value>,
required_capability: &str,
) -> bool {
let required_capability = required_capability.trim();
if required_capability.is_empty() {
return true;
}
let Some(capabilities) = capabilities else {
return false;
};
if let Some(object) = capabilities.as_object() {
return object.iter().any(|(key, value)| {
key.eq_ignore_ascii_case(required_capability)
&& match value {
serde_json::Value::Bool(value) => *value,
serde_json::Value::String(value) => value.eq_ignore_ascii_case("true"),
serde_json::Value::Number(value) => {
value.as_i64().is_some_and(|value| value > 0)
}
_ => false,
}
});
}
if let Some(items) = capabilities.as_array() {
return items.iter().any(|value| {
value
.as_str()
.is_some_and(|value| value.eq_ignore_ascii_case(required_capability))
});
}
false
}
pub(super) fn matches_model_mapping(pattern: &str, model_name: &str) -> bool {
let Ok(compiled) = Regex::new(&format!("^(?:{pattern})$")) else {
return false;
};
compiled.is_match(model_name)
}
pub(super) fn extract_global_priority_for_format(
raw: Option<&serde_json::Value>,
api_format: &str,
) -> Result<Option<i32>, DataLayerError> {
let Some(raw) = raw else {
return Ok(None);
};
let Some(object) = raw.as_object() else {
return Err(DataLayerError::UnexpectedValue(
"provider_api_keys.global_priority_by_format is not a JSON object".to_string(),
));
};
let Some(value) = object
.iter()
.find(|(key, _)| normalize_api_format(key) == api_format)
.map(|(_, value)| value)
else {
return Ok(None);
};
if let Some(value) = value.as_i64() {
return i32::try_from(value).map(Some).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"invalid provider_api_keys.global_priority_by_format value: {value}"
))
});
}
if let Some(value) = value.as_str() {
let value = value.trim().parse::<i32>().map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"invalid provider_api_keys.global_priority_by_format value: {value}"
))
})?;
return Ok(Some(value));
}
Err(DataLayerError::UnexpectedValue(
"provider_api_keys.global_priority_by_format contains a non-integer value".to_string(),
))
}
pub(super) fn normalize_api_format(value: &str) -> String {
value.trim().to_ascii_lowercase()
}

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

@@ -1,130 +1,34 @@
use std::collections::{BTreeMap, BTreeSet};
use aether_data::repository::candidates::StoredRequestCandidate;
use aether_data::repository::provider_catalog::{
StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_data::repository::quota::StoredProviderQuotaSnapshot;
use aether_wallet::{ProviderBillingType, ProviderQuotaSnapshot};
use crate::gateway::gateway_cache::SchedulerAffinityTarget;
use crate::gateway::gateway_data::StoredGatewayAuthApiKeySnapshot;
use crate::gateway::{AppState, GatewayError};
use super::super::health::{
count_recent_active_requests_for_provider, effective_provider_key_health_score,
is_candidate_in_recent_failure_cooldown, is_provider_key_circuit_open,
provider_key_health_bucket, provider_key_health_score, provider_key_rpm_allows_request_since,
ProviderKeyHealthBucket,
use aether_data::repository::provider_catalog::StoredProviderCatalogKey;
use aether_scheduler_core::{
build_provider_concurrent_limit_map, candidate_is_selectable_with_runtime_state,
reorder_candidates_by_scheduler_health as reorder_candidates_by_scheduler_health_in_core,
SchedulerAffinityTarget,
};
use super::{
compare_affinity_order, matches_affinity_target, GatewayMinimalCandidateSelectionCandidate,
};
use crate::data::auth::GatewayAuthApiKeySnapshot;
use crate::GatewayError;
use super::{GatewayMinimalCandidateSelectionCandidate, SchedulerRuntimeState};
pub(super) fn reorder_candidates_by_scheduler_health(
candidates: &mut [GatewayMinimalCandidateSelectionCandidate],
provider_key_rpm_states: &BTreeMap<String, StoredProviderCatalogKey>,
auth_snapshot: Option<&StoredGatewayAuthApiKeySnapshot>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
) {
let affinity_key = auth_snapshot
.map(|snapshot| snapshot.api_key_id.trim())
.filter(|value| !value.is_empty());
candidates.sort_by(|left, right| {
left.key_global_priority_for_format
.unwrap_or(i32::MAX)
.cmp(&right.key_global_priority_for_format.unwrap_or(i32::MAX))
.then_with(|| compare_provider_key_health_order(left, right, provider_key_rpm_states))
.then_with(|| compare_affinity_order(left, right, affinity_key))
.then(left.provider_priority.cmp(&right.provider_priority))
.then(left.key_internal_priority.cmp(&right.key_internal_priority))
.then(left.provider_id.cmp(&right.provider_id))
.then(left.endpoint_id.cmp(&right.endpoint_id))
.then(left.key_id.cmp(&right.key_id))
.then(
left.selected_provider_model_name
.cmp(&right.selected_provider_model_name),
)
});
reorder_candidates_by_scheduler_health_in_core(
candidates,
provider_key_rpm_states,
affinity_key,
);
}
fn compare_provider_key_health_order(
left: &GatewayMinimalCandidateSelectionCandidate,
right: &GatewayMinimalCandidateSelectionCandidate,
provider_key_rpm_states: &BTreeMap<String, StoredProviderCatalogKey>,
) -> std::cmp::Ordering {
let left_bucket = candidate_provider_key_health_bucket(left, provider_key_rpm_states);
let right_bucket = candidate_provider_key_health_bucket(right, provider_key_rpm_states);
right_bucket.cmp(&left_bucket).then_with(|| {
let left_score = candidate_provider_key_health_score(left, provider_key_rpm_states);
let right_score = candidate_provider_key_health_score(right, provider_key_rpm_states);
right_score
.partial_cmp(&left_score)
.unwrap_or(std::cmp::Ordering::Equal)
})
}
fn candidate_provider_key_health_bucket(
candidate: &GatewayMinimalCandidateSelectionCandidate,
provider_key_rpm_states: &BTreeMap<String, StoredProviderCatalogKey>,
) -> Option<ProviderKeyHealthBucket> {
provider_key_rpm_states
.get(&candidate.key_id)
.and_then(|key| provider_key_health_bucket(key, candidate.endpoint_api_format.as_str()))
}
fn candidate_provider_key_health_score(
candidate: &GatewayMinimalCandidateSelectionCandidate,
provider_key_rpm_states: &BTreeMap<String, StoredProviderCatalogKey>,
) -> f64 {
provider_key_rpm_states
.get(&candidate.key_id)
.and_then(|key| {
effective_provider_key_health_score(key, candidate.endpoint_api_format.as_str())
})
.unwrap_or(1.0)
}
pub(super) fn should_skip_provider_quota(
quota: &StoredProviderQuotaSnapshot,
now_unix_secs: u64,
) -> bool {
let snapshot = ProviderQuotaSnapshot {
provider_id: quota.provider_id.clone(),
billing_type: ProviderBillingType::parse(&quota.billing_type),
monthly_quota_usd: quota.monthly_quota_usd,
monthly_used_usd: quota.monthly_used_usd,
quota_reset_day: quota.quota_reset_day,
quota_last_reset_at_unix_secs: quota.quota_last_reset_at_unix_secs,
quota_expires_at_unix_secs: quota.quota_expires_at_unix_secs,
is_active: quota.is_active,
};
if !snapshot.is_active || snapshot.is_expired(now_unix_secs) {
return true;
}
match snapshot.billing_type {
ProviderBillingType::MonthlyQuota | ProviderBillingType::FreeTier => snapshot
.remaining_quota_usd()
.is_some_and(|remaining| remaining <= 0.0),
ProviderBillingType::PayAsYouGo | ProviderBillingType::Unknown => false,
}
}
fn is_candidate_cooled_down(
candidate: &GatewayMinimalCandidateSelectionCandidate,
recent_candidates: &[StoredRequestCandidate],
now_unix_secs: u64,
) -> bool {
is_candidate_in_recent_failure_cooldown(
recent_candidates,
candidate.provider_id.as_str(),
candidate.endpoint_id.as_str(),
candidate.key_id.as_str(),
now_unix_secs,
)
}
pub(super) use aether_scheduler_core::should_skip_provider_quota;
pub(super) async fn is_candidate_selectable(
candidate: &GatewayMinimalCandidateSelectionCandidate,
@@ -133,61 +37,29 @@ pub(super) async fn is_candidate_selectable(
provider_key_rpm_states: &BTreeMap<String, StoredProviderCatalogKey>,
now_unix_secs: u64,
cached_affinity_target: Option<&SchedulerAffinityTarget>,
state: &AppState,
state: &(impl SchedulerRuntimeState + ?Sized),
) -> Result<bool, GatewayError> {
let quota = state
let provider_quota_blocks_requests = state
.read_provider_quota_snapshot(&candidate.provider_id)
.await?;
if quota
.await?
.as_ref()
.is_some_and(|quota| should_skip_provider_quota(quota, now_unix_secs))
{
return Ok(false);
}
if is_candidate_cooled_down(candidate, recent_candidates, now_unix_secs) {
return Ok(false);
}
if provider_concurrent_limits
.get(&candidate.provider_id)
.is_some_and(|limit| {
count_recent_active_requests_for_provider(
recent_candidates,
candidate.provider_id.as_str(),
now_unix_secs,
) >= *limit
})
{
return Ok(false);
}
let is_cached_user =
cached_affinity_target.is_some_and(|target| matches_affinity_target(candidate, target));
if let Some(provider_key) = provider_key_rpm_states.get(&candidate.key_id) {
if is_provider_key_circuit_open(provider_key, candidate.endpoint_api_format.as_str()) {
return Ok(false);
}
if provider_key_health_score(provider_key, candidate.endpoint_api_format.as_str())
.is_some_and(|score| score <= 0.0)
{
return Ok(false);
}
let rpm_reset_at =
state.provider_key_rpm_reset_at(candidate.key_id.as_str(), now_unix_secs);
if !provider_key_rpm_allows_request_since(
provider_key,
recent_candidates,
now_unix_secs,
is_cached_user,
rpm_reset_at,
) {
return Ok(false);
}
}
.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(true)
Ok(candidate_is_selectable_with_runtime_state(
candidate,
recent_candidates,
provider_concurrent_limits,
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: &AppState,
state: &(impl SchedulerRuntimeState + ?Sized),
candidates: &[GatewayMinimalCandidateSelectionCandidate],
) -> Result<BTreeMap<String, usize>, GatewayError> {
let provider_ids = candidates
@@ -206,23 +78,8 @@ pub(super) async fn read_provider_concurrent_limits(
Ok(build_provider_concurrent_limit_map(providers))
}
fn build_provider_concurrent_limit_map(
providers: Vec<StoredProviderCatalogProvider>,
) -> BTreeMap<String, usize> {
providers
.into_iter()
.filter_map(|provider| {
provider
.concurrent_limit
.and_then(|limit| usize::try_from(limit).ok())
.filter(|limit| *limit > 0)
.map(|limit| (provider.id, limit))
})
.collect()
}
pub(super) async fn read_provider_key_rpm_states(
state: &AppState,
state: &(impl SchedulerRuntimeState + ?Sized),
candidates: &[GatewayMinimalCandidateSelectionCandidate],
) -> Result<BTreeMap<String, StoredProviderCatalogKey>, GatewayError> {
let key_ids = candidates

76
apps/aether-gateway/src/scheduler/candidate/state.rs Normal file
View File

@@ -0,0 +1,76 @@
use std::time::Duration;
use aether_data::repository::candidate_selection::StoredMinimalCandidateSelectionRow;
use aether_data::repository::candidates::StoredRequestCandidate;
use aether_data::repository::provider_catalog::{
StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_data::repository::quota::StoredProviderQuotaSnapshot;
use aether_data::DataLayerError;
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(
&self,
provider_id: &str,
) -> Result<Option<StoredProviderQuotaSnapshot>, GatewayError>;
async fn read_provider_catalog_providers_by_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogProvider>, GatewayError>;
async fn read_provider_catalog_keys_by_ids(
&self,
key_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKey>, GatewayError>;
async fn read_recent_request_candidates(
&self,
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(
&self,
cache_key: &str,
ttl: Duration,
) -> Option<SchedulerAffinityTarget>;
fn remember_scheduler_affinity_target(
&self,
cache_key: &str,
target: SchedulerAffinityTarget,
ttl: Duration,
max_entries: usize,
);
}

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

@@ -18,11 +18,12 @@ use super::{
build_scheduler_affinity_cache_key, candidate_affinity_hash, candidate_model_names,
matches_model_mapping, read_minimal_candidate_selection, resolve_provider_model_name,
resolve_requested_global_model_name, select_minimal_candidate, select_provider_model_name,
should_skip_provider_quota, GatewayMinimalCandidateSelectionCandidate,
StoredGatewayAuthApiKeySnapshot,
should_skip_provider_quota, GatewayAuthApiKeySnapshot,
GatewayMinimalCandidateSelectionCandidate,
};
use crate::gateway::gateway_cache::SchedulerAffinityTarget;
use crate::gateway::{AppState, GatewayDataState};
use crate::cache::SchedulerAffinityTarget;
use crate::data::GatewayDataState;
use crate::AppState;
fn sample_row() -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
@@ -457,7 +458,7 @@ async fn same_priority_candidates_are_distributed_by_affinity_key() {
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
let state = GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas);
let auth_snapshot = StoredGatewayAuthApiKeySnapshot {
let auth_snapshot = GatewayAuthApiKeySnapshot {
user_id: "user-1".to_string(),
username: "alice".to_string(),
email: None,
@@ -546,7 +547,7 @@ async fn read_minimal_candidate_selection_allows_resolved_global_model_in_auth_s
]));
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
let state = GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas);
let auth_snapshot = StoredGatewayAuthApiKeySnapshot {
let auth_snapshot = GatewayAuthApiKeySnapshot {
user_id: "user-1".to_string(),
username: "alice".to_string(),
email: None,
@@ -614,7 +615,7 @@ async fn reuses_cached_scheduler_affinity_candidate_before_sorted_fallback() {
GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas),
);
let auth_snapshot = StoredGatewayAuthApiKeySnapshot {
let auth_snapshot = GatewayAuthApiKeySnapshot {
user_id: "user-1".to_string(),
username: "alice".to_string(),
email: None,
@@ -799,7 +800,7 @@ async fn returns_none_when_auth_api_key_concurrent_limit_is_reached() {
),
);
let auth_snapshot = StoredGatewayAuthApiKeySnapshot {
let auth_snapshot = GatewayAuthApiKeySnapshot {
user_id: "user-1".to_string(),
username: "alice".to_string(),
email: None,
@@ -912,7 +913,7 @@ async fn selects_next_candidate_when_first_provider_key_rpm_slots_are_reserved_f
),
);
let auth_snapshot = StoredGatewayAuthApiKeySnapshot {
let auth_snapshot = GatewayAuthApiKeySnapshot {
user_id: "user-1".to_string(),
username: "alice".to_string(),
email: None,
@@ -1027,7 +1028,7 @@ async fn cached_affinity_candidate_can_use_reserved_provider_key_rpm_capacity()
),
);
let auth_snapshot = StoredGatewayAuthApiKeySnapshot {
let auth_snapshot = GatewayAuthApiKeySnapshot {
user_id: "user-1".to_string(),
username: "alice".to_string(),
email: None,

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

@@ -1,954 +1,9 @@
use aether_data::repository::candidates::{RequestCandidateStatus, StoredRequestCandidate};
use aether_data::repository::provider_catalog::StoredProviderCatalogKey;
const FAILURE_COOLDOWN_WINDOW_SECS: u64 = 60;
const FAILURE_COOLDOWN_THRESHOLD: usize = 2;
const ACTIVE_REQUEST_WINDOW_SECS: u64 = 300;
pub(crate) const PROVIDER_KEY_RPM_WINDOW_SECS: u64 = 60;
const PROBE_PHASE_REQUESTS: u32 = 100;
const PROBE_RESERVATION_RATIO: f64 = 0.1;
const STABLE_MIN_RESERVATION_RATIO: f64 = 0.1;
const STABLE_MAX_RESERVATION_RATIO: f64 = 0.35;
const SUCCESS_COUNT_FOR_FULL_CONFIDENCE: u32 = 50;
const COOLDOWN_HOURS_FOR_FULL_CONFIDENCE: f64 = 24.0;
const LOW_LOAD_THRESHOLD: f64 = 0.5;
const HIGH_LOAD_THRESHOLD: f64 = 0.8;
const ENFORCEMENT_CONFIDENCE_THRESHOLD: f64 = 0.6;
const HEALTH_DEGRADED_THRESHOLD: f64 = 0.8;
const HEALTH_LOW_THRESHOLD: f64 = 0.5;
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub(crate) enum ProviderKeyHealthBucket {
Low,
Degraded,
Healthy,
}
impl ProviderKeyHealthBucket {
fn from_score(score: f64) -> Self {
let score = score.clamp(0.0, 1.0);
if score < HEALTH_LOW_THRESHOLD {
return Self::Low;
}
if score < HEALTH_DEGRADED_THRESHOLD {
return Self::Degraded;
}
Self::Healthy
}
}
pub(crate) fn is_candidate_in_recent_failure_cooldown(
recent_candidates: &[StoredRequestCandidate],
provider_id: &str,
endpoint_id: &str,
key_id: &str,
now_unix_secs: u64,
) -> bool {
let mut recent_failures = 0usize;
for candidate in recent_candidates {
if candidate.provider_id.as_deref() != Some(provider_id)
|| candidate.endpoint_id.as_deref() != Some(endpoint_id)
|| candidate.key_id.as_deref() != Some(key_id)
{
continue;
}
let observed_at_unix_secs = candidate
.finished_at_unix_secs
.or(candidate.started_at_unix_secs)
.unwrap_or(candidate.created_at_unix_secs);
if now_unix_secs.saturating_sub(observed_at_unix_secs) > FAILURE_COOLDOWN_WINDOW_SECS {
continue;
}
match candidate.status {
RequestCandidateStatus::Success => return false,
RequestCandidateStatus::Failed | RequestCandidateStatus::Cancelled => {
recent_failures += 1;
if recent_failures >= FAILURE_COOLDOWN_THRESHOLD {
return true;
}
}
RequestCandidateStatus::Available
| RequestCandidateStatus::Unused
| RequestCandidateStatus::Pending
| RequestCandidateStatus::Streaming
| RequestCandidateStatus::Skipped => {}
}
}
false
}
pub(crate) fn count_recent_active_requests_for_provider(
recent_candidates: &[StoredRequestCandidate],
provider_id: &str,
now_unix_secs: u64,
) -> usize {
recent_candidates
.iter()
.filter(|candidate| candidate.provider_id.as_deref() == Some(provider_id))
.filter(|candidate| is_recently_active(candidate, now_unix_secs))
.count()
}
pub(crate) fn count_recent_active_requests_for_api_key(
recent_candidates: &[StoredRequestCandidate],
api_key_id: &str,
now_unix_secs: u64,
) -> usize {
recent_candidates
.iter()
.filter(|candidate| candidate.api_key_id.as_deref() == Some(api_key_id))
.filter(|candidate| is_recently_active(candidate, now_unix_secs))
.count()
}
pub(crate) fn effective_provider_key_rpm_limit(
key: &StoredProviderCatalogKey,
now_unix_secs: u64,
) -> Option<usize> {
if let Some(limit) = key.rpm_limit.filter(|limit| *limit > 0) {
return usize::try_from(limit).ok();
}
let learned_limit = key
.learned_rpm_limit
.filter(|limit| *limit > 0)
.and_then(|limit| usize::try_from(limit).ok())?;
if provider_key_reservation_confidence(key, now_unix_secs) < ENFORCEMENT_CONFIDENCE_THRESHOLD {
return None;
}
Some(learned_limit)
}
pub(crate) fn count_recent_rpm_requests_for_provider_key(
recent_candidates: &[StoredRequestCandidate],
key_id: &str,
now_unix_secs: u64,
) -> usize {
count_recent_rpm_requests_for_provider_key_since(recent_candidates, key_id, now_unix_secs, None)
}
pub(crate) fn count_recent_rpm_requests_for_provider_key_since(
recent_candidates: &[StoredRequestCandidate],
key_id: &str,
now_unix_secs: u64,
reset_after_unix_secs: Option<u64>,
) -> usize {
let mut attempted_count = 0usize;
let mut max_observed = 0usize;
for candidate in recent_candidates {
if candidate.key_id.as_deref() != Some(key_id) {
continue;
}
if !is_recent_rpm_observation(candidate, now_unix_secs) {
continue;
}
let observed_at_unix_secs = candidate
.started_at_unix_secs
.unwrap_or(candidate.created_at_unix_secs);
if reset_after_unix_secs.is_some_and(|reset_after| observed_at_unix_secs <= reset_after) {
continue;
}
attempted_count += 1;
max_observed = max_observed.max(candidate.concurrent_requests.unwrap_or_default() as usize);
}
max_observed.max(attempted_count)
}
pub(crate) fn provider_key_rpm_allows_request(
key: &StoredProviderCatalogKey,
recent_candidates: &[StoredRequestCandidate],
now_unix_secs: u64,
is_cached_user: bool,
) -> bool {
provider_key_rpm_allows_request_since(
key,
recent_candidates,
now_unix_secs,
is_cached_user,
None,
)
}
pub(crate) fn provider_key_rpm_allows_request_since(
key: &StoredProviderCatalogKey,
recent_candidates: &[StoredRequestCandidate],
now_unix_secs: u64,
is_cached_user: bool,
reset_after_unix_secs: Option<u64>,
) -> bool {
let Some(effective_limit) = effective_provider_key_rpm_limit(key, now_unix_secs) else {
return true;
};
if effective_limit == 0 {
return false;
}
let current_usage = count_recent_rpm_requests_for_provider_key_since(
recent_candidates,
key.id.as_str(),
now_unix_secs,
reset_after_unix_secs,
);
if is_cached_user {
return current_usage < effective_limit;
}
let available_for_new = available_provider_key_rpm_slots_for_new_user(
key,
current_usage,
effective_limit,
now_unix_secs,
);
current_usage < available_for_new
}
pub(crate) fn provider_key_health_score(
key: &StoredProviderCatalogKey,
api_format: &str,
) -> Option<f64> {
let score = key
.health_by_format
.as_ref()
.and_then(serde_json::Value::as_object)
.and_then(|values| values.get(api_format))
.and_then(serde_json::Value::as_object)
.and_then(|payload| payload.get("health_score"))
.and_then(json_value_as_f64)?;
Some(score.clamp(0.0, 1.0))
}
pub(crate) fn aggregate_provider_key_health_score(key: &StoredProviderCatalogKey) -> Option<f64> {
let health_by_format = key.health_by_format.as_ref()?.as_object()?;
let mut scores = Vec::new();
for payload in health_by_format.values() {
let Some(score) = payload
.as_object()
.and_then(|payload| payload.get("health_score"))
.and_then(json_value_as_f64)
else {
continue;
};
scores.push(score.clamp(0.0, 1.0));
}
scores.into_iter().reduce(f64::min)
}
pub(crate) fn effective_provider_key_health_score(
key: &StoredProviderCatalogKey,
api_format: &str,
) -> Option<f64> {
provider_key_health_score(key, api_format).or_else(|| aggregate_provider_key_health_score(key))
}
pub(crate) fn provider_key_health_bucket(
key: &StoredProviderCatalogKey,
api_format: &str,
) -> Option<ProviderKeyHealthBucket> {
effective_provider_key_health_score(key, api_format).map(ProviderKeyHealthBucket::from_score)
}
pub(crate) fn is_provider_key_circuit_open(
key: &StoredProviderCatalogKey,
api_format: &str,
) -> bool {
key.circuit_breaker_by_format
.as_ref()
.and_then(serde_json::Value::as_object)
.and_then(|values| values.get(api_format))
.and_then(serde_json::Value::as_object)
.and_then(|payload| payload.get("open"))
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
}
fn available_provider_key_rpm_slots_for_new_user(
key: &StoredProviderCatalogKey,
current_usage: usize,
effective_limit: usize,
now_unix_secs: u64,
) -> usize {
let reservation_ratio =
provider_key_dynamic_reservation_ratio(key, current_usage, effective_limit, now_unix_secs);
usize::max(
1,
(effective_limit as f64 * (1.0 - reservation_ratio)).floor() as usize,
)
}
fn provider_key_dynamic_reservation_ratio(
key: &StoredProviderCatalogKey,
current_usage: usize,
effective_limit: usize,
now_unix_secs: u64,
) -> f64 {
let total_requests = provider_key_total_requests(key);
if total_requests < PROBE_PHASE_REQUESTS {
return PROBE_RESERVATION_RATIO;
}
let confidence = provider_key_reservation_confidence(key, now_unix_secs);
let load_ratio = provider_key_load_ratio(current_usage, effective_limit);
if load_ratio < LOW_LOAD_THRESHOLD {
return STABLE_MIN_RESERVATION_RATIO;
}
if load_ratio < HIGH_LOAD_THRESHOLD {
let load_factor =
(load_ratio - LOW_LOAD_THRESHOLD) / (HIGH_LOAD_THRESHOLD - LOW_LOAD_THRESHOLD);
return STABLE_MIN_RESERVATION_RATIO
+ confidence
* load_factor
* (STABLE_MAX_RESERVATION_RATIO - STABLE_MIN_RESERVATION_RATIO);
}
STABLE_MIN_RESERVATION_RATIO
+ confidence * (STABLE_MAX_RESERVATION_RATIO - STABLE_MIN_RESERVATION_RATIO)
}
fn is_recently_active(candidate: &StoredRequestCandidate, now_unix_secs: u64) -> bool {
if candidate.finished_at_unix_secs.is_some() {
return false;
}
if !matches!(
candidate.status,
RequestCandidateStatus::Pending | RequestCandidateStatus::Streaming
) {
return false;
}
let observed_at_unix_secs = candidate
.started_at_unix_secs
.unwrap_or(candidate.created_at_unix_secs);
now_unix_secs.saturating_sub(observed_at_unix_secs) <= ACTIVE_REQUEST_WINDOW_SECS
}
fn is_recent_rpm_observation(candidate: &StoredRequestCandidate, now_unix_secs: u64) -> bool {
if !candidate
.status
.is_attempted(candidate.started_at_unix_secs)
{
return false;
}
let observed_at_unix_secs = candidate
.started_at_unix_secs
.unwrap_or(candidate.created_at_unix_secs);
now_unix_secs.saturating_sub(observed_at_unix_secs) <= PROVIDER_KEY_RPM_WINDOW_SECS
}
fn provider_key_total_requests(key: &StoredProviderCatalogKey) -> u32 {
let request_count = key.request_count.unwrap_or_default();
if request_count > 0 {
return request_count;
}
let history_count = key
.adjustment_history
.as_ref()
.and_then(serde_json::Value::as_array)
.map(|values| values.len() as u32 * 10)
.unwrap_or_default();
key.concurrent_429_count.unwrap_or_default()
+ key.rpm_429_count.unwrap_or_default()
+ key.success_count.unwrap_or_default()
+ history_count
}
fn provider_key_load_ratio(current_usage: usize, effective_limit: usize) -> f64 {
if effective_limit == 0 {
return 0.0;
}
(current_usage as f64 / effective_limit as f64).min(1.0)
}
fn provider_key_reservation_confidence(key: &StoredProviderCatalogKey, now_unix_secs: u64) -> f64 {
let request_count = key.request_count.unwrap_or_default() as f64;
let success_count = key.success_count.unwrap_or_default() as f64;
let success_score = if request_count >= SUCCESS_COUNT_FOR_FULL_CONFIDENCE as f64 {
let success_rate = if request_count > 0.0 {
success_count / request_count
} else {
0.0
};
success_rate * 0.4
} else if request_count > 0.0 {
let success_rate = success_count / request_count;
let progress_ratio = request_count / SUCCESS_COUNT_FOR_FULL_CONFIDENCE as f64;
success_rate * progress_ratio * 0.4
} else {
0.0
};
let cooldown_score = match key.last_429_at_unix_secs {
Some(last_429_at_unix_secs) => {
let hours_since_429 =
now_unix_secs.saturating_sub(last_429_at_unix_secs) as f64 / 3600.0;
(hours_since_429 / COOLDOWN_HOURS_FOR_FULL_CONFIDENCE).min(1.0) * 0.3
}
None => 0.3,
};
let stability_score = provider_key_stability_score(key);
(success_score + cooldown_score + stability_score).min(1.0)
}
fn provider_key_stability_score(key: &StoredProviderCatalogKey) -> f64 {
let Some(history) = key
.adjustment_history
.as_ref()
.and_then(serde_json::Value::as_array)
else {
return 0.15;
};
if history.len() < 3 {
return 0.15;
}
let recent = if history.len() > 5 {
&history[history.len() - 5..]
} else {
history.as_slice()
};
let limits = recent
.iter()
.filter_map(|entry| entry.get("new_limit"))
.filter_map(json_value_as_f64)
.collect::<Vec<_>>();
if limits.len() < 2 {
return 0.15;
}
let mean = limits.iter().sum::<f64>() / limits.len() as f64;
let variance = limits
.iter()
.map(|limit| {
let delta = *limit - mean;
delta * delta
})
.sum::<f64>()
/ (limits.len() as f64 - 1.0);
let stability_ratio = (1.0 - variance / 10.0).max(0.0);
stability_ratio * 0.3
}
fn json_value_as_f64(value: &serde_json::Value) -> Option<f64> {
value
.as_f64()
.or_else(|| value.as_i64().map(|raw| raw as f64))
.or_else(|| value.as_u64().map(|raw| raw as f64))
}
#[cfg(test)]
mod tests {
use aether_data::repository::candidates::{RequestCandidateStatus, StoredRequestCandidate};
use aether_data::repository::provider_catalog::StoredProviderCatalogKey;
use super::{
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,
};
fn stored_candidate(
id: &str,
status: RequestCandidateStatus,
created_at_unix_secs: i64,
) -> StoredRequestCandidate {
StoredRequestCandidate::new(
id.to_string(),
format!("req-{id}"),
None,
None,
None,
None,
0,
0,
Some("provider-a".to_string()),
Some("endpoint-a".to_string()),
Some("key-a".to_string()),
status,
None,
false,
None,
None,
None,
None,
None,
None,
None,
created_at_unix_secs,
Some(created_at_unix_secs),
Some(created_at_unix_secs),
)
.expect("candidate should build")
}
fn provider_catalog_key(id: &str) -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
id.to_string(),
"provider-a".to_string(),
"primary".to_string(),
"api_key".to_string(),
None,
true,
)
.expect("provider key should build")
}
#[test]
fn cooldown_triggers_after_two_recent_failures() {
let recent_candidates = vec![
stored_candidate("one", RequestCandidateStatus::Failed, 95),
stored_candidate("two", RequestCandidateStatus::Cancelled, 99),
];
assert!(is_candidate_in_recent_failure_cooldown(
&recent_candidates,
"provider-a",
"endpoint-a",
"key-a",
100,
));
}
#[test]
fn recent_success_clears_cooldown() {
let recent_candidates = vec![
stored_candidate("one", RequestCandidateStatus::Failed, 95),
stored_candidate("two", RequestCandidateStatus::Success, 99),
stored_candidate("three", RequestCandidateStatus::Cancelled, 98),
];
assert!(!is_candidate_in_recent_failure_cooldown(
&recent_candidates,
"provider-a",
"endpoint-a",
"key-a",
100,
));
}
#[test]
fn counts_only_recently_active_provider_requests() {
let recent_candidates = vec![
StoredRequestCandidate::new(
"one".to_string(),
"req-one".to_string(),
None,
Some("api-key-1".to_string()),
None,
None,
0,
0,
Some("provider-a".to_string()),
Some("endpoint-a".to_string()),
Some("key-a".to_string()),
RequestCandidateStatus::Pending,
None,
false,
None,
None,
None,
None,
None,
None,
None,
95,
Some(95),
None,
)
.expect("candidate should build"),
StoredRequestCandidate::new(
"two".to_string(),
"req-two".to_string(),
None,
Some("api-key-1".to_string()),
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,
96,
Some(96),
None,
)
.expect("candidate should build"),
StoredRequestCandidate::new(
"three".to_string(),
"req-three".to_string(),
None,
Some("api-key-1".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,
None,
None,
None,
None,
None,
None,
None,
97,
Some(97),
Some(98),
)
.expect("candidate should build"),
];
assert_eq!(
count_recent_active_requests_for_provider(&recent_candidates, "provider-a", 100),
2
);
assert_eq!(
count_recent_active_requests_for_api_key(&recent_candidates, "api-key-1", 100),
2
);
}
#[test]
fn fixed_provider_key_rpm_limit_takes_precedence() {
let key = provider_catalog_key("key-a").with_rate_limit_fields(
Some(120),
Some(80),
None,
None,
None,
None,
Some(10),
Some(10),
);
assert_eq!(effective_provider_key_rpm_limit(&key, 100), Some(120));
}
#[test]
fn learned_provider_key_rpm_limit_requires_confidence() {
let low_confidence = provider_catalog_key("key-a").with_rate_limit_fields(
None,
Some(80),
Some(0),
Some(0),
Some(99),
None,
Some(5),
Some(1),
);
assert_eq!(effective_provider_key_rpm_limit(&low_confidence, 100), None);
let high_confidence = provider_catalog_key("key-a").with_rate_limit_fields(
None,
Some(80),
Some(0),
Some(0),
None,
Some(serde_json::json!([
{"new_limit": 80},
{"new_limit": 81},
{"new_limit": 80},
])),
Some(120),
Some(118),
);
assert_eq!(
effective_provider_key_rpm_limit(&high_confidence, 100),
Some(80)
);
}
#[test]
fn counts_recent_provider_key_rpm_from_snapshot_or_recent_attempts() {
let recent_candidates = vec![
StoredRequestCandidate::new(
"one".to_string(),
"req-one".to_string(),
None,
None,
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(7),
None,
None,
95,
Some(95),
Some(96),
)
.expect("candidate should build"),
StoredRequestCandidate::new(
"two".to_string(),
"req-two".to_string(),
None,
None,
None,
None,
0,
0,
Some("provider-a".to_string()),
Some("endpoint-a".to_string()),
Some("key-a".to_string()),
RequestCandidateStatus::Failed,
None,
false,
Some(502),
None,
None,
Some(10),
None,
None,
None,
98,
Some(98),
Some(99),
)
.expect("candidate should build"),
];
assert_eq!(
count_recent_rpm_requests_for_provider_key(&recent_candidates, "key-a", 100),
7
);
}
#[test]
fn ignores_rpm_observations_before_reset_watermark() {
let recent_candidates = vec![
StoredRequestCandidate::new(
"one".to_string(),
"req-one".to_string(),
None,
None,
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(7),
None,
None,
95,
Some(95),
Some(96),
)
.expect("candidate should build"),
StoredRequestCandidate::new(
"two".to_string(),
"req-two".to_string(),
None,
None,
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(2),
None,
None,
99,
Some(99),
Some(100),
)
.expect("candidate should build"),
];
assert_eq!(
count_recent_rpm_requests_for_provider_key_since(
&recent_candidates,
"key-a",
100,
Some(98),
),
2
);
}
#[test]
fn provider_key_rpm_reserves_capacity_for_new_users() {
let key = provider_catalog_key("key-a").with_rate_limit_fields(
Some(10),
None,
None,
None,
None,
None,
Some(5),
Some(5),
);
let recent_candidates = vec![StoredRequestCandidate::new(
"one".to_string(),
"req-one".to_string(),
None,
None,
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")];
assert!(!provider_key_rpm_allows_request(
&key,
&recent_candidates,
100,
false,
));
assert!(provider_key_rpm_allows_request(
&key,
&recent_candidates,
100,
true,
));
assert!(provider_key_rpm_allows_request_since(
&key,
&recent_candidates,
100,
false,
Some(97),
));
}
#[test]
fn reads_provider_key_health_and_circuit_status_for_api_format() {
let key = provider_catalog_key("key-a").with_health_fields(
Some(serde_json::json!({
"openai:chat": {"health_score": 0.25},
"openai:responses": {"health_score": 0.75}
})),
Some(serde_json::json!({
"openai:chat": {"open": true},
"openai:responses": {"open": false}
})),
);
assert_eq!(provider_key_health_score(&key, "openai:chat"), Some(0.25));
assert_eq!(
provider_key_health_score(&key, "openai:responses"),
Some(0.75)
);
assert!(is_provider_key_circuit_open(&key, "openai:chat"));
assert!(!is_provider_key_circuit_open(&key, "openai:responses"));
}
#[test]
fn aggregates_provider_key_health_score_with_lower_bound_strategy() {
let key = provider_catalog_key("key-a").with_health_fields(
Some(serde_json::json!({
"openai:chat": {"health_score": 0.85},
"openai:responses": {"health_score": 0.45},
"claude:chat": {"health_score": 0.70}
})),
None,
);
assert_eq!(aggregate_provider_key_health_score(&key), Some(0.45));
assert_eq!(
effective_provider_key_health_score(&key, "gemini:chat"),
Some(0.45)
);
}
#[test]
fn classifies_provider_key_health_bucket_from_effective_score() {
let low = provider_catalog_key("key-low").with_health_fields(
Some(serde_json::json!({"openai:chat": {"health_score": 0.30}})),
None,
);
let degraded = provider_catalog_key("key-degraded").with_health_fields(
Some(serde_json::json!({"openai:chat": {"health_score": 0.65}})),
None,
);
let healthy = provider_catalog_key("key-healthy").with_health_fields(
Some(serde_json::json!({"openai:chat": {"health_score": 0.92}})),
None,
);
assert_eq!(
provider_key_health_bucket(&low, "openai:chat"),
Some(ProviderKeyHealthBucket::Low)
);
assert_eq!(
provider_key_health_bucket(&degraded, "openai:chat"),
Some(ProviderKeyHealthBucket::Degraded)
);
assert_eq!(
provider_key_health_bucket(&healthy, "openai:chat"),
Some(ProviderKeyHealthBucket::Healthy)
);
}
}
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,
};

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

@@ -1,6 +1,7 @@
mod candidate;
mod failover;
mod health;
mod request_candidate_state;
mod request_candidates;
mod route;
@@ -10,6 +11,9 @@ pub(crate) use candidate::{
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,
@@ -22,6 +26,7 @@ pub(crate) use health::{
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,

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

@@ -0,0 +1,19 @@
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>;
}

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

@@ -1,17 +1,24 @@
use std::time::{SystemTime, UNIX_EPOCH};
use aether_contracts::{ExecutionError, ExecutionPlan};
use aether_data::repository::candidates::{
RequestCandidateStatus, StoredRequestCandidate, UpsertRequestCandidateRecord,
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 crate::gateway::AppState;
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: &AppState,
state: &(impl SchedulerRequestCandidateRuntimeState + ?Sized),
plan: &ExecutionPlan,
report_context: Option<&Value>,
status: RequestCandidateStatus,
@@ -22,51 +29,25 @@ pub(crate) async fn record_local_request_candidate_status(
started_at_unix_secs: Option<u64>,
finished_at_unix_secs: Option<u64>,
) {
let Some(candidate_id) = plan
.candidate_id
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
else {
return;
};
let Some(metadata) = parse_report_context(report_context) else {
return;
};
let Some(candidate_index) = metadata.candidate_index else {
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(UpsertRequestCandidateRecord {
id: candidate_id.to_string(),
request_id: plan.request_id.clone(),
user_id: metadata.user_id,
api_key_id: metadata.api_key_id,
username: None,
api_key_name: None,
candidate_index,
retry_index: metadata.retry_index,
provider_id: Some(plan.provider_id.clone()),
endpoint_id: Some(plan.endpoint_id.clone()),
key_id: Some(plan.key_id.clone()),
status,
skip_reason: None,
is_cached: None,
status_code,
error_type,
error_message,
latency_ms,
concurrent_requests: None,
extra_data: None,
required_capabilities: None,
created_at_unix_secs: None,
started_at_unix_secs,
finished_at_unix_secs,
})
.await
{
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,
@@ -76,7 +57,7 @@ pub(crate) async fn record_local_request_candidate_status(
}
pub(crate) async fn record_report_request_candidate_status(
state: &AppState,
state: &(impl SchedulerRequestCandidateRuntimeState + ?Sized),
report_context: Option<&Value>,
status: RequestCandidateStatus,
status_code: Option<u16>,
@@ -89,48 +70,28 @@ pub(crate) async fn record_report_request_candidate_status(
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(),
);
let terminal_unix_secs = finished_at_unix_secs.unwrap_or_else(current_unix_secs);
let started_at_unix_secs = started_at_unix_secs
.or(slot.started_at_unix_secs)
.or_else(|| status.is_attempted(None).then_some(terminal_unix_secs));
let finished_at_unix_secs = finished_at_unix_secs
.or(slot.finished_at_unix_secs)
.or_else(|| is_terminal_candidate_status(status).then_some(terminal_unix_secs));
if let Err(err) = state
.upsert_request_candidate(UpsertRequestCandidateRecord {
id: slot.id,
request_id: slot.request_id.clone(),
user_id: slot.user_id,
api_key_id: slot.api_key_id,
username: None,
api_key_name: None,
candidate_index: slot.candidate_index,
retry_index: slot.retry_index,
provider_id: slot.provider_id,
endpoint_id: slot.endpoint_id,
key_id: slot.key_id,
status,
skip_reason: None,
is_cached: None,
status_code,
error_type,
error_message,
latency_ms,
concurrent_requests: None,
extra_data: slot.extra_data,
required_capabilities: None,
created_at_unix_secs: Some(slot.created_at_unix_secs),
started_at_unix_secs,
finished_at_unix_secs,
})
.await
{
if let Err(err) = state.upsert_request_candidate(record).await {
warn!(
request_id = %slot.request_id,
candidate_index = slot.candidate_index,
retry_index = slot.retry_index,
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"
);
@@ -138,7 +99,7 @@ pub(crate) async fn record_report_request_candidate_status(
}
pub(crate) async fn ensure_execution_request_candidate_slot(
state: &AppState,
state: &(impl SchedulerRequestCandidateRuntimeState + ?Sized),
plan: &mut ExecutionPlan,
report_context: &mut Option<Value>,
) {
@@ -154,80 +115,21 @@ pub(crate) async fn ensure_execution_request_candidate_slot(
return;
}
let mut context = report_context
.as_ref()
.and_then(Value::as_object)
.cloned()
.unwrap_or_default();
let request_id = context
.get("request_id")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(plan.request_id.as_str())
.to_string();
let candidate_index = context
.get("candidate_index")
.and_then(Value::as_u64)
.and_then(|value| u32::try_from(value).ok())
.unwrap_or(0);
let retry_index = context
.get("retry_index")
.and_then(Value::as_u64)
.and_then(|value| u32::try_from(value).ok())
.unwrap_or(0);
let generated_candidate_id = context
.get("candidate_id")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.unwrap_or_else(|| Uuid::new_v4().to_string());
let started_at_unix_secs = current_unix_secs();
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(UpsertRequestCandidateRecord {
id: generated_candidate_id.clone(),
request_id: request_id.clone(),
user_id: context
.get("user_id")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
api_key_id: context
.get("api_key_id")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
username: None,
api_key_name: None,
candidate_index,
retry_index,
provider_id: Some(plan.provider_id.clone()),
endpoint_id: Some(plan.endpoint_id.clone()),
key_id: Some(plan.key_id.clone()),
status: RequestCandidateStatus::Pending,
skip_reason: None,
is_cached: Some(false),
status_code: None,
error_type: None,
error_message: None,
latency_ms: None,
concurrent_requests: None,
extra_data: None,
required_capabilities: None,
created_at_unix_secs: Some(started_at_unix_secs),
started_at_unix_secs: Some(started_at_unix_secs),
finished_at_unix_secs: None,
})
.await
{
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"
@@ -237,54 +139,10 @@ pub(crate) async fn ensure_execution_request_candidate_slot(
};
plan.candidate_id = Some(candidate_id.clone());
context.insert("request_id".to_string(), Value::String(request_id));
context.insert("candidate_id".to_string(), Value::String(candidate_id));
context.insert(
"candidate_index".to_string(),
Value::Number(candidate_index.into()),
);
context.insert(
"provider_id".to_string(),
Value::String(plan.provider_id.clone()),
);
context.insert(
"endpoint_id".to_string(),
Value::String(plan.endpoint_id.clone()),
);
context.insert("key_id".to_string(), Value::String(plan.key_id.clone()));
*report_context = Some(Value::Object(context));
}
pub(crate) fn execution_error_details(
error: Option<&ExecutionError>,
body_json: Option<&Value>,
) -> (Option<String>, Option<String>) {
match error {
Some(error) => (
Some(format!("{:?}", error.kind)),
Some(error.message.trim().to_string()).filter(|value| !value.is_empty()),
),
None => (
None,
body_json
.and_then(extract_error_message)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
),
}
}
fn extract_error_message(body_json: &Value) -> Option<&str> {
body_json
.get("error")
.and_then(|error| {
error
.get("message")
.and_then(Value::as_str)
.or_else(|| error.as_str())
})
.or_else(|| body_json.get("message").and_then(Value::as_str))
*report_context = Some(finalize_execution_request_candidate_report_context(
seed.report_context,
&candidate_id,
));
}
pub(crate) fn current_unix_secs() -> u64 {
@@ -294,251 +152,22 @@ pub(crate) fn current_unix_secs() -> u64 {
.as_secs()
}
#[derive(Debug, Clone)]
struct RequestCandidateReportContext {
request_id: Option<String>,
candidate_id: Option<String>,
user_id: Option<String>,
api_key_id: Option<String>,
candidate_index: Option<u32>,
retry_index: u32,
provider_id: Option<String>,
endpoint_id: Option<String>,
key_id: Option<String>,
client_api_format: Option<String>,
provider_api_format: Option<String>,
}
fn parse_report_context(report_context: Option<&Value>) -> Option<RequestCandidateReportContext> {
let report_context = report_context?;
let retry_index = report_context
.get("retry_index")
.and_then(|value| value.as_u64())
.unwrap_or_default();
Some(RequestCandidateReportContext {
request_id: report_context
.get("request_id")
.and_then(|value| value.as_str())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
candidate_id: report_context
.get("candidate_id")
.and_then(|value| value.as_str())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
user_id: report_context
.get("user_id")
.and_then(|value| value.as_str())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
api_key_id: report_context
.get("api_key_id")
.and_then(|value| value.as_str())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
candidate_index: report_context
.get("candidate_index")
.and_then(|value| value.as_u64())
.map(|value| value as u32),
retry_index: retry_index as u32,
provider_id: report_context
.get("provider_id")
.and_then(|value| value.as_str())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
endpoint_id: report_context
.get("endpoint_id")
.and_then(|value| value.as_str())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
key_id: report_context
.get("key_id")
.and_then(|value| value.as_str())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
client_api_format: report_context
.get("client_api_format")
.and_then(|value| value.as_str())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
provider_api_format: report_context
.get("provider_api_format")
.and_then(|value| value.as_str())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
})
}
#[derive(Debug, Clone)]
struct ResolvedReportRequestCandidateSlot {
id: String,
request_id: String,
user_id: Option<String>,
api_key_id: Option<String>,
candidate_index: u32,
retry_index: u32,
provider_id: Option<String>,
endpoint_id: Option<String>,
key_id: Option<String>,
extra_data: Option<Value>,
created_at_unix_secs: u64,
started_at_unix_secs: Option<u64>,
finished_at_unix_secs: Option<u64>,
}
async fn resolve_report_request_candidate_slot(
state: &AppState,
state: &(impl SchedulerRequestCandidateRuntimeState + ?Sized),
report_context: Option<&Value>,
) -> Option<ResolvedReportRequestCandidateSlot> {
let metadata = parse_report_context(report_context)?;
) -> 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();
let matched_candidate = match_existing_report_candidate(&existing_candidates, &metadata);
let synthesized_extra_data = build_report_candidate_extra_data(&metadata);
let created_at_unix_secs = matched_candidate
.as_ref()
.map(|candidate| candidate.created_at_unix_secs)
.unwrap_or_else(current_unix_secs);
let candidate_index = matched_candidate
.as_ref()
.map(|candidate| candidate.candidate_index)
.or(metadata.candidate_index)
.unwrap_or_else(|| next_candidate_index(&existing_candidates));
let retry_index = matched_candidate
.as_ref()
.map(|candidate| candidate.retry_index)
.unwrap_or(metadata.retry_index);
Some(ResolvedReportRequestCandidateSlot {
id: matched_candidate
.as_ref()
.map(|candidate| candidate.id.clone())
.or(metadata.candidate_id)
.unwrap_or_else(|| Uuid::new_v4().to_string()),
request_id,
user_id: matched_candidate
.as_ref()
.and_then(|candidate| candidate.user_id.clone())
.or(metadata.user_id),
api_key_id: matched_candidate
.as_ref()
.and_then(|candidate| candidate.api_key_id.clone())
.or(metadata.api_key_id),
candidate_index,
retry_index,
provider_id: matched_candidate
.as_ref()
.and_then(|candidate| candidate.provider_id.clone())
.or(metadata.provider_id),
endpoint_id: matched_candidate
.as_ref()
.and_then(|candidate| candidate.endpoint_id.clone())
.or(metadata.endpoint_id),
key_id: matched_candidate
.as_ref()
.and_then(|candidate| candidate.key_id.clone())
.or(metadata.key_id),
extra_data: matched_candidate
.as_ref()
.and_then(|candidate| candidate.extra_data.clone())
.or(synthesized_extra_data),
created_at_unix_secs,
started_at_unix_secs: matched_candidate
.as_ref()
.and_then(|candidate| candidate.started_at_unix_secs),
finished_at_unix_secs: matched_candidate
.as_ref()
.and_then(|candidate| candidate.finished_at_unix_secs),
})
}
fn match_existing_report_candidate<'a>(
candidates: &'a [StoredRequestCandidate],
metadata: &RequestCandidateReportContext,
) -> Option<&'a StoredRequestCandidate> {
if let Some(candidate_id) = metadata.candidate_id.as_deref() {
if let Some(candidate) = candidates
.iter()
.find(|candidate| candidate.id == candidate_id)
{
return Some(candidate);
}
}
if let Some(candidate_index) = metadata.candidate_index {
if let Some(candidate) = candidates.iter().find(|candidate| {
candidate.candidate_index == candidate_index
&& candidate.retry_index == metadata.retry_index
}) {
return Some(candidate);
}
}
candidates
.iter()
.filter(|candidate| {
candidate.provider_id.as_deref() == metadata.provider_id.as_deref()
&& candidate.endpoint_id.as_deref() == metadata.endpoint_id.as_deref()
&& candidate.key_id.as_deref() == metadata.key_id.as_deref()
})
.max_by_key(|candidate| {
(
candidate.retry_index,
candidate.candidate_index,
candidate.created_at_unix_secs,
)
})
}
fn next_candidate_index(candidates: &[StoredRequestCandidate]) -> u32 {
candidates
.iter()
.map(|candidate| candidate.candidate_index)
.max()
.map(|value| value.saturating_add(1))
.unwrap_or_default()
}
fn build_report_candidate_extra_data(metadata: &RequestCandidateReportContext) -> Option<Value> {
let mut extra_data = serde_json::Map::new();
extra_data.insert("gateway_execution_runtime".to_string(), Value::Bool(true));
extra_data.insert("phase".to_string(), Value::String("3c_trial".to_string()));
if let Some(client_api_format) = metadata.client_api_format.clone() {
extra_data.insert(
"client_api_format".to_string(),
Value::String(client_api_format),
);
}
if let Some(provider_api_format) = metadata.provider_api_format.clone() {
extra_data.insert(
"provider_api_format".to_string(),
Value::String(provider_api_format),
);
}
(!extra_data.is_empty()).then_some(Value::Object(extra_data))
}
fn is_terminal_candidate_status(status: RequestCandidateStatus) -> bool {
matches!(
status,
RequestCandidateStatus::Unused
| RequestCandidateStatus::Success
| RequestCandidateStatus::Failed
| RequestCandidateStatus::Cancelled
| RequestCandidateStatus::Skipped
resolve_report_request_candidate_slot_from_candidates(
&existing_candidates,
metadata,
current_unix_secs(),
Uuid::new_v4().to_string(),
)
}
@@ -550,13 +179,14 @@ mod tests {
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;
use crate::gateway::gateway_data::GatewayDataState;
use crate::gateway::AppState;
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()
@@ -676,4 +306,72 @@ mod tests {
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));
}
}

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

@@ -1,4 +1,4 @@
use crate::gateway::GatewayControlDecision;
use crate::control::GatewayControlDecision;
pub(crate) fn resolve_execution_runtime_stream_plan_kind(
parts: &http::request::Parts,
@@ -291,7 +291,7 @@ mod tests {
resolve_execution_runtime_sync_plan_kind, supports_stream_scheduler_decision_kind,
supports_sync_scheduler_decision_kind,
};
use crate::gateway::GatewayControlDecision;
use crate::control::GatewayControlDecision;
fn sample_decision(route_family: &str, route_kind: &str) -> GatewayControlDecision {
GatewayControlDecision {