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

Unify pool score display as account health

Browse Source
This commit is contained in:
fawney19
2026-05-12 15:27:21 +08:00
parent fa73655134
commit 9c05b5f4e0
14 changed files with 401 additions and 518 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
apps/aether-gateway/src/ai_serving/mod.rs
View File

@@ -49,7 +49,7 @@ pub(crate) use self::planner::{
build_standard_sync_plan_from_decision, candidate_auth_channel_skip_reason,
extract_pool_sticky_session_token, maybe_build_stream_decision_payload,
maybe_build_stream_plan_payload, maybe_build_sync_decision_payload,
maybe_build_sync_plan_payload, planner_is_matching_stream_request,
maybe_build_sync_plan_payload, planner_is_matching_stream_request, provider_key_pool_score_id,
provider_key_pool_score_scope, read_candidate_transport_snapshot,
record_local_runtime_candidate_skip_reason,
set_local_openai_chat_execution_exhausted_diagnostic,

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

@@ -41,8 +41,9 @@ pub(crate) use self::plan_builders::{
build_standard_stream_plan_from_decision, build_standard_sync_plan_from_decision,
AiStreamAttempt, AiSyncAttempt,
};
pub(crate) use self::pool_scores::build_provider_key_pool_score_upsert;
pub(crate) use self::pool_scores::provider_key_pool_score_scope;
pub(crate) use self::pool_scores::{
build_provider_key_pool_score_upsert, provider_key_pool_score_id, provider_key_pool_score_scope,
};
pub(crate) use self::route::is_matching_stream_request as planner_is_matching_stream_request;
pub(crate) use self::runtime_miss::{
apply_local_runtime_candidate_terminal_reason, record_local_runtime_candidate_skip_reason,

49
apps/aether-gateway/src/ai_serving/planner/pool_scores.rs
View File

@@ -3,7 +3,7 @@ use aether_ai_serving::{
};
use aether_data_contracts::repository::pool_scores::{
PoolMemberIdentity, PoolMemberProbeStatus, PoolScoreScope, UpsertPoolMemberScore,
POOL_SCORE_SCOPE_KIND_MODEL,
POOL_SCORE_CAPABILITY_ACCOUNT, POOL_SCORE_SCOPE_KIND_ACCOUNT,
};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use serde_json::Value;
@@ -13,14 +13,12 @@ use crate::handlers::shared::{provider_key_health_summary, provider_key_status_s
pub(crate) fn build_provider_key_pool_score_upsert(
key: &StoredProviderCatalogKey,
provider_type: &str,
api_format: &str,
model_id: Option<&str>,
existing: Option<&aether_data_contracts::repository::pool_scores::StoredPoolMemberScore>,
now_unix_secs: u64,
score_rules: PoolMemberScoreRules,
) -> UpsertPoolMemberScore {
let identity = PoolMemberIdentity::provider_api_key(key.provider_id.clone(), key.id.clone());
let scope = provider_key_pool_score_scope(api_format, model_id);
let scope = provider_key_pool_score_scope();
let input = provider_key_score_input(
key,
provider_type,
@@ -54,17 +52,11 @@ pub(crate) fn build_provider_key_pool_score_upsert(
}
}
pub(crate) fn provider_key_pool_score_scope(
api_format: &str,
model_id: Option<&str>,
) -> PoolScoreScope {
pub(crate) fn provider_key_pool_score_scope() -> PoolScoreScope {
PoolScoreScope {
capability: api_format.trim().to_ascii_lowercase(),
scope_kind: POOL_SCORE_SCOPE_KIND_MODEL.to_string(),
scope_id: model_id
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
capability: POOL_SCORE_CAPABILITY_ACCOUNT.to_string(),
scope_kind: POOL_SCORE_SCOPE_KIND_ACCOUNT.to_string(),
scope_id: None,
}
}
@@ -133,8 +125,7 @@ fn provider_key_score_input(
.and_then(Value::as_bool)
.unwrap_or(false),
oauth_invalid_reason: key.oauth_invalid_reason.clone(),
circuit_open: any_circuit_open
|| circuit_open_for_scope(key.circuit_breaker_by_format.as_ref(), &scope),
circuit_open: any_circuit_open,
success_count: key.success_count.unwrap_or(0).into(),
error_count: key.error_count.unwrap_or(0).into(),
total_response_time_ms: key.total_response_time_ms.unwrap_or(0).into(),
@@ -150,21 +141,6 @@ fn provider_key_score_input(
}
}
fn circuit_open_for_scope(circuit_by_format: Option<&Value>, scope: &PoolScoreScope) -> bool {
let Some(formats) = circuit_by_format.and_then(Value::as_object) else {
return false;
};
let keys = api_format_lookup_keys(&scope.capability);
keys.iter().any(|key| {
formats
.get(key)
.and_then(Value::as_object)
.and_then(|value| value.get("open"))
.and_then(Value::as_bool)
.unwrap_or(false)
})
}
fn json_f64(value: &Value) -> Option<f64> {
value.as_f64().or_else(|| {
value
@@ -175,17 +151,6 @@ fn json_f64(value: &Value) -> Option<f64> {
})
}
fn api_format_lookup_keys(api_format: &str) -> Vec<String> {
let normalized = crate::ai_serving::normalize_api_format_alias(api_format);
let mut keys = crate::ai_serving::api_format_storage_aliases(&normalized);
if !keys.iter().any(|value| value == &normalized) {
keys.push(normalized);
}
keys.sort();
keys.dedup();
keys
}
fn stable_hash(bytes: &[u8]) -> u64 {
let mut hash = 0xcbf29ce484222325u64;
for byte in bytes {

5
apps/aether-gateway/src/dispatch/pool_scheduler.rs
View File

@@ -387,10 +387,7 @@ impl<'a> PoolKeyCursor<'a> {
}
async fn next_score_candidates(&mut self) -> Option<Vec<EligibleLocalExecutionCandidate>> {
let scope = provider_key_pool_score_scope(
self.group.candidate.endpoint_api_format.as_str(),
Some(self.group.candidate.model_id.as_str()),
);
let scope = provider_key_pool_score_scope();
let query = ListRankedPoolMembersQuery {
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
pool_id: self.group.candidate.provider_id.clone(),

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

@@ -6,6 +6,7 @@ use crate::handlers::admin::shared::{provider_key_status_snapshot_payload, unix_
use crate::provider_key_auth::{provider_key_auth_semantics, provider_key_effective_api_formats};
use aether_admin::provider::pool as admin_provider_pool_pure;
use aether_admin::provider::quota as admin_provider_quota_pure;
use aether_data_contracts::repository::pool_scores::StoredPoolMemberScore;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
};
@@ -906,6 +907,7 @@ pub(super) fn build_admin_pool_key_payload(
key: &StoredProviderCatalogKey,
runtime: &AdminProviderPoolRuntimeState,
pool_config: Option<AdminProviderPoolConfig>,
pool_score: Option<&StoredPoolMemberScore>,
codex_cycle_usage_by_code: Option<&BTreeMap<String, StoredProviderApiKeyWindowUsageSummary>>,
now_unix_secs: u64,
) -> serde_json::Value {
@@ -1087,6 +1089,34 @@ pub(super) fn build_admin_pool_key_payload(
payload.insert("status_snapshot".to_string(), status_snapshot);
payload.insert("quota_updated_at".to_string(), json!(quota_updated_at));
payload.insert("health_score".to_string(), json!(health_score));
payload.insert(
"pool_score".to_string(),
pool_score
.map(|score| {
json!({
"id": score.id.clone(),
"capability": score.capability.clone(),
"scope_kind": score.scope_kind.clone(),
"scope_id": score.scope_id.clone(),
"score": score.score,
"hard_state": score.hard_state.as_database(),
"score_version": score.score_version,
"score_reason": score.score_reason.clone(),
"last_ranked_at": score.last_ranked_at,
"last_scheduled_at": score.last_scheduled_at,
"last_success_at": score.last_success_at,
"last_failure_at": score.last_failure_at,
"failure_count": score.failure_count,
"last_probe_attempt_at": score.last_probe_attempt_at,
"last_probe_success_at": score.last_probe_success_at,
"last_probe_failure_at": score.last_probe_failure_at,
"probe_failure_count": score.probe_failure_count,
"probe_status": score.probe_status.as_database(),
"updated_at": score.updated_at,
})
})
.unwrap_or(serde_json::Value::Null),
);
payload.insert(
"circuit_breaker_open".to_string(),
json!(circuit_breaker_open),

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

@@ -7,9 +7,13 @@ use super::{
AdminPoolKeySortField, AdminProviderPoolRuntimeState, ProviderCatalogKeyListOrder,
ProviderCatalogKeyListQuery, ADMIN_POOL_PROVIDER_CATALOG_READER_UNAVAILABLE_DETAIL,
};
use crate::ai_serving::{provider_key_pool_score_id, provider_key_pool_score_scope};
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::GatewayError;
use aether_admin::provider::pool as admin_provider_pool_pure;
use aether_data_contracts::repository::pool_scores::{
GetPoolMemberScoresByIdsQuery, PoolMemberIdentity, StoredPoolMemberScore,
};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use aether_data_contracts::repository::usage::{
ProviderApiKeyWindowUsageRequest, StoredProviderApiKeyWindowUsageSummary,
@@ -56,6 +60,36 @@ fn admin_pool_current_unix_secs() -> u64 {
.unwrap_or(0)
}
async fn read_admin_pool_scores_by_key_id(
state: &AdminAppState<'_>,
provider_id: &str,
key_ids: &[String],
) -> Result<BTreeMap<String, StoredPoolMemberScore>, GatewayError> {
if key_ids.is_empty() {
return Ok(BTreeMap::new());
}
let score_scope = provider_key_pool_score_scope();
let score_ids = key_ids
.iter()
.map(|key_id| {
let identity =
PoolMemberIdentity::provider_api_key(provider_id.to_string(), key_id.clone());
provider_key_pool_score_id(&identity, &score_scope)
})
.collect::<Vec<_>>();
let scores = state
.app()
.data
.get_pool_member_scores_by_ids(&GetPoolMemberScoresByIdsQuery { ids: score_ids })
.await
.map_err(|err| GatewayError::Internal(format!("{err:?}")))?;
Ok(scores
.into_iter()
.map(|score| (score.member_id.clone(), score))
.collect::<BTreeMap<_, _>>())
}
fn admin_pool_codex_cycle_usage_request(
key: &StoredProviderCatalogKey,
window: &serde_json::Map<String, serde_json::Value>,
@@ -379,6 +413,9 @@ pub(super) async fn build_admin_pool_list_keys_response(
};
let key_ids = keys.iter().map(|key| key.id.clone()).collect::<Vec<_>>();
let pool_scores_by_key_id = read_admin_pool_scores_by_key_id(state, &provider.id, &key_ids)
.await
.unwrap_or_default();
let endpoints = state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider.id))
.await?;
@@ -414,6 +451,7 @@ pub(super) async fn build_admin_pool_list_keys_response(
&key,
&runtime,
pool_config.clone(),
pool_scores_by_key_id.get(&key.id),
codex_cycle_usage_by_key.get(&key.id),
now_unix_secs,
)

15
apps/aether-gateway/src/handlers/admin/provider/pool_admin/read_routes/scores.rs
View File

@@ -7,7 +7,7 @@ use crate::handlers::shared::query_param_value;
use crate::GatewayError;
use aether_data_contracts::repository::pool_scores::{
ListPoolMemberScoresQuery, PoolMemberHardState, PoolMemberProbeStatus,
POOL_KIND_PROVIDER_KEY_POOL, POOL_SCORE_SCOPE_KIND_MODEL,
POOL_KIND_PROVIDER_KEY_POOL, POOL_SCORE_CAPABILITY_ACCOUNT, POOL_SCORE_SCOPE_KIND_ACCOUNT,
};
use axum::{
body::Body,
@@ -49,9 +49,6 @@ pub(super) async fn build_admin_pool_scores_response(
}
};
let offset = page.saturating_sub(1).saturating_mul(page_size);
let api_format = query_param_value(query, "api_format")
.map(|value| crate::ai_serving::normalize_api_format_alias(value.as_str()));
let model_id = query_param_value(query, "model_id");
let hard_states = match parse_hard_state_filter(query) {
Ok(value) => value,
Err(message) => {
@@ -77,9 +74,9 @@ pub(super) async fn build_admin_pool_scores_response(
.list_pool_member_scores(&ListPoolMemberScoresQuery {
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
pool_id: provider_id.clone(),
capability: api_format.clone(),
scope_kind: Some(POOL_SCORE_SCOPE_KIND_MODEL.to_string()),
scope_id: model_id.clone(),
capability: Some(POOL_SCORE_CAPABILITY_ACCOUNT.to_string()),
scope_kind: Some(POOL_SCORE_SCOPE_KIND_ACCOUNT.to_string()),
scope_id: None,
hard_states,
probe_statuses,
offset,
@@ -146,8 +143,8 @@ pub(super) async fn build_admin_pool_scores_response(
"page": page,
"page_size": page_size,
"filters": {
"api_format": api_format,
"model_id": model_id,
"api_format": serde_json::Value::Null,
"model_id": serde_json::Value::Null,
"hard_state": query_param_value(query, "hard_state"),
"probe_status": query_param_value(query, "probe_status")
},

193
apps/aether-gateway/src/maintenance/runtime/pool_score_rebuild.rs
View File

@@ -1,7 +1,6 @@
use std::collections::BTreeMap;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use aether_data_contracts::repository::global_models::AdminProviderModelListQuery;
use aether_data_contracts::repository::pool_scores::GetPoolMemberScoresByIdsQuery;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
@@ -113,32 +112,11 @@ fn provider_offset_cursor_key(provider_id: &str) -> String {
format!("{POOL_SCORE_REBUILD_PROVIDER_OFFSET_PREFIX}:{provider_id}")
}
fn score_combo_indices(
flat_index: usize,
model_count: usize,
key_count: usize,
) -> (usize, usize, usize) {
let endpoint_stride = model_count.saturating_mul(key_count).max(1);
let endpoint_index = flat_index / endpoint_stride;
let remainder = flat_index % endpoint_stride;
let model_index = remainder / key_count.max(1);
let key_index = remainder % key_count.max(1);
(endpoint_index, model_index, key_index)
}
#[derive(Debug, Clone)]
struct ProviderScoreBuildItem {
endpoint_index: usize,
model_index: usize,
key_index: usize,
score_id: String,
}
pub(crate) async fn ensure_provider_key_pool_scores_for_keys(
state: &AppState,
provider: &StoredProviderCatalogProvider,
pool_config: &AdminProviderPoolConfig,
endpoints: &[StoredProviderCatalogEndpoint],
_endpoints: &[StoredProviderCatalogEndpoint],
keys: &[StoredProviderCatalogKey],
now_unix_secs: u64,
max_upserts: usize,
@@ -151,63 +129,28 @@ pub(crate) async fn ensure_provider_key_pool_scores_for_keys(
return Ok(0);
}
let endpoints = endpoints
.iter()
.filter(|endpoint| endpoint.is_active && !endpoint.api_format.trim().is_empty())
.collect::<Vec<_>>();
let keys = keys
.iter()
.filter(|key| key.is_active && key.provider_id == provider.id)
.collect::<Vec<_>>();
if endpoints.is_empty() || keys.is_empty() {
if keys.is_empty() {
return Ok(0);
}
let models = state
.list_admin_provider_models(&AdminProviderModelListQuery {
provider_id: provider.id.clone(),
is_active: Some(true),
offset: 0,
limit: 10_000,
})
.await?
let build_items = keys
.into_iter()
.filter(|model| model.is_available)
.take(max_upserts)
.map(|key| {
let draft = build_provider_key_pool_score_upsert(
key,
provider.provider_type.as_str(),
None,
now_unix_secs,
pool_config.score_rules,
);
(key, draft.id)
})
.collect::<Vec<_>>();
if models.is_empty() {
return Ok(0);
}
let max_items = endpoints
.len()
.saturating_mul(models.len())
.saturating_mul(keys.len())
.min(max_upserts);
let mut build_items = Vec::with_capacity(max_items);
'outer: for (endpoint_index, endpoint) in endpoints.iter().enumerate() {
for (model_index, model) in models.iter().enumerate() {
for (key_index, key) in keys.iter().enumerate() {
let draft = build_provider_key_pool_score_upsert(
key,
provider.provider_type.as_str(),
endpoint.api_format.trim(),
Some(model.id.as_str()),
None,
now_unix_secs,
pool_config.score_rules,
);
build_items.push(ProviderScoreBuildItem {
endpoint_index,
model_index,
key_index,
score_id: draft.id,
});
if build_items.len() >= max_upserts {
break 'outer;
}
}
}
}
if build_items.is_empty() {
return Ok(0);
}
@@ -217,7 +160,7 @@ pub(crate) async fn ensure_provider_key_pool_scores_for_keys(
.get_pool_member_scores_by_ids(&GetPoolMemberScoresByIdsQuery {
ids: build_items
.iter()
.map(|item| item.score_id.clone())
.map(|(_, score_id)| score_id.clone())
.collect(),
})
.await
@@ -234,18 +177,13 @@ pub(crate) async fn ensure_provider_key_pool_scores_for_keys(
.collect::<std::collections::BTreeSet<_>>();
let mut upserted = 0usize;
for item in &build_items {
if existing_score_ids.contains(&item.score_id) {
for (key, score_id) in &build_items {
if existing_score_ids.contains(score_id) {
continue;
}
let endpoint = endpoints[item.endpoint_index];
let model = &models[item.model_index];
let key = keys[item.key_index];
let upsert = build_provider_key_pool_score_upsert(
key,
provider.provider_type.as_str(),
endpoint.api_format.trim(),
Some(model.id.as_str()),
None,
now_unix_secs,
pool_config.score_rules,
@@ -292,19 +230,6 @@ pub(crate) async fn perform_pool_score_rebuild_once_with_config(
.iter()
.map(|(provider, _)| provider.id.clone())
.collect::<Vec<_>>();
let mut endpoints_by_provider = BTreeMap::new();
for endpoint in state
.list_provider_catalog_endpoints_by_provider_ids(&provider_ids)
.await?
.into_iter()
.filter(|endpoint| endpoint.is_active)
{
endpoints_by_provider
.entry(endpoint.provider_id.clone())
.or_insert_with(Vec::new)
.push(endpoint);
}
let mut keys_by_provider = BTreeMap::new();
for key in state
.list_provider_catalog_keys_by_provider_ids(&provider_ids)
@@ -315,6 +240,9 @@ pub(crate) async fn perform_pool_score_rebuild_once_with_config(
.or_insert_with(Vec::new)
.push(key);
}
for keys in keys_by_provider.values_mut() {
keys.sort_by(|left, right| left.id.cmp(&right.id));
}
let now = now_unix_secs();
let mut summary = PoolScoreRebuildRunSummary {
@@ -333,75 +261,39 @@ pub(crate) async fn perform_pool_score_rebuild_once_with_config(
}
last_provider_index = Some(provider_index);
let (provider, pool_config) = providers[provider_index].clone();
let endpoints = endpoints_by_provider
.remove(&provider.id)
.unwrap_or_default();
let keys = keys_by_provider.remove(&provider.id).unwrap_or_default();
if endpoints.is_empty() || keys.is_empty() {
continue;
}
let models = state
.list_admin_provider_models(&AdminProviderModelListQuery {
provider_id: provider.id.clone(),
is_active: Some(true),
offset: 0,
limit: 10_000,
})
.await?
let keys = keys
.into_iter()
.filter(|model| model.is_available)
.filter(|key| key.is_active)
.collect::<Vec<_>>();
if models.is_empty() {
continue;
}
let total_combinations = endpoints
.len()
.saturating_mul(models.len())
.saturating_mul(keys.len());
if total_combinations == 0 {
if keys.is_empty() {
continue;
}
let total_keys = keys.len();
let provider_cursor_key = provider_offset_cursor_key(&provider.id);
let provider_cursor =
load_runtime_usize(state, &provider_cursor_key).await % total_combinations.max(1);
let provider_cursor = load_runtime_usize(state, &provider_cursor_key).await % total_keys;
let remaining_budget = config
.max_upserts_per_tick
.saturating_sub(summary.scores_upserted);
let provider_budget = remaining_budget.min(total_combinations);
let provider_budget = remaining_budget.min(total_keys);
let mut build_items = Vec::with_capacity(provider_budget);
for offset in 0..provider_budget {
let flat_index = (provider_cursor + offset) % total_combinations;
let (endpoint_index, model_index, key_index) =
score_combo_indices(flat_index, models.len(), keys.len());
let endpoint = &endpoints[endpoint_index];
let api_format = endpoint.api_format.trim();
if api_format.is_empty() {
continue;
}
let model = &models[model_index];
let key_index = (provider_cursor + offset) % total_keys;
let key = &keys[key_index];
let draft = build_provider_key_pool_score_upsert(
key,
provider.provider_type.as_str(),
api_format,
Some(model.id.as_str()),
None,
now,
pool_config.score_rules,
);
build_items.push(ProviderScoreBuildItem {
endpoint_index,
model_index,
key_index,
score_id: draft.id,
});
build_items.push((key_index, draft.id));
}
if build_items.is_empty() {
store_runtime_usize(
state,
&provider_cursor_key,
(provider_cursor + provider_budget) % total_combinations,
(provider_cursor + provider_budget) % total_keys,
)
.await;
continue;
@@ -411,7 +303,7 @@ pub(crate) async fn perform_pool_score_rebuild_once_with_config(
.get_pool_member_scores_by_ids(&GetPoolMemberScoresByIdsQuery {
ids: build_items
.iter()
.map(|item| item.score_id.clone())
.map(|(_, score_id)| score_id.clone())
.collect(),
})
.await
@@ -428,19 +320,15 @@ pub(crate) async fn perform_pool_score_rebuild_once_with_config(
.collect::<BTreeMap<_, _>>();
let mut provider_upserts = 0usize;
summary.keys_seen = summary.keys_seen.saturating_add(keys.len());
for item in &build_items {
for (key_index, score_id) in &build_items {
if summary.scores_upserted >= config.max_upserts_per_tick {
break;
}
let endpoint = &endpoints[item.endpoint_index];
let model = &models[item.model_index];
let key = &keys[item.key_index];
let existing = existing_scores.get(&item.score_id);
let key = &keys[*key_index];
let existing = existing_scores.get(score_id);
let upsert = build_provider_key_pool_score_upsert(
key,
provider.provider_type.as_str(),
endpoint.api_format.trim(),
Some(model.id.as_str()),
existing,
now,
pool_config.score_rules,
@@ -459,7 +347,7 @@ pub(crate) async fn perform_pool_score_rebuild_once_with_config(
store_runtime_usize(
state,
&provider_cursor_key,
(provider_cursor + provider_budget) % total_combinations,
(provider_cursor + provider_budget) % total_keys,
)
.await;
if provider_upserts > 0 {
@@ -528,16 +416,3 @@ pub(crate) fn spawn_pool_score_rebuild_worker(
}
}))
}
#[cfg(test)]
mod tests {
use super::score_combo_indices;
#[test]
fn score_combo_indices_walks_endpoint_model_key_order() {
assert_eq!(score_combo_indices(0, 2, 3), (0, 0, 0));
assert_eq!(score_combo_indices(2, 2, 3), (0, 0, 2));
assert_eq!(score_combo_indices(3, 2, 3), (0, 1, 0));
assert_eq!(score_combo_indices(6, 2, 3), (1, 0, 0));
}
}

69
apps/aether-gateway/src/tests/control/admin/pool.rs
View File

@@ -5,8 +5,7 @@ use aether_data::repository::pool_scores::InMemoryPoolMemberScoreRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data::repository::usage::InMemoryUsageReadRepository;
use aether_data_contracts::repository::pool_scores::{
PoolMemberHardState, PoolMemberProbeStatus, StoredPoolMemberScore, POOL_KIND_PROVIDER_KEY_POOL,
POOL_MEMBER_KIND_PROVIDER_API_KEY, POOL_SCORE_SCOPE_KIND_MODEL,
PoolMemberHardState, PoolMemberIdentity, PoolMemberProbeStatus, StoredPoolMemberScore,
};
use aether_data_contracts::repository::provider_catalog::ProviderCatalogReadRepository;
use aether_data_contracts::repository::usage::StoredRequestUsageAudit;
@@ -20,6 +19,7 @@ use super::super::{
build_router_with_state, sample_endpoint, sample_key, sample_provider, start_server, AppState,
};
use crate::admin_api::{maybe_build_local_admin_pool_response, AdminAppState, AdminRequestContext};
use crate::ai_serving::{provider_key_pool_score_id, provider_key_pool_score_scope};
use crate::audit::AdminAuditEvent;
use crate::constants::{
GATEWAY_HEADER, TRUSTED_ADMIN_MANAGEMENT_TOKEN_ID_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER,
@@ -372,16 +372,19 @@ async fn gateway_handles_admin_pool_scores_locally_with_trusted_admin_principal(
Vec::new(),
vec![key.clone()],
));
let score_scope = provider_key_pool_score_scope();
let score_identity = PoolMemberIdentity::provider_api_key("provider-openai", "key-openai-a");
let score_id = provider_key_pool_score_id(&score_identity, &score_scope);
let pool_score_repository = Arc::new(InMemoryPoolMemberScoreRepository::seed(vec![
StoredPoolMemberScore {
id: "pms-provider-openai-key-openai-a-openai-chat-model-1".to_string(),
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
pool_id: "provider-openai".to_string(),
member_kind: POOL_MEMBER_KIND_PROVIDER_API_KEY.to_string(),
member_id: "key-openai-a".to_string(),
capability: "openai:chat".to_string(),
scope_kind: POOL_SCORE_SCOPE_KIND_MODEL.to_string(),
scope_id: Some("model-1".to_string()),
id: score_id,
pool_kind: score_identity.pool_kind.clone(),
pool_id: score_identity.pool_id.clone(),
member_kind: score_identity.member_kind.clone(),
member_id: score_identity.member_id.clone(),
capability: score_scope.capability.clone(),
scope_kind: score_scope.scope_kind.clone(),
scope_id: score_scope.scope_id.clone(),
score: 0.875,
hard_state: PoolMemberHardState::Available,
score_version: 1,
@@ -414,7 +417,7 @@ async fn gateway_handles_admin_pool_scores_locally_with_trusted_admin_principal(
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/pool/provider-openai/scores?api_format=openai:chat&model_id=model-1"
"{gateway_url}/api/admin/pool/provider-openai/scores"
))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
@@ -430,6 +433,9 @@ async fn gateway_handles_admin_pool_scores_locally_with_trusted_admin_principal(
Some(1)
);
assert_eq!(payload["items"][0]["member_id"], json!("key-openai-a"));
assert_eq!(payload["items"][0]["capability"], json!("account"));
assert_eq!(payload["items"][0]["scope_kind"], json!("account"));
assert_eq!(payload["items"][0]["scope_id"], serde_json::Value::Null);
assert_eq!(payload["items"][0]["key"]["name"], json!("score key"));
assert_eq!(payload["items"][0]["probe_status"], json!("ok"));
@@ -734,14 +740,46 @@ async fn gateway_handles_admin_pool_list_keys_locally_with_trusted_admin_princip
Vec::new(),
vec![primary_key, cooldown_key, inactive_key],
));
let score_scope = provider_key_pool_score_scope();
let score_identity = PoolMemberIdentity::provider_api_key("provider-openai", "key-openai-a");
let pool_score_repository = Arc::new(InMemoryPoolMemberScoreRepository::seed(vec![
StoredPoolMemberScore {
id: provider_key_pool_score_id(&score_identity, &score_scope),
pool_kind: score_identity.pool_kind.clone(),
pool_id: score_identity.pool_id.clone(),
member_kind: score_identity.member_kind.clone(),
member_id: score_identity.member_id.clone(),
capability: score_scope.capability.clone(),
scope_kind: score_scope.scope_kind.clone(),
scope_id: score_scope.scope_id.clone(),
score: 0.875,
hard_state: PoolMemberHardState::Available,
score_version: 1,
score_reason: json!({ "weights": { "manual_priority": 0.3 } }),
last_ranked_at: Some(1_700_000_000),
last_scheduled_at: Some(1_700_000_010),
last_success_at: Some(1_700_000_020),
last_failure_at: None,
failure_count: 0,
last_probe_attempt_at: Some(1_700_000_030),
last_probe_success_at: Some(1_700_000_040),
last_probe_failure_at: None,
probe_failure_count: 0,
probe_status: PoolMemberProbeStatus::Ok,
updated_at: 1_700_000_050,
},
]));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
)),
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
)
.with_pool_score_repository_for_tests(pool_score_repository),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
@@ -765,6 +803,9 @@ async fn gateway_handles_admin_pool_list_keys_locally_with_trusted_admin_princip
let keys = payload["keys"].as_array().expect("keys should be array");
assert_eq!(keys.len(), 2);
assert_eq!(keys[0]["key_name"], json!("alpha"));
assert_eq!(keys[0]["pool_score"]["score"], json!(0.875));
assert_eq!(keys[0]["pool_score"]["scope_kind"], json!("account"));
assert_eq!(keys[0]["pool_score"]["scope_id"], serde_json::Value::Null);
assert_eq!(keys[0]["scheduling_status"], json!("available"));
assert_eq!(keys[1]["key_name"], json!("beta"));
assert_eq!(keys[1]["scheduling_reason"], json!("available"));