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

Implement independent provider pool scheduling runtime

Browse Source
This commit is contained in:
fawney19
2026-04-17 23:05:49 +08:00
parent dd4641d618
commit 3321bb3ccc
38 changed files with 3222 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files

293
apps/aether-gateway/src/handlers/admin/provider/pool/config.rs
View File

@@ -1,4 +1,6 @@
use crate::handlers::admin::provider::shared::support::AdminProviderPoolConfig;
use crate::handlers::admin::provider::shared::support::{
AdminProviderPoolConfig, AdminProviderPoolSchedulingPreset, AdminProviderPoolUnschedulableRule,
};
use serde_json::{Map, Value};
fn json_u64(value: &Value) -> Option<u64> {
@@ -7,7 +9,115 @@ fn json_u64(value: &Value) -> Option<u64> {
.or_else(|| value.as_i64().and_then(|raw| u64::try_from(raw).ok()))
}
fn admin_provider_pool_lru_enabled(raw_pool_advanced: &Map<String, Value>) -> bool {
fn parse_pool_scheduling_presets(
raw_pool_advanced: &Map<String, Value>,
) -> Vec<AdminProviderPoolSchedulingPreset> {
let Some(presets) = raw_pool_advanced
.get("scheduling_presets")
.and_then(Value::as_array)
else {
return raw_pool_advanced
.get("lru_enabled")
.and_then(Value::as_bool)
.filter(|enabled| *enabled)
.map(|_| {
vec![AdminProviderPoolSchedulingPreset {
preset: "lru".to_string(),
enabled: true,
mode: None,
}]
})
.unwrap_or_default();
};
let mut normalized = Vec::new();
for item in presets {
if let Some(preset) = item
.as_str()
.map(str::trim)
.filter(|value| !value.is_empty())
{
normalized.push(AdminProviderPoolSchedulingPreset {
preset: preset.to_ascii_lowercase(),
enabled: true,
mode: None,
});
continue;
}
let Some(object) = item.as_object() else {
continue;
};
let Some(preset) = object
.get("preset")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
else {
continue;
};
normalized.push(AdminProviderPoolSchedulingPreset {
preset: preset.to_ascii_lowercase(),
enabled: object
.get("enabled")
.and_then(Value::as_bool)
.unwrap_or(true),
mode: object
.get("mode")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| value.to_ascii_lowercase()),
});
}
if normalized.is_empty()
&& raw_pool_advanced
.get("lru_enabled")
.and_then(Value::as_bool)
.unwrap_or(false)
{
normalized.push(AdminProviderPoolSchedulingPreset {
preset: "lru".to_string(),
enabled: true,
mode: None,
});
}
normalized
}
fn parse_pool_unschedulable_rules(
raw_pool_advanced: &Map<String, Value>,
) -> Vec<AdminProviderPoolUnschedulableRule> {
raw_pool_advanced
.get("unschedulable_rules")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter_map(|item| {
let object = item.as_object()?;
let keyword = object
.get("keyword")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())?;
Some(AdminProviderPoolUnschedulableRule {
keyword: keyword.to_string(),
duration_minutes: object
.get("duration_minutes")
.and_then(json_u64)
.filter(|value| *value > 0)
.unwrap_or(5),
})
})
.collect()
}
fn admin_provider_pool_lru_enabled(
raw_pool_advanced: &Map<String, Value>,
scheduling_presets: &[AdminProviderPoolSchedulingPreset],
) -> bool {
if let Some(explicit) = raw_pool_advanced
.get("lru_enabled")
.and_then(Value::as_bool)
@@ -15,56 +125,70 @@ fn admin_provider_pool_lru_enabled(raw_pool_advanced: &Map<String, Value>) -> bo
return explicit;
}
let Some(presets) = raw_pool_advanced
.get("scheduling_presets")
.and_then(Value::as_array)
else {
return false;
};
let Some(first) = presets.first() else {
return false;
};
if first.is_string() {
return raw_pool_advanced
.get("lru_enabled")
.and_then(Value::as_bool)
.unwrap_or(true);
}
presets.iter().filter_map(Value::as_object).any(|item| {
item.get("preset")
.and_then(Value::as_str)
.is_some_and(|preset| preset.eq_ignore_ascii_case("lru"))
&& item.get("enabled").and_then(Value::as_bool).unwrap_or(true)
})
scheduling_presets
.iter()
.any(|item| item.enabled && item.preset.eq_ignore_ascii_case("lru"))
}
pub(crate) fn admin_provider_pool_config(
provider: &aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider,
) -> Option<AdminProviderPoolConfig> {
let raw_pool_advanced = provider
.config
.as_ref()
admin_provider_pool_config_from_config_value(provider.config.as_ref())
}
pub(crate) fn admin_provider_pool_config_from_config_value(
config: Option<&serde_json::Value>,
) -> Option<AdminProviderPoolConfig> {
let raw_pool_advanced = config
.and_then(Value::as_object)
.and_then(|config| config.get("pool_advanced"))?;
let Some(pool_advanced) = raw_pool_advanced.as_object() else {
return Some(AdminProviderPoolConfig {
scheduling_presets: Vec::new(),
unschedulable_rules: Vec::new(),
lru_enabled: false,
skip_exhausted_accounts: false,
sticky_session_ttl_seconds: 3600,
latency_window_seconds: 3600,
latency_sample_limit: 50,
cost_window_seconds: 18_000,
cost_limit_per_key_tokens: None,
rate_limit_cooldown_seconds: 300,
overload_cooldown_seconds: 30,
health_policy_enabled: true,
stream_timeout_threshold: 3,
stream_timeout_window_seconds: 1800,
stream_timeout_cooldown_seconds: 300,
});
};
let scheduling_presets = parse_pool_scheduling_presets(pool_advanced);
let unschedulable_rules = parse_pool_unschedulable_rules(pool_advanced);
Some(AdminProviderPoolConfig {
lru_enabled: admin_provider_pool_lru_enabled(pool_advanced),
lru_enabled: admin_provider_pool_lru_enabled(pool_advanced, &scheduling_presets),
scheduling_presets,
unschedulable_rules,
skip_exhausted_accounts: pool_advanced
.get("skip_exhausted_accounts")
.and_then(Value::as_bool)
.unwrap_or(false),
sticky_session_ttl_seconds: pool_advanced
.get("sticky_session_ttl_seconds")
.and_then(json_u64)
.filter(|value| *value > 0)
.unwrap_or(3600),
latency_window_seconds: pool_advanced
.get("latency_window_seconds")
.and_then(json_u64)
.filter(|value| *value > 0)
.unwrap_or(3600),
latency_sample_limit: pool_advanced
.get("latency_sample_limit")
.and_then(json_u64)
.filter(|value| *value > 0)
.unwrap_or(50),
cost_window_seconds: pool_advanced
.get("cost_window_seconds")
.and_then(json_u64)
@@ -73,12 +197,41 @@ pub(crate) fn admin_provider_pool_config(
cost_limit_per_key_tokens: pool_advanced
.get("cost_limit_per_key_tokens")
.and_then(json_u64),
rate_limit_cooldown_seconds: pool_advanced
.get("rate_limit_cooldown_seconds")
.and_then(json_u64)
.filter(|value| *value > 0)
.unwrap_or(300),
overload_cooldown_seconds: pool_advanced
.get("overload_cooldown_seconds")
.and_then(json_u64)
.filter(|value| *value > 0)
.unwrap_or(30),
health_policy_enabled: pool_advanced
.get("health_policy_enabled")
.and_then(Value::as_bool)
.unwrap_or(true),
stream_timeout_threshold: pool_advanced
.get("stream_timeout_threshold")
.and_then(json_u64)
.filter(|value| *value > 0)
.unwrap_or(3),
stream_timeout_window_seconds: pool_advanced
.get("stream_timeout_window_seconds")
.and_then(json_u64)
.filter(|value| *value > 0)
.unwrap_or(1800),
stream_timeout_cooldown_seconds: pool_advanced
.get("stream_timeout_cooldown_seconds")
.and_then(json_u64)
.filter(|value| *value > 0)
.unwrap_or(300),
})
}
#[cfg(test)]
mod tests {
use super::admin_provider_pool_config;
use super::{admin_provider_pool_config, admin_provider_pool_config_from_config_value};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider;
use serde_json::json;
@@ -117,15 +270,91 @@ mod tests {
"pool_advanced": {
"skip_exhausted_accounts": true,
"lru_enabled": true,
"sticky_session_ttl_seconds": 600,
"latency_window_seconds": 900,
"latency_sample_limit": 75,
"cost_window_seconds": 7200,
"cost_limit_per_key_tokens": 12000
"cost_limit_per_key_tokens": 12000,
"rate_limit_cooldown_seconds": 420,
"overload_cooldown_seconds": 45,
"health_policy_enabled": false,
"stream_timeout_threshold": 4,
"stream_timeout_window_seconds": 900,
"stream_timeout_cooldown_seconds": 180
}
}));
let config = admin_provider_pool_config(&provider).expect("pool config should exist");
assert!(config.skip_exhausted_accounts);
assert!(config.lru_enabled);
assert_eq!(config.sticky_session_ttl_seconds, 600);
assert_eq!(config.latency_window_seconds, 900);
assert_eq!(config.latency_sample_limit, 75);
assert_eq!(config.cost_window_seconds, 7200);
assert_eq!(config.cost_limit_per_key_tokens, Some(12_000));
assert_eq!(config.rate_limit_cooldown_seconds, 420);
assert_eq!(config.overload_cooldown_seconds, 45);
assert!(!config.health_policy_enabled);
assert_eq!(config.stream_timeout_threshold, 4);
assert_eq!(config.stream_timeout_window_seconds, 900);
assert_eq!(config.stream_timeout_cooldown_seconds, 180);
}
#[test]
fn parses_pool_config_from_generic_config_value() {
let config = admin_provider_pool_config_from_config_value(Some(&json!({
"pool_advanced": {
"scheduling_presets": [{"preset": "lru", "enabled": true}],
"cost_limit_per_key_tokens": 4096
}
})))
.expect("pool config should parse");
assert!(config.lru_enabled);
assert_eq!(config.scheduling_presets.len(), 1);
assert_eq!(config.scheduling_presets[0].preset, "lru");
assert_eq!(config.cost_limit_per_key_tokens, Some(4096));
}
#[test]
fn parses_object_style_scheduling_presets_with_modes() {
let config = admin_provider_pool_config_from_config_value(Some(&json!({
"pool_advanced": {
"scheduling_presets": [
{"preset": "cache_affinity", "enabled": false},
{"preset": "plus_first", "enabled": true, "mode": "plus_only"}
]
}
})))
.expect("pool config should parse");
assert!(!config.lru_enabled);
assert_eq!(config.scheduling_presets.len(), 2);
assert_eq!(config.scheduling_presets[0].preset, "cache_affinity");
assert!(!config.scheduling_presets[0].enabled);
assert_eq!(config.scheduling_presets[1].preset, "plus_first");
assert_eq!(
config.scheduling_presets[1].mode.as_deref(),
Some("plus_only")
);
}
#[test]
fn parses_unschedulable_rules_from_pool_advanced() {
let config = admin_provider_pool_config_from_config_value(Some(&json!({
"pool_advanced": {
"unschedulable_rules": [
{"keyword": "suspended", "duration_minutes": 15},
{"keyword": "review_required"}
]
}
})))
.expect("pool config should parse");
assert_eq!(config.unschedulable_rules.len(), 2);
assert_eq!(config.unschedulable_rules[0].keyword, "suspended");
assert_eq!(config.unschedulable_rules[0].duration_minutes, 15);
assert_eq!(config.unschedulable_rules[1].keyword, "review_required");
assert_eq!(config.unschedulable_rules[1].duration_minutes, 5);
}
}

42
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/keys.rs
View File

@@ -4,6 +4,14 @@ pub(super) fn pool_sticky_pattern(keyspace: &RedisKeyspace, provider_id: &str) -
keyspace.key(&format!("ap:{provider_id}:sticky:*"))
}
pub(super) fn pool_sticky_key(
keyspace: &RedisKeyspace,
provider_id: &str,
session_token: &str,
) -> String {
keyspace.key(&format!("ap:{provider_id}:sticky:{session_token}"))
}
pub(super) fn pool_lru_key(keyspace: &RedisKeyspace, provider_id: &str) -> String {
keyspace.key(&format!("ap:{provider_id}:lru"))
}
@@ -24,6 +32,22 @@ pub(super) fn pool_cost_key(keyspace: &RedisKeyspace, provider_id: &str, key_id:
keyspace.key(&format!("ap:{provider_id}:cost:{key_id}"))
}
pub(super) fn pool_latency_key(
keyspace: &RedisKeyspace,
provider_id: &str,
key_id: &str,
) -> String {
keyspace.key(&format!("ap:{provider_id}:latency:{key_id}"))
}
pub(super) fn pool_stream_timeout_key(
keyspace: &RedisKeyspace,
provider_id: &str,
key_id: &str,
) -> String {
keyspace.key(&format!("ap:{provider_id}:stream_timeout:{key_id}"))
}
pub(super) fn parse_pool_cost_member(member: &str) -> u64 {
member
.rsplit_once(':')
@@ -31,6 +55,13 @@ pub(super) fn parse_pool_cost_member(member: &str) -> u64 {
.unwrap_or(0)
}
pub(super) fn parse_pool_latency_member(member: &str) -> u64 {
member
.rsplit_once(':')
.and_then(|(_, suffix)| suffix.parse::<u64>().ok())
.unwrap_or(0)
}
pub(super) fn pool_cooldown_keys(
keyspace: &RedisKeyspace,
provider_id: &str,
@@ -52,3 +83,14 @@ pub(super) fn pool_cost_keys(
.map(|key_id| pool_cost_key(keyspace, provider_id, key_id))
.collect()
}
pub(super) fn pool_latency_keys(
keyspace: &RedisKeyspace,
provider_id: &str,
key_ids: &[String],
) -> Vec<String> {
key_ids
.iter()
.map(|key_id| pool_latency_key(keyspace, provider_id, key_id))
.collect()
}

5
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/mod.rs
View File

@@ -2,6 +2,7 @@ mod keys;
mod mutations;
mod reads;
mod status;
mod writes;
pub(crate) use self::mutations::{
clear_admin_provider_pool_cooldown, reset_admin_provider_pool_cost,
@@ -11,3 +12,7 @@ pub(crate) use self::reads::{
read_admin_provider_pool_cooldown_key_ids, read_admin_provider_pool_runtime_state,
};
pub(crate) use self::status::build_admin_provider_pool_status_payload;
pub(crate) use self::writes::{
record_admin_provider_pool_error, record_admin_provider_pool_stream_timeout,
record_admin_provider_pool_success,
};

119
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/reads.rs
View File

@@ -1,6 +1,7 @@
use super::keys::{
parse_pool_cost_member, pool_cooldown_index_key, pool_cooldown_keys, pool_cost_keys,
pool_lru_key, pool_sticky_pattern,
parse_pool_cost_member, parse_pool_latency_member, pool_cooldown_index_key, pool_cooldown_key,
pool_cooldown_keys, pool_cost_keys, pool_latency_keys, pool_lru_key, pool_sticky_key,
pool_sticky_pattern,
};
use crate::handlers::admin::provider::shared::support::{
AdminProviderPoolConfig, AdminProviderPoolRuntimeState, ADMIN_PROVIDER_POOL_SCAN_BATCH,
@@ -11,6 +12,13 @@ use std::collections::BTreeMap;
use std::time::{SystemTime, UNIX_EPOCH};
use tracing::warn;
fn current_unix_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
}
async fn scan_redis_keys(
connection: &mut redis::aio::MultiplexedConnection,
pattern: &str,
@@ -77,7 +85,8 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
runner: &RedisKvRunner,
provider_id: &str,
key_ids: &[String],
pool_config: AdminProviderPoolConfig,
pool_config: &AdminProviderPoolConfig,
sticky_session_token: Option<&str>,
) -> AdminProviderPoolRuntimeState {
let mut runtime = AdminProviderPoolRuntimeState::default();
let Ok(mut connection) = runner.client().get_multiplexed_async_connection().await else {
@@ -87,6 +96,57 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
let keyspace = runner.keyspace().clone();
let cooldown_keys = pool_cooldown_keys(&keyspace, provider_id, key_ids);
let cost_keys = pool_cost_keys(&keyspace, provider_id, key_ids);
let latency_keys = pool_latency_keys(&keyspace, provider_id, key_ids);
if let Some(sticky_session_token) = sticky_session_token
.map(str::trim)
.filter(|value| !value.is_empty())
.filter(|_| pool_config.sticky_session_ttl_seconds > 0)
{
let sticky_key = pool_sticky_key(&keyspace, provider_id, sticky_session_token);
let sticky_bound_key_id = redis::cmd("GET")
.arg(&sticky_key)
.query_async::<Option<String>>(&mut connection)
.await
.unwrap_or_else(|err| {
warn!(
"gateway admin provider pool: failed to read sticky binding for provider {provider_id}: {:?}",
err
);
None
});
if let Some(bound_key_id) = sticky_bound_key_id {
let cooldown_key = pool_cooldown_key(&keyspace, provider_id, &bound_key_id);
runtime.sticky_bound_key_id = match redis::cmd("EXISTS")
.arg(&cooldown_key)
.query_async::<u64>(&mut connection)
.await
{
Ok(0) => {
let _: Result<bool, _> = redis::cmd("EXPIRE")
.arg(&sticky_key)
.arg(pool_config.sticky_session_ttl_seconds)
.query_async(&mut connection)
.await;
Some(bound_key_id)
}
Ok(_) => {
let _: Result<i64, _> = redis::cmd("DEL")
.arg(&sticky_key)
.query_async(&mut connection)
.await;
None
}
Err(err) => {
warn!(
"gateway admin provider pool: failed to validate sticky cooldown for provider {provider_id}: {:?}",
err
);
Some(bound_key_id)
}
};
}
}
let sticky_keys = match scan_redis_keys(
&mut connection,
@@ -174,11 +234,7 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
}
if !cost_keys.is_empty() {
let window_start = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
.saturating_sub(pool_config.cost_window_seconds);
let window_start = current_unix_secs().saturating_sub(pool_config.cost_window_seconds);
let mut cost_pipeline = redis::pipe();
for cost_key in &cost_keys {
cost_pipeline
@@ -208,7 +264,52 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
}
}
if pool_config.lru_enabled && !key_ids.is_empty() {
if !latency_keys.is_empty() {
let window_start = current_unix_secs().saturating_sub(pool_config.latency_window_seconds);
let mut latency_pipeline = redis::pipe();
for latency_key in &latency_keys {
latency_pipeline
.cmd("ZRANGEBYSCORE")
.arg(latency_key)
.arg(window_start)
.arg("+inf");
}
let members_by_key = latency_pipeline
.query_async::<Vec<Vec<String>>>(&mut connection)
.await
.unwrap_or_else(|err| {
warn!(
"gateway admin provider pool: failed to batch read latency windows for provider {provider_id}: {:?}",
err
);
vec![Vec::new(); latency_keys.len()]
});
for (key_id, members) in key_ids.iter().zip(members_by_key) {
let samples = members
.iter()
.map(|member| parse_pool_latency_member(member))
.filter(|value| *value > 0)
.collect::<Vec<_>>();
if samples.is_empty() {
continue;
}
let total = samples.iter().sum::<u64>() as f64;
let average = total / samples.len() as f64;
if average.is_finite() && average >= 0.0 {
runtime
.latency_avg_ms_by_key
.insert(key_id.clone(), average);
}
}
}
if (pool_config.lru_enabled
|| pool_config
.scheduling_presets
.iter()
.any(|item| item.enabled))
&& !key_ids.is_empty()
{
let mut command = redis::cmd("ZMSCORE");
command.arg(pool_lru_key(&keyspace, provider_id));
for key_id in key_ids {

11
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/status.rs
View File

@@ -36,8 +36,14 @@ pub(crate) async fn build_admin_provider_pool_status_payload(
let key_ids = keys.iter().map(|key| key.id.clone()).collect::<Vec<_>>();
let runtime = match state.redis_kv_runner() {
Some(runner) => {
read_admin_provider_pool_runtime_state(&runner, &provider.id, &key_ids, pool_config)
.await
read_admin_provider_pool_runtime_state(
&runner,
&provider.id,
&key_ids,
&pool_config,
None,
)
.await
}
None => AdminProviderPoolRuntimeState::default(),
};
@@ -56,6 +62,7 @@ pub(crate) async fn build_admin_provider_pool_status_payload(
"cost_window_usage": runtime.cost_window_usage_by_key.get(&key.id).copied().unwrap_or(0),
"cost_limit": pool_config.cost_limit_per_key_tokens,
"sticky_sessions": runtime.sticky_sessions_by_key.get(&key.id).copied().unwrap_or(0),
"latency_avg_ms": runtime.latency_avg_ms_by_key.get(&key.id).copied(),
"lru_score": runtime.lru_score_by_key.get(&key.id).copied(),
})
})

910
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/writes.rs Normal file
View File

@@ -0,0 +1,910 @@
use super::keys::{
pool_cooldown_index_key, pool_cooldown_key, pool_cost_key, pool_latency_key, pool_lru_key,
pool_sticky_key, pool_stream_timeout_key,
};
use crate::handlers::admin::provider::shared::support::{
AdminProviderPoolConfig, AdminProviderPoolUnschedulableRule,
};
use aether_data::redis::RedisKvRunner;
use regex::Regex;
use std::collections::BTreeMap;
use std::time::{SystemTime, UNIX_EPOCH};
use tracing::warn;
use uuid::Uuid;
const ACCOUNT_DISABLE_PATTERNS: &[&str] = &[
"organization has been disabled",
"organization_disabled",
"account has been disabled",
"account_disabled",
"account has been deactivated",
"account_deactivated",
"account deactivated",
];
const FORBIDDEN_ACCOUNT_PATTERNS: &[&str] = &[
"account suspended",
"account banned",
"account deactivated",
"subscription inactive",
"suspended",
"banned",
"deactivated",
];
fn current_unix_secs_f64() -> f64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs_f64()
}
fn enabled_pool_presets(pool_config: &AdminProviderPoolConfig) -> impl Iterator<Item = &str> {
pool_config
.scheduling_presets
.iter()
.filter(|item| item.enabled)
.map(|item| item.preset.as_str())
}
fn should_touch_lru(pool_config: &AdminProviderPoolConfig) -> bool {
pool_config.lru_enabled || enabled_pool_presets(pool_config).next().is_some()
}
fn should_record_latency(pool_config: &AdminProviderPoolConfig) -> bool {
enabled_pool_presets(pool_config).any(|preset| preset.eq_ignore_ascii_case("latency_first"))
}
fn oauth_cache_key(key_id: &str) -> String {
format!("provider_oauth_token_cache:{key_id}")
}
fn parse_retry_after_seconds(headers: Option<&BTreeMap<String, String>>) -> Option<u64> {
let raw = headers.and_then(|headers| {
headers
.iter()
.find(|(key, _)| key.eq_ignore_ascii_case("retry-after"))
.map(|(_, value)| value.trim())
.filter(|value| !value.is_empty())
})?;
let seconds = raw.parse::<u64>().ok()?;
Some(seconds.clamp(1, 3600))
}
fn parse_google_quota_duration_seconds(raw: &serde_json::Value) -> Option<u64> {
match raw {
serde_json::Value::Number(number) => {
let seconds = number.as_f64()?;
Some((seconds.max(1.0).floor() as u64).max(1))
}
serde_json::Value::String(text) => {
let trimmed = text.trim().to_ascii_lowercase();
if trimmed.is_empty() {
return None;
}
let regex = Regex::new(r"(\d+(?:\.\d+)?)([dhms])").ok()?;
let mut total_seconds = 0.0;
let mut matched = false;
for capture in regex.captures_iter(&trimmed) {
let amount = capture.get(1)?.as_str().parse::<f64>().ok()?;
let unit = capture.get(2)?.as_str();
matched = true;
total_seconds += match unit {
"d" => amount * 86_400.0,
"h" => amount * 3_600.0,
"m" => amount * 60.0,
"s" => amount,
_ => 0.0,
};
}
matched.then(|| (total_seconds.max(1.0).floor() as u64).max(1))
}
_ => None,
}
}
fn parse_google_quota_cooldown_seconds_at(
error_body: Option<&str>,
now_unix_secs: u64,
) -> Option<u64> {
let error_body = error_body
.map(str::trim)
.filter(|value| !value.is_empty())?;
let payload = serde_json::from_str::<serde_json::Value>(error_body).ok()?;
let error = payload.get("error")?.as_object()?;
if let Some(details) = error.get("details").and_then(serde_json::Value::as_array) {
for detail in details {
let Some(metadata) = detail
.get("metadata")
.and_then(serde_json::Value::as_object)
else {
continue;
};
if let Some(reset_at_text) = metadata
.get("quotaResetTimeStamp")
.or_else(|| metadata.get("quotaResetTimestamp"))
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
if let Ok(parsed) = chrono::DateTime::parse_from_rfc3339(reset_at_text) {
let reset_at = parsed.timestamp().max(0) as u64;
return Some(reset_at.saturating_sub(now_unix_secs).max(1));
}
}
if let Some(delay) = metadata
.get("quotaResetDelay")
.and_then(parse_google_quota_duration_seconds)
{
return Some(delay.max(1));
}
}
}
let message = error
.get("message")
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())?;
let regex = Regex::new(r"(?i)reset after\s+([^.,;]+)").ok()?;
let capture = regex.captures(message)?;
parse_google_quota_duration_seconds(&serde_json::Value::String(
capture.get(1)?.as_str().to_string(),
))
}
fn parse_google_quota_cooldown_seconds(error_body: Option<&str>) -> Option<u64> {
let now_unix_secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
parse_google_quota_cooldown_seconds_at(error_body, now_unix_secs)
}
fn extract_error_message(error_body: Option<&str>) -> String {
let Some(error_body) = error_body.map(str::trim).filter(|value| !value.is_empty()) else {
return String::new();
};
serde_json::from_str::<serde_json::Value>(error_body)
.ok()
.and_then(|value| {
value
.as_object()
.and_then(|object| object.get("error").or_else(|| object.get("message")))
.and_then(|error| match error {
serde_json::Value::Object(object) => object
.get("message")
.and_then(serde_json::Value::as_str)
.map(ToOwned::to_owned),
serde_json::Value::String(text) => Some(text.clone()),
_ => None,
})
})
.unwrap_or_else(|| error_body.chars().take(500).collect())
}
fn resolve_transient_cooldown_ttl(
status_code: u16,
retry_after_seconds: Option<u64>,
pool_config: &AdminProviderPoolConfig,
) -> u64 {
if matches!(status_code, 429 | 503) {
if let Some(retry_after_seconds) = retry_after_seconds {
return retry_after_seconds;
}
}
if status_code == 429 {
return pool_config.rate_limit_cooldown_seconds;
}
pool_config.overload_cooldown_seconds
}
async fn set_pool_cooldown(
runner: &RedisKvRunner,
provider_id: &str,
key_id: &str,
reason: &str,
ttl_seconds: u64,
) {
if ttl_seconds == 0 {
return;
}
let Ok(mut connection) = runner.client().get_multiplexed_async_connection().await else {
warn!(
"gateway admin provider pool: failed to connect redis to set cooldown for key {key_id}"
);
return;
};
let keyspace = runner.keyspace().clone();
let result: Result<(), _> = redis::pipe()
.cmd("SETEX")
.arg(pool_cooldown_key(&keyspace, provider_id, key_id))
.arg(ttl_seconds)
.arg(reason)
.ignore()
.cmd("SADD")
.arg(pool_cooldown_index_key(&keyspace, provider_id))
.arg(key_id)
.ignore()
.cmd("EXPIRE")
.arg(pool_cooldown_index_key(&keyspace, provider_id))
.arg(ttl_seconds.saturating_add(60))
.ignore()
.query_async(&mut connection)
.await;
if let Err(err) = result {
warn!(
"gateway admin provider pool: failed to set cooldown for provider {provider_id} key {key_id}: {:?}",
err
);
}
}
async fn invalidate_pool_oauth_cache(runner: &RedisKvRunner, key_id: &str) {
if let Err(err) = runner.del(&oauth_cache_key(key_id)).await {
warn!(
"gateway admin provider pool: failed to invalidate oauth cache for key {key_id}: {:?}",
err
);
}
}
fn matching_unschedulable_rule<'a>(
rules: &'a [AdminProviderPoolUnschedulableRule],
error_message: &str,
) -> Option<&'a AdminProviderPoolUnschedulableRule> {
rules.iter().find(|rule| {
let keyword = rule.keyword.trim().to_ascii_lowercase();
!keyword.is_empty() && error_message.contains(keyword.as_str())
})
}
pub(crate) async fn record_admin_provider_pool_success(
runner: &RedisKvRunner,
provider_id: &str,
key_id: &str,
pool_config: &AdminProviderPoolConfig,
sticky_session_token: Option<&str>,
tokens_used: u64,
ttfb_ms: Option<u64>,
) {
let Ok(mut connection) = runner.client().get_multiplexed_async_connection().await else {
warn!("gateway admin provider pool: failed to connect redis to record success for key {key_id}");
return;
};
let keyspace = runner.keyspace().clone();
let now = current_unix_secs_f64();
let mut pipeline = redis::pipe();
let mut has_commands = false;
if let Some(sticky_session_token) = sticky_session_token
.map(str::trim)
.filter(|value| !value.is_empty())
.filter(|_| pool_config.sticky_session_ttl_seconds > 0)
{
pipeline
.cmd("SETEX")
.arg(pool_sticky_key(
&keyspace,
provider_id,
sticky_session_token,
))
.arg(pool_config.sticky_session_ttl_seconds)
.arg(key_id)
.ignore();
has_commands = true;
}
if should_touch_lru(pool_config) {
pipeline
.cmd("ZADD")
.arg(pool_lru_key(&keyspace, provider_id))
.arg(now)
.arg(key_id)
.ignore();
has_commands = true;
}
if tokens_used > 0 && pool_config.cost_limit_per_key_tokens.is_some() {
let cost_key = pool_cost_key(&keyspace, provider_id, key_id);
let window_seconds = pool_config.cost_window_seconds.max(1);
let member = format!("{}:{tokens_used}", Uuid::new_v4().simple());
pipeline
.cmd("ZADD")
.arg(&cost_key)
.arg(now)
.arg(member)
.ignore()
.cmd("ZREMRANGEBYSCORE")
.arg(&cost_key)
.arg("-inf")
.arg(now - window_seconds as f64)
.ignore()
.cmd("EXPIRE")
.arg(&cost_key)
.arg(window_seconds.saturating_add(600))
.ignore();
has_commands = true;
}
if let Some(ttfb_ms) = ttfb_ms
.filter(|value| should_record_latency(pool_config))
.filter(|_| pool_config.latency_window_seconds > 0)
{
let latency_key = pool_latency_key(&keyspace, provider_id, key_id);
let window_seconds = pool_config.latency_window_seconds.max(1);
let sample_limit = pool_config.latency_sample_limit.max(1);
let member = format!("{}:{ttfb_ms}", Uuid::new_v4().simple());
pipeline
.cmd("ZADD")
.arg(&latency_key)
.arg(now)
.arg(member)
.ignore()
.cmd("ZREMRANGEBYSCORE")
.arg(&latency_key)
.arg("-inf")
.arg(now - window_seconds as f64)
.ignore()
.cmd("ZREMRANGEBYRANK")
.arg(&latency_key)
.arg(0)
.arg(-((sample_limit as i64) + 1))
.ignore()
.cmd("EXPIRE")
.arg(&latency_key)
.arg(window_seconds.saturating_add(600))
.ignore();
has_commands = true;
}
if !has_commands {
return;
}
let result: Result<(), _> = pipeline.query_async(&mut connection).await;
if let Err(err) = result {
warn!(
"gateway admin provider pool: failed to record success feedback for provider {provider_id} key {key_id}: {:?}",
err
);
}
}
pub(crate) async fn record_admin_provider_pool_error(
runner: &RedisKvRunner,
provider_id: &str,
key_id: &str,
pool_config: &AdminProviderPoolConfig,
status_code: u16,
error_body: Option<&str>,
response_headers: Option<&BTreeMap<String, String>>,
) {
if !pool_config.health_policy_enabled {
return;
}
let error_message = extract_error_message(error_body).to_ascii_lowercase();
if status_code == 401 {
invalidate_pool_oauth_cache(runner, key_id).await;
if ACCOUNT_DISABLE_PATTERNS
.iter()
.any(|pattern| error_message.contains(pattern))
{
set_pool_cooldown(runner, provider_id, key_id, "account_deactivated_401", 3600).await;
}
return;
}
if status_code == 402 {
set_pool_cooldown(runner, provider_id, key_id, "payment_required_402", 3600).await;
return;
}
if status_code == 403 {
let severe = FORBIDDEN_ACCOUNT_PATTERNS
.iter()
.any(|pattern| error_message.contains(pattern));
let ttl_seconds = if severe {
3600
} else {
pool_config.rate_limit_cooldown_seconds.max(300)
};
set_pool_cooldown(runner, provider_id, key_id, "forbidden_403", ttl_seconds).await;
return;
}
if status_code == 400 {
if let Some(pattern) = ACCOUNT_DISABLE_PATTERNS
.iter()
.find(|pattern| error_message.contains(**pattern))
{
set_pool_cooldown(
runner,
provider_id,
key_id,
&format!("account_disabled_400:{pattern}"),
3600,
)
.await;
return;
}
}
if let Some(rule) =
matching_unschedulable_rule(&pool_config.unschedulable_rules, &error_message)
{
let ttl_seconds = (rule.duration_minutes.max(1)).saturating_mul(60).max(60);
set_pool_cooldown(
runner,
provider_id,
key_id,
&format!("rule:{}", rule.keyword),
ttl_seconds,
)
.await;
return;
}
if status_code == 429 {
let ttl_seconds = resolve_transient_cooldown_ttl(
status_code,
parse_retry_after_seconds(response_headers)
.or_else(|| parse_google_quota_cooldown_seconds(error_body)),
pool_config,
);
set_pool_cooldown(runner, provider_id, key_id, "rate_limited_429", ttl_seconds).await;
return;
}
if status_code == 529 {
set_pool_cooldown(
runner,
provider_id,
key_id,
"overloaded_529",
pool_config.overload_cooldown_seconds,
)
.await;
return;
}
let transient_reason = match status_code {
408 => Some("request_timeout_408".to_string()),
409 => Some("conflict_409".to_string()),
423 => Some("locked_423".to_string()),
425 => Some("too_early_425".to_string()),
500 => Some("server_error_500".to_string()),
502 => Some("bad_gateway_502".to_string()),
503 => Some("service_unavailable_503".to_string()),
504 => Some("gateway_timeout_504".to_string()),
501 | 505..=599 => Some(format!("server_error_{status_code}")),
_ => None,
};
if let Some(reason) = transient_reason {
let ttl_seconds = resolve_transient_cooldown_ttl(
status_code,
parse_retry_after_seconds(response_headers),
pool_config,
);
set_pool_cooldown(runner, provider_id, key_id, &reason, ttl_seconds).await;
}
}
pub(crate) async fn record_admin_provider_pool_stream_timeout(
runner: &RedisKvRunner,
provider_id: &str,
key_id: &str,
pool_config: &AdminProviderPoolConfig,
) {
if !pool_config.health_policy_enabled || pool_config.stream_timeout_threshold == 0 {
return;
}
let Ok(mut connection) = runner.client().get_multiplexed_async_connection().await else {
warn!("gateway admin provider pool: failed to connect redis to record stream timeout for key {key_id}");
return;
};
let keyspace = runner.keyspace().clone();
let timeout_key = pool_stream_timeout_key(&keyspace, provider_id, key_id);
let now = current_unix_secs_f64();
let window_seconds = pool_config.stream_timeout_window_seconds.max(1);
let member = Uuid::new_v4().simple().to_string();
let results = redis::pipe()
.cmd("ZREMRANGEBYSCORE")
.arg(&timeout_key)
.arg("-inf")
.arg(now - window_seconds as f64)
.cmd("ZADD")
.arg(&timeout_key)
.arg(now)
.arg(member)
.cmd("ZCARD")
.arg(&timeout_key)
.cmd("EXPIRE")
.arg(&timeout_key)
.arg(window_seconds.saturating_add(60))
.query_async::<Vec<redis::Value>>(&mut connection)
.await;
let count = match results
.ok()
.and_then(|values| values.get(2).cloned())
.and_then(|value| redis::from_redis_value::<u64>(&value).ok())
{
Some(count) => count,
None => {
warn!(
"gateway admin provider pool: failed to compute stream timeout count for provider {provider_id} key {key_id}"
);
return;
}
};
if count >= pool_config.stream_timeout_threshold {
set_pool_cooldown(
runner,
provider_id,
key_id,
&format!("stream_timeout_x{count}"),
pool_config.stream_timeout_cooldown_seconds.max(1),
)
.await;
}
}
#[cfg(test)]
mod tests {
use super::{
parse_google_quota_cooldown_seconds_at, record_admin_provider_pool_error,
record_admin_provider_pool_stream_timeout, record_admin_provider_pool_success,
};
use crate::data::{GatewayDataConfig, GatewayDataState};
use crate::handlers::admin::provider::pool::runtime::reads::read_admin_provider_pool_runtime_state;
use crate::handlers::admin::provider::shared::support::{
AdminProviderPoolConfig, AdminProviderPoolSchedulingPreset,
AdminProviderPoolUnschedulableRule,
};
use crate::AppState;
use aether_testkit::ManagedRedisServer;
use std::collections::BTreeMap;
async fn start_managed_redis_or_skip() -> Option<ManagedRedisServer> {
match ManagedRedisServer::start().await {
Ok(server) => Some(server),
Err(err) if err.to_string().contains("No such file or directory") => {
eprintln!("skipping redis-backed pool runtime test: {err}");
None
}
Err(err) => panic!("redis server should start: {err}"),
}
}
fn sample_pool_config() -> AdminProviderPoolConfig {
AdminProviderPoolConfig {
scheduling_presets: vec![AdminProviderPoolSchedulingPreset {
preset: "latency_first".to_string(),
enabled: true,
mode: None,
}],
unschedulable_rules: Vec::new(),
lru_enabled: true,
skip_exhausted_accounts: false,
sticky_session_ttl_seconds: 120,
latency_window_seconds: 600,
latency_sample_limit: 10,
cost_window_seconds: 600,
cost_limit_per_key_tokens: Some(10_000),
rate_limit_cooldown_seconds: 300,
overload_cooldown_seconds: 30,
health_policy_enabled: true,
stream_timeout_threshold: 3,
stream_timeout_window_seconds: 1800,
stream_timeout_cooldown_seconds: 300,
}
}
fn build_runner_app(redis_url: &str, key_prefix: &str) -> AppState {
let data_state = GatewayDataState::from_config(
GatewayDataConfig::disabled().with_redis_url(redis_url, Some(key_prefix)),
)
.expect("data state should build");
AppState::new()
.expect("app state should build")
.with_data_state_for_tests(data_state)
}
#[test]
fn parses_google_quota_cooldown_from_reset_timestamp() {
let now_unix_secs = chrono::DateTime::parse_from_rfc3339("2026-04-17T10:00:00Z")
.expect("timestamp should parse")
.timestamp()
.max(0) as u64;
let cooldown = parse_google_quota_cooldown_seconds_at(
Some(
r#"{
"error": {
"message": "Quota exhausted.",
"details": [{
"metadata": {
"quotaResetTimeStamp": "2026-04-17T10:01:30Z"
}
}]
}
}"#,
),
now_unix_secs,
);
assert_eq!(cooldown, Some(90));
}
#[test]
fn parses_google_quota_cooldown_from_reset_delay_and_message() {
let now_unix_secs = chrono::DateTime::parse_from_rfc3339("2026-04-17T10:00:00Z")
.expect("timestamp should parse")
.timestamp()
.max(0) as u64;
let delay_cooldown = parse_google_quota_cooldown_seconds_at(
Some(
r#"{
"error": {
"message": "Quota exhausted.",
"details": [{
"metadata": {
"quotaResetDelay": "1h30m15s"
}
}]
}
}"#,
),
now_unix_secs,
);
let message_cooldown = parse_google_quota_cooldown_seconds_at(
Some(
r#"{
"error": {
"message": "Too many requests, reset after 45m."
}
}"#,
),
now_unix_secs,
);
assert_eq!(delay_cooldown, Some(5_415));
assert_eq!(message_cooldown, Some(2_700));
}
#[tokio::test]
async fn success_feedback_writes_sticky_lru_cost_and_latency() {
let Some(redis) = start_managed_redis_or_skip().await else {
return;
};
let app = build_runner_app(redis.redis_url(), "pool_runtime_success_feedback");
let runner = app.redis_kv_runner().expect("redis runner should exist");
let pool_config = sample_pool_config();
let key_ids = vec!["key-1".to_string()];
record_admin_provider_pool_success(
&runner,
"provider-1",
"key-1",
&pool_config,
Some("session-1"),
120,
Some(80),
)
.await;
let runtime = read_admin_provider_pool_runtime_state(
&runner,
"provider-1",
&key_ids,
&pool_config,
Some("session-1"),
)
.await;
assert_eq!(runtime.total_sticky_sessions, 1);
assert_eq!(runtime.sticky_bound_key_id.as_deref(), Some("key-1"));
assert_eq!(runtime.sticky_sessions_by_key.get("key-1"), Some(&1));
assert_eq!(runtime.cost_window_usage_by_key.get("key-1"), Some(&120));
assert_eq!(runtime.latency_avg_ms_by_key.get("key-1"), Some(&80.0));
assert!(runtime.lru_score_by_key.contains_key("key-1"));
}
#[tokio::test]
async fn error_feedback_respects_retry_after_for_rate_limits() {
let Some(redis) = start_managed_redis_or_skip().await else {
return;
};
let app = build_runner_app(redis.redis_url(), "pool_runtime_error_feedback");
let runner = app.redis_kv_runner().expect("redis runner should exist");
let pool_config = sample_pool_config();
let key_ids = vec!["key-2".to_string()];
record_admin_provider_pool_error(
&runner,
"provider-1",
"key-2",
&pool_config,
429,
Some(r#"{"error":{"message":"rate limited"}}"#),
Some(&BTreeMap::from([(
"Retry-After".to_string(),
"120".to_string(),
)])),
)
.await;
let runtime = read_admin_provider_pool_runtime_state(
&runner,
"provider-1",
&key_ids,
&pool_config,
None,
)
.await;
assert_eq!(
runtime
.cooldown_reason_by_key
.get("key-2")
.map(String::as_str),
Some("rate_limited_429")
);
assert!(runtime
.cooldown_ttl_by_key
.get("key-2")
.is_some_and(|ttl| *ttl <= 120 && *ttl >= 100));
}
#[tokio::test]
async fn error_feedback_uses_google_quota_cooldown_when_retry_after_missing() {
let Some(redis) = start_managed_redis_or_skip().await else {
return;
};
let app = build_runner_app(redis.redis_url(), "pool_runtime_google_quota_cooldown");
let runner = app.redis_kv_runner().expect("redis runner should exist");
let pool_config = sample_pool_config();
let key_ids = vec!["key-google-429".to_string()];
record_admin_provider_pool_error(
&runner,
"provider-1",
"key-google-429",
&pool_config,
429,
Some(
r#"{
"error": {
"message": "Quota exhausted. reset after 45s.",
"status": "RESOURCE_EXHAUSTED",
"details": [{
"metadata": {
"quotaResetDelay": "45s"
}
}]
}
}"#,
),
None,
)
.await;
let runtime = read_admin_provider_pool_runtime_state(
&runner,
"provider-1",
&key_ids,
&pool_config,
None,
)
.await;
assert_eq!(
runtime
.cooldown_reason_by_key
.get("key-google-429")
.map(String::as_str),
Some("rate_limited_429")
);
assert!(runtime
.cooldown_ttl_by_key
.get("key-google-429")
.is_some_and(|ttl| *ttl <= 45 && *ttl >= 30));
}
#[tokio::test]
async fn error_feedback_applies_unschedulable_rule_cooldown() {
let Some(redis) = start_managed_redis_or_skip().await else {
return;
};
let app = build_runner_app(redis.redis_url(), "pool_runtime_unschedulable_rule");
let runner = app.redis_kv_runner().expect("redis runner should exist");
let mut pool_config = sample_pool_config();
pool_config.unschedulable_rules = vec![AdminProviderPoolUnschedulableRule {
keyword: "review required".to_string(),
duration_minutes: 7,
}];
let key_ids = vec!["key-3".to_string()];
record_admin_provider_pool_error(
&runner,
"provider-1",
"key-3",
&pool_config,
418,
Some(r#"{"error":{"message":"manual review required before reuse"}}"#),
None,
)
.await;
let runtime = read_admin_provider_pool_runtime_state(
&runner,
"provider-1",
&key_ids,
&pool_config,
None,
)
.await;
assert_eq!(
runtime
.cooldown_reason_by_key
.get("key-3")
.map(String::as_str),
Some("rule:review required")
);
assert!(runtime
.cooldown_ttl_by_key
.get("key-3")
.is_some_and(|ttl| *ttl <= 420 && *ttl >= 380));
}
#[tokio::test]
async fn stream_timeout_policy_cools_down_after_threshold() {
let Some(redis) = start_managed_redis_or_skip().await else {
return;
};
let app = build_runner_app(redis.redis_url(), "pool_runtime_stream_timeout");
let runner = app.redis_kv_runner().expect("redis runner should exist");
let mut pool_config = sample_pool_config();
pool_config.stream_timeout_threshold = 2;
pool_config.stream_timeout_window_seconds = 300;
pool_config.stream_timeout_cooldown_seconds = 90;
let key_ids = vec!["key-4".to_string()];
record_admin_provider_pool_stream_timeout(&runner, "provider-1", "key-4", &pool_config)
.await;
record_admin_provider_pool_stream_timeout(&runner, "provider-1", "key-4", &pool_config)
.await;
let runtime = read_admin_provider_pool_runtime_state(
&runner,
"provider-1",
&key_ids,
&pool_config,
None,
)
.await;
assert_eq!(
runtime
.cooldown_reason_by_key
.get("key-4")
.map(String::as_str),
Some("stream_timeout_x2")
);
assert!(runtime
.cooldown_ttl_by_key
.get("key-4")
.is_some_and(|ttl| *ttl <= 90 && *ttl >= 70));
}
}

8
apps/aether-gateway/src/handlers/admin/provider/pool_admin/payloads.rs
View File

@@ -772,7 +772,9 @@ pub(super) fn build_admin_pool_key_payload(
let health_score = admin_pool_health_score(key);
let circuit_breaker_open = admin_pool_circuit_breaker_open(key);
let auth_semantics = provider_key_auth_semantics(key, provider_type);
let account_quota_exhausted = pool_config.is_some_and(|config| config.skip_exhausted_accounts)
let account_quota_exhausted = pool_config
.as_ref()
.is_some_and(|config| config.skip_exhausted_accounts)
&& admin_provider_pool_pure::admin_pool_key_account_quota_exhausted(key, provider_type);
let (scheduling_status, scheduling_reason, scheduling_label, scheduling_reasons) =
admin_pool_scheduling_payload(
@@ -988,7 +990,9 @@ pub(super) fn build_admin_pool_key_payload(
);
payload.insert(
"cost_limit".to_string(),
json!(pool_config.map(|config| config.cost_limit_per_key_tokens)),
json!(pool_config
.as_ref()
.map(|config| config.cost_limit_per_key_tokens)),
);
payload.insert(
"request_count".to_string(),

14
apps/aether-gateway/src/handlers/admin/provider/pool_admin/read_routes/keys.rs
View File

@@ -125,10 +125,16 @@ pub(super) async fn build_admin_pool_list_keys_response(
};
let key_ids = keys.iter().map(|key| key.id.clone()).collect::<Vec<_>>();
let runtime = match (state.redis_kv_runner(), pool_config) {
let runtime = match (state.redis_kv_runner(), pool_config.as_ref()) {
(Some(runner), Some(pool_config)) if !key_ids.is_empty() => {
read_admin_provider_pool_runtime_state(&runner, &provider.id, &key_ids, pool_config)
.await
read_admin_provider_pool_runtime_state(
&runner,
&provider.id,
&key_ids,
pool_config,
None,
)
.await
}
_ => AdminProviderPoolRuntimeState::default(),
};
@@ -145,7 +151,7 @@ pub(super) async fn build_admin_pool_list_keys_response(
&provider.provider_type,
&key,
&runtime,
pool_config,
pool_config.clone(),
);
if let Some(summary) = usage_summary_by_key_id.get(&key.id) {
if let Some(object) = payload.as_object_mut() {

28
apps/aether-gateway/src/handlers/admin/provider/shared/support.rs
View File

@@ -10,21 +10,47 @@ pub(crate) const ADMIN_PROVIDER_POOL_SCAN_BATCH: u64 = 200;
pub(crate) const ADMIN_PROVIDER_OAUTH_DATA_UNAVAILABLE_DETAIL: &str =
"Admin provider OAuth data unavailable";
#[derive(Debug, Clone, Copy)]
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct AdminProviderPoolSchedulingPreset {
pub(crate) preset: String,
pub(crate) enabled: bool,
pub(crate) mode: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct AdminProviderPoolUnschedulableRule {
pub(crate) keyword: String,
pub(crate) duration_minutes: u64,
}
#[derive(Debug, Clone)]
pub(crate) struct AdminProviderPoolConfig {
pub(crate) scheduling_presets: Vec<AdminProviderPoolSchedulingPreset>,
pub(crate) unschedulable_rules: Vec<AdminProviderPoolUnschedulableRule>,
pub(crate) lru_enabled: bool,
pub(crate) skip_exhausted_accounts: bool,
pub(crate) sticky_session_ttl_seconds: u64,
pub(crate) latency_window_seconds: u64,
pub(crate) latency_sample_limit: u64,
pub(crate) cost_window_seconds: u64,
pub(crate) cost_limit_per_key_tokens: Option<u64>,
pub(crate) rate_limit_cooldown_seconds: u64,
pub(crate) overload_cooldown_seconds: u64,
pub(crate) health_policy_enabled: bool,
pub(crate) stream_timeout_threshold: u64,
pub(crate) stream_timeout_window_seconds: u64,
pub(crate) stream_timeout_cooldown_seconds: u64,
}
#[derive(Debug, Default)]
pub(crate) struct AdminProviderPoolRuntimeState {
pub(crate) total_sticky_sessions: usize,
pub(crate) sticky_sessions_by_key: BTreeMap<String, usize>,
pub(crate) sticky_bound_key_id: Option<String>,
pub(crate) cooldown_reason_by_key: BTreeMap<String, String>,
pub(crate) cooldown_ttl_by_key: BTreeMap<String, u64>,
pub(crate) cost_window_usage_by_key: BTreeMap<String, u64>,
pub(crate) latency_avg_ms_by_key: BTreeMap<String, f64>,
pub(crate) lru_score_by_key: BTreeMap<String, f64>,
}

1
apps/aether-gateway/src/handlers/shared/mod.rs
View File

@@ -5,6 +5,7 @@ mod email_templates;
mod external_models;
mod normalize;
mod payloads;
pub(crate) mod provider_pool;
mod request_utils;
mod system_config_values;
mod usage_stats;

9
apps/aether-gateway/src/handlers/shared/provider_pool.rs Normal file
View File

@@ -0,0 +1,9 @@
pub(crate) use super::super::admin::provider::pool::config::admin_provider_pool_config_from_config_value;
pub(crate) use super::super::admin::provider::pool::runtime::{
read_admin_provider_pool_runtime_state, record_admin_provider_pool_error,
record_admin_provider_pool_stream_timeout, record_admin_provider_pool_success,
};
pub(crate) use super::super::admin::provider::shared::support::{
AdminProviderPoolConfig, AdminProviderPoolRuntimeState, AdminProviderPoolSchedulingPreset,
AdminProviderPoolUnschedulableRule, ADMIN_PROVIDER_POOL_SCAN_BATCH,
};